Policy Store: Check whether Policy is in use before dropping and support detach-all flag

integration-tests/src/main/java/org/apache/polaris/service/it/env/PolicyApi.java

@@ -28,6 +28,7 @@
 import org.apache.iceberg.catalog.Namespace;
 import org.apache.iceberg.rest.RESTUtil;
 import org.apache.polaris.core.policy.PolicyType;
+import org.apache.polaris.core.policy.exceptions.PolicyInUseException;
 import org.apache.polaris.service.types.ApplicablePolicy;
 import org.apache.polaris.service.types.AttachPolicyRequest;
 import org.apache.polaris.service.types.CreatePolicyRequest;
@@ -72,12 +73,24 @@ public List<PolicyIdentifier> listPolicies(String catalog, Namespace namespace,
   }
 
   public void dropPolicy(String catalog, PolicyIdentifier policyIdentifier) {
+    dropPolicy(catalog, policyIdentifier, null);
+  }
+
+  public void dropPolicy(String catalog, PolicyIdentifier policyIdentifier, Boolean detachAll) {
     String ns = RESTUtil.encodeNamespace(policyIdentifier.getNamespace());
+    Map<String, String> queryParams = new HashMap<>();
+    if (detachAll != null) {
+      queryParams.put("detach-all", detachAll.toString());
+    }
     try (Response res =
         request(
                 "polaris/v1/{cat}/namespaces/{ns}/policies/{policy}",
-                Map.of("cat", catalog, "ns", ns, "policy", policyIdentifier.getName()))
+                Map.of("cat", catalog, "ns", ns, "policy", policyIdentifier.getName()),
+                queryParams)
             .delete()) {
+      if (res.getStatus() == Response.Status.BAD_REQUEST.getStatusCode()) {
+        throw new PolicyInUseException("Policy in use");
+      }
       Assertions.assertThat(res.getStatus()).isEqualTo(Response.Status.NO_CONTENT.getStatusCode());
     }
   }

integration-tests/src/main/java/org/apache/polaris/service/it/test/PolarisPolicyServiceIntegrationTest.java

@@ -50,6 +50,7 @@
 import org.apache.polaris.core.catalog.PolarisCatalogHelpers;
 import org.apache.polaris.core.entity.CatalogEntity;
 import org.apache.polaris.core.policy.PredefinedPolicyTypes;
+import org.apache.polaris.core.policy.exceptions.PolicyInUseException;
 import org.apache.polaris.service.it.env.ClientCredentials;
 import org.apache.polaris.service.it.env.IcebergHelper;
 import org.apache.polaris.service.it.env.IntegrationTestsHelper;
@@ -268,8 +269,21 @@ public void testDropPolicy() {
         PredefinedPolicyTypes.DATA_COMPACTION,
         EXAMPLE_TABLE_MAINTENANCE_POLICY_CONTENT,
         "test policy");
-    policyApi.dropPolicy(currentCatalogName, NS1_P1);
+
+    PolicyAttachmentTarget catalogTarget =
+        PolicyAttachmentTarget.builder().setType(PolicyAttachmentTarget.TypeEnum.CATALOG).build();
+    policyApi.attachPolicy(currentCatalogName, NS1_P1, catalogTarget, Map.of());
+
+    // dropPolicy should fail because the policy is attached to the catalog
+    Assertions.assertThatThrownBy(() -> policyApi.dropPolicy(currentCatalogName, NS1_P1))
+        .isInstanceOf(PolicyInUseException.class);
+
+    // with detach-all=true, the policy and the attachment should be dropped
+    policyApi.dropPolicy(currentCatalogName, NS1_P1, true);
     Assertions.assertThat(policyApi.listPolicies(currentCatalogName, NS1)).hasSize(0);
+    // The policy mapping record should be dropped
+    Assertions.assertThat(policyApi.getApplicablePolicies(currentCatalogName, null, null, null))
+        .hasSize(0);
   }
 
   @Test

polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java

@@ -65,6 +65,7 @@
 import org.apache.polaris.core.persistence.dao.entity.ValidateAccessResult;
 import org.apache.polaris.core.policy.PolarisPolicyMappingRecord;
 import org.apache.polaris.core.policy.PolicyEntity;
+import org.apache.polaris.core.policy.PolicyMappingUtil;
 import org.apache.polaris.core.policy.PolicyType;
 import org.apache.polaris.core.storage.PolarisCredentialProperty;
 import org.apache.polaris.core.storage.PolarisStorageActions;
@@ -193,6 +194,27 @@ private void dropEntity(
         ms.loadAllGrantRecordsOnSecurable(callCtx, entity.getCatalogId(), entity.getId());
     ms.deleteAllEntityGrantRecords(callCtx, entity, grantsOnGrantee, grantsOnSecurable);
 
+    if (entity.getType() == PolarisEntityType.POLICY
+        || PolicyMappingUtil.isValidTargetEntityType(entity.getType(), entity.getSubType())) {
+      // Best-effort cleanup - for policy and potential target entities, drop all policy mapping
+      // records related
+      // TODO: Support some more formal garbage-collection mechanism similar to grant records case
+      // above
+      try {
+        final List<PolarisPolicyMappingRecord> mappingOnPolicy =
+            (entity.getType() == PolarisEntityType.POLICY)
+                ? ms.loadAllTargetsOnPolicy(callCtx, entity.getCatalogId(), entity.getId())
+                : List.of();
+        final List<PolarisPolicyMappingRecord> mappingOnTarget =
+            (entity.getType() == PolarisEntityType.POLICY)
+                ? List.of()
+                : ms.loadAllPoliciesOnTarget(callCtx, entity.getCatalogId(), entity.getId());
+        ms.deleteAllEntityPolicyMappingRecords(callCtx, entity, mappingOnTarget, mappingOnPolicy);
+      } catch (UnsupportedOperationException e) {
+        // Policy mapping persistence not implemented, but we should not block dropping entities
+      }
+    }
+
     // Now determine the set of entities on the other side of the grants we just removed. Grants
     // from/to these entities has been removed, hence we need to update the grant version of
     // each entity. Collect the id of each.
@@ -1179,6 +1201,18 @@ private void revokeGrantRecord(
           callCtx, null, refreshEntityToDrop.getCatalogId(), refreshEntityToDrop.getId())) {
         return new DropEntityResult(BaseResult.ReturnStatus.NAMESPACE_NOT_EMPTY, null);
       }
+    } else if (refreshEntityToDrop.getType() == PolarisEntityType.POLICY && !cleanup) {
+      try {
+        //  need to check if the policy is attached to any entity
+        List<PolarisPolicyMappingRecord> records =
+            ms.loadAllTargetsOnPolicy(
+                callCtx, refreshEntityToDrop.getCatalogId(), refreshEntityToDrop.getId());
+        if (!records.isEmpty()) {
+          return new DropEntityResult(BaseResult.ReturnStatus.POLICY_HAS_MAPPINGS, null);
+        }
+      } catch (UnsupportedOperationException e) {
+        // Policy mapping persistence not implemented, but we should not block dropping entities
+      }
     }
 
     // simply delete that entity. Will be removed from entities_active, added to the
@@ -1188,7 +1222,7 @@ private void revokeGrantRecord(
     // if cleanup, schedule a cleanup task for the entity. do this here, so that drop and scheduling
     // the cleanup task is transactional. Otherwise, we'll be unable to schedule the cleanup task
     // later
-    if (cleanup) {
+    if (cleanup && refreshEntityToDrop.getType() != PolarisEntityType.POLICY) {
       PolarisBaseEntity taskEntity =
           new PolarisEntity.Builder()
               .setId(ms.generateNewId(callCtx))

polaris-core/src/main/java/org/apache/polaris/core/persistence/dao/entity/BaseResult.java

@@ -117,6 +117,9 @@ public enum ReturnStatus {
 
     // policy mapping of same type already exists
     POLICY_MAPPING_OF_SAME_TYPE_ALREADY_EXISTS(15),
+
+    // policy has mappings and cannot be dropped
+    POLICY_HAS_MAPPINGS(16),
     ;
 
     // code for the enum

polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TransactionalMetaStoreManagerImpl.java

@@ -66,6 +66,7 @@
 import org.apache.polaris.core.persistence.dao.entity.ValidateAccessResult;
 import org.apache.polaris.core.policy.PolarisPolicyMappingRecord;
 import org.apache.polaris.core.policy.PolicyEntity;
+import org.apache.polaris.core.policy.PolicyMappingUtil;
 import org.apache.polaris.core.policy.PolicyType;
 import org.apache.polaris.core.storage.PolarisCredentialProperty;
 import org.apache.polaris.core.storage.PolarisStorageActions;
@@ -195,6 +196,28 @@ private void dropEntity(
       ms.writeEntityInCurrentTxn(callCtx, entityGrantChanged, false, originalEntity);
     }
 
+    if (entity.getType() == PolarisEntityType.POLICY
+        || PolicyMappingUtil.isValidTargetEntityType(entity.getType(), entity.getSubType())) {
+      // Best-effort cleanup - for policy and potential target entities, drop all policy mapping
+      // records related
+      try {
+        final List<PolarisPolicyMappingRecord> mappingOnPolicy =
+            (entity.getType() == PolarisEntityType.POLICY)
+                ? ms.loadAllTargetsOnPolicyInCurrentTxn(
+                    callCtx, entity.getCatalogId(), entity.getId())
+                : List.of();
+        final List<PolarisPolicyMappingRecord> mappingOnTarget =
+            (entity.getType() == PolarisEntityType.POLICY)
+                ? List.of()
+                : ms.loadAllPoliciesOnTargetInCurrentTxn(
+                    callCtx, entity.getCatalogId(), entity.getId());
+        ms.deleteAllEntityPolicyMappingRecordsInCurrentTxn(
+            callCtx, entity, mappingOnTarget, mappingOnPolicy);
+      } catch (UnsupportedOperationException e) {
+        // Policy mapping persistence not implemented, but we should not block dropping entities
+      }
+    }
+
     // remove the entity being dropped now
     ms.deleteEntityInCurrentTxn(callCtx, entity);
 
@@ -1361,6 +1384,18 @@ private void bootstrapPolarisService(
           callCtx, null, refreshEntityToDrop.getCatalogId(), refreshEntityToDrop.getId())) {
         return new DropEntityResult(BaseResult.ReturnStatus.NAMESPACE_NOT_EMPTY, null);
       }
+    } else if (refreshEntityToDrop.getType() == PolarisEntityType.POLICY && !cleanup) {
+      // need to check if the policy is attached to any entity
+      try {
+        List<PolarisPolicyMappingRecord> records =
+            ms.loadAllTargetsOnPolicyInCurrentTxn(
+                callCtx, refreshEntityToDrop.getCatalogId(), refreshEntityToDrop.getId());
+        if (!records.isEmpty()) {
+          return new DropEntityResult(BaseResult.ReturnStatus.POLICY_HAS_MAPPINGS, null);
+        }
+      } catch (UnsupportedOperationException e) {
+        // Policy mapping persistence not implemented, but we should not block dropping entities
+      }
     }
 
     // simply delete that entity. Will be removed from entities_active, added to the
@@ -1370,7 +1405,7 @@ private void bootstrapPolarisService(
     // if cleanup, schedule a cleanup task for the entity. do this here, so that drop and scheduling
     // the cleanup task is transactional. Otherwise, we'll be unable to schedule the cleanup task
     // later
-    if (cleanup) {
+    if (cleanup && refreshEntityToDrop.getType() != PolarisEntityType.POLICY) {
       PolarisBaseEntity taskEntity =
           new PolarisEntity.Builder()
               .setId(ms.generateNewIdInCurrentTxn(callCtx))

polaris-core/src/main/java/org/apache/polaris/core/persistence/transactional/TreeMapTransactionalPersistenceImpl.java

@@ -589,9 +589,9 @@ public void deleteAllEntityPolicyMappingRecordsInCurrentTxn(
 
     // also delete the other side. We need to delete these mapping one at a time versus doing a
     // range delete
-    mappingOnTarget.forEach(record -> this.store.getSlicePolicyMappingRecords().delete(record));
-    mappingOnPolicy.forEach(
+    mappingOnTarget.forEach(
         record -> this.store.getSlicePolicyMappingRecordsByPolicy().delete(record));
+    mappingOnPolicy.forEach(record -> this.store.getSlicePolicyMappingRecords().delete(record));
   }
 
   /** {@inheritDoc} */

polaris-core/src/main/java/org/apache/polaris/core/policy/PolicyMappingUtil.java

@@ -0,0 +1,57 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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.
+ */
+package org.apache.polaris.core.policy;
+
+import static org.apache.polaris.core.entity.PolarisEntityType.CATALOG;
+import static org.apache.polaris.core.entity.PolarisEntityType.NAMESPACE;
+import static org.apache.polaris.core.entity.PolarisEntityType.TABLE_LIKE;
+
+import java.util.Set;
+import org.apache.polaris.core.entity.PolarisEntitySubType;
+import org.apache.polaris.core.entity.PolarisEntityType;
+
+public class PolicyMappingUtil {
+
+  private static final Set<PolarisEntityType> ATTACHABLE_ENTITY_TYPES =
+      Set.of(CATALOG, NAMESPACE, TABLE_LIKE);
+
+  /**
+   * Checks if the given entity type and subtype are valid targets for policy attachment.
+   *
+   * <p>This method excludes non-attachable entities such as catalog-role and task. Each policy type
+   * may impose additional specific rules * to determine whether it can target a particular entity
+   * type or subtype.
+   */
+  public static boolean isValidTargetEntityType(
+      PolarisEntityType entityType, PolarisEntitySubType entitySubType) {
+    if (entityType == null) {
+      return false;
+    }
+
+    if (!ATTACHABLE_ENTITY_TYPES.contains(entityType)) {
+      return false;
+    }
+
+    if (entityType == TABLE_LIKE && entitySubType != PolarisEntitySubType.ICEBERG_TABLE) {
+      return false;
+    }
+
+    return true;
+  }
+}

polaris-core/src/main/java/org/apache/polaris/core/policy/exceptions/PolicyInUseException.java

@@ -0,0 +1,34 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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.
+ */
+package org.apache.polaris.core.policy.exceptions;
+
+import com.google.errorprone.annotations.FormatMethod;
+import org.apache.polaris.core.exceptions.PolarisException;
+
+public class PolicyInUseException extends PolarisException {
+  @FormatMethod
+  public PolicyInUseException(String message, Object... args) {
+    super(String.format(message, args));
+  }
+
+  @FormatMethod
+  public PolicyInUseException(Throwable cause, String message, Object... args) {
+    super(String.format(message, args), cause);
+  }
+}

polaris-core/src/main/java/org/apache/polaris/core/policy/validator/maintenance/BaseMaintenancePolicyValidator.java

@@ -18,40 +18,21 @@
  */
 package org.apache.polaris.core.policy.validator.maintenance;
 
-import static org.apache.polaris.core.entity.PolarisEntityType.CATALOG;
-import static org.apache.polaris.core.entity.PolarisEntityType.NAMESPACE;
-import static org.apache.polaris.core.entity.PolarisEntityType.TABLE_LIKE;
-
-import java.util.Set;
 import org.apache.polaris.core.entity.PolarisEntitySubType;
 import org.apache.polaris.core.entity.PolarisEntityType;
+import org.apache.polaris.core.policy.PolicyMappingUtil;
 import org.apache.polaris.core.policy.validator.InvalidPolicyException;
 import org.apache.polaris.core.policy.validator.PolicyValidator;
 
 public class BaseMaintenancePolicyValidator implements PolicyValidator {
   public static final BaseMaintenancePolicyValidator INSTANCE =
       new BaseMaintenancePolicyValidator();
 
-  private static final Set<PolarisEntityType> ATTACHABLE_ENTITY_TYPES =
-      Set.of(CATALOG, NAMESPACE, TABLE_LIKE);
-
   @Override
   public void validate(String content) throws InvalidPolicyException {}
 
   @Override
   public boolean canAttach(PolarisEntityType entityType, PolarisEntitySubType entitySubType) {
-    if (entityType == null) {
-      return false;
-    }
-
-    if (!ATTACHABLE_ENTITY_TYPES.contains(entityType)) {
-      return false;
-    }
-
-    if (entityType == TABLE_LIKE && entitySubType != PolarisEntitySubType.ICEBERG_TABLE) {
-      return false;
-    }
-
-    return true;
+    return PolicyMappingUtil.isValidTargetEntityType(entityType, entitySubType);
   }
 }

polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/BasePolarisMetaStoreManagerTest.java

@@ -292,6 +292,11 @@ protected void testPolicyMapping() {
     polarisTestMetaStoreManager.testPolicyMapping();
   }
 
+  @Test
+  protected void testPolicyMappingCleanup() {
+    polarisTestMetaStoreManager.testPolicyMappingCleanup();
+  }
+
   @Test
   protected void testLoadTasks() {
     for (int i = 0; i < 20; i++) {