feat(query): show grants support where/limit option (#15607)

* feat(query): support table_function show_grants 1. remove show grants plan, bind show grants stmt with rewrite query 2. support show grants on object <object_name> 3. add rewritekind::showgrants * drop show grants interpreter * show grants add show option * add test * add check in mget_table/db, if meta_len not eq table_names_len, should return err * modify setting enable_experimental_rbac_check comment
databendlabs · May 24, 2024 · 0604f10 · 0604f10
1 parent 56f0f17
commit 0604f10
Show file tree

Hide file tree

Showing 32 changed files with 1,140 additions and 428 deletions.
diff --git a/src/meta/api/src/schema_api_impl.rs b/src/meta/api/src/schema_api_impl.rs
@@ -190,13 +190,17 @@ use databend_common_meta_kvapi::kvapi;
 use databend_common_meta_kvapi::kvapi::DirName;
 use databend_common_meta_kvapi::kvapi::Key;
 use databend_common_meta_kvapi::kvapi::UpsertKVReq;
+use databend_common_meta_types::anyerror::AnyError;
 use databend_common_meta_types::protobuf as pb;
 use databend_common_meta_types::txn_op::Request;
 use databend_common_meta_types::txn_op_response::Response;
 use databend_common_meta_types::ConditionResult;
 use databend_common_meta_types::InvalidReply;
 use databend_common_meta_types::MatchSeq;
 use databend_common_meta_types::MatchSeqExt;
+use databend_common_meta_types::MetaAPIError;
+use databend_common_meta_types::MetaDataError;
+use databend_common_meta_types::MetaDataReadError;
 use databend_common_meta_types::MetaError;
 use databend_common_meta_types::MetaId;
 use databend_common_meta_types::MetaNetworkError;
@@ -2282,7 +2286,7 @@ impl<KV: kvapi::KVApi<Error = MetaError> + ?Sized> SchemaApi for KV {
 
         // Batch get all table-name by id
         let seq_names = self.mget_kv(&id_name_kv_keys).await?;
-        let mut table_names = Vec::with_capacity(id_name_kv_keys.len());
+        let mut table_names = Vec::with_capacity(table_ids.len());
 
         // None means table_name not found, maybe immutable table id. Ignore it
         for seq_name in seq_names.into_iter().flatten() {
@@ -2297,6 +2301,15 @@ impl<KV: kvapi::KVApi<Error = MetaError> + ?Sized> SchemaApi for KV {
         }
 
         let seq_metas = self.mget_kv(&meta_kv_keys).await?;
+        if seq_metas.len() != table_names.len() {
+            return Err(KVAppError::MetaError(MetaError::APIError(
+                MetaAPIError::DataError(MetaDataError::ReadError(MetaDataReadError::new(
+                    "mget_table_names_by_ids",
+                    "",
+                    &AnyError::error("The system is experiencing high load, please retry later"),
+                ))),
+            )));
+        }
         for (i, seq_meta_opt) in seq_metas.iter().enumerate() {
             if let Some(seq_meta) = seq_meta_opt {
                 let table_meta: TableMeta = deserialize_struct(&seq_meta.data)?;
@@ -2346,7 +2359,8 @@ impl<KV: kvapi::KVApi<Error = MetaError> + ?Sized> SchemaApi for KV {
 
         // Batch get all table-name by id
         let seq_names = self.mget_kv(&kv_keys).await?;
-        let mut db_names = Vec::with_capacity(kv_keys.len());
+        // If multi drop/create db the capacity may not same
+        let mut db_names = Vec::with_capacity(db_ids.len());
 
         // None means db_name not found, maybe immutable database id. Ignore it
         for seq_name in seq_names.into_iter().flatten() {
@@ -2361,6 +2375,15 @@ impl<KV: kvapi::KVApi<Error = MetaError> + ?Sized> SchemaApi for KV {
         }
 
         let seq_metas = self.mget_kv(&meta_kv_keys).await?;
+        if seq_metas.len() != db_names.len() {
+            return Err(KVAppError::MetaError(MetaError::APIError(
+                MetaAPIError::DataError(MetaDataError::ReadError(MetaDataReadError::new(
+                    "mget_table_names_by_ids",
+                    "",
+                    &AnyError::error("The system is experiencing high load, please retry later"),
+                ))),
+            )));
+        }
         for (i, seq_meta_opt) in seq_metas.iter().enumerate() {
             if let Some(seq_meta) = seq_meta_opt {
                 let db_meta: DatabaseMeta = deserialize_struct(&seq_meta.data)?;

diff --git a/src/query/ast/src/ast/format/ast_format.rs b/src/query/ast/src/ast/format/ast_format.rs
@@ -2236,7 +2236,11 @@ impl<'ast> Visitor<'ast> for AstFormatVisitor {
         self.children.push(node);
     }
 
-    fn visit_show_grant(&mut self, principal: &'ast Option<PrincipalIdentity>) {
+    fn visit_show_grant(
+        &mut self,
+        principal: &'ast Option<PrincipalIdentity>,
+        show_options: &'ast Option<ShowOptions>,
+    ) {
         let mut children = Vec::new();
         if let Some(principal) = &principal {
             let principal_name = match principal {
@@ -2246,6 +2250,18 @@ impl<'ast> Visitor<'ast> for AstFormatVisitor {
             let principal_format_ctx = AstFormatContext::new(principal_name);
             children.push(FormatTreeNode::new(principal_format_ctx));
         }
+        if let Some(show_options) = show_options {
+            if let Some(show_limit) = &show_options.show_limit {
+                self.visit_show_limit(show_limit);
+                children.push(self.children.pop().unwrap());
+            }
+            if let Some(limit) = show_options.limit {
+                let name = format!("Limit {}", limit);
+                let limit_format_ctx = AstFormatContext::new(name);
+                let node = FormatTreeNode::new(limit_format_ctx);
+                children.push(node);
+            }
+        }
         let name = "ShowGrant".to_string();
         let format_ctx = AstFormatContext::with_children(name, children.len());
         let node = FormatTreeNode::with_children(format_ctx, children);

diff --git a/src/query/ast/src/ast/statements/statement.rs b/src/query/ast/src/ast/statements/statement.rs
@@ -202,7 +202,9 @@ pub enum Statement {
     Grant(GrantStmt),
     ShowGrants {
         principal: Option<PrincipalIdentity>,
+        show_options: Option<ShowOptions>,
     },
+    ShowObjectPrivileges(ShowObjectPrivilegesStmt),
     Revoke(RevokeStmt),
 
     // UDF
@@ -605,13 +607,20 @@ impl Display for Statement {
                 write!(f, " '{role}'")?;
             }
             Statement::Grant(stmt) => write!(f, "{stmt}")?,
-            Statement::ShowGrants { principal } => {
+            Statement::ShowGrants {
+                principal,
+                show_options,
+            } => {
                 write!(f, "SHOW GRANTS")?;
                 if let Some(principal) = principal {
                     write!(f, " FOR")?;
                     write!(f, "{principal}")?;
                 }
+                if let Some(show_options) = show_options {
+                    write!(f, " {show_options}")?;
+                }
             }
+            Statement::ShowObjectPrivileges(stmt) => write!(f, "{stmt}")?,
             Statement::Revoke(stmt) => write!(f, "{stmt}")?,
             Statement::CreateUDF(stmt) => write!(f, "{stmt}")?,
             Statement::DropUDF {

diff --git a/src/query/ast/src/ast/statements/user.rs b/src/query/ast/src/ast/statements/user.rs
@@ -22,6 +22,7 @@ use crate::ast::write_comma_separated_list;
 use crate::ast::AuthType;
 use crate::ast::CreateOption;
 use crate::ast::PrincipalIdentity;
+use crate::ast::ShowOptions;
 use crate::ast::UserIdentity;
 use crate::ast::UserPrivilegeType;
 
@@ -135,6 +136,50 @@ impl Display for RevokeStmt {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
+pub struct ShowObjectPrivilegesStmt {
+    pub object: GrantObjectName,
+    pub show_option: Option<ShowOptions>,
+}
+
+#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
+pub enum GrantObjectName {
+    Database(#[drive(skip)] String),
+    Table(#[drive(skip)] Option<String>, #[drive(skip)] String),
+    UDF(#[drive(skip)] String),
+    Stage(#[drive(skip)] String),
+}
+
+impl Display for GrantObjectName {
+    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
+        match self {
+            GrantObjectName::Database(database_name) => {
+                write!(f, "DATABASE {database_name}")
+            }
+            GrantObjectName::Table(database_name, table_name) => {
+                if let Some(database_name) = database_name {
+                    write!(f, "TABLE {database_name}.{table_name}")
+                } else {
+                    write!(f, "TABLE {table_name}")
+                }
+            }
+            GrantObjectName::UDF(udf) => write!(f, " UDF {udf}"),
+            GrantObjectName::Stage(stage) => write!(f, " STAGE {stage}"),
+        }
+    }
+}
+
+impl Display for ShowObjectPrivilegesStmt {
+    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
+        write!(f, "SHOW GRANTS ON {}", self.object)?;
+
+        if let Some(show_option) = &self.show_option {
+            write!(f, " {show_option}")?;
+        }
+        Ok(())
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Drive, DriveMut)]
 pub enum AccountMgrSource {
     Role {
@@ -158,48 +203,12 @@ impl Display for AccountMgrSource {
                 write!(f, " ")?;
                 write_comma_separated_list(f, privileges.iter().map(|p| p.to_string()))?;
                 write!(f, " ON")?;
-                match level {
-                    AccountMgrLevel::Global => write!(f, " *.*")?,
-                    AccountMgrLevel::Database(database_name) => {
-                        if let Some(database_name) = database_name {
-                            write!(f, " {database_name}.*")?;
-                        } else {
-                            write!(f, " *")?;
-                        }
-                    }
-                    AccountMgrLevel::Table(database_name, table_name) => {
-                        if let Some(database_name) = database_name {
-                            write!(f, " {database_name}.{table_name}")?;
-                        } else {
-                            write!(f, " {table_name}")?;
-                        }
-                    }
-                    AccountMgrLevel::UDF(udf) => write!(f, " UDF {udf}")?,
-                    AccountMgrLevel::Stage(stage) => write!(f, " STAGE {stage}")?,
-                }
+                write!(f, " {}", level)?;
             }
             AccountMgrSource::ALL { level, .. } => {
                 write!(f, " ALL PRIVILEGES")?;
                 write!(f, " ON")?;
-                match level {
-                    AccountMgrLevel::Global => write!(f, " *.*")?,
-                    AccountMgrLevel::Database(database_name) => {
-                        if let Some(database_name) = database_name {
-                            write!(f, " {database_name}.*")?;
-                        } else {
-                            write!(f, " *")?;
-                        }
-                    }
-                    AccountMgrLevel::Table(database_name, table_name) => {
-                        if let Some(database_name) = database_name {
-                            write!(f, " {database_name}.{table_name}")?;
-                        } else {
-                            write!(f, " {table_name}")?;
-                        }
-                    }
-                    AccountMgrLevel::UDF(udf) => write!(f, " UDF {udf}")?,
-                    AccountMgrLevel::Stage(stage) => write!(f, " STAGE {stage}")?,
-                }
+                write!(f, " {}", level)?;
             }
         }
         Ok(())
@@ -215,6 +224,30 @@ pub enum AccountMgrLevel {
     Stage(#[drive(skip)] String),
 }
 
+impl Display for AccountMgrLevel {
+    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
+        match self {
+            AccountMgrLevel::Global => write!(f, " *.*"),
+            AccountMgrLevel::Database(database_name) => {
+                if let Some(database_name) = database_name {
+                    write!(f, " {database_name}.*")
+                } else {
+                    write!(f, " *")
+                }
+            }
+            AccountMgrLevel::Table(database_name, table_name) => {
+                if let Some(database_name) = database_name {
+                    write!(f, " {database_name}.{table_name}")
+                } else {
+                    write!(f, " {table_name}")
+                }
+            }
+            AccountMgrLevel::UDF(udf) => write!(f, " UDF {udf}"),
+            AccountMgrLevel::Stage(stage) => write!(f, " STAGE {stage}"),
+        }
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Drive, DriveMut)]
 pub enum SecondaryRolesOption {
     None,

diff --git a/src/query/ast/src/ast/visitors/visitor.rs b/src/query/ast/src/ast/visitors/visitor.rs
@@ -611,7 +611,14 @@ pub trait Visitor<'ast>: Sized {
 
     fn visit_grant(&mut self, _grant: &'ast GrantStmt) {}
 
-    fn visit_show_grant(&mut self, _principal: &'ast Option<PrincipalIdentity>) {}
+    fn visit_show_grant(
+        &mut self,
+        _principal: &'ast Option<PrincipalIdentity>,
+        _show_options: &'ast Option<ShowOptions>,
+    ) {
+    }
+
+    fn visit_show_object_priv(&mut self, _show: &'ast ShowObjectPrivilegesStmt) {}
 
     fn visit_revoke(&mut self, _revoke: &'ast RevokeStmt) {}
 

diff --git a/src/query/ast/src/ast/visitors/visitor_mut.rs b/src/query/ast/src/ast/visitors/visitor_mut.rs
@@ -623,7 +623,14 @@ pub trait VisitorMut: Sized {
 
     fn visit_grant(&mut self, _grant: &mut GrantStmt) {}
 
-    fn visit_show_grant(&mut self, _principal: &mut Option<PrincipalIdentity>) {}
+    fn visit_show_grant(
+        &mut self,
+        _principal: &mut Option<PrincipalIdentity>,
+        _show_options: &mut Option<ShowOptions>,
+    ) {
+    }
+
+    fn visit_show_object_priv(&mut self, _show: &mut ShowObjectPrivilegesStmt) {}
 
     fn visit_revoke(&mut self, _revoke: &mut RevokeStmt) {}
 

diff --git a/src/query/ast/src/ast/visitors/walk.rs b/src/query/ast/src/ast/visitors/walk.rs
@@ -503,7 +503,11 @@ pub fn walk_statement<'a, V: Visitor<'a>>(visitor: &mut V, statement: &'a Statem
             role_name,
         } => visitor.visit_drop_role(*if_exists, role_name),
         Statement::Grant(stmt) => visitor.visit_grant(stmt),
-        Statement::ShowGrants { principal } => visitor.visit_show_grant(principal),
+        Statement::ShowGrants {
+            principal,
+            show_options,
+        } => visitor.visit_show_grant(principal, show_options),
+        Statement::ShowObjectPrivileges(stmt) => visitor.visit_show_object_priv(stmt),
         Statement::Revoke(stmt) => visitor.visit_revoke(stmt),
         Statement::CreateUDF(stmt) => visitor.visit_create_udf(stmt),
         Statement::DropUDF {

diff --git a/src/query/ast/src/ast/visitors/walk_mut.rs b/src/query/ast/src/ast/visitors/walk_mut.rs
@@ -498,7 +498,11 @@ pub fn walk_statement_mut<V: VisitorMut>(visitor: &mut V, statement: &mut Statem
             role_name,
         } => visitor.visit_drop_role(*if_exists, role_name),
         Statement::Grant(stmt) => visitor.visit_grant(stmt),
-        Statement::ShowGrants { principal } => visitor.visit_show_grant(principal),
+        Statement::ShowGrants {
+            principal,
+            show_options,
+        } => visitor.visit_show_grant(principal, show_options),
+        Statement::ShowObjectPrivileges(stmt) => visitor.visit_show_object_priv(stmt),
         Statement::Revoke(stmt) => visitor.visit_revoke(stmt),
         Statement::CreateUDF(stmt) => visitor.visit_create_udf(stmt),
         Statement::DropUDF {