From 67498060d93922657f28ee8e98017b21b20fdefb Mon Sep 17 00:00:00 2001 From: Martin Nowak Date: Sun, 11 Jun 2023 13:17:06 +0200 Subject: [PATCH 1/2] Add support for DuckDB's CREATE MACRO statements - support scalar and table macros --- src/ast/mod.rs | 97 +++++++++++++++++++++++++++++++++++++++ src/keywords.rs | 1 + src/parser.rs | 60 ++++++++++++++++++++++++ src/tokenizer.rs | 4 ++ tests/sqlparser_duckdb.rs | 50 ++++++++++++++++++++ 5 files changed, 212 insertions(+) diff --git a/src/ast/mod.rs b/src/ast/mod.rs index 7f3b4742f..f64744ed6 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -1580,6 +1580,32 @@ pub enum Statement { params: CreateFunctionBody, }, /// ```sql + /// CREATE MACRO + /// ``` + /// + /// Supported variants: + /// 1. [DuckDB](https://duckdb.org/docs/sql/statements/create_macro) + CreateMacro { + or_replace: bool, + temporary: bool, + name: ObjectName, + args: Option>, + expr: Expr, + }, + /// ```sql + /// CREATE MACRO + /// ``` + /// + /// Supported variants: + /// 1. [DuckDB](https://duckdb.org/docs/sql/statements/create_macro) + CreateTableMacro { + or_replace: bool, + temporary: bool, + name: ObjectName, + args: Option>, + query: Query, + }, + /// ```sql /// CREATE STAGE /// ``` /// See @@ -2098,6 +2124,44 @@ impl fmt::Display for Statement { write!(f, "{params}")?; Ok(()) } + Statement::CreateMacro { + or_replace, + temporary, + name, + args, + expr, + } => { + write!( + f, + "CREATE {or_replace}{temp}MACRO {name}", + temp = if *temporary { "TEMPORARY " } else { "" }, + or_replace = if *or_replace { "OR REPLACE " } else { "" }, + )?; + if let Some(args) = args { + write!(f, "({})", display_comma_separated(args))?; + } + write!(f, " AS {expr}")?; + Ok(()) + } + Statement::CreateTableMacro { + or_replace, + temporary, + name, + args, + query, + } => { + write!( + f, + "CREATE {or_replace}{temp}MACRO {name} ", + temp = if *temporary { "TEMPORARY " } else { "" }, + or_replace = if *or_replace { "OR REPLACE " } else { "" }, + )?; + if let Some(args) = args { + write!(f, "({})", display_comma_separated(args))?; + } + write!(f, " AS TABLE {query}")?; + Ok(()) + } Statement::CreateView { name, or_replace, @@ -4304,6 +4368,39 @@ impl fmt::Display for CreateFunctionUsing { } } +/// DuckDB specific feature. +/// +/// See [Create Macro - DuckDB](https://duckdb.org/docs/sql/statements/create_macro) +/// for more details +#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))] +pub struct OperateMacroArg { + pub name: Ident, + pub default_expr: Option, +} + +impl OperateMacroArg { + /// Returns an argument with name. + pub fn new(name: &str) -> Self { + Self { + name: name.into(), + default_expr: None, + } + } +} + + +impl fmt::Display for OperateMacroArg { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "{}", self.name)?; + if let Some(default_expr) = &self.default_expr { + write!(f, " := {default_expr}")?; + } + Ok(()) + } +} + /// Schema possible naming variants ([1]). /// /// [1]: https://jakewheat.github.io/sql-overview/sql-2016-foundation-grammar.html#schema-definition diff --git a/src/keywords.rs b/src/keywords.rs index e73b89a98..663818c7e 100644 --- a/src/keywords.rs +++ b/src/keywords.rs @@ -347,6 +347,7 @@ define_keywords!( LOCKED, LOGIN, LOWER, + MACRO, MANAGEDLOCATION, MATCH, MATCHED, diff --git a/src/parser.rs b/src/parser.rs index b89077fb5..74e6caa9f 100644 --- a/src/parser.rs +++ b/src/parser.rs @@ -2346,6 +2346,8 @@ impl<'a> Parser<'a> { self.parse_create_external_table(or_replace) } else if self.parse_keyword(Keyword::FUNCTION) { self.parse_create_function(or_replace, temporary) + } else if self.parse_keyword(Keyword::MACRO) { + self.parse_create_macro(or_replace, temporary) } else if or_replace { self.expected( "[EXTERNAL] TABLE or [MATERIALIZED] VIEW or FUNCTION after CREATE OR REPLACE", @@ -2624,6 +2626,8 @@ impl<'a> Parser<'a> { return_type, params, }) + } else if dialect_of!(self is DuckDbDialect) { + self.parse_create_macro(or_replace, temporary) } else { self.prev_token(); self.expected("an object type after CREATE", self.peek_token()) @@ -2699,6 +2703,62 @@ impl<'a> Parser<'a> { } } + pub fn parse_create_macro( + &mut self, + or_replace: bool, + temporary: bool, + ) -> Result { + if dialect_of!(self is DuckDbDialect) { + let name = self.parse_object_name()?; + self.expect_token(&Token::LParen)?; + let args = if self.consume_token(&Token::RParen) { + self.prev_token(); + None + } else { + Some(self.parse_comma_separated(Parser::parse_macro_arg)?) + }; + + self.expect_token(&Token::RParen)?; + self.expect_keyword(Keyword::AS)?; + + if self.parse_keyword(Keyword::TABLE) { + Ok(Statement::CreateTableMacro { + or_replace, + temporary, + name, + args, + query: self.parse_query()?, + }) + } else { + Ok(Statement::CreateMacro { + or_replace, + temporary, + name, + args, + expr: self.parse_expr()?, + }) + } + } else { + self.prev_token(); + self.expected("an object type after CREATE", self.peek_token()) + } + } + + fn parse_macro_arg(&mut self) -> Result { + let name = self.parse_identifier()?; + + let default_expr = if self.consume_token(&Token::DuckAssignment) || self.consume_token(&Token::RArrow) + { + Some(self.parse_expr()?) + } else { + None + }; + Ok(OperateMacroArg { + name, + default_expr, + }) + } + pub fn parse_create_external_table( &mut self, or_replace: bool, diff --git a/src/tokenizer.rs b/src/tokenizer.rs index ffa1a96f2..257a3517e 100644 --- a/src/tokenizer.rs +++ b/src/tokenizer.rs @@ -114,6 +114,8 @@ pub enum Token { Colon, /// DoubleColon `::` (used for casting in postgresql) DoubleColon, + /// Assignment `:=` (used for keyword argument in DuckDB macros) + DuckAssignment, /// SemiColon `;` used as separator for COPY and payload SemiColon, /// Backslash `\` used in terminating the COPY payload with `\.` @@ -222,6 +224,7 @@ impl fmt::Display for Token { Token::Period => f.write_str("."), Token::Colon => f.write_str(":"), Token::DoubleColon => f.write_str("::"), + Token::DuckAssignment => f.write_str(":="), Token::SemiColon => f.write_str(";"), Token::Backslash => f.write_str("\\"), Token::LBracket => f.write_str("["), @@ -847,6 +850,7 @@ impl<'a> Tokenizer<'a> { chars.next(); match chars.peek() { Some(':') => self.consume_and_return(chars, Token::DoubleColon), + Some('=') => self.consume_and_return(chars, Token::DuckAssignment), _ => Ok(Some(Token::Colon)), } } diff --git a/tests/sqlparser_duckdb.rs b/tests/sqlparser_duckdb.rs index 1a4f04c33..84709cc85 100644 --- a/tests/sqlparser_duckdb.rs +++ b/tests/sqlparser_duckdb.rs @@ -68,3 +68,53 @@ fn test_select_wildcard_with_exclude() { fn parse_div_infix() { duckdb_and_generic().verified_stmt(r#"SELECT 5 // 2"#); } + +#[test] +fn test_create_macro() { + let _macro = duckdb().verified_stmt("CREATE MACRO schema.add(a, b) AS a + b"); + let expected = Statement::CreateMacro { + or_replace: false, + temporary: false, + name: ObjectName(vec![Ident::new("schema"), Ident::new("add")]), + args: Some(vec![OperateMacroArg::new("a"), OperateMacroArg::new("b")]), + expr: Expr::BinaryOp { + left: Box::new(Expr::Identifier(Ident::new("a"))), + op: BinaryOperator::Plus, + right: Box::new(Expr::Identifier(Ident::new("b"))), + }, + }; + assert_eq!(expected, _macro); +} + +#[test] +fn test_create_macro_default_args() { + let _macro = duckdb().verified_stmt("CREATE MACRO add_default(a, b := 5) AS a + b"); + let expected = Statement::CreateMacro { + or_replace: false, + temporary: false, + name: ObjectName(vec![Ident::new("add_default")]), + args: Some(vec![OperateMacroArg::new("a"), OperateMacroArg{ name: Ident::new("b"), default_expr: Some(Expr::Value(Value::Number("5".into(), false)))}]), + expr: Expr::BinaryOp { + left: Box::new(Expr::Identifier(Ident::new("a"))), + op: BinaryOperator::Plus, + right: Box::new(Expr::Identifier(Ident::new("b"))), + }, + }; + assert_eq!(expected, _macro); +} + +#[test] +fn test_create_table_macro() { + let query = "SELECT col1_value AS column1, col2_value AS column2 UNION ALL SELECT 'Hello' AS col1_value, 456 AS col2_value"; + let _macro = duckdb().verified_stmt( + &("CREATE OR REPLACE TEMPORARY MACRO dynamic_table (col1_value, col2_value) AS TABLE ".to_string() + query), + ); + let expected = Statement::CreateTableMacro { + or_replace: true, + temporary: true, + name: ObjectName(vec![Ident::new("dynamic_table")]), + args: Some(vec![OperateMacroArg::new("col1_value"), OperateMacroArg::new("col2_value")]), + query: duckdb().verified_query(query), + }; + assert_eq!(expected, _macro); +} From 088b2fe0759311801740d5ba490d7acb376a4e4a Mon Sep 17 00:00:00 2001 From: Martin Nowak Date: Mon, 19 Jun 2023 10:47:40 +0200 Subject: [PATCH 2/2] fixup! Add support for DuckDB's CREATE MACRO statements --- src/ast/mod.rs | 68 ++++++++++++++++----------------------- src/parser.rs | 51 ++++++++++++----------------- tests/sqlparser_duckdb.rs | 49 +++++++++++++++++++--------- 3 files changed, 82 insertions(+), 86 deletions(-) diff --git a/src/ast/mod.rs b/src/ast/mod.rs index f64744ed6..7afb576bc 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -1589,21 +1589,8 @@ pub enum Statement { or_replace: bool, temporary: bool, name: ObjectName, - args: Option>, - expr: Expr, - }, - /// ```sql - /// CREATE MACRO - /// ``` - /// - /// Supported variants: - /// 1. [DuckDB](https://duckdb.org/docs/sql/statements/create_macro) - CreateTableMacro { - or_replace: bool, - temporary: bool, - name: ObjectName, - args: Option>, - query: Query, + args: Option>, + definition: MacroDefinition, }, /// ```sql /// CREATE STAGE @@ -2129,7 +2116,7 @@ impl fmt::Display for Statement { temporary, name, args, - expr, + definition, } => { write!( f, @@ -2140,26 +2127,10 @@ impl fmt::Display for Statement { if let Some(args) = args { write!(f, "({})", display_comma_separated(args))?; } - write!(f, " AS {expr}")?; - Ok(()) - } - Statement::CreateTableMacro { - or_replace, - temporary, - name, - args, - query, - } => { - write!( - f, - "CREATE {or_replace}{temp}MACRO {name} ", - temp = if *temporary { "TEMPORARY " } else { "" }, - or_replace = if *or_replace { "OR REPLACE " } else { "" }, - )?; - if let Some(args) = args { - write!(f, "({})", display_comma_separated(args))?; + match definition { + MacroDefinition::Expr(expr) => write!(f, " AS {expr}")?, + MacroDefinition::Table(query) => write!(f, " AS TABLE {query}")?, } - write!(f, " AS TABLE {query}")?; Ok(()) } Statement::CreateView { @@ -4368,19 +4339,19 @@ impl fmt::Display for CreateFunctionUsing { } } -/// DuckDB specific feature. +/// `NAME = ` arguments for DuckDB macros /// /// See [Create Macro - DuckDB](https://duckdb.org/docs/sql/statements/create_macro) /// for more details #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)] #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))] -pub struct OperateMacroArg { +pub struct MacroArg { pub name: Ident, pub default_expr: Option, } -impl OperateMacroArg { +impl MacroArg { /// Returns an argument with name. pub fn new(name: &str) -> Self { Self { @@ -4390,8 +4361,7 @@ impl OperateMacroArg { } } - -impl fmt::Display for OperateMacroArg { +impl fmt::Display for MacroArg { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { write!(f, "{}", self.name)?; if let Some(default_expr) = &self.default_expr { @@ -4401,6 +4371,24 @@ impl fmt::Display for OperateMacroArg { } } +#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))] +pub enum MacroDefinition { + Expr(Expr), + Table(Query), +} + +impl fmt::Display for MacroDefinition { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + match self { + MacroDefinition::Expr(expr) => write!(f, "{expr}")?, + MacroDefinition::Table(query) => write!(f, "{query}")?, + } + Ok(()) + } +} + /// Schema possible naming variants ([1]). /// /// [1]: https://jakewheat.github.io/sql-overview/sql-2016-foundation-grammar.html#schema-definition diff --git a/src/parser.rs b/src/parser.rs index 74e6caa9f..92ad37859 100644 --- a/src/parser.rs +++ b/src/parser.rs @@ -2708,55 +2708,46 @@ impl<'a> Parser<'a> { or_replace: bool, temporary: bool, ) -> Result { - if dialect_of!(self is DuckDbDialect) { + if dialect_of!(self is DuckDbDialect | GenericDialect) { let name = self.parse_object_name()?; self.expect_token(&Token::LParen)?; let args = if self.consume_token(&Token::RParen) { self.prev_token(); None } else { - Some(self.parse_comma_separated(Parser::parse_macro_arg)?) + Some(self.parse_comma_separated(Parser::parse_macro_arg)?) }; self.expect_token(&Token::RParen)?; self.expect_keyword(Keyword::AS)?; - if self.parse_keyword(Keyword::TABLE) { - Ok(Statement::CreateTableMacro { - or_replace, - temporary, - name, - args, - query: self.parse_query()?, - }) - } else { - Ok(Statement::CreateMacro { - or_replace, - temporary, - name, - args, - expr: self.parse_expr()?, - }) - } + Ok(Statement::CreateMacro { + or_replace, + temporary, + name, + args, + definition: if self.parse_keyword(Keyword::TABLE) { + MacroDefinition::Table(self.parse_query()?) + } else { + MacroDefinition::Expr(self.parse_expr()?) + }, + }) } else { self.prev_token(); self.expected("an object type after CREATE", self.peek_token()) } } - fn parse_macro_arg(&mut self) -> Result { + fn parse_macro_arg(&mut self) -> Result { let name = self.parse_identifier()?; - let default_expr = if self.consume_token(&Token::DuckAssignment) || self.consume_token(&Token::RArrow) - { - Some(self.parse_expr()?) - } else { - None - }; - Ok(OperateMacroArg { - name, - default_expr, - }) + let default_expr = + if self.consume_token(&Token::DuckAssignment) || self.consume_token(&Token::RArrow) { + Some(self.parse_expr()?) + } else { + None + }; + Ok(MacroArg { name, default_expr }) } pub fn parse_create_external_table( diff --git a/tests/sqlparser_duckdb.rs b/tests/sqlparser_duckdb.rs index 84709cc85..bb60235a6 100644 --- a/tests/sqlparser_duckdb.rs +++ b/tests/sqlparser_duckdb.rs @@ -71,50 +71,67 @@ fn parse_div_infix() { #[test] fn test_create_macro() { - let _macro = duckdb().verified_stmt("CREATE MACRO schema.add(a, b) AS a + b"); + let macro_ = duckdb().verified_stmt("CREATE MACRO schema.add(a, b) AS a + b"); let expected = Statement::CreateMacro { or_replace: false, temporary: false, name: ObjectName(vec![Ident::new("schema"), Ident::new("add")]), - args: Some(vec![OperateMacroArg::new("a"), OperateMacroArg::new("b")]), - expr: Expr::BinaryOp { + args: Some(vec![MacroArg::new("a"), MacroArg::new("b")]), + definition: MacroDefinition::Expr(Expr::BinaryOp { left: Box::new(Expr::Identifier(Ident::new("a"))), op: BinaryOperator::Plus, right: Box::new(Expr::Identifier(Ident::new("b"))), - }, + }), }; - assert_eq!(expected, _macro); + assert_eq!(expected, macro_); } #[test] fn test_create_macro_default_args() { - let _macro = duckdb().verified_stmt("CREATE MACRO add_default(a, b := 5) AS a + b"); + let macro_ = duckdb().verified_stmt("CREATE MACRO add_default(a, b := 5) AS a + b"); let expected = Statement::CreateMacro { or_replace: false, temporary: false, name: ObjectName(vec![Ident::new("add_default")]), - args: Some(vec![OperateMacroArg::new("a"), OperateMacroArg{ name: Ident::new("b"), default_expr: Some(Expr::Value(Value::Number("5".into(), false)))}]), - expr: Expr::BinaryOp { + args: Some(vec![ + MacroArg::new("a"), + MacroArg { + name: Ident::new("b"), + default_expr: Some(Expr::Value(Value::Number( + #[cfg(not(feature = "bigdecimal"))] + 5.to_string(), + #[cfg(feature = "bigdecimal")] + bigdecimal::BigDecimal::from(5), + false, + ))), + }, + ]), + definition: MacroDefinition::Expr(Expr::BinaryOp { left: Box::new(Expr::Identifier(Ident::new("a"))), op: BinaryOperator::Plus, right: Box::new(Expr::Identifier(Ident::new("b"))), - }, + }), }; - assert_eq!(expected, _macro); + assert_eq!(expected, macro_); } #[test] fn test_create_table_macro() { let query = "SELECT col1_value AS column1, col2_value AS column2 UNION ALL SELECT 'Hello' AS col1_value, 456 AS col2_value"; - let _macro = duckdb().verified_stmt( - &("CREATE OR REPLACE TEMPORARY MACRO dynamic_table (col1_value, col2_value) AS TABLE ".to_string() + query), + let macro_ = duckdb().verified_stmt( + &("CREATE OR REPLACE TEMPORARY MACRO dynamic_table(col1_value, col2_value) AS TABLE " + .to_string() + + query), ); - let expected = Statement::CreateTableMacro { + let expected = Statement::CreateMacro { or_replace: true, temporary: true, name: ObjectName(vec![Ident::new("dynamic_table")]), - args: Some(vec![OperateMacroArg::new("col1_value"), OperateMacroArg::new("col2_value")]), - query: duckdb().verified_query(query), + args: Some(vec![ + MacroArg::new("col1_value"), + MacroArg::new("col2_value"), + ]), + definition: MacroDefinition::Table(duckdb().verified_query(query)), }; - assert_eq!(expected, _macro); + assert_eq!(expected, macro_); }