fix(meta): vacuum dropped table cleans policy refs (#19239)

Author: TCeason

src/meta/api/src/garbage_collection_api.rs

@@ -22,9 +22,13 @@ use databend_common_base::vec_ext::VecExt;
 use databend_common_meta_app::app_error::AppError;
 use databend_common_meta_app::app_error::CleanDbIdTableNamesFailed;
 use databend_common_meta_app::app_error::MarkDatabaseMetaAsGCInProgressFailed;
+use databend_common_meta_app::data_mask::MaskPolicyIdTableId;
+use databend_common_meta_app::data_mask::MaskPolicyTableIdIdent;
 use databend_common_meta_app::principal::AutoIncrementKey;
 use databend_common_meta_app::principal::OwnershipObject;
 use databend_common_meta_app::principal::TenantOwnershipObjectIdent;
+use databend_common_meta_app::row_access_policy::RowAccessPolicyTableIdIdent;
+use databend_common_meta_app::row_access_policy::row_access_policy_table_id_ident::RowAccessPolicyIdTableId;
 use databend_common_meta_app::schema::AutoIncrementStorageIdent;
 use databend_common_meta_app::schema::DBIdTableName;
 use databend_common_meta_app::schema::DatabaseId;
@@ -740,7 +744,6 @@ async fn remove_data_for_dropped_table(
     //     warn!("{}", err);
     //     return Ok(Err(err));
     // }
-
     txn_delete_exact(txn, table_id, seq_meta.seq);
 
     // Get id -> name mapping
@@ -823,6 +826,42 @@ async fn remove_data_for_dropped_table(
         }
     }
 
+    let tb_meta = &seq_meta.data;
+
+    // Clean up policy references for the dropped table.
+    // These records may have been orphaned if the table was dropped via DROP DATABASE.
+
+    // Delete masking policy references
+    {
+        let policy_ids: HashSet<u64> = tb_meta
+            .column_mask_policy_columns_ids
+            .values()
+            .map(|policy_map| policy_map.policy_id)
+            .collect();
+
+        txn.if_then.extend(policy_ids.into_iter().map(|policy_id| {
+            txn_op_del(&MaskPolicyTableIdIdent::new_generic(
+                tenant.clone(),
+                MaskPolicyIdTableId {
+                    policy_id,
+                    table_id: table_id.table_id,
+                },
+            ))
+        }));
+    }
+
+    // Delete row access policy reference
+    if let Some(policy_map) = &tb_meta.row_access_policy_columns_ids {
+        txn.if_then
+            .push(txn_op_del(&RowAccessPolicyTableIdIdent::new_generic(
+                tenant.clone(),
+                RowAccessPolicyIdTableId {
+                    policy_id: policy_map.policy_id,
+                    table_id: table_id.table_id,
+                },
+            )));
+    }
+
     Ok(Ok(()))
 }
 

src/query/ee/tests/it/storages/fuse/operations/vacuum.rs

@@ -1043,3 +1043,186 @@ async fn test_vacuum_dropped_table_clean_tag_refs() -> Result<()> {
 
     Ok(())
 }
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_vacuum_dropped_table_clean_policies() -> Result<()> {
+    use databend_common_meta_app::data_mask::MaskPolicyIdTableId;
+    use databend_common_meta_app::data_mask::MaskPolicyTableIdIdent;
+    use databend_common_meta_app::row_access_policy::RowAccessPolicyTableIdIdent;
+    use databend_common_meta_app::row_access_policy::row_access_policy_table_id_ident::RowAccessPolicyIdTableId;
+
+    // 1. Prepare local meta service
+    let meta = new_local_meta().await;
+    let endpoints = meta.endpoints.clone();
+
+    // Modify config to use local meta store
+    let mut ee_setup = EESetup::new();
+    let config = ee_setup.config_mut();
+    config.meta.endpoints = endpoints.clone();
+
+    // 2. Setup test fixture by using local meta store
+    let fixture = TestFixture::setup_with_custom(ee_setup).await?;
+
+    // Adjust retention period to 0, so that dropped tables will be vacuumed immediately
+    let session = fixture.default_session();
+    session.get_settings().set_data_retention_time_in_days(0)?;
+
+    // Enable row access policy feature
+    fixture
+        .execute_command("set global enable_experimental_row_access_policy = 1")
+        .await?;
+
+    // 3. Prepare test db and table
+    let ctx = fixture.new_query_ctx().await?;
+    let db_name = "test_vacuum_clean_policies";
+    let tbl_name = "t";
+    fixture
+        .execute_command(format!("create database {db_name}").as_str())
+        .await?;
+    fixture
+        .execute_command(
+            format!("create table {db_name}.{tbl_name} (id int, name string, email string)")
+                .as_str(),
+        )
+        .await?;
+
+    // 4. Create masking policies and row access policy
+    fixture
+        .execute_command(
+            "CREATE MASKING POLICY email_mask AS (val STRING) RETURNS STRING -> CASE WHEN current_role() = 'admin' THEN val ELSE '***' END",
+        )
+        .await?;
+
+    fixture
+        .execute_command(
+            "CREATE MASKING POLICY name_mask AS (val STRING) RETURNS STRING -> CASE WHEN current_role() = 'admin' THEN val ELSE 'REDACTED' END",
+        )
+        .await?;
+
+    fixture
+        .execute_command(
+            "CREATE ROW ACCESS POLICY row_filter AS (id int) RETURNS boolean -> current_role() = 'admin' OR id > 0",
+        )
+        .await?;
+
+    // 5. Apply policies to table columns
+    fixture
+        .execute_command(
+            format!(
+                "ALTER TABLE {db_name}.{tbl_name} MODIFY COLUMN email SET MASKING POLICY email_mask"
+            )
+            .as_str(),
+        )
+        .await?;
+
+    fixture
+        .execute_command(
+            format!(
+                "ALTER TABLE {db_name}.{tbl_name} MODIFY COLUMN name SET MASKING POLICY name_mask"
+            )
+            .as_str(),
+        )
+        .await?;
+
+    fixture
+        .execute_command(
+            format!("ALTER TABLE {db_name}.{tbl_name} ADD ROW ACCESS POLICY row_filter ON (id)")
+                .as_str(),
+        )
+        .await?;
+
+    // 6. Get table ID and verify policy references exist
+    let tenant = ctx.get_tenant();
+    let table = ctx
+        .get_default_catalog()?
+        .get_table(&tenant, db_name, tbl_name)
+        .await?;
+    let table_id = table.get_id();
+
+    // Get policy IDs from table metadata
+    let table_meta = table.get_table_info().meta.clone();
+
+    // Verify masking policy references exist
+    assert!(
+        !table_meta.column_mask_policy_columns_ids.is_empty(),
+        "masking policy references should exist after applying policies"
+    );
+
+    // Verify row access policy reference exists
+    assert!(
+        table_meta.row_access_policy_columns_ids.is_some(),
+        "row access policy reference should exist after applying policy"
+    );
+
+    let mask_policy_ids: Vec<u64> = table_meta
+        .column_mask_policy_columns_ids
+        .values()
+        .map(|policy_map| policy_map.policy_id)
+        .collect();
+
+    let row_policy_id = table_meta
+        .row_access_policy_columns_ids
+        .as_ref()
+        .unwrap()
+        .policy_id;
+
+    // Verify policy references in meta store
+    for policy_id in &mask_policy_ids {
+        let mask_policy_key =
+            MaskPolicyTableIdIdent::new_generic(tenant.clone(), MaskPolicyIdTableId {
+                policy_id: *policy_id,
+                table_id,
+            });
+        let v = meta.get_pb(&mask_policy_key).await?;
+        assert!(
+            v.is_some(),
+            "masking policy reference should exist in meta store"
+        );
+    }
+
+    let row_policy_key =
+        RowAccessPolicyTableIdIdent::new_generic(tenant.clone(), RowAccessPolicyIdTableId {
+            policy_id: row_policy_id,
+            table_id,
+        });
+    let v = meta.get_pb(&row_policy_key).await?;
+    assert!(
+        v.is_some(),
+        "row access policy reference should exist in meta store"
+    );
+
+    // 7. Drop database (this will mark both database and table as dropped)
+    fixture
+        .execute_command(format!("drop database {db_name}").as_str())
+        .await?;
+
+    // 8. Vacuum dropped tables
+    fixture.execute_command("vacuum drop table").await?;
+
+    // 9. Ensure that policy references are cleaned up
+    for policy_id in &mask_policy_ids {
+        let mask_policy_key =
+            MaskPolicyTableIdIdent::new_generic(tenant.clone(), MaskPolicyIdTableId {
+                policy_id: *policy_id,
+                table_id,
+            });
+        let v = meta.get_pb(&mask_policy_key).await?;
+        assert!(
+            v.is_none(),
+            "masking policy reference should be cleaned up after vacuum"
+        );
+    }
+
+    let row_policy_key =
+        RowAccessPolicyTableIdIdent::new_generic(tenant.clone(), RowAccessPolicyIdTableId {
+            policy_id: row_policy_id,
+            table_id,
+        });
+    let v = meta.get_pb(&row_policy_key).await?;
+    assert!(
+        v.is_none(),
+        "row access policy reference should be cleaned up after vacuum"
+    );
+
+    Ok(())
+}