diff --git a/extension/persistence/eclipselink/src/main/java/org/apache/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkStore.java b/extension/persistence/eclipselink/src/main/java/org/apache/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkStore.java index 9ebb4e8164..0988dcb7f6 100644 --- a/extension/persistence/eclipselink/src/main/java/org/apache/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkStore.java +++ b/extension/persistence/eclipselink/src/main/java/org/apache/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkStore.java @@ -420,7 +420,22 @@ void writeToPolicyMappingRecords( diagnosticServices.check(session != null, "session_is_null"); checkInitialized(); - session.persist(ModelPolicyMappingRecord.fromPolicyMappingRecord(mappingRecord)); + // TODO: combine existence check and write into one statement + ModelPolicyMappingRecord model = + lookupPolicyMappingRecord( + session, + mappingRecord.getTargetCatalogId(), + mappingRecord.getTargetId(), + mappingRecord.getPolicyTypeCode(), + mappingRecord.getPolicyCatalogId(), + mappingRecord.getPolicyId()); + if (model != null) { + model.update(mappingRecord); + } else { + model = ModelPolicyMappingRecord.fromPolicyMappingRecord(mappingRecord); + } + + session.persist(model); } void deleteFromPolicyMappingRecords( diff --git a/extension/persistence/jpa-model/src/main/java/org/apache/polaris/jpa/models/ModelPolicyMappingRecord.java b/extension/persistence/jpa-model/src/main/java/org/apache/polaris/jpa/models/ModelPolicyMappingRecord.java index c779758430..0448dec676 100644 --- a/extension/persistence/jpa-model/src/main/java/org/apache/polaris/jpa/models/ModelPolicyMappingRecord.java +++ b/extension/persistence/jpa-model/src/main/java/org/apache/polaris/jpa/models/ModelPolicyMappingRecord.java @@ -135,6 +135,15 @@ public ModelPolicyMappingRecord build() { } } + public void update(PolarisPolicyMappingRecord record) { + this.targetCatalogId = record.getTargetCatalogId(); + this.targetId = record.getTargetId(); + this.policyTypeCode = record.getPolicyTypeCode(); + this.policyCatalogId = record.getPolicyCatalogId(); + this.policyId = record.getPolicyId(); + this.parameters = record.getParameters(); + } + public static ModelPolicyMappingRecord fromPolicyMappingRecord( PolarisPolicyMappingRecord record) { if (record == null) return null; diff --git a/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/PolarisTestMetaStoreManager.java b/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/PolarisTestMetaStoreManager.java index b401d4efb3..8ab76e3c2c 100644 --- a/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/PolarisTestMetaStoreManager.java +++ b/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/PolarisTestMetaStoreManager.java @@ -995,9 +995,19 @@ void attachPolicyToTarget( PolarisBaseEntity target, List policyCatalogPath, PolicyEntity policy) { + attachPolicyToTarget(targetCatalogPath, target, policyCatalogPath, policy, null); + } + + void attachPolicyToTarget( + List targetCatalogPath, + PolarisBaseEntity target, + List policyCatalogPath, + PolicyEntity policy, + Map parameters) { polarisMetaStoreManager.attachPolicyToEntity( - polarisCallContext, targetCatalogPath, target, policyCatalogPath, policy, null); - ensurePolicyMappingRecordExists(target, policy); + polarisCallContext, targetCatalogPath, target, policyCatalogPath, policy, parameters); + + ensurePolicyMappingRecordExists(target, policy, parameters); } /** detach a policy from a target */ @@ -1016,8 +1026,10 @@ void detachPolicyFromTarget( * * @param target the target * @param policy the policy + * @param parameters the parameters */ - void ensurePolicyMappingRecordExists(PolarisBaseEntity target, PolicyEntity policy) { + void ensurePolicyMappingRecordExists( + PolarisBaseEntity target, PolicyEntity policy, Map parameters) { target = polarisMetaStoreManager .loadEntity( @@ -1042,7 +1054,7 @@ void ensurePolicyMappingRecordExists(PolarisBaseEntity target, PolicyEntity poli validateLoadedPolicyMappings(loadPolicyMappingsResult); checkPolicyMappingRecordExists( - loadPolicyMappingsResult.getPolicyMappingRecords(), target, policy); + loadPolicyMappingsResult.getPolicyMappingRecords(), target, policy, parameters); // also try load by specific type LoadPolicyMappingsResult loadPolicyMappingsResultByType = @@ -1050,7 +1062,7 @@ void ensurePolicyMappingRecordExists(PolarisBaseEntity target, PolicyEntity poli this.polarisCallContext, target, policy.getPolicyType()); validateLoadedPolicyMappings(loadPolicyMappingsResultByType); checkPolicyMappingRecordExists( - loadPolicyMappingsResultByType.getPolicyMappingRecords(), target, policy); + loadPolicyMappingsResultByType.getPolicyMappingRecords(), target, policy, parameters); } /** @@ -1134,8 +1146,9 @@ void validateLoadedPolicyMappings(LoadPolicyMappingsResult loadPolicyMappingReco void checkPolicyMappingRecordExists( List policyMappingRecords, PolarisBaseEntity target, - PolicyEntity policy) { - boolean exists = isPolicyMappingRecordExists(policyMappingRecords, target, policy); + PolicyEntity policy, + Map parameters) { + boolean exists = isPolicyMappingRecordExists(policyMappingRecords, target, policy, parameters); Assertions.assertThat(exists).isTrue(); } @@ -1178,6 +1191,32 @@ record -> return policyMappingCount == 1; } + /** + * Check if the policy mapping record exists and verify if the parameters also equals + * + * @param policyMappingRecords list of policy mapping records + * @param target the target + * @param policy the policy + * @param parameters the parameters + */ + boolean isPolicyMappingRecordExists( + List policyMappingRecords, + PolarisBaseEntity target, + PolicyEntity policy, + Map parameters) { + PolarisPolicyMappingRecord expected = + new PolarisPolicyMappingRecord( + target.getCatalogId(), + target.getId(), + policy.getCatalogId(), + policy.getId(), + policy.getPolicyTypeCode(), + parameters); + long policyMappingCount = + policyMappingRecords.stream().filter(record -> expected.equals(record)).count(); + return policyMappingCount == 1; + } + /** * Create a test catalog. This is a new catalog which will have the following objects (N is for a * namespace, T for a table, V for a view, R for a role, P for a principal, POL for a policy): @@ -2775,6 +2814,10 @@ void testPolicyMapping() { Assertions.assertThat(policyAttachmentResult.getReturnStatus()) .isEqualTo(BaseResult.ReturnStatus.POLICY_MAPPING_OF_SAME_TYPE_ALREADY_EXISTS); + // Attach the same policy to same target again should succeed and replace the existing one + attachPolicyToTarget( + List.of(catalog, N1, N1_N2), N1_N2_T1, List.of(catalog, N1), N1_P1, Map.of("test", "test")); + LoadPolicyMappingsResult loadPolicyMappingsResult = polarisMetaStoreManager.loadPoliciesOnEntityByType( polarisCallContext, N1_N2_T1, PredefinedPolicyTypes.DATA_COMPACTION);