Add support for DuckDB functions named arguments with assignment operator (#1195)

Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>
This commit is contained in:
gstvg 2024-04-06 13:46:36 -03:00 committed by GitHub
parent e747c9c2af
commit 14b33ac493
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 56 additions and 6 deletions

View file

@ -4584,6 +4584,8 @@ pub enum FunctionArgOperator {
Equals, Equals,
/// function(arg1 => value1) /// function(arg1 => value1)
RightArrow, RightArrow,
/// function(arg1 := value1)
Assignment,
} }
impl fmt::Display for FunctionArgOperator { impl fmt::Display for FunctionArgOperator {
@ -4591,6 +4593,7 @@ impl fmt::Display for FunctionArgOperator {
match self { match self {
FunctionArgOperator::Equals => f.write_str("="), FunctionArgOperator::Equals => f.write_str("="),
FunctionArgOperator::RightArrow => f.write_str("=>"), FunctionArgOperator::RightArrow => f.write_str("=>"),
FunctionArgOperator::Assignment => f.write_str(":="),
} }
} }
} }

View file

@ -3485,7 +3485,7 @@ impl<'a> Parser<'a> {
let name = self.parse_identifier(false)?; let name = self.parse_identifier(false)?;
let default_expr = let default_expr =
if self.consume_token(&Token::DuckAssignment) || self.consume_token(&Token::RArrow) { if self.consume_token(&Token::Assignment) || self.consume_token(&Token::RArrow) {
Some(self.parse_expr()?) Some(self.parse_expr()?)
} else { } else {
None None
@ -4183,7 +4183,7 @@ impl<'a> Parser<'a> {
self.next_token(); // Skip `DEFAULT` self.next_token(); // Skip `DEFAULT`
Some(DeclareAssignment::Default(Box::new(self.parse_expr()?))) Some(DeclareAssignment::Default(Box::new(self.parse_expr()?)))
} }
Token::DuckAssignment => { Token::Assignment => {
self.next_token(); // Skip `:=` self.next_token(); // Skip `:=`
Some(DeclareAssignment::DuckAssignment(Box::new( Some(DeclareAssignment::DuckAssignment(Box::new(
self.parse_expr()?, self.parse_expr()?,
@ -8602,6 +8602,19 @@ impl<'a> Parser<'a> {
arg, arg,
operator: FunctionArgOperator::Equals, operator: FunctionArgOperator::Equals,
}) })
} else if dialect_of!(self is DuckDbDialect | GenericDialect)
&& self.peek_nth_token(1) == Token::Assignment
{
let name = self.parse_identifier(false)?;
self.expect_token(&Token::Assignment)?;
let arg = self.parse_expr()?.into();
Ok(FunctionArg::Named {
name,
arg,
operator: FunctionArgOperator::Assignment,
})
} else { } else {
Ok(FunctionArg::Unnamed(self.parse_wildcard_expr()?.into())) Ok(FunctionArg::Unnamed(self.parse_wildcard_expr()?.into()))
} }

View file

@ -117,8 +117,8 @@ pub enum Token {
Colon, Colon,
/// DoubleColon `::` (used for casting in PostgreSQL) /// DoubleColon `::` (used for casting in PostgreSQL)
DoubleColon, DoubleColon,
/// Assignment `:=` (used for keyword argument in DuckDB macros) /// Assignment `:=` (used for keyword argument in DuckDB macros and some functions, and for variable declarations in DuckDB and Snowflake)
DuckAssignment, Assignment,
/// SemiColon `;` used as separator for COPY and payload /// SemiColon `;` used as separator for COPY and payload
SemiColon, SemiColon,
/// Backslash `\` used in terminating the COPY payload with `\.` /// Backslash `\` used in terminating the COPY payload with `\.`
@ -239,7 +239,7 @@ impl fmt::Display for Token {
Token::Period => f.write_str("."), Token::Period => f.write_str("."),
Token::Colon => f.write_str(":"), Token::Colon => f.write_str(":"),
Token::DoubleColon => f.write_str("::"), Token::DoubleColon => f.write_str("::"),
Token::DuckAssignment => f.write_str(":="), Token::Assignment => f.write_str(":="),
Token::SemiColon => f.write_str(";"), Token::SemiColon => f.write_str(";"),
Token::Backslash => f.write_str("\\"), Token::Backslash => f.write_str("\\"),
Token::LBracket => f.write_str("["), Token::LBracket => f.write_str("["),
@ -959,7 +959,7 @@ impl<'a> Tokenizer<'a> {
chars.next(); chars.next();
match chars.peek() { match chars.peek() {
Some(':') => self.consume_and_return(chars, Token::DoubleColon), Some(':') => self.consume_and_return(chars, Token::DoubleColon),
Some('=') => self.consume_and_return(chars, Token::DuckAssignment), Some('=') => self.consume_and_return(chars, Token::Assignment),
_ => Ok(Some(Token::Colon)), _ => Ok(Some(Token::Colon)),
} }
} }

View file

@ -333,3 +333,37 @@ fn test_duckdb_struct_literal() {
expr_from_projection(&select.projection[5]) expr_from_projection(&select.projection[5])
); );
} }
#[test]
fn test_duckdb_named_argument_function_with_assignment_operator() {
let sql = "SELECT FUN(a := '1', b := '2') FROM foo";
let select = duckdb_and_generic().verified_only_select(sql);
assert_eq!(
&Expr::Function(Function {
name: ObjectName(vec![Ident::new("FUN")]),
args: vec![
FunctionArg::Named {
name: Ident::new("a"),
arg: FunctionArgExpr::Expr(Expr::Value(Value::SingleQuotedString(
"1".to_owned()
))),
operator: FunctionArgOperator::Assignment
},
FunctionArg::Named {
name: Ident::new("b"),
arg: FunctionArgExpr::Expr(Expr::Value(Value::SingleQuotedString(
"2".to_owned()
))),
operator: FunctionArgOperator::Assignment
},
],
null_treatment: None,
filter: None,
over: None,
distinct: false,
special: false,
order_by: vec![],
}),
expr_from_projection(only(&select.projection))
);
}