SET statements: scope modifier for multiple assignments

Author: MohamedAbdeen21

src/ast/mod.rs

@@ -5727,13 +5727,14 @@ impl fmt::Display for SequenceOptions {
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
 pub struct SetAssignment {
+    pub scope: ContextModifier,
     pub name: ObjectName,
     pub value: Expr,
 }
 
 impl fmt::Display for SetAssignment {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "{} = {}", self.name, self.value)
+        write!(f, "{}{} = {}", self.scope, self.name, self.value)
     }
 }
 

src/dialect/generic.rs

@@ -159,4 +159,8 @@ impl Dialect for GenericDialect {
     fn supports_set_names(&self) -> bool {
         true
     }
+
+    fn supports_comma_separated_set_assignments(&self) -> bool {
+        true
+    }
 }

src/parser/mod.rs

@@ -11145,17 +11145,16 @@ impl<'a> Parser<'a> {
     }
 
     /// Parse a `SET ROLE` statement. Expects SET to be consumed already.
-    fn parse_set_role(&mut self, modifier: Option<Keyword>) -> Result<Statement, ParserError> {
+    fn parse_set_role(&mut self, modifier: ContextModifier) -> Result<Statement, ParserError> {
         self.expect_keyword_is(Keyword::ROLE)?;
-        let context_modifier = Self::keyword_to_modifier(modifier);
 
         let role_name = if self.parse_keyword(Keyword::NONE) {
             None
         } else {
             Some(self.parse_identifier()?)
         };
         Ok(Statement::Set(Set::SetRole {
-            context_modifier,
+            context_modifier: modifier,
             role_name,
         }))
     }
@@ -11191,46 +11190,52 @@ impl<'a> Parser<'a> {
         }
     }
 
-    fn parse_set_assignment(
-        &mut self,
-    ) -> Result<(OneOrManyWithParens<ObjectName>, Expr), ParserError> {
-        let variables = if self.dialect.supports_parenthesized_set_variables()
+    fn parse_context_modifier(&mut self) -> ContextModifier {
+        let modifier =
+            self.parse_one_of_keywords(&[Keyword::SESSION, Keyword::LOCAL, Keyword::GLOBAL]);
+
+        Self::keyword_to_modifier(modifier)
+    }
+
+    /// Parse a single SET statement assignment `var = expr`.
+    fn parse_set_assignment(&mut self) -> Result<SetAssignment, ParserError> {
+        let scope = self.parse_context_modifier();
+
+        let name = 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
+            // Parenthesized assignments are handled in the `parse_set` function after
+            // trying to parse list of assignments using this function.
+            // If a dialect supports both, and we find a LParen, we early exit from this function.
+            self.expected("Unparenthesized assignment", self.peek_token())?
         } else {
-            OneOrManyWithParens::One(self.parse_object_name(false)?)
+            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()?;
+        let value = self.parse_expr()?;
 
-        Ok((variables, values))
+        Ok(SetAssignment { scope, name, value })
     }
 
     fn parse_set(&mut self) -> Result<Statement, ParserError> {
-        let modifier = self.parse_one_of_keywords(&[
-            Keyword::SESSION,
-            Keyword::LOCAL,
-            Keyword::HIVEVAR,
-            Keyword::GLOBAL,
-        ]);
-
-        if let Some(Keyword::HIVEVAR) = modifier {
+        let hivevar = self.parse_keyword(Keyword::HIVEVAR);
+
+        // Modifier is either HIVEVAR: or a ContextModifier (LOCAL, SESSION, etc), not both
+        let scope = if !hivevar {
+            self.parse_context_modifier()
+        } else {
+            ContextModifier::None
+        };
+
+        if hivevar {
             self.expect_token(&Token::Colon)?;
         }
 
-        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(scope))? {
             return Ok(set_role_stmt);
         }
 
@@ -11240,8 +11245,8 @@ impl<'a> Parser<'a> {
         {
             if self.consume_token(&Token::Eq) || self.parse_keyword(Keyword::TO) {
                 return Ok(Set::SingleAssignment {
-                    scope: Self::keyword_to_modifier(modifier),
-                    hivevar: modifier == Some(Keyword::HIVEVAR),
+                    scope,
+                    hivevar,
                     variable: ObjectName::from(vec!["TIMEZONE".into()]),
                     values: self.parse_set_values(false)?,
                 }
@@ -11251,7 +11256,7 @@ impl<'a> Parser<'a> {
                 // the assignment operator. It's originally PostgreSQL specific,
                 // but we allow it for all the dialects
                 return Ok(Set::SetTimeZone {
-                    local: modifier == Some(Keyword::LOCAL),
+                    local: scope == ContextModifier::Local,
                     value: self.parse_expr()?,
                 }
                 .into());
@@ -11299,41 +11304,26 @@ impl<'a> Parser<'a> {
         }
 
         if self.dialect.supports_comma_separated_set_assignments() {
+            if scope != ContextModifier::None {
+                self.prev_token();
+            }
+
             if let Some(assignments) = self
                 .maybe_parse(|parser| parser.parse_comma_separated(Parser::parse_set_assignment))?
             {
                 return if assignments.len() > 1 {
-                    let assignments = assignments
-                        .into_iter()
-                        .map(|(var, val)| match var {
-                            OneOrManyWithParens::One(v) => Ok(SetAssignment {
-                                name: v,
-                                value: val,
-                            }),
-                            OneOrManyWithParens::Many(_) => {
-                                self.expected("List of single identifiers", self.peek_token())
-                            }
-                        })
-                        .collect::<Result<_, _>>()?;
-
                     Ok(Set::MultipleAssignments { assignments }.into())
                 } else {
-                    let (vars, values): (Vec<_>, Vec<_>) = assignments.into_iter().unzip();
-
-                    let variable = match vars.into_iter().next() {
-                        Some(OneOrManyWithParens::One(v)) => Ok(v),
-                        Some(OneOrManyWithParens::Many(_)) => self.expected(
-                            "Single assignment or list of assignments",
-                            self.peek_token(),
-                        ),
-                        None => self.expected("At least one identifier", self.peek_token()),
-                    }?;
+                    let SetAssignment { scope, name, value } =
+                        assignments.into_iter().next().ok_or_else(|| {
+                            ParserError::ParserError("Expected at least one assignment".to_string())
+                        })?;
 
                     Ok(Set::SingleAssignment {
-                        scope: Self::keyword_to_modifier(modifier),
-                        hivevar: modifier == Some(Keyword::HIVEVAR),
-                        variable,
-                        values,
+                        scope,
+                        hivevar,
+                        variable: name,
+                        values: vec![value],
                     }
                     .into())
                 };
@@ -11358,8 +11348,8 @@ impl<'a> Parser<'a> {
         if self.consume_token(&Token::Eq) || self.parse_keyword(Keyword::TO) {
             let stmt = match variables {
                 OneOrManyWithParens::One(var) => Set::SingleAssignment {
-                    scope: Self::keyword_to_modifier(modifier),
-                    hivevar: modifier == Some(Keyword::HIVEVAR),
+                    scope,
+                    hivevar,
                     variable: var,
                     values: self.parse_set_values(false)?,
                 },

tests/sqlparser_common.rs

@@ -14856,10 +14856,12 @@ fn parse_multiple_set_statements() -> Result<(), ParserError> {
                 assignments,
                 vec![
                     SetAssignment {
+                        scope: ContextModifier::None,
                         name: ObjectName::from(vec!["@a".into()]),
                         value: Expr::value(number("1"))
                     },
                     SetAssignment {
+                        scope: ContextModifier::None,
                         name: ObjectName::from(vec!["b".into()]),
                         value: Expr::value(number("2"))
                     }
@@ -14869,6 +14871,39 @@ fn parse_multiple_set_statements() -> Result<(), ParserError> {
         _ => panic!("Expected SetVariable with 2 variables and 2 values"),
     };
 
+    let stmt = dialects.verified_stmt("SET GLOBAL @a = 1, SESSION b = 2, LOCAL c = 3, d = 4");
+
+    match stmt {
+        Statement::Set(Set::MultipleAssignments { assignments }) => {
+            assert_eq!(
+                assignments,
+                vec![
+                    SetAssignment {
+                        scope: ContextModifier::Global,
+                        name: ObjectName::from(vec!["@a".into()]),
+                        value: Expr::value(number("1"))
+                    },
+                    SetAssignment {
+                        scope: ContextModifier::Session,
+                        name: ObjectName::from(vec!["b".into()]),
+                        value: Expr::value(number("2"))
+                    },
+                    SetAssignment {
+                        scope: ContextModifier::Local,
+                        name: ObjectName::from(vec!["c".into()]),
+                        value: Expr::value(number("3"))
+                    },
+                    SetAssignment {
+                        scope: ContextModifier::None,
+                        name: ObjectName::from(vec!["d".into()]),
+                        value: Expr::value(number("4"))
+                    }
+                ]
+            );
+        }
+        _ => panic!("Expected SetVariable with 2 variables and 2 values"),
+    };
+
     Ok(())
 }