Skip to content

Add support for DuckDB's CREATE MACRO statements #897

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Jun 21, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
85 changes: 85 additions & 0 deletions src/ast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1580,6 +1580,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>
Expand Down Expand Up @@ -2098,6 +2111,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,
Expand Down Expand Up @@ -4304,6 +4339,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
Expand Down
1 change: 1 addition & 0 deletions src/keywords.rs
Original file line number Diff line number Diff line change
Expand Up @@ -347,6 +347,7 @@ define_keywords!(
LOCKED,
LOGIN,
LOWER,
MACRO,
MANAGEDLOCATION,
MATCH,
MATCHED,
Expand Down
51 changes: 51 additions & 0 deletions src/parser.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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())
Expand Down Expand Up @@ -2699,6 +2703,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,
Expand Down
4 changes: 4 additions & 0 deletions src/tokenizer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 `\.`
Expand Down Expand Up @@ -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("["),
Expand Down Expand Up @@ -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)),
}
}
Expand Down
67 changes: 67 additions & 0 deletions tests/sqlparser_duckdb.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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_);
}