feat(query):masking policy support rbac (#18982)

* feat(query):masking policy support rbac

* add some comment

* add some unit test

Author: TCeason

Cargo.lock

@@ -48,6 +48,17 @@ pub trait DatamaskApi: Send + Sync {
         name_ident: &DataMaskNameIdent,
     ) -> Result<Option<SeqV<DatamaskMeta>>, MetaError>;
 
+    async fn get_data_mask_id(
+        &self,
+        name_ident: &DataMaskNameIdent,
+    ) -> Result<Option<SeqV<DataMaskId>>, MetaError>;
+
+    async fn get_data_mask_name_by_id(
+        &self,
+        tenant: &Tenant,
+        policy_id: u64,
+    ) -> Result<Option<String>, MetaError>;
+
     async fn get_data_mask_by_id(
         &self,
         tenant: &Tenant,

src/meta/api/src/data_mask_api.rs

@@ -17,7 +17,9 @@ use databend_common_meta_app::data_mask::CreateDatamaskReply;
 use databend_common_meta_app::data_mask::CreateDatamaskReq;
 use databend_common_meta_app::data_mask::DataMaskId;
 use databend_common_meta_app::data_mask::DataMaskIdIdent;
+use databend_common_meta_app::data_mask::DataMaskIdToNameIdent;
 use databend_common_meta_app::data_mask::DataMaskNameIdent;
+use databend_common_meta_app::data_mask::DataMaskNameIdentRaw;
 use databend_common_meta_app::data_mask::DatamaskMeta;
 use databend_common_meta_app::data_mask::MaskPolicyIdTableId;
 use databend_common_meta_app::data_mask::MaskPolicyTableIdIdent;
@@ -46,6 +48,7 @@ use crate::txn_condition_util::txn_cond_eq_keys_with_prefix;
 use crate::txn_condition_util::txn_cond_eq_seq;
 use crate::txn_core_util::send_txn;
 use crate::txn_core_util::txn_delete_exact;
+use crate::txn_op_builder_util::txn_op_del;
 use crate::txn_op_builder_util::txn_op_put_pb;
 
 /// DatamaskApi is implemented upon kvapi::KVApi.
@@ -77,6 +80,8 @@ impl<KV: kvapi::KVApi<Error = MetaError>> DatamaskApi for KV {
 
         let id = DataMaskId::new(masking_policy_id);
         let id_ident = DataMaskIdIdent::new_generic(name_ident.tenant(), id);
+        let id_to_name_ident = DataMaskIdToNameIdent::new_generic(name_ident.tenant(), id);
+        let name_raw = DataMaskNameIdentRaw::from(name_ident.clone());
         let id_list_key = MaskPolicyTableIdListIdent::new_from(name_ident.clone());
 
         debug!(
@@ -91,9 +96,11 @@ impl<KV: kvapi::KVApi<Error = MetaError>> DatamaskApi for KV {
             txn.condition.push(txn_cond_eq_seq(name_ident, 0));
             txn.condition
                 .push(txn_cond_eq_seq(&row_access_name_ident, 0));
+
             txn.if_then.extend(vec![
-                txn_op_put_pb(name_ident, &id, None)?,  // name -> db_id
+                txn_op_put_pb(name_ident, &id, None)?, // name -> masking_policy_id
                 txn_op_put_pb(&id_ident, &meta, None)?, // id -> meta
+                txn_op_put_pb(&id_to_name_ident, &name_raw, None)?, // id -> name
                 // TODO: Tentative retention for compatibility MaskPolicyTableIdListIdent related logic. It can be directly deleted later
                 txn_op_put_pb(&id_list_key, &id_list, None)?, // data mask name -> id_list
             ]);
@@ -159,7 +166,9 @@ impl<KV: kvapi::KVApi<Error = MetaError>> DatamaskApi for KV {
             }
 
             // No references - drop the policy
-            let id_ident = seq_id.data.into_t_ident(tenant);
+            let id_ident = seq_id.data.into_t_ident(tenant.clone());
+            let id_to_name_ident =
+                DataMaskIdToNameIdent::new_generic(tenant, DataMaskId::new(policy_id));
             let mut txn = TxnRequest::default();
 
             // Ensure no new references were created
@@ -168,6 +177,7 @@ impl<KV: kvapi::KVApi<Error = MetaError>> DatamaskApi for KV {
 
             txn_delete_exact(&mut txn, name_ident, seq_id.seq);
             txn_delete_exact(&mut txn, &id_ident, seq_meta.seq);
+            txn.if_then.push(txn_op_del(&id_to_name_ident));
 
             // TODO: Tentative retention for compatibility. Can be deleted later.
             clear_table_column_mask_policy(self, name_ident, &mut txn).await?;
@@ -191,6 +201,26 @@ impl<KV: kvapi::KVApi<Error = MetaError>> DatamaskApi for KV {
         Ok(res.map(|(_, seq_meta)| seq_meta))
     }
 
+    async fn get_data_mask_id(
+        &self,
+        name_ident: &DataMaskNameIdent,
+    ) -> Result<Option<SeqV<DataMaskId>>, MetaError> {
+        self.get_pb(name_ident).await
+    }
+
+    async fn get_data_mask_name_by_id(
+        &self,
+        tenant: &Tenant,
+        policy_id: u64,
+    ) -> Result<Option<String>, MetaError> {
+        let ident = DataMaskIdToNameIdent::new_generic(tenant.clone(), DataMaskId::new(policy_id));
+        let seq_meta = self.get_pb(&ident).await?;
+
+        debug!(ident :% =(&ident); "get_data_mask_name_by_id");
+
+        Ok(seq_meta.map(|s| s.data.data_mask_name().to_string()))
+    }
+
     async fn get_data_mask_by_id(
         &self,
         tenant: &Tenant,

src/meta/api/src/data_mask_api_impl.rs

@@ -19,7 +19,7 @@
 
 extern crate databend_common_meta_types;
 pub mod catalog_api;
-mod data_mask_api;
+pub mod data_mask_api;
 mod data_mask_api_impl;
 pub mod data_retention_util;
 mod database_api;

src/meta/api/src/lib.rs

@@ -51,6 +51,12 @@ pub trait RowAccessPolicyApi: Send + Sync {
         name_ident: &RowAccessPolicyNameIdent,
     ) -> Result<Option<(SeqV<RowAccessPolicyId>, SeqV<RowAccessPolicyMeta>)>, MetaError>;
 
+    async fn get_row_access_policy_name_by_id(
+        &self,
+        tenant: &Tenant,
+        policy_id: u64,
+    ) -> Result<Option<String>, MetaError>;
+
     async fn get_row_access_policy_by_id(
         &self,
         tenant: &Tenant,

src/meta/api/src/row_access_policy_api.rs

@@ -20,8 +20,10 @@ use databend_common_meta_app::row_access_policy::CreateRowAccessPolicyReply;
 use databend_common_meta_app::row_access_policy::CreateRowAccessPolicyReq;
 use databend_common_meta_app::row_access_policy::RowAccessPolicyId;
 use databend_common_meta_app::row_access_policy::RowAccessPolicyIdIdent;
+use databend_common_meta_app::row_access_policy::RowAccessPolicyIdToNameIdent;
 use databend_common_meta_app::row_access_policy::RowAccessPolicyMeta;
 use databend_common_meta_app::row_access_policy::RowAccessPolicyNameIdent;
+use databend_common_meta_app::row_access_policy::RowAccessPolicyNameIdentRaw;
 use databend_common_meta_app::row_access_policy::RowAccessPolicyTableIdIdent;
 use databend_common_meta_app::tenant::Tenant;
 use databend_common_meta_app::tenant_key::errors::ExistError;
@@ -44,6 +46,7 @@ use crate::txn_condition_util::txn_cond_eq_keys_with_prefix;
 use crate::txn_condition_util::txn_cond_eq_seq;
 use crate::txn_core_util::send_txn;
 use crate::txn_core_util::txn_delete_exact;
+use crate::txn_op_builder_util::txn_op_del;
 use crate::txn_op_builder_util::txn_op_put_pb;
 
 /// RowAccessPolicyApi is implemented upon kvapi::KVApi.
@@ -73,8 +76,12 @@ impl<KV: kvapi::KVApi<Error = MetaError>> RowAccessPolicyApi for KV {
         // Create row policy by inserting these record:
         // name -> id
         // id -> policy
+        // id -> name
 
         let id_ident = RowAccessPolicyIdIdent::new_generic(name_ident.tenant(), policy_id);
+        let id_to_name_ident =
+            RowAccessPolicyIdToNameIdent::new_generic(name_ident.tenant(), policy_id);
+        let name_raw = RowAccessPolicyNameIdentRaw::from(name_ident.clone());
 
         debug!(
             id :? =(&id_ident),
@@ -89,6 +96,7 @@ impl<KV: kvapi::KVApi<Error = MetaError>> RowAccessPolicyApi for KV {
             txn.if_then.extend(vec![
                 txn_op_put_pb(name_ident, &policy_id, None)?, // name -> policy_id
                 txn_op_put_pb(&id_ident, &meta, None)?,       // id -> meta
+                txn_op_put_pb(&id_to_name_ident, &name_raw, None)?, // id -> name
             ]);
         }
 
@@ -151,6 +159,10 @@ impl<KV: kvapi::KVApi<Error = MetaError>> RowAccessPolicyApi for KV {
 
             // No references - drop the policy
             let id_ident = seq_id.data.into_t_ident(tenant);
+            let id_to_name_ident = RowAccessPolicyIdToNameIdent::new_generic(
+                tenant.clone(),
+                RowAccessPolicyId::new(policy_id),
+            );
             let mut txn = TxnRequest::default();
 
             // Ensure no new references were created
@@ -159,6 +171,7 @@ impl<KV: kvapi::KVApi<Error = MetaError>> RowAccessPolicyApi for KV {
 
             txn_delete_exact(&mut txn, name_ident, seq_id.seq);
             txn_delete_exact(&mut txn, &id_ident, seq_meta.seq);
+            txn.if_then.push(txn_op_del(&id_to_name_ident));
 
             let (succ, _responses) = send_txn(self, txn).await?;
             if succ {
@@ -179,6 +192,24 @@ impl<KV: kvapi::KVApi<Error = MetaError>> RowAccessPolicyApi for KV {
         Ok(res)
     }
 
+    async fn get_row_access_policy_name_by_id(
+        &self,
+        tenant: &Tenant,
+        policy_id: u64,
+    ) -> Result<Option<String>, MetaError> {
+        let ident = RowAccessPolicyIdToNameIdent::new_generic(
+            tenant.clone(),
+            RowAccessPolicyId::new(policy_id),
+        );
+        let seq_meta = self.get_pb(&ident).await?;
+
+        debug!(ident :% =(&ident); "get_row_access_policy_name_by_id");
+
+        let name = seq_meta.map(|s| s.data.row_access_name().to_string());
+
+        Ok(name)
+    }
+
     async fn get_row_access_policy_by_id(
         &self,
         tenant: &Tenant,

src/meta/api/src/row_access_policy_api_impl.rs

@@ -0,0 +1,81 @@
+// 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::data_mask::DataMaskId;
+use crate::tenant_key::ident::TIdent;
+use crate::tenant_key::raw::TIdentRaw;
+
+/// Tenantless key mapping a masking policy id back to its name.
+/// This enables reverse lookup for SHOW GRANTS and ownership listings.
+pub type DataMaskIdToNameIdent = TIdent<Resource, DataMaskId>;
+/// Raw form of [`DataMaskIdToNameIdent`] used for serde/protobuf.
+pub type DataMaskIdToNameIdentRaw = TIdentRaw<Resource, DataMaskId>;
+
+pub use kvapi_impl::Resource;
+
+impl DataMaskIdToNameIdent {
+    pub fn data_mask_id(&self) -> DataMaskId {
+        *self.name()
+    }
+}
+
+impl DataMaskIdToNameIdentRaw {
+    pub fn data_mask_id(&self) -> DataMaskId {
+        *self.name()
+    }
+}
+
+mod kvapi_impl {
+
+    use databend_common_meta_kvapi::kvapi;
+
+    use crate::data_mask::data_mask_id_to_name_ident::DataMaskIdToNameIdent;
+    use crate::data_mask::data_mask_name_ident::DataMaskNameIdentRaw;
+    use crate::tenant_key::resource::TenantResource;
+
+    pub struct Resource;
+    impl TenantResource for Resource {
+        const PREFIX: &'static str = "__fd_datamask_id_to_name";
+        const TYPE: &'static str = "DataMaskIdToNameIdent";
+        const HAS_TENANT: bool = true;
+        type ValueType = DataMaskNameIdentRaw;
+    }
+
+    impl kvapi::Value for DataMaskNameIdentRaw {
+        type KeyType = DataMaskIdToNameIdent;
+        fn dependency_keys(&self, _key: &Self::KeyType) -> impl IntoIterator<Item = String> {
+            []
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use databend_common_meta_kvapi::kvapi::Key;
+
+    use crate::data_mask::DataMaskId;
+    use crate::data_mask::DataMaskIdToNameIdent;
+    use crate::tenant::Tenant;
+
+    #[test]
+    fn test_data_mask_id_ident() {
+        let tenant = Tenant::new_literal("dummy");
+        let ident = DataMaskIdToNameIdent::new_generic(tenant, DataMaskId::new(3));
+
+        let key = ident.to_string_key();
+        assert_eq!(key, "__fd_datamask_id_to_name/dummy/3");
+
+        assert_eq!(ident, DataMaskIdToNameIdent::from_str_key(&key).unwrap());
+    }
+}

src/meta/app/src/data_mask/data_mask_id_to_name_ident.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 pub mod data_mask_id_ident;
+pub mod data_mask_id_to_name_ident;
 pub mod data_mask_name_ident;
 pub mod mask_policy_policy_table_id_ident;
 pub mod mask_policy_table_id_list_ident;
@@ -23,7 +24,10 @@ use chrono::DateTime;
 use chrono::Utc;
 pub use data_mask_id_ident::DataMaskId;
 pub use data_mask_id_ident::DataMaskIdIdent;
+pub use data_mask_id_to_name_ident::DataMaskIdToNameIdent;
+pub use data_mask_id_to_name_ident::DataMaskIdToNameIdentRaw;
 pub use data_mask_name_ident::DataMaskNameIdent;
+pub use data_mask_name_ident::DataMaskNameIdentRaw;
 pub use mask_policy_policy_table_id_ident::MaskPolicyIdTableId;
 pub use mask_policy_policy_table_id_ident::MaskPolicyTableIdIdent;
 pub use mask_policy_table_id_list_ident::MaskPolicyTableIdListIdent;

src/meta/app/src/data_mask/mod.rs

@@ -30,6 +30,8 @@ use databend_common_meta_kvapi::kvapi::KeyCodec;
 /// - `table-by-catalog-id/<catalog>/<table_id>`
 /// - `stage-by-name/<stage_name>`
 /// - `udf-by-name/<udf_name>`
+/// - `procedure-by-id/<procedure_id>`
+/// - `masking-policy-by-id/<policy_id>`
 #[derive(serde::Serialize, serde::Deserialize, Clone, Debug, Eq, PartialEq, Hash)]
 pub enum OwnershipObject {
     /// used on the fuse databases
@@ -68,6 +70,10 @@ pub enum OwnershipObject {
     Procedure {
         procedure_id: u64,
     },
+
+    MaskingPolicy {
+        policy_id: u64,
+    },
 }
 
 impl OwnershipObject {
@@ -100,6 +106,9 @@ impl fmt::Display for OwnershipObject {
             OwnershipObject::Connection { name } => write!(f, "CONNECTION {name}"),
             OwnershipObject::Sequence { name } => write!(f, "SEQUENCE {name}"),
             OwnershipObject::Procedure { procedure_id } => write!(f, "PROCEDURE {procedure_id}"),
+            OwnershipObject::MaskingPolicy { policy_id } => {
+                write!(f, "MASKING POLICY {policy_id}")
+            }
         }
     }
 }
@@ -145,6 +154,9 @@ impl KeyCodec for OwnershipObject {
             OwnershipObject::Procedure { procedure_id } => {
                 b.push_raw("procedure-by-id").push_u64(*procedure_id)
             }
+            OwnershipObject::MaskingPolicy { policy_id } => {
+                b.push_raw("masking-policy-by-id").push_u64(*policy_id)
+            }
         }
     }
 
@@ -207,9 +219,13 @@ impl KeyCodec for OwnershipObject {
                 let procedure_id = p.next_u64()?;
                 Ok(OwnershipObject::Procedure { procedure_id })
             }
+            "masking-policy-by-id" => {
+                let policy_id = p.next_u64()?;
+                Ok(OwnershipObject::MaskingPolicy { policy_id })
+            }
             _ => 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|warehouse-by-id|connection-by-name"
+                expect: "database-by-id|database-by-catalog-id|table-by-id|table-by-catalog-id|stage-by-name|udf-by-name|warehouse-by-id|connection-by-name|masking-policy-by-id"
                     .to_string(),
                 got: q.to_string(),
             }),

src/meta/app/src/principal/ownership_object.rs

@@ -271,6 +271,23 @@ mod tests {
             let parsed = TenantOwnershipObjectIdent::from_str_key(&key).unwrap();
             assert_eq!(role_grantee, parsed);
         }
+
+        // masking policy
+        {
+            let role_grantee = TenantOwnershipObjectIdent::new_unchecked(
+                Tenant::new_literal("tenant_mask"),
+                OwnershipObject::MaskingPolicy { policy_id: 99 },
+            );
+
+            let key = role_grantee.to_string_key();
+            assert_eq!(
+                "__fd_object_owners/tenant_mask/masking-policy-by-id/99",
+                key
+            );
+
+            let parsed = TenantOwnershipObjectIdent::from_str_key(&key).unwrap();
+            assert_eq!(role_grantee, parsed);
+        }
     }
 
     #[test]