Extend snowflake grant options support (apache#1794)

yoavcloud · ayman-sigma · commit 2abde6db3c2f · 2025-04-10T14:47:57.000-07:00
diff --git a/src/ast/mod.rs b/src/ast/mod.rs
@@ -6097,10 +6097,10 @@ pub enum Action {
     ManageReleases,
     ManageVersions,
     Modify {
-        modify_type: ActionModifyType,
+        modify_type: Option<ActionModifyType>,
     },
     Monitor {
-        monitor_type: ActionMonitorType,
+        monitor_type: Option<ActionMonitorType>,
     },
     Operate,
     OverrideShareRestrictions,
@@ -6133,7 +6133,7 @@ impl fmt::Display for Action {
         match self {
             Action::AddSearchOptimization => f.write_str("ADD SEARCH OPTIMIZATION")?,
             Action::Apply { apply_type } => write!(f, "APPLY {apply_type}")?,
-            Action::ApplyBudget => f.write_str("APPLY BUDGET")?,
+            Action::ApplyBudget => f.write_str("APPLYBUDGET")?,
             Action::AttachListing => f.write_str("ATTACH LISTING")?,
             Action::AttachPolicy => f.write_str("ATTACH POLICY")?,
             Action::Audit => f.write_str("AUDIT")?,
@@ -6161,8 +6161,18 @@ impl fmt::Display for Action {
             Action::Manage { manage_type } => write!(f, "MANAGE {manage_type}")?,
             Action::ManageReleases => f.write_str("MANAGE RELEASES")?,
             Action::ManageVersions => f.write_str("MANAGE VERSIONS")?,
-            Action::Modify { modify_type } => write!(f, "MODIFY {modify_type}")?,
-            Action::Monitor { monitor_type } => write!(f, "MONITOR {monitor_type}")?,
+            Action::Modify { modify_type } => {
+                write!(f, "MODIFY")?;
+                if let Some(modify_type) = modify_type {
+                    write!(f, " {modify_type}")?;
+                }
+            }
+            Action::Monitor { monitor_type } => {
+                write!(f, "MONITOR")?;
+                if let Some(monitor_type) = monitor_type {
+                    write!(f, " {monitor_type}")?
+                }
+            }
             Action::Operate => f.write_str("OPERATE")?,
             Action::OverrideShareRestrictions => f.write_str("OVERRIDE SHARE RESTRICTIONS")?,
             Action::Ownership => f.write_str("OWNERSHIP")?,
@@ -6480,6 +6490,20 @@ pub enum GrantObjects {
     Warehouses(Vec<ObjectName>),
     /// Grant privileges on specific integrations
     Integrations(Vec<ObjectName>),
+    /// Grant privileges on resource monitors
+    ResourceMonitors(Vec<ObjectName>),
+    /// Grant privileges on users
+    Users(Vec<ObjectName>),
+    /// Grant privileges on compute pools
+    ComputePools(Vec<ObjectName>),
+    /// Grant privileges on connections
+    Connections(Vec<ObjectName>),
+    /// Grant privileges on failover groups
+    FailoverGroup(Vec<ObjectName>),
+    /// Grant privileges on replication group
+    ReplicationGroup(Vec<ObjectName>),
+    /// Grant privileges on external volumes
+    ExternalVolumes(Vec<ObjectName>),
 }
 
 impl fmt::Display for GrantObjects {
@@ -6520,6 +6544,27 @@ impl fmt::Display for GrantObjects {
                     display_comma_separated(schemas)
                 )
             }
+            GrantObjects::ResourceMonitors(objects) => {
+                write!(f, "RESOURCE MONITOR {}", display_comma_separated(objects))
+            }
+            GrantObjects::Users(objects) => {
+                write!(f, "USER {}", display_comma_separated(objects))
+            }
+            GrantObjects::ComputePools(objects) => {
+                write!(f, "COMPUTE POOL {}", display_comma_separated(objects))
+            }
+            GrantObjects::Connections(objects) => {
+                write!(f, "CONNECTION {}", display_comma_separated(objects))
+            }
+            GrantObjects::FailoverGroup(objects) => {
+                write!(f, "FAILOVER GROUP {}", display_comma_separated(objects))
+            }
+            GrantObjects::ReplicationGroup(objects) => {
+                write!(f, "REPLICATION GROUP {}", display_comma_separated(objects))
+            }
+            GrantObjects::ExternalVolumes(objects) => {
+                write!(f, "EXTERNAL VOLUME {}", display_comma_separated(objects))
+            }
         }
     }
 }
diff --git a/src/keywords.rs b/src/keywords.rs
@@ -738,6 +738,7 @@ define_keywords!(
     REPLICATION,
     RESET,
     RESOLVE,
+    RESOURCE,
     RESPECT,
     RESTART,
     RESTRICT,
diff --git a/src/parser/mod.rs b/src/parser/mod.rs
@@ -12885,6 +12885,26 @@ impl<'a> Parser<'a> {
                 Some(GrantObjects::AllSequencesInSchema {
                     schemas: self.parse_comma_separated(|p| p.parse_object_name(false))?,
                 })
+            } else if self.parse_keywords(&[Keyword::RESOURCE, Keyword::MONITOR]) {
+                Some(GrantObjects::ResourceMonitors(self.parse_comma_separated(
+                    |p| p.parse_object_name_with_wildcards(false, true),
+                )?))
+            } else if self.parse_keywords(&[Keyword::COMPUTE, Keyword::POOL]) {
+                Some(GrantObjects::ComputePools(self.parse_comma_separated(
+                    |p| p.parse_object_name_with_wildcards(false, true),
+                )?))
+            } else if self.parse_keywords(&[Keyword::FAILOVER, Keyword::GROUP]) {
+                Some(GrantObjects::FailoverGroup(self.parse_comma_separated(
+                    |p| p.parse_object_name_with_wildcards(false, true),
+                )?))
+            } else if self.parse_keywords(&[Keyword::REPLICATION, Keyword::GROUP]) {
+                Some(GrantObjects::ReplicationGroup(self.parse_comma_separated(
+                    |p| p.parse_object_name_with_wildcards(false, true),
+                )?))
+            } else if self.parse_keywords(&[Keyword::EXTERNAL, Keyword::VOLUME]) {
+                Some(GrantObjects::ExternalVolumes(self.parse_comma_separated(
+                    |p| p.parse_object_name_with_wildcards(false, true),
+                )?))
             } else {
                 let object_type = self.parse_one_of_keywords(&[
                     Keyword::SEQUENCE,
@@ -12898,6 +12918,8 @@ impl<'a> Parser<'a> {
                     Keyword::VIEW,
                     Keyword::WAREHOUSE,
                     Keyword::INTEGRATION,
+                    Keyword::USER,
+                    Keyword::CONNECTION,
                 ]);
                 let objects =
                     self.parse_comma_separated(|p| p.parse_object_name_with_wildcards(false, true));
@@ -12908,6 +12930,8 @@ impl<'a> Parser<'a> {
                     Some(Keyword::WAREHOUSE) => Some(GrantObjects::Warehouses(objects?)),
                     Some(Keyword::INTEGRATION) => Some(GrantObjects::Integrations(objects?)),
                     Some(Keyword::VIEW) => Some(GrantObjects::Views(objects?)),
+                    Some(Keyword::USER) => Some(GrantObjects::Users(objects?)),
+                    Some(Keyword::CONNECTION) => Some(GrantObjects::Connections(objects?)),
                     Some(Keyword::TABLE) | None => Some(GrantObjects::Tables(objects?)),
                     _ => unreachable!(),
                 }
@@ -12993,10 +13017,10 @@ impl<'a> Parser<'a> {
             let manage_type = self.parse_action_manage_type()?;
             Ok(Action::Manage { manage_type })
         } else if self.parse_keyword(Keyword::MODIFY) {
-            let modify_type = self.parse_action_modify_type()?;
+            let modify_type = self.parse_action_modify_type();
             Ok(Action::Modify { modify_type })
         } else if self.parse_keyword(Keyword::MONITOR) {
-            let monitor_type = self.parse_action_monitor_type()?;
+            let monitor_type = self.parse_action_monitor_type();
             Ok(Action::Monitor { monitor_type })
         } else if self.parse_keyword(Keyword::OPERATE) {
             Ok(Action::Operate)
@@ -13137,29 +13161,29 @@ impl<'a> Parser<'a> {
         }
     }
 
-    fn parse_action_modify_type(&mut self) -> Result<ActionModifyType, ParserError> {
+    fn parse_action_modify_type(&mut self) -> Option<ActionModifyType> {
         if self.parse_keywords(&[Keyword::LOG, Keyword::LEVEL]) {
-            Ok(ActionModifyType::LogLevel)
+            Some(ActionModifyType::LogLevel)
         } else if self.parse_keywords(&[Keyword::TRACE, Keyword::LEVEL]) {
-            Ok(ActionModifyType::TraceLevel)
+            Some(ActionModifyType::TraceLevel)
         } else if self.parse_keywords(&[Keyword::SESSION, Keyword::LOG, Keyword::LEVEL]) {
-            Ok(ActionModifyType::SessionLogLevel)
+            Some(ActionModifyType::SessionLogLevel)
         } else if self.parse_keywords(&[Keyword::SESSION, Keyword::TRACE, Keyword::LEVEL]) {
-            Ok(ActionModifyType::SessionTraceLevel)
+            Some(ActionModifyType::SessionTraceLevel)
         } else {
-            self.expected("GRANT MODIFY type", self.peek_token())
+            None
         }
     }
 
-    fn parse_action_monitor_type(&mut self) -> Result<ActionMonitorType, ParserError> {
+    fn parse_action_monitor_type(&mut self) -> Option<ActionMonitorType> {
         if self.parse_keyword(Keyword::EXECUTION) {
-            Ok(ActionMonitorType::Execution)
+            Some(ActionMonitorType::Execution)
         } else if self.parse_keyword(Keyword::SECURITY) {
-            Ok(ActionMonitorType::Security)
+            Some(ActionMonitorType::Security)
         } else if self.parse_keyword(Keyword::USAGE) {
-            Ok(ActionMonitorType::Usage)
+            Some(ActionMonitorType::Usage)
         } else {
-            self.expected("GRANT MONITOR type", self.peek_token())
+            None
         }
     }
 
diff --git a/tests/sqlparser_snowflake.rs b/tests/sqlparser_snowflake.rs
@@ -3364,7 +3364,7 @@ fn test_timetravel_at_before() {
 }
 
 #[test]
-fn test_grant_account_privileges() {
+fn test_grant_account_global_privileges() {
     let privileges = vec![
         "ALL",
         "ALL PRIVILEGES",
@@ -3469,6 +3469,43 @@ fn test_grant_account_privileges() {
     }
 }
 
+#[test]
+fn test_grant_account_object_privileges() {
+    let privileges = vec![
+        "ALL",
+        "ALL PRIVILEGES",
+        "APPLYBUDGET",
+        "MODIFY",
+        "MONITOR",
+        "USAGE",
+        "OPERATE",
+    ];
+
+    let objects_types = vec![
+        "USER",
+        "RESOURCE MONITOR",
+        "WAREHOUSE",
+        "COMPUTE POOL",
+        "DATABASE",
+        "INTEGRATION",
+        "CONNECTION",
+        "FAILOVER GROUP",
+        "REPLICATION GROUP",
+        "EXTERNAL VOLUME",
+    ];
+
+    let with_grant_options = vec!["", " WITH GRANT OPTION"];
+
+    for t in &objects_types {
+        for p in &privileges {
+            for wgo in &with_grant_options {
+                let sql = format!("GRANT {p} ON {t} obj1 TO ROLE role1{wgo}");
+                snowflake_and_generic().verified_stmt(&sql);
+            }
+        }
+    }
+}
+
 #[test]
 fn test_grant_role_to() {
     snowflake_and_generic().verified_stmt("GRANT ROLE r1 TO ROLE r2");
