Resolved bug in parse_function_arg

Author: LucaCappelletti94

src/parser/mod.rs

@@ -5199,12 +5199,22 @@ impl<'a> Parser<'a> {
 
         // parse: [ argname ] argtype
         let mut name = None;
+        let next_token = self.peek_token();
         let mut data_type = self.parse_data_type()?;
-        if let DataType::Custom(n, _) = &data_type {
-            // the first token is actually a name
-            match n.0[0].clone() {
-                ObjectNamePart::Identifier(ident) => name = Some(ident),
-            }
+
+        // It may appear that the first token can be converted into a known
+        // type, but this could also be a collision as some types are only
+        // present in some dialects and therefore some type reserved keywords
+        // may be freely used as argument names in other dialects.
+
+        // To check whether the first token is a name or a type, we need to
+        // peek the next token, which if it is another type keyword, then the
+        // first token is a name and not a type in itself.
+        let potential_tokens = [Token::Eq, Token::RParen, Token::Comma];
+        if !self.peek_keyword(Keyword::DEFAULT)
+            && !potential_tokens.contains(&self.peek_token().token)
+        {
+            name = Some(Ident::new(next_token.to_string()));
             data_type = self.parse_data_type()?;
         }
 

tests/sqlparser_postgres.rs

@@ -4098,6 +4098,174 @@ fn parse_update_in_with_subquery() {
     pg_and_generic().verified_stmt(r#"WITH "result" AS (UPDATE "Hero" SET "name" = 'Captain America', "number_of_movies" = "number_of_movies" + 1 WHERE "secret_identity" = 'Sam Wilson' RETURNING "id", "name", "secret_identity", "number_of_movies") SELECT * FROM "result""#);
 }
 
+#[test]
+fn parser_create_function_with_args() {
+    let sql1 = r#"CREATE OR REPLACE FUNCTION check_strings_different(str1 VARCHAR, str2 VARCHAR) RETURNS BOOLEAN LANGUAGE plpgsql AS $$
+BEGIN
+    IF str1 <> str2 THEN
+        RETURN TRUE;
+    ELSE
+        RETURN FALSE;
+    END IF;
+END;
+$$"#;
+
+    assert_eq!(
+        pg_and_generic().verified_stmt(sql1),
+        Statement::CreateFunction(CreateFunction {
+            or_alter: false,
+            or_replace: true,
+            temporary: false,
+            name: ObjectName::from(vec![Ident::new("check_strings_different")]),
+            args: Some(vec![
+                OperateFunctionArg::with_name(
+                    "str1",
+                    DataType::Varchar(None),
+                ),
+                OperateFunctionArg::with_name(
+                    "str2",
+                    DataType::Varchar(None),
+                ),
+            ]),
+            return_type: Some(DataType::Boolean),
+            language: Some("plpgsql".into()),
+            behavior: None,
+            called_on_null: None,
+            parallel: None,
+            function_body: Some(CreateFunctionBody::AsBeforeOptions(Expr::Value(
+                (Value::DollarQuotedString(DollarQuotedString {value: "\nBEGIN\n    IF str1 <> str2 THEN\n        RETURN TRUE;\n    ELSE\n        RETURN FALSE;\n    END IF;\nEND;\n".to_owned(), tag: None})).with_empty_span()
+            ))),
+            if_not_exists: false,
+            using: None,
+            determinism_specifier: None,
+            options: None,
+            remote_connection: None,
+        })
+    );
+
+    let sql2 = r#"CREATE OR REPLACE FUNCTION check_not_zero(int1 INT) RETURNS BOOLEAN LANGUAGE plpgsql AS $$
+BEGIN
+    IF int1 <> 0 THEN
+        RETURN TRUE;
+    ELSE
+        RETURN FALSE;
+    END IF;
+END;
+$$"#;
+    assert_eq!(
+        pg_and_generic().verified_stmt(sql2),
+        Statement::CreateFunction(CreateFunction {
+            or_alter: false,
+            or_replace: true,
+            temporary: false,
+            name: ObjectName::from(vec![Ident::new("check_not_zero")]),
+            args: Some(vec![
+                OperateFunctionArg::with_name(
+                    "int1",
+                    DataType::Int(None)
+                )
+            ]),
+            return_type: Some(DataType::Boolean),
+            language: Some("plpgsql".into()),
+            behavior: None,
+            called_on_null: None,
+            parallel: None,
+            function_body: Some(CreateFunctionBody::AsBeforeOptions(Expr::Value(
+                (Value::DollarQuotedString(DollarQuotedString {value: "\nBEGIN\n    IF int1 <> 0 THEN\n        RETURN TRUE;\n    ELSE\n        RETURN FALSE;\n    END IF;\nEND;\n".to_owned(), tag: None})).with_empty_span()
+            ))),
+            if_not_exists: false,
+            using: None,
+            determinism_specifier: None,
+            options: None,
+            remote_connection: None,
+        })
+    );
+
+    let sql3 = r#"CREATE OR REPLACE FUNCTION check_values_different(a INT, b INT) RETURNS BOOLEAN LANGUAGE plpgsql AS $$
+BEGIN
+    IF a <> b THEN
+        RETURN TRUE;
+    ELSE
+        RETURN FALSE;
+    END IF;
+END;
+$$"#;
+    assert_eq!(
+        pg_and_generic().verified_stmt(sql3),
+        Statement::CreateFunction(CreateFunction {
+            or_alter: false,
+            or_replace: true,
+            temporary: false,
+            name: ObjectName::from(vec![Ident::new("check_values_different")]),
+            args: Some(vec![
+                OperateFunctionArg::with_name(
+                    "a",
+                    DataType::Int(None)
+                ),
+                OperateFunctionArg::with_name(
+                    "b",
+                    DataType::Int(None)
+                ),
+            ]),
+            return_type: Some(DataType::Boolean),
+            language: Some("plpgsql".into()),
+            behavior: None,
+            called_on_null: None,
+            parallel: None,
+            function_body: Some(CreateFunctionBody::AsBeforeOptions(Expr::Value(
+                (Value::DollarQuotedString(DollarQuotedString {value: "\nBEGIN\n    IF a <> b THEN\n        RETURN TRUE;\n    ELSE\n        RETURN FALSE;\n    END IF;\nEND;\n".to_owned(), tag: None})).with_empty_span()
+            ))),
+            if_not_exists: false,
+            using: None,
+            determinism_specifier: None,
+            options: None,
+            remote_connection: None,
+        })
+    );
+
+    let sql4 = r#"CREATE OR REPLACE FUNCTION check_values_different(int1 INT, int2 INT) RETURNS BOOLEAN LANGUAGE plpgsql AS $$
+BEGIN
+    IF int1 <> int2 THEN
+        RETURN TRUE;
+    ELSE
+        RETURN FALSE;
+    END IF;
+END;
+$$"#;
+    assert_eq!(
+        pg_and_generic().verified_stmt(sql4),
+        Statement::CreateFunction(CreateFunction {
+            or_alter: false,
+            or_replace: true,
+            temporary: false,
+            name: ObjectName::from(vec![Ident::new("check_values_different")]),
+            args: Some(vec![
+                OperateFunctionArg::with_name(
+                    "int1",
+                    DataType::Int(None)
+                ),
+                OperateFunctionArg::with_name(
+                    "int2",
+                    DataType::Int(None)
+                ),
+            ]),
+            return_type: Some(DataType::Boolean),
+            language: Some("plpgsql".into()),
+            behavior: None,
+            called_on_null: None,
+            parallel: None,
+            function_body: Some(CreateFunctionBody::AsBeforeOptions(Expr::Value(
+                (Value::DollarQuotedString(DollarQuotedString {value: "\nBEGIN\n    IF int1 <> int2 THEN\n        RETURN TRUE;\n    ELSE\n        RETURN FALSE;\n    END IF;\nEND;\n".to_owned(), tag: None})).with_empty_span()
+            ))),
+            if_not_exists: false,
+            using: None,
+            determinism_specifier: None,
+            options: None,
+            remote_connection: None,
+        })
+    );
+}
+
 #[test]
 fn parse_create_function() {
     let sql = "CREATE FUNCTION add(INTEGER, INTEGER) RETURNS INTEGER LANGUAGE SQL IMMUTABLE STRICT PARALLEL SAFE AS 'select $1 + $2;'";