Keep the quoting and span of a function argument name

Author: LucaCappelletti94

src/parser/mod.rs

@@ -6268,14 +6268,14 @@ impl<'a> Parser<'a> {
         }
 
         if let Some(next_data_type) = self.maybe_parse(parse_data_type_no_default)? {
-            let token = self.token_at(data_type_idx);
+            let token = self.token_at(data_type_idx).clone();
 
             // We ensure that the token is a `Word` token, and not other special tokens.
-            if !matches!(token.token, Token::Word(_)) {
-                return self.expected("a name or type", token.clone());
+            match token.token {
+                Token::Word(word) => name = Some(word.into_ident(token.span)),
+                _ => return self.expected("a name or type", token),
             }
 
-            name = Some(Ident::new(token.to_string()));
             data_type = next_data_type;
         }
 
@@ -6332,12 +6332,12 @@ impl<'a> Parser<'a> {
         }
 
         if let Some(next_data_type) = self.maybe_parse(parse_data_type_for_aggregate_arg)? {
-            let token = self.token_at(data_type_idx);
-            if !matches!(token.token, Token::Word(_)) {
-                return self.expected("a name or type", token.clone());
+            let token = self.token_at(data_type_idx).clone();
+            match token.token {
+                Token::Word(word) => name = Some(word.into_ident(token.span)),
+                _ => return self.expected("a name or type", token),
             }
 
-            name = Some(Ident::new(token.to_string()));
             data_type = next_data_type;
         }
 

tests/sqlparser_mssql.rs

@@ -2925,3 +2925,19 @@ fn parse_mssql_money_constants() {
         expr_from_projection(only(&select.projection)),
     );
 }
+
+#[test]
+fn parse_bracket_quoted_function_argument_name() {
+    let Statement::DropFunction(drop) = ms().verified_stmt("DROP FUNCTION f([Role] INT)") else {
+        panic!("expected a DROP FUNCTION statement");
+    };
+    assert_eq!(
+        drop.func_desc[0].args,
+        Some(vec![OperateFunctionArg {
+            mode: None,
+            name: Some(Ident::with_quote('[', "Role")),
+            data_type: Int(None),
+            default_expr: None,
+        }])
+    );
+}

tests/sqlparser_postgres.rs

@@ -25,8 +25,8 @@ mod test_utils;
 use helpers::attached_token::AttachedToken;
 use sqlparser::ast::*;
 use sqlparser::dialect::{Dialect, GenericDialect, MySqlDialect, PostgreSqlDialect, SQLiteDialect};
-use sqlparser::parser::ParserError;
-use sqlparser::tokenizer::Span;
+use sqlparser::parser::{Parser, ParserError};
+use sqlparser::tokenizer::{Location, Span};
 use test_utils::*;
 
 #[test]
@@ -9678,3 +9678,143 @@ fn parse_right_deep_join_chain() {
     // NATURAL JOIN followed by a constrained join must stay left-associative.
     pg().verified_stmt("SELECT * FROM t0 NATURAL JOIN t1 INNER JOIN t2 ON true");
 }
+
+#[test]
+fn parse_quoted_function_argument_name() {
+    let statement = pg_and_generic().verified_stmt(
+        r#"CREATE FUNCTION is_member("Role" TEXT) RETURNS BOOLEAN LANGUAGE SQL AS 'SELECT true'"#,
+    );
+    let Statement::CreateFunction(function) = statement else {
+        panic!("expected a CREATE FUNCTION statement");
+    };
+    assert_eq!(
+        function.args,
+        Some(vec![OperateFunctionArg {
+            mode: None,
+            name: Some(Ident::with_quote('"', "Role")),
+            data_type: DataType::Text,
+            default_expr: None,
+        }])
+    );
+
+    // An embedded quote is doubled in the input and belongs to the name once.
+    let statement = pg_and_generic().verified_stmt(
+        r#"CREATE FUNCTION f("we""ird" INT) RETURNS BOOLEAN LANGUAGE SQL RETURN true"#,
+    );
+    let Statement::CreateFunction(function) = statement else {
+        panic!("expected a CREATE FUNCTION statement");
+    };
+    assert_eq!(
+        function.args,
+        Some(vec![OperateFunctionArg {
+            mode: None,
+            name: Some(Ident::with_quote('"', r#"we"ird"#)),
+            data_type: DataType::Int(None),
+            default_expr: None,
+        }])
+    );
+
+    // A name outside ASCII is quoted for the same reason and survives the same way.
+    pg_and_generic().verified_stmt(
+        r#"CREATE FUNCTION f("Rôle" TEXT) RETURNS BOOLEAN LANGUAGE SQL RETURN true"#,
+    );
+}
+
+#[test]
+fn parse_quoted_function_argument_name_span() {
+    let sql =
+        r#"CREATE FUNCTION is_member("Role" TEXT) RETURNS BOOLEAN LANGUAGE SQL AS 'SELECT true'"#;
+    // Parsed directly rather than through the test helpers, which tokenize
+    // without locations.
+    let mut statements = Parser::parse_sql(&PostgreSqlDialect {}, sql).unwrap();
+    let Some(Statement::CreateFunction(function)) = statements.pop() else {
+        panic!("expected a CREATE FUNCTION statement");
+    };
+    let name = function.args.as_ref().unwrap()[0].name.as_ref().unwrap();
+    assert_eq!(
+        name.span,
+        Span::new(Location::new(1, 27), Location::new(1, 33)),
+        "the span covers the quoted name in the input"
+    );
+}
+
+#[test]
+fn parse_function_argument_modes_and_defaults_keep_quoted_names() {
+    let sql = r#"CREATE FUNCTION f(IN "A" INT = 1, OUT "B" TEXT, INOUT "C" BOOLEAN, VARIADIC "D" INT[]) RETURNS INT LANGUAGE SQL AS 'x'"#;
+    let Statement::CreateFunction(function) = pg_and_generic().verified_stmt(sql) else {
+        panic!("expected a CREATE FUNCTION statement");
+    };
+    assert_eq!(
+        function.args,
+        Some(vec![
+            OperateFunctionArg {
+                mode: Some(ArgMode::In),
+                name: Some(Ident::with_quote('"', "A")),
+                data_type: DataType::Int(None),
+                default_expr: Some(Expr::Value(
+                    Value::Number("1".parse().unwrap(), false).with_empty_span()
+                )),
+            },
+            OperateFunctionArg {
+                mode: Some(ArgMode::Out),
+                name: Some(Ident::with_quote('"', "B")),
+                data_type: DataType::Text,
+                default_expr: None,
+            },
+            OperateFunctionArg {
+                mode: Some(ArgMode::InOut),
+                name: Some(Ident::with_quote('"', "C")),
+                data_type: DataType::Boolean,
+                default_expr: None,
+            },
+            OperateFunctionArg {
+                mode: Some(ArgMode::Variadic),
+                name: Some(Ident::with_quote('"', "D")),
+                data_type: DataType::Array(ArrayElemTypeDef::SquareBracket(
+                    Box::new(DataType::Int(None)),
+                    None
+                )),
+                default_expr: None,
+            },
+        ])
+    );
+
+    // The `DEFAULT` spelling of the same argument list renders as `=`.
+    pg_and_generic().one_statement_parses_to(
+        r#"CREATE FUNCTION f("A" INT DEFAULT 1) RETURNS INT LANGUAGE SQL AS 'x'"#,
+        r#"CREATE FUNCTION f("A" INT = 1) RETURNS INT LANGUAGE SQL AS 'x'"#,
+    );
+}
+
+#[test]
+fn parse_quoted_argument_names_in_function_signatures() {
+    pg_and_generic().verified_stmt(r#"DROP FUNCTION f("Role" TEXT)"#);
+    pg_and_generic().verified_stmt(r#"DROP PROCEDURE p("Role" TEXT)"#);
+    pg_and_generic().verified_stmt(r#"ALTER FUNCTION f("Role" TEXT) RENAME TO g"#);
+    pg_and_generic().verified_stmt(r#"ALTER AGGREGATE my_agg("Role" TEXT) RENAME TO other_agg"#);
+
+    // An aggregate's `ORDER BY` arguments are parsed by the same reader as its
+    // direct ones.
+    let sql = r#"ALTER AGGREGATE my_agg("A" INT ORDER BY "B" INT) OWNER TO some_role"#;
+    let Statement::AlterFunction(alter) = pg_and_generic().verified_stmt(sql) else {
+        panic!("expected an ALTER AGGREGATE statement");
+    };
+    assert_eq!(
+        alter.function.args,
+        Some(vec![OperateFunctionArg {
+            mode: None,
+            name: Some(Ident::with_quote('"', "A")),
+            data_type: DataType::Int(None),
+            default_expr: None,
+        }])
+    );
+    assert_eq!(
+        alter.aggregate_order_by,
+        Some(vec![OperateFunctionArg {
+            mode: None,
+            name: Some(Ident::with_quote('"', "B")),
+            data_type: DataType::Int(None),
+            default_expr: None,
+        }])
+    );
+}