refactor: optimize list some OwnershipObjects (#18455)

Author: TCeason

src/meta/app/src/principal/tenant_ownership_object_ident.rs

@@ -105,6 +105,7 @@ mod kvapi_impl {
 
 #[cfg(test)]
 mod tests {
+    use databend_common_meta_kvapi::kvapi;
     use databend_common_meta_kvapi::kvapi::Key;
 
     use crate::principal::OwnershipObject;
@@ -272,6 +273,21 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_ownership_seq_list_key() {
+        use databend_common_meta_kvapi::kvapi::Key;
+        let obj = OwnershipObject::Sequence {
+            name: "seq1".to_string(),
+        };
+
+        let ident = TenantOwnershipObjectIdent::new(Tenant::new_literal("tenant1"), obj);
+        let dir_name = kvapi::DirName::new(ident);
+        assert_eq!(
+            dir_name.to_string_key(),
+            "__fd_object_owners/tenant1/sequence-by-name"
+        );
+    }
+
     #[test]
     fn test_tenant_ownership_object_with_key_space() {
         // TODO(xp): implement this test

src/query/catalog/src/table_context.rs

@@ -57,6 +57,7 @@ use databend_common_storage::StageFileInfo;
 use databend_common_storage::StageFilesInfo;
 use databend_common_storage::StorageMetrics;
 use databend_common_users::GrantObjectVisibilityChecker;
+use databend_common_users::Object;
 use databend_storages_common_session::SessionState;
 use databend_storages_common_session::TxnManagerRef;
 use databend_storages_common_table_meta::meta::Location;
@@ -234,6 +235,7 @@ pub trait TableContext: Send + Sync {
     async fn get_visibility_checker(
         &self,
         ignore_ownership: bool,
+        object: Object,
     ) -> Result<GrantObjectVisibilityChecker>;
     fn get_fuse_version(&self) -> String;
     fn get_format_settings(&self) -> Result<FormatSettings>;

src/query/management/src/role/role_api.rs

@@ -34,6 +34,12 @@ pub trait RoleApi: Sync + Send {
 
     async fn list_ownerships(&self) -> Result<Vec<SeqV<OwnershipInfo>>>;
 
+    async fn list_udf_ownerships(&self) -> Result<Vec<OwnershipInfo>>;
+    async fn list_stage_ownerships(&self) -> Result<Vec<OwnershipInfo>>;
+    async fn list_seq_ownerships(&self) -> Result<Vec<OwnershipInfo>>;
+    async fn list_connection_ownerships(&self) -> Result<Vec<OwnershipInfo>>;
+    async fn list_warehouse_ownerships(&self) -> Result<Vec<OwnershipInfo>>;
+
     /// General role update.
     ///
     /// It fetches the role that matches the specified seq number, update it in place, then write it back with the seq it sees.

src/query/management/src/role/role_mgr.rs

@@ -16,6 +16,7 @@ use std::sync::Arc;
 
 use databend_common_base::base::tokio::sync::Mutex;
 use databend_common_exception::ErrorCode;
+use databend_common_meta_api::kv_pb_api::KVPbApi;
 use databend_common_meta_api::reply::unpack_txn_reply;
 use databend_common_meta_api::txn_backoff::txn_backoff;
 use databend_common_meta_api::txn_cond_seq;
@@ -35,6 +36,7 @@ use databend_common_meta_app::KeyWithTenant;
 use databend_common_meta_cache::Cache;
 use databend_common_meta_client::ClientHandle;
 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::ListKVReply;
 use databend_common_meta_kvapi::kvapi::UpsertKVReply;
@@ -48,6 +50,7 @@ use databend_common_meta_types::TxnRequest;
 use databend_common_meta_types::UpsertKV;
 use enumflags2::make_bitflags;
 use fastrace::func_name;
+use futures::TryStreamExt;
 use log::debug;
 use log::error;
 use log::info;
@@ -298,6 +301,80 @@ impl RoleApi for RoleMgr {
         Ok(r)
     }
 
+    #[async_backtrace::framed]
+    #[fastrace::trace]
+    async fn list_udf_ownerships(&self) -> databend_common_exception::Result<Vec<OwnershipInfo>> {
+        let obj = OwnershipObject::UDF {
+            name: "foo".to_string(),
+        };
+
+        let ident = TenantOwnershipObjectIdent::new(self.tenant.clone(), obj);
+        let dir_name = DirName::new(ident);
+        let values = self.kv_api.list_pb_values(&dir_name).await?;
+        let udfs = values.try_collect().await?;
+        Ok(udfs)
+    }
+
+    #[async_backtrace::framed]
+    #[fastrace::trace]
+    async fn list_stage_ownerships(&self) -> databend_common_exception::Result<Vec<OwnershipInfo>> {
+        let obj = OwnershipObject::Stage {
+            name: "s1".to_string(),
+        };
+
+        let ident = TenantOwnershipObjectIdent::new(self.tenant.clone(), obj);
+        let dir_name = DirName::new(ident);
+        let values = self.kv_api.list_pb_values(&dir_name).await?;
+        let stages = values.try_collect().await?;
+        Ok(stages)
+    }
+
+    #[async_backtrace::framed]
+    #[fastrace::trace]
+    async fn list_seq_ownerships(&self) -> databend_common_exception::Result<Vec<OwnershipInfo>> {
+        let obj = OwnershipObject::Sequence {
+            name: "seq1".to_string(),
+        };
+
+        let ident = TenantOwnershipObjectIdent::new(self.tenant.clone(), obj);
+        let dir_name = DirName::new(ident);
+        let values = self.kv_api.list_pb_values(&dir_name).await?;
+        let seqs = values.try_collect().await?;
+        Ok(seqs)
+    }
+
+    #[async_backtrace::framed]
+    #[fastrace::trace]
+    async fn list_connection_ownerships(
+        &self,
+    ) -> databend_common_exception::Result<Vec<OwnershipInfo>> {
+        let obj = OwnershipObject::Connection {
+            name: "con".to_string(),
+        };
+
+        let ident = TenantOwnershipObjectIdent::new(self.tenant.clone(), obj);
+        let dir_name = DirName::new(ident);
+        let values = self.kv_api.list_pb_values(&dir_name).await?;
+        let conns = values.try_collect().await?;
+        Ok(conns)
+    }
+
+    #[async_backtrace::framed]
+    #[fastrace::trace]
+    async fn list_warehouse_ownerships(
+        &self,
+    ) -> databend_common_exception::Result<Vec<OwnershipInfo>> {
+        let obj = OwnershipObject::Warehouse {
+            id: "w".to_string(),
+        };
+
+        let ident = TenantOwnershipObjectIdent::new(self.tenant.clone(), obj);
+        let dir_name = DirName::new(ident);
+        let values = self.kv_api.list_pb_values(&dir_name).await?;
+        let ws = values.try_collect().await?;
+        Ok(ws)
+    }
+
     /// General role update.
     ///
     /// It fetch the role that matches the specified seq number, update it in place, then write it back with the seq it sees.

src/query/service/src/interpreters/interpreter_connection_show.rs

@@ -19,6 +19,7 @@ use databend_common_exception::Result;
 use databend_common_expression::types::StringType;
 use databend_common_expression::DataBlock;
 use databend_common_expression::FromData;
+use databend_common_users::Object;
 use databend_common_users::UserApiProvider;
 use log::debug;
 
@@ -64,7 +65,10 @@ impl Interpreter for ShowConnectionsInterpreter {
             .get_settings()
             .get_enable_experimental_connection_privilege_check()?
         {
-            let visibility_checker = self.ctx.get_visibility_checker(false).await?;
+            let visibility_checker = self
+                .ctx
+                .get_visibility_checker(false, Object::Connection)
+                .await?;
             formats.retain(|c| visibility_checker.check_connection_visibility(&c.name));
         }
 

src/query/service/src/interpreters/interpreter_show_warehouses.rs

@@ -24,6 +24,7 @@ use databend_common_expression::Scalar;
 use databend_common_license::license::Feature;
 use databend_common_license::license_manager::LicenseManagerSwitch;
 use databend_common_management::WarehouseInfo;
+use databend_common_users::Object;
 use databend_enterprise_resources_management::ResourcesManagement;
 
 use crate::interpreters::Interpreter;
@@ -63,7 +64,10 @@ impl Interpreter for ShowWarehousesInterpreter {
         let mut warehouses_status =
             ColumnBuilder::with_capacity(&DataType::String, warehouses.len());
 
-        let visibility_checker = self.ctx.get_visibility_checker(false).await?;
+        let visibility_checker = self
+            .ctx
+            .get_visibility_checker(false, Object::Warehouse)
+            .await?;
         for warehouse in warehouses {
             match warehouse {
                 WarehouseInfo::SelfManaged(name) => {

src/query/service/src/pipelines/processors/transforms/transform_async_function.rs

@@ -26,6 +26,7 @@ use databend_common_meta_app::schema::GetSequenceNextValueReq;
 use databend_common_meta_app::schema::SequenceIdent;
 use databend_common_pipeline_transforms::processors::AsyncTransform;
 use databend_common_storages_fuse::TableContext;
+use databend_common_users::Object;
 
 use crate::pipelines::processors::transforms::transform_dictionary::DictionaryOperator;
 use crate::sessions::QueryContext;
@@ -171,7 +172,7 @@ impl TransformAsyncFunction {
                             .get_settings()
                             .get_enable_experimental_sequence_privilege_check()?
                         {
-                            Some(ctx.get_visibility_checker(false).await?)
+                            Some(ctx.get_visibility_checker(false, Object::Sequence).await?)
                         } else {
                             None
                         };

src/query/service/src/servers/http/v1/catalog/get_database_table.rs

@@ -18,6 +18,7 @@ use databend_common_ast::parser::Dialect;
 use databend_common_catalog::catalog::CatalogManager;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
+use databend_common_users::Object;
 use poem::error::InternalServerError;
 use poem::error::NotFound;
 use poem::error::Result as PoemResult;
@@ -58,7 +59,10 @@ async fn handle(
     table: String,
 ) -> Result<GetDatabaseTableResponse> {
     let tenant = ctx.session.get_current_tenant();
-    let visibility_checker = ctx.session.get_visibility_checker(false).await?;
+    let visibility_checker = ctx
+        .session
+        .get_visibility_checker(false, Object::All)
+        .await?;
 
     let catalog = CatalogManager::instance().get_default_catalog(Default::default())?;
 

src/query/service/src/servers/http/v1/catalog/list_database_table_fields.rs

@@ -15,6 +15,7 @@
 use databend_common_catalog::catalog::CatalogManager;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
+use databend_common_users::Object;
 use poem::error::InternalServerError;
 use poem::error::NotFound;
 use poem::error::Result as PoemResult;
@@ -47,7 +48,10 @@ async fn handle(
     table: String,
 ) -> Result<ListDatabaseTableFieldsResponse> {
     let tenant = ctx.session.get_current_tenant();
-    let visibility_checker = ctx.session.get_visibility_checker(false).await?;
+    let visibility_checker = ctx
+        .session
+        .get_visibility_checker(false, Object::All)
+        .await?;
 
     let catalog = CatalogManager::instance().get_default_catalog(Default::default())?;
     let db = catalog.get_database(&tenant, &database).await?;

src/query/service/src/servers/http/v1/catalog/list_database_tables.rs

@@ -17,6 +17,7 @@ use chrono::Utc;
 use databend_common_catalog::catalog::CatalogManager;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
+use databend_common_users::Object;
 use poem::error::InternalServerError;
 use poem::error::NotFound;
 use poem::error::Result as PoemResult;
@@ -50,7 +51,10 @@ pub struct TableInfo {
 #[async_backtrace::framed]
 async fn handle(ctx: &HttpQueryContext, database: String) -> Result<ListDatabaseTablesResponse> {
     let tenant = ctx.session.get_current_tenant();
-    let visibility_checker = ctx.session.get_visibility_checker(false).await?;
+    let visibility_checker = ctx
+        .session
+        .get_visibility_checker(false, Object::All)
+        .await?;
 
     let catalog = CatalogManager::instance().get_default_catalog(Default::default())?;
     let db = catalog.get_database(&tenant, &database).await?;