From 7a3a5c9f4b7f69891a889dd20cefbf0437387aa6 Mon Sep 17 00:00:00 2001 From: Prashant Date: Fri, 25 Apr 2025 13:42:14 -0700 Subject: [PATCH 1/8] policy store - 1 --- .../jdbc/JdbcBasePersistenceImpl.java | 137 ++++++++++++++- .../relational/jdbc/QueryGenerator.java | 31 +++- .../jdbc/models/ModelPolicyMappingRecord.java | 165 ++++++++++++++++++ .../src/main/resources/h2/schema-v1.sql | 14 ++ ...anagerWithJdbcBasePersistenceImplTest.java | 9 - .../relational/jdbc/QueryGeneratorTest.java | 8 +- 6 files changed, 342 insertions(+), 22 deletions(-) create mode 100644 extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/models/ModelPolicyMappingRecord.java diff --git a/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java b/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java index ea99ed98da..22c4559599 100644 --- a/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java +++ b/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java @@ -47,11 +47,13 @@ import org.apache.polaris.core.persistence.IntegrationPersistence; import org.apache.polaris.core.persistence.PrincipalSecretsGenerator; import org.apache.polaris.core.persistence.RetryOnConcurrencyException; +import org.apache.polaris.core.policy.PolarisPolicyMappingRecord; import org.apache.polaris.core.storage.PolarisStorageConfigurationInfo; import org.apache.polaris.core.storage.PolarisStorageIntegration; import org.apache.polaris.core.storage.PolarisStorageIntegrationProvider; import org.apache.polaris.extension.persistence.relational.jdbc.models.ModelEntity; import org.apache.polaris.extension.persistence.relational.jdbc.models.ModelGrantRecord; +import org.apache.polaris.extension.persistence.relational.jdbc.models.ModelPolicyMappingRecord; import org.apache.polaris.extension.persistence.relational.jdbc.models.ModelPrincipalAuthenticationData; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -220,7 +222,7 @@ public void deleteEntity(@Nonnull PolarisCallContext callCtx, @Nonnull PolarisBa public void deleteFromGrantRecords( @Nonnull PolarisCallContext callCtx, @Nonnull PolarisGrantRecord grantRec) { ModelGrantRecord modelGrantRecord = ModelGrantRecord.fromGrantRecord(grantRec); - String query = generateDeleteQuery(modelGrantRecord, ModelGrantRecord.class, realmId); + String query = generateDeleteQuery(modelGrantRecord, realmId); try { datasourceOperations.executeUpdate(query); } catch (SQLException e) { @@ -701,6 +703,139 @@ public void deletePrincipalSecrets( } } + @Override + public void writeToPolicyMappingRecords( + @Nonnull PolarisCallContext callCtx, @Nonnull PolarisPolicyMappingRecord record) { + ModelPolicyMappingRecord modelPolicyMappingRecord = + ModelPolicyMappingRecord.fromPolicyMappingRecord(record); + String query = generateInsertQuery(modelPolicyMappingRecord, realmId); + try { + datasourceOperations.executeUpdate(query); + } catch (SQLException e) { + throw new RuntimeException( + String.format("Failed to write to policy mapping records due to %s", e.getMessage()), e); + } + } + + @Override + public void deleteFromPolicyMappingRecords( + @Nonnull PolarisCallContext callCtx, @Nonnull PolarisPolicyMappingRecord record) { + ModelPolicyMappingRecord modelPolicyMappingRecord = + ModelPolicyMappingRecord.fromPolicyMappingRecord(record); + String query = generateDeleteQuery(modelPolicyMappingRecord, realmId); + try { + datasourceOperations.executeUpdate(query); + } catch (SQLException e) { + throw new RuntimeException( + String.format("Failed to write to policy records due to %s", e.getMessage()), e); + } + } + + @Override + public void deleteAllEntityPolicyMappingRecords( + @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisEntityCore entity, + @Nonnull List mappingOnTarget, + @Nonnull List mappingOnPolicy) { + try { + datasourceOperations.executeUpdate( + generateDeleteQueryForEntityPolicyMappingRecords(entity, realmId)); + } catch (SQLException e) { + throw new RuntimeException( + String.format("Failed to delete policy mapping records due to %s", e.getMessage()), e); + } + } + + @Nullable + @Override + public PolarisPolicyMappingRecord lookupPolicyMappingRecord( + @Nonnull PolarisCallContext callCtx, + long targetCatalogId, + long targetId, + int policyTypeCode, + long policyCatalogId, + long policyId) { + Map params = + Map.of( + "target_catalog_id", + targetCatalogId, + "target_id", + targetId, + "policy_type_code", + policyTypeCode, + "policy_id", + policyId, + "policy_catalog_id", + policyCatalogId, + "realm_id", + realmId); + String query = generateSelectQuery(ModelPolicyMappingRecord.class, params); + List results = fetchPolicyMappingRecords(query); + if (results.size() > 1) { + throw new IllegalStateException( + String.format( + "More than 1 policy %s for a given policy mapping", results.getFirst())); + } + return results.size() == 1 ? results.getFirst() : null; + } + + @Nonnull + @Override + public List loadPoliciesOnTargetByType( + @Nonnull PolarisCallContext callCtx, + long targetCatalogId, + long targetId, + int policyTypeCode) { + Map params = + Map.of( + "target_catalog_id", + targetCatalogId, + "target_id", + targetId, + "policy_type_code", + policyTypeCode, + "realm_id", + realmId); + String query = generateSelectQuery(ModelPolicyMappingRecord.class, params); + return fetchPolicyMappingRecords(query); + } + + @Nonnull + @Override + public List loadAllPoliciesOnTarget( + @Nonnull PolarisCallContext callCtx, long targetCatalogId, long targetId) { + Map params = + Map.of("target_catalog_id", targetCatalogId, "target_id", targetId, "realm_id", realmId); + String query = generateSelectQuery(ModelPolicyMappingRecord.class, params); + return fetchPolicyMappingRecords(query); + } + + @Nonnull + @Override + public List loadAllTargetsOnPolicy( + @Nonnull PolarisCallContext callCtx, long policyCatalogId, long policyId) { + Map params = + Map.of("policy_catalog_id", policyCatalogId, "policy_id", policyId, "realm_id", realmId); + String query = generateSelectQuery(ModelPolicyMappingRecord.class, params); + return fetchPolicyMappingRecords(query); + } + + private List fetchPolicyMappingRecords(String query) { + try { + List results = + datasourceOperations.executeSelect( + query, + ModelPolicyMappingRecord.class, + ModelPolicyMappingRecord::toPolicyMappingRecord, + null, + Integer.MAX_VALUE); + return results == null ? Collections.emptyList() : results; + } catch (SQLException e) { + throw new RuntimeException( + String.format("Failed to retrieve policy mapping records %s", e.getMessage()), e); + } + } + @Nullable @Override public diff --git a/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/QueryGenerator.java b/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/QueryGenerator.java index 07e07badb9..aff9c061ed 100644 --- a/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/QueryGenerator.java +++ b/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/QueryGenerator.java @@ -26,10 +26,7 @@ import java.util.Map; import org.apache.polaris.core.entity.PolarisEntityCore; import org.apache.polaris.core.entity.PolarisEntityId; -import org.apache.polaris.extension.persistence.relational.jdbc.models.Converter; -import org.apache.polaris.extension.persistence.relational.jdbc.models.ModelEntity; -import org.apache.polaris.extension.persistence.relational.jdbc.models.ModelGrantRecord; -import org.apache.polaris.extension.persistence.relational.jdbc.models.ModelPrincipalAuthenticationData; +import org.apache.polaris.extension.persistence.relational.jdbc.models.*; public class QueryGenerator { @@ -59,6 +56,26 @@ public static String generateDeleteQueryForEntityGrantRecords( return generateDeleteQuery(ModelGrantRecord.class, whereClause); } + public static String generateDeleteQueryForEntityPolicyMappingRecords( + @Nonnull PolarisEntityCore entity, @Nonnull String realmId) { + String targetCondition = + String.format( + "target_id = %s AND target_catalog_id = %s", entity.getId(), entity.getCatalogId()); + String sourceCondition = + String.format( + "policy_id = %s AND policy_catalog_id = %s", entity.getId(), entity.getCatalogId()); + + String whereClause = + " WHERE (" + + targetCondition + + " OR " + + sourceCondition + + ") AND realm_id = '" + + realmId + + "'"; + return generateDeleteQuery(ModelGrantRecord.class, whereClause); + } + public static String generateSelectQueryWithEntityIds( @Nonnull String realmId, @Nonnull List entityIds) { if (entityIds.isEmpty()) { @@ -127,8 +144,8 @@ public static String generateDeleteAll(@Nonnull Class entityClass, @Nonnull S } public static String generateDeleteQuery( - @Nonnull Converter entity, @Nonnull Class entityClass, @Nonnull String realmId) { - String tableName = getTableName(entityClass); + @Nonnull Converter entity, @Nonnull String realmId) { + String tableName = getTableName(entity.getClass()); Map objMap = entity.toMap(); objMap.put("realm_id", realmId); String whereConditions = generateWhereClause(objMap); @@ -186,6 +203,8 @@ public static String getTableName(@Nonnull Class entityClass) { tableName = "GRANT_RECORDS"; } else if (entityClass.equals(ModelPrincipalAuthenticationData.class)) { tableName = "PRINCIPAL_AUTHENTICATION_DATA"; + } else if (entityClass.equals(ModelPolicyMappingRecord.class)) { + tableName = "POLICY_MAPPING_RECORD"; } else { throw new IllegalArgumentException("Unsupported entity class: " + entityClass.getName()); } diff --git a/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/models/ModelPolicyMappingRecord.java b/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/models/ModelPolicyMappingRecord.java new file mode 100644 index 0000000000..0913a3570e --- /dev/null +++ b/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/models/ModelPolicyMappingRecord.java @@ -0,0 +1,165 @@ +/* + * 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.extension.persistence.relational.jdbc.models; + +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.HashMap; +import java.util.Map; +import org.apache.polaris.core.policy.PolarisPolicyMappingRecord; + +public class ModelPolicyMappingRecord implements Converter { + // id of the catalog where target entity resides + private long targetCatalogId; + + // id of the target entity + private long targetId; + + // id associated to the policy type + private int policyTypeCode; + + // id of the catalog where the policy entity resides + private long policyCatalogId; + + // id of the policy + private long policyId; + + // additional parameters of the mapping + private String parameters; + + public long getTargetCatalogId() { + return targetCatalogId; + } + + public long getTargetId() { + return targetId; + } + + public int getPolicyTypeCode() { + return policyTypeCode; + } + + public long getPolicyCatalogId() { + return policyCatalogId; + } + + public long getPolicyId() { + return policyId; + } + + public String getParameters() { + return parameters; + } + + public static ModelPolicyMappingRecord.Builder builder() { + return new ModelPolicyMappingRecord.Builder(); + } + + public static final class Builder { + private final ModelPolicyMappingRecord policyMappingRecord; + + private Builder() { + policyMappingRecord = new ModelPolicyMappingRecord(); + } + + public Builder targetCatalogId(long targetCatalogId) { + policyMappingRecord.targetCatalogId = targetCatalogId; + return this; + } + + public Builder targetId(long targetId) { + policyMappingRecord.targetId = targetId; + return this; + } + + public Builder policyTypeCode(int policyTypeCode) { + policyMappingRecord.policyTypeCode = policyTypeCode; + return this; + } + + public Builder policyCatalogId(long policyCatalogId) { + policyMappingRecord.policyCatalogId = policyCatalogId; + return this; + } + + public Builder policyId(long policyId) { + policyMappingRecord.policyId = policyId; + return this; + } + + public Builder parameters(String parameters) { + policyMappingRecord.parameters = parameters; + return this; + } + + public ModelPolicyMappingRecord build() { + return policyMappingRecord; + } + } + + public static ModelPolicyMappingRecord fromPolicyMappingRecord( + PolarisPolicyMappingRecord record) { + if (record == null) return null; + + return ModelPolicyMappingRecord.builder() + .targetCatalogId(record.getTargetCatalogId()) + .targetId(record.getTargetId()) + .policyTypeCode(record.getPolicyTypeCode()) + .policyCatalogId(record.getPolicyCatalogId()) + .policyId(record.getPolicyId()) + .parameters(record.getParameters()) + .build(); + } + + public static PolarisPolicyMappingRecord toPolicyMappingRecord(ModelPolicyMappingRecord model) { + if (model == null) return null; + + return new PolarisPolicyMappingRecord( + model.getTargetCatalogId(), + model.getTargetId(), + model.getPolicyCatalogId(), + model.getPolicyId(), + model.getPolicyTypeCode(), + model.getParameters()); + } + + @Override + public ModelPolicyMappingRecord fromResultSet(ResultSet rs) throws SQLException { + return ModelPolicyMappingRecord.builder() + .targetCatalogId(rs.getObject("target_catalog_id", Long.class)) + .targetId(rs.getObject("target_id", Long.class)) + .policyTypeCode(rs.getObject("policy_type_code", Integer.class)) + .policyCatalogId(rs.getObject("policy_catalog_id", Long.class)) + .policyId(rs.getObject("policy_id", Long.class)) + .parameters(rs.getString("parameters")) + .build(); + } + + @Override + public Map toMap() { + Map map = new HashMap<>(); + map.put("target_catalog_id", targetCatalogId); + map.put("target_id", targetId); + map.put("policy_type_code", policyTypeCode); + map.put("policy_catalog_id", policyCatalogId); + map.put("policy_id", policyId); + map.put("parameters", parameters); + return map; + } +} diff --git a/extension/persistence/relational-jdbc/src/main/resources/h2/schema-v1.sql b/extension/persistence/relational-jdbc/src/main/resources/h2/schema-v1.sql index 47ff53b14c..cacff3209e 100644 --- a/extension/persistence/relational-jdbc/src/main/resources/h2/schema-v1.sql +++ b/extension/persistence/relational-jdbc/src/main/resources/h2/schema-v1.sql @@ -91,3 +91,17 @@ CREATE TABLE IF NOT EXISTS principal_authentication_data ( ); COMMENT ON TABLE principal_authentication_data IS 'authentication data for client'; + +DROP TABLE IF EXISTS policy_mapping_record; +CREATE TABLE IF NOT EXISTS policy_mapping_record ( + realm_id TEXT NOT NULL, + target_catalog_id BIGINT NOT NULL, + target_id BIGINT NOT NULL, + policy_type_code INTEGER NOT NULL, + policy_catalog_id BIGINT NOT NULL, + policy_id BIGINT NOT NULL, + parameters TEXT NOT NULL DEFAULT '{}', + 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_catalog_id, policy_id, target_catalog_id, target_id); diff --git a/extension/persistence/relational-jdbc/src/test/java/org/apache/polaris/extension/persistence/impl/relational/jdbc/AtomicMetastoreManagerWithJdbcBasePersistenceImplTest.java b/extension/persistence/relational-jdbc/src/test/java/org/apache/polaris/extension/persistence/impl/relational/jdbc/AtomicMetastoreManagerWithJdbcBasePersistenceImplTest.java index 72ef15913b..92d31d8343 100644 --- a/extension/persistence/relational-jdbc/src/test/java/org/apache/polaris/extension/persistence/impl/relational/jdbc/AtomicMetastoreManagerWithJdbcBasePersistenceImplTest.java +++ b/extension/persistence/relational-jdbc/src/test/java/org/apache/polaris/extension/persistence/impl/relational/jdbc/AtomicMetastoreManagerWithJdbcBasePersistenceImplTest.java @@ -19,7 +19,6 @@ package org.apache.polaris.extension.persistence.impl.relational.jdbc; import static org.apache.polaris.core.persistence.PrincipalSecretsGenerator.RANDOM_SECRETS; -import static org.junit.jupiter.api.Assertions.assertThrows; import java.sql.SQLException; import java.time.ZoneId; @@ -35,7 +34,6 @@ import org.apache.polaris.extension.persistence.relational.jdbc.DatasourceOperations; import org.apache.polaris.extension.persistence.relational.jdbc.JdbcBasePersistenceImpl; import org.h2.jdbcx.JdbcConnectionPool; -import org.junit.jupiter.api.Test; import org.mockito.Mockito; public class AtomicMetastoreManagerWithJdbcBasePersistenceImplTest @@ -69,11 +67,4 @@ protected PolarisTestMetaStoreManager createPolarisTestMetaStoreManager() { new PolarisConfigurationStore() {}, timeSource.withZone(ZoneId.systemDefault()))); } - - @Override - @Test - protected void testPolicyMapping() { - // TODO: add Policy support. - assertThrows(UnsupportedOperationException.class, super::testPolicyMapping); - } } diff --git a/extension/persistence/relational-jdbc/src/test/java/org/apache/polaris/extension/persistence/impl/relational/jdbc/QueryGeneratorTest.java b/extension/persistence/relational-jdbc/src/test/java/org/apache/polaris/extension/persistence/impl/relational/jdbc/QueryGeneratorTest.java index e0e86e3018..653487b472 100644 --- a/extension/persistence/relational-jdbc/src/test/java/org/apache/polaris/extension/persistence/impl/relational/jdbc/QueryGeneratorTest.java +++ b/extension/persistence/relational-jdbc/src/test/java/org/apache/polaris/extension/persistence/impl/relational/jdbc/QueryGeneratorTest.java @@ -151,9 +151,7 @@ void testGenerateDeleteQuery_byObject() { ModelEntity entityToDelete = ModelEntity.builder().name("test").entityVersion(1).build(); String expectedQuery = "DELETE FROM POLARIS_SCHEMA.ENTITIES WHERE entity_version = 1 AND to_purge_timestamp = 0 AND realm_id = 'testRealm' AND internal_properties = '{}' AND catalog_id = 0 AND purge_timestamp = 0 AND sub_type_code = 0 AND create_timestamp = 0 AND last_update_timestamp = 0 AND parent_id = 0 AND name = 'test' AND id = 0 AND drop_timestamp = 0 AND properties = '{}' AND grant_records_version = 0 AND type_code = 0"; - assertEquals( - expectedQuery, - QueryGenerator.generateDeleteQuery(entityToDelete, ModelEntity.class, REALM_ID)); + assertEquals(expectedQuery, QueryGenerator.generateDeleteQuery(entityToDelete, REALM_ID)); } @Test @@ -161,9 +159,7 @@ void testGenerateDeleteQuery_byObject_nullValue() { ModelEntity entityToDelete = ModelEntity.builder().name("test").dropTimestamp(0L).build(); String expectedQuery = "DELETE FROM POLARIS_SCHEMA.ENTITIES WHERE entity_version = 0 AND to_purge_timestamp = 0 AND realm_id = 'testRealm' AND internal_properties = '{}' AND catalog_id = 0 AND purge_timestamp = 0 AND sub_type_code = 0 AND create_timestamp = 0 AND last_update_timestamp = 0 AND parent_id = 0 AND name = 'test' AND id = 0 AND drop_timestamp = 0 AND properties = '{}' AND grant_records_version = 0 AND type_code = 0"; - assertEquals( - expectedQuery, - QueryGenerator.generateDeleteQuery(entityToDelete, ModelEntity.class, REALM_ID)); + assertEquals(expectedQuery, QueryGenerator.generateDeleteQuery(entityToDelete, REALM_ID)); } @Test From 03472a801f15c444b47e80993f36d1a2a006ae46 Mon Sep 17 00:00:00 2001 From: Prashant Date: Fri, 25 Apr 2025 18:00:51 -0700 Subject: [PATCH 2/8] [JDBC] : Support Policy --- .../jdbc/JdbcBasePersistenceImpl.java | 54 ++++++++++++++----- .../relational/jdbc/QueryGenerator.java | 8 ++- .../src/main/resources/postgres/schema-v1.sql | 14 +++++ 3 files changed, 61 insertions(+), 15 deletions(-) diff --git a/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java b/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java index 22c4559599..e3a8ec7bf9 100644 --- a/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java +++ b/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java @@ -20,6 +20,7 @@ import static org.apache.polaris.extension.persistence.relational.jdbc.QueryGenerator.*; +import com.google.common.base.Preconditions; import jakarta.annotation.Nonnull; import jakarta.annotation.Nullable; import java.sql.SQLException; @@ -45,9 +46,11 @@ import org.apache.polaris.core.persistence.BasePersistence; import org.apache.polaris.core.persistence.EntityAlreadyExistsException; import org.apache.polaris.core.persistence.IntegrationPersistence; +import org.apache.polaris.core.persistence.PolicyMappingAlreadyExistsException; import org.apache.polaris.core.persistence.PrincipalSecretsGenerator; import org.apache.polaris.core.persistence.RetryOnConcurrencyException; import org.apache.polaris.core.policy.PolarisPolicyMappingRecord; +import org.apache.polaris.core.policy.PolicyType; import org.apache.polaris.core.storage.PolarisStorageConfigurationInfo; import org.apache.polaris.core.storage.PolarisStorageIntegration; import org.apache.polaris.core.storage.PolarisStorageIntegrationProvider; @@ -706,11 +709,37 @@ public void deletePrincipalSecrets( @Override public void writeToPolicyMappingRecords( @Nonnull PolarisCallContext callCtx, @Nonnull PolarisPolicyMappingRecord record) { - ModelPolicyMappingRecord modelPolicyMappingRecord = - ModelPolicyMappingRecord.fromPolicyMappingRecord(record); - String query = generateInsertQuery(modelPolicyMappingRecord, realmId); try { - datasourceOperations.executeUpdate(query); + datasourceOperations.runWithinTransaction( + statement -> { + PolicyType policyType = PolicyType.fromCode(record.getPolicyTypeCode()); + Preconditions.checkArgument( + policyType != null, "Invalid policy type code: %s", record.getPolicyTypeCode()); + + if (policyType.isInheritable()) { + List existingRecords = + loadPoliciesOnTargetByType( + callCtx, + record.getTargetCatalogId(), + record.getTargetId(), + record.getPolicyTypeCode()); + if (existingRecords.size() > 1) { + throw new PolicyMappingAlreadyExistsException(existingRecords.get(0)); + } else if (existingRecords.size() == 1) { + PolarisPolicyMappingRecord existingRecord = existingRecords.get(0); + if (existingRecord.getPolicyCatalogId() != record.getPolicyCatalogId() + || existingRecord.getPolicyId() != record.getPolicyId()) { + throw new PolicyMappingAlreadyExistsException(existingRecord); + } + } + + ModelPolicyMappingRecord modelPolicyMappingRecord = + ModelPolicyMappingRecord.fromPolicyMappingRecord(record); + String query = generateInsertQuery(modelPolicyMappingRecord, realmId); + statement.executeUpdate(query); + } + return true; + }); } catch (SQLException e) { throw new RuntimeException( String.format("Failed to write to policy mapping records due to %s", e.getMessage()), e); @@ -773,8 +802,7 @@ public PolarisPolicyMappingRecord lookupPolicyMappingRecord( List results = fetchPolicyMappingRecords(query); if (results.size() > 1) { throw new IllegalStateException( - String.format( - "More than 1 policy %s for a given policy mapping", results.getFirst())); + String.format("More than 1 policy %s for a given policy mapping", results.getFirst())); } return results.size() == 1 ? results.getFirst() : null; } @@ -823,16 +851,16 @@ public List loadAllTargetsOnPolicy( private List fetchPolicyMappingRecords(String query) { try { List results = - datasourceOperations.executeSelect( - query, - ModelPolicyMappingRecord.class, - ModelPolicyMappingRecord::toPolicyMappingRecord, - null, - Integer.MAX_VALUE); + datasourceOperations.executeSelect( + query, + ModelPolicyMappingRecord.class, + ModelPolicyMappingRecord::toPolicyMappingRecord, + null, + Integer.MAX_VALUE); return results == null ? Collections.emptyList() : results; } catch (SQLException e) { throw new RuntimeException( - String.format("Failed to retrieve policy mapping records %s", e.getMessage()), e); + String.format("Failed to retrieve policy mapping records %s", e.getMessage()), e); } } diff --git a/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/QueryGenerator.java b/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/QueryGenerator.java index aff9c061ed..7d0b5ec928 100644 --- a/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/QueryGenerator.java +++ b/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/QueryGenerator.java @@ -26,7 +26,11 @@ import java.util.Map; import org.apache.polaris.core.entity.PolarisEntityCore; import org.apache.polaris.core.entity.PolarisEntityId; -import org.apache.polaris.extension.persistence.relational.jdbc.models.*; +import org.apache.polaris.extension.persistence.relational.jdbc.models.Converter; +import org.apache.polaris.extension.persistence.relational.jdbc.models.ModelEntity; +import org.apache.polaris.extension.persistence.relational.jdbc.models.ModelGrantRecord; +import org.apache.polaris.extension.persistence.relational.jdbc.models.ModelPolicyMappingRecord; +import org.apache.polaris.extension.persistence.relational.jdbc.models.ModelPrincipalAuthenticationData; public class QueryGenerator { @@ -73,7 +77,7 @@ public static String generateDeleteQueryForEntityPolicyMappingRecords( + ") AND realm_id = '" + realmId + "'"; - return generateDeleteQuery(ModelGrantRecord.class, whereClause); + return generateDeleteQuery(ModelPolicyMappingRecord.class, whereClause); } public static String generateSelectQueryWithEntityIds( diff --git a/extension/persistence/relational-jdbc/src/main/resources/postgres/schema-v1.sql b/extension/persistence/relational-jdbc/src/main/resources/postgres/schema-v1.sql index 12da8f7a84..ea375b357a 100644 --- a/extension/persistence/relational-jdbc/src/main/resources/postgres/schema-v1.sql +++ b/extension/persistence/relational-jdbc/src/main/resources/postgres/schema-v1.sql @@ -91,3 +91,17 @@ CREATE TABLE IF NOT EXISTS principal_authentication_data ( ); COMMENT ON TABLE principal_authentication_data IS 'authentication data for client'; + +DROP TABLE IF EXISTS policy_mapping_record; +CREATE TABLE IF NOT EXISTS policy_mapping_record ( + realm_id TEXT NOT NULL, + target_catalog_id BIGINT NOT NULL, + target_id BIGINT NOT NULL, + policy_type_code INTEGER NOT NULL, + policy_catalog_id BIGINT NOT NULL, + policy_id BIGINT NOT NULL, + parameters JSONB NOT NULL DEFAULT '{}'::JSONB, + 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_catalog_id, policy_id, target_catalog_id, target_id); From 63460eeacf24bbad55ea0cd95469ca03bc5824d3 Mon Sep 17 00:00:00 2001 From: Prashant Date: Sat, 26 Apr 2025 00:31:03 -0700 Subject: [PATCH 3/8] Add update query --- .../jdbc/JdbcBasePersistenceImpl.java | 28 +++++++++++++++---- .../src/main/resources/h2/schema-v1.sql | 2 +- .../src/main/resources/postgres/schema-v1.sql | 2 +- 3 files changed, 24 insertions(+), 8 deletions(-) diff --git a/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java b/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java index e3a8ec7bf9..3653565049 100644 --- a/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java +++ b/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java @@ -716,6 +716,9 @@ public void writeToPolicyMappingRecords( Preconditions.checkArgument( policyType != null, "Invalid policy type code: %s", record.getPolicyTypeCode()); + String query = + generateInsertQuery( + ModelPolicyMappingRecord.fromPolicyMappingRecord(record), realmId); if (policyType.isInheritable()) { List existingRecords = loadPoliciesOnTargetByType( @@ -724,18 +727,31 @@ public void writeToPolicyMappingRecords( record.getTargetId(), record.getPolicyTypeCode()); if (existingRecords.size() > 1) { - throw new PolicyMappingAlreadyExistsException(existingRecords.get(0)); + throw new PolicyMappingAlreadyExistsException(existingRecords.getFirst()); } else if (existingRecords.size() == 1) { - PolarisPolicyMappingRecord existingRecord = existingRecords.get(0); + PolarisPolicyMappingRecord existingRecord = existingRecords.getFirst(); if (existingRecord.getPolicyCatalogId() != record.getPolicyCatalogId() || existingRecord.getPolicyId() != record.getPolicyId()) { throw new PolicyMappingAlreadyExistsException(existingRecord); } + Map 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); + query = + generateUpdateQuery( + ModelPolicyMappingRecord.fromPolicyMappingRecord(record), updateClause); } - - ModelPolicyMappingRecord modelPolicyMappingRecord = - ModelPolicyMappingRecord.fromPolicyMappingRecord(record); - String query = generateInsertQuery(modelPolicyMappingRecord, realmId); statement.executeUpdate(query); } return true; diff --git a/extension/persistence/relational-jdbc/src/main/resources/h2/schema-v1.sql b/extension/persistence/relational-jdbc/src/main/resources/h2/schema-v1.sql index cacff3209e..f97140f33c 100644 --- a/extension/persistence/relational-jdbc/src/main/resources/h2/schema-v1.sql +++ b/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_catalog_id, policy_id, target_catalog_id, target_id); +CREATE INDEX IF NOT EXISTS idx_policy_mapping_record ON policy_mapping_record (realm_id, target_catalog_id, target_id, policy_type_code, policy_catalog_id, policy_id); diff --git a/extension/persistence/relational-jdbc/src/main/resources/postgres/schema-v1.sql b/extension/persistence/relational-jdbc/src/main/resources/postgres/schema-v1.sql index ea375b357a..ce2cdd7763 100644 --- a/extension/persistence/relational-jdbc/src/main/resources/postgres/schema-v1.sql +++ b/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_catalog_id, policy_id, target_catalog_id, target_id); +CREATE INDEX IF NOT EXISTS idx_policy_mapping_record ON policy_mapping_record (realm_id, target_catalog_id, target_id, policy_type_code, policy_catalog_id, policy_id); From bd3531e32dda34b0c0fb8ec71bdc5a90ef64aa15 Mon Sep 17 00:00:00 2001 From: Prashant Date: Sun, 27 Apr 2025 19:26:40 -0700 Subject: [PATCH 4/8] flip the index of policy --- .../persistence/relational/jdbc/JdbcBasePersistenceImpl.java | 5 ++++- .../relational-jdbc/src/main/resources/h2/schema-v1.sql | 2 +- .../src/main/resources/postgres/schema-v1.sql | 2 +- 3 files changed, 6 insertions(+), 3 deletions(-) diff --git a/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java b/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java index 3653565049..caa8384805 100644 --- a/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java +++ b/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java @@ -253,7 +253,10 @@ public void deleteAll(@Nonnull PolarisCallContext callCtx) { try { datasourceOperations.executeUpdate(generateDeleteAll(ModelEntity.class, realmId)); datasourceOperations.executeUpdate(generateDeleteAll(ModelGrantRecord.class, realmId)); - datasourceOperations.executeUpdate(generateDeleteAll(ModelEntity.class, realmId)); + datasourceOperations.executeUpdate( + generateDeleteAll(ModelPrincipalAuthenticationData.class, realmId)); + datasourceOperations.executeUpdate( + generateDeleteAll(ModelPolicyMappingRecord.class, realmId)); } catch (SQLException e) { throw new RuntimeException( String.format("Failed to delete all due to %s", e.getMessage()), e); diff --git a/extension/persistence/relational-jdbc/src/main/resources/h2/schema-v1.sql b/extension/persistence/relational-jdbc/src/main/resources/h2/schema-v1.sql index f97140f33c..1625e1c16f 100644 --- a/extension/persistence/relational-jdbc/src/main/resources/h2/schema-v1.sql +++ b/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, 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); diff --git a/extension/persistence/relational-jdbc/src/main/resources/postgres/schema-v1.sql b/extension/persistence/relational-jdbc/src/main/resources/postgres/schema-v1.sql index ce2cdd7763..90a9356e17 100644 --- a/extension/persistence/relational-jdbc/src/main/resources/postgres/schema-v1.sql +++ b/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, 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); From f93e5a8cf1fd32e4848bd2a9b7c665d84b50bc34 Mon Sep 17 00:00:00 2001 From: Prashant Date: Sun, 27 Apr 2025 19:55:10 -0700 Subject: [PATCH 5/8] Add quarkus service integ test for policy --- .../jdbc/JdbcQuarkusPolicyServiceIT.java | 28 +++++++++++++++++++ 1 file changed, 28 insertions(+) create mode 100644 quarkus/service/src/intTest/java/org/apache/polaris/service/quarkus/it/relational/jdbc/JdbcQuarkusPolicyServiceIT.java diff --git a/quarkus/service/src/intTest/java/org/apache/polaris/service/quarkus/it/relational/jdbc/JdbcQuarkusPolicyServiceIT.java b/quarkus/service/src/intTest/java/org/apache/polaris/service/quarkus/it/relational/jdbc/JdbcQuarkusPolicyServiceIT.java new file mode 100644 index 0000000000..5907fe5515 --- /dev/null +++ b/quarkus/service/src/intTest/java/org/apache/polaris/service/quarkus/it/relational/jdbc/JdbcQuarkusPolicyServiceIT.java @@ -0,0 +1,28 @@ +/* + * 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.service.quarkus.it.relational.jdbc; + +import io.quarkus.test.junit.QuarkusIntegrationTest; +import io.quarkus.test.junit.TestProfile; +import org.apache.polaris.service.it.test.PolarisPolicyServiceIntegrationTest; +import org.apache.polaris.test.commons.RelationalJdbcProfile; + +@TestProfile(RelationalJdbcProfile.class) +@QuarkusIntegrationTest +public class JdbcQuarkusPolicyServiceIT extends PolarisPolicyServiceIntegrationTest {} From 0cf173fcd5b834e5b368e27e6db82fda52730a99 Mon Sep 17 00:00:00 2001 From: Prashant Date: Mon, 28 Apr 2025 00:04:37 -0700 Subject: [PATCH 6/8] wrap delete all in transaction --- .../relational/jdbc/JdbcBasePersistenceImpl.java | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java b/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java index caa8384805..930b3a38a6 100644 --- a/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java +++ b/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java @@ -251,12 +251,15 @@ public void deleteAllEntityGrantRecords( @Override public void deleteAll(@Nonnull PolarisCallContext callCtx) { try { - datasourceOperations.executeUpdate(generateDeleteAll(ModelEntity.class, realmId)); - datasourceOperations.executeUpdate(generateDeleteAll(ModelGrantRecord.class, realmId)); - datasourceOperations.executeUpdate( - generateDeleteAll(ModelPrincipalAuthenticationData.class, realmId)); - datasourceOperations.executeUpdate( - generateDeleteAll(ModelPolicyMappingRecord.class, realmId)); + datasourceOperations.runWithinTransaction( + statement -> { + statement.executeUpdate(generateDeleteAll(ModelEntity.class, realmId)); + statement.executeUpdate(generateDeleteAll(ModelGrantRecord.class, realmId)); + statement.executeUpdate( + generateDeleteAll(ModelPrincipalAuthenticationData.class, realmId)); + statement.executeUpdate(generateDeleteAll(ModelPolicyMappingRecord.class, realmId)); + return true; + }); } catch (SQLException e) { throw new RuntimeException( String.format("Failed to delete all due to %s", e.getMessage()), e); From 60c7de8eb1c39bf6705d535fd69430b35dfee5db Mon Sep 17 00:00:00 2001 From: Prashant Date: Mon, 28 Apr 2025 11:08:06 -0700 Subject: [PATCH 7/8] Address yufei feedback --- .../jdbc/JdbcBasePersistenceImpl.java | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java b/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java index 930b3a38a6..9a552bc989 100644 --- a/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java +++ b/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java @@ -721,8 +721,7 @@ public void writeToPolicyMappingRecords( PolicyType policyType = PolicyType.fromCode(record.getPolicyTypeCode()); Preconditions.checkArgument( policyType != null, "Invalid policy type code: %s", record.getPolicyTypeCode()); - - String query = + String insertQuery = generateInsertQuery( ModelPolicyMappingRecord.fromPolicyMappingRecord(record), realmId); if (policyType.isInheritable()) { @@ -754,11 +753,15 @@ public void writeToPolicyMappingRecords( record.getPolicyCatalogId(), "realm_id", realmId); - query = + String updateQuery = generateUpdateQuery( ModelPolicyMappingRecord.fromPolicyMappingRecord(record), updateClause); + statement.executeUpdate(updateQuery); + } else { + statement.executeUpdate(insertQuery); } - statement.executeUpdate(query); + } else { + statement.executeUpdate(insertQuery); } return true; }); @@ -771,8 +774,7 @@ public void writeToPolicyMappingRecords( @Override public void deleteFromPolicyMappingRecords( @Nonnull PolarisCallContext callCtx, @Nonnull PolarisPolicyMappingRecord record) { - ModelPolicyMappingRecord modelPolicyMappingRecord = - ModelPolicyMappingRecord.fromPolicyMappingRecord(record); + var modelPolicyMappingRecord = ModelPolicyMappingRecord.fromPolicyMappingRecord(record); String query = generateDeleteQuery(modelPolicyMappingRecord, realmId); try { datasourceOperations.executeUpdate(query); @@ -822,10 +824,7 @@ public PolarisPolicyMappingRecord lookupPolicyMappingRecord( realmId); String query = generateSelectQuery(ModelPolicyMappingRecord.class, params); List results = fetchPolicyMappingRecords(query); - if (results.size() > 1) { - throw new IllegalStateException( - String.format("More than 1 policy %s for a given policy mapping", results.getFirst())); - } + Preconditions.checkState(results.size() <= 1, "More than one policy mapping records found"); return results.size() == 1 ? results.getFirst() : null; } From 9a2860beebb64ab973784b453062588a5ace930b Mon Sep 17 00:00:00 2001 From: Prashant Date: Mon, 28 Apr 2025 11:23:41 -0700 Subject: [PATCH 8/8] refactor to different function --- .../jdbc/JdbcBasePersistenceImpl.java | 85 +++++++++++-------- .../src/main/resources/h2/schema-v1.sql | 2 +- .../src/main/resources/postgres/schema-v1.sql | 2 +- 3 files changed, 50 insertions(+), 39 deletions(-) diff --git a/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java b/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java index 9a552bc989..5ffce813f9 100644 --- a/extension/persistence/relational-jdbc/src/main/java/org/apache/polaris/extension/persistence/relational/jdbc/JdbcBasePersistenceImpl.java +++ b/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 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 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 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 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) { diff --git a/extension/persistence/relational-jdbc/src/main/resources/h2/schema-v1.sql b/extension/persistence/relational-jdbc/src/main/resources/h2/schema-v1.sql index 1625e1c16f..cda0352efb 100644 --- a/extension/persistence/relational-jdbc/src/main/resources/h2/schema-v1.sql +++ b/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); diff --git a/extension/persistence/relational-jdbc/src/main/resources/postgres/schema-v1.sql b/extension/persistence/relational-jdbc/src/main/resources/postgres/schema-v1.sql index 90a9356e17..25acd2ed0f 100644 --- a/extension/persistence/relational-jdbc/src/main/resources/postgres/schema-v1.sql +++ b/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);