setting multiple variables at once (MySQL)

Author: MohamedAbdeen21

src/dialect/hive.rs

@@ -44,10 +44,6 @@ impl Dialect for HiveDialect {
         true
     }
 
-    fn supports_set_multiple_values(&self) -> bool {
-        true
-    }
-
     fn supports_numeric_prefix(&self) -> bool {
         true
     }

src/dialect/mod.rs

@@ -352,12 +352,6 @@ pub trait Dialect: Debug + Any {
         false
     }
 
-    /// Returns true if the dialect supports multiple values in a SET expression
-    /// e.g. `SET OFFSETS SELECT, FROM, ORDER, TABLE, PROCEDURE, EXECUTE ON`
-    fn supports_set_multiple_values(&self) -> bool {
-        false
-    }
-
     /// Returns true if the dialects supports specifying null treatment
     /// as part of a window function's parameter list as opposed
     /// to after the parameter list.

src/dialect/mssql.rs

@@ -58,10 +58,6 @@ impl Dialect for MsSqlDialect {
         true
     }
 
-    fn supports_set_multiple_values(&self) -> bool {
-        true
-    }
-
     fn supports_try_convert(&self) -> bool {
         true
     }

src/keywords.rs

@@ -173,7 +173,6 @@ define_keywords!(
     CHANNEL,
     CHAR,
     CHARACTER,
-    CHARACTERISTIC,
     CHARACTERS,
     CHARACTER_LENGTH,
     CHARSET,
@@ -558,7 +557,6 @@ define_keywords!(
     MULTISET,
     MUTATION,
     NAME,
-    NAMES,
     NANOSECOND,
     NANOSECONDS,
     NATIONAL,

src/parser/mod.rs

@@ -10992,17 +10992,66 @@ impl<'a> Parser<'a> {
         }
     }
 
-    pub fn parse_set(&mut self) -> Result<Statement, ParserError> {
+    fn parse_set_assignment(
+        &mut self,
+    ) -> Result<(OneOrManyWithParens<ObjectName>, Expr), ParserError> {
+        let variables = if self.dialect.supports_parenthesized_set_variables()
+            && self.consume_token(&Token::LParen)
+        {
+            let vars = OneOrManyWithParens::Many(
+                self.parse_comma_separated(|parser: &mut Parser<'a>| parser.parse_identifier())?
+                    .into_iter()
+                    .map(|ident| ObjectName::from(vec![ident]))
+                    .collect(),
+            );
+            self.expect_token(&Token::RParen)?;
+            vars
+        } else {
+            OneOrManyWithParens::One(self.parse_object_name(false)?)
+        };
+
+        if !(self.consume_token(&Token::Eq) || self.parse_keyword(Keyword::TO)) {
+            return self.expected("assignment operator", self.peek_token());
+        }
+
+        let values = self.parse_expr()?;
+
+        Ok((variables, values))
+    }
+
+    fn parse_set(&mut self) -> Result<Statement, ParserError> {
         let modifier =
             self.parse_one_of_keywords(&[Keyword::SESSION, Keyword::LOCAL, Keyword::HIVEVAR]);
+
         if let Some(Keyword::HIVEVAR) = modifier {
             self.expect_token(&Token::Colon)?;
-        } else if let Some(set_role_stmt) =
-            self.maybe_parse(|parser| parser.parse_set_role(modifier))?
-        {
+        }
+
+        if let Some(set_role_stmt) = self.maybe_parse(|parser| parser.parse_set_role(modifier))? {
             return Ok(set_role_stmt);
         }
 
+        if self.dialect.supports_comma_separated_set_assignments() {
+            if let Ok(v) = self
+                .try_parse(|parser| Ok(parser.parse_comma_separated(Parser::parse_set_assignment)?))
+            {
+                let (variables, values): (Vec<_>, Vec<_>) = v.into_iter().unzip();
+
+                let variables = if variables.len() == 1 {
+                    variables.into_iter().next().unwrap()
+                } else {
+                    OneOrManyWithParens::Many(variables.into_iter().flatten().map(|v| v).collect())
+                };
+
+                return Ok(Statement::SetVariable {
+                    local: modifier == Some(Keyword::LOCAL),
+                    hivevar: modifier == Some(Keyword::HIVEVAR),
+                    variables,
+                    value: values,
+                });
+            }
+        }
+
         let variables = if self.parse_keywords(&[Keyword::TIME, Keyword::ZONE]) {
             OneOrManyWithParens::One(ObjectName::from(vec!["TIMEZONE".into()]))
         } else if self.dialect.supports_parenthesized_set_variables()
@@ -11022,7 +11071,7 @@ impl<'a> Parser<'a> {
 
         if self.consume_token(&Token::Eq) || self.parse_keyword(Keyword::TO) {
             let parenthesized_assignment = matches!(&variables, OneOrManyWithParens::Many(_));
-            let values = self.parse_set_values(parenthesized_assignment);
+            let values = self.parse_set_values(parenthesized_assignment)?;
 
             return Ok(Statement::SetVariable {
                 local: modifier == Some(Keyword::LOCAL),

t.sql

@@ -8618,10 +8618,10 @@ fn parse_set_variable() {
             "SET (a) = ((SELECT 22 FROM tbl1, (SELECT 1 FROM tbl2)))",
             "SET (a) = ((SELECT 22 FROM tbl1, (SELECT 1 FROM tbl2)))",
         ),
-        // (
-        //     "SET (a, b) = ((SELECT 22 FROM tbl1, (SELECT 1 FROM tbl2)), SELECT 33 FROM tbl3)",
-        //     "SET (a, b) = ((SELECT 22 FROM tbl1, (SELECT 1 FROM tbl2)), (SELECT 33 FROM tbl3))",
-        // ),
+        (
+            "SET (a, b) = ((SELECT 22 FROM tbl1, (SELECT 1 FROM tbl2)), SELECT 33 FROM tbl3)",
+            "SET (a, b) = ((SELECT 22 FROM tbl1, (SELECT 1 FROM tbl2)), (SELECT 33 FROM tbl3))",
+        ),
     ] {
         multi_variable_dialects.one_statement_parses_to(sql, canonical);
     }
@@ -14653,8 +14653,8 @@ fn parse_multiple_set_statements() -> Result<(), ParserError> {
         Statement::SetVariable {
             variables, value, ..
         } => {
-            assert_eq!(variables.len(), 2);
             assert_eq!(value.len(), 2);
+            assert_eq!(variables.len(), 2);
         }
         _ => assert!(false, "Expected SetVariable with 2 variables and 2 values"),
     };