Skip to content

Commit cfe22b7

Browse files
singhpk234singhpk234
authored
[JDBC] Support Policy (#1468)
1 parent a84b5c4 commit cfe22b7

File tree

8 files changed

+447
-21
lines changed

8 files changed

+447
-21
lines changed

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

Lines changed: 199 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020

2121
import static org.apache.polaris.extension.persistence.relational.jdbc.QueryGenerator.*;
2222

23+
import com.google.common.base.Preconditions;
2324
import jakarta.annotation.Nonnull;
2425
import jakarta.annotation.Nullable;
2526
import java.sql.SQLException;
@@ -45,13 +46,17 @@
4546
import org.apache.polaris.core.persistence.BasePersistence;
4647
import org.apache.polaris.core.persistence.EntityAlreadyExistsException;
4748
import org.apache.polaris.core.persistence.IntegrationPersistence;
49+
import org.apache.polaris.core.persistence.PolicyMappingAlreadyExistsException;
4850
import org.apache.polaris.core.persistence.PrincipalSecretsGenerator;
4951
import org.apache.polaris.core.persistence.RetryOnConcurrencyException;
52+
import org.apache.polaris.core.policy.PolarisPolicyMappingRecord;
53+
import org.apache.polaris.core.policy.PolicyType;
5054
import org.apache.polaris.core.storage.PolarisStorageConfigurationInfo;
5155
import org.apache.polaris.core.storage.PolarisStorageIntegration;
5256
import org.apache.polaris.core.storage.PolarisStorageIntegrationProvider;
5357
import org.apache.polaris.extension.persistence.relational.jdbc.models.ModelEntity;
5458
import org.apache.polaris.extension.persistence.relational.jdbc.models.ModelGrantRecord;
59+
import org.apache.polaris.extension.persistence.relational.jdbc.models.ModelPolicyMappingRecord;
5560
import org.apache.polaris.extension.persistence.relational.jdbc.models.ModelPrincipalAuthenticationData;
5661
import org.slf4j.Logger;
5762
import org.slf4j.LoggerFactory;
@@ -220,7 +225,7 @@ public void deleteEntity(@Nonnull PolarisCallContext callCtx, @Nonnull PolarisBa
220225
public void deleteFromGrantRecords(
221226
@Nonnull PolarisCallContext callCtx, @Nonnull PolarisGrantRecord grantRec) {
222227
ModelGrantRecord modelGrantRecord = ModelGrantRecord.fromGrantRecord(grantRec);
223-
String query = generateDeleteQuery(modelGrantRecord, ModelGrantRecord.class, realmId);
228+
String query = generateDeleteQuery(modelGrantRecord, realmId);
224229
try {
225230
datasourceOperations.executeUpdate(query);
226231
} catch (SQLException e) {
@@ -246,9 +251,15 @@ public void deleteAllEntityGrantRecords(
246251
@Override
247252
public void deleteAll(@Nonnull PolarisCallContext callCtx) {
248253
try {
249-
datasourceOperations.executeUpdate(generateDeleteAll(ModelEntity.class, realmId));
250-
datasourceOperations.executeUpdate(generateDeleteAll(ModelGrantRecord.class, realmId));
251-
datasourceOperations.executeUpdate(generateDeleteAll(ModelEntity.class, realmId));
254+
datasourceOperations.runWithinTransaction(
255+
statement -> {
256+
statement.executeUpdate(generateDeleteAll(ModelEntity.class, realmId));
257+
statement.executeUpdate(generateDeleteAll(ModelGrantRecord.class, realmId));
258+
statement.executeUpdate(
259+
generateDeleteAll(ModelPrincipalAuthenticationData.class, realmId));
260+
statement.executeUpdate(generateDeleteAll(ModelPolicyMappingRecord.class, realmId));
261+
return true;
262+
});
252263
} catch (SQLException e) {
253264
throw new RuntimeException(
254265
String.format("Failed to delete all due to %s", e.getMessage()), e);
@@ -701,6 +712,190 @@ public void deletePrincipalSecrets(
701712
}
702713
}
703714

715+
@Override
716+
public void writeToPolicyMappingRecords(
717+
@Nonnull PolarisCallContext callCtx, @Nonnull PolarisPolicyMappingRecord record) {
718+
try {
719+
datasourceOperations.runWithinTransaction(
720+
statement -> {
721+
PolicyType policyType = PolicyType.fromCode(record.getPolicyTypeCode());
722+
Preconditions.checkArgument(
723+
policyType != null, "Invalid policy type code: %s", record.getPolicyTypeCode());
724+
String insertPolicyMappingQuery =
725+
generateInsertQuery(
726+
ModelPolicyMappingRecord.fromPolicyMappingRecord(record), realmId);
727+
if (policyType.isInheritable()) {
728+
return handleInheritablePolicy(callCtx, record, insertPolicyMappingQuery, statement);
729+
} else {
730+
statement.executeUpdate(insertPolicyMappingQuery);
731+
}
732+
return true;
733+
});
734+
} catch (SQLException e) {
735+
throw new RuntimeException(
736+
String.format("Failed to write to policy mapping records due to %s", e.getMessage()), e);
737+
}
738+
}
739+
740+
private boolean handleInheritablePolicy(
741+
@Nonnull PolarisCallContext callCtx,
742+
@Nonnull PolarisPolicyMappingRecord record,
743+
@Nonnull String insertQuery,
744+
Statement statement)
745+
throws SQLException {
746+
List<PolarisPolicyMappingRecord> existingRecords =
747+
loadPoliciesOnTargetByType(
748+
callCtx, record.getTargetCatalogId(), record.getTargetId(), record.getPolicyTypeCode());
749+
if (existingRecords.size() > 1) {
750+
throw new PolicyMappingAlreadyExistsException(existingRecords.getFirst());
751+
} else if (existingRecords.size() == 1) {
752+
PolarisPolicyMappingRecord existingRecord = existingRecords.getFirst();
753+
if (existingRecord.getPolicyCatalogId() != record.getPolicyCatalogId()
754+
|| existingRecord.getPolicyId() != record.getPolicyId()) {
755+
// Only one policy of the same type can be attached to an entity when the policy is
756+
// inheritable.
757+
throw new PolicyMappingAlreadyExistsException(existingRecord);
758+
}
759+
Map<String, Object> updateClause =
760+
Map.of(
761+
"target_catalog_id",
762+
record.getTargetCatalogId(),
763+
"target_id",
764+
record.getTargetId(),
765+
"policy_type_code",
766+
record.getPolicyTypeCode(),
767+
"policy_id",
768+
record.getPolicyId(),
769+
"policy_catalog_id",
770+
record.getPolicyCatalogId(),
771+
"realm_id",
772+
realmId);
773+
// In case of the mapping exist, update the policy mapping with the new parameters.
774+
String updateQuery =
775+
generateUpdateQuery(
776+
ModelPolicyMappingRecord.fromPolicyMappingRecord(record), updateClause);
777+
statement.executeUpdate(updateQuery);
778+
} else {
779+
// record doesn't exist do an insert.
780+
statement.executeUpdate(insertQuery);
781+
}
782+
return true;
783+
}
784+
785+
@Override
786+
public void deleteFromPolicyMappingRecords(
787+
@Nonnull PolarisCallContext callCtx, @Nonnull PolarisPolicyMappingRecord record) {
788+
var modelPolicyMappingRecord = ModelPolicyMappingRecord.fromPolicyMappingRecord(record);
789+
String query = generateDeleteQuery(modelPolicyMappingRecord, realmId);
790+
try {
791+
datasourceOperations.executeUpdate(query);
792+
} catch (SQLException e) {
793+
throw new RuntimeException(
794+
String.format("Failed to write to policy records due to %s", e.getMessage()), e);
795+
}
796+
}
797+
798+
@Override
799+
public void deleteAllEntityPolicyMappingRecords(
800+
@Nonnull PolarisCallContext callCtx,
801+
@Nonnull PolarisEntityCore entity,
802+
@Nonnull List<PolarisPolicyMappingRecord> mappingOnTarget,
803+
@Nonnull List<PolarisPolicyMappingRecord> mappingOnPolicy) {
804+
try {
805+
datasourceOperations.executeUpdate(
806+
generateDeleteQueryForEntityPolicyMappingRecords(entity, realmId));
807+
} catch (SQLException e) {
808+
throw new RuntimeException(
809+
String.format("Failed to delete policy mapping records due to %s", e.getMessage()), e);
810+
}
811+
}
812+
813+
@Nullable
814+
@Override
815+
public PolarisPolicyMappingRecord lookupPolicyMappingRecord(
816+
@Nonnull PolarisCallContext callCtx,
817+
long targetCatalogId,
818+
long targetId,
819+
int policyTypeCode,
820+
long policyCatalogId,
821+
long policyId) {
822+
Map<String, Object> params =
823+
Map.of(
824+
"target_catalog_id",
825+
targetCatalogId,
826+
"target_id",
827+
targetId,
828+
"policy_type_code",
829+
policyTypeCode,
830+
"policy_id",
831+
policyId,
832+
"policy_catalog_id",
833+
policyCatalogId,
834+
"realm_id",
835+
realmId);
836+
String query = generateSelectQuery(ModelPolicyMappingRecord.class, params);
837+
List<PolarisPolicyMappingRecord> results = fetchPolicyMappingRecords(query);
838+
Preconditions.checkState(results.size() <= 1, "More than one policy mapping records found");
839+
return results.size() == 1 ? results.getFirst() : null;
840+
}
841+
842+
@Nonnull
843+
@Override
844+
public List<PolarisPolicyMappingRecord> loadPoliciesOnTargetByType(
845+
@Nonnull PolarisCallContext callCtx,
846+
long targetCatalogId,
847+
long targetId,
848+
int policyTypeCode) {
849+
Map<String, Object> params =
850+
Map.of(
851+
"target_catalog_id",
852+
targetCatalogId,
853+
"target_id",
854+
targetId,
855+
"policy_type_code",
856+
policyTypeCode,
857+
"realm_id",
858+
realmId);
859+
String query = generateSelectQuery(ModelPolicyMappingRecord.class, params);
860+
return fetchPolicyMappingRecords(query);
861+
}
862+
863+
@Nonnull
864+
@Override
865+
public List<PolarisPolicyMappingRecord> loadAllPoliciesOnTarget(
866+
@Nonnull PolarisCallContext callCtx, long targetCatalogId, long targetId) {
867+
Map<String, Object> params =
868+
Map.of("target_catalog_id", targetCatalogId, "target_id", targetId, "realm_id", realmId);
869+
String query = generateSelectQuery(ModelPolicyMappingRecord.class, params);
870+
return fetchPolicyMappingRecords(query);
871+
}
872+
873+
@Nonnull
874+
@Override
875+
public List<PolarisPolicyMappingRecord> loadAllTargetsOnPolicy(
876+
@Nonnull PolarisCallContext callCtx, long policyCatalogId, long policyId) {
877+
Map<String, Object> params =
878+
Map.of("policy_catalog_id", policyCatalogId, "policy_id", policyId, "realm_id", realmId);
879+
String query = generateSelectQuery(ModelPolicyMappingRecord.class, params);
880+
return fetchPolicyMappingRecords(query);
881+
}
882+
883+
private List<PolarisPolicyMappingRecord> fetchPolicyMappingRecords(String query) {
884+
try {
885+
List<PolarisPolicyMappingRecord> results =
886+
datasourceOperations.executeSelect(
887+
query,
888+
ModelPolicyMappingRecord.class,
889+
ModelPolicyMappingRecord::toPolicyMappingRecord,
890+
null,
891+
Integer.MAX_VALUE);
892+
return results == null ? Collections.emptyList() : results;
893+
} catch (SQLException e) {
894+
throw new RuntimeException(
895+
String.format("Failed to retrieve policy mapping records %s", e.getMessage()), e);
896+
}
897+
}
898+
704899
@Nullable
705900
@Override
706901
public <T extends PolarisStorageConfigurationInfo>

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

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
import org.apache.polaris.extension.persistence.relational.jdbc.models.Converter;
3030
import org.apache.polaris.extension.persistence.relational.jdbc.models.ModelEntity;
3131
import org.apache.polaris.extension.persistence.relational.jdbc.models.ModelGrantRecord;
32+
import org.apache.polaris.extension.persistence.relational.jdbc.models.ModelPolicyMappingRecord;
3233
import org.apache.polaris.extension.persistence.relational.jdbc.models.ModelPrincipalAuthenticationData;
3334

3435
public class QueryGenerator {
@@ -59,6 +60,26 @@ public static String generateDeleteQueryForEntityGrantRecords(
5960
return generateDeleteQuery(ModelGrantRecord.class, whereClause);
6061
}
6162

63+
public static String generateDeleteQueryForEntityPolicyMappingRecords(
64+
@Nonnull PolarisEntityCore entity, @Nonnull String realmId) {
65+
String targetCondition =
66+
String.format(
67+
"target_id = %s AND target_catalog_id = %s", entity.getId(), entity.getCatalogId());
68+
String sourceCondition =
69+
String.format(
70+
"policy_id = %s AND policy_catalog_id = %s", entity.getId(), entity.getCatalogId());
71+
72+
String whereClause =
73+
" WHERE ("
74+
+ targetCondition
75+
+ " OR "
76+
+ sourceCondition
77+
+ ") AND realm_id = '"
78+
+ realmId
79+
+ "'";
80+
return generateDeleteQuery(ModelPolicyMappingRecord.class, whereClause);
81+
}
82+
6283
public static String generateSelectQueryWithEntityIds(
6384
@Nonnull String realmId, @Nonnull List<PolarisEntityId> entityIds) {
6485
if (entityIds.isEmpty()) {
@@ -127,8 +148,8 @@ public static String generateDeleteAll(@Nonnull Class<?> entityClass, @Nonnull S
127148
}
128149

129150
public static <T> String generateDeleteQuery(
130-
@Nonnull Converter<T> entity, @Nonnull Class<?> entityClass, @Nonnull String realmId) {
131-
String tableName = getTableName(entityClass);
151+
@Nonnull Converter<T> entity, @Nonnull String realmId) {
152+
String tableName = getTableName(entity.getClass());
132153
Map<String, Object> objMap = entity.toMap();
133154
objMap.put("realm_id", realmId);
134155
String whereConditions = generateWhereClause(objMap);
@@ -186,6 +207,8 @@ public static String getTableName(@Nonnull Class<?> entityClass) {
186207
tableName = "GRANT_RECORDS";
187208
} else if (entityClass.equals(ModelPrincipalAuthenticationData.class)) {
188209
tableName = "PRINCIPAL_AUTHENTICATION_DATA";
210+
} else if (entityClass.equals(ModelPolicyMappingRecord.class)) {
211+
tableName = "POLICY_MAPPING_RECORD";
189212
} else {
190213
throw new IllegalArgumentException("Unsupported entity class: " + entityClass.getName());
191214
}

0 commit comments

Comments
 (0)