PostgreSQL: Parse ROUTINE and ALL PROCEDURES | ROUTINES IN SCHEMA grant targets

fuzz/fuzz_seeds/postgres_2761_grant_on_routine_table.sql

@@ -0,0 +1 @@
+GRANT SELECT ON routine, function.t TO r

src/ast/mod.rs

@@ -7779,6 +7779,16 @@ pub enum GrantObjects {
         /// The target schema names.
         schemas: Vec<ObjectName>,
     },
+    /// Grant privileges on `ALL PROCEDURES IN SCHEMA <schema_name> [, ...]`
+    AllProceduresInSchema {
+        /// The target schema names.
+        schemas: Vec<ObjectName>,
+    },
+    /// Grant privileges on `ALL ROUTINES IN SCHEMA <schema_name> [, ...]`
+    AllRoutinesInSchema {
+        /// The target schema names.
+        schemas: Vec<ObjectName>,
+    },
     /// Grant privileges on `FUTURE SCHEMAS IN DATABASE <database_name> [, ...]`
     FutureSchemasInDatabase {
         /// The target database names.
@@ -7847,6 +7857,13 @@ pub enum GrantObjects {
     /// For example:
     /// `GRANT EXECUTE ON FUNCTION foo(p_id uuid), bar() TO role1`
     Functions(Vec<FunctionDesc>),
+    /// Grant privileges on routines, which name functions and procedures alike.
+    ///
+    /// For example:
+    /// `GRANT EXECUTE ON ROUTINE foo(p_id uuid), bar() TO role1`
+    ///
+    /// [PostgreSQL](https://www.postgresql.org/docs/current/sql-grant.html)
+    Routines(Vec<FunctionDesc>),
 }
 
 impl fmt::Display for GrantObjects {
@@ -7915,6 +7932,20 @@ impl fmt::Display for GrantObjects {
                     display_comma_separated(schemas)
                 )
             }
+            GrantObjects::AllProceduresInSchema { schemas } => {
+                write!(
+                    f,
+                    "ALL PROCEDURES IN SCHEMA {}",
+                    display_comma_separated(schemas)
+                )
+            }
+            GrantObjects::AllRoutinesInSchema { schemas } => {
+                write!(
+                    f,
+                    "ALL ROUTINES IN SCHEMA {}",
+                    display_comma_separated(schemas)
+                )
+            }
             GrantObjects::FutureSchemasInDatabase { databases } => {
                 write!(
                     f,
@@ -7984,6 +8015,9 @@ impl fmt::Display for GrantObjects {
             GrantObjects::Functions(functions) => {
                 write!(f, "FUNCTION {}", display_comma_separated(functions))
             }
+            GrantObjects::Routines(routines) => {
+                write!(f, "ROUTINE {}", display_comma_separated(routines))
+            }
         }
     }
 }

src/keywords.rs

@@ -845,6 +845,7 @@ define_keywords!(
     PRIVILEGES,
     PROC,
     PROCEDURE,
+    PROCEDURES,
     PROCESSLIST,
     PROFILE,
     PROGRAM,
@@ -934,6 +935,8 @@ define_keywords!(
     ROLLBACK,
     ROLLUP,
     ROOT,
+    ROUTINE,
+    ROUTINES,
     ROW,
     ROWGROUPSIZE,
     ROWID,

src/parser/mod.rs

@@ -18666,6 +18666,20 @@ impl<'a> Parser<'a> {
         Ok(values)
     }
 
+    /// Whether the keyword at the next token introduces a grant object type. A table named after
+    /// the keyword is followed by `TO`, `FROM`, `,` or `.`, as in `GRANT SELECT ON function TO r`.
+    fn peek_grant_object_type(&self) -> bool {
+        !matches!(
+            &self.peek_nth_token_ref(1).token,
+            Token::Comma
+                | Token::Period
+                | Token::Word(Word {
+                    keyword: Keyword::TO | Keyword::FROM,
+                    ..
+                })
+        )
+    }
+
     /// Parse privileges and optional target objects for GRANT/DENY/REVOKE statements.
     pub fn parse_grant_deny_revoke_privileges_objects(
         &mut self,
@@ -18722,6 +18736,24 @@ impl<'a> Parser<'a> {
                 Some(GrantObjects::AllFunctionsInSchema {
                     schemas: self.parse_comma_separated(|p| p.parse_object_name(false))?,
                 })
+            } else if self.parse_keywords(&[
+                Keyword::ALL,
+                Keyword::PROCEDURES,
+                Keyword::IN,
+                Keyword::SCHEMA,
+            ]) {
+                Some(GrantObjects::AllProceduresInSchema {
+                    schemas: self.parse_comma_separated(|p| p.parse_object_name(false))?,
+                })
+            } else if self.parse_keywords(&[
+                Keyword::ALL,
+                Keyword::ROUTINES,
+                Keyword::IN,
+                Keyword::SCHEMA,
+            ]) {
+                Some(GrantObjects::AllRoutinesInSchema {
+                    schemas: self.parse_comma_separated(|p| p.parse_object_name(false))?,
+                })
             } else if self.parse_keywords(&[
                 Keyword::FUTURE,
                 Keyword::SCHEMAS,
@@ -18807,29 +18839,37 @@ impl<'a> Parser<'a> {
                 Some(GrantObjects::ExternalVolumes(
                     self.parse_comma_separated(|p| p.parse_object_name(false))?,
                 ))
-            } else if self.parse_keyword(Keyword::FUNCTION) {
+            } else if self.peek_grant_object_type() && self.parse_keyword(Keyword::FUNCTION) {
                 Some(GrantObjects::Functions(
                     self.parse_comma_separated(Parser::parse_function_desc)?,
                 ))
-            } else if self.parse_keyword(Keyword::PROCEDURE) {
+            } else if self.peek_grant_object_type() && self.parse_keyword(Keyword::PROCEDURE) {
                 Some(GrantObjects::Procedures(
                     self.parse_comma_separated(Parser::parse_function_desc)?,
                 ))
+            } else if self.peek_grant_object_type() && self.parse_keyword(Keyword::ROUTINE) {
+                Some(GrantObjects::Routines(
+                    self.parse_comma_separated(Parser::parse_function_desc)?,
+                ))
             } else {
-                let object_type = self.parse_one_of_keywords(&[
-                    Keyword::SEQUENCE,
-                    Keyword::DATABASE,
-                    Keyword::SCHEMA,
-                    Keyword::TABLE,
-                    Keyword::VIEW,
-                    Keyword::WAREHOUSE,
-                    Keyword::INTEGRATION,
-                    Keyword::VIEW,
-                    Keyword::WAREHOUSE,
-                    Keyword::INTEGRATION,
-                    Keyword::USER,
-                    Keyword::CONNECTION,
-                ]);
+                let object_type = if self.peek_grant_object_type() {
+                    self.parse_one_of_keywords(&[
+                        Keyword::SEQUENCE,
+                        Keyword::DATABASE,
+                        Keyword::SCHEMA,
+                        Keyword::TABLE,
+                        Keyword::VIEW,
+                        Keyword::WAREHOUSE,
+                        Keyword::INTEGRATION,
+                        Keyword::VIEW,
+                        Keyword::WAREHOUSE,
+                        Keyword::INTEGRATION,
+                        Keyword::USER,
+                        Keyword::CONNECTION,
+                    ])
+                } else {
+                    None
+                };
                 let objects =
                     self.parse_comma_separated(|p| p.parse_object_name_inner(false, true));
                 match object_type {

tests/sqlparser_common.rs

@@ -22744,3 +22744,107 @@ fn display_create_procedure_parameter_list_propagates_write_errors() {
     let mut writer = RejectingWriter(" (");
     assert_eq!(write!(writer, "{statement}"), Err(fmt::Error));
 }
+
+#[test]
+fn parse_grant_routines() {
+    let objects = |sql: &str| match verified_stmt(sql) {
+        Statement::Grant(grant) => grant.objects,
+        Statement::Revoke(revoke) => revoke.objects,
+        stmt => panic!("unexpected statement: {stmt:?}"),
+    };
+    let name = |name: &str| ObjectName::from(vec![Ident::new(name)]);
+
+    assert_eq!(
+        objects("GRANT EXECUTE ON ROUTINE f(INT) TO bob"),
+        Some(GrantObjects::Routines(vec![FunctionDesc {
+            name: name("f"),
+            args: Some(vec![OperateFunctionArg::unnamed(DataType::Int(None))]),
+        }]))
+    );
+    assert_eq!(
+        objects("REVOKE ALL ON ROUTINE f, g() FROM bob"),
+        Some(GrantObjects::Routines(vec![
+            FunctionDesc {
+                name: name("f"),
+                args: None,
+            },
+            FunctionDesc {
+                name: name("g"),
+                args: Some(vec![]),
+            },
+        ]))
+    );
+    assert_eq!(
+        objects("GRANT EXECUTE ON ALL PROCEDURES IN SCHEMA s TO bob"),
+        Some(GrantObjects::AllProceduresInSchema {
+            schemas: vec![name("s")],
+        })
+    );
+    assert_eq!(
+        objects("REVOKE EXECUTE ON ALL ROUTINES IN SCHEMA s, t FROM bob"),
+        Some(GrantObjects::AllRoutinesInSchema {
+            schemas: vec![name("s"), name("t")],
+        })
+    );
+
+    verified_stmt("REVOKE ALL ON ROUTINE public.f(p_id UUID), g() FROM bob");
+    verified_stmt("REVOKE EXECUTE ON ALL PROCEDURES IN SCHEMA s, t FROM bob");
+    verified_stmt("GRANT EXECUTE ON ALL ROUTINES IN SCHEMA s TO bob");
+    assert_eq!(
+        objects("GRANT EXECUTE ON ROUTINE routine TO bob"),
+        Some(GrantObjects::Routines(vec![FunctionDesc {
+            name: name("routine"),
+            args: None,
+        }]))
+    );
+
+    for (sql, err) in [
+        ("GRANT EXECUTE ON ROUTINE f(INT", "Expected: ), found: EOF"),
+        (
+            "GRANT EXECUTE ON ALL PROCEDURES IN SCHEMA",
+            "Expected: identifier, found: EOF",
+        ),
+        (
+            "GRANT EXECUTE ON ALL ROUTINES IN SCHEMA",
+            "Expected: identifier, found: EOF",
+        ),
+    ] {
+        assert_eq!(
+            all_dialects().parse_sql_statements(sql),
+            Err(ParserError::ParserError(err.to_string())),
+            "{sql}"
+        );
+    }
+}
+
+#[test]
+fn parse_grant_on_tables_named_after_object_types() {
+    for name in [
+        "function",
+        "procedure",
+        "routine",
+        "sequence",
+        "database",
+        "schema",
+        "view",
+        "warehouse",
+        "integration",
+        "user",
+        "connection",
+    ] {
+        let sql = format!("GRANT SELECT ON {name} TO r");
+        match verified_stmt(&sql) {
+            Statement::Grant(grant) => assert_eq!(
+                grant.objects,
+                Some(GrantObjects::Tables(vec![ObjectName::from(vec![
+                    Ident::new(name)
+                ])])),
+                "{sql}"
+            ),
+            stmt => panic!("unexpected statement: {stmt:?}"),
+        }
+        one_statement_parses_to(&format!("GRANT SELECT ON TABLE {name} TO r"), &sql);
+        verified_stmt(&format!("REVOKE SELECT ON {name}, t FROM r"));
+        verified_stmt(&format!("GRANT SELECT ON {name}.t TO r"));
+    }
+}