SET statements: scope modifier for multiple assignments (#1772)

Author: MohamedAbdeen21

src/ast/mod.rs

@@ -2638,7 +2638,7 @@ pub enum Set {
     /// SQL Standard-style
     /// SET a = 1;
     SingleAssignment {
-        scope: ContextModifier,
+        scope: Option<ContextModifier>,
         hivevar: bool,
         variable: ObjectName,
         values: Vec<Expr>,
@@ -2668,7 +2668,7 @@ pub enum Set {
     /// [4]: https://docs.oracle.com/cd/B19306_01/server.102/b14200/statements_10004.htm
     SetRole {
         /// Non-ANSI optional identifier to inform if the role is defined inside the current session (`SESSION`) or transaction (`LOCAL`).
-        context_modifier: ContextModifier,
+        context_modifier: Option<ContextModifier>,
         /// Role name. If NONE is specified, then the current role name is removed.
         role_name: Option<Ident>,
     },
@@ -2720,7 +2720,13 @@ impl Display for Set {
                 role_name,
             } => {
                 let role_name = role_name.clone().unwrap_or_else(|| Ident::new("NONE"));
-                write!(f, "SET {context_modifier}ROLE {role_name}")
+                write!(
+                    f,
+                    "SET {modifier}ROLE {role_name}",
+                    modifier = context_modifier
+                        .map(|m| format!("{}", m))
+                        .unwrap_or_default()
+                )
             }
             Self::SetSessionParam(kind) => write!(f, "SET {kind}"),
             Self::SetTransaction {
@@ -2775,7 +2781,7 @@ impl Display for Set {
                 write!(
                     f,
                     "SET {}{}{} = {}",
-                    scope,
+                    scope.map(|s| format!("{}", s)).unwrap_or_default(),
                     if *hivevar { "HIVEVAR:" } else { "" },
                     variable,
                     display_comma_separated(values)
@@ -5736,13 +5742,20 @@ impl fmt::Display for SequenceOptions {
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
 pub struct SetAssignment {
+    pub scope: Option<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.map(|s| format!("{}", s)).unwrap_or_default(),
+            self.name,
+            self.value
+        )
     }
 }
 
@@ -7969,8 +7982,6 @@ impl fmt::Display for FlushLocation {
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
 pub enum ContextModifier {
-    /// No context defined. Each dialect defines the default in this scenario.
-    None,
     /// `LOCAL` identifier, usually related to transactional states.
     Local,
     /// `SESSION` identifier
@@ -7982,9 +7993,6 @@ pub enum ContextModifier {
 impl fmt::Display for ContextModifier {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
-            Self::None => {
-                write!(f, "")
-            }
             Self::Local => {
                 write!(f, "LOCAL ")
             }

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

@@ -1822,12 +1822,12 @@ impl<'a> Parser<'a> {
         })
     }
 
-    fn keyword_to_modifier(k: Option<Keyword>) -> ContextModifier {
+    fn keyword_to_modifier(k: Keyword) -> Option<ContextModifier> {
         match k {
-            Some(Keyword::LOCAL) => ContextModifier::Local,
-            Some(Keyword::GLOBAL) => ContextModifier::Global,
-            Some(Keyword::SESSION) => ContextModifier::Session,
-            _ => ContextModifier::None,
+            Keyword::LOCAL => Some(ContextModifier::Local),
+            Keyword::GLOBAL => Some(ContextModifier::Global),
+            Keyword::SESSION => Some(ContextModifier::Session),
+            _ => None,
         }
     }
 
@@ -11157,17 +11157,19 @@ 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: Option<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,
         }))
     }
@@ -11203,46 +11205,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) -> Option<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 {
+            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);
         }
 
@@ -11252,8 +11260,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)?,
                 }
@@ -11263,7 +11271,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 == Some(ContextModifier::Local),
                     value: self.parse_expr()?,
                 }
                 .into());
@@ -11311,41 +11319,26 @@ impl<'a> Parser<'a> {
         }
 
         if self.dialect.supports_comma_separated_set_assignments() {
+            if scope.is_some() {
+                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())
                 };
@@ -11370,8 +11363,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

@@ -8635,7 +8635,7 @@ fn parse_set_variable() {
             variable,
             values,
         }) => {
-            assert_eq!(scope, ContextModifier::None);
+            assert_eq!(scope, None);
             assert!(!hivevar);
             assert_eq!(variable, ObjectName::from(vec!["SOMETHING".into()]));
             assert_eq!(
@@ -8655,7 +8655,7 @@ fn parse_set_variable() {
             variable,
             values,
         }) => {
-            assert_eq!(scope, ContextModifier::Global);
+            assert_eq!(scope, Some(ContextModifier::Global));
             assert!(!hivevar);
             assert_eq!(variable, ObjectName::from(vec!["VARIABLE".into()]));
             assert_eq!(
@@ -8747,7 +8747,7 @@ fn parse_set_role_as_variable() {
             variable,
             values,
         }) => {
-            assert_eq!(scope, ContextModifier::None);
+            assert_eq!(scope, None);
             assert!(!hivevar);
             assert_eq!(variable, ObjectName::from(vec!["role".into()]));
             assert_eq!(
@@ -8794,7 +8794,7 @@ fn parse_set_time_zone() {
             variable,
             values,
         }) => {
-            assert_eq!(scope, ContextModifier::None);
+            assert_eq!(scope, None);
             assert!(!hivevar);
             assert_eq!(variable, ObjectName::from(vec!["TIMEZONE".into()]));
             assert_eq!(
@@ -14859,10 +14859,12 @@ fn parse_multiple_set_statements() -> Result<(), ParserError> {
                 assignments,
                 vec![
                     SetAssignment {
+                        scope: None,
                         name: ObjectName::from(vec!["@a".into()]),
                         value: Expr::value(number("1"))
                     },
                     SetAssignment {
+                        scope: None,
                         name: ObjectName::from(vec!["b".into()]),
                         value: Expr::value(number("2"))
                     }
@@ -14872,6 +14874,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: Some(ContextModifier::Global),
+                        name: ObjectName::from(vec!["@a".into()]),
+                        value: Expr::value(number("1"))
+                    },
+                    SetAssignment {
+                        scope: Some(ContextModifier::Session),
+                        name: ObjectName::from(vec!["b".into()]),
+                        value: Expr::value(number("2"))
+                    },
+                    SetAssignment {
+                        scope: Some(ContextModifier::Local),
+                        name: ObjectName::from(vec!["c".into()]),
+                        value: Expr::value(number("3"))
+                    },
+                    SetAssignment {
+                        scope: None,
+                        name: ObjectName::from(vec!["d".into()]),
+                        value: Expr::value(number("4"))
+                    }
+                ]
+            );
+        }
+        _ => panic!("Expected MultipleAssignments with 4 scoped variables and 4 values"),
+    };
+
     Ok(())
 }