Add support for DuckDB's CREATE MACRO statements (apache#897)

Author: MartinNowak

src/ast/mod.rs

@@ -1697,6 +1697,19 @@ 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<Vec<MacroArg>>,
+        definition: MacroDefinition,
+    },
+    /// ```sql
     /// CREATE STAGE
     /// ```
     /// See <https://docs.snowflake.com/en/sql-reference/sql/create-stage>
@@ -2288,6 +2301,28 @@ impl fmt::Display for Statement {
                 write!(f, "{params}")?;
                 Ok(())
             }
+            Statement::CreateMacro {
+                or_replace,
+                temporary,
+                name,
+                args,
+                definition,
+            } => {
+                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}")?,
+                }
+                Ok(())
+            }
             Statement::CreateView {
                 name,
                 or_replace,
@@ -4585,6 +4620,56 @@ impl fmt::Display for CreateFunctionUsing {
     }
 }
 
+/// `NAME = <EXPR>` 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 MacroArg {
+    pub name: Ident,
+    pub default_expr: Option<Expr>,
+}
+
+impl MacroArg {
+    /// Returns an argument with name.
+    pub fn new(name: &str) -> Self {
+        Self {
+            name: name.into(),
+            default_expr: None,
+        }
+    }
+}
+
+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 {
+            write!(f, " := {default_expr}")?;
+        }
+        Ok(())
+    }
+}
+
+#[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

src/keywords.rs

@@ -351,6 +351,7 @@ define_keywords!(
     LOCKED,
     LOGIN,
     LOWER,
+    MACRO,
     MANAGEDLOCATION,
     MAPPING,
     MATCH,

src/parser.rs

@@ -2350,6 +2350,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",
@@ -2632,6 +2634,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())
@@ -2707,6 +2711,53 @@ impl<'a> Parser<'a> {
         }
     }
 
+    pub fn parse_create_macro(
+        &mut self,
+        or_replace: bool,
+        temporary: bool,
+    ) -> Result<Statement, ParserError> {
+        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)?)
+            };
+
+            self.expect_token(&Token::RParen)?;
+            self.expect_keyword(Keyword::AS)?;
+
+            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<MacroArg, ParserError> {
+        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(MacroArg { name, default_expr })
+    }
+
     pub fn parse_create_external_table(
         &mut self,
         or_replace: bool,

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)),
                     }
                 }

tests/sqlparser_duckdb.rs

@@ -68,3 +68,70 @@ 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![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_);
+}
+
+#[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![
+            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_);
+}
+
+#[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::CreateMacro {
+        or_replace: true,
+        temporary: true,
+        name: ObjectName(vec![Ident::new("dynamic_table")]),
+        args: Some(vec![
+            MacroArg::new("col1_value"),
+            MacroArg::new("col2_value"),
+        ]),
+        definition: MacroDefinition::Table(duckdb().verified_query(query)),
+    };
+    assert_eq!(expected, macro_);
+}