Ensure writeToPolicyMappingRecord update existing record if primary key equals in EclipseLink Persistence Impl (#1469)

* update PolicyMappingRecord if not exists

* update test

* add TODO

Author: HonahX

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(

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;

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

@@ -1001,9 +1001,19 @@ void attachPolicyToTarget(
       PolarisBaseEntity target,
       List<PolarisEntityCore> policyCatalogPath,
       PolicyEntity policy) {
+    attachPolicyToTarget(targetCatalogPath, target, policyCatalogPath, policy, null);
+  }
+
+  void attachPolicyToTarget(
+      List<PolarisEntityCore> targetCatalogPath,
+      PolarisBaseEntity target,
+      List<PolarisEntityCore> policyCatalogPath,
+      PolicyEntity policy,
+      Map<String, String> 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 */
@@ -1022,8 +1032,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<String, String> parameters) {
     target =
         polarisMetaStoreManager
             .loadEntity(
@@ -1048,15 +1060,15 @@ 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 =
         polarisMetaStoreManager.loadPoliciesOnEntityByType(
             this.polarisCallContext, target, policy.getPolicyType());
     validateLoadedPolicyMappings(loadPolicyMappingsResultByType);
     checkPolicyMappingRecordExists(
-        loadPolicyMappingsResultByType.getPolicyMappingRecords(), target, policy);
+        loadPolicyMappingsResultByType.getPolicyMappingRecords(), target, policy, parameters);
   }
 
   /**
@@ -1140,8 +1152,9 @@ void validateLoadedPolicyMappings(LoadPolicyMappingsResult loadPolicyMappingReco
   void checkPolicyMappingRecordExists(
       List<PolarisPolicyMappingRecord> policyMappingRecords,
       PolarisBaseEntity target,
-      PolicyEntity policy) {
-    boolean exists = isPolicyMappingRecordExists(policyMappingRecords, target, policy);
+      PolicyEntity policy,
+      Map<String, String> parameters) {
+    boolean exists = isPolicyMappingRecordExists(policyMappingRecords, target, policy, parameters);
     Assertions.assertThat(exists).isTrue();
   }
 
@@ -1184,6 +1197,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<PolarisPolicyMappingRecord> policyMappingRecords,
+      PolarisBaseEntity target,
+      PolicyEntity policy,
+      Map<String, String> 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):
@@ -2781,6 +2820,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);