From 5364da51a40a8695da7d62faf084a9884cb3ebb2 Mon Sep 17 00:00:00 2001 From: Cora Sutton Date: Sat, 12 Sep 2026 19:23:57 +0000 Subject: [PATCH] DuckDB: Support SET VARIABLE --- src/ast/mod.rs | 12 ++ src/dialect/duckdb.rs | 4 + src/dialect/mod.rs | 5 + src/parser/mod.rs | 31 ++++ tests/sqlparser_duckdb.rs | 317 +++++++++++++++++++++++++++++++++++++- 5 files changed, 366 insertions(+), 3 deletions(-) diff --git a/src/ast/mod.rs b/src/ast/mod.rs index 20058b83ab..01988f8732 100644 --- a/src/ast/mod.rs +++ b/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 ({}) = ({})", diff --git a/src/dialect/duckdb.rs b/src/dialect/duckdb.rs index 2e3673bc4c..e985ec1edf 100644 --- a/src/dialect/duckdb.rs +++ b/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 diff --git a/src/dialect/mod.rs b/src/dialect/mod.rs index ff83a4da61..81430e1304 100644 --- a/src/dialect/mod.rs +++ b/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. /// diff --git a/src/parser/mod.rs b/src/parser/mod.rs index 5edc437145..6ae39f4e1d 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -15730,6 +15730,37 @@ impl<'a> Parser<'a> { } fn parse_set(&mut self) -> Result { + 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 diff --git a/tests/sqlparser_duckdb.rs b/tests/sqlparser_duckdb.rs index a338ef7a82..70b879c1c1 100644 --- a/tests/sqlparser_duckdb.rs +++ b/tests/sqlparser_duckdb.rs @@ -19,12 +19,12 @@ mod test_utils; use helpers::attached_token::AttachedToken; -use sqlparser::tokenizer::Span; +use sqlparser::tokenizer::{Location, Span}; use test_utils::*; use sqlparser::ast::*; -use sqlparser::dialect::{DuckDbDialect, GenericDialect}; -use sqlparser::parser::ParserError; +use sqlparser::dialect::{Dialect, DuckDbDialect, GenericDialect}; +use sqlparser::parser::{Parser, ParserError}; fn duckdb() -> TestedDialects { TestedDialects::new(vec![Box::new(DuckDbDialect {})]) @@ -910,3 +910,314 @@ fn test_duckdb_lambda_function() { let sql_transform = "SELECT list_transform([1, 2, 3], lambda x : x * 2)"; duckdb().verified_stmt(sql_transform); } + +#[test] +fn parse_duckdb_set_variable() { + assert!(DuckDbDialect {}.supports_set_variable_statement()); + for (name, parts) in [ + ("threshold", vec![Ident::new("threshold")]), + ( + "config.threshold", + vec!["config".into(), "threshold".into()], + ), + ( + r#""config"."threshold""#, + vec![ + Ident::with_quote('"', "config"), + Ident::with_quote('"', "threshold"), + ], + ), + ( + r#"config."threshold".value"#, + vec![ + "config".into(), + Ident::with_quote('"', "threshold"), + "value".into(), + ], + ), + ( + r#""my variable""#, + vec![Ident::with_quote('"', "my variable")], + ), + ("VARIABLE", vec!["VARIABLE".into()]), + (r#""TO""#, vec![Ident::with_quote('"', "TO")]), + ] { + let expected = Statement::Set(Set::SetVariable { + variable: ObjectName::from(parts), + value: Expr::value(number("42")), + }); + for separator in ["=", "TO"] { + assert_eq!( + duckdb().one_statement_parses_to( + &format!("SET VARIABLE {name} {separator} 42"), + &format!("SET VARIABLE {name} = 42"), + ), + expected + ); + } + } + let statements = + Parser::parse_sql(&DuckDbDialect {}, r#"SET VARIABLE config."threshold" = 42"#).unwrap(); + let Statement::Set(Set::SetVariable { variable, .. }) = &statements[0] else { + panic!("Expected SET VARIABLE"); + }; + assert_eq!( + variable + .0 + .iter() + .map(|part| part.as_ident().unwrap().span) + .collect::>(), + vec![ + Span::new(Location::new(1, 14), Location::new(1, 20)), + Span::new(Location::new(1, 21), Location::new(1, 32)), + ] + ); +} + +#[test] +fn parse_duckdb_set_variable_values() { + for value in [ + "1 + 2 * 3", + "'hello'", + "true", + "NULL", + "DEFAULT", + "DATE '2024-06-10'", + "TIMESTAMP '2024-06-10 12:34:56'", + "[1, 2, 3]", + "MAP {'a': 1, 'b': 2}", + "{'answer': 42}", + "(SELECT max(n) FROM (VALUES (1), (2)) AS t (n))", + ] { + let expected = Statement::Set(Set::SetVariable { + variable: ObjectName::from(vec!["threshold".into()]), + value: duckdb().verified_expr(value), + }); + for separator in ["=", "TO"] { + assert_eq!( + duckdb().one_statement_parses_to( + &format!("SET VARIABLE threshold {separator} {value}"), + &format!("SET VARIABLE threshold = {value}"), + ), + expected + ); + } + } + assert_eq!( + duckdb().verified_stmt("SET VARIABLE threshold = 1 + 2 * 3"), + Statement::Set(Set::SetVariable { + variable: ObjectName::from(vec!["threshold".into()]), + value: Expr::BinaryOp { + left: Box::new(Expr::value(number("1"))), + op: BinaryOperator::Plus, + right: Box::new(Expr::BinaryOp { + left: Box::new(Expr::value(number("2"))), + op: BinaryOperator::Multiply, + right: Box::new(Expr::value(number("3"))), + }), + }, + }) + ); +} + +#[test] +fn parse_duckdb_set_variable_comments_and_script() { + for sql in [ + "set /* a */ variable /* b */ config /* c */ . /* d */ threshold /* e */ to /* f */ 42; -- end", + r#"SET -- a +VARIABLE -- b +config -- c +. -- d +threshold -- e += -- f +42 -- end"#, + ] { + assert_eq!( + duckdb().one_statement_parses_to(sql, "SET VARIABLE config.threshold = 42"), + Statement::Set(Set::SetVariable { + variable: ObjectName::from(vec!["config".into(), "threshold".into()]), + value: Expr::value(number("42")), + }) + ); + } + let statements = duckdb().statements_parse_to( + r#"SET VARIABLE org_id = 42; +SELECT getvariable('org_id') AS org_id;"#, + "SET VARIABLE org_id = 42; SELECT getvariable('org_id') AS org_id", + ); + assert_eq!(statements.len(), 2); + assert_eq!( + statements[0], + Statement::Set(Set::SetVariable { + variable: ObjectName::from(vec!["org_id".into()]), + value: Expr::value(number("42")), + }) + ); + assert_eq!( + statements[1], + duckdb().verified_stmt("SELECT getvariable('org_id') AS org_id") + ); +} + +#[test] +fn parse_duckdb_set_variable_as_ordinary_set_name() { + for (name, parts) in [ + ("VARIABLE", vec![Ident::new("VARIABLE")]), + ("VARIABLE.foo", vec!["VARIABLE".into(), "foo".into()]), + ( + r#"VARIABLE."foo""#, + vec!["VARIABLE".into(), Ident::with_quote('"', "foo")], + ), + ( + "VARIABLE.foo.bar", + vec!["VARIABLE".into(), "foo".into(), "bar".into()], + ), + ] { + for separator in ["=", "TO"] { + assert_eq!( + duckdb().one_statement_parses_to( + &format!("SET {name} {separator} 42"), + &format!("SET {name} = 42"), + ), + Statement::Set(Set::SingleAssignment { + scope: None, + hivevar: false, + variable: ObjectName::from(parts.clone()), + values: vec![Expr::value(number("42"))], + }) + ); + } + } + for sql in [ + r#"SET VARIABLE /* a */ . /* b */ "foo" TO 42"#, + r#"SET VARIABLE -- a +. -- b +"foo" = 42"#, + ] { + assert_eq!( + duckdb().one_statement_parses_to(sql, r#"SET VARIABLE."foo" = 42"#), + Statement::Set(Set::SingleAssignment { + scope: None, + hivevar: false, + variable: ObjectName::from(vec!["VARIABLE".into(), Ident::with_quote('"', "foo")]), + values: vec![Expr::value(number("42"))], + }) + ); + } +} + +#[test] +fn parse_duckdb_set_variable_rejects_string_names() { + for name in [ + "'threshold'", + "'threshold'.value", + "config.'threshold'", + "config.'threshold'.value", + "config /* a */ . /* b */ 'threshold'", + ] { + for separator in ["=", "TO"] { + let sql = format!("SET VARIABLE {name} {separator} 42"); + assert_eq!( + duckdb().parse_sql_statements(&sql).unwrap_err(), + ParserError::ParserError("Expected: identifier, found: 'threshold'".to_string()), + "{sql}" + ); + } + } +} + +#[test] +fn parse_duckdb_set_variable_rejects_malformed_syntax() { + for (sql, expected) in [ + ("SET VARIABLE", "Expected: identifier, found: EOF"), + ("SET VARIABLE 1 = 42", "Expected: identifier, found: 1"), + ( + "SET VARIABLE threshold", + "Expected: equals sign or TO, found: EOF", + ), + ( + "SET VARIABLE threshold 42", + "Expected: equals sign or TO, found: 42", + ), + ( + "SET VARIABLE threshold =", + "Expected: an expression, found: EOF", + ), + ( + "SET VARIABLE threshold TO", + "Expected: an expression, found: EOF", + ), + ( + "SET VARIABLE threshold = ;", + "Expected: an expression, found: ;", + ), + ( + "SET VARIABLE threshold = 1 +", + "Expected: an expression, found: EOF", + ), + ( + "SET VARIABLE threshold. = 42", + "Expected: identifier, found: =", + ), + ( + "SET VARIABLE threshold..value = 42", + "Expected: identifier, found: .", + ), + ( + "SET VARIABLE threshold.", + "Expected: identifier, found: EOF", + ), + ( + "SET VARIABLE threshold = 1, 2", + "Expected: end of statement, found: ,", + ), + ( + "SET VARIABLE threshold = 1, other = 2", + "Expected: end of statement, found: ,", + ), + ( + "SET VARIABLE threshold = 42 extra", + "Expected: end of statement, found: extra", + ), + ] { + assert_eq!( + duckdb().parse_sql_statements(sql).unwrap_err(), + ParserError::ParserError(expected.to_string()), + "{sql}" + ); + } +} + +#[test] +fn parse_duckdb_set_variable_is_dialect_specific() { + let unsupported = all_dialects_where(|d| !d.supports_set_variable_statement()); + for dialect in unsupported.dialects { + let dialect_name = format!("{dialect:?}"); + let supports_session_params = dialect.supports_set_stmt_without_operator(); + let tested = TestedDialects::new(vec![dialect]); + for sql in [ + "SET VARIABLE threshold = 42", + "SET VARIABLE /* name */ threshold = 42", + ] { + if supports_session_params { + assert_eq!( + tested.one_statement_parses_to(sql, "SET VARIABLE threshold = 42"), + Statement::Set(Set::SetSessionParam(SetSessionParamKind::Generic( + SetSessionParamGeneric { + names: vec!["VARIABLE".to_string()], + value: "threshold = 42".to_string(), + }, + ))) + ); + } else { + assert_eq!( + tested.parse_sql_statements(sql).unwrap_err(), + ParserError::ParserError( + "Expected: equals sign or TO, found: threshold".to_string() + ), + "{dialect_name}: {sql}" + ); + } + } + } +}