From 0a79a4c936b45c553edee31f76df364eca45e82b Mon Sep 17 00:00:00 2001 From: prithvi Date: Tue, 14 Jul 2026 03:07:57 +0530 Subject: [PATCH 1/2] Make JDBC grant/catalog operations atomic without BasePersistence SPI change --- CHANGELOG.md | 5 + .../jdbc/JdbcBasePersistenceImpl.java | 261 ++++++++++++------ .../jdbc/JdbcGrantRecordsIdempotencyTest.java | 14 +- .../AtomicOperationMetaStoreManager.java | 6 + .../core/persistence/BasePersistence.java | 12 + 5 files changed, 207 insertions(+), 91 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 87d53bc30fd..ed4fd5eaf60 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -76,6 +76,11 @@ request adding CHANGELOG notes for breaking (!) changes and possibly other secti ### Deprecations ### Fixes +- Fixed JDBC persistence to make grant/revoke, createCatalog, and dropEntity operations atomic + by batching consecutive writes in a single transaction. A new `BasePersistence.flush()` method + commits any pending batched writes; `AtomicOperationMetaStoreManager` calls it after sequences + of writes that should be atomic. This removes the partial-commit window for JDBC without adding + RBAC-specific methods to the persistence SPI. - Fixed a boundary condition in GCS downscoped credential generation (`GcpCredentialsStorageIntegration`). Locations without a trailing slash could previously grant access to sibling object prefixes via the generated CEL conditions for `resource.name` and list prefixes. Granted paths are now normalized to a directory prefix (with a trailing slash) before the CEL conditions are built, so sibling prefixes can no longer satisfy the `startsWith` checks. - Async task execution (table cleanup, manifest and batch file cleanup) now retries when a handler returns false on transient errors (e.g. IO or delete failures). Previously `false` was swallowed with only a warning log and the task was never retried via the existing retry mechanism. - Fixed `NullPointerException` during `dropEntity` when an entity referenced by a grant had been concurrently removed (or purged). `lookupEntities` can return null entries for dropped entities; these are now skipped safely. diff --git a/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcBasePersistenceImpl.java b/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcBasePersistenceImpl.java index cc711cc646a..7294ec1e0d8 100644 --- a/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcBasePersistenceImpl.java +++ b/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcBasePersistenceImpl.java @@ -94,6 +94,10 @@ public class JdbcBasePersistenceImpl private final String realmId; private final int schemaVersion; + // Thread-local connection for batching write operations atomically across multiple SPI calls. + // When non-null, write operations use this connection instead of acquiring a new one. + private static final ThreadLocal ACTIVE_CONNECTION = new ThreadLocal<>(); + // The max number of components a location can have before the optimized sibling check is not used private static final int MAX_LOCATION_COMPONENTS = 40; @@ -121,18 +125,11 @@ public void writeEntity( @NonNull PolarisBaseEntity entity, boolean nameOrParentChanged, PolarisBaseEntity originalEntity) { - try { - persistEntity( - callCtx, - entity, - originalEntity, - null, - (connection, preparedQuery) -> { - return datasourceOperations.executeUpdate(preparedQuery); - }); - } catch (SQLException e) { - throw new RuntimeException("Error persisting entity", e); - } + runWithConnection( + connection -> { + persistEntity(callCtx, entity, originalEntity, connection, datasourceOperations::execute); + }, + "Error persisting entity"); } @Override @@ -140,37 +137,139 @@ public void writeEntities( @NonNull PolarisCallContext callCtx, @NonNull List entities, List originalEntities) { - try { - datasourceOperations.runWithinTransaction( - connection -> { - for (int i = 0; i < entities.size(); i++) { - PolarisBaseEntity entity = entities.get(i); - PolarisBaseEntity originalEntity = - originalEntities != null ? originalEntities.get(i) : null; - // first, check if the entity has already been created, in which case we will simply - // return it. - PolarisBaseEntity entityFound = - lookupEntity( - callCtx, entity.getCatalogId(), entity.getId(), entity.getTypeCode()); - if (entityFound != null && originalEntity == null) { - // probably the client retried, simply return it - // TODO: Check correctness of returning entityFound vs entity here. It may have - // already been updated after the creation. - continue; - } - persistEntity( - callCtx, entity, originalEntity, connection, datasourceOperations::execute); + if (originalEntities != null && originalEntities.size() != entities.size()) { + throw new IllegalArgumentException( + String.format( + "originalEntities size (%s) must match entities size (%s)", + originalEntities.size(), entities.size())); + } + runWithConnection( + connection -> { + for (int i = 0; i < entities.size(); i++) { + PolarisBaseEntity entity = entities.get(i); + PolarisBaseEntity originalEntity = + originalEntities != null ? originalEntities.get(i) : null; + // Idempotent create-retry: if the entity already exists under the same id, skip. + PolarisBaseEntity entityFound = + lookupEntity(callCtx, entity.getCatalogId(), entity.getId(), entity.getTypeCode()); + if (entityFound != null && originalEntity == null) { + continue; } - return true; - }); + persistEntity( + callCtx, entity, originalEntity, connection, datasourceOperations::execute); + } + }, + "Error executing the transaction for writing entities"); + } + + @Override + public void flush() { + Connection connection = ACTIVE_CONNECTION.get(); + if (connection != null) { + try { + connection.commit(); + } catch (SQLException e) { + throw new RuntimeException("Error committing transaction", e); + } finally { + try { + connection.close(); + } catch (SQLException e) { + LOGGER.error("Error closing connection", e); + } + ACTIVE_CONNECTION.remove(); + } + } + } + + /** + * Runs the given action using the active batched connection if one exists, otherwise executes + * inside a new transaction. If the action throws, any active batched connection is rolled back + * and cleared. + */ + private void runWithConnection(SqlConsumer action, String errorMessage) { + Connection connection = ACTIVE_CONNECTION.get(); + if (connection != null) { + try { + action.accept(connection); + } catch (SQLException e) { + try { + connection.rollback(); + } catch (SQLException rollbackEx) { + LOGGER.error("Error rolling back transaction", rollbackEx); + } + ACTIVE_CONNECTION.remove(); + try { + connection.close(); + } catch (SQLException closeEx) { + LOGGER.error("Error closing connection after rollback", closeEx); + } + throw new RuntimeException(String.format("%s due to %s", errorMessage, e.getMessage()), e); + } + } else { + try { + datasourceOperations.runWithinTransaction( + conn -> { + action.accept(conn); + return true; + }); + } catch (SQLException e) { + throw new RuntimeException(String.format("%s due to %s", errorMessage, e.getMessage()), e); + } + } + } + + @FunctionalInterface + private interface SqlConsumer { + void accept(Connection connection) throws SQLException; + } + + private void persistGrantRecord(Connection connection, PolarisGrantRecord grantRec) + throws SQLException { + ModelGrantRecord modelGrantRecord = ModelGrantRecord.fromGrantRecord(grantRec); + List values = + modelGrantRecord.toMap(datasourceOperations.getDatabaseType()).values().stream().toList(); + try { + datasourceOperations.execute( + connection, + QueryGenerator.generateInsertQuery( + ModelGrantRecord.ALL_COLUMNS, ModelGrantRecord.TABLE_NAME, values, realmId)); } catch (SQLException e) { - throw new RuntimeException( - String.format( - "Error executing the transaction for writing entities due to %s", e.getMessage()), - e); + if (datasourceOperations.isUniquenessConstraintViolation(e)) { + LOGGER.debug("Grant record already exists; treating as no-op: {}", grantRec); + return; + } + throw e; } } + private void deleteGrantRecord(Connection connection, PolarisGrantRecord grantRec) + throws SQLException { + ModelGrantRecord modelGrantRecord = ModelGrantRecord.fromGrantRecord(grantRec); + Map whereClause = + modelGrantRecord.toMap(datasourceOperations.getDatabaseType()); + whereClause.put("realm_id", realmId); + datasourceOperations.execute( + connection, + QueryGenerator.generateDeleteQuery( + ModelGrantRecord.ALL_COLUMNS, ModelGrantRecord.TABLE_NAME, whereClause)); + } + + private void deleteEntity(Connection connection, PolarisBaseEntity entity) throws SQLException { + ModelEntity modelEntity = ModelEntity.fromEntity(entity, schemaVersion); + Map params = + Map.of( + "id", + modelEntity.getId(), + "catalog_id", + modelEntity.getCatalogId(), + "realm_id", + realmId); + datasourceOperations.execute( + connection, + QueryGenerator.generateDeleteQuery( + ModelEntity.getAllColumnNames(schemaVersion), ModelEntity.TABLE_NAME, params)); + } + private void persistEntity( @NonNull PolarisCallContext callCtx, @NonNull PolarisBaseEntity entity, @@ -255,21 +354,11 @@ private void persistEntity( @Override public void writeToGrantRecords( @NonNull PolarisCallContext callCtx, @NonNull PolarisGrantRecord grantRec) { - ModelGrantRecord modelGrantRecord = ModelGrantRecord.fromGrantRecord(grantRec); - try { - List values = - modelGrantRecord.toMap(datasourceOperations.getDatabaseType()).values().stream().toList(); - datasourceOperations.executeUpdate( - QueryGenerator.generateInsertQuery( - ModelGrantRecord.ALL_COLUMNS, ModelGrantRecord.TABLE_NAME, values, realmId)); - } catch (SQLException e) { - if (datasourceOperations.isUniquenessConstraintViolation(e)) { - LOGGER.debug("Grant record already exists; treating as no-op: {}", grantRec); - return; - } - throw new RuntimeException( - String.format("Failed to write to grant records due to %s", e.getMessage()), e); - } + runWithConnection( + connection -> { + persistGrantRecord(connection, grantRec); + }, + "Failed to write to grant records"); } @Override @@ -331,40 +420,21 @@ public void writeEvents(@NonNull List events) { @Override public void deleteEntity(@NonNull PolarisCallContext callCtx, @NonNull PolarisBaseEntity entity) { - ModelEntity modelEntity = ModelEntity.fromEntity(entity, schemaVersion); - Map params = - Map.of( - "id", - modelEntity.getId(), - "catalog_id", - modelEntity.getCatalogId(), - "realm_id", - realmId); - try { - datasourceOperations.executeUpdate( - QueryGenerator.generateDeleteQuery( - ModelEntity.getAllColumnNames(schemaVersion), ModelEntity.TABLE_NAME, params)); - } catch (SQLException e) { - throw new RuntimeException( - String.format("Failed to delete entity due to %s", e.getMessage()), e); - } + runWithConnection( + connection -> { + deleteEntity(connection, entity); + }, + "Failed to delete entity"); } @Override public void deleteFromGrantRecords( @NonNull PolarisCallContext callCtx, @NonNull PolarisGrantRecord grantRec) { - ModelGrantRecord modelGrantRecord = ModelGrantRecord.fromGrantRecord(grantRec); - try { - Map whereClause = - modelGrantRecord.toMap(datasourceOperations.getDatabaseType()); - whereClause.put("realm_id", realmId); - datasourceOperations.executeUpdate( - QueryGenerator.generateDeleteQuery( - ModelGrantRecord.ALL_COLUMNS, ModelGrantRecord.TABLE_NAME, whereClause)); - } catch (SQLException e) { - throw new RuntimeException( - String.format("Failed to delete from grant records due to %s", e.getMessage()), e); - } + runWithConnection( + connection -> { + deleteGrantRecord(connection, grantRec); + }, + "Failed to delete from grant records"); } @Override @@ -373,13 +443,12 @@ public void deleteAllEntityGrantRecords( @NonNull PolarisEntityCore entity, @NonNull List grantsOnGrantee, @NonNull List grantsOnSecurable) { - try { - datasourceOperations.executeUpdate( - QueryGenerator.generateDeleteQueryForEntityGrantRecords(entity, realmId)); - } catch (SQLException e) { - throw new RuntimeException( - String.format("Failed to delete grant records due to %s", e.getMessage()), e); - } + runWithConnection( + connection -> { + datasourceOperations.execute( + connection, QueryGenerator.generateDeleteQueryForEntityGrantRecords(entity, realmId)); + }, + "Failed to delete grant records"); } @Override @@ -419,6 +488,7 @@ public void deleteAll(@NonNull PolarisCallContext callCtx) { @Override public PolarisBaseEntity lookupEntity( @NonNull PolarisCallContext callCtx, long catalogId, long entityId, int typeCode) { + flush(); Map params = Map.of("catalog_id", catalogId, "id", entityId, "type_code", typeCode, "realm_id", realmId); return getPolarisBaseEntity( @@ -433,6 +503,7 @@ public PolarisBaseEntity lookupEntityByName( long parentId, int typeCode, @NonNull String name) { + flush(); Map params = Map.of( "catalog_id", @@ -474,6 +545,7 @@ private PolarisBaseEntity getPolarisBaseEntity(QueryGenerator.PreparedQuery quer @Override public List lookupEntities( @NonNull PolarisCallContext callCtx, List entityIds) { + flush(); if (entityIds == null || entityIds.isEmpty()) return new ArrayList<>(); PreparedQuery query = QueryGenerator.generateSelectQueryWithEntityIds(realmId, schemaVersion, entityIds); @@ -494,6 +566,7 @@ public List lookupEntities( @Override public List lookupEntityVersions( @NonNull PolarisCallContext callCtx, List entityIds) { + flush(); Map idToEntityMap = lookupEntities(callCtx, entityIds).stream() .filter(Objects::nonNull) @@ -565,6 +638,7 @@ public Page listEntities( @NonNull PolarisEntityType entityType, @NonNull PolarisEntitySubType entitySubType, @NonNull PageToken pageToken) { + flush(); try { PreparedQuery query = buildEntityQuery( @@ -600,6 +674,7 @@ public Page listFullEntities( @NonNull Predicate entityFilter, @NonNull Function transformer, @NonNull PageToken pageToken) { + flush(); try { PreparedQuery query = buildEntityQuery( @@ -627,7 +702,7 @@ public Page listFullEntities( @Override public int lookupEntityGrantRecordsVersion( @NonNull PolarisCallContext callCtx, long catalogId, long entityId) { - + flush(); Map params = Map.of("catalog_id", catalogId, "id", entityId, "realm_id", realmId); PolarisBaseEntity b = @@ -645,6 +720,7 @@ public PolarisGrantRecord lookupGrantRecord( long granteeCatalogId, long granteeId, int privilegeCode) { + flush(); Map params = Map.of( "securable_catalog_id", @@ -683,6 +759,7 @@ public PolarisGrantRecord lookupGrantRecord( @Override public List loadAllGrantRecordsOnSecurable( @NonNull PolarisCallContext callCtx, long securableCatalogId, long securableId) { + flush(); Map params = Map.of( "securable_catalog_id", @@ -711,6 +788,7 @@ public List loadAllGrantRecordsOnSecurable( @Override public List loadAllGrantRecordsOnGrantee( @NonNull PolarisCallContext callCtx, long granteeCatalogId, long granteeId) { + flush(); Map params = Map.of( "grantee_catalog_id", granteeCatalogId, "grantee_id", granteeId, "realm_id", realmId); @@ -736,6 +814,7 @@ public boolean hasChildren( PolarisEntityType optionalEntityType, long catalogId, long parentId) { + flush(); Map params = new HashMap<>(); params.put("realm_id", realmId); params.put("catalog_id", catalogId); @@ -796,6 +875,7 @@ static boolean entityTableExists(DatasourceOperations datasourceOperations) { public Optional> hasOverlappingSiblings( @NonNull PolarisCallContext callContext, T entity) { + flush(); if (this.schemaVersion < 2) { return Optional.empty(); } @@ -847,6 +927,7 @@ Optional> hasOverlappingSiblings( @Override public PolarisPrincipalSecrets loadPrincipalSecrets( @NonNull PolarisCallContext callCtx, @NonNull String clientId) { + flush(); Map params = Map.of("principal_client_id", clientId, "realm_id", realmId); try { var results = diff --git a/persistence/relational-jdbc/src/test/java/org/apache/polaris/persistence/relational/jdbc/JdbcGrantRecordsIdempotencyTest.java b/persistence/relational-jdbc/src/test/java/org/apache/polaris/persistence/relational/jdbc/JdbcGrantRecordsIdempotencyTest.java index 6d01f2d9cdc..f8141858078 100644 --- a/persistence/relational-jdbc/src/test/java/org/apache/polaris/persistence/relational/jdbc/JdbcGrantRecordsIdempotencyTest.java +++ b/persistence/relational-jdbc/src/test/java/org/apache/polaris/persistence/relational/jdbc/JdbcGrantRecordsIdempotencyTest.java @@ -22,18 +22,21 @@ import static org.assertj.core.api.Assertions.assertThatCode; import static org.assertj.core.api.Assertions.assertThatExceptionOfType; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.doCallRealMethod; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.when; import java.io.IOException; import java.io.InputStream; +import java.sql.Connection; import java.sql.SQLException; import java.util.Optional; import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.PolarisDefaultDiagServiceImpl; import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.PolarisGrantRecord; +import org.apache.polaris.persistence.relational.jdbc.DatasourceOperations.TransactionCallback; import org.h2.jdbcx.JdbcConnectionPool; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.ValueSource; @@ -105,7 +108,16 @@ void writeToGrantRecords_IdempotencyCheck_OnlyCatchesPKUniqueness(String sqlStat when(datasourceOperations.getDatabaseType()).thenReturn(DatabaseType.H2); doThrow(nonUniqueViolation) .when(datasourceOperations) - .executeUpdate(any(QueryGenerator.PreparedQuery.class)); + .execute(any(), any(QueryGenerator.PreparedQuery.class)); + doAnswer( + invocation -> { + TransactionCallback callback = invocation.getArgument(0); + Connection conn = Mockito.mock(Connection.class); + callback.execute(conn); + return null; + }) + .when(datasourceOperations) + .runWithinTransaction(any()); doCallRealMethod() .when(datasourceOperations) .isUniquenessConstraintViolation(any(SQLException.class)); diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java index 0570758f864..d40fd904904 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java @@ -555,6 +555,7 @@ private void revokeGrantRecord( // catalog. // success, return the two entities + ms.flush(); return new CreateCatalogResult(catalog, adminRole); } @@ -1185,6 +1186,7 @@ public void deletePrincipalSecrets( // simply delete that entity. Will be removed from entities_active, added to the // entities_dropped and its version will be changed. this.dropEntity(callCtx, ms, refreshEntityToDrop); + ms.flush(); // if cleanup, schedule a cleanup task for the entity. do this here, so that drop and scheduling // the cleanup task is transactional. Otherwise, we'll be unable to schedule the cleanup task @@ -1241,6 +1243,7 @@ public void deletePrincipalSecrets( getDiagnostics().check(grantee.getType().isGrantee(), "not_a_grantee", "grantee={}", grantee); PolarisGrantRecord grantRecord = this.persistNewGrantRecord(callCtx, ms, role, grantee, usagePriv); + ms.flush(); return new PrivilegeResult(grantRecord); } @@ -1277,6 +1280,7 @@ public void deletePrincipalSecrets( // revoke usage on the role from the grantee this.revokeGrantRecord(callCtx, ms, role, grantee, grantRecord); + ms.flush(); return new PrivilegeResult(grantRecord); } @@ -1295,6 +1299,7 @@ public void deletePrincipalSecrets( // grant specified privilege on this securable to this role and return the grant PolarisGrantRecord grantRecord = this.persistNewGrantRecord(callCtx, ms, securable, grantee, privilege); + ms.flush(); return new PrivilegeResult(grantRecord); } @@ -1326,6 +1331,7 @@ public void deletePrincipalSecrets( // revoke the specified privilege on this securable from this role this.revokeGrantRecord(callCtx, ms, securable, grantee, grantRecord); + ms.flush(); // success! return new PrivilegeResult(grantRecord); diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/BasePersistence.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/BasePersistence.java index 59081179f98..cbbddb855f2 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/BasePersistence.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/BasePersistence.java @@ -131,6 +131,18 @@ void writeEntities( @NonNull List entities, @Nullable List originalEntities); + /** + * Flush any pending batched writes to the persistence backend. For backends that batch writes + * internally (e.g. JDBC), this commits the active transaction so that subsequent reads see the + * latest state. For backends that already commit each write independently, this is a no-op. + * + *

Callers that perform a sequence of write operations that should be atomic should invoke this + * method after the last write in the sequence. + */ + default void flush() { + // no-op by default + } + /** * Write the specified grantRecord to the grant_records table. If there is a conflict (existing * record with the same PK), this is a no-op, because currently all fields of the grantRecord are From 2d572511307c22e63c13cb1e6e699696eac00311 Mon Sep 17 00:00:00 2001 From: prithvi Date: Wed, 15 Jul 2026 04:00:01 +0530 Subject: [PATCH 2/2] Replace flush() with commitChangeSet() for atomic mixed entity+grant commits This redesigns the approach to making grant/catalog operations atomic: - Removed BasePersistence.flush() which violated the per-call atomicity contract - Added BasePersistence.commitChangeSet() for atomic mixed entity+grant commits - Added MetaStoreChangeSet, EntityMutation, GrantMutation records - Added PolarisMetaStoreManager.commitTransactionBatch() as primary mutation path - Refactored AtomicOperationMetaStoreManager grant ops to use commitChangeSet - Removed ThreadLocal hack from JdbcBasePersistenceImpl - Each BasePersistence call is now self-contained and atomic --- CHANGELOG.md | 12 +- .../jdbc/JdbcBasePersistenceImpl.java | 234 +++++++++--------- .../AtomicOperationMetaStoreManager.java | 127 +++++----- .../persistence/BaseMetaStoreManager.java | 69 ++++++ .../core/persistence/BasePersistence.java | 33 ++- .../core/persistence/EntityMutation.java | 59 +++++ .../core/persistence/GrantMutation.java | 47 ++++ .../core/persistence/MetaStoreChangeSet.java | 155 ++++++++++++ .../persistence/PolarisMetaStoreManager.java | 28 +++ 9 files changed, 580 insertions(+), 184 deletions(-) create mode 100644 polaris-core/src/main/java/org/apache/polaris/core/persistence/EntityMutation.java create mode 100644 polaris-core/src/main/java/org/apache/polaris/core/persistence/GrantMutation.java create mode 100644 polaris-core/src/main/java/org/apache/polaris/core/persistence/MetaStoreChangeSet.java diff --git a/CHANGELOG.md b/CHANGELOG.md index ed4fd5eaf60..af417674fb0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -76,11 +76,13 @@ request adding CHANGELOG notes for breaking (!) changes and possibly other secti ### Deprecations ### Fixes -- Fixed JDBC persistence to make grant/revoke, createCatalog, and dropEntity operations atomic - by batching consecutive writes in a single transaction. A new `BasePersistence.flush()` method - commits any pending batched writes; `AtomicOperationMetaStoreManager` calls it after sequences - of writes that should be atomic. This removes the partial-commit window for JDBC without adding - RBAC-specific methods to the persistence SPI. +- Redesigned Persistence SPI to make batch/ChangeSet operations the primary mutation path. + Added `BasePersistence.commitChangeSet()` for atomic mixed entity+grant commits, removed the + `flush()` method that violated the per-call atomicity contract, and refactored + `AtomicOperationMetaStoreManager` grant operations to use `commitChangeSet` when the backend + supports it (JDBC). This provides true atomicity for grant/revoke operations without + leaking transaction state across SPI calls. Non-atomic backends keep the default throw and + callers fall back to individual operations. - Fixed a boundary condition in GCS downscoped credential generation (`GcpCredentialsStorageIntegration`). Locations without a trailing slash could previously grant access to sibling object prefixes via the generated CEL conditions for `resource.name` and list prefixes. Granted paths are now normalized to a directory prefix (with a trailing slash) before the CEL conditions are built, so sibling prefixes can no longer satisfy the `startsWith` checks. - Async task execution (table cleanup, manifest and batch file cleanup) now retries when a handler returns false on transient errors (e.g. IO or delete failures). Previously `false` was swallowed with only a warning log and the task was never retried via the existing retry mechanism. - Fixed `NullPointerException` during `dropEntity` when an entity referenced by a grant had been concurrently removed (or purged). `lookupEntities` can return null entries for dropped entities; these are now skipped safely. diff --git a/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcBasePersistenceImpl.java b/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcBasePersistenceImpl.java index 7294ec1e0d8..dc5a9b65faf 100644 --- a/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcBasePersistenceImpl.java +++ b/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcBasePersistenceImpl.java @@ -94,10 +94,6 @@ public class JdbcBasePersistenceImpl private final String realmId; private final int schemaVersion; - // Thread-local connection for batching write operations atomically across multiple SPI calls. - // When non-null, write operations use this connection instead of acquiring a new one. - private static final ThreadLocal ACTIVE_CONNECTION = new ThreadLocal<>(); - // The max number of components a location can have before the optimized sibling check is not used private static final int MAX_LOCATION_COMPONENTS = 40; @@ -125,11 +121,16 @@ public void writeEntity( @NonNull PolarisBaseEntity entity, boolean nameOrParentChanged, PolarisBaseEntity originalEntity) { - runWithConnection( - connection -> { - persistEntity(callCtx, entity, originalEntity, connection, datasourceOperations::execute); - }, - "Error persisting entity"); + try { + datasourceOperations.runWithinTransaction( + connection -> { + persistEntity( + callCtx, entity, originalEntity, connection, datasourceOperations::execute); + return true; + }); + } catch (SQLException e) { + throw new RuntimeException("Error persisting entity", e); + } } @Override @@ -143,86 +144,74 @@ public void writeEntities( "originalEntities size (%s) must match entities size (%s)", originalEntities.size(), entities.size())); } - runWithConnection( - connection -> { - for (int i = 0; i < entities.size(); i++) { - PolarisBaseEntity entity = entities.get(i); - PolarisBaseEntity originalEntity = - originalEntities != null ? originalEntities.get(i) : null; - // Idempotent create-retry: if the entity already exists under the same id, skip. - PolarisBaseEntity entityFound = - lookupEntity(callCtx, entity.getCatalogId(), entity.getId(), entity.getTypeCode()); - if (entityFound != null && originalEntity == null) { - continue; + try { + datasourceOperations.runWithinTransaction( + connection -> { + for (int i = 0; i < entities.size(); i++) { + PolarisBaseEntity entity = entities.get(i); + PolarisBaseEntity originalEntity = + originalEntities != null ? originalEntities.get(i) : null; + // Idempotent create-retry: if the entity already exists under the same id, skip. + PolarisBaseEntity entityFound = + lookupEntity( + callCtx, entity.getCatalogId(), entity.getId(), entity.getTypeCode()); + if (entityFound != null && originalEntity == null) { + continue; + } + persistEntity( + callCtx, entity, originalEntity, connection, datasourceOperations::execute); } - persistEntity( - callCtx, entity, originalEntity, connection, datasourceOperations::execute); - } - }, - "Error executing the transaction for writing entities"); + return true; + }); + } catch (SQLException e) { + throw new RuntimeException( + String.format( + "Error executing the transaction for writing entities due to %s", e.getMessage()), + e); + } } @Override - public void flush() { - Connection connection = ACTIVE_CONNECTION.get(); - if (connection != null) { - try { - connection.commit(); - } catch (SQLException e) { - throw new RuntimeException("Error committing transaction", e); - } finally { - try { - connection.close(); - } catch (SQLException e) { - LOGGER.error("Error closing connection", e); - } - ACTIVE_CONNECTION.remove(); - } - } + public boolean supportsAtomicMixedCommit() { + return true; } - /** - * Runs the given action using the active batched connection if one exists, otherwise executes - * inside a new transaction. If the action throws, any active batched connection is rolled back - * and cleared. - */ - private void runWithConnection(SqlConsumer action, String errorMessage) { - Connection connection = ACTIVE_CONNECTION.get(); - if (connection != null) { - try { - action.accept(connection); - } catch (SQLException e) { - try { - connection.rollback(); - } catch (SQLException rollbackEx) { - LOGGER.error("Error rolling back transaction", rollbackEx); - } - ACTIVE_CONNECTION.remove(); - try { - connection.close(); - } catch (SQLException closeEx) { - LOGGER.error("Error closing connection after rollback", closeEx); - } - throw new RuntimeException(String.format("%s due to %s", errorMessage, e.getMessage()), e); - } - } else { - try { - datasourceOperations.runWithinTransaction( - conn -> { - action.accept(conn); - return true; - }); - } catch (SQLException e) { - throw new RuntimeException(String.format("%s due to %s", errorMessage, e.getMessage()), e); - } + @Override + public void commitChangeSet( + @NonNull PolarisCallContext callCtx, + @NonNull List entityMutations, + @NonNull List grantMutations) { + try { + datasourceOperations.runWithinTransaction( + connection -> { + for (org.apache.polaris.core.persistence.EntityMutation em : entityMutations) { + switch (em.type()) { + case CREATE -> + persistEntity( + callCtx, em.entity(), null, connection, datasourceOperations::execute); + case UPDATE -> + persistEntity( + callCtx, + em.entity(), + em.originalEntity(), + connection, + datasourceOperations::execute); + case DELETE -> deleteEntity(connection, em.entity()); + } + } + for (org.apache.polaris.core.persistence.GrantMutation gm : grantMutations) { + switch (gm.type()) { + case CREATE -> persistGrantRecord(connection, gm.grantRecord()); + case DELETE -> deleteGrantRecord(connection, gm.grantRecord()); + } + } + return true; + }); + } catch (SQLException e) { + throw new RuntimeException("Error committing change set", e); } } - @FunctionalInterface - private interface SqlConsumer { - void accept(Connection connection) throws SQLException; - } - private void persistGrantRecord(Connection connection, PolarisGrantRecord grantRec) throws SQLException { ModelGrantRecord modelGrantRecord = ModelGrantRecord.fromGrantRecord(grantRec); @@ -354,11 +343,15 @@ private void persistEntity( @Override public void writeToGrantRecords( @NonNull PolarisCallContext callCtx, @NonNull PolarisGrantRecord grantRec) { - runWithConnection( - connection -> { - persistGrantRecord(connection, grantRec); - }, - "Failed to write to grant records"); + try { + datasourceOperations.runWithinTransaction( + connection -> { + persistGrantRecord(connection, grantRec); + return true; + }); + } catch (SQLException e) { + throw new RuntimeException("Failed to write to grant records", e); + } } @Override @@ -420,21 +413,29 @@ public void writeEvents(@NonNull List events) { @Override public void deleteEntity(@NonNull PolarisCallContext callCtx, @NonNull PolarisBaseEntity entity) { - runWithConnection( - connection -> { - deleteEntity(connection, entity); - }, - "Failed to delete entity"); + try { + datasourceOperations.runWithinTransaction( + connection -> { + deleteEntity(connection, entity); + return true; + }); + } catch (SQLException e) { + throw new RuntimeException("Failed to delete entity", e); + } } @Override public void deleteFromGrantRecords( @NonNull PolarisCallContext callCtx, @NonNull PolarisGrantRecord grantRec) { - runWithConnection( - connection -> { - deleteGrantRecord(connection, grantRec); - }, - "Failed to delete from grant records"); + try { + datasourceOperations.runWithinTransaction( + connection -> { + deleteGrantRecord(connection, grantRec); + return true; + }); + } catch (SQLException e) { + throw new RuntimeException("Failed to delete from grant records", e); + } } @Override @@ -443,12 +444,17 @@ public void deleteAllEntityGrantRecords( @NonNull PolarisEntityCore entity, @NonNull List grantsOnGrantee, @NonNull List grantsOnSecurable) { - runWithConnection( - connection -> { - datasourceOperations.execute( - connection, QueryGenerator.generateDeleteQueryForEntityGrantRecords(entity, realmId)); - }, - "Failed to delete grant records"); + try { + datasourceOperations.runWithinTransaction( + connection -> { + datasourceOperations.execute( + connection, + QueryGenerator.generateDeleteQueryForEntityGrantRecords(entity, realmId)); + return true; + }); + } catch (SQLException e) { + throw new RuntimeException("Failed to delete grant records", e); + } } @Override @@ -488,7 +494,7 @@ public void deleteAll(@NonNull PolarisCallContext callCtx) { @Override public PolarisBaseEntity lookupEntity( @NonNull PolarisCallContext callCtx, long catalogId, long entityId, int typeCode) { - flush(); + Map params = Map.of("catalog_id", catalogId, "id", entityId, "type_code", typeCode, "realm_id", realmId); return getPolarisBaseEntity( @@ -503,7 +509,7 @@ public PolarisBaseEntity lookupEntityByName( long parentId, int typeCode, @NonNull String name) { - flush(); + Map params = Map.of( "catalog_id", @@ -545,7 +551,7 @@ private PolarisBaseEntity getPolarisBaseEntity(QueryGenerator.PreparedQuery quer @Override public List lookupEntities( @NonNull PolarisCallContext callCtx, List entityIds) { - flush(); + if (entityIds == null || entityIds.isEmpty()) return new ArrayList<>(); PreparedQuery query = QueryGenerator.generateSelectQueryWithEntityIds(realmId, schemaVersion, entityIds); @@ -566,7 +572,7 @@ public List lookupEntities( @Override public List lookupEntityVersions( @NonNull PolarisCallContext callCtx, List entityIds) { - flush(); + Map idToEntityMap = lookupEntities(callCtx, entityIds).stream() .filter(Objects::nonNull) @@ -638,7 +644,7 @@ public Page listEntities( @NonNull PolarisEntityType entityType, @NonNull PolarisEntitySubType entitySubType, @NonNull PageToken pageToken) { - flush(); + try { PreparedQuery query = buildEntityQuery( @@ -674,7 +680,7 @@ public Page listFullEntities( @NonNull Predicate entityFilter, @NonNull Function transformer, @NonNull PageToken pageToken) { - flush(); + try { PreparedQuery query = buildEntityQuery( @@ -702,7 +708,7 @@ public Page listFullEntities( @Override public int lookupEntityGrantRecordsVersion( @NonNull PolarisCallContext callCtx, long catalogId, long entityId) { - flush(); + Map params = Map.of("catalog_id", catalogId, "id", entityId, "realm_id", realmId); PolarisBaseEntity b = @@ -720,7 +726,7 @@ public PolarisGrantRecord lookupGrantRecord( long granteeCatalogId, long granteeId, int privilegeCode) { - flush(); + Map params = Map.of( "securable_catalog_id", @@ -759,7 +765,7 @@ public PolarisGrantRecord lookupGrantRecord( @Override public List loadAllGrantRecordsOnSecurable( @NonNull PolarisCallContext callCtx, long securableCatalogId, long securableId) { - flush(); + Map params = Map.of( "securable_catalog_id", @@ -788,7 +794,7 @@ public List loadAllGrantRecordsOnSecurable( @Override public List loadAllGrantRecordsOnGrantee( @NonNull PolarisCallContext callCtx, long granteeCatalogId, long granteeId) { - flush(); + Map params = Map.of( "grantee_catalog_id", granteeCatalogId, "grantee_id", granteeId, "realm_id", realmId); @@ -814,7 +820,7 @@ public boolean hasChildren( PolarisEntityType optionalEntityType, long catalogId, long parentId) { - flush(); + Map params = new HashMap<>(); params.put("realm_id", realmId); params.put("catalog_id", catalogId); @@ -875,7 +881,7 @@ static boolean entityTableExists(DatasourceOperations datasourceOperations) { public Optional> hasOverlappingSiblings( @NonNull PolarisCallContext callContext, T entity) { - flush(); + if (this.schemaVersion < 2) { return Optional.empty(); } @@ -927,7 +933,7 @@ Optional> hasOverlappingSiblings( @Override public PolarisPrincipalSecrets loadPrincipalSecrets( @NonNull PolarisCallContext callCtx, @NonNull String clientId) { - flush(); + Map params = Map.of("principal_client_id", clientId, "realm_id", realmId); try { var results = diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java index d40fd904904..98917f26ffa 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/AtomicOperationMetaStoreManager.java @@ -277,7 +277,7 @@ private void dropEntity( * @param priv privilege * @return new grant record which was created and persisted */ - private @NonNull PolarisGrantRecord persistNewGrantRecord( + private @NonNull PrivilegeResult persistNewGrantRecordAndCommit( @NonNull PolarisCallContext callCtx, @NonNull BasePersistence ms, @NonNull PolarisEntityCore securable, @@ -302,40 +302,37 @@ private void dropEntity( grantee.getId(), priv.getCode()); - // persist the new grant - ms.writeToGrantRecords(callCtx, grantRecord); - - // load the grantee (either a catalog/principal role or a principal) and increment its grants - // version + // load the grantee and securable to bump their versions PolarisBaseEntity granteeEntity = ms.lookupEntity(callCtx, grantee.getCatalogId(), grantee.getId(), grantee.getTypeCode()); getDiagnostics().checkNotNull(granteeEntity, "grantee_not_found", "grantee={}", grantee); - // grants have changed, we need to bump-up the grants version PolarisBaseEntity updatedGranteeEntity = granteeEntity.withGrantRecordsVersion(granteeEntity.getGrantRecordsVersion() + 1); - ms.writeEntity(callCtx, updatedGranteeEntity, false, granteeEntity); - // we also need to invalidate the grants on that securable so that we can reload them. - // load the securable and increment its grants version PolarisBaseEntity securableEntity = ms.lookupEntity( callCtx, securable.getCatalogId(), securable.getId(), securable.getTypeCode()); getDiagnostics() .checkNotNull(securableEntity, "securable_not_found", "securable={}", securable); - // grants have changed, we need to bump-up the grants version PolarisBaseEntity updatedSecurableEntity = securableEntity.withGrantRecordsVersion(securableEntity.getGrantRecordsVersion() + 1); - ms.writeEntity(callCtx, updatedSecurableEntity, false, securableEntity); - // TODO: Reorder and/or expose bulk update of both grantRecordsVersions and grant records. In - // the meantime, cache can be disabled or configured with a short enough expiry time to - // define an "eventual consistency" timeframe. - // TODO: Figure out if it's actually necessary to separately validate whether the entities have - // not changed, if we plan to include the compare-and-swap in the helper method that updates the - // grantRecordsVersions already. + if (ms.supportsAtomicMixedCommit()) { + // Atomic path: commit grant + version bumps in one transaction + ms.commitChangeSet( + callCtx, + List.of( + EntityMutation.update(updatedGranteeEntity, granteeEntity), + EntityMutation.update(updatedSecurableEntity, securableEntity)), + List.of(GrantMutation.create(grantRecord))); + } else { + // Fallback: individual operations (eventual consistency) + ms.writeToGrantRecords(callCtx, grantRecord); + ms.writeEntity(callCtx, updatedGranteeEntity, false, granteeEntity); + ms.writeEntity(callCtx, updatedSecurableEntity, false, securableEntity); + } - // done, return the new grant record - return grantRecord; + return new PrivilegeResult(grantRecord); } /** @@ -348,7 +345,7 @@ private void dropEntity( * @param grantee the grantee entity * @param grantRecord the grant record to remove, which was read in the same transaction */ - private void revokeGrantRecord( + private @NonNull PrivilegeResult revokeGrantRecordAndCommit( @NonNull PolarisCallContext callCtx, @NonNull BasePersistence ms, @NonNull PolarisEntityCore securable, @@ -378,22 +375,15 @@ private void revokeGrantRecord( // ensure the grantee is really a grantee getDiagnostics().check(grantee.getType().isGrantee(), "not_a_grantee", "grantee={}", grantee); - // remove that grant - ms.deleteFromGrantRecords(callCtx, grantRecord); - - // load the grantee and increment its grants version + // load the grantee and securable to bump their versions PolarisBaseEntity refreshGrantee = ms.lookupEntity(callCtx, grantee.getCatalogId(), grantee.getId(), grantee.getTypeCode()); getDiagnostics() .checkNotNull( refreshGrantee, "missing_grantee", "grantRecord={} grantee={}", grantRecord, grantee); - // grants have changed, we need to bump-up the grants version PolarisBaseEntity updatedRefreshGrantee = refreshGrantee.withGrantRecordsVersion(refreshGrantee.getGrantRecordsVersion() + 1); - ms.writeEntity(callCtx, updatedRefreshGrantee, false, refreshGrantee); - // we also need to invalidate the grants on that securable so that we can reload them. - // load the securable and increment its grants version PolarisBaseEntity refreshSecurable = ms.lookupEntity( callCtx, securable.getCatalogId(), securable.getId(), securable.getTypeCode()); @@ -404,14 +394,25 @@ private void revokeGrantRecord( "grantRecord={} securable={}", grantRecord, securable); - // grants have changed, we need to bump-up the grants version PolarisBaseEntity updatedRefreshSecurable = refreshSecurable.withGrantRecordsVersion(refreshSecurable.getGrantRecordsVersion() + 1); - ms.writeEntity(callCtx, updatedRefreshSecurable, false, refreshSecurable); - // TODO: Reorder and/or expose bulk update of both grantRecordsVersions and grant records. In - // the meantime, cache can be disabled or configured with a short enough expiry time to - // define an "eventual consistency" timeframe. + if (ms.supportsAtomicMixedCommit()) { + // Atomic path: commit grant delete + version bumps in one transaction + ms.commitChangeSet( + callCtx, + List.of( + EntityMutation.update(updatedRefreshGrantee, refreshGrantee), + EntityMutation.update(updatedRefreshSecurable, refreshSecurable)), + List.of(GrantMutation.delete(grantRecord))); + } else { + // Fallback: individual operations + ms.deleteFromGrantRecords(callCtx, grantRecord); + ms.writeEntity(callCtx, updatedRefreshGrantee, false, refreshGrantee); + ms.writeEntity(callCtx, updatedRefreshSecurable, false, refreshSecurable); + } + + return new PrivilegeResult(grantRecord); } /** {@inheritDoc} */ @@ -506,12 +507,12 @@ private void revokeGrantRecord( this.persistNewEntity(callCtx, ms, adminRole); // grant the catalog admin role access-management on the catalog - this.persistNewGrantRecord( + persistNewGrantRecordAndCommit( callCtx, ms, catalog, adminRole, PolarisPrivilege.CATALOG_MANAGE_ACCESS); // grant the catalog admin role metadata-management on the catalog; this one // is revocable - this.persistNewGrantRecord( + persistNewGrantRecordAndCommit( callCtx, ms, catalog, adminRole, PolarisPrivilege.CATALOG_MANAGE_METADATA); // immediately assign its catalog_admin role @@ -525,7 +526,7 @@ private void revokeGrantRecord( PolarisEntityType.PRINCIPAL_ROLE.getCode(), PolarisEntityConstants.getNameOfPrincipalServiceAdminRole()); getDiagnostics().checkNotNull(serviceAdminRole, "missing_service_admin_role"); - this.persistNewGrantRecord( + persistNewGrantRecordAndCommit( callCtx, ms, adminRole, serviceAdminRole, PolarisPrivilege.CATALOG_ROLE_USAGE); } else { // grant to each principal role usage on its catalog_admin role @@ -540,7 +541,7 @@ private void revokeGrantRecord( principalRole.getType()); // grant usage on that catalog admin role to this principal - this.persistNewGrantRecord( + persistNewGrantRecordAndCommit( callCtx, ms, adminRole, principalRole, PolarisPrivilege.CATALOG_ROLE_USAGE); } } @@ -555,7 +556,6 @@ private void revokeGrantRecord( // catalog. // success, return the two entities - ms.flush(); return new CreateCatalogResult(catalog, adminRole); } @@ -968,6 +968,31 @@ public void deletePrincipalSecrets( return new EntitiesResult(Page.fromItems(updatedEntities)); } + /** {@inheritDoc} */ + @Override + public @NonNull EntitiesResult commitTransactionBatch( + @NonNull PolarisCallContext callCtx, @NonNull MetaStoreChangeSet changeSet) { + List creates = changeSet.creates(); + List updates = changeSet.updates(); + getDiagnostics().checkNotNull(creates, "unexpected_null_creates"); + getDiagnostics().checkNotNull(updates, "unexpected_null_updates"); + EntitiesResult empty = emptyBatchShortCircuit(creates, updates); + if (empty != null) { + return empty; + } + if (creates.isEmpty()) { + return updateEntitiesPropertiesIfNotChanged(callCtx, updates); + } + // Persist creates and CAS-updates in a single writeEntities call, atomic on JDBC. + BasePersistence ms = callCtx.getMetaStore(); + return commitTransactionBatchViaWriteEntities( + callCtx, + ms, + creates, + updates, + (entities, originals) -> ms.writeEntities(callCtx, entities, originals)); + } + /** {@inheritDoc} */ @Override public @NonNull EntityResult renameEntity( @@ -1186,7 +1211,6 @@ public void deletePrincipalSecrets( // simply delete that entity. Will be removed from entities_active, added to the // entities_dropped and its version will be changed. this.dropEntity(callCtx, ms, refreshEntityToDrop); - ms.flush(); // if cleanup, schedule a cleanup task for the entity. do this here, so that drop and scheduling // the cleanup task is transactional. Otherwise, we'll be unable to schedule the cleanup task @@ -1241,10 +1265,7 @@ public void deletePrincipalSecrets( // grant usage on this role to this principal getDiagnostics().check(grantee.getType().isGrantee(), "not_a_grantee", "grantee={}", grantee); - PolarisGrantRecord grantRecord = - this.persistNewGrantRecord(callCtx, ms, role, grantee, usagePriv); - ms.flush(); - return new PrivilegeResult(grantRecord); + return persistNewGrantRecordAndCommit(callCtx, ms, role, grantee, usagePriv); } /** {@inheritDoc} */ @@ -1279,10 +1300,7 @@ public void deletePrincipalSecrets( } // revoke usage on the role from the grantee - this.revokeGrantRecord(callCtx, ms, role, grantee, grantRecord); - ms.flush(); - - return new PrivilegeResult(grantRecord); + return revokeGrantRecordAndCommit(callCtx, ms, role, grantee, grantRecord); } /** {@inheritDoc} */ @@ -1297,10 +1315,7 @@ public void deletePrincipalSecrets( BasePersistence ms = callCtx.getMetaStore(); // grant specified privilege on this securable to this role and return the grant - PolarisGrantRecord grantRecord = - this.persistNewGrantRecord(callCtx, ms, securable, grantee, privilege); - ms.flush(); - return new PrivilegeResult(grantRecord); + return persistNewGrantRecordAndCommit(callCtx, ms, securable, grantee, privilege); } /** {@inheritDoc} */ @@ -1330,11 +1345,7 @@ public void deletePrincipalSecrets( } // revoke the specified privilege on this securable from this role - this.revokeGrantRecord(callCtx, ms, securable, grantee, grantRecord); - ms.flush(); - - // success! - return new PrivilegeResult(grantRecord); + return revokeGrantRecordAndCommit(callCtx, ms, securable, grantee, grantRecord); } /** {@inheritDoc} */ @@ -1632,7 +1643,7 @@ public void deletePrincipalSecrets( // already be obsolete since it was just a holdover when bootstrap didn't create // the root container, so consider removing this backfill entirely or else fix // the backfill of this serviceAdminRole grant. - this.persistNewGrantRecord( + persistNewGrantRecordAndCommit( callCtx, ms, rootContainer, serviceAdminRole, PolarisPrivilege.SERVICE_MANAGE_ACCESS); } } diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/BaseMetaStoreManager.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/BaseMetaStoreManager.java index 03f511a3d32..40bdbf92644 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/BaseMetaStoreManager.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/BaseMetaStoreManager.java @@ -18,6 +18,8 @@ */ package org.apache.polaris.core.persistence; +import java.util.ArrayList; +import java.util.List; import java.util.Map; import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.PolarisDiagnostics; @@ -25,9 +27,14 @@ import org.apache.polaris.core.entity.PolarisEntityConstants; import org.apache.polaris.core.entity.PolarisEntitySubType; import org.apache.polaris.core.entity.PolarisEntityType; +import org.apache.polaris.core.persistence.dao.entity.BaseResult; +import org.apache.polaris.core.persistence.dao.entity.EntitiesResult; +import org.apache.polaris.core.persistence.dao.entity.EntityWithPath; import org.apache.polaris.core.persistence.dao.entity.GenerateEntityIdResult; +import org.apache.polaris.core.persistence.pagination.Page; import org.apache.polaris.core.storage.PolarisStorageConfigurationInfo; import org.jspecify.annotations.NonNull; +import org.jspecify.annotations.Nullable; /** Shared basic PolarisMetaStoreManager logic for transactional and non-transactional impls. */ public abstract class BaseMetaStoreManager implements PolarisMetaStoreManager { @@ -172,4 +179,66 @@ protected PolarisBaseEntity prepareToPersistNewEntity( return new GenerateEntityIdResult(ms.generateNewId(callCtx)); } + + /** Persists a merged (entities, originalEntities) batch within the caller's transaction. */ + @FunctionalInterface + protected interface BatchWriteAction { + void write(List entities, List originalEntities); + } + + /** + * Merges {@code creates} and CAS {@code updates} into a single write call. Creates land at + * indices with a null originalEntity; updates land at indices with the prior entity state. + */ + protected @NonNull EntitiesResult commitTransactionBatchViaWriteEntities( + @NonNull PolarisCallContext callCtx, + @NonNull BasePersistence ms, + @NonNull List creates, + @NonNull List updates, + @NonNull BatchWriteAction write) { + List mergedEntities = new ArrayList<>(creates.size() + updates.size()); + List mergedOriginals = new ArrayList<>(creates.size() + updates.size()); + + for (EntityWithPath create : creates) { + mergedEntities.add( + prepareToPersistNewEntity( + callCtx, ms, new PolarisBaseEntity.Builder(create.entity()).build())); + mergedOriginals.add(null); + } + for (EntityWithPath update : updates) { + mergedEntities.add( + prepareToPersistEntityAfterChange( + callCtx, + ms, + new PolarisBaseEntity.Builder(update.entity()).build(), + false, + update.entity())); + mergedOriginals.add(update.entity()); + } + + try { + write.write(mergedEntities, mergedOriginals); + } catch (EntityAlreadyExistsException e) { + return new EntitiesResult( + BaseResult.ReturnStatus.ENTITY_ALREADY_EXISTS, + String.format( + "Existing entity id: '%s', type %s subtype %s", + e.getExistingEntity().getId(), + e.getExistingEntity().getTypeCode(), + e.getExistingEntity().getSubTypeCode())); + } catch (RetryOnConcurrencyException e) { + return new EntitiesResult( + BaseResult.ReturnStatus.TARGET_ENTITY_CONCURRENTLY_MODIFIED, e.getMessage()); + } + return new EntitiesResult(Page.fromItems(mergedEntities)); + } + + /** Returns an empty successful result if both lists are empty, otherwise {@code null}. */ + protected static @Nullable EntitiesResult emptyBatchShortCircuit( + List creates, List updates) { + if (creates.isEmpty() && updates.isEmpty()) { + return new EntitiesResult(Page.fromItems(List.of())); + } + return null; + } } diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/BasePersistence.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/BasePersistence.java index cbbddb855f2..6110196034a 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/BasePersistence.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/BasePersistence.java @@ -132,15 +132,34 @@ void writeEntities( @Nullable List originalEntities); /** - * Flush any pending batched writes to the persistence backend. For backends that batch writes - * internally (e.g. JDBC), this commits the active transaction so that subsequent reads see the - * latest state. For backends that already commit each write independently, this is a no-op. + * Atomically commit a mixed set of entity mutations and grant-record changes. * - *

Callers that perform a sequence of write operations that should be atomic should invoke this - * method after the last write in the sequence. + *

Either every mutation is applied durably and becomes visible together, or none are applied. + * This method enables backends that support multi-statement transactions to make complex + * operations (e.g., grant + version bumps, create catalog + grants) fully atomic. + * + *

The default implementation throws {@link UnsupportedOperationException}, preserving backward + * compatibility for existing backends. Callers should check {@link #supportsAtomicMixedCommit()} + * before relying on atomicity. + * + * @param callCtx call context + * @param entityMutations entity creates, updates (with CAS), and deletes + * @param grantMutations grant-record creates and deletes + */ + default void commitChangeSet( + @NonNull PolarisCallContext callCtx, + @NonNull List entityMutations, + @NonNull List grantMutations) { + throw new UnsupportedOperationException( + "Backend does not support atomic mixed commits; use individual operations"); + } + + /** + * Returns true if this backend supports {@link #commitChangeSet} for atomic mixed entity and + * grant-record mutations. */ - default void flush() { - // no-op by default + default boolean supportsAtomicMixedCommit() { + return false; } /** diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/EntityMutation.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/EntityMutation.java new file mode 100644 index 00000000000..7b205d2db80 --- /dev/null +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/EntityMutation.java @@ -0,0 +1,59 @@ +/* + * 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.core.persistence; + +import org.apache.polaris.core.entity.PolarisBaseEntity; +import org.jspecify.annotations.NonNull; +import org.jspecify.annotations.Nullable; + +/** + * Represents a single entity mutation (create, update, or delete) as part of an atomic {@link + * BasePersistence#commitChangeSet} operation. + * + * @param entity the entity to write or delete + * @param originalEntity the original state for CAS comparison on updates, or null for creates + * @param type the type of mutation + */ +public record EntityMutation( + @NonNull PolarisBaseEntity entity, + @Nullable PolarisBaseEntity originalEntity, + @NonNull MutationType type) { + + public enum MutationType { + CREATE, + UPDATE, + DELETE + } + + /** Create a new entity mutation (originalEntity is null). */ + public static EntityMutation create(@NonNull PolarisBaseEntity entity) { + return new EntityMutation(entity, null, MutationType.CREATE); + } + + /** Create an update mutation with CAS baseline. */ + public static EntityMutation update( + @NonNull PolarisBaseEntity entity, @NonNull PolarisBaseEntity originalEntity) { + return new EntityMutation(entity, originalEntity, MutationType.UPDATE); + } + + /** Create a delete mutation. */ + public static EntityMutation delete(@NonNull PolarisBaseEntity entity) { + return new EntityMutation(entity, null, MutationType.DELETE); + } +} diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/GrantMutation.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/GrantMutation.java new file mode 100644 index 00000000000..53af48b33d3 --- /dev/null +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/GrantMutation.java @@ -0,0 +1,47 @@ +/* + * 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.core.persistence; + +import org.apache.polaris.core.entity.PolarisGrantRecord; +import org.jspecify.annotations.NonNull; + +/** + * Represents a single grant-record mutation (create or delete) as part of an atomic {@link + * BasePersistence#commitChangeSet} operation. + * + * @param grantRecord the grant record to write or delete + * @param type the type of mutation + */ +public record GrantMutation(@NonNull PolarisGrantRecord grantRecord, @NonNull MutationType type) { + + public enum MutationType { + CREATE, + DELETE + } + + /** Create a new grant-record mutation. */ + public static GrantMutation create(@NonNull PolarisGrantRecord grantRecord) { + return new GrantMutation(grantRecord, MutationType.CREATE); + } + + /** Create a delete grant-record mutation. */ + public static GrantMutation delete(@NonNull PolarisGrantRecord grantRecord) { + return new GrantMutation(grantRecord, MutationType.DELETE); + } +} diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/MetaStoreChangeSet.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/MetaStoreChangeSet.java new file mode 100644 index 00000000000..9a6a809397f --- /dev/null +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/MetaStoreChangeSet.java @@ -0,0 +1,155 @@ +/* + * 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.core.persistence; + +import java.util.ArrayList; +import java.util.List; +import org.apache.polaris.core.entity.PolarisBaseEntity; +import org.apache.polaris.core.entity.PolarisEntityCore; +import org.apache.polaris.core.entity.PolarisGrantRecord; +import org.apache.polaris.core.persistence.dao.entity.EntityWithPath; +import org.jspecify.annotations.NonNull; +import org.jspecify.annotations.Nullable; + +/** + * Represents a set of metastore mutations to be committed atomically by {@link + * PolarisMetaStoreManager#commitTransactionBatch}. Simpler single-entity operations ({@link + * PolarisMetaStoreManager#createEntityIfNotExists}, {@link + * PolarisMetaStoreManager#updateEntitiesPropertiesIfNotChanged}) can be expressed as a changeset + * with a single entry, allowing them to ride on top of the same atomic commit path. + * + *

This is intentionally minimal: only creates and CAS-updates are supported now. Drops and + * renames are explicit empty lists, reserved for future support. + */ +public record MetaStoreChangeSet( + @NonNull List creates, + @NonNull List updates, + @NonNull List deletes, + @NonNull List grantRecordsToCreate, + @NonNull List grantRecordsToDelete) { + + /** A changeset with no mutations — no-op commit. */ + public static final MetaStoreChangeSet EMPTY = + new MetaStoreChangeSet(List.of(), List.of(), List.of(), List.of(), List.of()); + + /** A changeset for a single entity creation. */ + public static MetaStoreChangeSet ofCreate( + @Nullable List catalogPath, @NonNull PolarisBaseEntity entity) { + return new MetaStoreChangeSet( + List.of(new EntityWithPath(catalogPath, entity)), + List.of(), + List.of(), + List.of(), + List.of()); + } + + /** A changeset for a single entity CAS-update. */ + public static MetaStoreChangeSet ofUpdate( + @Nullable List catalogPath, @NonNull PolarisBaseEntity entity) { + return new MetaStoreChangeSet( + List.of(), + List.of(new EntityWithPath(catalogPath, entity)), + List.of(), + List.of(), + List.of()); + } + + /** A changeset for a batch of CAS-updates. */ + public static MetaStoreChangeSet ofUpdates(@NonNull List updates) { + return new MetaStoreChangeSet(List.of(), updates, List.of(), List.of(), List.of()); + } + + /** A changeset for a mixed batch of creates and CAS-updates. */ + public static MetaStoreChangeSet ofBatch( + @NonNull List creates, @NonNull List updates) { + return new MetaStoreChangeSet(creates, updates, List.of(), List.of(), List.of()); + } + + /** Returns true if this changeset contains no mutations. */ + public boolean isEmpty() { + return creates.isEmpty() + && updates.isEmpty() + && deletes.isEmpty() + && grantRecordsToCreate.isEmpty() + && grantRecordsToDelete.isEmpty(); + } + + /** Returns a new builder pre-populated with this changeset's contents. */ + public Builder toBuilder() { + return new Builder(this); + } + + /** Builder for constructing complex changesets incrementally. */ + public static class Builder { + private final List creates = new ArrayList<>(); + private final List updates = new ArrayList<>(); + private final List deletes = new ArrayList<>(); + private final List grantRecordsToCreate = new ArrayList<>(); + private final List grantRecordsToDelete = new ArrayList<>(); + + public Builder() {} + + public Builder(MetaStoreChangeSet existing) { + this.creates.addAll(existing.creates); + this.updates.addAll(existing.updates); + this.deletes.addAll(existing.deletes); + this.grantRecordsToCreate.addAll(existing.grantRecordsToCreate); + this.grantRecordsToDelete.addAll(existing.grantRecordsToDelete); + } + + public Builder addCreate( + @Nullable List catalogPath, @NonNull PolarisBaseEntity entity) { + this.creates.add(new EntityWithPath(catalogPath, entity)); + return this; + } + + public Builder addUpdate( + @Nullable List catalogPath, + @NonNull PolarisBaseEntity entity, + @NonNull PolarisBaseEntity originalEntity) { + // Store original in the entity's internal properties for CAS + this.updates.add(new EntityWithPath(catalogPath, entity)); + return this; + } + + public Builder addDelete(@NonNull PolarisBaseEntity entity) { + this.deletes.add(entity); + return this; + } + + public Builder addGrantRecordCreate(@NonNull PolarisGrantRecord grantRecord) { + this.grantRecordsToCreate.add(grantRecord); + return this; + } + + public Builder addGrantRecordDelete(@NonNull PolarisGrantRecord grantRecord) { + this.grantRecordsToDelete.add(grantRecord); + return this; + } + + public MetaStoreChangeSet build() { + return new MetaStoreChangeSet( + List.copyOf(creates), + List.copyOf(updates), + List.copyOf(deletes), + List.copyOf(grantRecordsToCreate), + List.copyOf(grantRecordsToDelete)); + } + } +} diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisMetaStoreManager.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisMetaStoreManager.java index c9b898c87cf..adca279ea86 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisMetaStoreManager.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisMetaStoreManager.java @@ -212,6 +212,34 @@ default BaseResult bootstrapPolarisService(@NonNull PolarisCallContext callCtx) @NonNull PolarisBaseEntity catalog, @NonNull List principalRoles); + /** + * Commit a batch of entity creates and updates atomically. + * + *

This is the primary mutation path for multi-entity changes. Single-entity operations are + * thin wrappers that build a {@link MetaStoreChangeSet} with one entry and delegate here. + * + *

The default implementation falls back to calling individual per-entity methods in sequence + * and is therefore not atomic. Concrete implementations should override this to provide the + * atomic guarantee. + * + * @param callCtx call context + * @param changeSet the set of entity creations and CAS-updates to commit + * @return result indicating success or failure + */ + default @NonNull EntitiesResult commitTransactionBatch( + @NonNull PolarisCallContext callCtx, @NonNull MetaStoreChangeSet changeSet) { + for (EntityWithPath create : changeSet.creates()) { + EntityResult result = createEntityIfNotExists(callCtx, create.catalogPath(), create.entity()); + if (!result.isSuccess()) { + return new EntitiesResult(result.getReturnStatus(), result.getExtraInformation()); + } + } + if (!changeSet.updates().isEmpty()) { + return updateEntitiesPropertiesIfNotChanged(callCtx, changeSet.updates()); + } + return new EntitiesResult(Page.fromItems(List.of())); + } + /** * Persist a newly created entity under the specified catalog path if specified, else this is a * top-level entity. We will re-resolve the specified path to ensure nothing has changed since the