Skip to content

Commit 75f18ec

Browse files
authored
Add support for DuckDB's CREATE MACRO statements (apache#897)
1 parent 2296de2 commit 75f18ec

File tree

5 files changed

+208
-0
lines changed

5 files changed

+208
-0
lines changed

src/ast/mod.rs

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1580,6 +1580,19 @@ pub enum Statement {
15801580
params: CreateFunctionBody,
15811581
},
15821582
/// ```sql
1583+
/// CREATE MACRO
1584+
/// ```
1585+
///
1586+
/// Supported variants:
1587+
/// 1. [DuckDB](https://duckdb.org/docs/sql/statements/create_macro)
1588+
CreateMacro {
1589+
or_replace: bool,
1590+
temporary: bool,
1591+
name: ObjectName,
1592+
args: Option<Vec<MacroArg>>,
1593+
definition: MacroDefinition,
1594+
},
1595+
/// ```sql
15831596
/// CREATE STAGE
15841597
/// ```
15851598
/// See <https://docs.snowflake.com/en/sql-reference/sql/create-stage>
@@ -2098,6 +2111,28 @@ impl fmt::Display for Statement {
20982111
write!(f, "{params}")?;
20992112
Ok(())
21002113
}
2114+
Statement::CreateMacro {
2115+
or_replace,
2116+
temporary,
2117+
name,
2118+
args,
2119+
definition,
2120+
} => {
2121+
write!(
2122+
f,
2123+
"CREATE {or_replace}{temp}MACRO {name}",
2124+
temp = if *temporary { "TEMPORARY " } else { "" },
2125+
or_replace = if *or_replace { "OR REPLACE " } else { "" },
2126+
)?;
2127+
if let Some(args) = args {
2128+
write!(f, "({})", display_comma_separated(args))?;
2129+
}
2130+
match definition {
2131+
MacroDefinition::Expr(expr) => write!(f, " AS {expr}")?,
2132+
MacroDefinition::Table(query) => write!(f, " AS TABLE {query}")?,
2133+
}
2134+
Ok(())
2135+
}
21012136
Statement::CreateView {
21022137
name,
21032138
or_replace,
@@ -4304,6 +4339,56 @@ impl fmt::Display for CreateFunctionUsing {
43044339
}
43054340
}
43064341

4342+
/// `NAME = <EXPR>` arguments for DuckDB macros
4343+
///
4344+
/// See [Create Macro - DuckDB](https://duckdb.org/docs/sql/statements/create_macro)
4345+
/// for more details
4346+
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
4347+
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
4348+
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
4349+
pub struct MacroArg {
4350+
pub name: Ident,
4351+
pub default_expr: Option<Expr>,
4352+
}
4353+
4354+
impl MacroArg {
4355+
/// Returns an argument with name.
4356+
pub fn new(name: &str) -> Self {
4357+
Self {
4358+
name: name.into(),
4359+
default_expr: None,
4360+
}
4361+
}
4362+
}
4363+
4364+
impl fmt::Display for MacroArg {
4365+
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
4366+
write!(f, "{}", self.name)?;
4367+
if let Some(default_expr) = &self.default_expr {
4368+
write!(f, " := {default_expr}")?;
4369+
}
4370+
Ok(())
4371+
}
4372+
}
4373+
4374+
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
4375+
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
4376+
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
4377+
pub enum MacroDefinition {
4378+
Expr(Expr),
4379+
Table(Query),
4380+
}
4381+
4382+
impl fmt::Display for MacroDefinition {
4383+
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
4384+
match self {
4385+
MacroDefinition::Expr(expr) => write!(f, "{expr}")?,
4386+
MacroDefinition::Table(query) => write!(f, "{query}")?,
4387+
}
4388+
Ok(())
4389+
}
4390+
}
4391+
43074392
/// Schema possible naming variants ([1]).
43084393
///
43094394
/// [1]: https://jakewheat.github.io/sql-overview/sql-2016-foundation-grammar.html#schema-definition

src/keywords.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -347,6 +347,7 @@ define_keywords!(
347347
LOCKED,
348348
LOGIN,
349349
LOWER,
350+
MACRO,
350351
MANAGEDLOCATION,
351352
MATCH,
352353
MATCHED,

src/parser.rs

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2346,6 +2346,8 @@ impl<'a> Parser<'a> {
23462346
self.parse_create_external_table(or_replace)
23472347
} else if self.parse_keyword(Keyword::FUNCTION) {
23482348
self.parse_create_function(or_replace, temporary)
2349+
} else if self.parse_keyword(Keyword::MACRO) {
2350+
self.parse_create_macro(or_replace, temporary)
23492351
} else if or_replace {
23502352
self.expected(
23512353
"[EXTERNAL] TABLE or [MATERIALIZED] VIEW or FUNCTION after CREATE OR REPLACE",
@@ -2624,6 +2626,8 @@ impl<'a> Parser<'a> {
26242626
return_type,
26252627
params,
26262628
})
2629+
} else if dialect_of!(self is DuckDbDialect) {
2630+
self.parse_create_macro(or_replace, temporary)
26272631
} else {
26282632
self.prev_token();
26292633
self.expected("an object type after CREATE", self.peek_token())
@@ -2699,6 +2703,53 @@ impl<'a> Parser<'a> {
26992703
}
27002704
}
27012705

2706+
pub fn parse_create_macro(
2707+
&mut self,
2708+
or_replace: bool,
2709+
temporary: bool,
2710+
) -> Result<Statement, ParserError> {
2711+
if dialect_of!(self is DuckDbDialect | GenericDialect) {
2712+
let name = self.parse_object_name()?;
2713+
self.expect_token(&Token::LParen)?;
2714+
let args = if self.consume_token(&Token::RParen) {
2715+
self.prev_token();
2716+
None
2717+
} else {
2718+
Some(self.parse_comma_separated(Parser::parse_macro_arg)?)
2719+
};
2720+
2721+
self.expect_token(&Token::RParen)?;
2722+
self.expect_keyword(Keyword::AS)?;
2723+
2724+
Ok(Statement::CreateMacro {
2725+
or_replace,
2726+
temporary,
2727+
name,
2728+
args,
2729+
definition: if self.parse_keyword(Keyword::TABLE) {
2730+
MacroDefinition::Table(self.parse_query()?)
2731+
} else {
2732+
MacroDefinition::Expr(self.parse_expr()?)
2733+
},
2734+
})
2735+
} else {
2736+
self.prev_token();
2737+
self.expected("an object type after CREATE", self.peek_token())
2738+
}
2739+
}
2740+
2741+
fn parse_macro_arg(&mut self) -> Result<MacroArg, ParserError> {
2742+
let name = self.parse_identifier()?;
2743+
2744+
let default_expr =
2745+
if self.consume_token(&Token::DuckAssignment) || self.consume_token(&Token::RArrow) {
2746+
Some(self.parse_expr()?)
2747+
} else {
2748+
None
2749+
};
2750+
Ok(MacroArg { name, default_expr })
2751+
}
2752+
27022753
pub fn parse_create_external_table(
27032754
&mut self,
27042755
or_replace: bool,

src/tokenizer.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -114,6 +114,8 @@ pub enum Token {
114114
Colon,
115115
/// DoubleColon `::` (used for casting in postgresql)
116116
DoubleColon,
117+
/// Assignment `:=` (used for keyword argument in DuckDB macros)
118+
DuckAssignment,
117119
/// SemiColon `;` used as separator for COPY and payload
118120
SemiColon,
119121
/// Backslash `\` used in terminating the COPY payload with `\.`
@@ -222,6 +224,7 @@ impl fmt::Display for Token {
222224
Token::Period => f.write_str("."),
223225
Token::Colon => f.write_str(":"),
224226
Token::DoubleColon => f.write_str("::"),
227+
Token::DuckAssignment => f.write_str(":="),
225228
Token::SemiColon => f.write_str(";"),
226229
Token::Backslash => f.write_str("\\"),
227230
Token::LBracket => f.write_str("["),
@@ -847,6 +850,7 @@ impl<'a> Tokenizer<'a> {
847850
chars.next();
848851
match chars.peek() {
849852
Some(':') => self.consume_and_return(chars, Token::DoubleColon),
853+
Some('=') => self.consume_and_return(chars, Token::DuckAssignment),
850854
_ => Ok(Some(Token::Colon)),
851855
}
852856
}

tests/sqlparser_duckdb.rs

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,3 +68,70 @@ fn test_select_wildcard_with_exclude() {
6868
fn parse_div_infix() {
6969
duckdb_and_generic().verified_stmt(r#"SELECT 5 // 2"#);
7070
}
71+
72+
#[test]
73+
fn test_create_macro() {
74+
let macro_ = duckdb().verified_stmt("CREATE MACRO schema.add(a, b) AS a + b");
75+
let expected = Statement::CreateMacro {
76+
or_replace: false,
77+
temporary: false,
78+
name: ObjectName(vec![Ident::new("schema"), Ident::new("add")]),
79+
args: Some(vec![MacroArg::new("a"), MacroArg::new("b")]),
80+
definition: MacroDefinition::Expr(Expr::BinaryOp {
81+
left: Box::new(Expr::Identifier(Ident::new("a"))),
82+
op: BinaryOperator::Plus,
83+
right: Box::new(Expr::Identifier(Ident::new("b"))),
84+
}),
85+
};
86+
assert_eq!(expected, macro_);
87+
}
88+
89+
#[test]
90+
fn test_create_macro_default_args() {
91+
let macro_ = duckdb().verified_stmt("CREATE MACRO add_default(a, b := 5) AS a + b");
92+
let expected = Statement::CreateMacro {
93+
or_replace: false,
94+
temporary: false,
95+
name: ObjectName(vec![Ident::new("add_default")]),
96+
args: Some(vec![
97+
MacroArg::new("a"),
98+
MacroArg {
99+
name: Ident::new("b"),
100+
default_expr: Some(Expr::Value(Value::Number(
101+
#[cfg(not(feature = "bigdecimal"))]
102+
5.to_string(),
103+
#[cfg(feature = "bigdecimal")]
104+
bigdecimal::BigDecimal::from(5),
105+
false,
106+
))),
107+
},
108+
]),
109+
definition: MacroDefinition::Expr(Expr::BinaryOp {
110+
left: Box::new(Expr::Identifier(Ident::new("a"))),
111+
op: BinaryOperator::Plus,
112+
right: Box::new(Expr::Identifier(Ident::new("b"))),
113+
}),
114+
};
115+
assert_eq!(expected, macro_);
116+
}
117+
118+
#[test]
119+
fn test_create_table_macro() {
120+
let query = "SELECT col1_value AS column1, col2_value AS column2 UNION ALL SELECT 'Hello' AS col1_value, 456 AS col2_value";
121+
let macro_ = duckdb().verified_stmt(
122+
&("CREATE OR REPLACE TEMPORARY MACRO dynamic_table(col1_value, col2_value) AS TABLE "
123+
.to_string()
124+
+ query),
125+
);
126+
let expected = Statement::CreateMacro {
127+
or_replace: true,
128+
temporary: true,
129+
name: ObjectName(vec![Ident::new("dynamic_table")]),
130+
args: Some(vec![
131+
MacroArg::new("col1_value"),
132+
MacroArg::new("col2_value"),
133+
]),
134+
definition: MacroDefinition::Table(duckdb().verified_query(query)),
135+
};
136+
assert_eq!(expected, macro_);
137+
}

0 commit comments

Comments
 (0)