Add support for granting privileges to procedures and functions in Snowflake (apache#1930)

Author: yoavcloud

src/ast/mod.rs

@@ -7052,6 +7052,25 @@ pub enum GrantObjects {
     ReplicationGroup(Vec<ObjectName>),
     /// Grant privileges on external volumes
     ExternalVolumes(Vec<ObjectName>),
+    /// Grant privileges on a procedure. In dialects that
+    /// support overloading, the argument types must be specified.
+    ///
+    /// For example:
+    /// `GRANT USAGE ON PROCEDURE foo(varchar) TO ROLE role1`
+    Procedure {
+        name: ObjectName,
+        arg_types: Vec<DataType>,
+    },
+
+    /// Grant privileges on a function. In dialects that
+    /// support overloading, the argument types must be specified.
+    ///
+    /// For example:
+    /// `GRANT USAGE ON FUNCTION foo(varchar) TO ROLE role1`
+    Function {
+        name: ObjectName,
+        arg_types: Vec<DataType>,
+    },
 }
 
 impl fmt::Display for GrantObjects {
@@ -7176,6 +7195,20 @@ impl fmt::Display for GrantObjects {
             GrantObjects::ExternalVolumes(objects) => {
                 write!(f, "EXTERNAL VOLUME {}", display_comma_separated(objects))
             }
+            GrantObjects::Procedure { name, arg_types } => {
+                write!(f, "PROCEDURE {name}")?;
+                if !arg_types.is_empty() {
+                    write!(f, "({})", display_comma_separated(arg_types))?;
+                }
+                Ok(())
+            }
+            GrantObjects::Function { name, arg_types } => {
+                write!(f, "FUNCTION {name}")?;
+                if !arg_types.is_empty() {
+                    write!(f, "({})", display_comma_separated(arg_types))?;
+                }
+                Ok(())
+            }
         }
     }
 }

src/parser/mod.rs

@@ -14093,6 +14093,8 @@ impl<'a> Parser<'a> {
                     Keyword::INTEGRATION,
                     Keyword::USER,
                     Keyword::CONNECTION,
+                    Keyword::PROCEDURE,
+                    Keyword::FUNCTION,
                 ]);
                 let objects =
                     self.parse_comma_separated(|p| p.parse_object_name_inner(false, true));
@@ -14105,6 +14107,13 @@ impl<'a> Parser<'a> {
                     Some(Keyword::VIEW) => Some(GrantObjects::Views(objects?)),
                     Some(Keyword::USER) => Some(GrantObjects::Users(objects?)),
                     Some(Keyword::CONNECTION) => Some(GrantObjects::Connections(objects?)),
+                    kw @ (Some(Keyword::PROCEDURE) | Some(Keyword::FUNCTION)) => {
+                        if let Some(name) = objects?.first() {
+                            self.parse_grant_procedure_or_function(name, &kw)?
+                        } else {
+                            self.expected("procedure or function name", self.peek_token())?
+                        }
+                    }
                     Some(Keyword::TABLE) | None => Some(GrantObjects::Tables(objects?)),
                     _ => unreachable!(),
                 }
@@ -14116,6 +14125,31 @@ impl<'a> Parser<'a> {
         Ok((privileges, objects))
     }
 
+    fn parse_grant_procedure_or_function(
+        &mut self,
+        name: &ObjectName,
+        kw: &Option<Keyword>,
+    ) -> Result<Option<GrantObjects>, ParserError> {
+        let arg_types = if self.consume_token(&Token::LParen) {
+            let list = self.parse_comma_separated0(Self::parse_data_type, Token::RParen)?;
+            self.expect_token(&Token::RParen)?;
+            list
+        } else {
+            vec![]
+        };
+        match kw {
+            Some(Keyword::PROCEDURE) => Ok(Some(GrantObjects::Procedure {
+                name: name.clone(),
+                arg_types,
+            })),
+            Some(Keyword::FUNCTION) => Ok(Some(GrantObjects::Function {
+                name: name.clone(),
+                arg_types,
+            })),
+            _ => self.expected("procedure or function keywords", self.peek_token())?,
+        }
+    }
+
     pub fn parse_grant_permission(&mut self) -> Result<Action, ParserError> {
         fn parse_columns(parser: &mut Parser) -> Result<Option<Vec<Ident>>, ParserError> {
             let columns = parser.parse_parenthesized_column_list(Optional, false)?;

tests/sqlparser_common.rs

@@ -9461,6 +9461,8 @@ fn parse_grant() {
     verified_stmt("GRANT SELECT ON FUTURE VIEWS IN SCHEMA db1.sc1 TO ROLE role1");
     verified_stmt("GRANT SELECT ON FUTURE MATERIALIZED VIEWS IN SCHEMA db1.sc1 TO ROLE role1");
     verified_stmt("GRANT SELECT ON FUTURE SEQUENCES IN SCHEMA db1.sc1 TO ROLE role1");
+    verified_stmt("GRANT USAGE ON PROCEDURE db1.sc1.foo(INT) TO ROLE role1");
+    verified_stmt("GRANT USAGE ON FUNCTION db1.sc1.foo(INT) TO ROLE role1");
 }
 
 #[test]