Add support for sql variable inside query in snowflake and mysql dialect

see :
https://docs.snowflake.com/en/sql-reference/session-variables.html
https://dev.mysql.com/doc/refman/8.0/en/user-variables.html

Author: eyalsatori

src/ast/mod.rs

@@ -235,6 +235,10 @@ pub enum Expr {
     Subquery(Box<Query>),
     /// The `LISTAGG` function `SELECT LISTAGG(...) WITHIN GROUP (ORDER BY ...)`
     ListAgg(ListAgg),
+    /// Embed variable inside a query is supported by some databases:
+    /// - Mysql: https://dev.mysql.com/doc/refman/8.0/en/user-variables.html
+    /// - Snowflake: https://docs.snowflake.com/en/sql-reference/session-variables.html
+    SqlVariable { prefix: char, name: Ident },
 }
 
 impl fmt::Display for Expr {
@@ -315,6 +319,7 @@ impl fmt::Display for Expr {
             Expr::Exists(s) => write!(f, "EXISTS ({})", s),
             Expr::Subquery(s) => write!(f, "({})", s),
             Expr::ListAgg(listagg) => write!(f, "{}", listagg),
+            Expr::SqlVariable { prefix, name } => write!(f, "{}{}", prefix, name),
         }
     }
 }

src/parser.rs

@@ -312,6 +312,16 @@ impl<'a> Parser<'a> {
                 self.expect_token(&Token::RParen)?;
                 Ok(expr)
             }
+            Token::Dollar if dialect_of!(self is SnowflakeDialect) => {
+                // Snowflake user defined variables starts with $
+                let name = self.parse_identifier()?;
+                Ok(Expr::SqlVariable { prefix: '$', name })
+            }
+            Token::At if dialect_of!(self is MySqlDialect) => {
+                // Mysql user defined variables starts with @
+                let name = self.parse_identifier()?;
+                Ok(Expr::SqlVariable { prefix: '@', name })
+            }
             unexpected => self.expected("an expression", unexpected),
         }?;
 

src/tokenizer.rs

@@ -101,6 +101,10 @@ pub enum Token {
     RBrace,
     /// Right Arrow `=>`
     RArrow,
+    /// Dollar sign `$`
+    Dollar,
+    /// At sign `@`
+    At,
 }
 
 impl fmt::Display for Token {
@@ -142,6 +146,8 @@ impl fmt::Display for Token {
             Token::LBrace => f.write_str("{"),
             Token::RBrace => f.write_str("}"),
             Token::RArrow => f.write_str("=>"),
+            Token::Dollar => f.write_str("$"),
+            Token::At => f.write_str("@"),
         }
     }
 }
@@ -448,6 +454,8 @@ impl<'a> Tokenizer<'a> {
                 '^' => self.consume_and_return(chars, Token::Caret),
                 '{' => self.consume_and_return(chars, Token::LBrace),
                 '}' => self.consume_and_return(chars, Token::RBrace),
+                '$' => self.consume_and_return(chars, Token::Dollar),
+                '@' => self.consume_and_return(chars, Token::At),
                 other => self.consume_and_return(chars, Token::Char(other)),
             },
             None => Ok(None),

tests/sqlparser_mysql.rs

@@ -152,6 +152,39 @@ fn parse_quote_identifiers() {
     }
 }
 
+#[test]
+fn test_query_with_variable_name() {
+    let sql = "SELECT @var1";
+    let select = mysql().verified_only_select(sql);
+
+    assert_eq!(
+        only(select.projection),
+        SelectItem::UnnamedExpr(Expr::SqlVariable {
+            prefix: '@',
+            name: Ident::new("var1")
+        },)
+    );
+
+    let sql = "SELECT c1 FROM t1 WHERE num BETWEEN @min AND @max";
+    let select = mysql().verified_only_select(sql);
+
+    assert_eq!(
+        select.selection.unwrap(),
+        Expr::Between {
+            expr: Box::new(Expr::Identifier("num".into())),
+            low: Box::new(Expr::SqlVariable {
+                prefix: '@',
+                name: Ident::new("min")
+            }),
+            high: Box::new(Expr::SqlVariable {
+                prefix: '@',
+                name: Ident::new("max")
+            }),
+            negated: false,
+        }
+    );
+}
+
 fn mysql() -> TestedDialects {
     TestedDialects {
         dialects: vec![Box::new(MySqlDialect {})],

tests/sqlparser_snowflake.rs

@@ -24,6 +24,45 @@ fn test_snowflake_create_table() {
     }
 }
 
+#[test]
+fn test_query_with_variable_name() {
+    let sql = "SELECT $var1";
+    let select = snowflake().verified_only_select(sql);
+
+    assert_eq!(
+        only(select.projection),
+        SelectItem::UnnamedExpr(Expr::SqlVariable {
+            prefix: '$',
+            name: Ident::new("var1")
+        },)
+    );
+
+    let sql = "SELECT c1 FROM t1 WHERE num BETWEEN $min AND $max";
+    let select = snowflake().verified_only_select(sql);
+
+    assert_eq!(
+        select.selection.unwrap(),
+        Expr::Between {
+            expr: Box::new(Expr::Identifier("num".into())),
+            low: Box::new(Expr::SqlVariable {
+                prefix: '$',
+                name: Ident::new("min")
+            }),
+            high: Box::new(Expr::SqlVariable {
+                prefix: '$',
+                name: Ident::new("max")
+            }),
+            negated: false,
+        }
+    );
+}
+
+fn snowflake() -> TestedDialects {
+    TestedDialects {
+        dialects: vec![Box::new(SnowflakeDialect {})],
+    }
+}
+
 fn snowflake_and_generic() -> TestedDialects {
     TestedDialects {
         dialects: vec![Box::new(SnowflakeDialect {}), Box::new(GenericDialect {})],