refactor: Add RoleGrantee to provide meta-service key for RoleMgr (#1…

…4752)

Encapsulate the logic of encoding/decoding role-grantee key into
`RoleGrantee`. `RoleGrantee` is a standard `kvapi::Key` implementation
so that the export program could rely on the `kvapi::Key` API to
traverse meta-data to fetch all data that belongs to a tenant.

Legacy issue:

- For `OwnershipObject::Database{ catalog_name, db_id }`, `catalog_name`
  is not encoded into the key, we have to be compatible with this. The
  solution is to encode a `"default"` catalog with
  `database-by-id/<db-id>` prefix, and encode a non-default catalog with
  `database-by-catalog-id/<catalog>/<db-id>` prefix.

  The same for the `OwnershipObject::Table{ catalog_name, ... }`.

- For `OwnershipObject::Table{ catalog_name, db_id, table_id }`, `db_id`
  is not encoded into the key. Thus `kvapi::Key::from_str_key()` can not
  be implemented because `db_id` is absent.

- In this version, only `"default"` catalog is allowed. Using other
  catalog will panic. This restrict will be removed if it is confirmed
  there is no other catalog in our meta-data.

---

- Part of #14738

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

@@ -52,6 +52,7 @@ pub use user_defined_function::UserDefinedFunction;
 pub use user_grant::GrantEntry;
 pub use user_grant::GrantObject;
 pub use user_grant::OwnershipObject;
+pub use user_grant::TenantOwnershipObject;
 pub use user_grant::UserGrantSet;
 pub use user_identity::UserIdentity;
 pub use user_info::UserInfo;

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

@@ -16,10 +16,13 @@ use std::collections::HashSet;
 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;
 use crate::principal::UserPrivilegeType;
+use crate::tenant::Tenant;
 
 /// [`OwnershipObject`] is used to maintain the grant object that support rename by id. Using ID over name
 /// have many benefits, it can avoid lost privileges after the object get renamed.
@@ -48,6 +51,170 @@ pub enum OwnershipObject {
     },
 }
 
+impl OwnershipObject {
+    // # Legacy issue: Catalog is not encoded into key.
+    //
+    // But at that time, only `"default"` catalog is used.
+    // Thus we follow the same rule, if catalog is `"default"`, we don't encode it.
+    //
+    // This issue is introduced in https://github.com/drmingdrmer/databend/blob/7681763dc54306e55b5e0326af0510292d244be3/src/query/management/src/role/role_mgr.rs#L86
+    const DEFAULT_CATALOG: &'static str = "default";
+
+    /// Build key with the provided KeyBuilder as sub part of key used to access meta-service
+    pub(crate) fn build_key(&self, b: kvapi::KeyBuilder) -> kvapi::KeyBuilder {
+        match self {
+            OwnershipObject::Database {
+                catalog_name,
+                db_id,
+            } => {
+                // Legacy issue: Catalog is not encoded into key.
+                if catalog_name == Self::DEFAULT_CATALOG {
+                    b.push_raw("database-by-id").push_u64(*db_id)
+                } else {
+                    b.push_raw("database-by-catalog-id")
+                        .push_str(catalog_name)
+                        .push_u64(*db_id)
+                }
+            }
+            OwnershipObject::Table {
+                catalog_name,
+                db_id,
+                table_id,
+            } => {
+                // TODO(flaneur): db_id is not encoded into key. Thus such key can not be decoded.
+                let _ = db_id;
+
+                // Legacy issue: Catalog is not encoded into key.
+                if catalog_name == Self::DEFAULT_CATALOG {
+                    b.push_raw("table-by-id").push_u64(*table_id)
+                } else {
+                    b.push_raw("table-by-catalog-id")
+                        .push_str(catalog_name)
+                        .push_u64(*table_id)
+                }
+            }
+            OwnershipObject::Stage { name } => b.push_raw("stage-by-name").push_str(name),
+            OwnershipObject::UDF { name } => b.push_raw("udf-by-name").push_str(name),
+        }
+    }
+
+    /// Parse encoded key and return the OwnershipObject
+    pub(crate) fn parse_key(p: &mut kvapi::KeyParser) -> Result<Self, kvapi::KeyError> {
+        let q = p.next_raw()?;
+        match q {
+            "database-by-id" => {
+                let db_id = p.next_u64()?;
+                Ok(OwnershipObject::Database {
+                    catalog_name: Self::DEFAULT_CATALOG.to_string(),
+                    db_id,
+                })
+            }
+            "database-by-catalog-id" => {
+                let catalog_name = p.next_str()?;
+                let db_id = p.next_u64()?;
+                Ok(OwnershipObject::Database {
+                    catalog_name,
+                    db_id,
+                })
+            }
+            "table-by-id" => {
+                let table_id = p.next_u64()?;
+                Ok(OwnershipObject::Table {
+                    catalog_name: Self::DEFAULT_CATALOG.to_string(),
+                    db_id: 0,
+                    table_id,
+                })
+            }
+            "table-by-catalog-id" => {
+                let catalog_name = p.next_str()?;
+                let table_id = p.next_u64()?;
+                Ok(OwnershipObject::Table {
+                    catalog_name,
+                    db_id: 0, // string key does not contain db_id
+                    table_id,
+                })
+            }
+            "stage-by-name" => {
+                let name = p.next_str()?;
+                Ok(OwnershipObject::Stage { name })
+            }
+            "udf-by-name" => {
+                let name = p.next_str()?;
+                Ok(OwnershipObject::UDF { name })
+            }
+            _ => Err(kvapi::KeyError::InvalidSegment {
+                i: p.index(),
+                expect: "database-by-id|database-by-catalog-id|table-by-id|table-by-catalog-id|stage-by-name|udf-by-name"
+                    .to_string(),
+                got: q.to_string(),
+            }),
+        }
+        //
+    }
+}
+
+/// The meta-service key of object whose ownership to grant.
+///
+/// It could be a tenant's database, a tenant's table etc.
+/// It is in form of `__fd_object_owners/<tenant>/<object>`.
+/// where `<object>` could be:
+/// - `database-by-id/<db_id>`
+/// - `database-by-catalog-id/<catalog>/<db_id>`
+/// - `table-by-id/<table_id>`
+/// - `table-by-catalog-id/<catalog>/<table_id>`
+/// - `stage-by-name/<stage_name>`
+/// - `udf-by-name/<udf_name>`
+#[derive(Clone, Debug, Eq, PartialEq, Hash)]
+pub struct TenantOwnershipObject {
+    tenant: Tenant,
+    object: OwnershipObject,
+}
+
+impl TenantOwnershipObject {
+    pub fn new(tenant: Tenant, subject: OwnershipObject) -> Self {
+        // Legacy issue: Assert compatibility: No other catalog should be used.
+        // Assertion is disabled.
+        // Instead, accessing field `object` is disallowed.
+        // match &subject {
+        //     OwnershipObject::Database { catalog_name, .. } => {
+        //         assert_eq!(catalog_name, OwnershipObject::DEFAULT_CATALOG);
+        //     }
+        //     OwnershipObject::Table { catalog_name, .. } => {
+        //         assert_eq!(catalog_name, OwnershipObject::DEFAULT_CATALOG);
+        //     }
+        //     OwnershipObject::Stage { .. } => {}
+        //     OwnershipObject::UDF { .. } => {}
+        // }
+
+        Self::new_unchecked(tenant, subject)
+    }
+
+    pub(crate) fn new_unchecked(tenant: Tenant, object: OwnershipObject) -> Self {
+        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.
+    // pub fn subject(&self) -> &OwnershipObject {
+    //     &self.object
+    // }
+}
+
 #[derive(serde::Serialize, serde::Deserialize, Clone, Debug, Eq, PartialEq, Hash)]
 pub enum GrantObject {
     Global,
@@ -315,3 +482,172 @@ impl fmt::Display for UserGrantSet {
         write!(f, "ROLES: {:?}", self.roles())
     }
 }
+
+mod kvapi_key_impl {
+    use databend_common_meta_kvapi::kvapi;
+    use databend_common_meta_kvapi::kvapi::KeyError;
+
+    use crate::principal::user_grant::TenantOwnershipObject;
+    use crate::principal::OwnershipInfo;
+    use crate::principal::OwnershipObject;
+    use crate::tenant::Tenant;
+
+    impl kvapi::Key for TenantOwnershipObject {
+        const PREFIX: &'static str = "__fd_object_owners";
+        type ValueType = OwnershipInfo;
+
+        fn parent(&self) -> Option<String> {
+            Some(self.tenant.to_string_key())
+        }
+
+        fn to_string_key(&self) -> String {
+            let b = kvapi::KeyBuilder::new_prefixed(Self::PREFIX).push_str(&self.tenant.tenant);
+            self.object.build_key(b).done()
+        }
+
+        fn from_str_key(s: &str) -> Result<Self, KeyError> {
+            let mut p = kvapi::KeyParser::new_prefixed(s, Self::PREFIX)?;
+
+            let tenant = p.next_str()?;
+            let subject = OwnershipObject::parse_key(&mut p)?;
+            p.done()?;
+
+            Ok(TenantOwnershipObject {
+                tenant: Tenant::new(tenant),
+                object: subject,
+            })
+        }
+    }
+
+    impl kvapi::Value for OwnershipInfo {
+        fn dependency_keys(&self) -> impl IntoIterator<Item = String> {
+            []
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use databend_common_meta_kvapi::kvapi::Key;
+
+    use crate::principal::user_grant::TenantOwnershipObject;
+    use crate::principal::OwnershipObject;
+    use crate::tenant::Tenant;
+
+    #[test]
+    fn test_role_grantee_as_kvapi_key() {
+        // db with default catalog
+        {
+            let role_grantee = TenantOwnershipObject::new_unchecked(
+                Tenant::new("test"),
+                OwnershipObject::Database {
+                    catalog_name: "default".to_string(),
+                    db_id: 1,
+                },
+            );
+
+            let key = role_grantee.to_string_key();
+            assert_eq!("__fd_object_owners/test/database-by-id/1", key);
+
+            let parsed = TenantOwnershipObject::from_str_key(&key).unwrap();
+            assert_eq!(role_grantee, parsed);
+        }
+
+        // db with catalog
+        {
+            let role_grantee = TenantOwnershipObject::new_unchecked(
+                Tenant::new("test"),
+                OwnershipObject::Database {
+                    catalog_name: "cata/foo".to_string(),
+                    db_id: 1,
+                },
+            );
+
+            let key = role_grantee.to_string_key();
+            assert_eq!(
+                "__fd_object_owners/test/database-by-catalog-id/cata%2ffoo/1",
+                key
+            );
+
+            let parsed = TenantOwnershipObject::from_str_key(&key).unwrap();
+            assert_eq!(role_grantee, parsed);
+        }
+
+        // table with default catalog
+        {
+            let obj =
+                TenantOwnershipObject::new_unchecked(Tenant::new("test"), OwnershipObject::Table {
+                    catalog_name: "default".to_string(),
+                    db_id: 1,
+                    table_id: 2,
+                });
+
+            let key = obj.to_string_key();
+            assert_eq!("__fd_object_owners/test/table-by-id/2", key);
+
+            let parsed = TenantOwnershipObject::from_str_key(&key).unwrap();
+            assert_eq!(
+                TenantOwnershipObject::new_unchecked(Tenant::new("test"), OwnershipObject::Table {
+                    catalog_name: "default".to_string(),
+                    db_id: 0, // db_id is not encoded into key
+                    table_id: 2,
+                }),
+                parsed
+            );
+        }
+
+        // table with catalog
+        {
+            let obj =
+                TenantOwnershipObject::new_unchecked(Tenant::new("test"), OwnershipObject::Table {
+                    catalog_name: "cata/foo".to_string(),
+                    db_id: 1,
+                    table_id: 2,
+                });
+
+            let key = obj.to_string_key();
+            assert_eq!(
+                "__fd_object_owners/test/table-by-catalog-id/cata%2ffoo/2",
+                key
+            );
+
+            let parsed = TenantOwnershipObject::from_str_key(&key).unwrap();
+            assert_eq!(
+                TenantOwnershipObject::new_unchecked(Tenant::new("test"), OwnershipObject::Table {
+                    catalog_name: "cata/foo".to_string(),
+                    db_id: 0, // db_id is not encoded into key
+                    table_id: 2,
+                }),
+                parsed
+            );
+        }
+
+        // stage
+        {
+            let role_grantee =
+                TenantOwnershipObject::new_unchecked(Tenant::new("test"), OwnershipObject::Stage {
+                    name: "foo".to_string(),
+                });
+
+            let key = role_grantee.to_string_key();
+            assert_eq!("__fd_object_owners/test/stage-by-name/foo", key);
+
+            let parsed = TenantOwnershipObject::from_str_key(&key).unwrap();
+            assert_eq!(role_grantee, parsed);
+        }
+
+        // udf
+        {
+            let role_grantee =
+                TenantOwnershipObject::new_unchecked(Tenant::new("test"), OwnershipObject::UDF {
+                    name: "foo".to_string(),
+                });
+
+            let key = role_grantee.to_string_key();
+            assert_eq!("__fd_object_owners/test/udf-by-name/foo", key);
+
+            let parsed = TenantOwnershipObject::from_str_key(&key).unwrap();
+            assert_eq!(role_grantee, parsed);
+        }
+    }
+}

src/meta/kvapi/src/kvapi/key_parser.rs

@@ -51,6 +51,13 @@ impl<'s> KeyParser<'s> {
         Ok(s)
     }
 
+    /// Get the index of the last returned element.
+    ///
+    /// If no element is returned, it will panic.
+    pub fn index(&self) -> usize {
+        self.i - 1
+    }
+
     /// Pop the next element in raw `&str`, without unescaping or decoding.
     ///
     /// If there is no more element, it returns KeyError::WrongNumberOfSegments.

src/meta/types/src/non_empty.rs

@@ -26,7 +26,7 @@ impl<'a> NonEmptyStr<'a> {
         Ok(NonEmptyStr { non_empty: s })
     }
 
-    fn get(&self) -> &str {
+    pub fn get(&self) -> &str {
         self.non_empty
     }
 }