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

This commit is contained in:
dawg 2023-06-21 21:12:58 +02:00 committed by GitHub
parent 2296de2bc4
commit 75f18ecfda
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 208 additions and 0 deletions

View file

@ -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>
@ -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,
@ -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

View file

@ -347,6 +347,7 @@ define_keywords!(
LOCKED,
LOGIN,
LOWER,
MACRO,
MANAGEDLOCATION,
MATCH,
MATCHED,

View file

@ -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,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,

View file

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

View file

@ -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_);
}