DuckDB: Support SET VARIABLE

src/ast/mod.rs

@@ -3279,6 +3279,15 @@ impl Display for FromTable {
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
 /// Variants for the `SET` family of statements.
 pub enum Set {
+    /// ```sql
+    /// SET VARIABLE variable_name = expression
+    /// ```
+    SetVariable {
+        /// Variable name to assign.
+        variable: ObjectName,
+        /// Value assigned to the variable.
+        value: Expr,
+    },
     /// SQL Standard-style
     /// SET a = 1;
     /// `SET var = value` (standard SQL-style assignment).
@@ -3383,6 +3392,9 @@ pub enum Set {
 impl Display for Set {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
+            Self::SetVariable { variable, value } => {
+                write!(f, "SET VARIABLE {variable} = {value}")
+            }
             Self::ParenthesizedAssignments { variables, values } => write!(
                 f,
                 "SET ({}) = ({})",

src/dialect/duckdb.rs

@@ -80,6 +80,10 @@ impl Dialect for DuckDbDialect {
         true
     }
 
+    fn supports_set_variable_statement(&self) -> bool {
+        true
+    }
+
     /// Returns true if this dialect allows the `EXTRACT` function to use single quotes in the part being extracted.
     fn allow_extract_single_quotes(&self) -> bool {
         true

src/dialect/mod.rs

@@ -549,6 +549,11 @@ pub trait Dialect: Debug + Any {
         false
     }
 
+    /// Returns true if the dialect supports `SET VARIABLE name = expression`.
+    fn supports_set_variable_statement(&self) -> bool {
+        false
+    }
+
     /// Returns true if the dialect supports multiple `SET` statements
     /// in a single statement.
     ///

src/parser/mod.rs

@@ -15730,6 +15730,37 @@ impl<'a> Parser<'a> {
     }
 
     fn parse_set(&mut self) -> Result<Statement, ParserError> {
+        let variable_is_ordinary_assignment = self.peek_keyword(Keyword::VARIABLE)
+            && (matches!(self.peek_nth_token_ref(1).token, Token::Eq | Token::Period)
+                || matches!(
+                    &self.peek_nth_token_ref(1).token,
+                    Token::Word(word) if word.keyword == Keyword::TO
+                ));
+        if self.dialect.supports_set_variable_statement()
+            && !variable_is_ordinary_assignment
+            && self.parse_keyword(Keyword::VARIABLE)
+        {
+            let mut parts = vec![];
+            loop {
+                let token = self.next_token();
+                match token.token {
+                    Token::Word(word) if matches!(word.quote_style, None | Some('"')) => {
+                        parts.push(word.into_ident(token.span));
+                    }
+                    _ => return self.expected("identifier", token),
+                }
+                if !self.consume_token(&Token::Period) {
+                    break;
+                }
+            }
+            let variable = ObjectName::from(parts);
+            if !(self.consume_token(&Token::Eq) || self.parse_keyword(Keyword::TO)) {
+                return self.expected_ref("equals sign or TO", self.peek_token_ref());
+            }
+            let value = self.parse_expr()?;
+            return Ok(Set::SetVariable { variable, value }.into());
+        }
+
         let hivevar = self.parse_keyword(Keyword::HIVEVAR);
 
         // Modifier is either HIVEVAR: or a ContextModifier (LOCAL, SESSION, etc), not both