refactor: use NonEmptyString to access UDFMgr. (#14844)

* refactor: use NonEmptyString to access UDFMgr.

No need to return TenantIsEmpty error anymore

* fixup: tolerate time diff in timeout test

Author: drmingdrmer

Cargo.lock

@@ -25,6 +25,7 @@ databend-common-expression = { path = "../query/expression" }
 databend-common-license = { path = "../common/license" }
 databend-common-meta-app = { path = "../meta/app" }
 databend-common-meta-embedded = { path = "../meta/embedded" }
+databend-common-meta-types = { path = "../meta/types" }
 databend-common-users = { path = "../query/users" }
 databend-query = { path = "../query/service", features = [
     "simd",

src/bendpy/Cargo.toml

@@ -19,6 +19,7 @@ use databend_common_exception::Result;
 use databend_common_meta_app::principal::GrantObject;
 use databend_common_meta_app::principal::UserInfo;
 use databend_common_meta_app::principal::UserPrivilegeSet;
+use databend_common_meta_types::NonEmptyString;
 use databend_common_users::UserApiProvider;
 use databend_query::sessions::QueryContext;
 use databend_query::sessions::Session;
@@ -55,12 +56,14 @@ impl PySessionContext {
                 uuid::Uuid::new_v4().to_string()
             };
 
+            let tenant = NonEmptyString::new(tenant).unwrap();
+
             let config = GlobalConfig::instance();
             UserApiProvider::try_create_simple(config.meta.to_meta_grpc_client_conf(), &tenant)
                 .await
                 .unwrap();
 
-            session.set_current_tenant(tenant.to_owned());
+            session.set_current_tenant(tenant.to_string());
 
             let mut user = UserInfo::new_no_auth("root", "%");
             user.grants.grant_privileges(

src/bendpy/src/context.rs

@@ -305,7 +305,9 @@ impl kvapi::TestSuite {
                 assert_eq!(v2.seq, 3);
                 assert_eq!(v2.data, b("v2"));
                 let v2_meta = v2.meta.unwrap();
-                assert_eq!(v2_meta.get_expire_at_ms().unwrap() / 1000, now_sec + 10);
+                let expire_at_sec = v2_meta.get_expire_at_ms().unwrap() / 1000;
+                let want = now_sec + 10;
+                assert!((want..want + 2).contains(&expire_at_sec));
             }
         }
 

src/meta/kvapi/src/kvapi/test_suite.rs

@@ -103,7 +103,7 @@ impl GlobalServices {
         UserApiProvider::init(
             config.meta.to_meta_grpc_client_conf(),
             config.query.idm.clone(),
-            config.query.tenant_id.as_str(),
+            &config.query.tenant_id,
             config.query.tenant_quota.clone(),
         )
         .await?;

src/query/service/src/global_services.rs

@@ -26,6 +26,7 @@ use databend_common_meta_app::principal::StageType;
 use databend_common_meta_app::principal::UserGrantSet;
 use databend_common_meta_app::principal::UserPrivilegeSet;
 use databend_common_meta_app::principal::UserPrivilegeType;
+use databend_common_meta_types::NonEmptyString;
 use databend_common_sql::optimizer::get_udf_names;
 use databend_common_sql::plans::InsertInputSource;
 use databend_common_sql::plans::PresignAction;
@@ -507,7 +508,7 @@ impl AccessChecker for PrivilegeAccess {
                             ObjectId::Table(db_id, table_id) => { (db_id, Some(table_id)) }
                             ObjectId::Database(db_id) => { (db_id, None) }
                         };
-                        let has_priv = has_priv(tenant.as_str(), database, None, db_id, table_id, grant_set).await?;
+                        let has_priv = has_priv(&tenant, database, None, db_id, table_id, grant_set).await?;
                         return if has_priv {
                             Ok(())
                         } else {
@@ -526,7 +527,7 @@ impl AccessChecker for PrivilegeAccess {
                             ObjectId::Table(db_id, table_id) => { (db_id, Some(table_id)) }
                             ObjectId::Database(db_id) => { (db_id, None) }
                         };
-                        let has_priv = has_priv(tenant.as_str(), database, None, db_id, table_id, grant_set).await?;
+                        let has_priv = has_priv(&tenant, database, None, db_id, table_id, grant_set).await?;
                         return if has_priv {
                             Ok(())
                         } else {
@@ -545,7 +546,7 @@ impl AccessChecker for PrivilegeAccess {
                             ObjectId::Table(db_id, table_id) => { (db_id, Some(table_id)) }
                             ObjectId::Database(db_id) => { (db_id, None) }
                         };
-                        let has_priv = has_priv(tenant.as_str(), database, Some(table), db_id, table_id, grant_set).await?;
+                        let has_priv = has_priv(&tenant, database, Some(table), db_id, table_id, grant_set).await?;
                         return if has_priv {
                             Ok(())
                         } else {
@@ -628,7 +629,7 @@ impl AccessChecker for PrivilegeAccess {
                     ObjectId::Table(db_id, table_id) => { (db_id, Some(table_id)) }
                     ObjectId::Database(db_id) => { (db_id, None) }
                 };
-                let has_priv = has_priv(tenant.as_str(), &plan.database, None, db_id, None, grant_set).await?;
+                let has_priv = has_priv(&tenant, &plan.database, None, db_id, None, grant_set).await?;
 
                 return if has_priv {
                     Ok(())
@@ -1002,7 +1003,7 @@ impl AccessChecker for PrivilegeAccess {
 
 // TODO(liyz): replace it with verify_access
 async fn has_priv(
-    tenant: &str,
+    tenant: &NonEmptyString,
     db_name: &str,
     table_name: Option<&str>,
     db_id: u64,

src/query/service/src/interpreters/access/privilege_access.rs

@@ -80,10 +80,7 @@ pub async fn validate_grant_object_exists(
             }
         }
         GrantObject::UDF(udf) => {
-            if !UserApiProvider::instance()
-                .exists_udf(tenant.as_str(), udf)
-                .await?
-            {
+            if !UserApiProvider::instance().exists_udf(&tenant, udf).await? {
                 return Err(databend_common_exception::ErrorCode::UnknownStage(format!(
                     "udf {udf} not exists"
                 )));

src/query/service/src/interpreters/common/grant.rs

@@ -112,7 +112,7 @@ impl Interpreter for CreateDatabaseInterpreter {
         let quota_api = UserApiProvider::instance().get_tenant_quota_api_client(&tenant)?;
         let quota = quota_api.get_quota(MatchSeq::GE(0)).await?.data;
         let catalog = self.ctx.get_catalog(&self.plan.catalog).await?;
-        let databases = catalog.list_databases(&tenant).await?;
+        let databases = catalog.list_databases(tenant.as_str()).await?;
         if quota.max_databases != 0 && databases.len() >= quota.max_databases as usize {
             return Err(ErrorCode::TenantQuotaExceeded(format!(
                 "Max databases quota exceeded {}",
@@ -121,7 +121,7 @@ impl Interpreter for CreateDatabaseInterpreter {
         };
         // if create from other tenant, check from share endpoint
         if let Some(ref share_name) = self.plan.meta.from_share {
-            self.check_create_database_from_share(&tenant, share_name)
+            self.check_create_database_from_share(&tenant.to_string(), share_name)
                 .await?;
         }
 
@@ -130,7 +130,7 @@ impl Interpreter for CreateDatabaseInterpreter {
 
         // Grant ownership as the current role. The above create_db_req.meta.owner could be removed in
         // the future.
-        let role_api = UserApiProvider::instance().get_role_api_client(&tenant)?;
+        let role_api = UserApiProvider::instance().role_api(&tenant);
         if let Some(current_role) = self.ctx.get_current_role() {
             role_api
                 .grant_ownership(

src/query/service/src/interpreters/interpreter_database_create.rs

@@ -58,14 +58,14 @@ impl Interpreter for DropDatabaseInterpreter {
             .get_database(tenant.as_str(), &self.plan.database)
             .await;
         if let Ok(db) = db {
-            let role_api = UserApiProvider::instance().get_role_api_client(tenant.as_str())?;
+            let role_api = UserApiProvider::instance().role_api(&tenant);
             let owner_object = OwnershipObject::Database {
                 catalog_name: self.plan.catalog.clone(),
                 db_id: db.get_db_info().ident.db_id,
             };
 
             role_api.revoke_ownership(&owner_object).await?;
-            RoleCacheManager::instance().invalidate_cache(tenant.as_str());
+            RoleCacheManager::instance().invalidate_cache(&tenant);
         }
 
         // actual drop database

src/query/service/src/interpreters/interpreter_database_drop.rs

@@ -21,6 +21,7 @@ use databend_common_meta_app::principal::OwnershipObject;
 use databend_common_meta_app::principal::PrincipalIdentity;
 use databend_common_meta_app::principal::UserPrivilegeSet;
 use databend_common_meta_app::principal::UserPrivilegeType::Ownership;
+use databend_common_meta_types::NonEmptyString;
 use databend_common_sql::plans::GrantPrivilegePlan;
 use databend_common_users::RoleCacheManager;
 use databend_common_users::UserApiProvider;
@@ -111,7 +112,7 @@ impl GrantPrivilegeInterpreter {
     async fn grant_ownership(
         &self,
         ctx: &Arc<QueryContext>,
-        tenant: &str,
+        tenant: &NonEmptyString,
         owner_object: &OwnershipObject,
         new_role: &str,
     ) -> Result<()> {
@@ -199,7 +200,7 @@ impl Interpreter for GrantPrivilegeInterpreter {
                         .convert_to_ownerobject(tenant.as_str(), &plan.on, plan.on.catalog())
                         .await?;
                     if self.ctx.get_current_role().is_some() {
-                        self.grant_ownership(&self.ctx, tenant.as_str(), &owner_object, &role)
+                        self.grant_ownership(&self.ctx, &tenant, &owner_object, &role)
                             .await?;
                     } else {
                         return Err(databend_common_exception::ErrorCode::UnknownRole(
@@ -208,9 +209,9 @@ impl Interpreter for GrantPrivilegeInterpreter {
                     }
                 } else {
                     user_mgr
-                        .grant_privileges_to_role(tenant.as_str(), &role, plan.on, plan.priv_types)
+                        .grant_privileges_to_role(&tenant, &role, plan.on, plan.priv_types)
                         .await?;
-                    RoleCacheManager::instance().invalidate_cache(tenant.as_str());
+                    RoleCacheManager::instance().invalidate_cache(&tenant);
                 }
             }
         }