refactor to different function

Author: singhpk234

extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java

@@ -721,47 +721,13 @@ public void writeToPolicyMappingRecords(
             PolicyType policyType = PolicyType.fromCode(record.getPolicyTypeCode());
             Preconditions.checkArgument(
                 policyType != null, "Invalid policy type code: %s", record.getPolicyTypeCode());
-            String insertQuery =
+            String insertPolicyMappingQuery =
                 generateInsertQuery(
                     ModelPolicyMappingRecord.fromPolicyMappingRecord(record), realmId);
             if (policyType.isInheritable()) {
-              List<PolarisPolicyMappingRecord> existingRecords =
-                  loadPoliciesOnTargetByType(
-                      callCtx,
-                      record.getTargetCatalogId(),
-                      record.getTargetId(),
-                      record.getPolicyTypeCode());
-              if (existingRecords.size() > 1) {
-                throw new PolicyMappingAlreadyExistsException(existingRecords.getFirst());
-              } else if (existingRecords.size() == 1) {
-                PolarisPolicyMappingRecord existingRecord = existingRecords.getFirst();
-                if (existingRecord.getPolicyCatalogId() != record.getPolicyCatalogId()
-                    || existingRecord.getPolicyId() != record.getPolicyId()) {
-                  throw new PolicyMappingAlreadyExistsException(existingRecord);
-                }
-                Map<String, Object> updateClause =
-                    Map.of(
-                        "target_catalog_id",
-                        record.getTargetCatalogId(),
-                        "target_id",
-                        record.getTargetId(),
-                        "policy_type_code",
-                        record.getPolicyTypeCode(),
-                        "policy_id",
-                        record.getPolicyId(),
-                        "policy_catalog_id",
-                        record.getPolicyCatalogId(),
-                        "realm_id",
-                        realmId);
-                String updateQuery =
-                    generateUpdateQuery(
-                        ModelPolicyMappingRecord.fromPolicyMappingRecord(record), updateClause);
-                statement.executeUpdate(updateQuery);
-              } else {
-                statement.executeUpdate(insertQuery);
-              }
+              return handleInheritablePolicy(callCtx, record, insertPolicyMappingQuery, statement);
             } else {
-              statement.executeUpdate(insertQuery);
+              statement.executeUpdate(insertPolicyMappingQuery);
             }
             return true;
           });
@@ -771,6 +737,51 @@ public void writeToPolicyMappingRecords(
     }
   }
 
+  private boolean handleInheritablePolicy(
+      @Nonnull PolarisCallContext callCtx,
+      @Nonnull PolarisPolicyMappingRecord record,
+      @Nonnull String insertQuery,
+      Statement statement)
+      throws SQLException {
+    List<PolarisPolicyMappingRecord> existingRecords =
+        loadPoliciesOnTargetByType(
+            callCtx, record.getTargetCatalogId(), record.getTargetId(), record.getPolicyTypeCode());
+    if (existingRecords.size() > 1) {
+      throw new PolicyMappingAlreadyExistsException(existingRecords.getFirst());
+    } else if (existingRecords.size() == 1) {
+      PolarisPolicyMappingRecord existingRecord = existingRecords.getFirst();
+      if (existingRecord.getPolicyCatalogId() != record.getPolicyCatalogId()
+          || existingRecord.getPolicyId() != record.getPolicyId()) {
+        // Only one policy of the same type can be attached to an entity when the policy is
+        // inheritable.
+        throw new PolicyMappingAlreadyExistsException(existingRecord);
+      }
+      Map<String, Object> updateClause =
+          Map.of(
+              "target_catalog_id",
+              record.getTargetCatalogId(),
+              "target_id",
+              record.getTargetId(),
+              "policy_type_code",
+              record.getPolicyTypeCode(),
+              "policy_id",
+              record.getPolicyId(),
+              "policy_catalog_id",
+              record.getPolicyCatalogId(),
+              "realm_id",
+              realmId);
+      // In case of the mapping exist, update the policy mapping with the new parameters.
+      String updateQuery =
+          generateUpdateQuery(
+              ModelPolicyMappingRecord.fromPolicyMappingRecord(record), updateClause);
+      statement.executeUpdate(updateQuery);
+    } else {
+      // record doesn't exist do an insert.
+      statement.executeUpdate(insertQuery);
+    }
+    return true;
+  }
+
   @Override
   public void deleteFromPolicyMappingRecords(
       @Nonnull PolarisCallContext callCtx, @Nonnull PolarisPolicyMappingRecord record) {

extension/persistence/relational-jdbc/src/main/resources/h2/schema-v1.sql

@@ -104,4 +104,4 @@ CREATE TABLE IF NOT EXISTS policy_mapping_record (
     PRIMARY KEY (realm_id, target_catalog_id, target_id, policy_type_code, policy_catalog_id, policy_id)
 );
 
-CREATE INDEX IF NOT EXISTS idx_policy_mapping_record ON policy_mapping_record (realm_id, policy_type_code, policy_catalog_id, policy_id, policy_type_code, target_catalog_id, target_id);
+CREATE INDEX IF NOT EXISTS idx_policy_mapping_record ON policy_mapping_record (realm_id, policy_type_code, policy_catalog_id, policy_id, target_catalog_id, target_id);

extension/persistence/relational-jdbc/src/main/resources/postgres/schema-v1.sql

@@ -104,4 +104,4 @@ CREATE TABLE IF NOT EXISTS policy_mapping_record (
     PRIMARY KEY (realm_id, target_catalog_id, target_id, policy_type_code, policy_catalog_id, policy_id)
 );
 
-CREATE INDEX IF NOT EXISTS idx_policy_mapping_record ON policy_mapping_record (realm_id, policy_type_code, policy_catalog_id, policy_id, policy_type_code, target_catalog_id, target_id);
+CREATE INDEX IF NOT EXISTS idx_policy_mapping_record ON policy_mapping_record (realm_id, policy_type_code, policy_catalog_id, policy_id, target_catalog_id, target_id);