feat: Support expression in SET statement (#7)

Author: vasilev-alex

src/ast/mod.rs

@@ -2159,32 +2159,15 @@ impl fmt::Display for ShowStatementFilter {
     }
 }
 
-#[derive(Debug, Clone, PartialEq, Eq, Hash)]
-#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
-pub enum SetVariableValue {
-    Ident(Ident),
-    Literal(Value),
-}
-
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 pub struct SetVariableKeyValue {
     pub key: Ident,
-    pub value: Vec<SetVariableValue>,
+    pub value: Vec<Expr>,
     pub local: bool,
     pub hivevar: bool,
 }
 
-impl fmt::Display for SetVariableValue {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        use SetVariableValue::*;
-        match self {
-            Ident(ident) => write!(f, "{}", ident),
-            Literal(literal) => write!(f, "{}", literal),
-        }
-    }
-}
-
 /// Sqlite specific syntax
 ///
 /// https://sqlite.org/lang_conflict.html

src/dialect/mysql.rs

@@ -24,6 +24,7 @@ impl Dialect for MySqlDialect {
             || ('A'..='Z').contains(&ch)
             || ch == '_'
             || ch == '$'
+            || ch == '@'
             || ('\u{0080}'..='\u{ffff}').contains(&ch)
     }
 

src/parser.rs

@@ -3052,8 +3052,8 @@ impl<'a> Parser<'a> {
             let mut values = vec![];
 
             loop {
-                let value = if let Ok(value) = self.parse_value() {
-                    SetVariableValue::Literal(value)
+                let value = if let Ok(expr) = self.parse_expr() {
+                    expr
                 } else {
                     self.expected("variable value", self.peek_token())?
                 };
@@ -3082,8 +3082,8 @@ impl<'a> Parser<'a> {
             let mut values = vec![];
 
             if self.consume_token(&Token::Eq) || self.parse_keyword(Keyword::TO) {
-                let value = if let Ok(value) = self.parse_value() {
-                    SetVariableValue::Literal(value)
+                let value = if let Ok(expr) = self.parse_expr() {
+                    expr
                 } else {
                     self.expected("variable value", self.peek_token())?
                 };

tests/sqlparser_mysql.rs

@@ -131,41 +131,37 @@ fn parse_set_transaction() {
 
 #[test]
 fn parse_set_variables() {
-    let stmt = mysql_and_generic().verified_stmt("SET autocommit = 1, sql_mode = 'test'");
-
     assert_eq!(
-        stmt,
+        mysql_and_generic().verified_stmt("SET autocommit = 1, sql_mode = 'test'"),
         Statement::SetVariable {
             key_values: [
                 SetVariableKeyValue {
                     local: false,
                     hivevar: false,
                     key: "autocommit".into(),
-                    value: vec![SetVariableValue::Literal(number("1"))],
+                    value: vec![Expr::Value(Value::Number("1".into(), false))],
                 },
                 SetVariableKeyValue {
                     local: false,
                     hivevar: false,
                     key: "sql_mode".into(),
-                    value: vec![SetVariableValue::Literal(Value::SingleQuotedString(
-                        "test".into()
-                    ))],
+                    value: vec![Expr::Value(Value::SingleQuotedString("test".into()))],
                 }
             ]
             .to_vec()
         }
     );
 
-    let stmt = mysql_and_generic().verified_stmt("SET LOCAL autocommit = 1");
+    mysql_and_generic().verified_stmt("SET sql_mode = CONCAT(@@sql_mode, ',STRICT_TRANS_TABLES')");
 
     assert_eq!(
-        stmt,
+        mysql_and_generic().verified_stmt("SET LOCAL autocommit = 1"),
         Statement::SetVariable {
             key_values: [SetVariableKeyValue {
                 local: true,
                 hivevar: false,
                 key: "autocommit".into(),
-                value: vec![SetVariableValue::Literal(number("1"))],
+                value: vec![Expr::Value(Value::Number("1".into(), false))],
             },]
             .to_vec()
         }

tests/sqlparser_postgres.rs

@@ -466,7 +466,10 @@ fn parse_set() {
         Statement::SetVariable {
             key_values: [SetVariableKeyValue {
                 key: "a".into(),
-                value: vec![SetVariableValue::Ident("b".into())],
+                value: vec![Expr::Identifier(Ident {
+                    value: "b".into(),
+                    quote_style: None
+                })],
                 local: false,
                 hivevar: false,
             }]
@@ -482,9 +485,7 @@ fn parse_set() {
                 local: false,
                 hivevar: false,
                 key: "a".into(),
-                value: vec![SetVariableValue::Literal(Value::SingleQuotedString(
-                    "b".into()
-                ))],
+                value: vec![Expr::Value(Value::SingleQuotedString("b".into()))],
             }]
             .to_vec()
         }
@@ -498,7 +499,7 @@ fn parse_set() {
                 local: false,
                 hivevar: false,
                 key: "a".into(),
-                value: vec![SetVariableValue::Literal(number("0"))],
+                value: vec![Expr::Value(Value::Number("0".into(), false))],
             }]
             .to_vec()
         }
@@ -512,7 +513,10 @@ fn parse_set() {
                 local: false,
                 hivevar: false,
                 key: "a".into(),
-                value: vec![SetVariableValue::Ident("DEFAULT".into())],
+                value: vec![Expr::Identifier(Ident {
+                    value: "DEFAULT".into(),
+                    quote_style: None
+                })],
             }]
             .to_vec()
         }
@@ -526,7 +530,10 @@ fn parse_set() {
                 local: true,
                 hivevar: false,
                 key: "a".into(),
-                value: vec![SetVariableValue::Ident("b".into())],
+                value: vec![Expr::Identifier(Ident {
+                    value: "b".into(),
+                    quote_style: None
+                })],
             }]
             .to_vec()
         }