From 14b33ac493b9c2214487f3d389d4287e038f8fc0 Mon Sep 17 00:00:00 2001
From: gstvg <28798827+gstvg@users.noreply.github.com>
Date: Sat, 6 Apr 2024 13:46:36 -0300
Subject: [PATCH] Add support for DuckDB functions named arguments with
 assignment operator (#1195)

Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>
---
 src/ast/mod.rs            |  3 +++
 src/parser/mod.rs         | 17 +++++++++++++++--
 src/tokenizer.rs          |  8 ++++----
 tests/sqlparser_duckdb.rs | 34 ++++++++++++++++++++++++++++++++++
 4 files changed, 56 insertions(+), 6 deletions(-)

diff --git a/src/ast/mod.rs b/src/ast/mod.rs
index 7818dacd..a378b58b 100644
--- a/src/ast/mod.rs
+++ b/src/ast/mod.rs
@@ -4584,6 +4584,8 @@ pub enum FunctionArgOperator {
     Equals,
     /// function(arg1 => value1)
     RightArrow,
+    /// function(arg1 := value1)
+    Assignment,
 }
 
 impl fmt::Display for FunctionArgOperator {
@@ -4591,6 +4593,7 @@ impl fmt::Display for FunctionArgOperator {
         match self {
             FunctionArgOperator::Equals => f.write_str("="),
             FunctionArgOperator::RightArrow => f.write_str("=>"),
+            FunctionArgOperator::Assignment => f.write_str(":="),
         }
     }
 }
diff --git a/src/parser/mod.rs b/src/parser/mod.rs
index 2a5e9567..a3d7a7cf 100644
--- a/src/parser/mod.rs
+++ b/src/parser/mod.rs
@@ -3485,7 +3485,7 @@ impl<'a> Parser<'a> {
         let name = self.parse_identifier(false)?;
 
         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()?)
             } else {
                 None
@@ -4183,7 +4183,7 @@ impl<'a> Parser<'a> {
                 self.next_token(); // Skip `DEFAULT`
                 Some(DeclareAssignment::Default(Box::new(self.parse_expr()?)))
             }
-            Token::DuckAssignment => {
+            Token::Assignment => {
                 self.next_token(); // Skip `:=`
                 Some(DeclareAssignment::DuckAssignment(Box::new(
                     self.parse_expr()?,
@@ -8602,6 +8602,19 @@ impl<'a> Parser<'a> {
                 arg,
                 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 {
             Ok(FunctionArg::Unnamed(self.parse_wildcard_expr()?.into()))
         }
diff --git a/src/tokenizer.rs b/src/tokenizer.rs
index a1a2eae2..e31fccca 100644
--- a/src/tokenizer.rs
+++ b/src/tokenizer.rs
@@ -117,8 +117,8 @@ pub enum Token {
     Colon,
     /// DoubleColon `::` (used for casting in PostgreSQL)
     DoubleColon,
-    /// Assignment `:=` (used for keyword argument in DuckDB macros)
-    DuckAssignment,
+    /// Assignment `:=` (used for keyword argument in DuckDB macros and some functions, and for variable declarations in DuckDB and Snowflake)
+    Assignment,
     /// SemiColon `;` used as separator for COPY and payload
     SemiColon,
     /// Backslash `\` used in terminating the COPY payload with `\.`
@@ -239,7 +239,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::Assignment => f.write_str(":="),
             Token::SemiColon => f.write_str(";"),
             Token::Backslash => f.write_str("\\"),
             Token::LBracket => f.write_str("["),
@@ -959,7 +959,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),
+                        Some('=') => self.consume_and_return(chars, Token::Assignment),
                         _ => Ok(Some(Token::Colon)),
                     }
                 }
diff --git a/tests/sqlparser_duckdb.rs b/tests/sqlparser_duckdb.rs
index a29d4008..e41109d9 100644
--- a/tests/sqlparser_duckdb.rs
+++ b/tests/sqlparser_duckdb.rs
@@ -333,3 +333,37 @@ fn test_duckdb_struct_literal() {
         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))
+    );
+}