feat(query): support create/drop/desc row policy (#18490)

* feat(query):support create/drop/desc row policy

* fix conversation

* Introduce Experimental Feature Gate

enable_experimental_row_access_policy default disable

* RowAccessPolicyTableIdList -> RowAccessPolicyTableIdIdent

__fd_row_access_policy_apply_table_id/tenant/<policy_id>/<table_id>

* fix conversation

* optimize Error

Author: TCeason

Cargo.lock

@@ -200,6 +200,7 @@ databend-enterprise-hilbert-clustering = { path = "src/query/ee_features/hilbert
 databend-enterprise-meta = { path = "src/meta/ee" }
 databend-enterprise-query = { path = "src/query/ee" }
 databend-enterprise-resources-management = { path = "src/query/ee_features/resources_management" }
+databend-enterprise-row-access-policy-feature = { path = "src/query/ee_features/row_access_policy" }
 databend-enterprise-storage-encryption = { path = "src/query/ee_features/storage_encryption" }
 databend-enterprise-storage-quota = { path = "src/query/ee_features/storage_quota" }
 databend-enterprise-stream-handler = { path = "src/query/ee_features/stream_handler" }

Cargo.toml

@@ -283,6 +283,8 @@ build_exceptions! {
     UnmatchMaskPolicyReturnType(1121),
     /// Empty share endpoint config
     EmptyShareEndpointConfig(1130),
+    /// Unknown row policy
+    UnknownRowAccessPolicy(1131),
 }
 
 // Sequence Errors [1124-1126, 3101]
@@ -519,6 +521,8 @@ build_exceptions! {
     CommitTableMetaError(2322),
     /// Create as drop table without drop time
     CreateAsDropTableWithoutDropTime(2323),
+    /// Row Policy already exists
+    RowAccessPolicyAlreadyExists(2324),
 }
 
 // Stage and Connection Errors [2501-2505, 2510-2512]

src/common/exception/src/exception_code.rs

@@ -89,6 +89,8 @@ pub enum Feature {
     MaxNodeQuota(usize),
     #[serde(alias = "max_cpu_quota", alias = "MAX_CPU_QUOTA")]
     MaxCpuQuota(usize),
+    #[serde(alias = "row_access_policy", alias = "ROW_ACCESS_POLICY")]
+    RowAccessPolicy,
     #[serde(other)]
     Unknown,
 }
@@ -141,6 +143,7 @@ impl fmt::Display for Feature {
             Feature::SystemHistory => write!(f, "system_history"),
             Feature::VectorIndex => write!(f, "vector_index"),
             Feature::PrivateTask => write!(f, "private_task"),
+            Feature::RowAccessPolicy => write!(f, "row_access_policy"),
             Feature::Unknown => write!(f, "unknown"),
             Feature::MaxCpuQuota(v) => write!(f, "max_cpu_quota({})", v),
             Feature::MaxNodeQuota(v) => write!(f, "max_node_quota({})", v),
@@ -191,6 +194,7 @@ impl Feature {
             | (Feature::LicenseInfo, Feature::LicenseInfo)
             | (Feature::Stream, Feature::Stream)
             | (Feature::DataMask, Feature::DataMask)
+            | (Feature::RowAccessPolicy, Feature::RowAccessPolicy)
             | (Feature::InvertedIndex, Feature::InvertedIndex)
             | (Feature::VirtualColumn, Feature::VirtualColumn)
             | (Feature::AttacheTable, Feature::AttacheTable)
@@ -396,6 +400,11 @@ mod tests {
             serde_json::from_str::<Feature>("{\"MaxNodeQuota\": 1}").unwrap()
         );
 
+        assert_eq!(
+            Feature::RowAccessPolicy,
+            serde_json::from_str::<Feature>("\"RowAccessPolicy\"").unwrap()
+        );
+
         assert_eq!(
             Feature::Unknown,
             serde_json::from_str::<Feature>("\"ssss\"").unwrap()
@@ -433,11 +442,12 @@ mod tests {
                 Feature::WorkloadGroup,
                 Feature::SystemHistory,
                 Feature::PrivateTask,
+                Feature::RowAccessPolicy,
             ]),
         };
 
         assert_eq!(
-            "LicenseInfo{ type: enterprise, org: databend, tenants: [databend_tenant,foo], features: [aggregate_index,amend_table,attach_table,compute_quota(threads_num: 1, memory_usage: 1),computed_column,data_mask,hilbert_clustering,inverted_index,license_info,ngram_index,private_task,storage_encryption,storage_quota(storage_usage: 1),stream,system_history,vacuum,virtual_column,workload_group] }",
+            "LicenseInfo{ type: enterprise, org: databend, tenants: [databend_tenant,foo], features: [aggregate_index,amend_table,attach_table,compute_quota(threads_num: 1, memory_usage: 1),computed_column,data_mask,hilbert_clustering,inverted_index,license_info,ngram_index,private_task,row_access_policy,storage_encryption,storage_quota(storage_usage: 1),stream,system_history,vacuum,virtual_column,workload_group] }",
             license_info.to_string()
         );
     }

src/common/license/src/license.rs

@@ -36,10 +36,13 @@ pub mod txn_backoff;
 pub mod util;
 
 pub mod crud;
+mod row_access_policy_api;
+mod row_access_policy_api_impl;
 mod sequence_api_impl;
 pub(crate) mod sequence_nextval_impl;
 
 pub use data_mask_api::DatamaskApi;
+pub use row_access_policy_api::RowAccessPolicyApi;
 pub use schema_api::SchemaApi;
 pub use schema_api_test_suite::SchemaApiTestSuite;
 pub use sequence_api::SequenceApi;

src/meta/api/src/lib.rs

@@ -0,0 +1,48 @@
+// 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 databend_common_meta_app::row_access_policy::row_access_policy_name_ident;
+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::RowAccessPolicyMeta;
+use databend_common_meta_app::row_access_policy::RowAccessPolicyNameIdent;
+use databend_common_meta_app::tenant_key::errors::ExistError;
+use databend_common_meta_types::MetaError;
+use databend_common_meta_types::SeqV;
+
+use crate::meta_txn_error::MetaTxnError;
+
+#[async_trait::async_trait]
+pub trait RowAccessPolicyApi: Send + Sync {
+    async fn create_row_access(
+        &self,
+        req: CreateRowAccessPolicyReq,
+    ) -> Result<
+        Result<CreateRowAccessPolicyReply, ExistError<row_access_policy_name_ident::Resource>>,
+        MetaTxnError,
+    >;
+
+    /// On success, returns the dropped id and row policy.
+    /// Returning None, means nothing is removed.
+    async fn drop_row_access(
+        &self,
+        name_ident: &RowAccessPolicyNameIdent,
+    ) -> Result<Option<(SeqV<RowAccessPolicyId>, SeqV<RowAccessPolicyMeta>)>, MetaTxnError>;
+
+    async fn get_row_access(
+        &self,
+        name_ident: &RowAccessPolicyNameIdent,
+    ) -> Result<Option<SeqV<RowAccessPolicyMeta>>, MetaError>;
+}

src/meta/api/src/row_access_policy_api.rs

@@ -0,0 +1,167 @@
+// 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 databend_common_meta_app::id_generator::IdGenerator;
+use databend_common_meta_app::row_access_policy::row_access_policy_name_ident;
+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::RowAccessPolicyMeta;
+use databend_common_meta_app::row_access_policy::RowAccessPolicyNameIdent;
+use databend_common_meta_app::tenant_key::errors::ExistError;
+use databend_common_meta_app::KeyWithTenant;
+use databend_common_meta_kvapi::kvapi;
+use databend_common_meta_types::MetaError;
+use databend_common_meta_types::SeqV;
+use databend_common_meta_types::TxnRequest;
+use fastrace::func_name;
+use log::debug;
+
+use crate::fetch_id;
+use crate::kv_pb_api::KVPbApi;
+use crate::meta_txn_error::MetaTxnError;
+use crate::row_access_policy_api::RowAccessPolicyApi;
+use crate::send_txn;
+use crate::txn_backoff::txn_backoff;
+use crate::txn_cond_eq_seq;
+use crate::util::txn_delete_exact;
+use crate::util::txn_op_put_pb;
+
+/// RowAccessPolicyApi is implemented upon kvapi::KVApi.
+/// Thus every type that impl kvapi::KVApi impls RowAccessPolicyApi.
+#[tonic::async_trait]
+impl<KV: kvapi::KVApi<Error = MetaError>> RowAccessPolicyApi for KV {
+    async fn create_row_access(
+        &self,
+        req: CreateRowAccessPolicyReq,
+    ) -> Result<
+        Result<CreateRowAccessPolicyReply, ExistError<row_access_policy_name_ident::Resource>>,
+        MetaTxnError,
+    > {
+        debug!(req :? =(&req); "RowAccessPolicyApi: {}", func_name!());
+
+        let name_ident = &req.name;
+
+        let id = fetch_id(self, IdGenerator::row_access_id()).await?;
+        let mut trials = txn_backoff(None, func_name!());
+        let id = loop {
+            trials.next().unwrap()?.await;
+
+            let mut txn = TxnRequest::default();
+
+            let res = self.get_id_and_value(name_ident).await?;
+            debug!(res :? = res, name_key :? =(name_ident); "create_row_access");
+
+            let mut curr_seq = 0;
+
+            if let Some((seq_id, seq_meta)) = res {
+                if req.can_replace {
+                    let id_ident = seq_id.data.into_t_ident(name_ident.tenant());
+
+                    txn_delete_exact(&mut txn, &id_ident, seq_meta.seq);
+
+                    // TODO(eason): need to remove row policy from table meta
+
+                    curr_seq = seq_id.seq;
+                } else {
+                    return Ok(Err(name_ident.exist_error(func_name!())));
+                }
+            }
+
+            // Create row policy by inserting these record:
+            // name -> id
+            // id -> policy
+
+            let id = RowAccessPolicyId::new(id);
+            let id_ident = RowAccessPolicyIdIdent::new_generic(name_ident.tenant(), id);
+
+            debug!(
+                id :? =(&id_ident),
+                name_key :? =(name_ident);
+                "new RowAccessPolicy id"
+            );
+
+            {
+                let meta: RowAccessPolicyMeta = req.row_access_policy_meta.clone();
+                txn.condition.push(txn_cond_eq_seq(name_ident, curr_seq));
+                txn.if_then.extend(vec![
+                    txn_op_put_pb(name_ident, &id, None)?,  // name -> policy_id
+                    txn_op_put_pb(&id_ident, &meta, None)?, // id -> meta
+                ]);
+
+                let (succ, _responses) = send_txn(self, txn).await?;
+
+                debug!(
+                    name :? =(name_ident),
+                    id :? =(&id_ident),
+                    succ = succ;
+                    "create_row_access"
+                );
+
+                if succ {
+                    break id;
+                }
+            }
+        };
+
+        Ok(Ok(CreateRowAccessPolicyReply { id: *id }))
+    }
+
+    async fn drop_row_access(
+        &self,
+        name_ident: &RowAccessPolicyNameIdent,
+    ) -> Result<Option<(SeqV<RowAccessPolicyId>, SeqV<RowAccessPolicyMeta>)>, MetaTxnError> {
+        debug!(name_ident :? =(name_ident); "RowAccessPolicyApi: {}", func_name!());
+
+        let mut trials = txn_backoff(None, func_name!());
+        loop {
+            trials.next().unwrap()?.await;
+
+            let mut txn = TxnRequest::default();
+
+            let res = self.get_id_and_value(name_ident).await?;
+            debug!(res :? = res, name_key :? =(name_ident); "{}", func_name!());
+
+            let Some((seq_id, seq_meta)) = res else {
+                return Ok(None);
+            };
+
+            let id_ident = seq_id.data.into_t_ident(name_ident.tenant());
+
+            txn_delete_exact(&mut txn, name_ident, seq_id.seq);
+            txn_delete_exact(&mut txn, &id_ident, seq_meta.seq);
+
+            // TODO(eason): need to remove row policy from table meta
+
+            let (succ, _responses) = send_txn(self, txn).await?;
+            debug!(succ = succ;"{}", func_name!());
+
+            if succ {
+                return Ok(Some((seq_id, seq_meta)));
+            }
+        }
+    }
+
+    async fn get_row_access(
+        &self,
+        name_ident: &RowAccessPolicyNameIdent,
+    ) -> Result<Option<SeqV<RowAccessPolicyMeta>>, MetaError> {
+        debug!(req :? =(&name_ident); "RowAccessPolicyApi: {}", func_name!());
+
+        let res = self.get_id_and_value(name_ident).await?;
+
+        Ok(res.map(|(_, seq_meta)| seq_meta))
+    }
+}

src/meta/api/src/row_access_policy_api_impl.rs

@@ -20,6 +20,7 @@ use databend_common_meta_types::MatchSeq;
 use crate::data_mask::data_mask_name_ident;
 use crate::principal::procedure_name_ident;
 use crate::principal::ProcedureIdentity;
+use crate::row_access_policy::row_access_policy_name_ident;
 use crate::schema::catalog_name_ident;
 use crate::schema::dictionary_name_ident;
 use crate::schema::index_name_ident;
@@ -1088,6 +1089,9 @@ pub enum AppError {
     #[error(transparent)]
     UnknownDataMask(#[from] UnknownError<data_mask_name_ident::Resource>),
 
+    #[error(transparent)]
+    UnknownRowAccessPolicy(#[from] UnknownError<row_access_policy_name_ident::Resource>),
+
     #[error(transparent)]
     UnmatchColumnDataType(#[from] UnmatchColumnDataType),
 
@@ -1467,12 +1471,6 @@ impl AppErrorMessage for IndexColumnIdNotFound {
     }
 }
 
-impl AppErrorMessage for UnknownDatamask {
-    fn message(&self) -> String {
-        format!("Datamask '{}' does not exists", self.name)
-    }
-}
-
 impl AppErrorMessage for UnmatchColumnDataType {
     fn message(&self) -> String {
         format!(
@@ -1627,6 +1625,9 @@ impl From<AppError> for ErrorCode {
 
             AppError::DatamaskAlreadyExists(err) => ErrorCode::DatamaskAlreadyExists(err.message()),
             AppError::UnknownDataMask(err) => ErrorCode::UnknownDatamask(err.message()),
+            AppError::UnknownRowAccessPolicy(err) => {
+                ErrorCode::UnknownRowAccessPolicy(err.message())
+            }
 
             AppError::UnmatchColumnDataType(err) => ErrorCode::UnmatchColumnDataType(err.message()),
             AppError::UnmatchMaskPolicyReturnType(err) => {