refactor: Add RoleIdent to define kvapi::Key to access roles in meta-service

src/meta/app/src/key_with_tenant.rs

@@ -0,0 +1,37 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Define behaviors of a `kvapi::Key` that contains a tenant.
+
+use databend_common_meta_kvapi::kvapi;
+
+use crate::tenant::Tenant;
+
+/// Define behaviors of a `kvapi::Key` that contains a tenant.
+pub trait KeyWithTenant: kvapi::Key {
+    /// Return the tenant this key belongs to.
+    fn tenant(&self) -> &Tenant;
+
+    /// Return a encoded key prefix for listing keys of this kind that belong to the tenant.
+    ///
+    /// It is in form of `<__PREFIX>/<tenant>/`.
+    /// The trailing `/` is important for exclude tenants with prefix same as this tenant.
+    fn tenant_prefix(&self) -> String {
+        kvapi::KeyBuilder::new_prefixed(Self::PREFIX)
+            .push_str(self.tenant().name())
+            // Add trailing "/"
+            .push_raw("")
+            .done()
+    }
+}

src/meta/app/src/lib.rs

@@ -30,3 +30,6 @@ pub mod schema;
 pub mod share;
 pub mod storage;
 pub mod tenant;
+
+mod key_with_tenant;
+pub use key_with_tenant::KeyWithTenant;

src/meta/app/src/principal/mod.rs

@@ -20,6 +20,7 @@ mod network_policy;
 mod ownership_info;
 mod password_policy;
 mod principal_identity;
+mod role_ident;
 mod role_info;
 mod user_auth;
 mod user_defined_file_format;
@@ -38,6 +39,7 @@ pub use network_policy::NetworkPolicy;
 pub use ownership_info::OwnershipInfo;
 pub use password_policy::PasswordPolicy;
 pub use principal_identity::PrincipalIdentity;
+pub use role_ident::RoleIdent;
 pub use role_info::RoleInfo;
 pub use role_info::RoleInfoSerdeError;
 pub use user_auth::AuthInfo;

src/meta/app/src/principal/role_ident.rs

@@ -0,0 +1,103 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use crate::tenant::Tenant;
+
+/// The identifier of a role.
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct RoleIdent {
+    tenant: Tenant,
+    role_name: String,
+}
+
+impl RoleIdent {
+    pub fn new(tenant: Tenant, role_name: impl ToString) -> RoleIdent {
+        RoleIdent {
+            tenant,
+            role_name: role_name.to_string(),
+        }
+    }
+
+    pub fn role_name(&self) -> &str {
+        &self.role_name
+    }
+}
+
+mod kvapi_key_impl {
+    use databend_common_meta_kvapi::kvapi;
+
+    use crate::principal::role_ident::RoleIdent;
+    use crate::principal::RoleInfo;
+    use crate::tenant::Tenant;
+    use crate::KeyWithTenant;
+
+    impl kvapi::Key for RoleIdent {
+        const PREFIX: &'static str = "__fd_roles";
+        type ValueType = RoleInfo;
+
+        fn parent(&self) -> Option<String> {
+            Some(self.tenant.to_string_key())
+        }
+
+        fn to_string_key(&self) -> String {
+            kvapi::KeyBuilder::new_prefixed(Self::PREFIX)
+                .push_str(self.tenant.name())
+                .push_str(&self.role_name)
+                .done()
+        }
+
+        fn from_str_key(s: &str) -> Result<Self, kvapi::KeyError> {
+            let mut p = kvapi::KeyParser::new_prefixed(s, Self::PREFIX)?;
+            let tenant = p.next_str()?;
+            let role_name = p.next_str()?;
+            p.done()?;
+
+            Ok(RoleIdent::new(Tenant::new(tenant), role_name))
+        }
+    }
+
+    impl kvapi::Value for RoleInfo {
+        fn dependency_keys(&self) -> impl IntoIterator<Item = String> {
+            []
+        }
+    }
+
+    impl KeyWithTenant for RoleIdent {
+        fn tenant(&self) -> &Tenant {
+            &self.tenant
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use databend_common_meta_kvapi::kvapi::Key;
+
+    use crate::principal::role_ident::RoleIdent;
+    use crate::tenant::Tenant;
+    use crate::KeyWithTenant;
+
+    #[test]
+    fn test_role_ident_tenant_prefix() {
+        let r = RoleIdent::new(Tenant::new("tenant"), "role");
+        assert_eq!("__fd_roles/tenant/", r.tenant_prefix());
+    }
+
+    #[test]
+    fn test_role_ident_key() {
+        let r = RoleIdent::new(Tenant::new("tenant"), "role");
+        assert_eq!("__fd_roles/tenant/role", r.to_string_key());
+        assert_eq!(Ok(r), RoleIdent::from_str_key("__fd_roles/tenant/role"));
+    }
+}

src/meta/app/src/principal/user_grant.rs

@@ -17,7 +17,6 @@ use std::fmt;
 use std::ops;
 
 use databend_common_meta_kvapi::kvapi;
-use databend_common_meta_kvapi::kvapi::Key;
 use enumflags2::BitFlags;
 
 use crate::principal::UserPrivilegeSet;
@@ -193,20 +192,6 @@ impl TenantOwnershipObject {
         TenantOwnershipObject { tenant, object }
     }
 
-    /// Return a encoded key prefix for listing keys belongs to the tenant.
-    ///
-    /// It is in form of `__fd_object_owners/<tenant>/`.
-    /// The trailing `/` is important for exclude tenants with prefix same as this tenant.
-    pub fn tenant_prefix(&self) -> String {
-        kvapi::KeyBuilder::new_prefixed(Self::PREFIX)
-            .push_str(&self.tenant.tenant)
-            .done()
-    }
-
-    pub fn tenant(&self) -> &Tenant {
-        &self.tenant
-    }
-
     // Because the encoded string key does not contain all of the field:
     // Catalog and db_id is missing, so we can not rebuild the complete key from string key.
     // Thus it is not allow to access this field.
@@ -491,6 +476,7 @@ mod kvapi_key_impl {
     use crate::principal::OwnershipInfo;
     use crate::principal::OwnershipObject;
     use crate::tenant::Tenant;
+    use crate::KeyWithTenant;
 
     impl kvapi::Key for TenantOwnershipObject {
         const PREFIX: &'static str = "__fd_object_owners";
@@ -524,6 +510,12 @@ mod kvapi_key_impl {
             []
         }
     }
+
+    impl KeyWithTenant for TenantOwnershipObject {
+        fn tenant(&self) -> &Tenant {
+            &self.tenant
+        }
+    }
 }
 
 #[cfg(test)]
@@ -533,6 +525,17 @@ mod tests {
     use crate::principal::user_grant::TenantOwnershipObject;
     use crate::principal::OwnershipObject;
     use crate::tenant::Tenant;
+    use crate::KeyWithTenant;
+
+    #[test]
+    fn test_tenant_ownership_object_tenant_prefix() {
+        let r = TenantOwnershipObject::new(Tenant::new("tenant"), OwnershipObject::Database {
+            catalog_name: "cat".to_string(),
+            db_id: 1,
+        });
+
+        assert_eq!("__fd_object_owners/tenant/", r.tenant_prefix());
+    }
 
     #[test]
     fn test_role_grantee_as_kvapi_key() {

src/meta/app/src/tenant/tenant.rs

@@ -27,6 +27,10 @@ impl Tenant {
             tenant: tenant.to_string(),
         }
     }
+
+    pub fn name(&self) -> &str {
+        &self.tenant
+    }
 }
 
 mod kvapi_key_impl {

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

@@ -23,10 +23,12 @@ use databend_common_meta_app::app_error::TxnRetryMaxTimes;
 use databend_common_meta_app::principal::GrantObject;
 use databend_common_meta_app::principal::OwnershipInfo;
 use databend_common_meta_app::principal::OwnershipObject;
+use databend_common_meta_app::principal::RoleIdent;
 use databend_common_meta_app::principal::RoleInfo;
 use databend_common_meta_app::principal::TenantOwnershipObject;
 use databend_common_meta_app::principal::UserPrivilegeType;
 use databend_common_meta_app::tenant::Tenant;
+use databend_common_meta_app::KeyWithTenant;
 use databend_common_meta_kvapi::kvapi;
 use databend_common_meta_kvapi::kvapi::Key;
 use databend_common_meta_kvapi::kvapi::UpsertKVReply;
@@ -46,28 +48,23 @@ use crate::role::role_api::RoleApi;
 use crate::serde::check_and_upgrade_to_pb;
 use crate::serialize_struct;
 
-static ROLE_API_KEY_PREFIX: &str = "__fd_roles";
-
 static TXN_MAX_RETRY_TIMES: u32 = 5;
 
 static BUILTIN_ROLE_ACCOUNT_ADMIN: &str = "account_admin";
 
 pub struct RoleMgr {
     kv_api: Arc<dyn kvapi::KVApi<Error = MetaError> + Send + Sync>,
     tenant: NonEmptyString,
-    role_prefix: String,
 }
 
 impl RoleMgr {
     pub fn create(
         kv_api: Arc<dyn kvapi::KVApi<Error = MetaError> + Send + Sync>,
         tenant: NonEmptyStr,
     ) -> Self {
-        let t = tenant.get();
         RoleMgr {
             kv_api,
             tenant: tenant.into(),
-            role_prefix: format!("{}/{}", ROLE_API_KEY_PREFIX, t),
         }
     }
 
@@ -77,7 +74,7 @@ impl RoleMgr {
         role_info: &RoleInfo,
         seq: MatchSeq,
     ) -> Result<u64, ErrorCode> {
-        let key = self.make_role_key(role_info.identity());
+        let key = self.role_key(role_info.identity());
         let value = serialize_struct(role_info, ErrorCode::IllegalUserInfoFormat, || "")?;
 
         let res = self
@@ -100,8 +97,7 @@ impl RoleMgr {
         value: Vec<u8>,
         seq: MatchSeq,
     ) -> Result<UpsertKVReply, MetaError> {
-        let kv_api = self.kv_api.clone();
-        kv_api
+        self.kv_api
             .upsert_kv(UpsertKVReq::new(&key, seq, Operation::Update(value), None))
             .await
     }
@@ -123,8 +119,14 @@ impl RoleMgr {
         grantee.tenant_prefix()
     }
 
-    fn make_role_key(&self, role: &str) -> String {
-        format!("{}/{}", self.role_prefix, role)
+    fn role_key(&self, role: &str) -> String {
+        let r = RoleIdent::new(Tenant::new(self.tenant.as_str()), role.to_string());
+        r.to_string_key()
+    }
+
+    fn role_prefix(&self) -> String {
+        let r = RoleIdent::new(Tenant::new(self.tenant.as_str()), "dummy".to_string());
+        r.tenant_prefix()
     }
 }
 
@@ -134,7 +136,7 @@ impl RoleApi for RoleMgr {
     #[minitrace::trace]
     async fn add_role(&self, role_info: RoleInfo) -> databend_common_exception::Result<u64> {
         let match_seq = MatchSeq::Exact(0);
-        let key = self.make_role_key(role_info.identity());
+        let key = self.role_key(role_info.identity());
         let value = serialize_struct(&role_info, ErrorCode::IllegalUserInfoFormat, || "")?;
 
         let upsert_kv = self.kv_api.upsert_kv(UpsertKVReq::new(
@@ -154,7 +156,7 @@ impl RoleApi for RoleMgr {
     #[async_backtrace::framed]
     #[minitrace::trace]
     async fn get_role(&self, role: &String, seq: MatchSeq) -> Result<SeqV<RoleInfo>, ErrorCode> {
-        let key = self.make_role_key(role);
+        let key = self.role_key(role);
         let res = self.kv_api.get_kv(&key).await?;
         let seq_value =
             res.ok_or_else(|| ErrorCode::UnknownRole(format!("Role '{}' does not exist.", role)))?;
@@ -183,7 +185,7 @@ impl RoleApi for RoleMgr {
     #[async_backtrace::framed]
     #[minitrace::trace]
     async fn get_roles(&self) -> Result<Vec<SeqV<RoleInfo>>, ErrorCode> {
-        let role_prefix = self.role_prefix.clone();
+        let role_prefix = self.role_prefix();
         let values = self.kv_api.prefix_list_kv(role_prefix.as_str()).await?;
 
         let mut r = vec![];
@@ -207,8 +209,10 @@ impl RoleApi for RoleMgr {
     #[minitrace::trace]
     async fn get_ownerships(&self) -> Result<Vec<SeqV<OwnershipInfo>>, ErrorCode> {
         let object_owner_prefix = self.ownership_object_prefix();
-        let kv_api = self.kv_api.clone();
-        let values = kv_api.prefix_list_kv(object_owner_prefix.as_str()).await?;
+        let values = self
+            .kv_api
+            .prefix_list_kv(object_owner_prefix.as_str())
+            .await?;
 
         let mut r = vec![];
         for (key, val) in values {
@@ -284,7 +288,7 @@ impl RoleApi for RoleMgr {
         if let Some(old_role) = old_role {
             // BUILTIN role or Dropped role may get err, no need to revoke
             if let Ok(seqv) = self.get_role(&old_role.to_owned(), MatchSeq::GE(1)).await {
-                let old_key = self.make_role_key(&old_role);
+                let old_key = self.role_key(&old_role);
                 let old_seq = seqv.seq;
                 let mut old_role_info = seqv.data;
                 old_role_info.grants.revoke_privileges(
@@ -301,7 +305,7 @@ impl RoleApi for RoleMgr {
 
         // account_admin has all privilege, no need to grant ownership.
         if new_role != BUILTIN_ROLE_ACCOUNT_ADMIN {
-            let new_key = self.make_role_key(new_role);
+            let new_key = self.role_key(new_role);
             let SeqV {
                 seq: new_seq,
                 data: mut new_role_info,
@@ -383,7 +387,7 @@ impl RoleApi for RoleMgr {
 
         if let Some(role) = role {
             if let Ok(seqv) = self.get_role(&role.to_owned(), MatchSeq::GE(1)).await {
-                let old_key = self.make_role_key(&role);
+                let old_key = self.role_key(&role);
                 let grant_object = convert_to_grant_obj(object);
                 let old_seq = seqv.seq;
                 let mut old_role_info = seqv.data;
@@ -425,7 +429,7 @@ impl RoleApi for RoleMgr {
     #[async_backtrace::framed]
     #[minitrace::trace]
     async fn drop_role(&self, role: String, seq: MatchSeq) -> Result<(), ErrorCode> {
-        let key = self.make_role_key(&role);
+        let key = self.role_key(&role);
 
         let res = self
             .kv_api