From b84f4624db8d0bd5b8920b0df719bcc15666008f Mon Sep 17 00:00:00 2001 From: Alexandre Dutra Date: Tue, 14 Jan 2025 12:14:19 +0100 Subject: [PATCH] Remove CallContext and its ThreadLocal usage (#589) --- ...pseLinkPolarisMetaStoreManagerFactory.java | 29 +- ...olarisEclipseLinkMetaStoreSessionImpl.java | 182 +-- ...olarisEclipseLinkMetaStoreManagerTest.java | 25 +- .../polaris/core/PolarisCallContext.java | 84 -- .../polaris/core/PolarisConfiguration.java | 20 - .../core/PolarisConfigurationStore.java | 32 +- .../polaris/core/auth/PolarisAuthorizer.java | 3 + .../core/auth/PolarisAuthorizerImpl.java | 7 +- .../core/auth/PolarisGrantManager.java | 38 +- .../core/auth/PolarisSecretsManager.java | 10 +- .../polaris/core/context/CallContext.java | 158 --- .../polaris/core/context/RealmContext.java | 6 + .../polaris/core/entity/TaskEntity.java | 29 +- .../LocalPolarisMetaStoreManagerFactory.java | 77 +- .../persistence/PolarisEntityManager.java | 25 +- .../persistence/PolarisEntityResolver.java | 77 +- .../persistence/PolarisMetaStoreManager.java | 68 +- .../PolarisMetaStoreManagerImpl.java | 1190 +++++++---------- .../persistence/PolarisMetaStoreSession.java | 142 +- .../persistence/PolarisObjectMapperUtil.java | 24 +- .../PolarisTreeMapMetaStoreSessionImpl.java | 178 +-- .../core/persistence/PolarisTreeMapStore.java | 27 +- .../TransactionWorkspaceMetaStoreManager.java | 209 ++- .../core/persistence/cache/EntityCache.java | 57 +- .../persistence/cache/PolarisRemoteCache.java | 20 +- .../resolver/PolarisResolutionManifest.java | 16 +- .../core/persistence/resolver/Resolver.java | 25 +- .../storage/InMemoryStorageIntegration.java | 35 +- .../core/storage/PolarisCredentialVendor.java | 10 +- .../PolarisStorageConfigurationInfo.java | 19 +- .../storage/PolarisStorageIntegration.java | 8 +- .../aws/AwsCredentialsStorageIntegration.java | 16 +- .../AzureCredentialsStorageIntegration.java | 14 +- .../storage/cache/StorageCredentialCache.java | 39 +- .../cache/StorageCredentialCacheKey.java | 20 +- .../gcp/GcpCredentialsStorageIntegration.java | 9 +- .../core/persistence/EntityCacheTest.java | 73 +- .../PolarisTreeMapMetaStoreManagerTest.java | 15 +- .../core/persistence/ResolverTest.java | 33 +- .../InMemoryStorageIntegrationTest.java | 119 +- .../cache/StorageCredentialCacheTest.java | 147 +- .../PolarisConfigurationStoreTest.java | 14 +- .../AwsCredentialsStorageIntegrationTest.java | 45 +- ...AzureCredentialStorageIntegrationTest.java | 35 +- .../GcpCredentialsStorageIntegrationTest.java | 28 +- .../BasePolarisMetaStoreManagerTest.java | 274 ++-- .../PolarisTestMetaStoreManager.java | 181 +-- .../quarkus/config/QuarkusProducers.java | 57 +- .../quarkus/task/QuarkusTaskExecutorImpl.java | 25 +- .../polaris/service/quarkus/TestServices.java | 61 +- .../admin/PolarisAdminServiceAuthzTest.java | 5 +- .../quarkus/admin/PolarisAuthzTestBase.java | 118 +- .../quarkus/auth/JWTRSAKeyPairTest.java | 19 +- .../auth/JWTSymmetricKeyGeneratorTest.java | 37 +- .../catalog/BasePolarisCatalogTest.java | 191 +-- .../catalog/BasePolarisCatalogViewTest.java | 47 +- ...PolarisCatalogHandlerWrapperAuthzTest.java | 108 +- .../PolarisPassthroughResolutionView.java | 18 +- .../config/DefaultConfigurationStoreTest.java | 53 +- .../ManifestFileCleanupTaskHandlerTest.java | 662 +++++---- .../task/TableCleanupTaskHandlerTest.java | 1053 ++++++++------- .../test/PolarisIntegrationTestFixture.java | 42 +- .../test/PolarisIntegrationTestHelper.java | 8 - .../service/admin/PolarisAdminService.java | 184 ++- .../service/admin/PolarisServiceImpl.java | 61 +- .../auth/BasePolarisAuthenticator.java | 21 +- .../auth/DefaultActiveRolesProvider.java | 30 +- .../service/auth/DefaultOAuth2ApiService.java | 14 +- .../auth/DefaultPolarisAuthenticator.java | 19 +- .../polaris/service/auth/JWTBroker.java | 24 +- .../polaris/service/auth/JWTRSAKeyPair.java | 4 +- .../service/auth/JWTRSAKeyPairFactory.java | 1 + .../service/auth/JWTSymmetricKeyBroker.java | 4 +- .../service/auth/JWTSymmetricKeyFactory.java | 1 + .../service/auth/NoneTokenBrokerFactory.java | 7 +- ...InlineBearerTokenPolarisAuthenticator.java | 13 +- .../service/auth/TestOAuth2ApiService.java | 21 +- .../polaris/service/auth/TokenBroker.java | 16 +- .../service/catalog/BasePolarisCatalog.java | 218 ++- .../catalog/IcebergCatalogAdapter.java | 336 +++-- .../catalog/PolarisCatalogHandlerWrapper.java | 365 +++-- .../config/DefaultConfigurationStore.java | 22 +- .../config/RealmEntityManagerFactory.java | 9 +- .../context/CallContextCatalogFactory.java | 4 +- .../service/context/CallContextResolver.java | 29 - .../context/DefaultCallContextResolver.java | 69 - .../PolarisCallContextCatalogFactory.java | 50 +- ...nMemoryPolarisMetaStoreManagerFactory.java | 20 +- ...PolarisStorageIntegrationProviderImpl.java | 24 +- .../task/ManifestFileCleanupTaskHandler.java | 28 +- .../service/task/TableCleanupTaskHandler.java | 85 +- .../polaris/service/task/TaskExecutor.java | 4 +- .../service/task/TaskExecutorImpl.java | 63 +- .../service/task/TaskFileIOSupplier.java | 20 +- .../polaris/service/task/TaskHandler.java | 3 +- .../auth/BasePolarisAuthenticatorTest.java | 21 +- 96 files changed, 3714 insertions(+), 4472 deletions(-) delete mode 100644 polaris-core/src/main/java/org/apache/polaris/core/PolarisCallContext.java delete mode 100644 polaris-core/src/main/java/org/apache/polaris/core/context/CallContext.java delete mode 100644 service/common/src/main/java/org/apache/polaris/service/context/CallContextResolver.java delete mode 100644 service/common/src/main/java/org/apache/polaris/service/context/DefaultCallContextResolver.java diff --git a/extension/persistence/eclipselink/src/main/java/org/apache/polaris/extension/persistence/impl/eclipselink/EclipseLinkPolarisMetaStoreManagerFactory.java b/extension/persistence/eclipselink/src/main/java/org/apache/polaris/extension/persistence/impl/eclipselink/EclipseLinkPolarisMetaStoreManagerFactory.java index f7a119761..ade392be0 100644 --- a/extension/persistence/eclipselink/src/main/java/org/apache/polaris/extension/persistence/impl/eclipselink/EclipseLinkPolarisMetaStoreManagerFactory.java +++ b/extension/persistence/eclipselink/src/main/java/org/apache/polaris/extension/persistence/impl/eclipselink/EclipseLinkPolarisMetaStoreManagerFactory.java @@ -23,6 +23,8 @@ import jakarta.enterprise.context.ApplicationScoped; import jakarta.inject.Inject; import java.nio.file.Path; +import java.time.Clock; +import org.apache.polaris.core.PolarisConfigurationStore; import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.persistence.LocalPolarisMetaStoreManagerFactory; @@ -40,8 +42,24 @@ public class EclipseLinkPolarisMetaStoreManagerFactory extends LocalPolarisMetaStoreManagerFactory { - @Inject EclipseLinkConfiguration eclipseLinkConfiguration; - @Inject PolarisStorageIntegrationProvider storageIntegrationProvider; + private final EclipseLinkConfiguration eclipseLinkConfiguration; + private final PolarisStorageIntegrationProvider storageIntegrationProvider; + + public EclipseLinkPolarisMetaStoreManagerFactory() { + this(null, null, null, null, null); + } + + @Inject + public EclipseLinkPolarisMetaStoreManagerFactory( + EclipseLinkConfiguration eclipseLinkConfiguration, + PolarisStorageIntegrationProvider storageIntegrationProvider, + PolarisConfigurationStore configurationStore, + PolarisDiagnostics diagnostics, + Clock clock) { + super(configurationStore, diagnostics, clock); + this.eclipseLinkConfiguration = eclipseLinkConfiguration; + this.storageIntegrationProvider = storageIntegrationProvider; + } @Override protected PolarisEclipseLinkStore createBackingStore(@Nonnull PolarisDiagnostics diagnostics) { @@ -50,14 +68,17 @@ protected PolarisEclipseLinkStore createBackingStore(@Nonnull PolarisDiagnostics @Override protected PolarisMetaStoreSession createMetaStoreSession( - @Nonnull PolarisEclipseLinkStore store, @Nonnull RealmContext realmContext) { + @Nonnull PolarisEclipseLinkStore store, + @Nonnull RealmContext realmContext, + @Nonnull PolarisDiagnostics diagnostics) { return new PolarisEclipseLinkMetaStoreSessionImpl( store, storageIntegrationProvider, realmContext, configurationFile(), persistenceUnitName(), - secretsGenerator(realmContext)); + secretsGenerator(realmContext), + diagnostics); } private String configurationFile() { diff --git a/extension/persistence/eclipselink/src/main/java/org/apache/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkMetaStoreSessionImpl.java b/extension/persistence/eclipselink/src/main/java/org/apache/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkMetaStoreSessionImpl.java index 7950075c9..6a5670137 100644 --- a/extension/persistence/eclipselink/src/main/java/org/apache/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkMetaStoreSessionImpl.java +++ b/extension/persistence/eclipselink/src/main/java/org/apache/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkMetaStoreSessionImpl.java @@ -37,7 +37,7 @@ import java.util.function.Predicate; import java.util.function.Supplier; import java.util.stream.Collectors; -import org.apache.polaris.core.PolarisCallContext; +import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.PolarisBaseEntity; import org.apache.polaris.core.entity.PolarisChangeTrackingVersions; @@ -81,6 +81,7 @@ public class PolarisEclipseLinkMetaStoreSessionImpl implements PolarisMetaStoreS private final PolarisEclipseLinkStore store; private final PolarisStorageIntegrationProvider storageIntegrationProvider; private final PrincipalSecretsGenerator secretsGenerator; + private final PolarisDiagnostics diagnostics; /** * Create a meta store session against provided realm. Each realm has its own database. @@ -97,7 +98,9 @@ public PolarisEclipseLinkMetaStoreSessionImpl( @Nonnull RealmContext realmContext, @Nullable String confFile, @Nullable String persistenceUnitName, - @Nonnull PrincipalSecretsGenerator secretsGenerator) { + @Nonnull PrincipalSecretsGenerator secretsGenerator, + @Nonnull PolarisDiagnostics diagnostics) { + this.diagnostics = diagnostics; LOGGER.debug( "Creating EclipseLink Meta Store Session for realm {}", realmContext.getRealmIdentifier()); emf = createEntityManagerFactory(realmContext, confFile, persistenceUnitName); @@ -143,9 +146,8 @@ static void clearEntityManagerFactories() { /** {@inheritDoc} */ @Override - public T runInTransaction( - @Nonnull PolarisCallContext callCtx, @Nonnull Supplier transactionCode) { - callCtx.getDiagServices().check(localSession.get() == null, "cannot nest transaction"); + public T runInTransaction(@Nonnull Supplier transactionCode) { + diagnostics.check(localSession.get() == null, "cannot nest transaction"); try (EntityManager session = emf.createEntityManager()) { localSession.set(session); @@ -186,9 +188,8 @@ public T runInTransaction( /** {@inheritDoc} */ @Override - public void runActionInTransaction( - @Nonnull PolarisCallContext callCtx, @Nonnull Runnable transactionCode) { - callCtx.getDiagServices().check(localSession.get() == null, "cannot nest transaction"); + public void runActionInTransaction(@Nonnull Runnable transactionCode) { + diagnostics.check(localSession.get() == null, "cannot nest transaction"); try (EntityManager session = emf.createEntityManager()) { localSession.set(session); @@ -221,43 +222,39 @@ public void runActionInTransaction( /** {@inheritDoc} */ @Override - public T runInReadTransaction( - @Nonnull PolarisCallContext callCtx, @Nonnull Supplier transactionCode) { + public T runInReadTransaction(@Nonnull Supplier transactionCode) { // EclipseLink doesn't support readOnly transaction - return runInTransaction(callCtx, transactionCode); + return runInTransaction(transactionCode); } /** {@inheritDoc} */ @Override - public void runActionInReadTransaction( - @Nonnull PolarisCallContext callCtx, @Nonnull Runnable transactionCode) { + public void runActionInReadTransaction(@Nonnull Runnable transactionCode) { // EclipseLink doesn't support readOnly transaction - runActionInTransaction(callCtx, transactionCode); + runActionInTransaction(transactionCode); } /** * @return new unique entity identifier */ @Override - public long generateNewId(@Nonnull PolarisCallContext callCtx) { + public long generateNewId() { // This function can be called within a transaction or out of transaction. // If called out of transaction, create a new transaction, otherwise run in current transaction return localSession.get() != null ? this.store.getNextSequence(localSession.get()) - : runInReadTransaction(callCtx, () -> generateNewId(callCtx)); + : runInReadTransaction(this::generateNewId); } /** {@inheritDoc} */ @Override - public void writeToEntities( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisBaseEntity entity) { + public void writeToEntities(@Nonnull PolarisBaseEntity entity) { this.store.writeToEntities(localSession.get(), entity); } /** {@inheritDoc} */ @Override public void persistStorageIntegrationIfNeeded( - @Nonnull PolarisCallContext callContext, @Nonnull PolarisBaseEntity entity, @Nullable PolarisStorageIntegration storageIntegration) { // not implemented for eclipselink store @@ -265,40 +262,35 @@ public void persistStorageIntegratio /** {@inheritDoc} */ @Override - public void writeToEntitiesActive( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisBaseEntity entity) { + public void writeToEntitiesActive(@Nonnull PolarisBaseEntity entity) { // write it this.store.writeToEntitiesActive(localSession.get(), entity); } /** {@inheritDoc} */ @Override - public void writeToEntitiesDropped( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisBaseEntity entity) { + public void writeToEntitiesDropped(@Nonnull PolarisBaseEntity entity) { // write it this.store.writeToEntitiesDropped(localSession.get(), entity); } /** {@inheritDoc} */ @Override - public void writeToEntitiesChangeTracking( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisBaseEntity entity) { + public void writeToEntitiesChangeTracking(@Nonnull PolarisBaseEntity entity) { // write it this.store.writeToEntitiesChangeTracking(localSession.get(), entity); } /** {@inheritDoc} */ @Override - public void writeToGrantRecords( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisGrantRecord grantRec) { + public void writeToGrantRecords(@Nonnull PolarisGrantRecord grantRec) { // write it this.store.writeToGrantRecords(localSession.get(), grantRec); } /** {@inheritDoc} */ @Override - public void deleteFromEntities( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisEntityCore entity) { + public void deleteFromEntities(@Nonnull PolarisEntityCore entity) { // delete it this.store.deleteFromEntities(localSession.get(), entity.getCatalogId(), entity.getId()); @@ -306,16 +298,14 @@ public void deleteFromEntities( /** {@inheritDoc} */ @Override - public void deleteFromEntitiesActive( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisEntityCore entity) { + public void deleteFromEntitiesActive(@Nonnull PolarisEntityCore entity) { // delete it this.store.deleteFromEntitiesActive(localSession.get(), new PolarisEntitiesActiveKey(entity)); } /** {@inheritDoc} */ @Override - public void deleteFromEntitiesDropped( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisBaseEntity entity) { + public void deleteFromEntitiesDropped(@Nonnull PolarisBaseEntity entity) { // delete it this.store.deleteFromEntitiesDropped(localSession.get(), entity.getCatalogId(), entity.getId()); } @@ -323,27 +313,23 @@ public void deleteFromEntitiesDropped( /** * {@inheritDoc} * - * @param callCtx * @param entity entity record to delete */ @Override - public void deleteFromEntitiesChangeTracking( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisEntityCore entity) { + public void deleteFromEntitiesChangeTracking(@Nonnull PolarisEntityCore entity) { // delete it this.store.deleteFromEntitiesChangeTracking(localSession.get(), entity); } /** {@inheritDoc} */ @Override - public void deleteFromGrantRecords( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisGrantRecord grantRec) { + public void deleteFromGrantRecords(@Nonnull PolarisGrantRecord grantRec) { this.store.deleteFromGrantRecords(localSession.get(), grantRec); } /** {@inheritDoc} */ @Override public void deleteAllEntityGrantRecords( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisEntityCore entity, @Nonnull List grantsOnGrantee, @Nonnull List grantsOnSecurable) { @@ -352,20 +338,18 @@ public void deleteAllEntityGrantRecords( /** {@inheritDoc} */ @Override - public void deleteAll(@Nonnull PolarisCallContext callCtx) { + public void deleteAll() { this.store.deleteAll(localSession.get()); } /** {@inheritDoc} */ @Override - public @Nullable PolarisBaseEntity lookupEntity( - @Nonnull PolarisCallContext callCtx, long catalogId, long entityId) { + public @Nullable PolarisBaseEntity lookupEntity(long catalogId, long entityId) { return ModelEntity.toEntity(this.store.lookupEntity(localSession.get(), catalogId, entityId)); } @Override - public @Nonnull List lookupEntities( - @Nonnull PolarisCallContext callCtx, List entityIds) { + public @Nonnull List lookupEntities(List entityIds) { return this.store.lookupEntities(localSession.get(), entityIds).stream() .map(ModelEntity::toEntity) .toList(); @@ -373,8 +357,7 @@ public void deleteAll(@Nonnull PolarisCallContext callCtx) { /** {@inheritDoc} */ @Override - public int lookupEntityVersion( - @Nonnull PolarisCallContext callCtx, long catalogId, long entityId) { + public int lookupEntityVersion(long catalogId, long entityId) { ModelEntity model = this.store.lookupEntity(localSession.get(), catalogId, entityId); return model == null ? 0 : model.getEntityVersion(); } @@ -382,7 +365,7 @@ public int lookupEntityVersion( /** {@inheritDoc} */ @Override public @Nonnull List lookupEntityVersions( - @Nonnull PolarisCallContext callCtx, List entityIds) { + List entityIds) { Map idToEntityMap = this.store.lookupEntities(localSession.get(), entityIds).stream() .collect( @@ -405,7 +388,7 @@ public int lookupEntityVersion( @Override @Nullable public PolarisEntityActiveRecord lookupEntityActive( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisEntitiesActiveKey entityActiveKey) { + @Nonnull PolarisEntitiesActiveKey entityActiveKey) { // lookup the active entity slice return ModelEntityActive.toEntityActive( this.store.lookupEntityActive(localSession.get(), entityActiveKey)); @@ -415,34 +398,26 @@ public PolarisEntityActiveRecord lookupEntityActive( @Override @Nonnull public List lookupEntityActiveBatch( - @Nonnull PolarisCallContext callCtx, @Nonnull List entityActiveKeys) { // now build a list to quickly verify that nothing has changed - return entityActiveKeys.stream() - .map(entityActiveKey -> this.lookupEntityActive(callCtx, entityActiveKey)) - .collect(Collectors.toList()); + return entityActiveKeys.stream().map(this::lookupEntityActive).collect(Collectors.toList()); } /** {@inheritDoc} */ @Override public @Nonnull List listActiveEntities( - @Nonnull PolarisCallContext callCtx, - long catalogId, - long parentId, - @Nonnull PolarisEntityType entityType) { - return listActiveEntities(callCtx, catalogId, parentId, entityType, Predicates.alwaysTrue()); + long catalogId, long parentId, @Nonnull PolarisEntityType entityType) { + return listActiveEntities(catalogId, parentId, entityType, Predicates.alwaysTrue()); } @Override public @Nonnull List listActiveEntities( - @Nonnull PolarisCallContext callCtx, long catalogId, long parentId, @Nonnull PolarisEntityType entityType, @Nonnull Predicate entityFilter) { // full range scan under the parent for that type return listActiveEntities( - callCtx, catalogId, parentId, entityType, @@ -460,7 +435,6 @@ public List lookupEntityActiveBatch( @Override public @Nonnull List listActiveEntities( - @Nonnull PolarisCallContext callCtx, long catalogId, long parentId, @Nonnull PolarisEntityType entityType, @@ -481,10 +455,7 @@ public List lookupEntityActiveBatch( /** {@inheritDoc} */ @Override public boolean hasChildren( - @Nonnull PolarisCallContext callContext, - @Nullable PolarisEntityType entityType, - long catalogId, - long parentId) { + @Nullable PolarisEntityType entityType, long catalogId, long parentId) { // check if it has children return this.store.countActiveChildEntities(localSession.get(), catalogId, parentId, entityType) > 0; @@ -492,8 +463,7 @@ public boolean hasChildren( /** {@inheritDoc} */ @Override - public int lookupEntityGrantRecordsVersion( - @Nonnull PolarisCallContext callCtx, long catalogId, long entityId) { + public int lookupEntityGrantRecordsVersion(long catalogId, long entityId) { ModelEntityChangeTracking entity = this.store.lookupEntityChangeTracking(localSession.get(), catalogId, entityId); @@ -504,7 +474,6 @@ public int lookupEntityGrantRecordsVersion( /** {@inheritDoc} */ @Override public @Nullable PolarisGrantRecord lookupGrantRecord( - @Nonnull PolarisCallContext callCtx, long securableCatalogId, long securableId, long granteeCatalogId, @@ -524,7 +493,7 @@ public int lookupEntityGrantRecordsVersion( /** {@inheritDoc} */ @Override public @Nonnull List loadAllGrantRecordsOnSecurable( - @Nonnull PolarisCallContext callCtx, long securableCatalogId, long securableId) { + long securableCatalogId, long securableId) { // now fetch all grants for this securable return this.store .lookupAllGrantRecordsOnSecurable(localSession.get(), securableCatalogId, securableId) @@ -536,7 +505,7 @@ public int lookupEntityGrantRecordsVersion( /** {@inheritDoc} */ @Override public @Nonnull List loadAllGrantRecordsOnGrantee( - @Nonnull PolarisCallContext callCtx, long granteeCatalogId, long granteeId) { + long granteeCatalogId, long granteeId) { // now fetch all grants assigned to this grantee return this.store .lookupGrantRecordsOnGrantee(localSession.get(), granteeCatalogId, granteeId) @@ -547,8 +516,7 @@ public int lookupEntityGrantRecordsVersion( /** {@inheritDoc} */ @Override - public @Nullable PolarisPrincipalSecrets loadPrincipalSecrets( - @Nonnull PolarisCallContext callCtx, @Nonnull String clientId) { + public @Nullable PolarisPrincipalSecrets loadPrincipalSecrets(@Nonnull String clientId) { return ModelPrincipalSecrets.toPrincipalSecrets( this.store.lookupPrincipalSecrets(localSession.get(), clientId)); } @@ -556,7 +524,7 @@ public int lookupEntityGrantRecordsVersion( /** {@inheritDoc} */ @Override public @Nonnull PolarisPrincipalSecrets generateNewPrincipalSecrets( - @Nonnull PolarisCallContext callCtx, @Nonnull String principalName, long principalId) { + @Nonnull String principalName, long principalId) { // ensure principal client id is unique PolarisPrincipalSecrets principalSecrets; ModelPrincipalSecrets lookupPrincipalSecrets; @@ -580,11 +548,7 @@ public int lookupEntityGrantRecordsVersion( /** {@inheritDoc} */ @Override public @Nonnull PolarisPrincipalSecrets rotatePrincipalSecrets( - @Nonnull PolarisCallContext callCtx, - @Nonnull String clientId, - long principalId, - boolean reset, - @Nonnull String oldSecretHash) { + @Nonnull String clientId, long principalId, boolean reset, @Nonnull String oldSecretHash) { // load the existing secrets PolarisPrincipalSecrets principalSecrets = @@ -592,24 +556,20 @@ public int lookupEntityGrantRecordsVersion( this.store.lookupPrincipalSecrets(localSession.get(), clientId)); // should be found - callCtx - .getDiagServices() - .checkNotNull( - principalSecrets, - "cannot_find_secrets", - "client_id={} principalId={}", - clientId, - principalId); + diagnostics.checkNotNull( + principalSecrets, + "cannot_find_secrets", + "client_id={} principalId={}", + clientId, + principalId); // ensure principal id is matching - callCtx - .getDiagServices() - .check( - principalId == principalSecrets.getPrincipalId(), - "principal_id_mismatch", - "expectedId={} id={}", - principalId, - principalSecrets.getPrincipalId()); + diagnostics.check( + principalId == principalSecrets.getPrincipalId(), + "principal_id_mismatch", + "expectedId={} id={}", + principalId, + principalSecrets.getPrincipalId()); // rotate the secrets principalSecrets.rotateSecrets(oldSecretHash); @@ -626,31 +586,26 @@ public int lookupEntityGrantRecordsVersion( /** {@inheritDoc} */ @Override - public void deletePrincipalSecrets( - @Nonnull PolarisCallContext callCtx, @Nonnull String clientId, long principalId) { + public void deletePrincipalSecrets(@Nonnull String clientId, long principalId) { // load the existing secrets ModelPrincipalSecrets principalSecrets = this.store.lookupPrincipalSecrets(localSession.get(), clientId); // should be found - callCtx - .getDiagServices() - .checkNotNull( - principalSecrets, - "cannot_find_secrets", - "client_id={} principalId={}", - clientId, - principalId); + diagnostics.checkNotNull( + principalSecrets, + "cannot_find_secrets", + "client_id={} principalId={}", + clientId, + principalId); // ensure principal id is matching - callCtx - .getDiagServices() - .check( - principalId == principalSecrets.getPrincipalId(), - "principal_id_mismatch", - "expectedId={} id={}", - principalId, - principalSecrets.getPrincipalId()); + diagnostics.check( + principalId == principalSecrets.getPrincipalId(), + "principal_id_mismatch", + "expectedId={} id={}", + principalId, + principalSecrets.getPrincipalId()); // delete these secrets this.store.deletePrincipalSecrets(localSession.get(), clientId); @@ -660,7 +615,6 @@ public void deletePrincipalSecrets( @Override public @Nullable PolarisStorageIntegration createStorageIntegration( - @Nonnull PolarisCallContext callCtx, long catalogId, long entityId, PolarisStorageConfigurationInfo polarisStorageConfigurationInfo) { @@ -672,9 +626,9 @@ PolarisStorageIntegration createStorageIntegration( @Override public @Nullable PolarisStorageIntegration loadPolarisStorageIntegration( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisBaseEntity entity) { + @Nonnull PolarisBaseEntity entity) { PolarisStorageConfigurationInfo storageConfig = - PolarisMetaStoreManagerImpl.readStorageConfiguration(callCtx, entity); + PolarisMetaStoreManagerImpl.readStorageConfiguration(diagnostics, entity); return storageIntegrationProvider.getStorageIntegrationForConfig(storageConfig); } diff --git a/extension/persistence/eclipselink/src/test/java/org/apache/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkMetaStoreManagerTest.java b/extension/persistence/eclipselink/src/test/java/org/apache/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkMetaStoreManagerTest.java index cd9b08590..147f76e22 100644 --- a/extension/persistence/eclipselink/src/test/java/org/apache/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkMetaStoreManagerTest.java +++ b/extension/persistence/eclipselink/src/test/java/org/apache/polaris/extension/persistence/impl/eclipselink/PolarisEclipseLinkMetaStoreManagerTest.java @@ -35,13 +35,14 @@ import java.util.Comparator; import java.util.Objects; import java.util.stream.Stream; -import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.PolarisConfigurationStore; import org.apache.polaris.core.PolarisDefaultDiagServiceImpl; import org.apache.polaris.core.PolarisDiagnostics; +import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.PolarisPrincipalSecrets; import org.apache.polaris.core.persistence.BasePolarisMetaStoreManagerTest; import org.apache.polaris.core.persistence.PolarisMetaStoreManagerImpl; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; import org.apache.polaris.core.persistence.PolarisTestMetaStoreManager; import org.apache.polaris.jpa.models.ModelPrincipalSecrets; import org.junit.jupiter.api.AfterAll; @@ -100,16 +101,18 @@ static void deleteConfFiles() throws IOException { protected PolarisTestMetaStoreManager createPolarisTestMetaStoreManager() { PolarisDiagnostics diagServices = new PolarisDefaultDiagServiceImpl(); PolarisEclipseLinkStore store = new PolarisEclipseLinkStore(diagServices); - PolarisEclipseLinkMetaStoreSessionImpl session = + RealmContext realmContext = () -> "realm"; + PolarisMetaStoreSession session = new PolarisEclipseLinkMetaStoreSessionImpl( - store, Mockito.mock(), () -> "realm", null, "polaris", RANDOM_SECRETS); + store, Mockito.mock(), realmContext, null, "polaris", RANDOM_SECRETS, diagServices); return new PolarisTestMetaStoreManager( - new PolarisMetaStoreManagerImpl(), - new PolarisCallContext( - session, + new PolarisMetaStoreManagerImpl( + realmContext, diagServices, new PolarisConfigurationStore() {}, - timeSource.withZone(ZoneId.systemDefault()))); + timeSource.withZone(ZoneId.systemDefault())), + session, + diagServices); } @ParameterizedTest @@ -123,7 +126,13 @@ void testCreateStoreSession(String confFile, boolean success) { try { var session = new PolarisEclipseLinkMetaStoreSessionImpl( - store, Mockito.mock(), () -> "realm", confFile, "polaris", RANDOM_SECRETS); + store, + Mockito.mock(), + () -> "realm", + confFile, + "polaris", + RANDOM_SECRETS, + diagServices); assertNotNull(session); assertTrue(success); } catch (Exception e) { diff --git a/polaris-core/src/main/java/org/apache/polaris/core/PolarisCallContext.java b/polaris-core/src/main/java/org/apache/polaris/core/PolarisCallContext.java deleted file mode 100644 index abe598714..000000000 --- a/polaris-core/src/main/java/org/apache/polaris/core/PolarisCallContext.java +++ /dev/null @@ -1,84 +0,0 @@ -/* - * 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; - -import jakarta.annotation.Nonnull; -import java.time.Clock; -import java.time.ZoneId; -import org.apache.polaris.core.persistence.PolarisMetaStoreSession; - -/** - * The Call context is allocated each time a new REST request is processed. It contains instances of - * low-level services required to process that request - */ -public class PolarisCallContext { - - // meta store which is used to persist Polaris entity metadata - private final PolarisMetaStoreSession metaStore; - - // diag services - private final PolarisDiagnostics diagServices; - - private final PolarisConfigurationStore configurationStore; - - private final Clock clock; - - public PolarisCallContext( - @Nonnull PolarisMetaStoreSession metaStore, - @Nonnull PolarisDiagnostics diagServices, - @Nonnull PolarisConfigurationStore configurationStore, - @Nonnull Clock clock) { - this.metaStore = metaStore; - this.diagServices = diagServices; - this.configurationStore = configurationStore; - this.clock = clock; - } - - public PolarisCallContext( - @Nonnull PolarisMetaStoreSession metaStore, @Nonnull PolarisDiagnostics diagServices) { - this.metaStore = metaStore; - this.diagServices = diagServices; - this.configurationStore = new PolarisConfigurationStore() {}; - this.clock = Clock.system(ZoneId.systemDefault()); - } - - public static PolarisCallContext copyOf(PolarisCallContext original) { - return new PolarisCallContext( - original.getMetaStore(), - original.getDiagServices(), - original.getConfigurationStore(), - original.getClock()); - } - - public PolarisMetaStoreSession getMetaStore() { - return metaStore; - } - - public PolarisDiagnostics getDiagServices() { - return diagServices; - } - - public PolarisConfigurationStore getConfigurationStore() { - return configurationStore; - } - - public Clock getClock() { - return clock; - } -} diff --git a/polaris-core/src/main/java/org/apache/polaris/core/PolarisConfiguration.java b/polaris-core/src/main/java/org/apache/polaris/core/PolarisConfiguration.java index c57d58da7..c8e2bae90 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/PolarisConfiguration.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/PolarisConfiguration.java @@ -22,7 +22,6 @@ import java.util.List; import java.util.Optional; import org.apache.polaris.core.admin.model.StorageConfigInfo; -import org.apache.polaris.core.context.CallContext; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -101,25 +100,6 @@ public PolarisConfiguration build() { } } - /** - * Returns the value of a `PolarisConfiguration`, or the default if it cannot be loaded. This - * method does not need to be used when a `CallContext` is already available - */ - public static T loadConfig(PolarisConfiguration configuration) { - var callContext = CallContext.getCurrentContext(); - if (callContext == null) { - LOGGER.warn( - String.format( - "Unable to load current call context; using %s = %s", - configuration.key, configuration.defaultValue)); - return configuration.defaultValue; - } - return callContext - .getPolarisCallContext() - .getConfigurationStore() - .getConfiguration(callContext.getPolarisCallContext(), configuration); - } - public static Builder builder() { return new Builder<>(); } diff --git a/polaris-core/src/main/java/org/apache/polaris/core/PolarisConfigurationStore.java b/polaris-core/src/main/java/org/apache/polaris/core/PolarisConfigurationStore.java index a3c902206..369821c08 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/PolarisConfigurationStore.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/PolarisConfigurationStore.java @@ -23,6 +23,7 @@ import jakarta.annotation.Nullable; import java.util.ArrayList; import java.util.List; +import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.CatalogEntity; /** @@ -33,12 +34,12 @@ public interface PolarisConfigurationStore { /** * Retrieve the current value for a configuration key. May be null if not set. * - * @param ctx the current call context + * @param the type of the configuration value + * @param realmContext the realm context to check for overrides; may be null. * @param configName the name of the configuration key to check * @return the current value set for the configuration key or null if not set - * @param the type of the configuration value */ - default @Nullable T getConfiguration(PolarisCallContext ctx, String configName) { + default @Nullable T getConfiguration(@Nullable RealmContext realmContext, String configName) { return null; } @@ -46,16 +47,16 @@ public interface PolarisConfigurationStore { * Retrieve the current value for a configuration key. If not set, return the non-null default * value. * - * @param ctx the current call context + * @param the type of the configuration value + * @param realmContext the realm context to check for overrides; may be null. * @param configName the name of the configuration key to check * @param defaultValue the default value if the configuration key has no value * @return the current value or the supplied default value - * @param the type of the configuration value */ default @Nonnull T getConfiguration( - PolarisCallContext ctx, String configName, @Nonnull T defaultValue) { + @Nullable RealmContext realmContext, String configName, @Nonnull T defaultValue) { Preconditions.checkNotNull(defaultValue, "Cannot pass null as a default value"); - T configValue = getConfiguration(ctx, configName); + T configValue = getConfiguration(realmContext, configName); return configValue != null ? configValue : defaultValue; } @@ -86,13 +87,14 @@ public interface PolarisConfigurationStore { /** * Retrieve the current value for a configuration. * - * @param ctx the current call context + * @param the type of the configuration value + * @param realmContext the realm context to check for overrides; may be null. * @param config the configuration to load * @return the current value set for the configuration key or null if not set - * @param the type of the configuration value */ - default @Nonnull T getConfiguration(PolarisCallContext ctx, PolarisConfiguration config) { - T result = getConfiguration(ctx, config.key, config.defaultValue); + default @Nonnull T getConfiguration( + @Nullable RealmContext realmContext, PolarisConfiguration config) { + T result = getConfiguration(realmContext, config.key, config.defaultValue); return tryCast(config, result); } @@ -100,21 +102,21 @@ public interface PolarisConfigurationStore { * Retrieve the current value for a configuration, overriding with a catalog config if it is * present. * - * @param ctx the current call context + * @param the type of the configuration value + * @param realmContext the realm context to check for overrides; may be null. * @param catalogEntity the catalog to check for an override * @param config the configuration to load * @return the current value set for the configuration key or null if not set - * @param the type of the configuration value */ default @Nonnull T getConfiguration( - PolarisCallContext ctx, + @Nullable RealmContext realmContext, @Nonnull CatalogEntity catalogEntity, PolarisConfiguration config) { if (config.hasCatalogConfig() && catalogEntity.getPropertiesAsMap().containsKey(config.catalogConfig())) { return tryCast(config, catalogEntity.getPropertiesAsMap().get(config.catalogConfig())); } else { - return getConfiguration(ctx, config); + return getConfiguration(realmContext, config); } } } diff --git a/polaris-core/src/main/java/org/apache/polaris/core/auth/PolarisAuthorizer.java b/polaris-core/src/main/java/org/apache/polaris/core/auth/PolarisAuthorizer.java index 31e69b083..adb5cf7b8 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/auth/PolarisAuthorizer.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/auth/PolarisAuthorizer.java @@ -22,6 +22,7 @@ import jakarta.annotation.Nullable; import java.util.List; import java.util.Set; +import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.PolarisBaseEntity; import org.apache.polaris.core.persistence.PolarisResolvedPathWrapper; @@ -29,6 +30,7 @@ public interface PolarisAuthorizer { void authorizeOrThrow( + @Nonnull RealmContext realmContext, @Nonnull AuthenticatedPolarisPrincipal authenticatedPrincipal, @Nonnull Set activatedEntities, @Nonnull PolarisAuthorizableOperation authzOp, @@ -36,6 +38,7 @@ void authorizeOrThrow( @Nullable PolarisResolvedPathWrapper secondary); void authorizeOrThrow( + @Nonnull RealmContext realmContext, @Nonnull AuthenticatedPolarisPrincipal authenticatedPrincipal, @Nonnull Set activatedEntities, @Nonnull PolarisAuthorizableOperation authzOp, diff --git a/polaris-core/src/main/java/org/apache/polaris/core/auth/PolarisAuthorizerImpl.java b/polaris-core/src/main/java/org/apache/polaris/core/auth/PolarisAuthorizerImpl.java index 96f449acb..a8410584f 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/auth/PolarisAuthorizerImpl.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/auth/PolarisAuthorizerImpl.java @@ -100,7 +100,7 @@ import org.apache.iceberg.exceptions.ForbiddenException; import org.apache.polaris.core.PolarisConfiguration; import org.apache.polaris.core.PolarisConfigurationStore; -import org.apache.polaris.core.context.CallContext; +import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.PolarisBaseEntity; import org.apache.polaris.core.entity.PolarisEntityConstants; import org.apache.polaris.core.entity.PolarisEntityCore; @@ -487,12 +487,14 @@ public boolean matchesOrIsSubsumedBy( @Override public void authorizeOrThrow( + @Nonnull RealmContext realmContext, @Nonnull AuthenticatedPolarisPrincipal authenticatedPrincipal, @Nonnull Set activatedEntities, @Nonnull PolarisAuthorizableOperation authzOp, @Nullable PolarisResolvedPathWrapper target, @Nullable PolarisResolvedPathWrapper secondary) { authorizeOrThrow( + realmContext, authenticatedPrincipal, activatedEntities, authzOp, @@ -502,6 +504,7 @@ public void authorizeOrThrow( @Override public void authorizeOrThrow( + @Nonnull RealmContext realmContext, @Nonnull AuthenticatedPolarisPrincipal authenticatedPrincipal, @Nonnull Set activatedEntities, @Nonnull PolarisAuthorizableOperation authzOp, @@ -509,7 +512,7 @@ public void authorizeOrThrow( @Nullable List secondaries) { boolean enforceCredentialRotationRequiredState = featureConfig.getConfiguration( - CallContext.getCurrentContext().getPolarisCallContext(), + realmContext, PolarisConfiguration.ENFORCE_PRINCIPAL_CREDENTIAL_ROTATION_REQUIRED_CHECKING); if (enforceCredentialRotationRequiredState && authenticatedPrincipal diff --git a/polaris-core/src/main/java/org/apache/polaris/core/auth/PolarisGrantManager.java b/polaris-core/src/main/java/org/apache/polaris/core/auth/PolarisGrantManager.java index 35291a2a2..c90b8879a 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/auth/PolarisGrantManager.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/auth/PolarisGrantManager.java @@ -26,12 +26,12 @@ import java.util.List; import java.util.Map; import java.util.stream.Collectors; -import org.apache.polaris.core.PolarisCallContext; 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.entity.PolarisPrivilege; import org.apache.polaris.core.persistence.BaseResult; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; /** Manage grants for Polaris entities. */ public interface PolarisGrantManager { @@ -39,7 +39,7 @@ public interface PolarisGrantManager { * Grant usage on a role to a grantee, for example granting usage on a catalog role to a principal * role or granting a principal role to a principal. * - * @param callCtx call context + * @param session the metastore session * @param catalog if the role is a catalog role, the caller needs to pass-in the catalog entity * which was used to resolve that granted. Else null. * @param role resolved catalog or principal role @@ -49,7 +49,7 @@ public interface PolarisGrantManager { */ @Nonnull PrivilegeResult grantUsageOnRoleToGrantee( - @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisMetaStoreSession session, @Nullable PolarisEntityCore catalog, @Nonnull PolarisEntityCore role, @Nonnull PolarisEntityCore grantee); @@ -58,7 +58,7 @@ PrivilegeResult grantUsageOnRoleToGrantee( * Revoke usage on a role (a catalog or a principal role) from a grantee (e.g. a principal role or * a principal). * - * @param callCtx call context + * @param session the metastore session * @param catalog if the granted is a catalog role, the caller needs to pass-in the catalog entity * which was used to resolve that role. Else null should be passed-in. * @param role a catalog/principal role as resolved by the caller @@ -69,7 +69,7 @@ PrivilegeResult grantUsageOnRoleToGrantee( */ @Nonnull PrivilegeResult revokeUsageOnRoleFromGrantee( - @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisMetaStoreSession session, @Nullable PolarisEntityCore catalog, @Nonnull PolarisEntityCore role, @Nonnull PolarisEntityCore grantee); @@ -77,7 +77,7 @@ PrivilegeResult revokeUsageOnRoleFromGrantee( /** * Grant a privilege on a catalog securable to a grantee. * - * @param callCtx call context + * @param session the metastore session * @param grantee resolved role, the grantee * @param catalogPath path to that entity, cannot be null or empty unless securable is top-level * @param securable securable entity, must have been resolved by the client. Can be the catalog @@ -88,7 +88,7 @@ PrivilegeResult revokeUsageOnRoleFromGrantee( */ @Nonnull PrivilegeResult grantPrivilegeOnSecurableToRole( - @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisMetaStoreSession session, @Nonnull PolarisEntityCore grantee, @Nullable List catalogPath, @Nonnull PolarisEntityCore securable, @@ -97,7 +97,7 @@ PrivilegeResult grantPrivilegeOnSecurableToRole( /** * Revoke a privilege on a catalog securable from a grantee. * - * @param callCtx call context + * @param session the metastore session * @param grantee resolved role, the grantee * @param catalogPath path to that entity, cannot be null or empty unless securable is top-level * @param securable securable entity, must have been resolved by the client. Can be the catalog @@ -109,7 +109,7 @@ PrivilegeResult grantPrivilegeOnSecurableToRole( */ @Nonnull PrivilegeResult revokePrivilegeOnSecurableFromRole( - @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisMetaStoreSession session, @Nonnull PolarisEntityCore grantee, @Nullable List catalogPath, @Nonnull PolarisEntityCore securable, @@ -118,21 +118,21 @@ PrivilegeResult revokePrivilegeOnSecurableFromRole( /** * This method should be used by the Polaris app to cache all grant records on a securable. * - * @param callCtx call context + * @param session the metastore session * @param securable the securable entity * @return the list of grants and the version of the grant records. We will return * ENTITY_NOT_FOUND if the securable cannot be found */ @Nonnull default LoadGrantsResult loadGrantsOnSecurable( - @Nonnull PolarisCallContext callCtx, PolarisBaseEntity securable) { - return loadGrantsOnSecurable(callCtx, securable.getCatalogId(), securable.getId()); + @Nonnull PolarisMetaStoreSession session, PolarisBaseEntity securable) { + return loadGrantsOnSecurable(session, securable.getCatalogId(), securable.getId()); } /** * This method should be used by the Polaris app to cache all grant records on a securable. * - * @param callCtx call context + * @param session the metastore session * @param securableCatalogId id of the catalog this securable belongs to * @param securableId id of the securable * @return the list of grants and the version of the grant records. We will return @@ -140,28 +140,28 @@ default LoadGrantsResult loadGrantsOnSecurable( */ @Nonnull LoadGrantsResult loadGrantsOnSecurable( - @Nonnull PolarisCallContext callCtx, long securableCatalogId, long securableId); + @Nonnull PolarisMetaStoreSession session, long securableCatalogId, long securableId); /** * This method should be used by the Polaris app to load all grants made to a grantee, either a * role or a principal. * - * @param callCtx call context + * @param session the metastore session * @param grantee the grantee entity * @return the list of grants and the version of the grant records. We will return NULL if the * grantee does not exist */ @Nonnull default LoadGrantsResult loadGrantsToGrantee( - @Nonnull PolarisCallContext callCtx, PolarisBaseEntity grantee) { - return loadGrantsToGrantee(callCtx, grantee.getCatalogId(), grantee.getId()); + @Nonnull PolarisMetaStoreSession session, PolarisBaseEntity grantee) { + return loadGrantsToGrantee(session, grantee.getCatalogId(), grantee.getId()); } /** * This method should be used by the Polaris app to load all grants made to a grantee, either a * role or a principal. * - * @param callCtx call context + * @param session the metastore session * @param granteeCatalogId id of the catalog this grantee belongs to * @param granteeId id of the grantee * @return the list of grants and the version of the grant records. We will return NULL if the @@ -169,7 +169,7 @@ default LoadGrantsResult loadGrantsToGrantee( */ @Nonnull LoadGrantsResult loadGrantsToGrantee( - PolarisCallContext callCtx, long granteeCatalogId, long granteeId); + @Nonnull PolarisMetaStoreSession session, long granteeCatalogId, long granteeId); /** Result of a grant/revoke privilege call */ class PrivilegeResult extends BaseResult { diff --git a/polaris-core/src/main/java/org/apache/polaris/core/auth/PolarisSecretsManager.java b/polaris-core/src/main/java/org/apache/polaris/core/auth/PolarisSecretsManager.java index d63368fd9..bd396fe89 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/auth/PolarisSecretsManager.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/auth/PolarisSecretsManager.java @@ -22,27 +22,27 @@ import com.fasterxml.jackson.annotation.JsonProperty; import jakarta.annotation.Nonnull; import jakarta.annotation.Nullable; -import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.entity.PolarisPrincipalSecrets; import org.apache.polaris.core.persistence.BaseResult; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; /** Manages secrets for Polaris principals. */ public interface PolarisSecretsManager { /** * Load the principal secrets given the client_id. * - * @param callCtx call context + * @param session the metastore session * @param clientId principal client id * @return the secrets associated to that principal, including the entity id of the principal */ @Nonnull PrincipalSecretsResult loadPrincipalSecrets( - @Nonnull PolarisCallContext callCtx, @Nonnull String clientId); + @Nonnull PolarisMetaStoreSession session, @Nonnull String clientId); /** * Rotate secrets * - * @param callCtx call context + * @param session the metastore session * @param clientId principal client id * @param principalId id of the principal * @param reset true if the principal's secrets should be disabled and replaced with a one-time @@ -53,7 +53,7 @@ PrincipalSecretsResult loadPrincipalSecrets( */ @Nonnull PrincipalSecretsResult rotatePrincipalSecrets( - @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisMetaStoreSession session, @Nonnull String clientId, long principalId, boolean reset, diff --git a/polaris-core/src/main/java/org/apache/polaris/core/context/CallContext.java b/polaris-core/src/main/java/org/apache/polaris/core/context/CallContext.java deleted file mode 100644 index 71c457720..000000000 --- a/polaris-core/src/main/java/org/apache/polaris/core/context/CallContext.java +++ /dev/null @@ -1,158 +0,0 @@ -/* - * 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.context; - -import jakarta.annotation.Nonnull; -import java.io.IOException; -import java.util.HashMap; -import java.util.Map; -import java.util.stream.Collectors; -import org.apache.iceberg.io.CloseableGroup; -import org.apache.polaris.core.PolarisCallContext; -import org.apache.polaris.core.PolarisDiagnostics; -import org.apache.polaris.core.auth.AuthenticatedPolarisPrincipal; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -/** - * Stores elements associated with an individual REST request such as RealmContext, caller - * identity/role, authn/authz, etc. This class is distinct from RealmContext because implementations - * may need to first independently resolve a RealmContext before resolving the identity/role - * elements of the CallContext that reside exclusively within the resolved Realm. For example, the - * principal/role entities may be defined within a Realm-specific persistence layer, and the - * underlying nature of the persistence layer may differ between different realms. - */ -public interface CallContext extends AutoCloseable { - InheritableThreadLocal CURRENT_CONTEXT = new InheritableThreadLocal<>(); - - // For requests that make use of a Catalog instance, this holds the instance that was - // created, scoped to the current call context. - String REQUEST_PATH_CATALOG_INSTANCE_KEY = "REQUEST_PATH_CATALOG_INSTANCE"; - - // Authenticator filters should populate this field alongside resolving a SecurityContext. - // Value type: AuthenticatedPolarisPrincipal - String AUTHENTICATED_PRINCIPAL = "AUTHENTICATED_PRINCIPAL"; - String CLOSEABLES = "closeables"; - - static CallContext setCurrentContext(CallContext context) { - CURRENT_CONTEXT.set(context); - return context; - } - - static CallContext getCurrentContext() { - return CURRENT_CONTEXT.get(); - } - - static PolarisDiagnostics getDiagnostics() { - return CURRENT_CONTEXT.get().getPolarisCallContext().getDiagServices(); - } - - static AuthenticatedPolarisPrincipal getAuthenticatedPrincipal() { - return (AuthenticatedPolarisPrincipal) - CallContext.getCurrentContext().contextVariables().get(CallContext.AUTHENTICATED_PRINCIPAL); - } - - static void unsetCurrentContext() { - CURRENT_CONTEXT.remove(); - } - - static CallContext of( - final RealmContext realmContext, final PolarisCallContext polarisCallContext) { - Map map = new HashMap<>(); - return new CallContext() { - @Override - public RealmContext getRealmContext() { - return realmContext; - } - - @Override - public PolarisCallContext getPolarisCallContext() { - return polarisCallContext; - } - - @Override - public Map contextVariables() { - return map; - } - }; - } - - /** - * Copy the {@link CallContext}. {@link #contextVariables()} will be copied except for {@link - * #closeables()}. The original {@link #contextVariables()} map is untouched and {@link - * #closeables()} in the original {@link CallContext} should be closed along with the {@link - * CallContext}. - */ - static CallContext copyOf(CallContext base) { - String realmId = base.getRealmContext().getRealmIdentifier(); - RealmContext realmContext = () -> realmId; - PolarisCallContext polarisCallContext = PolarisCallContext.copyOf(base.getPolarisCallContext()); - Map contextVariables = - base.contextVariables().entrySet().stream() - .filter(e -> !e.getKey().equals(CLOSEABLES)) - .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); - return new CallContext() { - @Override - public RealmContext getRealmContext() { - return realmContext; - } - - @Override - public PolarisCallContext getPolarisCallContext() { - return polarisCallContext; - } - - @Override - public Map contextVariables() { - return contextVariables; - } - }; - } - - RealmContext getRealmContext(); - - /** - * @return the inner context used for delegating services - */ - PolarisCallContext getPolarisCallContext(); - - Map contextVariables(); - - default @Nonnull CloseableGroup closeables() { - return (CloseableGroup) - contextVariables().computeIfAbsent(CLOSEABLES, key -> new CloseableGroup()); - } - - @Override - default void close() { - if (CURRENT_CONTEXT.get() == this) { - unsetCurrentContext(); - CloseableGroup closeables = closeables(); - try { - closeables.close(); - } catch (IOException e) { - Logger logger = LoggerFactory.getLogger(CallContext.class); - logger - .atWarn() - .addKeyValue("closeableGroup", closeables) - .log("Unable to close closeable group", e); - } - } - } -} diff --git a/polaris-core/src/main/java/org/apache/polaris/core/context/RealmContext.java b/polaris-core/src/main/java/org/apache/polaris/core/context/RealmContext.java index 9a6f24224..cce27972d 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/context/RealmContext.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/context/RealmContext.java @@ -24,5 +24,11 @@ * prod), and/or account. */ public interface RealmContext { + + static RealmContext copyOf(RealmContext original) { + String realmIdentifier = original.getRealmIdentifier(); + return () -> realmIdentifier; + } + String getRealmIdentifier(); } diff --git a/polaris-core/src/main/java/org/apache/polaris/core/entity/TaskEntity.java b/polaris-core/src/main/java/org/apache/polaris/core/entity/TaskEntity.java index 5a0add8d4..65d2d761a 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/entity/TaskEntity.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/entity/TaskEntity.java @@ -18,8 +18,7 @@ */ package org.apache.polaris.core.entity; -import org.apache.polaris.core.PolarisCallContext; -import org.apache.polaris.core.context.CallContext; +import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.persistence.PolarisObjectMapperUtil; /** @@ -39,18 +38,14 @@ public static TaskEntity of(PolarisBaseEntity polarisEntity) { } } - public T readData(Class klass) { - PolarisCallContext polarisCallContext = CallContext.getCurrentContext().getPolarisCallContext(); + public T readData(PolarisDiagnostics diagnostics, Class klass) { return PolarisObjectMapperUtil.deserialize( - polarisCallContext, getPropertiesAsMap().get(PolarisTaskConstants.TASK_DATA), klass); + diagnostics, getPropertiesAsMap().get(PolarisTaskConstants.TASK_DATA), klass); } - public AsyncTaskType getTaskType() { - PolarisCallContext polarisCallContext = CallContext.getCurrentContext().getPolarisCallContext(); + public AsyncTaskType getTaskType(PolarisDiagnostics diagnostics) { return PolarisObjectMapperUtil.deserialize( - polarisCallContext, - getPropertiesAsMap().get(PolarisTaskConstants.TASK_TYPE), - AsyncTaskType.class); + diagnostics, getPropertiesAsMap().get(PolarisTaskConstants.TASK_TYPE), AsyncTaskType.class); } public static class Builder extends PolarisEntity.BaseBuilder { @@ -65,21 +60,15 @@ public Builder(TaskEntity original) { super(original); } - public Builder withTaskType(AsyncTaskType taskType) { - PolarisCallContext polarisCallContext = - CallContext.getCurrentContext().getPolarisCallContext(); + public Builder withTaskType(PolarisDiagnostics diagnostics, AsyncTaskType taskType) { properties.put( - PolarisTaskConstants.TASK_TYPE, - PolarisObjectMapperUtil.serialize(polarisCallContext, taskType)); + PolarisTaskConstants.TASK_TYPE, PolarisObjectMapperUtil.serialize(diagnostics, taskType)); return this; } - public Builder withData(Object data) { - PolarisCallContext polarisCallContext = - CallContext.getCurrentContext().getPolarisCallContext(); + public Builder withData(PolarisDiagnostics diagnostics, Object data) { properties.put( - PolarisTaskConstants.TASK_DATA, - PolarisObjectMapperUtil.serialize(polarisCallContext, data)); + PolarisTaskConstants.TASK_DATA, PolarisObjectMapperUtil.serialize(diagnostics, data)); return this; } diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/LocalPolarisMetaStoreManagerFactory.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/LocalPolarisMetaStoreManagerFactory.java index 6c9294418..bdea45ed7 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/LocalPolarisMetaStoreManagerFactory.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/LocalPolarisMetaStoreManagerFactory.java @@ -19,15 +19,14 @@ package org.apache.polaris.core.persistence; import jakarta.annotation.Nonnull; +import java.time.Clock; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.function.Supplier; -import org.apache.polaris.core.PolarisCallContext; -import org.apache.polaris.core.PolarisDefaultDiagServiceImpl; +import org.apache.polaris.core.PolarisConfigurationStore; import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.auth.PolarisSecretsManager.PrincipalSecretsResult; -import org.apache.polaris.core.context.CallContext; import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.PolarisEntity; import org.apache.polaris.core.entity.PolarisEntityConstants; @@ -47,22 +46,33 @@ public abstract class LocalPolarisMetaStoreManagerFactory implements MetaStoreManagerFactory { - final Map metaStoreManagerMap = new HashMap<>(); - final Map storageCredentialCacheMap = new HashMap<>(); - final Map entityCacheMap = new HashMap<>(); - final Map backingStoreMap = new HashMap<>(); - final Map> sessionSupplierMap = new HashMap<>(); - protected final PolarisDiagnostics diagServices = new PolarisDefaultDiagServiceImpl(); - private static final Logger LOGGER = LoggerFactory.getLogger(LocalPolarisMetaStoreManagerFactory.class); + private final Map metaStoreManagerMap = new HashMap<>(); + private final Map storageCredentialCacheMap = new HashMap<>(); + private final Map entityCacheMap = new HashMap<>(); + private final Map> sessionSupplierMap = new HashMap<>(); + + private final PolarisConfigurationStore configurationStore; + private final PolarisDiagnostics diagnostics; + private final Clock clock; + private boolean bootstrap; + protected LocalPolarisMetaStoreManagerFactory( + PolarisConfigurationStore configurationStore, PolarisDiagnostics diagnostics, Clock clock) { + this.configurationStore = configurationStore; + this.diagnostics = diagnostics; + this.clock = clock; + } + protected abstract StoreType createBackingStore(@Nonnull PolarisDiagnostics diagnostics); protected abstract PolarisMetaStoreSession createMetaStoreSession( - @Nonnull StoreType store, @Nonnull RealmContext realmContext); + @Nonnull StoreType store, + @Nonnull RealmContext realmContext, + @Nonnull PolarisDiagnostics diagnostics); protected PrincipalSecretsGenerator secretsGenerator(RealmContext realmContext) { if (bootstrap) { @@ -73,13 +83,13 @@ protected PrincipalSecretsGenerator secretsGenerator(RealmContext realmContext) } private void initializeForRealm(RealmContext realmContext) { - final StoreType backingStore = createBackingStore(diagServices); - backingStoreMap.put(realmContext.getRealmIdentifier(), backingStore); + final StoreType backingStore = createBackingStore(diagnostics); sessionSupplierMap.put( realmContext.getRealmIdentifier(), - () -> createMetaStoreSession(backingStore, realmContext)); + () -> createMetaStoreSession(backingStore, realmContext, diagnostics)); - PolarisMetaStoreManager metaStoreManager = new PolarisMetaStoreManagerImpl(); + PolarisMetaStoreManager metaStoreManager = + new PolarisMetaStoreManagerImpl(realmContext, diagnostics, configurationStore, clock); metaStoreManagerMap.put(realmContext.getRealmIdentifier(), metaStoreManager); } @@ -112,11 +122,9 @@ public void purgeRealms(List realms) { PolarisMetaStoreManager metaStoreManager = getOrCreateMetaStoreManager(() -> realm); PolarisMetaStoreSession session = getOrCreateSessionSupplier(() -> realm).get(); - PolarisCallContext callContext = new PolarisCallContext(session, diagServices); - metaStoreManager.purge(callContext); + metaStoreManager.purge(session); storageCredentialCacheMap.remove(realm); - backingStoreMap.remove(realm); sessionSupplierMap.remove(realm); metaStoreManagerMap.remove(realm); } @@ -152,7 +160,8 @@ public synchronized StorageCredentialCache getOrCreateStorageCredentialCache( RealmContext realmContext) { if (!storageCredentialCacheMap.containsKey(realmContext.getRealmIdentifier())) { storageCredentialCacheMap.put( - realmContext.getRealmIdentifier(), new StorageCredentialCache()); + realmContext.getRealmIdentifier(), + new StorageCredentialCache(diagnostics, configurationStore)); } return storageCredentialCacheMap.get(realmContext.getRealmIdentifier()); @@ -162,7 +171,8 @@ public synchronized StorageCredentialCache getOrCreateStorageCredentialCache( public synchronized EntityCache getOrCreateEntityCache(RealmContext realmContext) { if (!entityCacheMap.containsKey(realmContext.getRealmIdentifier())) { PolarisMetaStoreManager metaStoreManager = getOrCreateMetaStoreManager(realmContext); - entityCacheMap.put(realmContext.getRealmIdentifier(), new EntityCache(metaStoreManager)); + entityCacheMap.put( + realmContext.getRealmIdentifier(), new EntityCache(metaStoreManager, diagnostics)); } return entityCacheMap.get(realmContext.getRealmIdentifier()); @@ -175,16 +185,12 @@ public synchronized EntityCache getOrCreateEntityCache(RealmContext realmContext */ private PrincipalSecretsResult bootstrapServiceAndCreatePolarisPrincipalForRealm( RealmContext realmContext, PolarisMetaStoreManager metaStoreManager) { - // While bootstrapping we need to act as a fake privileged context since the real - // CallContext hasn't even been resolved yet. - PolarisCallContext polarisContext = - new PolarisCallContext( - sessionSupplierMap.get(realmContext.getRealmIdentifier()).get(), diagServices); - CallContext.setCurrentContext(CallContext.of(realmContext, polarisContext)); + PolarisMetaStoreSession metaStoreSession = + sessionSupplierMap.get(realmContext.getRealmIdentifier()).get(); PolarisMetaStoreManager.EntityResult preliminaryRootPrincipalLookup = metaStoreManager.readEntityByName( - polarisContext, + metaStoreSession, null, PolarisEntityType.PRINCIPAL, PolarisEntitySubType.NULL_SUBTYPE, @@ -197,11 +203,11 @@ private PrincipalSecretsResult bootstrapServiceAndCreatePolarisPrincipalForRealm throw new IllegalArgumentException(overrideMessage); } - metaStoreManager.bootstrapPolarisService(polarisContext); + metaStoreManager.bootstrapPolarisService(metaStoreSession); PolarisMetaStoreManager.EntityResult rootPrincipalLookup = metaStoreManager.readEntityByName( - polarisContext, + metaStoreSession, null, PolarisEntityType.PRINCIPAL, PolarisEntitySubType.NULL_SUBTYPE, @@ -209,14 +215,14 @@ private PrincipalSecretsResult bootstrapServiceAndCreatePolarisPrincipalForRealm PolarisPrincipalSecrets secrets = metaStoreManager .loadPrincipalSecrets( - polarisContext, + metaStoreSession, PolarisEntity.of(rootPrincipalLookup.getEntity()) .getInternalPropertiesAsMap() .get(PolarisEntityConstants.getClientIdPropertyName())) .getPrincipalSecrets(); PrincipalSecretsResult rotatedSecrets = metaStoreManager.rotatePrincipalSecrets( - polarisContext, + metaStoreSession, secrets.getPrincipalClientId(), secrets.getPrincipalId(), false, @@ -233,14 +239,13 @@ private PrincipalSecretsResult bootstrapServiceAndCreatePolarisPrincipalForRealm */ private void checkPolarisServiceBootstrappedForRealm( RealmContext realmContext, PolarisMetaStoreManager metaStoreManager) { - PolarisCallContext polarisContext = - new PolarisCallContext( - sessionSupplierMap.get(realmContext.getRealmIdentifier()).get(), diagServices); - CallContext.setCurrentContext(CallContext.of(realmContext, polarisContext)); + + PolarisMetaStoreSession metaStoreSession = + sessionSupplierMap.get(realmContext.getRealmIdentifier()).get(); PolarisMetaStoreManager.EntityResult rootPrincipalLookup = metaStoreManager.readEntityByName( - polarisContext, + metaStoreSession, null, PolarisEntityType.PRINCIPAL, PolarisEntitySubType.NULL_SUBTYPE, diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisEntityManager.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisEntityManager.java index 5cc36a841..50abaa5e7 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisEntityManager.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisEntityManager.java @@ -22,7 +22,7 @@ import jakarta.annotation.Nullable; import jakarta.ws.rs.core.SecurityContext; import java.util.List; -import org.apache.polaris.core.context.CallContext; +import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.entity.PolarisEntity; import org.apache.polaris.core.entity.PolarisEntityConstants; import org.apache.polaris.core.entity.PolarisEntitySubType; @@ -40,10 +40,11 @@ * id and name resolution mechanics around PolarisEntities. */ public class PolarisEntityManager { + private final PolarisMetaStoreManager metaStoreManager; private final EntityCache entityCache; - private final StorageCredentialCache credentialCache; + private final PolarisDiagnostics diagnostics; // Lazily instantiated only a single time per entity manager. private ResolvedPolarisEntity implicitResolvedRootContainerEntity = null; @@ -56,18 +57,21 @@ public class PolarisEntityManager { public PolarisEntityManager( @Nonnull PolarisMetaStoreManager metaStoreManager, @Nonnull StorageCredentialCache credentialCache, - @Nonnull EntityCache entityCache) { + @Nonnull EntityCache entityCache, + @Nonnull PolarisDiagnostics diagnostics) { + this.diagnostics = diagnostics; this.metaStoreManager = metaStoreManager; this.credentialCache = credentialCache; this.entityCache = entityCache; } public Resolver prepareResolver( - @Nonnull CallContext callContext, + @Nonnull PolarisMetaStoreSession metaStoreSession, @Nonnull SecurityContext securityContext, @Nullable String referenceCatalogName) { return new Resolver( - callContext.getPolarisCallContext(), + metaStoreSession, + diagnostics, metaStoreManager, securityContext, entityCache, @@ -75,13 +79,14 @@ public Resolver prepareResolver( } public PolarisResolutionManifest prepareResolutionManifest( - @Nonnull CallContext callContext, + @Nonnull PolarisMetaStoreSession metaStoreSession, @Nonnull SecurityContext securityContext, @Nullable String referenceCatalogName) { PolarisResolutionManifest manifest = - new PolarisResolutionManifest(callContext, this, securityContext, referenceCatalogName); + new PolarisResolutionManifest( + metaStoreSession, diagnostics, this, securityContext, referenceCatalogName); manifest.setSimulatedResolvedRootContainerEntity( - getSimulatedResolvedRootContainerEntity(callContext)); + getSimulatedResolvedRootContainerEntity(metaStoreSession)); return manifest; } @@ -90,7 +95,7 @@ public PolarisResolutionManifest prepareResolutionManifest( * parent container of all things in this realm. */ private synchronized ResolvedPolarisEntity getSimulatedResolvedRootContainerEntity( - CallContext callContext) { + PolarisMetaStoreSession metaStoreSession) { if (implicitResolvedRootContainerEntity == null) { // For now, the root container is only implicit and doesn't exist in the entity store, and // only @@ -104,7 +109,7 @@ private synchronized ResolvedPolarisEntity getSimulatedResolvedRootContainerEnti PolarisEntity.of( metaStoreManager .readEntityByName( - callContext.getPolarisCallContext(), + metaStoreSession, null, PolarisEntityType.PRINCIPAL_ROLE, PolarisEntitySubType.NULL_SUBTYPE, diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisEntityResolver.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisEntityResolver.java index 3d795d735..d0ac3afbf 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisEntityResolver.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisEntityResolver.java @@ -24,7 +24,6 @@ import java.util.Iterator; import java.util.List; import java.util.stream.Collectors; -import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.entity.PolarisBaseEntity; import org.apache.polaris.core.entity.PolarisEntitiesActiveKey; @@ -67,7 +66,7 @@ public class PolarisEntityResolver { *

The resolver will ensure that none of the entities which are passed in have been dropped or * were renamed or moved. * - * @param callCtx call context + * @param diagnostics the diagnostics service * @param ms meta store in read mode * @param catalogPath path within the catalog. The first element MUST be a catalog entity. * @param resolvedEntity optional entity to resolve under that catalog path. If a non-null value @@ -77,36 +76,32 @@ public class PolarisEntityResolver { * or a principal can be specified here */ PolarisEntityResolver( - @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisDiagnostics diagnostics, @Nonnull PolarisMetaStoreSession ms, @Nullable List catalogPath, @Nullable PolarisEntityCore resolvedEntity, @Nullable List otherTopLevelEntities) { // cache diagnostics services - this.diagnostics = callCtx.getDiagServices(); + this.diagnostics = diagnostics; // validate path if one was specified if (catalogPath != null) { // cannot be an empty list - callCtx.getDiagServices().check(!catalogPath.isEmpty(), "catalogPath_cannot_be_empty"); + diagnostics.check(!catalogPath.isEmpty(), "catalogPath_cannot_be_empty"); // first in the path should be the catalog - callCtx - .getDiagServices() - .check( - catalogPath.get(0).getTypeCode() == PolarisEntityType.CATALOG.getCode(), - "entity_is_not_catalog", - "entity={}", - this); + diagnostics.check( + catalogPath.get(0).getTypeCode() == PolarisEntityType.CATALOG.getCode(), + "entity_is_not_catalog", + "entity={}", + this); } else if (resolvedEntity != null) { // if an entity is specified without any path, it better be a top-level entity - callCtx - .getDiagServices() - .check( - resolvedEntity.getType().isTopLevel(), - "not_top_level_entity", - "resolvedEntity={}", - resolvedEntity); + diagnostics.check( + resolvedEntity.getType().isTopLevel(), + "not_top_level_entity", + "resolvedEntity={}", + resolvedEntity); } // validate the otherTopLevelCatalogEntities list. Must be top-level catalog entities @@ -114,24 +109,20 @@ public class PolarisEntityResolver { // ensure all entities are top-level for (PolarisEntityCore topLevelCatalogEntityDto : otherTopLevelEntities) { // top-level (catalog or account) and is catalog, catalog path must be specified - callCtx - .getDiagServices() - .check( - topLevelCatalogEntityDto.isTopLevel() - || (topLevelCatalogEntityDto.getType().getParentType() - == PolarisEntityType.CATALOG - && catalogPath != null), - "not_top_level_or_missing_catalog_path", - "entity={} catalogPath={}", - topLevelCatalogEntityDto, - catalogPath); + diagnostics.check( + topLevelCatalogEntityDto.isTopLevel() + || (topLevelCatalogEntityDto.getType().getParentType() == PolarisEntityType.CATALOG + && catalogPath != null), + "not_top_level_or_missing_catalog_path", + "entity={} catalogPath={}", + topLevelCatalogEntityDto, + catalogPath); } } // call the resolution logic this.isSuccess = - this.resolveEntitiesIfNeeded( - callCtx, ms, catalogPath, resolvedEntity, otherTopLevelEntities); + this.resolveEntitiesIfNeeded(ms, catalogPath, resolvedEntity, otherTopLevelEntities); // process result if (!this.isSuccess) { @@ -150,50 +141,50 @@ public class PolarisEntityResolver { /** * Constructor for the resolver, when we only need to resolve a path * - * @param callCtx call context + * @param diagnostics the diagnostics service * @param ms meta store in read mode * @param catalogPath input path, can be null or empty list if the entity is a top-level entity * like a catalog. */ PolarisEntityResolver( - @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisDiagnostics diagnostics, @Nonnull PolarisMetaStoreSession ms, @Nullable List catalogPath) { - this(callCtx, ms, catalogPath, null, null); + this(diagnostics, ms, catalogPath, null, null); } /** * Constructor for the resolver, when we only need to resolve a path * - * @param callCtx call context + * @param diagnostics the diagnostics service * @param ms meta store in read mode * @param catalogPath input path, can be null or empty list if the entity is a top-level entity * like a catalog. * @param resolvedEntityDto resolved entity DTO */ PolarisEntityResolver( - @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisDiagnostics diagnostics, @Nonnull PolarisMetaStoreSession ms, @Nullable List catalogPath, PolarisEntityCore resolvedEntityDto) { - this(callCtx, ms, catalogPath, resolvedEntityDto, null); + this(diagnostics, ms, catalogPath, resolvedEntityDto, null); } /** * Constructor for the resolver, when we only need to resolve a path * - * @param callCtx call context + * @param diagnostics the diagnostics service * @param ms meta store in read mode * @param catalogPath input path, can be null or empty list if the entity is a top-level entity * like a catalog. * @param entity Polaris base entity */ PolarisEntityResolver( - @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisDiagnostics diagnostics, @Nonnull PolarisMetaStoreSession ms, @Nullable List catalogPath, @Nonnull PolarisBaseEntity entity) { - this(callCtx, ms, catalogPath, new PolarisEntityCore(entity), null); + this(diagnostics, ms, catalogPath, new PolarisEntityCore(entity), null); } /** @@ -225,7 +216,6 @@ long getCatalogIdOrNull() { /** * Ensure all specified entities are still active, have not been renamed or re-parented. * - * @param callCtx call context * @param ms meta store in read mode * @param catalogPath path within the catalog. The first element MUST be a catalog. Null or empty * for top-level entities like catalog @@ -238,7 +228,6 @@ long getCatalogIdOrNull() { * @return true if all entities have been resolved successfully */ private boolean resolveEntitiesIfNeeded( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisMetaStoreSession ms, @Nullable List catalogPath, @Nullable PolarisEntityCore resolvedEntity, @@ -283,7 +272,7 @@ private boolean resolveEntitiesIfNeeded( // now lookup all these entities by name Iterator activeRecordIt = - ms.lookupEntityActiveBatch(callCtx, entityActiveKeys).iterator(); + ms.lookupEntityActiveBatch(entityActiveKeys).iterator(); // now validate if there was a change and if yes, re-resolve again for (PolarisEntityCore resolveEntity : toResolve) { 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 69652cb87..d65cca86e 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 @@ -25,7 +25,6 @@ import jakarta.annotation.Nullable; import java.util.List; import java.util.Map; -import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.auth.PolarisGrantManager; import org.apache.polaris.core.auth.PolarisSecretsManager; import org.apache.polaris.core.entity.PolarisBaseEntity; @@ -52,11 +51,11 @@ public interface PolarisMetaStoreManager * Bootstrap the Polaris service, creating the root catalog, root principal, and associated * service admin role. Will fail if the service has already been bootstrapped. * - * @param callCtx call context + * @param session the metastore session * @return the result of the bootstrap attempt */ @Nonnull - BaseResult bootstrapPolarisService(@Nonnull PolarisCallContext callCtx); + BaseResult bootstrapPolarisService(@Nonnull PolarisMetaStoreSession session); /** * Purge all metadata associated with the Polaris service, resetting the metastore to the state it @@ -66,11 +65,11 @@ public interface PolarisMetaStoreManager * *

This will destroy whatever Polaris metadata exists in the metastore * - * @param callCtx call context + * @param session the metastore session * @return always success or unexpected error */ @Nonnull - BaseResult purge(@Nonnull PolarisCallContext callCtx); + BaseResult purge(@Nonnull PolarisMetaStoreSession session); /** the return for an entity lookup call */ class EntityResult extends BaseResult { @@ -153,7 +152,7 @@ public PolarisBaseEntity getEntity() { * catalog like a namespace, a role, a table like entity, or a principal. If the entity is inside * a catalog, the parameter catalogPath must be specified * - * @param callCtx call context + * @param session the metastore session * @param catalogPath path inside a catalog to that entity, rooted by the catalog. If null, the * entity being resolved is a top-level account entity like a catalog. * @param entityType entity type @@ -166,7 +165,7 @@ public PolarisBaseEntity getEntity() { */ @Nonnull PolarisMetaStoreManager.EntityResult readEntityByName( - @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisMetaStoreSession session, @Nullable List catalogPath, @Nonnull PolarisEntityType entityType, @Nonnull PolarisEntitySubType entitySubType, @@ -218,7 +217,7 @@ public List getEntities() { * List all entities of the specified type under the specified catalogPath. If the catalogPath is * null, listed entities will be top-level entities like catalogs. * - * @param callCtx call context + * @param session the metastore session * @param catalogPath path inside a catalog. If null or empty, the entities to list are top-level, * like catalogs * @param entityType entity type @@ -228,7 +227,7 @@ public List getEntities() { */ @Nonnull ListEntitiesResult listEntities( - @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisMetaStoreSession session, @Nullable List catalogPath, @Nonnull PolarisEntityType entityType, @Nonnull PolarisEntitySubType entitySubType); @@ -279,11 +278,11 @@ public Long getId() { * Generate a new unique id that can be used by the Polaris client when it needs to create a new * entity * - * @param callCtx call context + * @param session the metastore session * @return the newly created id, not expected to fail */ @Nonnull - GenerateEntityIdResult generateNewEntityId(@Nonnull PolarisCallContext callCtx); + GenerateEntityIdResult generateNewEntityId(@Nonnull PolarisMetaStoreSession session); /** the return the result of a create-principal method */ class CreatePrincipalResult extends BaseResult { @@ -343,14 +342,14 @@ public PolarisPrincipalSecrets getPrincipalSecrets() { * Create a new principal. This not only creates the new principal entity but also generates a * client_id/secret pair for this new principal. * - * @param callCtx call context + * @param session the metastore session * @param principal the principal entity to create * @return the client_id/secret for the new principal which was created. Will return * ENTITY_ALREADY_EXISTS if the principal already exists */ @Nonnull CreatePrincipalResult createPrincipal( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisBaseEntity principal); + @Nonnull PolarisMetaStoreSession session, @Nonnull PolarisBaseEntity principal); /** the return the result of a create-catalog method */ class CreateCatalogResult extends BaseResult { @@ -412,7 +411,7 @@ public PolarisBaseEntity getCatalogAdminRole() { * role required to admin this catalog. If inline storage integration property is provided, create * a storage integration. * - * @param callCtx call context + * @param session the metastore session * @param catalog the catalog entity to create * @param principalRoles once the catalog has been created, list of principal roles to grant its * catalog_admin role to. If no principal role is specified, we will grant the catalog_admin @@ -421,7 +420,7 @@ public PolarisBaseEntity getCatalogAdminRole() { */ @Nonnull CreateCatalogResult createCatalog( - @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisMetaStoreSession session, @Nonnull PolarisBaseEntity catalog, @Nonnull List principalRoles); @@ -433,7 +432,7 @@ CreateCatalogResult createCatalog( * cannot be resolved, we will return null. And of course if another entity exists with the same * name, we will fail and also return null. * - * @param callCtx call context + * @param session the metastore session * @param catalogPath path inside a catalog. If null, the entity to persist is assumed to be * top-level. * @param entity entity to write @@ -444,7 +443,7 @@ CreateCatalogResult createCatalog( */ @Nonnull EntityResult createEntityIfNotExists( - @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisMetaStoreSession session, @Nullable List catalogPath, @Nonnull PolarisBaseEntity entity); @@ -499,7 +498,7 @@ public List getEntities() { * will be persisted. And of course if any entity conflicts with an existing entity with the same * name, we will fail all entities and also return null. * - * @param callCtx call context + * @param session the metastore session * @param catalogPath path inside a catalog. If null, the entity to persist is assumed to be * top-level. * @param entities batch of entities to write @@ -510,7 +509,7 @@ public List getEntities() { */ @Nonnull EntitiesResult createEntitiesIfNotExist( - @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisMetaStoreSession session, @Nullable List catalogPath, @Nonnull List entities); @@ -519,14 +518,14 @@ EntitiesResult createEntitiesIfNotExist( * has not changed. If this is not the case we will return false. Else we will update both the * internal and visible properties and return true * - * @param callCtx call context + * @param session the metastore session * @param catalogPath path to that entity. Could be null if this entity is top-level * @param entity entity to update, cannot be null * @return the entity we updated or null if the client should retry */ @Nonnull EntityResult updateEntityPropertiesIfNotChanged( - @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisMetaStoreSession session, @Nullable List catalogPath, @Nonnull PolarisBaseEntity entity); @@ -556,23 +555,23 @@ public EntityWithPath( } /** - * This works exactly like {@link #updateEntityPropertiesIfNotChanged(PolarisCallContext, List, - * PolarisBaseEntity)} but allows to operate on multiple entities at once. Just loop through the - * list, calling each entity update and return null if any of those fail. + * This works exactly like {@link #updateEntityPropertiesIfNotChanged(PolarisMetaStoreSession, + * List, PolarisBaseEntity)} but allows to operate on multiple entities at once. Just loop through + * the list, calling each entity update and return null if any of those fail. * - * @param callCtx call context + * @param session the metastore session * @param entities the set of entities to update * @return list of all entities we updated or null if the client should retry because one update * failed */ @Nonnull EntitiesResult updateEntitiesPropertiesIfNotChanged( - @Nonnull PolarisCallContext callCtx, @Nonnull List entities); + @Nonnull PolarisMetaStoreSession session, @Nonnull List entities); /** * Rename an entity, potentially re-parenting it. * - * @param callCtx call context + * @param session the metastore session * @param catalogPath path to that entity. Could be an empty list of the entity is a catalog. * @param entityToRename entity to rename. This entity should have been resolved by the client * @param newCatalogPath if not null, new catalog path @@ -583,7 +582,7 @@ EntitiesResult updateEntitiesPropertiesIfNotChanged( */ @Nonnull EntityResult renameEntity( - @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisMetaStoreSession session, @Nullable List catalogPath, @Nonnull PolarisEntityCore entityToRename, @Nullable List newCatalogPath, @@ -650,7 +649,7 @@ public boolean isEntityUnDroppable() { /** * Drop the specified entity assuming it exists * - * @param callCtx call context + * @param session the metastore session * @param catalogPath path to that entity. Could be an empty list of the entity is a catalog. * @param entityToDrop entity to drop, must have been resolved by the client * @param cleanupProperties if not null, properties that will be persisted with the cleanup task @@ -661,7 +660,7 @@ public boolean isEntityUnDroppable() { */ @Nonnull DropEntityResult dropEntityIfExists( - @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisMetaStoreSession session, @Nullable List catalogPath, @Nonnull PolarisEntityCore entityToDrop, @Nullable Map cleanupProperties, @@ -671,21 +670,22 @@ DropEntityResult dropEntityIfExists( * Load the entity from backend store. Will return NULL if the entity does not exist, i.e. has * been purged. The entity being loaded might have been dropped * - * @param callCtx call context + * @param session the metastore session * @param entityCatalogId id of the catalog for that entity * @param entityId the id of the entity to load */ @Nonnull - EntityResult loadEntity(@Nonnull PolarisCallContext callCtx, long entityCatalogId, long entityId); + EntityResult loadEntity( + @Nonnull PolarisMetaStoreSession session, long entityCatalogId, long entityId); /** * Fetch a list of tasks to be completed. Tasks * - * @param callCtx call context + * @param session the metastore session * @param executorId executor id * @param limit limit * @return list of tasks to be completed */ @Nonnull - EntitiesResult loadTasks(@Nonnull PolarisCallContext callCtx, String executorId, int limit); + EntitiesResult loadTasks(@Nonnull PolarisMetaStoreSession session, String executorId, int limit); } diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisMetaStoreManagerImpl.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisMetaStoreManagerImpl.java index b90c7e51c..3566fc5db 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisMetaStoreManagerImpl.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisMetaStoreManagerImpl.java @@ -25,6 +25,7 @@ import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; import jakarta.annotation.Nonnull; import jakarta.annotation.Nullable; +import java.time.Clock; import java.util.ArrayList; import java.util.EnumMap; import java.util.HashMap; @@ -35,7 +36,9 @@ import java.util.Set; import java.util.function.Function; import java.util.stream.Collectors; -import org.apache.polaris.core.PolarisCallContext; +import org.apache.polaris.core.PolarisConfigurationStore; +import org.apache.polaris.core.PolarisDiagnostics; +import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.AsyncTaskType; import org.apache.polaris.core.entity.PolarisBaseEntity; import org.apache.polaris.core.entity.PolarisChangeTrackingVersions; @@ -55,6 +58,7 @@ import org.apache.polaris.core.storage.PolarisStorageActions; import org.apache.polaris.core.storage.PolarisStorageConfigurationInfo; import org.apache.polaris.core.storage.PolarisStorageIntegration; +import org.jetbrains.annotations.NotNull; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -72,20 +76,33 @@ public class PolarisMetaStoreManagerImpl implements PolarisMetaStoreManager { /** use synchronous drop for entities */ private static final boolean USE_SYNCHRONOUS_DROP = true; + private final RealmContext realmContext; + private final PolarisDiagnostics diagnostics; + private final PolarisConfigurationStore configurationStore; + private final Clock clock; + + public PolarisMetaStoreManagerImpl( + RealmContext realmContext, + PolarisDiagnostics diagnostics, + PolarisConfigurationStore configurationStore, + Clock clock) { + this.realmContext = realmContext; + this.diagnostics = diagnostics; + this.configurationStore = configurationStore; + this.clock = clock; + } + /** * Lookup an entity by its name * - * @param callCtx call context * @param ms meta store * @param entityActiveKey lookup key * @return the entity if it exists, null otherwise */ private @Nullable PolarisBaseEntity lookupEntityByName( - @Nonnull PolarisCallContext callCtx, - @Nonnull PolarisMetaStoreSession ms, - @Nonnull PolarisEntitiesActiveKey entityActiveKey) { + @Nonnull PolarisMetaStoreSession ms, @Nonnull PolarisEntitiesActiveKey entityActiveKey) { // ensure that the entity exists - PolarisEntityActiveRecord entityActiveRecord = ms.lookupEntityActive(callCtx, entityActiveKey); + PolarisEntityActiveRecord entityActiveRecord = ms.lookupEntityActive(entityActiveKey); // if not found, return null if (entityActiveRecord == null) { @@ -94,11 +111,9 @@ public class PolarisMetaStoreManagerImpl implements PolarisMetaStoreManager { // lookup the entity, should be there PolarisBaseEntity entity = - ms.lookupEntity(callCtx, entityActiveRecord.getCatalogId(), entityActiveRecord.getId()); - callCtx - .getDiagServices() - .checkNotNull( - entity, "unexpected_not_found_entity", "entityActiveRecord={}", entityActiveRecord); + ms.lookupEntity(entityActiveRecord.getCatalogId(), entityActiveRecord.getId()); + diagnostics.checkNotNull( + entity, "unexpected_not_found_entity", "entityActiveRecord={}", entityActiveRecord); // return it now return entity; @@ -107,21 +122,19 @@ public class PolarisMetaStoreManagerImpl implements PolarisMetaStoreManager { /** * Write this entity to the meta store. * - * @param callCtx call context * @param ms meta store in read/write mode * @param entity entity to persist * @param writeToActive if true, write it to active */ private void writeEntity( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisMetaStoreSession ms, @Nonnull PolarisBaseEntity entity, boolean writeToActive) { - ms.writeToEntities(callCtx, entity); - ms.writeToEntitiesChangeTracking(callCtx, entity); + ms.writeToEntities(entity); + ms.writeToEntitiesChangeTracking(entity); if (writeToActive) { - ms.writeToEntitiesActive(callCtx, entity); + ms.writeToEntitiesActive(entity); } } @@ -129,53 +142,39 @@ private void writeEntity( * Persist the specified new entity. Persist will write this entity in the ENTITIES, in the * ENTITIES_ACTIVE and finally in the ENTITIES_CHANGE_TRACKING tables * - * @param callCtx call context * @param ms meta store in read/write mode * @param entity entity we need a DPO for */ private void persistNewEntity( - @Nonnull PolarisCallContext callCtx, - @Nonnull PolarisMetaStoreSession ms, - @Nonnull PolarisBaseEntity entity) { + @Nonnull PolarisMetaStoreSession ms, @Nonnull PolarisBaseEntity entity) { // validate the entity type and subtype - callCtx.getDiagServices().checkNotNull(entity, "unexpected_null_entity"); - callCtx - .getDiagServices() - .checkNotNull(entity.getName(), "unexpected_null_name", "entity={}", entity); + diagnostics.checkNotNull(entity, "unexpected_null_entity"); + diagnostics.checkNotNull(entity.getName(), "unexpected_null_name", "entity={}", entity); PolarisEntityType type = PolarisEntityType.fromCode(entity.getTypeCode()); - callCtx.getDiagServices().checkNotNull(type, "unknown_type", "entity={}", entity); + diagnostics.checkNotNull(type, "unknown_type", "entity={}", entity); PolarisEntitySubType subType = PolarisEntitySubType.fromCode(entity.getSubTypeCode()); - callCtx.getDiagServices().checkNotNull(subType, "unexpected_null_subType", "entity={}", entity); - callCtx - .getDiagServices() - .check( - subType.getParentType() == null || subType.getParentType() == type, - "invalid_subtype", - "type={} subType={}", - type, - subType); + diagnostics.checkNotNull(subType, "unexpected_null_subType", "entity={}", entity); + diagnostics.check( + subType.getParentType() == null || subType.getParentType() == type, + "invalid_subtype", + "type={} subType={}", + type, + subType); // if top-level entity, its parent should be the account - callCtx - .getDiagServices() - .check( - !type.isTopLevel() || entity.getParentId() == PolarisEntityConstants.getRootEntityId(), - "top_level_parent_should_be_account", - "entity={}", - entity); + diagnostics.check( + !type.isTopLevel() || entity.getParentId() == PolarisEntityConstants.getRootEntityId(), + "top_level_parent_should_be_account", + "entity={}", + entity); // id should not be null - callCtx - .getDiagServices() - .check( - entity.getId() != 0 || type == PolarisEntityType.ROOT, - "id_not_set", - "entity={}", - entity); + diagnostics.check( + entity.getId() != 0 || type == PolarisEntityType.ROOT, "id_not_set", "entity={}", entity); // creation timestamp must be filled - callCtx.getDiagServices().check(entity.getCreateTimestamp() != 0, "null_create_timestamp"); + diagnostics.check(entity.getCreateTimestamp() != 0, "null_create_timestamp"); // this is the first change entity.setLastUpdateTimestamp(entity.getCreateTimestamp()); @@ -186,7 +185,7 @@ private void persistNewEntity( entity.setToPurgeTimestamp(0); // write it - this.writeEntity(callCtx, ms, entity, true); + this.writeEntity(ms, entity, true); } /** @@ -194,45 +193,34 @@ private void persistNewEntity( * increment the entity version and persist it back to the ENTITIES and ENTITIES_CHANGE_TRACKING * tables * - * @param callCtx call context * @param ms meta store * @param entity the entity which has been changed * @return the entity with its version and lastUpdateTimestamp updated */ private @Nonnull PolarisBaseEntity persistEntityAfterChange( - @Nonnull PolarisCallContext callCtx, - @Nonnull PolarisMetaStoreSession ms, - @Nonnull PolarisBaseEntity entity) { + @Nonnull PolarisMetaStoreSession ms, @Nonnull PolarisBaseEntity entity) { // validate the entity type and subtype - callCtx.getDiagServices().checkNotNull(entity, "unexpected_null_entity"); - callCtx - .getDiagServices() - .checkNotNull(entity.getName(), "unexpected_null_name", "entity={}", entity); + diagnostics.checkNotNull(entity, "unexpected_null_entity"); + diagnostics.checkNotNull(entity.getName(), "unexpected_null_name", "entity={}", entity); PolarisEntityType type = entity.getType(); - callCtx.getDiagServices().checkNotNull(type, "unexpected_null_type", "entity={}", entity); + diagnostics.checkNotNull(type, "unexpected_null_type", "entity={}", entity); PolarisEntitySubType subType = entity.getSubType(); - callCtx.getDiagServices().checkNotNull(subType, "unexpected_null_subType", "entity={}", entity); - callCtx - .getDiagServices() - .check( - subType.getParentType() == null || subType.getParentType() == type, - "invalid_subtype", - "type={} subType={} entity={}", - type, - subType, - entity); + diagnostics.checkNotNull(subType, "unexpected_null_subType", "entity={}", entity); + diagnostics.check( + subType.getParentType() == null || subType.getParentType() == type, + "invalid_subtype", + "type={} subType={} entity={}", + type, + subType, + entity); // entity should not have been dropped - callCtx - .getDiagServices() - .check(entity.getDropTimestamp() == 0, "entity_dropped", "entity={}", entity); + diagnostics.check(entity.getDropTimestamp() == 0, "entity_dropped", "entity={}", entity); // creation timestamp must be filled long createTimestamp = entity.getCreateTimestamp(); - callCtx - .getDiagServices() - .check(createTimestamp != 0, "null_create_timestamp", "entity={}", entity); + diagnostics.check(createTimestamp != 0, "null_create_timestamp", "entity={}", entity); // ensure time is not moving backward... long now = System.currentTimeMillis(); @@ -245,7 +233,7 @@ private void persistNewEntity( entity.setEntityVersion(entity.getEntityVersion() + 1); // persist it to the various slices - this.writeEntity(callCtx, ms, entity, false); + this.writeEntity(ms, entity, false); // return it return entity; @@ -261,24 +249,20 @@ private void persistNewEntity( * - we will fully delete the entity from persistence store * * - * @param callCtx call context * @param ms meta store * @param entity the entity being dropped */ - private void dropEntity( - @Nonnull PolarisCallContext callCtx, - @Nonnull PolarisMetaStoreSession ms, - @Nonnull PolarisBaseEntity entity) { + private void dropEntity(@Nonnull PolarisMetaStoreSession ms, @Nonnull PolarisBaseEntity entity) { // validate the entity type and subtype - callCtx.getDiagServices().checkNotNull(entity, "unexpected_null_dpo"); - callCtx.getDiagServices().checkNotNull(entity.getName(), "unexpected_null_name"); + diagnostics.checkNotNull(entity, "unexpected_null_dpo"); + diagnostics.checkNotNull(entity.getName(), "unexpected_null_name"); // creation timestamp must be filled - callCtx.getDiagServices().check(entity.getDropTimestamp() == 0, "already_dropped"); + diagnostics.check(entity.getDropTimestamp() == 0, "already_dropped"); // delete it from active slice - ms.deleteFromEntitiesActive(callCtx, entity); + ms.deleteFromEntitiesActive(entity); // for now drop all entities synchronously if (USE_SYNCHRONOUS_DROP) { @@ -287,11 +271,11 @@ private void dropEntity( // delete ALL grant records to (if the entity is a grantee) and from that entity final List grantsOnGrantee = (entity.getType().isGrantee()) - ? ms.loadAllGrantRecordsOnGrantee(callCtx, entity.getCatalogId(), entity.getId()) + ? ms.loadAllGrantRecordsOnGrantee(entity.getCatalogId(), entity.getId()) : List.of(); final List grantsOnSecurable = - ms.loadAllGrantRecordsOnSecurable(callCtx, entity.getCatalogId(), entity.getId()); - ms.deleteAllEntityGrantRecords(callCtx, entity, grantsOnGrantee, grantsOnSecurable); + ms.loadAllGrantRecordsOnSecurable(entity.getCatalogId(), entity.getId()); + ms.deleteAllEntityGrantRecords(entity, grantsOnGrantee, grantsOnSecurable); // Now determine the set of entities on the other side of the grants we just removed. Grants // from/to these entities has been removed, hence we need to update the grant version of @@ -307,29 +291,27 @@ private void dropEntity( new PolarisEntityId(gr.getGranteeCatalogId(), gr.getGranteeId()))); // Bump up the grant version of these entities - List entities = - ms.lookupEntities(callCtx, new ArrayList<>(entityIdsGrantChanged)); + List entities = ms.lookupEntities(new ArrayList<>(entityIdsGrantChanged)); for (PolarisBaseEntity entityGrantChanged : entities) { entityGrantChanged.setGrantRecordsVersion(entityGrantChanged.getGrantRecordsVersion() + 1); - ms.writeToEntities(callCtx, entityGrantChanged); - ms.writeToEntitiesChangeTracking(callCtx, entityGrantChanged); + ms.writeToEntities(entityGrantChanged); + ms.writeToEntitiesChangeTracking(entityGrantChanged); } // remove the entity being dropped now - ms.deleteFromEntities(callCtx, entity); - ms.deleteFromEntitiesChangeTracking(callCtx, entity); + ms.deleteFromEntities(entity); + ms.deleteFromEntitiesChangeTracking(entity); // if it is a principal, we also need to drop the secrets if (entity.getType() == PolarisEntityType.PRINCIPAL) { // get internal properties - Map properties = - this.deserializeProperties(callCtx, entity.getInternalProperties()); + Map properties = this.deserializeProperties(entity.getInternalProperties()); // get client_id String clientId = properties.get(PolarisEntityConstants.getClientIdPropertyName()); // delete it from the secret slice - ms.deletePrincipalSecrets(callCtx, clientId, entity.getId()); + ms.deletePrincipalSecrets(clientId, entity.getId()); } } else { @@ -345,9 +327,9 @@ private void dropEntity( entity.setEntityVersion(entity.getEntityVersion() + 1); // write to the dropped slice and to purge slice - ms.writeToEntities(callCtx, entity); - ms.writeToEntitiesDropped(callCtx, entity); - ms.writeToEntitiesChangeTracking(callCtx, entity); + ms.writeToEntities(entity); + ms.writeToEntitiesDropped(entity); + ms.writeToEntitiesChangeTracking(entity); } } @@ -355,7 +337,6 @@ private void dropEntity( * Create and persist a new grant record. This will at the same time invalidate the grant records * of the grantee and the securable if the grantee is a catalog role * - * @param callCtx call context * @param ms meta store in read/write mode * @param securable securable * @param grantee grantee, either a catalog role, a principal role or a principal @@ -363,21 +344,19 @@ private void dropEntity( * @return new grant record which was created and persisted */ private @Nonnull PolarisGrantRecord persistNewGrantRecord( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisMetaStoreSession ms, @Nonnull PolarisEntityCore securable, @Nonnull PolarisEntityCore grantee, @Nonnull PolarisPrivilege priv) { // validate non null arguments - callCtx.getDiagServices().checkNotNull(securable, "unexpected_null_securable"); - callCtx.getDiagServices().checkNotNull(grantee, "unexpected_null_grantee"); - callCtx.getDiagServices().checkNotNull(priv, "unexpected_null_priv"); + diagnostics.checkNotNull(securable, "unexpected_null_securable"); + diagnostics.checkNotNull(grantee, "unexpected_null_grantee"); + diagnostics.checkNotNull(priv, "unexpected_null_priv"); // ensure that this entity is indeed a grantee like entity - callCtx - .getDiagServices() - .check(grantee.getType().isGrantee(), "entity_must_be_grantee", "entity={}", grantee); + diagnostics.check( + grantee.getType().isGrantee(), "entity_must_be_grantee", "entity={}", grantee); // create new grant record PolarisGrantRecord grantRecord = @@ -389,31 +368,26 @@ private void dropEntity( priv.getCode()); // persist the new grant - ms.writeToGrantRecords(callCtx, grantRecord); + ms.writeToGrantRecords(grantRecord); // load the grantee (either a catalog/principal role or a principal) and increment its grants // version - PolarisBaseEntity granteeEntity = - ms.lookupEntity(callCtx, grantee.getCatalogId(), grantee.getId()); - callCtx - .getDiagServices() - .checkNotNull(granteeEntity, "grantee_not_found", "grantee={}", grantee); + PolarisBaseEntity granteeEntity = ms.lookupEntity(grantee.getCatalogId(), grantee.getId()); + diagnostics.checkNotNull(granteeEntity, "grantee_not_found", "grantee={}", grantee); // grants have changed, we need to bump-up the grants version granteeEntity.setGrantRecordsVersion(granteeEntity.getGrantRecordsVersion() + 1); - this.writeEntity(callCtx, ms, granteeEntity, false); + this.writeEntity(ms, granteeEntity, false); // 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()); - callCtx - .getDiagServices() - .checkNotNull(securableEntity, "securable_not_found", "securable={}", securable); + ms.lookupEntity(securable.getCatalogId(), securable.getId()); + diagnostics.checkNotNull(securableEntity, "securable_not_found", "securable={}", securable); // grants have changed, we need to bump-up the grants version securableEntity.setGrantRecordsVersion(securableEntity.getGrantRecordsVersion() + 1); - this.writeEntity(callCtx, ms, securableEntity, false); + this.writeEntity(ms, securableEntity, false); // done, return the new grant record return grantRecord; @@ -423,84 +397,70 @@ private void dropEntity( * Delete the specified grant record from the GRANT_RECORDS table. This will at the same time * invalidate the grant records of the grantee and the securable if the grantee is a role * - * @param callCtx call context * @param ms meta store * @param securable the securable entity * @param grantee the grantee entity * @param grantRecord the grant record to remove, which was read in the same transaction */ private void revokeGrantRecord( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisMetaStoreSession ms, @Nonnull PolarisEntityCore securable, @Nonnull PolarisEntityCore grantee, @Nonnull PolarisGrantRecord grantRecord) { // validate securable - callCtx - .getDiagServices() - .check( - securable.getCatalogId() == grantRecord.getSecurableCatalogId() - && securable.getId() == grantRecord.getSecurableId(), - "securable_mismatch", - "securable={} grantRec={}", - securable, - grantRecord); + diagnostics.check( + securable.getCatalogId() == grantRecord.getSecurableCatalogId() + && securable.getId() == grantRecord.getSecurableId(), + "securable_mismatch", + "securable={} grantRec={}", + securable, + grantRecord); // validate grantee - callCtx - .getDiagServices() - .check( - grantee.getCatalogId() == grantRecord.getGranteeCatalogId() - && grantee.getId() == grantRecord.getGranteeId(), - "grantee_mismatch", - "grantee={} grantRec={}", - grantee, - grantRecord); + diagnostics.check( + grantee.getCatalogId() == grantRecord.getGranteeCatalogId() + && grantee.getId() == grantRecord.getGranteeId(), + "grantee_mismatch", + "grantee={} grantRec={}", + grantee, + grantRecord); // ensure the grantee is really a grantee - callCtx - .getDiagServices() - .check(grantee.getType().isGrantee(), "not_a_grantee", "grantee={}", grantee); + diagnostics.check(grantee.getType().isGrantee(), "not_a_grantee", "grantee={}", grantee); // remove that grant - ms.deleteFromGrantRecords(callCtx, grantRecord); + ms.deleteFromGrantRecords(grantRecord); // load the grantee and increment its grants version - PolarisBaseEntity refreshGrantee = - ms.lookupEntity(callCtx, grantee.getCatalogId(), grantee.getId()); - callCtx - .getDiagServices() - .checkNotNull( - refreshGrantee, "missing_grantee", "grantRecord={} grantee={}", grantRecord, grantee); + PolarisBaseEntity refreshGrantee = ms.lookupEntity(grantee.getCatalogId(), grantee.getId()); + diagnostics.checkNotNull( + refreshGrantee, "missing_grantee", "grantRecord={} grantee={}", grantRecord, grantee); // grants have changed, we need to bump-up the grants version refreshGrantee.setGrantRecordsVersion(refreshGrantee.getGrantRecordsVersion() + 1); - this.writeEntity(callCtx, ms, refreshGrantee, false); + this.writeEntity(ms, refreshGrantee, false); // 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()); - callCtx - .getDiagServices() - .checkNotNull( - refreshSecurable, - "missing_securable", - "grantRecord={} securable={}", - grantRecord, - securable); + ms.lookupEntity(securable.getCatalogId(), securable.getId()); + diagnostics.checkNotNull( + refreshSecurable, + "missing_securable", + "grantRecord={} securable={}", + grantRecord, + securable); // grants have changed, we need to bump-up the grants version refreshSecurable.setGrantRecordsVersion(refreshSecurable.getGrantRecordsVersion() + 1); - this.writeEntity(callCtx, ms, refreshSecurable, false); + this.writeEntity(ms, refreshSecurable, false); } /** * Create a new catalog. This not only creates the new catalog entity but also the initial admin * role required to admin this catalog. * - * @param callCtx call context * @param ms meta store in read/write mode * @param catalog the catalog entity to create * @param integration the storage integration that should be attached to the catalog. If null, do @@ -511,29 +471,25 @@ private void revokeGrantRecord( * @return the catalog we just created and its associated admin catalog role or error if we failed * to */ - private @Nonnull CreateCatalogResult createCatalog( - @Nonnull PolarisCallContext callCtx, + private @Nonnull CreateCatalogResult doCreateCatalog( @Nonnull PolarisMetaStoreSession ms, @Nonnull PolarisBaseEntity catalog, @Nullable PolarisStorageIntegration integration, @Nonnull List principalRoles) { // validate input - callCtx.getDiagServices().checkNotNull(catalog, "unexpected_null_catalog"); + diagnostics.checkNotNull(catalog, "unexpected_null_catalog"); // check if that catalog has already been created - PolarisBaseEntity refreshCatalog = - ms.lookupEntity(callCtx, catalog.getCatalogId(), catalog.getId()); + PolarisBaseEntity refreshCatalog = ms.lookupEntity(catalog.getCatalogId(), catalog.getId()); // if found, probably a retry, simply return the previously created catalog if (refreshCatalog != null) { // if found, ensure it is indeed a catalog - callCtx - .getDiagServices() - .check( - refreshCatalog.getTypeCode() == PolarisEntityType.CATALOG.getCode(), - "not_a_catalog", - "catalog={}", - catalog); + diagnostics.check( + refreshCatalog.getTypeCode() == PolarisEntityType.CATALOG.getCode(), + "not_a_catalog", + "catalog={}", + catalog); // lookup catalog admin role, should exist PolarisEntitiesActiveKey adminRoleKey = @@ -542,13 +498,11 @@ private void revokeGrantRecord( refreshCatalog.getId(), PolarisEntityType.CATALOG_ROLE.getCode(), PolarisEntityConstants.getNameOfCatalogAdminRole()); - PolarisBaseEntity catalogAdminRole = this.lookupEntityByName(callCtx, ms, adminRoleKey); + PolarisBaseEntity catalogAdminRole = this.lookupEntityByName(ms, adminRoleKey); // if found, ensure not null - callCtx - .getDiagServices() - .checkNotNull( - catalogAdminRole, "catalog_admin_role_not_found", "catalog={}", refreshCatalog); + diagnostics.checkNotNull( + catalogAdminRole, "catalog_admin_role_not_found", "catalog={}", refreshCatalog); // done, return the existing catalog return new CreateCatalogResult(refreshCatalog, catalogAdminRole); @@ -561,20 +515,20 @@ private void revokeGrantRecord( PolarisEntityConstants.getRootEntityId(), PolarisEntityType.CATALOG.getCode(), catalog.getName()); - PolarisEntityActiveRecord otherCatalogRecord = ms.lookupEntityActive(callCtx, catalogNameKey); + PolarisEntityActiveRecord otherCatalogRecord = ms.lookupEntityActive(catalogNameKey); // if it exists, this is an error, the client should retry if (otherCatalogRecord != null) { return new CreateCatalogResult(BaseResult.ReturnStatus.ENTITY_ALREADY_EXISTS, null); } - ms.persistStorageIntegrationIfNeeded(callCtx, catalog, integration); + ms.persistStorageIntegrationIfNeeded(catalog, integration); // now create and persist new catalog entity - this.persistNewEntity(callCtx, ms, catalog); + this.persistNewEntity(ms, catalog); // create the catalog admin role for this new catalog - long adminRoleId = ms.generateNewId(callCtx); + long adminRoleId = ms.generateNewId(); PolarisBaseEntity adminRole = new PolarisBaseEntity( catalog.getId(), @@ -583,16 +537,14 @@ private void revokeGrantRecord( PolarisEntitySubType.NULL_SUBTYPE, catalog.getId(), PolarisEntityConstants.getNameOfCatalogAdminRole()); - this.persistNewEntity(callCtx, ms, adminRole); + this.persistNewEntity(ms, adminRole); // grant the catalog admin role access-management on the catalog - this.persistNewGrantRecord( - callCtx, ms, catalog, adminRole, PolarisPrivilege.CATALOG_MANAGE_ACCESS); + this.persistNewGrantRecord(ms, catalog, adminRole, PolarisPrivilege.CATALOG_MANAGE_ACCESS); // grant the catalog admin role metadata-management on the catalog; this one // is revocable - this.persistNewGrantRecord( - callCtx, ms, catalog, adminRole, PolarisPrivilege.CATALOG_MANAGE_METADATA); + this.persistNewGrantRecord(ms, catalog, adminRole, PolarisPrivilege.CATALOG_MANAGE_METADATA); // immediately assign its catalog_admin role if (principalRoles.isEmpty()) { @@ -603,27 +555,24 @@ private void revokeGrantRecord( PolarisEntityConstants.getRootEntityId(), PolarisEntityType.PRINCIPAL_ROLE.getCode(), PolarisEntityConstants.getNameOfPrincipalServiceAdminRole()); - PolarisBaseEntity serviceAdminRole = - this.lookupEntityByName(callCtx, ms, serviceAdminRoleKey); - callCtx.getDiagServices().checkNotNull(serviceAdminRole, "missing_service_admin_role"); + PolarisBaseEntity serviceAdminRole = this.lookupEntityByName(ms, serviceAdminRoleKey); + diagnostics.checkNotNull(serviceAdminRole, "missing_service_admin_role"); this.persistNewGrantRecord( - callCtx, ms, adminRole, serviceAdminRole, PolarisPrivilege.CATALOG_ROLE_USAGE); + ms, adminRole, serviceAdminRole, PolarisPrivilege.CATALOG_ROLE_USAGE); } else { // grant to each principal role usage on its catalog_admin role for (PolarisEntityCore principalRole : principalRoles) { // validate not null and really a principal role - callCtx.getDiagServices().checkNotNull(principalRole, "null principal role"); - callCtx - .getDiagServices() - .check( - principalRole.getTypeCode() == PolarisEntityType.PRINCIPAL_ROLE.getCode(), - "not_principal_role", - "type={}", - principalRole.getType()); + diagnostics.checkNotNull(principalRole, "null principal role"); + diagnostics.check( + principalRole.getTypeCode() == PolarisEntityType.PRINCIPAL_ROLE.getCode(), + "not_principal_role", + "type={}", + principalRole.getType()); // grant usage on that catalog admin role to this principal this.persistNewGrantRecord( - callCtx, ms, adminRole, principalRole, PolarisPrivilege.CATALOG_ROLE_USAGE); + ms, adminRole, principalRole, PolarisPrivilege.CATALOG_ROLE_USAGE); } } @@ -634,11 +583,9 @@ private void revokeGrantRecord( /** * Bootstrap Polaris catalog service * - * @param callCtx call context * @param ms meta store in read/write mode */ - private void bootstrapPolarisService( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisMetaStoreSession ms) { + private void doBootstrapPolarisService(@Nonnull PolarisMetaStoreSession ms) { // Create a root container entity that can represent the securable for any top-level grants. PolarisBaseEntity rootContainer = @@ -649,12 +596,12 @@ private void bootstrapPolarisService( PolarisEntitySubType.NULL_SUBTYPE, PolarisEntityConstants.getRootEntityId(), PolarisEntityConstants.getRootContainerName()); - this.persistNewEntity(callCtx, ms, rootContainer); + this.persistNewEntity(ms, rootContainer); // Now bootstrap the service by creating the root principal and the service_admin principal // role. The principal role will be granted to that root principal and the root catalog admin // of the root catalog will be granted to that principal role. - long rootPrincipalId = ms.generateNewId(callCtx); + long rootPrincipalId = ms.generateNewId(); PolarisBaseEntity rootPrincipal = new PolarisBaseEntity( PolarisEntityConstants.getNullId(), @@ -665,10 +612,10 @@ private void bootstrapPolarisService( PolarisEntityConstants.getRootPrincipalName()); // create this principal - this.createPrincipal(callCtx, ms, rootPrincipal); + this.doCreatePrincipal(ms, rootPrincipal); // now create the account admin principal role - long serviceAdminPrincipalRoleId = ms.generateNewId(callCtx); + long serviceAdminPrincipalRoleId = ms.generateNewId(); PolarisBaseEntity serviceAdminPrincipalRole = new PolarisBaseEntity( PolarisEntityConstants.getNullId(), @@ -677,46 +624,34 @@ private void bootstrapPolarisService( PolarisEntitySubType.NULL_SUBTYPE, PolarisEntityConstants.getRootEntityId(), PolarisEntityConstants.getNameOfPrincipalServiceAdminRole()); - this.persistNewEntity(callCtx, ms, serviceAdminPrincipalRole); + this.persistNewEntity(ms, serviceAdminPrincipalRole); // we also need to grant usage on the account-admin principal to the principal this.persistNewGrantRecord( - callCtx, - ms, - serviceAdminPrincipalRole, - rootPrincipal, - PolarisPrivilege.PRINCIPAL_ROLE_USAGE); + ms, serviceAdminPrincipalRole, rootPrincipal, PolarisPrivilege.PRINCIPAL_ROLE_USAGE); // grant SERVICE_MANAGE_ACCESS on the rootContainer to the serviceAdminPrincipalRole this.persistNewGrantRecord( - callCtx, - ms, - rootContainer, - serviceAdminPrincipalRole, - PolarisPrivilege.SERVICE_MANAGE_ACCESS); + ms, rootContainer, serviceAdminPrincipalRole, PolarisPrivilege.SERVICE_MANAGE_ACCESS); } /** {@inheritDoc} */ @Override - public @Nonnull BaseResult bootstrapPolarisService(@Nonnull PolarisCallContext callCtx) { - // get meta store we should be using - PolarisMetaStoreSession ms = callCtx.getMetaStore(); - + public @Nonnull BaseResult bootstrapPolarisService(@NotNull PolarisMetaStoreSession session) { // run operation in a read/write transaction - ms.runActionInTransaction(callCtx, () -> this.bootstrapPolarisService(callCtx, ms)); + session.runActionInTransaction(() -> this.doBootstrapPolarisService(session)); // all good return new BaseResult(BaseResult.ReturnStatus.SUCCESS); } @Override - public @Nonnull BaseResult purge(@Nonnull PolarisCallContext callCtx) { + public @Nonnull BaseResult purge(@NotNull PolarisMetaStoreSession session) { // get meta store we should be using - PolarisMetaStoreSession ms = callCtx.getMetaStore(); // run operation in a read/write transaction LOGGER.warn("Deleting all metadata in the metastore..."); - ms.runActionInTransaction(callCtx, () -> ms.deleteAll(callCtx)); + session.runActionInTransaction(session::deleteAll); LOGGER.warn("Finished deleting all metadata in the metastore"); // all good @@ -724,18 +659,17 @@ private void bootstrapPolarisService( } /** - * See {@link #readEntityByName(PolarisCallContext, List, PolarisEntityType, PolarisEntitySubType, - * String)} + * See {@link #readEntityByName(PolarisMetaStoreSession, List, PolarisEntityType, + * PolarisEntitySubType, String)} */ - private @Nonnull PolarisMetaStoreManager.EntityResult readEntityByName( - @Nonnull PolarisCallContext callCtx, + private @Nonnull PolarisMetaStoreManager.EntityResult doReadEntityByName( @Nonnull PolarisMetaStoreSession ms, @Nullable List catalogPath, @Nonnull PolarisEntityType entityType, @Nonnull PolarisEntitySubType entitySubType, @Nonnull String name) { // first resolve again the catalogPath to that entity - PolarisEntityResolver resolver = new PolarisEntityResolver(callCtx, ms, catalogPath); + PolarisEntityResolver resolver = new PolarisEntityResolver(diagnostics, ms, catalogPath); // return if we failed to resolve if (resolver.isFailure()) { @@ -746,7 +680,7 @@ private void bootstrapPolarisService( PolarisEntitiesActiveKey entityActiveKey = new PolarisEntitiesActiveKey( resolver.getCatalogIdOrNull(), resolver.getParentId(), entityType.getCode(), name); - PolarisBaseEntity entity = this.lookupEntityByName(callCtx, ms, entityActiveKey); + PolarisBaseEntity entity = this.lookupEntityByName(ms, entityActiveKey); // if found, check if subType really matches if (entity != null @@ -764,30 +698,29 @@ private void bootstrapPolarisService( /** {@inheritDoc} */ @Override public @Nonnull PolarisMetaStoreManager.EntityResult readEntityByName( - @Nonnull PolarisCallContext callCtx, + PolarisMetaStoreSession ms, @Nullable List catalogPath, @Nonnull PolarisEntityType entityType, @Nonnull PolarisEntitySubType entitySubType, @Nonnull String name) { // get meta store we should be using - PolarisMetaStoreSession ms = callCtx.getMetaStore(); // run operation in a read/write transaction return ms.runInReadTransaction( - callCtx, () -> readEntityByName(callCtx, ms, catalogPath, entityType, entitySubType, name)); + () -> doReadEntityByName(ms, catalogPath, entityType, entitySubType, name)); } /** - * See {@link #listEntities(PolarisCallContext, List, PolarisEntityType, PolarisEntitySubType)} + * See {@link #listEntities(PolarisMetaStoreSession, List, PolarisEntityType, + * PolarisEntitySubType)} */ - private @Nonnull ListEntitiesResult listEntities( - @Nonnull PolarisCallContext callCtx, + private @Nonnull ListEntitiesResult doListEntities( @Nonnull PolarisMetaStoreSession ms, @Nullable List catalogPath, @Nonnull PolarisEntityType entityType, @Nonnull PolarisEntitySubType entitySubType) { // first resolve again the catalogPath to that entity - PolarisEntityResolver resolver = new PolarisEntityResolver(callCtx, ms, catalogPath); + PolarisEntityResolver resolver = new PolarisEntityResolver(diagnostics, ms, catalogPath); // return if we failed to resolve if (resolver.isFailure()) { @@ -796,8 +729,7 @@ private void bootstrapPolarisService( // return list of active entities List toreturnList = - ms.listActiveEntities( - callCtx, resolver.getCatalogIdOrNull(), resolver.getParentId(), entityType); + ms.listActiveEntities(resolver.getCatalogIdOrNull(), resolver.getParentId(), entityType); // prune the returned list with only entities matching the entity subtype if (entitySubType != PolarisEntitySubType.ANY_SUBTYPE) { @@ -814,25 +746,18 @@ private void bootstrapPolarisService( /** {@inheritDoc} */ @Override public @Nonnull ListEntitiesResult listEntities( - @Nonnull PolarisCallContext callCtx, + PolarisMetaStoreSession ms, @Nullable List catalogPath, @Nonnull PolarisEntityType entityType, @Nonnull PolarisEntitySubType entitySubType) { - // get meta store we should be using - PolarisMetaStoreSession ms = callCtx.getMetaStore(); - - // run operation in a read transaction return ms.runInReadTransaction( - callCtx, () -> listEntities(callCtx, ms, catalogPath, entityType, entitySubType)); + () -> doListEntities(ms, catalogPath, entityType, entitySubType)); } /** {@inheritDoc} */ @Override - public @Nonnull GenerateEntityIdResult generateNewEntityId(@Nonnull PolarisCallContext callCtx) { - // get meta store we should be using - PolarisMetaStoreSession ms = callCtx.getMetaStore(); - - return new GenerateEntityIdResult(ms.generateNewId(callCtx)); + public @Nonnull GenerateEntityIdResult generateNewEntityId(PolarisMetaStoreSession ms) { + return new GenerateEntityIdResult(ms.generateNewId()); } /** @@ -841,14 +766,14 @@ private void bootstrapPolarisService( * @param properties a map of key/value pairs * @return a String, the JSON representation of the map */ - public String serializeProperties(PolarisCallContext callCtx, Map properties) { + public String serializeProperties(Map properties) { String jsonString = null; try { // Deserialize the JSON string to a Map jsonString = MAPPER.writeValueAsString(properties); } catch (JsonProcessingException ex) { - callCtx.getDiagServices().fail("got_json_processing_exception", "ex={}", ex); + diagnostics.fail("got_json_processing_exception", "ex={}", ex); } return jsonString; @@ -860,80 +785,67 @@ public String serializeProperties(PolarisCallContext callCtx, Map deserializeProperties(PolarisCallContext callCtx, String properties) { + public Map deserializeProperties(String properties) { Map retProperties = null; try { // Deserialize the JSON string to a Map retProperties = MAPPER.readValue(properties, new TypeReference<>() {}); } catch (JsonMappingException ex) { - callCtx.getDiagServices().fail("got_json_mapping_exception", "ex={}", ex); + diagnostics.fail("got_json_mapping_exception", "ex={}", ex); } catch (JsonProcessingException ex) { - callCtx.getDiagServices().fail("got_json_processing_exception", "ex={}", ex); + diagnostics.fail("got_json_processing_exception", "ex={}", ex); } return retProperties; } - /** {@link #createPrincipal(PolarisCallContext, PolarisBaseEntity)} */ - private @Nonnull CreatePrincipalResult createPrincipal( - @Nonnull PolarisCallContext callCtx, - @Nonnull PolarisMetaStoreSession ms, - @Nonnull PolarisBaseEntity principal) { + /** {@link #createPrincipal(PolarisMetaStoreSession, PolarisBaseEntity)} */ + private @Nonnull CreatePrincipalResult doCreatePrincipal( + @Nonnull PolarisMetaStoreSession ms, @Nonnull PolarisBaseEntity principal) { // validate input - callCtx.getDiagServices().checkNotNull(principal, "unexpected_null_principal"); + diagnostics.checkNotNull(principal, "unexpected_null_principal"); // check if that catalog has already been created PolarisBaseEntity refreshPrincipal = - ms.lookupEntity(callCtx, principal.getCatalogId(), principal.getId()); + ms.lookupEntity(principal.getCatalogId(), principal.getId()); // if found, probably a retry, simply return the previously created principal if (refreshPrincipal != null) { // if found, ensure it is indeed a principal - callCtx - .getDiagServices() - .check( - principal.getTypeCode() == PolarisEntityType.PRINCIPAL.getCode(), - "not_a_principal", - "principal={}", - principal); + diagnostics.check( + principal.getTypeCode() == PolarisEntityType.PRINCIPAL.getCode(), + "not_a_principal", + "principal={}", + principal); // get internal properties Map properties = - this.deserializeProperties(callCtx, refreshPrincipal.getInternalProperties()); + this.deserializeProperties(refreshPrincipal.getInternalProperties()); // get client_id String clientId = properties.get(PolarisEntityConstants.getClientIdPropertyName()); // should not be null - callCtx - .getDiagServices() - .checkNotNull( - clientId, - "null_client_id", - "properties={}", - refreshPrincipal.getInternalProperties()); + diagnostics.checkNotNull( + clientId, "null_client_id", "properties={}", refreshPrincipal.getInternalProperties()); // ensure non null and non empty - callCtx - .getDiagServices() - .check( - !clientId.isEmpty(), - "empty_client_id", - "properties={}", - refreshPrincipal.getInternalProperties()); + diagnostics.check( + !clientId.isEmpty(), + "empty_client_id", + "properties={}", + refreshPrincipal.getInternalProperties()); // get the main and secondary secrets for that client - PolarisPrincipalSecrets principalSecrets = ms.loadPrincipalSecrets(callCtx, clientId); + PolarisPrincipalSecrets principalSecrets = ms.loadPrincipalSecrets(clientId); // should not be null - callCtx - .getDiagServices() - .checkNotNull( - principalSecrets, - "missing_principal_secrets", - "clientId={} principal={}", - clientId, - refreshPrincipal); + diagnostics.checkNotNull( + principalSecrets, + "missing_principal_secrets", + "clientId={} principal={}", + clientId, + refreshPrincipal); // done, return the newly created principal return new CreatePrincipalResult(refreshPrincipal, principalSecrets); @@ -946,8 +858,7 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str PolarisEntityConstants.getRootEntityId(), PolarisEntityType.PRINCIPAL.getCode(), principal.getName()); - PolarisEntityActiveRecord otherPrincipalRecord = - ms.lookupEntityActive(callCtx, principalNameKey); + PolarisEntityActiveRecord otherPrincipalRecord = ms.lookupEntityActive(principalNameKey); // if it exists, this is an error, the client should retry if (otherPrincipalRecord != null) { @@ -956,18 +867,18 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str // generate new secrets for this principal PolarisPrincipalSecrets principalSecrets = - ms.generateNewPrincipalSecrets(callCtx, principal.getName(), principal.getId()); + ms.generateNewPrincipalSecrets(principal.getName(), principal.getId()); // generate properties - Map internalProperties = getInternalPropertyMap(callCtx, principal); + Map internalProperties = getInternalPropertyMap(principal); internalProperties.put( PolarisEntityConstants.getClientIdPropertyName(), principalSecrets.getPrincipalClientId()); // remember client id - principal.setInternalProperties(this.serializeProperties(callCtx, internalProperties)); + principal.setInternalProperties(this.serializeProperties(internalProperties)); // now create and persist new catalog entity - this.persistNewEntity(callCtx, ms, principal); + this.persistNewEntity(ms, principal); // success, return the two entities return new CreatePrincipalResult(principal, principalSecrets); @@ -976,39 +887,30 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str /** {@inheritDoc} */ @Override public @Nonnull CreatePrincipalResult createPrincipal( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisBaseEntity principal) { - // get metastore we should be using - PolarisMetaStoreSession ms = callCtx.getMetaStore(); - - // need to run inside a read/write transaction - return ms.runInTransaction(callCtx, () -> this.createPrincipal(callCtx, ms, principal)); + PolarisMetaStoreSession ms, @Nonnull PolarisBaseEntity principal) { + return ms.runInTransaction(() -> this.doCreatePrincipal(ms, principal)); } - /** See {@link #loadPrincipalSecrets(PolarisCallContext, String)} */ - private @Nullable PolarisPrincipalSecrets loadPrincipalSecrets( - @Nonnull PolarisCallContext callCtx, PolarisMetaStoreSession ms, @Nonnull String clientId) { - return ms.loadPrincipalSecrets(callCtx, clientId); + /** See {@link #loadPrincipalSecrets(PolarisMetaStoreSession, String)} */ + private @Nullable PolarisPrincipalSecrets doLoadPrincipalSecrets( + PolarisMetaStoreSession ms, @Nonnull String clientId) { + return ms.loadPrincipalSecrets(clientId); } /** {@inheritDoc} */ @Override public @Nonnull PrincipalSecretsResult loadPrincipalSecrets( - @Nonnull PolarisCallContext callCtx, @Nonnull String clientId) { - // get metastore we should be using - PolarisMetaStoreSession ms = callCtx.getMetaStore(); - - // need to run inside a read/write transaction + PolarisMetaStoreSession ms, @Nonnull String clientId) { PolarisPrincipalSecrets secrets = - ms.runInTransaction(callCtx, () -> this.loadPrincipalSecrets(callCtx, ms, clientId)); + ms.runInTransaction(() -> this.doLoadPrincipalSecrets(ms, clientId)); return (secrets == null) ? new PrincipalSecretsResult(BaseResult.ReturnStatus.ENTITY_NOT_FOUND, null) : new PrincipalSecretsResult(secrets); } - /** See {@link #} */ - private @Nullable PolarisPrincipalSecrets rotatePrincipalSecrets( - @Nonnull PolarisCallContext callCtx, + private @Nullable PolarisPrincipalSecrets doRotatePrincipalSecrets( + @Nonnull PolarisDiagnostics diagnostics, @Nonnull PolarisMetaStoreSession ms, @Nonnull String clientId, long principalId, @@ -1016,7 +918,7 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str @Nonnull String oldSecretHash) { // if not found, the principal must have been dropped EntityResult loadEntityResult = - loadEntity(callCtx, ms, PolarisEntityConstants.getNullId(), principalId); + doLoadEntity(ms, PolarisEntityConstants.getNullId(), principalId); if (loadEntityResult.getReturnStatus() != BaseResult.ReturnStatus.SUCCESS) { return null; } @@ -1024,7 +926,7 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str PolarisBaseEntity principal = loadEntityResult.getEntity(); Map internalProps = PolarisObjectMapperUtil.deserializeProperties( - callCtx, + diagnostics, principal.getInternalProperties() == null ? "{}" : principal.getInternalProperties()); boolean doReset = @@ -1033,7 +935,7 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str PolarisEntityConstants.PRINCIPAL_CREDENTIAL_ROTATION_REQUIRED_STATE) != null; PolarisPrincipalSecrets secrets = - ms.rotatePrincipalSecrets(callCtx, clientId, principalId, doReset, oldSecretHash); + ms.rotatePrincipalSecrets(clientId, principalId, doReset, oldSecretHash); if (reset && !internalProps.containsKey( @@ -1041,16 +943,16 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str internalProps.put( PolarisEntityConstants.PRINCIPAL_CREDENTIAL_ROTATION_REQUIRED_STATE, "true"); principal.setInternalProperties( - PolarisObjectMapperUtil.serializeProperties(callCtx, internalProps)); + PolarisObjectMapperUtil.serializeProperties(diagnostics, internalProps)); principal.setEntityVersion(principal.getEntityVersion() + 1); - writeEntity(callCtx, ms, principal, true); + writeEntity(ms, principal, true); } else if (internalProps.containsKey( PolarisEntityConstants.PRINCIPAL_CREDENTIAL_ROTATION_REQUIRED_STATE)) { internalProps.remove(PolarisEntityConstants.PRINCIPAL_CREDENTIAL_ROTATION_REQUIRED_STATE); principal.setInternalProperties( - PolarisObjectMapperUtil.serializeProperties(callCtx, internalProps)); + PolarisObjectMapperUtil.serializeProperties(diagnostics, internalProps)); principal.setEntityVersion(principal.getEntityVersion() + 1); - writeEntity(callCtx, ms, principal, true); + writeEntity(ms, principal, true); } return secrets; } @@ -1058,21 +960,19 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str /** {@inheritDoc} */ @Override public @Nonnull PrincipalSecretsResult rotatePrincipalSecrets( - @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisMetaStoreSession ms, @Nonnull String clientId, long principalId, boolean reset, @Nonnull String oldSecretHash) { // get metastore we should be using - PolarisMetaStoreSession ms = callCtx.getMetaStore(); // need to run inside a read/write transaction PolarisPrincipalSecrets secrets = ms.runInTransaction( - callCtx, () -> - this.rotatePrincipalSecrets( - callCtx, ms, clientId, principalId, reset, oldSecretHash)); + this.doRotatePrincipalSecrets( + diagnostics, ms, clientId, principalId, reset, oldSecretHash)); return (secrets == null) ? new PrincipalSecretsResult(BaseResult.ReturnStatus.ENTITY_NOT_FOUND, null) @@ -1082,13 +982,10 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str /** {@inheritDoc} */ @Override public @Nonnull CreateCatalogResult createCatalog( - @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisMetaStoreSession ms, @Nonnull PolarisBaseEntity catalog, @Nonnull List principalRoles) { - // get metastore we should be using - PolarisMetaStoreSession ms = callCtx.getMetaStore(); - - Map internalProp = getInternalPropertyMap(callCtx, catalog); + Map internalProp = getInternalPropertyMap(catalog); String integrationIdentifierOrId = internalProp.get(PolarisEntityConstants.getStorageIntegrationIdentifierPropertyName()); String storageConfigInfoStr = @@ -1099,41 +996,38 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str if (storageConfigInfoStr != null && integrationIdentifierOrId == null) { integration = ms.createStorageIntegration( - callCtx, catalog.getCatalogId(), catalog.getId(), - PolarisStorageConfigurationInfo.deserialize( - callCtx.getDiagServices(), storageConfigInfoStr)); + PolarisStorageConfigurationInfo.deserialize(diagnostics, storageConfigInfoStr)); } else { integration = null; } // need to run inside a read/write transaction return ms.runInTransaction( - callCtx, () -> this.createCatalog(callCtx, ms, catalog, integration, principalRoles)); + () -> this.doCreateCatalog(ms, catalog, integration, principalRoles)); } - /** {@link #createEntityIfNotExists(PolarisCallContext, List, PolarisBaseEntity)} */ - private @Nonnull EntityResult createEntityIfNotExists( - @Nonnull PolarisCallContext callCtx, + /** {@link #createEntityIfNotExists(PolarisMetaStoreSession, List, PolarisBaseEntity)} */ + private @Nonnull EntityResult doCreateEntityIfNotExists( @Nonnull PolarisMetaStoreSession ms, @Nullable List catalogPath, @Nonnull PolarisBaseEntity entity) { // entity cannot be null - callCtx.getDiagServices().checkNotNull(entity, "unexpected_null_entity"); + diagnostics.checkNotNull(entity, "unexpected_null_entity"); // entity name must be specified - callCtx.getDiagServices().checkNotNull(entity.getName(), "unexpected_null_entity_name"); + diagnostics.checkNotNull(entity.getName(), "unexpected_null_entity_name"); // first, check if the entity has already been created, in which case we will simply return it - PolarisBaseEntity entityFound = ms.lookupEntity(callCtx, entity.getCatalogId(), entity.getId()); + PolarisBaseEntity entityFound = ms.lookupEntity(entity.getCatalogId(), entity.getId()); if (entityFound != null) { // probably the client retried, simply return it return new EntityResult(entityFound); } // first resolve again the catalogPath - PolarisEntityResolver resolver = new PolarisEntityResolver(callCtx, ms, catalogPath); + PolarisEntityResolver resolver = new PolarisEntityResolver(diagnostics, ms, catalogPath); // return if we failed to resolve if (resolver.isFailure()) { @@ -1147,14 +1041,14 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str entity.getParentId(), entity.getType().getCode(), entity.getName()); - PolarisEntityActiveRecord entityActiveRecord = ms.lookupEntityActive(callCtx, entityActiveKey); + PolarisEntityActiveRecord entityActiveRecord = ms.lookupEntityActive(entityActiveKey); if (entityActiveRecord != null) { return new EntityResult( BaseResult.ReturnStatus.ENTITY_ALREADY_EXISTS, entityActiveRecord.getSubTypeCode()); } // persist that new entity - this.persistNewEntity(callCtx, ms, entity); + this.persistNewEntity(ms, entity); // done, return that newly created entity return new EntityResult(entity); @@ -1163,33 +1057,28 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str /** {@inheritDoc} */ @Override public @Nonnull EntityResult createEntityIfNotExists( - @Nonnull PolarisCallContext callCtx, + PolarisMetaStoreSession ms, @Nullable List catalogPath, @Nonnull PolarisBaseEntity entity) { // get metastore we should be using - PolarisMetaStoreSession ms = callCtx.getMetaStore(); // need to run inside a read/write transaction - return ms.runInTransaction( - callCtx, () -> this.createEntityIfNotExists(callCtx, ms, catalogPath, entity)); + return ms.runInTransaction(() -> this.doCreateEntityIfNotExists(ms, catalogPath, entity)); } @Override public @Nonnull EntitiesResult createEntitiesIfNotExist( - @Nonnull PolarisCallContext callCtx, + PolarisMetaStoreSession ms, @Nullable List catalogPath, @Nonnull List entities) { // get metastore we should be using - PolarisMetaStoreSession ms = callCtx.getMetaStore(); // need to run inside a read/write transaction return ms.runInTransaction( - callCtx, () -> { List createdEntities = new ArrayList<>(entities.size()); for (PolarisBaseEntity entity : entities) { - EntityResult entityCreateResult = - createEntityIfNotExists(callCtx, ms, catalogPath, entity); + EntityResult entityCreateResult = doCreateEntityIfNotExists(ms, catalogPath, entity); // abort everything if error if (entityCreateResult.getReturnStatus() != BaseResult.ReturnStatus.SUCCESS) { ms.rollback(); @@ -1203,18 +1092,19 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str } /** - * See {@link #updateEntityPropertiesIfNotChanged(PolarisCallContext, List, PolarisBaseEntity)} + * See {@link #updateEntityPropertiesIfNotChanged(PolarisMetaStoreSession, List, + * PolarisBaseEntity)} */ - private @Nonnull EntityResult updateEntityPropertiesIfNotChanged( - @Nonnull PolarisCallContext callCtx, + private @Nonnull EntityResult doUpdateEntityPropertiesIfNotChanged( @Nonnull PolarisMetaStoreSession ms, @Nullable List catalogPath, @Nonnull PolarisBaseEntity entity) { // entity cannot be null - callCtx.getDiagServices().checkNotNull(entity, "unexpected_null_entity"); + diagnostics.checkNotNull(entity, "unexpected_null_entity"); // re-resolve everything including that entity - PolarisEntityResolver resolver = new PolarisEntityResolver(callCtx, ms, catalogPath, entity); + PolarisEntityResolver resolver = + new PolarisEntityResolver(diagnostics, ms, catalogPath, entity); // if resolution failed, return false if (resolver.isFailure()) { @@ -1222,11 +1112,8 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str } // lookup the entity, cannot be null - PolarisBaseEntity entityRefreshed = - ms.lookupEntity(callCtx, entity.getCatalogId(), entity.getId()); - callCtx - .getDiagServices() - .checkNotNull(entityRefreshed, "unexpected_entity_not_found", "entity={}", entity); + PolarisBaseEntity entityRefreshed = ms.lookupEntity(entity.getCatalogId(), entity.getId()); + diagnostics.checkNotNull(entityRefreshed, "unexpected_entity_not_found", "entity={}", entity); // check that the version of the entity has not changed at all to avoid concurrent updates if (entityRefreshed.getEntityVersion() != entity.getEntityVersion()) { @@ -1239,31 +1126,25 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str // persist this entity after changing it. This will update the version and update the last // updated time. Because the entity version is changed, we will update the change tracking table - PolarisBaseEntity persistedEntity = this.persistEntityAfterChange(callCtx, ms, entityRefreshed); + PolarisBaseEntity persistedEntity = this.persistEntityAfterChange(ms, entityRefreshed); return new EntityResult(persistedEntity); } /** {@inheritDoc} */ @Override public @Nonnull EntityResult updateEntityPropertiesIfNotChanged( - @Nonnull PolarisCallContext callCtx, + PolarisMetaStoreSession ms, @Nullable List catalogPath, @Nonnull PolarisBaseEntity entity) { - // get metastore we should be using - PolarisMetaStoreSession ms = callCtx.getMetaStore(); - - // need to run inside a read/write transaction return ms.runInTransaction( - callCtx, () -> this.updateEntityPropertiesIfNotChanged(callCtx, ms, catalogPath, entity)); + () -> this.doUpdateEntityPropertiesIfNotChanged(ms, catalogPath, entity)); } - /** See {@link #updateEntitiesPropertiesIfNotChanged(PolarisCallContext, List)} */ - private @Nonnull EntitiesResult updateEntitiesPropertiesIfNotChanged( - @Nonnull PolarisCallContext callCtx, - @Nonnull PolarisMetaStoreSession ms, - @Nonnull List entities) { + /** See {@link #updateEntitiesPropertiesIfNotChanged(PolarisMetaStoreSession, List)} */ + private @Nonnull EntitiesResult doUpdateEntitiesPropertiesIfNotChanged( + @Nonnull PolarisMetaStoreSession ms, @Nonnull List entities) { // ensure that the entities list is not null - callCtx.getDiagServices().checkNotNull(entities, "unexpected_null_entities"); + diagnostics.checkNotNull(entities, "unexpected_null_entities"); // list of all updated entities List updatedEntities = new ArrayList<>(entities.size()); @@ -1272,8 +1153,8 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str for (EntityWithPath entityWithPath : entities) { // update that entity, abort if it fails EntityResult updatedEntityResult = - this.updateEntityPropertiesIfNotChanged( - callCtx, ms, entityWithPath.getCatalogPath(), entityWithPath.getEntity()); + this.doUpdateEntityPropertiesIfNotChanged( + ms, entityWithPath.getCatalogPath(), entityWithPath.getEntity()); // if failed, rollback and return the last error if (updatedEntityResult.getReturnStatus() != BaseResult.ReturnStatus.SUCCESS) { @@ -1293,21 +1174,15 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str /** {@inheritDoc} */ @Override public @Nonnull EntitiesResult updateEntitiesPropertiesIfNotChanged( - @Nonnull PolarisCallContext callCtx, @Nonnull List entities) { - // get metastore we should be using - PolarisMetaStoreSession ms = callCtx.getMetaStore(); - - // need to run inside a read/write transaction - return ms.runInTransaction( - callCtx, () -> this.updateEntitiesPropertiesIfNotChanged(callCtx, ms, entities)); + PolarisMetaStoreSession ms, @Nonnull List entities) { + return ms.runInTransaction(() -> this.doUpdateEntitiesPropertiesIfNotChanged(ms, entities)); } /** - * See {@link PolarisMetaStoreManager#renameEntity(PolarisCallContext, List, PolarisEntityCore, - * List, PolarisEntity)} + * See {@link PolarisMetaStoreManager#renameEntity(PolarisMetaStoreSession, List, + * PolarisEntityCore, List, PolarisEntity)} */ - private @Nonnull EntityResult renameEntity( - @Nonnull PolarisCallContext callCtx, + private @Nonnull EntityResult doRenameEntity( @Nonnull PolarisMetaStoreSession ms, @Nullable List catalogPath, @Nonnull PolarisEntityCore entityToRename, @@ -1315,16 +1190,14 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str @Nonnull PolarisBaseEntity renamedEntity) { // entity and new name cannot be null - callCtx.getDiagServices().checkNotNull(entityToRename, "unexpected_null_entityToRename"); - callCtx.getDiagServices().checkNotNull(renamedEntity, "unexpected_null_renamedEntity"); + diagnostics.checkNotNull(entityToRename, "unexpected_null_entityToRename"); + diagnostics.checkNotNull(renamedEntity, "unexpected_null_renamedEntity"); // if a new catalog path is specified (i.e. re-parent operation), a catalog path should be // specified too - callCtx - .getDiagServices() - .check( - (newCatalogPath == null) || (catalogPath != null), - "newCatalogPath_specified_without_catalogPath"); + diagnostics.check( + (newCatalogPath == null) || (catalogPath != null), + "newCatalogPath_specified_without_catalogPath"); // null is shorthand for saying the path isn't changing if (newCatalogPath == null) { @@ -1333,7 +1206,7 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str // re-resolve everything including that entity PolarisEntityResolver resolver = - new PolarisEntityResolver(callCtx, ms, catalogPath, entityToRename); + new PolarisEntityResolver(diagnostics, ms, catalogPath, entityToRename); // if resolution failed, return false if (resolver.isFailure()) { @@ -1342,7 +1215,7 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str // find the entity to rename PolarisBaseEntity refreshEntityToRename = - ms.lookupEntity(callCtx, entityToRename.getCatalogId(), entityToRename.getId()); + ms.lookupEntity(entityToRename.getCatalogId(), entityToRename.getId()); // if this entity was not found, return failure. Not expected here because it was // resolved successfully (see above) @@ -1362,7 +1235,7 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str // re-resolve the new catalog path if this entity is going to be moved if (newCatalogPath != null) { - resolver = new PolarisEntityResolver(callCtx, ms, newCatalogPath); + resolver = new PolarisEntityResolver(diagnostics, ms, newCatalogPath); // if resolution failed, return false if (resolver.isFailure()) { @@ -1378,14 +1251,14 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str refreshEntityToRename.getTypeCode(), renamedEntity.getName()); // if this entity already exists, this is an error - PolarisEntityActiveRecord entityActiveRecord = ms.lookupEntityActive(callCtx, entityActiveKey); + PolarisEntityActiveRecord entityActiveRecord = ms.lookupEntityActive(entityActiveKey); if (entityActiveRecord != null) { return new EntityResult( BaseResult.ReturnStatus.ENTITY_ALREADY_EXISTS, entityActiveRecord.getSubTypeCode()); } // all good, delete the existing entity from the active slice - ms.deleteFromEntitiesActive(callCtx, refreshEntityToRename); + ms.deleteFromEntitiesActive(refreshEntityToRename); // change its name now refreshEntityToRename.setName(renamedEntity.getName()); @@ -1398,52 +1271,44 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str } // persist back to the active slice with its new name and parent - ms.writeToEntitiesActive(callCtx, refreshEntityToRename); + ms.writeToEntitiesActive(refreshEntityToRename); // persist the entity after change. This wil update the lastUpdateTimestamp and bump up the // version PolarisBaseEntity renamedEntityToReturn = - this.persistEntityAfterChange(callCtx, ms, refreshEntityToRename); + this.persistEntityAfterChange(ms, refreshEntityToRename); return new EntityResult(renamedEntityToReturn); } /** {@inheritDoc} */ @Override public @Nonnull EntityResult renameEntity( - @Nonnull PolarisCallContext callCtx, + PolarisMetaStoreSession ms, @Nullable List catalogPath, @Nonnull PolarisEntityCore entityToRename, @Nullable List newCatalogPath, @Nonnull PolarisEntity renamedEntity) { - // get metastore we should be using - PolarisMetaStoreSession ms = callCtx.getMetaStore(); - - // need to run inside a read/write transaction return ms.runInTransaction( - callCtx, - () -> - this.renameEntity( - callCtx, ms, catalogPath, entityToRename, newCatalogPath, renamedEntity)); + () -> this.doRenameEntity(ms, catalogPath, entityToRename, newCatalogPath, renamedEntity)); } /** * See * - *

{@link #dropEntityIfExists(PolarisCallContext, List, PolarisEntityCore, Map, boolean)} + *

{@link #dropEntityIfExists(PolarisMetaStoreSession, List, PolarisEntityCore, Map, boolean)} */ - private @Nonnull DropEntityResult dropEntityIfExists( - @Nonnull PolarisCallContext callCtx, + private @Nonnull DropEntityResult doDropEntityIfExists( @Nonnull PolarisMetaStoreSession ms, @Nullable List catalogPath, @Nonnull PolarisEntityCore entityToDrop, @Nullable Map cleanupProperties, boolean cleanup) { // entity cannot be null - callCtx.getDiagServices().checkNotNull(entityToDrop, "unexpected_null_entity"); + diagnostics.checkNotNull(entityToDrop, "unexpected_null_entity"); // re-resolve everything including that entity PolarisEntityResolver resolver = - new PolarisEntityResolver(callCtx, ms, catalogPath, entityToDrop); + new PolarisEntityResolver(diagnostics, ms, catalogPath, entityToDrop); // if resolution failed, return false if (resolver.isFailure()) { @@ -1452,7 +1317,7 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str // first find the entity to drop PolarisBaseEntity refreshEntityToDrop = - ms.lookupEntity(callCtx, entityToDrop.getCatalogId(), entityToDrop.getId()); + ms.lookupEntity(entityToDrop.getCatalogId(), entityToDrop.getId()); // if this entity was not found, return failure if (refreshEntityToDrop == null) { @@ -1471,14 +1336,13 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str long catalogId = refreshEntityToDrop.getId(); // if not all namespaces are dropped, we cannot drop this catalog - if (ms.hasChildren(callCtx, PolarisEntityType.NAMESPACE, catalogId, catalogId)) { + if (ms.hasChildren(PolarisEntityType.NAMESPACE, catalogId, catalogId)) { return new DropEntityResult(BaseResult.ReturnStatus.NAMESPACE_NOT_EMPTY, null); } // get the list of catalog roles, at most 2 List catalogRoles = ms.listActiveEntities( - callCtx, catalogId, catalogId, PolarisEntityType.CATALOG_ROLE, @@ -1494,18 +1358,17 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str // if 1, drop the last catalog role. Should be the catalog admin role but don't validate this if (!catalogRoles.isEmpty()) { // drop the last catalog role in that catalog, should be the admin catalog role - this.dropEntity(callCtx, ms, catalogRoles.get(0)); + this.dropEntity(ms, catalogRoles.get(0)); } } else if (refreshEntityToDrop.getType() == PolarisEntityType.NAMESPACE) { - if (ms.hasChildren( - callCtx, null, refreshEntityToDrop.getCatalogId(), refreshEntityToDrop.getId())) { + if (ms.hasChildren(null, refreshEntityToDrop.getCatalogId(), refreshEntityToDrop.getId())) { return new DropEntityResult(BaseResult.ReturnStatus.NAMESPACE_NOT_EMPTY, null); } } // 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); + this.dropEntity(ms, refreshEntityToDrop); // 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 @@ -1513,25 +1376,26 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str if (cleanup) { PolarisBaseEntity taskEntity = new PolarisEntity.Builder() - .setId(generateNewEntityId(callCtx).getId()) + .setId(generateNewEntityId(ms).getId()) .setCatalogId(0L) .setName("entityCleanup_" + entityToDrop.getId()) .setType(PolarisEntityType.TASK) .setSubType(PolarisEntitySubType.NULL_SUBTYPE) - .setCreateTimestamp(callCtx.getClock().millis()) + .setCreateTimestamp(clock.millis()) .build(); Map properties = new HashMap<>(); properties.put( PolarisTaskConstants.TASK_TYPE, String.valueOf(AsyncTaskType.ENTITY_CLEANUP_SCHEDULER.typeCode())); - properties.put("data", PolarisObjectMapperUtil.serialize(callCtx, refreshEntityToDrop)); - taskEntity.setProperties(PolarisObjectMapperUtil.serializeProperties(callCtx, properties)); + properties.put("data", PolarisObjectMapperUtil.serialize(diagnostics, refreshEntityToDrop)); + taskEntity.setProperties( + PolarisObjectMapperUtil.serializeProperties(diagnostics, properties)); if (cleanupProperties != null) { taskEntity.setInternalProperties( - PolarisObjectMapperUtil.serializeProperties(callCtx, cleanupProperties)); + PolarisObjectMapperUtil.serializeProperties(diagnostics, cleanupProperties)); } - createEntityIfNotExists(callCtx, ms, null, taskEntity); + doCreateEntityIfNotExists(ms, null, taskEntity); return new DropEntityResult(taskEntity.getId()); } @@ -1542,27 +1406,22 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str /** {@inheritDoc} */ @Override public @Nonnull DropEntityResult dropEntityIfExists( - @Nonnull PolarisCallContext callCtx, + PolarisMetaStoreSession ms, @Nullable List catalogPath, @Nonnull PolarisEntityCore entityToDrop, @Nullable Map cleanupProperties, boolean cleanup) { // get metastore we should be using - PolarisMetaStoreSession ms = callCtx.getMetaStore(); // need to run inside a read/write transaction return ms.runInTransaction( - callCtx, - () -> - this.dropEntityIfExists( - callCtx, ms, catalogPath, entityToDrop, cleanupProperties, cleanup)); + () -> this.doDropEntityIfExists(ms, catalogPath, entityToDrop, cleanupProperties, cleanup)); } /** * Resolve the arguments of granting/revoking a usage grant between a role (catalog or principal * role) and a grantee (either a principal role or a principal) * - * @param callCtx call context * @param ms meta store in read/write mode * @param catalog if the role is a catalog role, the caller needs to pass-in the catalog entity * which was used to resolve that role. Else null. @@ -1571,36 +1430,31 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str * @return resolver for the specified entities */ private @Nonnull PolarisEntityResolver resolveRoleToGranteeUsageGrant( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisMetaStoreSession ms, @Nullable PolarisEntityCore catalog, @Nonnull PolarisEntityCore role, @Nonnull PolarisEntityCore grantee) { // validate the grantee input - callCtx.getDiagServices().checkNotNull(grantee, "unexpected_null_grantee"); - callCtx - .getDiagServices() - .check(grantee.getType().isGrantee(), "not_a_grantee", "grantee={}", grantee); + diagnostics.checkNotNull(grantee, "unexpected_null_grantee"); + diagnostics.check(grantee.getType().isGrantee(), "not_a_grantee", "grantee={}", grantee); // validate role - callCtx.getDiagServices().checkNotNull(role, "unexpected_null_role"); + diagnostics.checkNotNull(role, "unexpected_null_role"); // role should be a catalog or a principal role boolean isCatalogRole = role.getTypeCode() == PolarisEntityType.CATALOG_ROLE.getCode(); boolean isPrincipalRole = role.getTypeCode() == PolarisEntityType.PRINCIPAL_ROLE.getCode(); - callCtx.getDiagServices().check(isCatalogRole || isPrincipalRole, "not_a_role"); + diagnostics.check(isCatalogRole || isPrincipalRole, "not_a_role"); // if the role is a catalog role, ensure a catalog is specified and // vice-versa, catalog should be null if the role is a principal role - callCtx - .getDiagServices() - .check( - (catalog == null && isPrincipalRole) || (catalog != null && isCatalogRole), - "catalog_mismatch", - "catalog={} role={}", - catalog, - role); + diagnostics.check( + (catalog == null && isPrincipalRole) || (catalog != null && isCatalogRole), + "catalog_mismatch", + "catalog={} role={}", + catalog, + role); // re-resolve now all these entities List otherTopLevelEntities = new ArrayList<>(2); @@ -1609,7 +1463,7 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str // ensure these entities have not changed return new PolarisEntityResolver( - callCtx, ms, catalog != null ? List.of(catalog) : null, null, otherTopLevelEntities); + diagnostics, ms, catalog != null ? List.of(catalog) : null, null, otherTopLevelEntities); } /** @@ -1621,34 +1475,30 @@ public Map deserializeProperties(PolarisCallContext callCtx, Str * @return a resolver for the role, the catalog path and the securable */ private PolarisEntityResolver resolveSecurableToRoleGrant( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisMetaStoreSession ms, @Nonnull PolarisEntityCore grantee, @Nullable List catalogPath, @Nonnull PolarisEntityCore securable) { // validate role input - callCtx.getDiagServices().checkNotNull(grantee, "unexpected_null_grantee"); - callCtx - .getDiagServices() - .check(grantee.getType().isGrantee(), "not_grantee_type", "grantee={}", grantee); + diagnostics.checkNotNull(grantee, "unexpected_null_grantee"); + diagnostics.check(grantee.getType().isGrantee(), "not_grantee_type", "grantee={}", grantee); // securable must be supplied - callCtx.getDiagServices().checkNotNull(securable, "unexpected_null_securable"); + diagnostics.checkNotNull(securable, "unexpected_null_securable"); if (securable.getCatalogId() > 0) { // catalogPath must be supplied if the securable has a catalogId - callCtx.getDiagServices().checkNotNull(catalogPath, "unexpected_null_catalogPath"); + diagnostics.checkNotNull(catalogPath, "unexpected_null_catalogPath"); } // re-resolve now all these entities - return new PolarisEntityResolver(callCtx, ms, catalogPath, securable, List.of(grantee)); + return new PolarisEntityResolver(diagnostics, ms, catalogPath, securable, List.of(grantee)); } /** - * See {@link #grantUsageOnRoleToGrantee(PolarisCallContext, PolarisEntityCore, PolarisEntityCore, - * PolarisEntityCore)} + * See {@link #grantUsageOnRoleToGrantee(PolarisMetaStoreSession, PolarisEntityCore, + * PolarisEntityCore, PolarisEntityCore)} */ - private @Nonnull PrivilegeResult grantUsageOnRoleToGrantee( - @Nonnull PolarisCallContext callCtx, + private @Nonnull PrivilegeResult doGrantUsageOnRoleToGrantee( @Nonnull PolarisMetaStoreSession ms, @Nullable PolarisEntityCore catalog, @Nonnull PolarisEntityCore role, @@ -1656,7 +1506,7 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( // ensure these entities have not changed PolarisEntityResolver resolver = - this.resolveRoleToGranteeUsageGrant(callCtx, ms, catalog, role, grantee); + this.resolveRoleToGranteeUsageGrant(ms, catalog, role, grantee); // if failure to resolve, let the caller know if (resolver.isFailure()) { @@ -1670,35 +1520,26 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( : PolarisPrivilege.PRINCIPAL_ROLE_USAGE; // grant usage on this role to this principal - callCtx - .getDiagServices() - .check(grantee.getType().isGrantee(), "not_a_grantee", "grantee={}", grantee); - PolarisGrantRecord grantRecord = - this.persistNewGrantRecord(callCtx, ms, role, grantee, usagePriv); + diagnostics.check(grantee.getType().isGrantee(), "not_a_grantee", "grantee={}", grantee); + PolarisGrantRecord grantRecord = this.persistNewGrantRecord(ms, role, grantee, usagePriv); return new PrivilegeResult(grantRecord); } /** {@inheritDoc} */ @Override public @Nonnull PrivilegeResult grantUsageOnRoleToGrantee( - @Nonnull PolarisCallContext callCtx, + PolarisMetaStoreSession ms, @Nullable PolarisEntityCore catalog, @Nonnull PolarisEntityCore role, @Nonnull PolarisEntityCore grantee) { - // get metastore we should be using - PolarisMetaStoreSession ms = callCtx.getMetaStore(); - - // need to run inside a read/write transaction - return ms.runInTransaction( - callCtx, () -> this.grantUsageOnRoleToGrantee(callCtx, ms, catalog, role, grantee)); + return ms.runInTransaction(() -> this.doGrantUsageOnRoleToGrantee(ms, catalog, role, grantee)); } /** - * See {@link #revokeUsageOnRoleFromGrantee(PolarisCallContext, PolarisEntityCore, + * See {@link #revokeUsageOnRoleFromGrantee(PolarisMetaStoreSession, PolarisEntityCore, * PolarisEntityCore, PolarisEntityCore)} */ - private @Nonnull PrivilegeResult revokeUsageOnRoleFromGrantee( - @Nonnull PolarisCallContext callCtx, + private @Nonnull PrivilegeResult doRevokeUsageOnRoleFromGrantee( @Nonnull PolarisMetaStoreSession ms, @Nullable PolarisEntityCore catalog, @Nonnull PolarisEntityCore role, @@ -1706,7 +1547,7 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( // ensure these entities have not changed PolarisEntityResolver resolver = - this.resolveRoleToGranteeUsageGrant(callCtx, ms, catalog, role, grantee); + this.resolveRoleToGranteeUsageGrant(ms, catalog, role, grantee); // if failure to resolve, let the caller know if (resolver.isFailure()) { @@ -1722,7 +1563,6 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( // first, ensure that this privilege has been granted PolarisGrantRecord grantRecord = ms.lookupGrantRecord( - callCtx, role.getCatalogId(), role.getId(), grantee.getCatalogId(), @@ -1735,7 +1575,7 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( } // revoke usage on the role from the grantee - this.revokeGrantRecord(callCtx, ms, role, grantee, grantRecord); + this.revokeGrantRecord(ms, role, grantee, grantRecord); return new PrivilegeResult(grantRecord); } @@ -1743,24 +1583,19 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( /** {@inheritDoc} */ @Override public @Nonnull PrivilegeResult revokeUsageOnRoleFromGrantee( - @Nonnull PolarisCallContext callCtx, + PolarisMetaStoreSession ms, @Nullable PolarisEntityCore catalog, @Nonnull PolarisEntityCore role, @Nonnull PolarisEntityCore grantee) { - // get metastore we should be using - PolarisMetaStoreSession ms = callCtx.getMetaStore(); - - // need to run inside a read/write transaction return ms.runInTransaction( - callCtx, () -> this.revokeUsageOnRoleFromGrantee(callCtx, ms, catalog, role, grantee)); + () -> this.doRevokeUsageOnRoleFromGrantee(ms, catalog, role, grantee)); } /** - * See {@link #grantPrivilegeOnSecurableToRole(PolarisCallContext, PolarisEntityCore, List, + * See {@link #grantPrivilegeOnSecurableToRole(PolarisMetaStoreSession, PolarisEntityCore, List, * PolarisEntityCore, PolarisPrivilege)} */ - private @Nonnull PrivilegeResult grantPrivilegeOnSecurableToRole( - @Nonnull PolarisCallContext callCtx, + private @Nonnull PrivilegeResult doGrantPrivilegeOnSecurableToRole( @Nonnull PolarisMetaStoreSession ms, @Nonnull PolarisEntityCore grantee, @Nullable List catalogPath, @@ -1769,7 +1604,7 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( // re-resolve now all these entities PolarisEntityResolver resolver = - this.resolveSecurableToRoleGrant(callCtx, ms, grantee, catalogPath, securable); + this.resolveSecurableToRoleGrant(ms, grantee, catalogPath, securable); // if failure to resolve, let the caller know if (resolver.isFailure()) { @@ -1777,36 +1612,28 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( } // grant specified privilege on this securable to this role and return the grant - PolarisGrantRecord grantRecord = - this.persistNewGrantRecord(callCtx, ms, securable, grantee, priv); + PolarisGrantRecord grantRecord = this.persistNewGrantRecord(ms, securable, grantee, priv); return new PrivilegeResult(grantRecord); } /** {@inheritDoc} */ @Override public @Nonnull PrivilegeResult grantPrivilegeOnSecurableToRole( - @Nonnull PolarisCallContext callCtx, + PolarisMetaStoreSession ms, @Nonnull PolarisEntityCore grantee, @Nullable List catalogPath, @Nonnull PolarisEntityCore securable, @Nonnull PolarisPrivilege privilege) { - // get metastore we should be using - PolarisMetaStoreSession ms = callCtx.getMetaStore(); - - // need to run inside a read/write transaction return ms.runInTransaction( - callCtx, () -> - this.grantPrivilegeOnSecurableToRole( - callCtx, ms, grantee, catalogPath, securable, privilege)); + this.doGrantPrivilegeOnSecurableToRole(ms, grantee, catalogPath, securable, privilege)); } /** - * See {@link #revokePrivilegeOnSecurableFromRole(PolarisCallContext, PolarisEntityCore, List, - * PolarisEntityCore, PolarisPrivilege)} + * See {@link #revokePrivilegeOnSecurableFromRole(PolarisMetaStoreSession, PolarisEntityCore, + * List, PolarisEntityCore, PolarisPrivilege)} */ - private @Nonnull PrivilegeResult revokePrivilegeOnSecurableFromRole( - @Nonnull PolarisCallContext callCtx, + private @Nonnull PrivilegeResult doRevokePrivilegeOnSecurableFromRole( @Nonnull PolarisMetaStoreSession ms, @Nonnull PolarisEntityCore grantee, @Nullable List catalogPath, @@ -1815,7 +1642,7 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( // re-resolve now all these entities PolarisEntityResolver resolver = - this.resolveSecurableToRoleGrant(callCtx, ms, grantee, catalogPath, securable); + this.resolveSecurableToRoleGrant(ms, grantee, catalogPath, securable); // if failure to resolve, let the caller know if (resolver.isFailure()) { @@ -1825,7 +1652,6 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( // lookup the grants records to find this grant PolarisGrantRecord grantRecord = ms.lookupGrantRecord( - callCtx, securable.getCatalogId(), securable.getId(), grantee.getCatalogId(), @@ -1838,7 +1664,7 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( } // revoke the specified privilege on this securable from this role - this.revokeGrantRecord(callCtx, ms, securable, grantee, grantRecord); + this.revokeGrantRecord(ms, securable, grantee, grantRecord); // success! return new PrivilegeResult(grantRecord); @@ -1847,32 +1673,23 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( /** {@inheritDoc} */ @Override public @Nonnull PrivilegeResult revokePrivilegeOnSecurableFromRole( - @Nonnull PolarisCallContext callCtx, + PolarisMetaStoreSession ms, @Nonnull PolarisEntityCore grantee, @Nullable List catalogPath, @Nonnull PolarisEntityCore securable, @Nonnull PolarisPrivilege privilege) { - // get metastore we should be using - PolarisMetaStoreSession ms = callCtx.getMetaStore(); - - // need to run inside a read/write transaction return ms.runInTransaction( - callCtx, () -> - this.revokePrivilegeOnSecurableFromRole( - callCtx, ms, grantee, catalogPath, securable, privilege)); + this.doRevokePrivilegeOnSecurableFromRole( + ms, grantee, catalogPath, securable, privilege)); } - /** {@link #loadGrantsOnSecurable(PolarisCallContext, long, long)} */ - private @Nonnull LoadGrantsResult loadGrantsOnSecurable( - @Nonnull PolarisCallContext callCtx, - @Nonnull PolarisMetaStoreSession ms, - long securableCatalogId, - long securableId) { + /** {@link #loadGrantsOnSecurable(PolarisMetaStoreSession, long, long)} */ + private @Nonnull LoadGrantsResult doLoadGrantsOnSecurable( + @Nonnull PolarisMetaStoreSession ms, long securableCatalogId, long securableId) { // lookup grants version for this securable entity - int grantsVersion = - ms.lookupEntityGrantRecordsVersion(callCtx, securableCatalogId, securableId); + int grantsVersion = ms.lookupEntityGrantRecordsVersion(securableCatalogId, securableId); // return null if securable does not exists if (grantsVersion == 0) { @@ -1881,7 +1698,7 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( // now fetch all grants for this securable final List returnGrantRecords = - ms.loadAllGrantRecordsOnSecurable(callCtx, securableCatalogId, securableId); + ms.loadAllGrantRecordsOnSecurable(securableCatalogId, securableId); // find all unique grantees List entityIds = @@ -1892,7 +1709,7 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( grantRecord.getGranteeCatalogId(), grantRecord.getGranteeId())) .distinct() .collect(Collectors.toList()); - List entities = ms.lookupEntities(callCtx, entityIds); + List entities = ms.lookupEntities(entityIds); // done, return the list of grants and their version return new LoadGrantsResult( @@ -1904,24 +1721,17 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( /** {@inheritDoc} */ @Override public @Nonnull LoadGrantsResult loadGrantsOnSecurable( - @Nonnull PolarisCallContext callCtx, long securableCatalogId, long securableId) { - // get metastore we should be using - PolarisMetaStoreSession ms = callCtx.getMetaStore(); - - // need to run inside a read transaction + PolarisMetaStoreSession ms, long securableCatalogId, long securableId) { return ms.runInReadTransaction( - callCtx, () -> this.loadGrantsOnSecurable(callCtx, ms, securableCatalogId, securableId)); + () -> this.doLoadGrantsOnSecurable(ms, securableCatalogId, securableId)); } - /** {@link #loadGrantsToGrantee(PolarisCallContext, long, long)} */ - public @Nonnull LoadGrantsResult loadGrantsToGrantee( - @Nonnull PolarisCallContext callCtx, - @Nonnull PolarisMetaStoreSession ms, - long granteeCatalogId, - long granteeId) { + /** {@link #loadGrantsToGrantee(PolarisMetaStoreSession, long, long)} */ + private @Nonnull LoadGrantsResult doLoadGrantsToGrantee( + @Nonnull PolarisMetaStoreSession ms, long granteeCatalogId, long granteeId) { // lookup grants version for this grantee entity - int grantsVersion = ms.lookupEntityGrantRecordsVersion(callCtx, granteeCatalogId, granteeId); + int grantsVersion = ms.lookupEntityGrantRecordsVersion(granteeCatalogId, granteeId); // return null if grantee does not exists if (grantsVersion == 0) { @@ -1930,7 +1740,7 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( // now fetch all grants for this grantee final List returnGrantRecords = - ms.loadAllGrantRecordsOnGrantee(callCtx, granteeCatalogId, granteeId); + ms.loadAllGrantRecordsOnGrantee(granteeCatalogId, granteeId); // find all unique securables List entityIds = @@ -1941,7 +1751,7 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( grantRecord.getSecurableCatalogId(), grantRecord.getSecurableId())) .distinct() .collect(Collectors.toList()); - List entities = ms.lookupEntities(callCtx, entityIds); + List entities = ms.lookupEntities(entityIds); // done, return the list of grants and their version return new LoadGrantsResult( @@ -1953,45 +1763,30 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( /** {@inheritDoc} */ @Override public @Nonnull LoadGrantsResult loadGrantsToGrantee( - @Nonnull PolarisCallContext callCtx, long granteeCatalogId, long granteeId) { - // get metastore we should be using - PolarisMetaStoreSession ms = callCtx.getMetaStore(); - - // need to run inside a read transaction + PolarisMetaStoreSession ms, long granteeCatalogId, long granteeId) { return ms.runInReadTransaction( - callCtx, () -> this.loadGrantsToGrantee(callCtx, ms, granteeCatalogId, granteeId)); + () -> this.doLoadGrantsToGrantee(ms, granteeCatalogId, granteeId)); } - /** {@link PolarisMetaStoreManager#loadEntitiesChangeTracking(PolarisCallContext, List)} */ - private @Nonnull ChangeTrackingResult loadEntitiesChangeTracking( - @Nonnull PolarisCallContext callCtx, - @Nonnull PolarisMetaStoreSession ms, - @Nonnull List entityIds) { - List changeTracking = - ms.lookupEntityVersions(callCtx, entityIds); + /** {@link PolarisMetaStoreManager#loadEntitiesChangeTracking(PolarisMetaStoreSession, List)} */ + private @Nonnull ChangeTrackingResult doLoadEntitiesChangeTracking( + @Nonnull PolarisMetaStoreSession ms, @Nonnull List entityIds) { + List changeTracking = ms.lookupEntityVersions(entityIds); return new ChangeTrackingResult(changeTracking); } /** {@inheritDoc} */ @Override public @Nonnull ChangeTrackingResult loadEntitiesChangeTracking( - @Nonnull PolarisCallContext callCtx, @Nonnull List entityIds) { - // get metastore we should be using - PolarisMetaStoreSession ms = callCtx.getMetaStore(); - - // need to run inside a read transaction - return ms.runInReadTransaction( - callCtx, () -> this.loadEntitiesChangeTracking(callCtx, ms, entityIds)); + PolarisMetaStoreSession ms, @Nonnull List entityIds) { + return ms.runInReadTransaction(() -> this.doLoadEntitiesChangeTracking(ms, entityIds)); } - /** Refer to {@link #loadEntity(PolarisCallContext, long, long)} */ - private @Nonnull EntityResult loadEntity( - @Nonnull PolarisCallContext callCtx, - @Nonnull PolarisMetaStoreSession ms, - long entityCatalogId, - long entityId) { + /** Refer to {@link #loadEntity(PolarisMetaStoreSession, long, long)} */ + private @Nonnull EntityResult doLoadEntity( + @Nonnull PolarisMetaStoreSession ms, long entityCatalogId, long entityId) { // this is an easy one - PolarisBaseEntity entity = ms.lookupEntity(callCtx, entityCatalogId, entityId); + PolarisBaseEntity entity = ms.lookupEntity(entityCatalogId, entityId); return (entity != null) ? new EntityResult(entity) : new EntityResult(BaseResult.ReturnStatus.ENTITY_NOT_FOUND, null); @@ -2000,26 +1795,17 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( /** {@inheritDoc} */ @Override public @Nonnull EntityResult loadEntity( - @Nonnull PolarisCallContext callCtx, long entityCatalogId, long entityId) { - // get metastore we should be using - PolarisMetaStoreSession ms = callCtx.getMetaStore(); - - // need to run inside a read transaction - return ms.runInReadTransaction( - callCtx, () -> this.loadEntity(callCtx, ms, entityCatalogId, entityId)); + PolarisMetaStoreSession ms, long entityCatalogId, long entityId) { + return ms.runInReadTransaction(() -> this.doLoadEntity(ms, entityCatalogId, entityId)); } - /** Refer to {@link #loadTasks(PolarisCallContext, String, int)} */ - private @Nonnull EntitiesResult loadTasks( - @Nonnull PolarisCallContext callCtx, - @Nonnull PolarisMetaStoreSession ms, - String executorId, - int limit) { + /** Refer to {@link #loadTasks(PolarisMetaStoreSession, String, int)} */ + private @Nonnull EntitiesResult doLoadTasks( + @Nonnull PolarisMetaStoreSession ms, String executorId, int limit) { // find all available tasks List availableTasks = ms.listActiveEntities( - callCtx, PolarisEntityConstants.getRootEntityId(), PolarisEntityConstants.getRootEntityId(), PolarisEntityType.TASK, @@ -2028,49 +1814,45 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( PolarisObjectMapperUtil.TaskExecutionState taskState = PolarisObjectMapperUtil.parseTaskState(entity); long taskAgeTimeout = - callCtx - .getConfigurationStore() - .getConfiguration( - callCtx, - PolarisTaskConstants.TASK_TIMEOUT_MILLIS_CONFIG, - PolarisTaskConstants.TASK_TIMEOUT_MILLIS); + configurationStore.getConfiguration( + realmContext, + PolarisTaskConstants.TASK_TIMEOUT_MILLIS_CONFIG, + PolarisTaskConstants.TASK_TIMEOUT_MILLIS); return taskState == null || taskState.executor == null - || callCtx.getClock().millis() - taskState.lastAttemptStartTime > taskAgeTimeout; + || clock.millis() - taskState.lastAttemptStartTime > taskAgeTimeout; }, Function.identity()); availableTasks.forEach( task -> { Map properties = - PolarisObjectMapperUtil.deserializeProperties(callCtx, task.getProperties()); + PolarisObjectMapperUtil.deserializeProperties(diagnostics, task.getProperties()); properties.put(PolarisTaskConstants.LAST_ATTEMPT_EXECUTOR_ID, executorId); properties.put( - PolarisTaskConstants.LAST_ATTEMPT_START_TIME, - String.valueOf(callCtx.getClock().millis())); + PolarisTaskConstants.LAST_ATTEMPT_START_TIME, String.valueOf(clock.millis())); properties.put( PolarisTaskConstants.ATTEMPT_COUNT, String.valueOf( Integer.parseInt(properties.getOrDefault(PolarisTaskConstants.ATTEMPT_COUNT, "0")) + 1)); task.setEntityVersion(task.getEntityVersion() + 1); - task.setProperties(PolarisObjectMapperUtil.serializeProperties(callCtx, properties)); - writeEntity(callCtx, ms, task, false); + task.setProperties(PolarisObjectMapperUtil.serializeProperties(diagnostics, properties)); + writeEntity(ms, task, false); }); return new EntitiesResult(availableTasks); } @Override public @Nonnull EntitiesResult loadTasks( - @Nonnull PolarisCallContext callCtx, String executorId, int limit) { - PolarisMetaStoreSession ms = callCtx.getMetaStore(); - return ms.runInTransaction(callCtx, () -> this.loadTasks(callCtx, ms, executorId, limit)); + PolarisMetaStoreSession ms, String executorId, int limit) { + return ms.runInTransaction(() -> this.doLoadTasks(ms, executorId, limit)); } /** {@inheritDoc} */ @Override public @Nonnull ScopedCredentialsResult getSubscopedCredsForEntity( - @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisMetaStoreSession ms, long catalogId, long entityId, boolean allowListOperation, @@ -2078,15 +1860,12 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( @Nonnull Set allowedWriteLocations) { // get meta store session we should be using - PolarisMetaStoreSession ms = callCtx.getMetaStore(); - callCtx - .getDiagServices() - .check( - !allowedReadLocations.isEmpty() || !allowedWriteLocations.isEmpty(), - "allowed_locations_to_subscope_is_required"); + diagnostics.check( + !allowedReadLocations.isEmpty() || !allowedWriteLocations.isEmpty(), + "allowed_locations_to_subscope_is_required"); // reload the entity, error out if not found - EntityResult reloadedEntity = loadEntity(callCtx, catalogId, entityId); + EntityResult reloadedEntity = loadEntity(ms, catalogId, entityId); if (reloadedEntity.getReturnStatus() != BaseResult.ReturnStatus.SUCCESS) { return new ScopedCredentialsResult( reloadedEntity.getReturnStatus(), reloadedEntity.getExtraInformation()); @@ -2094,24 +1873,23 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( // get storage integration PolarisStorageIntegration storageIntegration = - ms.loadPolarisStorageIntegration(callCtx, reloadedEntity.getEntity()); + ms.loadPolarisStorageIntegration(reloadedEntity.getEntity()); // cannot be null - callCtx - .getDiagServices() - .checkNotNull( - storageIntegration, - "storage_integration_not_exists", - "catalogId={}, entityId={}", - catalogId, - entityId); + diagnostics.checkNotNull( + storageIntegration, + "storage_integration_not_exists", + "catalogId={}, entityId={}", + catalogId, + entityId); PolarisStorageConfigurationInfo storageConfigurationInfo = - readStorageConfiguration(callCtx, reloadedEntity.getEntity()); + readStorageConfiguration(diagnostics, reloadedEntity.getEntity()); try { EnumMap creds = storageIntegration.getSubscopedCreds( - callCtx.getDiagServices(), + realmContext, + diagnostics, storageConfigurationInfo, allowListOperation, allowedReadLocations, @@ -2126,20 +1904,17 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( /** {@inheritDoc} */ @Override public @Nonnull ValidateAccessResult validateAccessToLocations( - @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisMetaStoreSession metaStoreSession, long catalogId, long entityId, @Nonnull Set actions, @Nonnull Set locations) { // get meta store we should be using - PolarisMetaStoreSession ms = callCtx.getMetaStore(); - callCtx - .getDiagServices() - .check( - !actions.isEmpty() && !locations.isEmpty(), - "locations_and_operations_privileges_are_required"); + diagnostics.check( + !actions.isEmpty() && !locations.isEmpty(), + "locations_and_operations_privileges_are_required"); // reload the entity, error out if not found - EntityResult reloadedEntity = loadEntity(callCtx, catalogId, entityId); + EntityResult reloadedEntity = loadEntity(metaStoreSession, catalogId, entityId); if (reloadedEntity.getReturnStatus() != BaseResult.ReturnStatus.SUCCESS) { return new ValidateAccessResult( reloadedEntity.getReturnStatus(), reloadedEntity.getExtraInformation()); @@ -2147,79 +1922,69 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( // get storage integration, expect not null PolarisStorageIntegration storageIntegration = - ms.loadPolarisStorageIntegration(callCtx, reloadedEntity.getEntity()); - callCtx - .getDiagServices() - .checkNotNull( - storageIntegration, - "storage_integration_not_exists", - "catalogId={}, entityId={}", - catalogId, - entityId); + metaStoreSession.loadPolarisStorageIntegration(reloadedEntity.getEntity()); + diagnostics.checkNotNull( + storageIntegration, + "storage_integration_not_exists", + "catalogId={}, entityId={}", + catalogId, + entityId); // validate access PolarisStorageConfigurationInfo storageConfigurationInfo = - readStorageConfiguration(callCtx, reloadedEntity.getEntity()); + readStorageConfiguration(diagnostics, reloadedEntity.getEntity()); Map validateLocationAccess = storageIntegration - .validateAccessToLocations(storageConfigurationInfo, actions, locations) + .validateAccessToLocations(realmContext, storageConfigurationInfo, actions, locations) .entrySet() .stream() .collect( Collectors.toMap( Map.Entry::getKey, - e -> PolarisObjectMapperUtil.serialize(callCtx, e.getValue()))); + e -> PolarisObjectMapperUtil.serialize(diagnostics, e.getValue()))); // done, return result return new ValidateAccessResult(validateLocationAccess); } public static PolarisStorageConfigurationInfo readStorageConfiguration( - @Nonnull PolarisCallContext callCtx, PolarisBaseEntity reloadedEntity) { + PolarisDiagnostics diagnostics, PolarisBaseEntity reloadedEntity) { Map propMap = PolarisObjectMapperUtil.deserializeProperties( - callCtx, reloadedEntity.getInternalProperties()); + diagnostics, reloadedEntity.getInternalProperties()); String storageConfigInfoStr = propMap.get(PolarisEntityConstants.getStorageConfigInfoPropertyName()); - callCtx - .getDiagServices() - .check( - storageConfigInfoStr != null, - "missing_storage_configuration_info", - "catalogId={}, entityId={}", - reloadedEntity.getCatalogId(), - reloadedEntity.getId()); - return PolarisStorageConfigurationInfo.deserialize( - callCtx.getDiagServices(), storageConfigInfoStr); + diagnostics.check( + storageConfigInfoStr != null, + "missing_storage_configuration_info", + "catalogId={}, entityId={}", + reloadedEntity.getCatalogId(), + reloadedEntity.getId()); + return PolarisStorageConfigurationInfo.deserialize(diagnostics, storageConfigInfoStr); } /** * Get the internal property map for an entity * - * @param callCtx the polaris call context * @param entity the target entity * @return a map of string representing the internal properties */ - public Map getInternalPropertyMap( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisBaseEntity entity) { + public Map getInternalPropertyMap(@Nonnull PolarisBaseEntity entity) { String internalPropStr = entity.getInternalProperties(); Map res = new HashMap<>(); if (internalPropStr == null) { return res; } - return deserializeProperties(callCtx, internalPropStr); + return deserializeProperties(internalPropStr); } - /** {@link #loadCachedEntryById(PolarisCallContext, long, long)} */ - private @Nonnull CachedEntryResult loadCachedEntryById( - @Nonnull PolarisCallContext callCtx, - @Nonnull PolarisMetaStoreSession ms, - long entityCatalogId, - long entityId) { + /** {@link #loadCachedEntryById(PolarisMetaStoreSession, long, long)} */ + private @Nonnull CachedEntryResult doLoadCachedEntryById( + @Nonnull PolarisMetaStoreSession ms, long entityCatalogId, long entityId) { // load that entity - PolarisBaseEntity entity = ms.lookupEntity(callCtx, entityCatalogId, entityId); + PolarisBaseEntity entity = ms.lookupEntity(entityCatalogId, entityId); // if entity not found, return null if (entity == null) { @@ -2229,11 +1994,10 @@ public Map getInternalPropertyMap( // load the grant records final List grantRecords; if (entity.getType().isGrantee()) { - grantRecords = - new ArrayList<>(ms.loadAllGrantRecordsOnGrantee(callCtx, entityCatalogId, entityId)); - grantRecords.addAll(ms.loadAllGrantRecordsOnSecurable(callCtx, entityCatalogId, entityId)); + grantRecords = new ArrayList<>(ms.loadAllGrantRecordsOnGrantee(entityCatalogId, entityId)); + grantRecords.addAll(ms.loadAllGrantRecordsOnSecurable(entityCatalogId, entityId)); } else { - grantRecords = ms.loadAllGrantRecordsOnSecurable(callCtx, entityCatalogId, entityId); + grantRecords = ms.loadAllGrantRecordsOnSecurable(entityCatalogId, entityId); } // return the result @@ -2243,18 +2007,14 @@ public Map getInternalPropertyMap( /** {@inheritDoc} */ @Override public @Nonnull CachedEntryResult loadCachedEntryById( - @Nonnull PolarisCallContext callCtx, long entityCatalogId, long entityId) { - // get metastore we should be using - PolarisMetaStoreSession ms = callCtx.getMetaStore(); - - // need to run inside a read transaction - return ms.runInReadTransaction( - callCtx, () -> this.loadCachedEntryById(callCtx, ms, entityCatalogId, entityId)); + PolarisMetaStoreSession ms, long entityCatalogId, long entityId) { + return ms.runInReadTransaction(() -> this.doLoadCachedEntryById(ms, entityCatalogId, entityId)); } - /** {@link #loadCachedEntryById(PolarisCallContext, long, long)} */ - private @Nonnull CachedEntryResult loadCachedEntryByName( - @Nonnull PolarisCallContext callCtx, + /** + * {@link #loadCachedEntryByName(PolarisMetaStoreSession, long, long, PolarisEntityType, String)} + */ + private @Nonnull CachedEntryResult doLoadCachedEntryByName( @Nonnull PolarisMetaStoreSession ms, long entityCatalogId, long parentId, @@ -2264,7 +2024,7 @@ public Map getInternalPropertyMap( // load that entity PolarisEntitiesActiveKey entityActiveKey = new PolarisEntitiesActiveKey(entityCatalogId, parentId, entityType.getCode(), entityName); - PolarisBaseEntity entity = this.lookupEntityByName(callCtx, ms, entityActiveKey); + PolarisBaseEntity entity = this.lookupEntityByName(ms, entityActiveKey); // null if entity not found if (entity == null) { @@ -2275,12 +2035,10 @@ public Map getInternalPropertyMap( final List grantRecords; if (entity.getType().isGrantee()) { grantRecords = - new ArrayList<>( - ms.loadAllGrantRecordsOnGrantee(callCtx, entityCatalogId, entity.getId())); - grantRecords.addAll( - ms.loadAllGrantRecordsOnSecurable(callCtx, entityCatalogId, entity.getId())); + new ArrayList<>(ms.loadAllGrantRecordsOnGrantee(entityCatalogId, entity.getId())); + grantRecords.addAll(ms.loadAllGrantRecordsOnSecurable(entityCatalogId, entity.getId())); } else { - grantRecords = ms.loadAllGrantRecordsOnSecurable(callCtx, entityCatalogId, entity.getId()); + grantRecords = ms.loadAllGrantRecordsOnSecurable(entityCatalogId, entity.getId()); } // return the result @@ -2290,27 +2048,21 @@ public Map getInternalPropertyMap( /** {@inheritDoc} */ @Override public @Nonnull CachedEntryResult loadCachedEntryByName( - @Nonnull PolarisCallContext callCtx, + PolarisMetaStoreSession ms, long entityCatalogId, long parentId, @Nonnull PolarisEntityType entityType, @Nonnull String entityName) { - // get metastore we should be using - PolarisMetaStoreSession ms = callCtx.getMetaStore(); - - // need to run inside a read transaction CachedEntryResult result = ms.runInReadTransaction( - callCtx, () -> - this.loadCachedEntryByName( - callCtx, ms, entityCatalogId, parentId, entityType, entityName)); + this.doLoadCachedEntryByName( + ms, entityCatalogId, parentId, entityType, entityName)); if (PolarisEntityConstants.getRootContainerName().equals(entityName) && entityType == PolarisEntityType.ROOT && !result.isSuccess()) { // Backfill rootContainer if needed. ms.runActionInTransaction( - callCtx, () -> { PolarisBaseEntity rootContainer = new PolarisBaseEntity( @@ -2320,8 +2072,7 @@ public Map getInternalPropertyMap( PolarisEntitySubType.NULL_SUBTYPE, PolarisEntityConstants.getRootEntityId(), PolarisEntityConstants.getRootContainerName()); - EntityResult backfillResult = - this.createEntityIfNotExists(callCtx, ms, null, rootContainer); + EntityResult backfillResult = this.createEntityIfNotExists(ms, null, rootContainer); if (backfillResult.isSuccess()) { PolarisEntitiesActiveKey serviceAdminRoleKey = new PolarisEntitiesActiveKey( @@ -2329,15 +2080,10 @@ public Map getInternalPropertyMap( 0L, PolarisEntityType.PRINCIPAL_ROLE.getCode(), PolarisEntityConstants.getNameOfPrincipalServiceAdminRole()); - PolarisBaseEntity serviceAdminRole = - this.lookupEntityByName(callCtx, ms, serviceAdminRoleKey); + PolarisBaseEntity serviceAdminRole = this.lookupEntityByName(ms, serviceAdminRoleKey); if (serviceAdminRole != null) { this.persistNewGrantRecord( - callCtx, - ms, - rootContainer, - serviceAdminRole, - PolarisPrivilege.SERVICE_MANAGE_ACCESS); + ms, rootContainer, serviceAdminRole, PolarisPrivilege.SERVICE_MANAGE_ACCESS); } } }); @@ -2345,17 +2091,15 @@ public Map getInternalPropertyMap( // Redo the lookup in a separate read transaction. result = ms.runInReadTransaction( - callCtx, () -> - this.loadCachedEntryByName( - callCtx, ms, entityCatalogId, parentId, entityType, entityName)); + this.doLoadCachedEntryByName( + ms, entityCatalogId, parentId, entityType, entityName)); } return result; } /** {@inheritDoc} */ - private @Nonnull CachedEntryResult refreshCachedEntity( - @Nonnull PolarisCallContext callCtx, + private @Nonnull CachedEntryResult doRefreshCachedEntity( @Nonnull PolarisMetaStoreSession ms, int entityVersion, int entityGrantRecordsVersion, @@ -2365,8 +2109,7 @@ public Map getInternalPropertyMap( // load version information PolarisChangeTrackingVersions entityVersions = - ms.lookupEntityVersions(callCtx, List.of(new PolarisEntityId(entityCatalogId, entityId))) - .get(0); + ms.lookupEntityVersions(List.of(new PolarisEntityId(entityCatalogId, entityId))).get(0); // if null, the entity has been purged if (entityVersions == null) { @@ -2376,7 +2119,7 @@ public Map getInternalPropertyMap( // load the entity if something changed final PolarisBaseEntity entity; if (entityVersion != entityVersions.getEntityVersion()) { - entity = ms.lookupEntity(callCtx, entityCatalogId, entityId); + entity = ms.lookupEntity(entityCatalogId, entityId); // if not found, return null if (entity == null) { @@ -2391,11 +2134,10 @@ public Map getInternalPropertyMap( final List grantRecords; if (entityVersions.getGrantRecordsVersion() != entityGrantRecordsVersion) { if (entityType.isGrantee()) { - grantRecords = - new ArrayList<>(ms.loadAllGrantRecordsOnGrantee(callCtx, entityCatalogId, entityId)); - grantRecords.addAll(ms.loadAllGrantRecordsOnSecurable(callCtx, entityCatalogId, entityId)); + grantRecords = new ArrayList<>(ms.loadAllGrantRecordsOnGrantee(entityCatalogId, entityId)); + grantRecords.addAll(ms.loadAllGrantRecordsOnSecurable(entityCatalogId, entityId)); } else { - grantRecords = ms.loadAllGrantRecordsOnSecurable(callCtx, entityCatalogId, entityId); + grantRecords = ms.loadAllGrantRecordsOnSecurable(entityCatalogId, entityId); } } else { grantRecords = null; @@ -2408,21 +2150,15 @@ public Map getInternalPropertyMap( /** {@inheritDoc} */ @Override public @Nonnull CachedEntryResult refreshCachedEntity( - @Nonnull PolarisCallContext callCtx, + PolarisMetaStoreSession ms, int entityVersion, int entityGrantRecordsVersion, @Nonnull PolarisEntityType entityType, long entityCatalogId, long entityId) { - // get metastore we should be using - PolarisMetaStoreSession ms = callCtx.getMetaStore(); - - // need to run inside a read transaction return ms.runInReadTransaction( - callCtx, () -> - this.refreshCachedEntity( - callCtx, + this.doRefreshCachedEntity( ms, entityVersion, entityGrantRecordsVersion, diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisMetaStoreSession.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisMetaStoreSession.java index 124ec1d7a..d8884d31e 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisMetaStoreSession.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisMetaStoreSession.java @@ -24,7 +24,6 @@ import java.util.function.Function; import java.util.function.Predicate; import java.util.function.Supplier; -import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.entity.PolarisBaseEntity; import org.apache.polaris.core.entity.PolarisChangeTrackingVersions; import org.apache.polaris.core.entity.PolarisEntitiesActiveKey; @@ -55,10 +54,9 @@ public interface PolarisMetaStoreSession { * error. The result of the supplier lambda is returned if success, else the error will be * re-thrown. * - * @param callCtx call context * @param transactionCode code of the transaction being executed, a supplier lambda */ - T runInTransaction(@Nonnull PolarisCallContext callCtx, @Nonnull Supplier transactionCode); + T runInTransaction(@Nonnull Supplier transactionCode); /** * Run the specified transaction code (a runnable lambda type) in a database read/write @@ -66,11 +64,9 @@ public interface PolarisMetaStoreSession { * the transaction will be committed, else the transaction will be automatically rolled-back on * error. * - * @param callCtx call context * @param transactionCode code of the transaction being executed, a runnable lambda */ - void runActionInTransaction( - @Nonnull PolarisCallContext callCtx, @Nonnull Runnable transactionCode); + void runActionInTransaction(@Nonnull Runnable transactionCode); /** * Run the specified transaction code (a Supplier lambda type) in a database read transaction. If @@ -78,140 +74,115 @@ void runActionInTransaction( * will be committed, else the transaction will be automatically rolled-back on error. The result * of the supplier lambda is returned if success, else the error will be re-thrown. * - * @param callCtx call context * @param transactionCode code of the transaction being executed, a supplier lambda */ - T runInReadTransaction( - @Nonnull PolarisCallContext callCtx, @Nonnull Supplier transactionCode); + T runInReadTransaction(@Nonnull Supplier transactionCode); /** * Run the specified transaction code (a runnable lambda type) in a database read transaction. If * the code of the transaction does not throw any exception and returns normally, the transaction * will be committed, else the transaction will be automatically rolled-back on error. * - * @param callCtx call context * @param transactionCode code of the transaction being executed, a runnable lambda */ - void runActionInReadTransaction( - @Nonnull PolarisCallContext callCtx, @Nonnull Runnable transactionCode); + void runActionInReadTransaction(@Nonnull Runnable transactionCode); /** - * @param callCtx call context * @return new unique entity identifier */ - long generateNewId(@Nonnull PolarisCallContext callCtx); + long generateNewId(); /** * Write the base entity to the entities table. If there is a conflict (existing record with the * same id), all attributes of the new record will replace the existing one. * - * @param callCtx call context * @param entity entity record to write, potentially replacing an existing entity record with the * same key */ - void writeToEntities(@Nonnull PolarisCallContext callCtx, @Nonnull PolarisBaseEntity entity); + void writeToEntities(@Nonnull PolarisBaseEntity entity); /** * Write the base entity to the entities_active table. If there is a conflict (existing record * with the same PK), all attributes of the new record will replace the existing one. * - * @param callCtx call context * @param entity entity record to write, potentially replacing an existing entity record with the * same key */ - void writeToEntitiesActive( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisBaseEntity entity); + void writeToEntitiesActive(@Nonnull PolarisBaseEntity entity); /** * Write the base entity to the entities_dropped table. If there is a conflict (existing record * with the same PK), all attributes of the new record will replace the existing one. * - * @param callCtx call context * @param entity entity record to write, potentially replacing an existing entity record with the * same key */ - void writeToEntitiesDropped( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisBaseEntity entity); + void writeToEntitiesDropped(@Nonnull PolarisBaseEntity entity); /** * Write the base entity to the entities change tracking table. If there is a conflict (existing * record with the same id), all attributes of the new record will replace the existing one. * - * @param callCtx call context * @param entity entity record to write, potentially replacing an existing entity record with the * same key */ - void writeToEntitiesChangeTracking( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisBaseEntity entity); + void writeToEntitiesChangeTracking(@Nonnull PolarisBaseEntity entity); /** * Write the specified grantRecord to the grant_records table. If there is a conflict (existing * record with the same PK), all attributes of the new record will replace the existing one. * - * @param callCtx call context * @param grantRec entity record to write, potentially replacing an existing entity record with * the same key */ - void writeToGrantRecords( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisGrantRecord grantRec); + void writeToGrantRecords(@Nonnull PolarisGrantRecord grantRec); /** * Delete the base entity from the entities table. * - * @param callCtx call context * @param entity entity record to delete */ - void deleteFromEntities(@Nonnull PolarisCallContext callCtx, @Nonnull PolarisEntityCore entity); + void deleteFromEntities(@Nonnull PolarisEntityCore entity); /** * Delete the base entity from the entities_active table. * - * @param callCtx call context * @param entity entity record to delete */ - void deleteFromEntitiesActive( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisEntityCore entity); + void deleteFromEntitiesActive(@Nonnull PolarisEntityCore entity); /** * Delete the base entity to the entities_dropped table * - * @param callCtx call context * @param entity entity record to delete */ - void deleteFromEntitiesDropped( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisBaseEntity entity); + void deleteFromEntitiesDropped(@Nonnull PolarisBaseEntity entity); /** * Delete the base entity from the entities change tracking table * - * @param callCtx call context * @param entity entity record to delete */ - void deleteFromEntitiesChangeTracking( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisEntityCore entity); + void deleteFromEntitiesChangeTracking(@Nonnull PolarisEntityCore entity); /** * Delete the specified grantRecord to the grant_records table. * - * @param callCtx call context * @param grantRec entity record to delete. */ - void deleteFromGrantRecords( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisGrantRecord grantRec); + void deleteFromGrantRecords(@Nonnull PolarisGrantRecord grantRec); /** * Delete the all grant records in the grant_records table for the specified entity. This method * will delete all grant records on that securable entity and also all grants to that grantee * entity assuming that the entity is a grantee (catalog role, principal role or principal). * - * @param callCtx call context * @param entity entity whose grant records to and from should be deleted * @param grantsOnGrantee all grants to that grantee entity. Empty list if that entity is not a * grantee * @param grantsOnSecurable all grants on that securable entity */ void deleteAllEntityGrantRecords( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisEntityCore entity, @Nonnull List grantsOnGrantee, @Nonnull List grantsOnSecurable); @@ -219,71 +190,60 @@ void deleteAllEntityGrantRecords( /** * Delete Polaris entity and grant record metadata from all tables. This is used during metadata * bootstrap to reset all tables to their original state - * - * @param callCtx call context */ - void deleteAll(@Nonnull PolarisCallContext callCtx); + void deleteAll(); /** * Lookup an entity given its catalog id (which can be NULL_ID for top-level entities) and its * unique id. * - * @param callCtx call context * @param catalogId catalog id or NULL_ID * @param entityId unique entity id * @return NULL if the entity was not found, else the base entity. */ @Nullable - PolarisBaseEntity lookupEntity( - @Nonnull PolarisCallContext callCtx, long catalogId, long entityId); + PolarisBaseEntity lookupEntity(long catalogId, long entityId); /** * Lookup a set of entities given their catalog id/entity id unique identifier * - * @param callCtx call context * @param entityIds list of entity ids * @return list of polaris base entities, parallel to the input list of ids. An entity in the list * will be null if the corresponding entity could not be found. */ @Nonnull - List lookupEntities( - @Nonnull PolarisCallContext callCtx, List entityIds); + List lookupEntities(List entityIds); /** * Lookup in the entities_change_tracking table the current version of an entity given its catalog * id (which can be NULL_ID for top-level entities) and its unique id. Will return 0 if the entity * does not exist. * - * @param callCtx call context * @param catalogId catalog id or NULL_ID * @param entityId unique entity id * @return current version for that entity or 0 if entity was not found. */ - int lookupEntityVersion(@Nonnull PolarisCallContext callCtx, long catalogId, long entityId); + int lookupEntityVersion(long catalogId, long entityId); /** * Get change tracking versions for all specified entity ids. * - * @param callCtx call context * @param entityIds list of entity id * @return list parallel to the input list of entity versions. If an entity cannot be found, the * corresponding element in the list will be null */ @Nonnull - List lookupEntityVersions( - @Nonnull PolarisCallContext callCtx, List entityIds); + List lookupEntityVersions(List entityIds); /** * Lookup in the entities_active table to determine if the specified entity exists. Return the * result of that lookup * - * @param callCtx call context * @param entityActiveKey key in the ENTITIES_ACTIVE table * @return null if the specified entity does not exist or has been dropped. */ @Nullable - PolarisEntityActiveRecord lookupEntityActive( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisEntitiesActiveKey entityActiveKey); + PolarisEntityActiveRecord lookupEntityActive(@Nonnull PolarisEntitiesActiveKey entityActiveKey); /** * Lookup in the entities_active table to determine if the specified set of entities exist. Return @@ -294,12 +254,11 @@ PolarisEntityActiveRecord lookupEntityActive( */ @Nonnull List lookupEntityActiveBatch( - @Nonnull PolarisCallContext callCtx, List entityActiveKeys); + List entityActiveKeys); /** * List all active entities of the specified type which are child entities of the specified parent * - * @param callCtx call context * @param catalogId catalog id for that entity, NULL_ID if the entity is top-level * @param parentId id of the parent, can be the special 0 value representing the root entity * @param entityType type of entities to list @@ -307,15 +266,11 @@ List lookupEntityActiveBatch( */ @Nonnull List listActiveEntities( - @Nonnull PolarisCallContext callCtx, - long catalogId, - long parentId, - @Nonnull PolarisEntityType entityType); + long catalogId, long parentId, @Nonnull PolarisEntityType entityType); /** * List active entities where some predicate returns true * - * @param callCtx call context * @param catalogId catalog id for that entity, NULL_ID if the entity is top-level * @param parentId id of the parent, can be the special 0 value representing the root entity * @param entityType type of entities to list @@ -325,7 +280,6 @@ List listActiveEntities( */ @Nonnull List listActiveEntities( - @Nonnull PolarisCallContext callCtx, long catalogId, long parentId, @Nonnull PolarisEntityType entityType, @@ -335,7 +289,6 @@ List listActiveEntities( * List active entities where some predicate returns true and transform the entities with a * function * - * @param callCtx call context * @param catalogId catalog id for that entity, NULL_ID if the entity is top-level * @param parentId id of the parent, can be the special 0 value representing the root entity * @param entityType type of entities to list @@ -348,7 +301,6 @@ List listActiveEntities( */ @Nonnull List listActiveEntities( - @Nonnull PolarisCallContext callCtx, long catalogId, long parentId, @Nonnull PolarisEntityType entityType, @@ -361,18 +313,15 @@ List listActiveEntities( * entity. That version is changed everytime a grant record is added or removed on a base * securable or added to a grantee. * - * @param callCtx call context * @param catalogId catalog id or NULL_ID * @param entityId unique entity id * @return current grant records version for that entity. */ - int lookupEntityGrantRecordsVersion( - @Nonnull PolarisCallContext callCtx, long catalogId, long entityId); + int lookupEntityGrantRecordsVersion(long catalogId, long entityId); /** * Lookup the specified grant record from the grant_records table. Return NULL if not found * - * @param callCtx call context * @param securableCatalogId catalog id of the securable entity, NULL_ID if the entity is * top-level * @param securableId id of the securable entity @@ -383,7 +332,6 @@ int lookupEntityGrantRecordsVersion( */ @Nullable PolarisGrantRecord lookupGrantRecord( - @Nonnull PolarisCallContext callCtx, long securableCatalogId, long securableId, long granteeCatalogId, @@ -393,7 +341,6 @@ PolarisGrantRecord lookupGrantRecord( /** * Get all grant records on the specified securable entity. * - * @param callCtx call context * @param securableCatalogId catalog id of the securable entity, NULL_ID if the entity is * top-level * @param securableId id of the securable entity @@ -401,47 +348,41 @@ PolarisGrantRecord lookupGrantRecord( */ @Nonnull List loadAllGrantRecordsOnSecurable( - @Nonnull PolarisCallContext callCtx, long securableCatalogId, long securableId); + long securableCatalogId, long securableId); /** * Get all grant records granted to the specified grantee entity. * - * @param callCtx call context * @param granteeCatalogId catalog id of the grantee entity, NULL_ID if the entity is top-level * @param granteeId id of the grantee entity * @return the list of grant records for the specified grantee */ @Nonnull - List loadAllGrantRecordsOnGrantee( - @Nonnull PolarisCallContext callCtx, long granteeCatalogId, long granteeId); + List loadAllGrantRecordsOnGrantee(long granteeCatalogId, long granteeId); /** * Allows to retrieve to the secrets of a principal given its unique client id * - * @param callCtx call context * @param clientId principal client id * @return the secrets */ @Nullable - PolarisPrincipalSecrets loadPrincipalSecrets( - @Nonnull PolarisCallContext callCtx, @Nonnull String clientId); + PolarisPrincipalSecrets loadPrincipalSecrets(@Nonnull String clientId); /** * generate and store a client id and associated secrets for a newly created principal entity * - * @param callCtx call context * @param principalName name of the principal * @param principalId principal id */ @Nonnull PolarisPrincipalSecrets generateNewPrincipalSecrets( - @Nonnull PolarisCallContext callCtx, @Nonnull String principalName, long principalId); + @Nonnull String principalName, long principalId); /** * Rotate the secrets of a principal entity, i.e. make the specified main secrets the secondary * and generate a new main secret * - * @param callCtx call context * @param clientId principal client id * @param principalId principal id * @param reset true if the principal secrets should be disabled and replaced with a one-time @@ -450,26 +391,19 @@ PolarisPrincipalSecrets generateNewPrincipalSecrets( */ @Nullable PolarisPrincipalSecrets rotatePrincipalSecrets( - @Nonnull PolarisCallContext callCtx, - @Nonnull String clientId, - long principalId, - boolean reset, - @Nonnull String oldSecretHash); + @Nonnull String clientId, long principalId, boolean reset, @Nonnull String oldSecretHash); /** * When dropping a principal, we also need to drop the secrets of that principal * - * @param callCtx the call context * @param clientId principal client id * @param principalId the id of the principal whose secrets are dropped */ - void deletePrincipalSecrets( - @Nonnull PolarisCallContext callCtx, @Nonnull String clientId, long principalId); + void deletePrincipalSecrets(@Nonnull String clientId, long principalId); /** * Create an in-memory storage integration * - * @param callCtx the polaris calllctx * @param catalogId the catalog id * @param entityId the entity id * @param polarisStorageConfigurationInfo the storage configuration information @@ -477,7 +411,6 @@ void deletePrincipalSecrets( */ @Nullable PolarisStorageIntegration createStorageIntegration( - @Nonnull PolarisCallContext callCtx, long catalogId, long entityId, PolarisStorageConfigurationInfo polarisStorageConfigurationInfo); @@ -485,31 +418,25 @@ PolarisStorageIntegration createS /** * Persist a storage integration in the metastore * - * @param callContext the polaris call context * @param entity the entity of the object * @param storageIntegration the storage integration to persist */ void persistStorageIntegrationIfNeeded( - @Nonnull PolarisCallContext callContext, - @Nonnull PolarisBaseEntity entity, - @Nullable PolarisStorageIntegration storageIntegration); + @Nonnull PolarisBaseEntity entity, @Nullable PolarisStorageIntegration storageIntegration); /** * Load the polaris storage integration for a polaris entity (Catalog,Namespace,Table,View) * - * @param callContext the polaris call context * @param entity the polaris entity * @return a polaris storage integration */ @Nullable - PolarisStorageIntegration loadPolarisStorageIntegration( - @Nonnull PolarisCallContext callContext, @Nonnull PolarisBaseEntity entity); + PolarisStorageIntegration loadPolarisStorageIntegration(@Nonnull PolarisBaseEntity entity); /** * Check if the specified parent entity has children. * - * @param callContext the polaris call context * @param optionalEntityType if not null, only check for the specified type, else check for all * types of children entities * @param catalogId id of the catalog @@ -517,10 +444,7 @@ PolarisStorageIntegration loadPolarisStorageIntegration( * @return true if the parent entity has children */ boolean hasChildren( - @Nonnull PolarisCallContext callContext, - @Nullable PolarisEntityType optionalEntityType, - long catalogId, - long parentId); + @Nullable PolarisEntityType optionalEntityType, long catalogId, long parentId); /** Rollback the current transaction */ void rollback(); diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisObjectMapperUtil.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisObjectMapperUtil.java index a91d2bbab..39125dc68 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisObjectMapperUtil.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisObjectMapperUtil.java @@ -30,7 +30,7 @@ import java.io.IOException; import java.util.Map; import org.apache.iceberg.rest.RESTSerializers; -import org.apache.polaris.core.PolarisCallContext; +import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.entity.PolarisBaseEntity; import org.apache.polaris.core.entity.PolarisTaskConstants; import org.slf4j.Logger; @@ -57,33 +57,33 @@ private static ObjectMapper configureMapper() { * @return a String, the JSON representation of the map */ public static String serializeProperties( - PolarisCallContext callCtx, Map properties) { + PolarisDiagnostics diagnostics, Map properties) { String jsonString = null; try { // Deserialize the JSON string to a Map jsonString = MAPPER.writeValueAsString(properties); } catch (JsonProcessingException ex) { - callCtx.getDiagServices().fail("got_json_processing_exception", ex.getMessage()); + diagnostics.fail("got_json_processing_exception", ex.getMessage()); } return jsonString; } - public static String serialize(PolarisCallContext callCtx, Object object) { + public static String serialize(PolarisDiagnostics diagnostics, Object object) { try { return MAPPER.writeValueAsString(object); } catch (JsonProcessingException e) { - callCtx.getDiagServices().fail("got_json_processing_exception", e.getMessage()); + diagnostics.fail("got_json_processing_exception", e.getMessage()); } return ""; } - public static T deserialize(PolarisCallContext callCtx, String text, Class klass) { + public static T deserialize(PolarisDiagnostics diagnostics, String text, Class klass) { try { return MAPPER.readValue(text, klass); } catch (JsonProcessingException e) { - callCtx.getDiagServices().fail("got_json_processing_exception", e.getMessage()); + diagnostics.fail("got_json_processing_exception", e.getMessage()); } return null; } @@ -95,20 +95,16 @@ public static T deserialize(PolarisCallContext callCtx, String text, Class deserializeProperties( - PolarisCallContext callCtx, String properties) { + PolarisDiagnostics diagnostics, String properties) { Map retProperties = null; try { // Deserialize the JSON string to a Map retProperties = MAPPER.readValue(properties, new TypeReference<>() {}); } catch (JsonMappingException ex) { - callCtx - .getDiagServices() - .fail("got_json_mapping_exception", "properties={}, ex={}", properties, ex); + diagnostics.fail("got_json_mapping_exception", "properties={}, ex={}", properties, ex); } catch (JsonProcessingException ex) { - callCtx - .getDiagServices() - .fail("got_json_processing_exception", "properties={}, ex={}", properties, ex); + diagnostics.fail("got_json_processing_exception", "properties={}, ex={}", properties, ex); } return retProperties; diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisTreeMapMetaStoreSessionImpl.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisTreeMapMetaStoreSessionImpl.java index bea205415..d999e662e 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisTreeMapMetaStoreSessionImpl.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisTreeMapMetaStoreSessionImpl.java @@ -26,7 +26,7 @@ import java.util.function.Predicate; import java.util.function.Supplier; import java.util.stream.Collectors; -import org.apache.polaris.core.PolarisCallContext; +import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.entity.PolarisBaseEntity; import org.apache.polaris.core.entity.PolarisChangeTrackingVersions; import org.apache.polaris.core.entity.PolarisEntitiesActiveKey; @@ -46,65 +46,63 @@ public class PolarisTreeMapMetaStoreSessionImpl implements PolarisMetaStoreSessi private final PolarisTreeMapStore store; private final PolarisStorageIntegrationProvider storageIntegrationProvider; private final PrincipalSecretsGenerator secretsGenerator; + private final PolarisDiagnostics diagnostics; public PolarisTreeMapMetaStoreSessionImpl( @Nonnull PolarisTreeMapStore store, @Nonnull PolarisStorageIntegrationProvider storageIntegrationProvider, - @Nonnull PrincipalSecretsGenerator secretsGenerator) { + @Nonnull PrincipalSecretsGenerator secretsGenerator, + @Nonnull PolarisDiagnostics diagnostics) { // init store this.store = store; this.storageIntegrationProvider = storageIntegrationProvider; this.secretsGenerator = secretsGenerator; + this.diagnostics = diagnostics; } /** {@inheritDoc} */ @Override - public T runInTransaction( - @Nonnull PolarisCallContext callCtx, @Nonnull Supplier transactionCode) { + public T runInTransaction(@Nonnull Supplier transactionCode) { // run transaction on our underlying store - return store.runInTransaction(callCtx, transactionCode); + return store.runInTransaction(transactionCode); } /** {@inheritDoc} */ @Override - public void runActionInTransaction( - @Nonnull PolarisCallContext callCtx, @Nonnull Runnable transactionCode) { + public void runActionInTransaction(@Nonnull Runnable transactionCode) { // run transaction on our underlying store - store.runActionInTransaction(callCtx, transactionCode); + store.runActionInTransaction(transactionCode); } /** {@inheritDoc} */ @Override - public T runInReadTransaction( - @Nonnull PolarisCallContext callCtx, @Nonnull Supplier transactionCode) { + public T runInReadTransaction(@Nonnull Supplier transactionCode) { // run transaction on our underlying store - return store.runInReadTransaction(callCtx, transactionCode); + return store.runInReadTransaction(transactionCode); } /** {@inheritDoc} */ @Override - public void runActionInReadTransaction( - @Nonnull PolarisCallContext callCtx, @Nonnull Runnable transactionCode) { + public void runActionInReadTransaction(@Nonnull Runnable transactionCode) { // run transaction on our underlying store - store.runActionInReadTransaction(callCtx, transactionCode); + store.runActionInReadTransaction(transactionCode); } /** * @return new unique entity identifier */ @Override - public long generateNewId(@Nonnull PolarisCallContext callCtx) { + public long generateNewId() { return this.store.getNextSequence(); } /** {@inheritDoc} */ @Override - public void writeToEntities( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisBaseEntity entity) { + public void writeToEntities(@Nonnull PolarisBaseEntity entity) { // write it this.store.getSliceEntities().write(entity); } @@ -112,7 +110,6 @@ public void writeToEntities( /** {@inheritDoc} */ @Override public void persistStorageIntegrationIfNeeded( - @Nonnull PolarisCallContext callContext, @Nonnull PolarisBaseEntity entity, @Nullable PolarisStorageIntegration storageIntegration) { // not implemented for in-memory store @@ -120,16 +117,14 @@ public void persistStorageIntegratio /** {@inheritDoc} */ @Override - public void writeToEntitiesActive( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisBaseEntity entity) { + public void writeToEntitiesActive(@Nonnull PolarisBaseEntity entity) { // write it this.store.getSliceEntitiesActive().write(entity); } /** {@inheritDoc} */ @Override - public void writeToEntitiesDropped( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisBaseEntity entity) { + public void writeToEntitiesDropped(@Nonnull PolarisBaseEntity entity) { // write it this.store.getSliceEntitiesDropped().write(entity); this.store.getSliceEntitiesDroppedToPurge().write(entity); @@ -137,16 +132,14 @@ public void writeToEntitiesDropped( /** {@inheritDoc} */ @Override - public void writeToEntitiesChangeTracking( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisBaseEntity entity) { + public void writeToEntitiesChangeTracking(@Nonnull PolarisBaseEntity entity) { // write it this.store.getSliceEntitiesChangeTracking().write(entity); } /** {@inheritDoc} */ @Override - public void writeToGrantRecords( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisGrantRecord grantRec) { + public void writeToGrantRecords(@Nonnull PolarisGrantRecord grantRec) { // write it this.store.getSliceGrantRecords().write(grantRec); this.store.getSliceGrantRecordsByGrantee().write(grantRec); @@ -154,8 +147,7 @@ public void writeToGrantRecords( /** {@inheritDoc} */ @Override - public void deleteFromEntities( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisEntityCore entity) { + public void deleteFromEntities(@Nonnull PolarisEntityCore entity) { // delete it this.store.getSliceEntities().delete(this.store.buildEntitiesKey(entity)); @@ -163,16 +155,14 @@ public void deleteFromEntities( /** {@inheritDoc} */ @Override - public void deleteFromEntitiesActive( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisEntityCore entity) { + public void deleteFromEntitiesActive(@Nonnull PolarisEntityCore entity) { // delete it this.store.getSliceEntitiesActive().delete(this.store.buildEntitiesActiveKey(entity)); } /** {@inheritDoc} */ @Override - public void deleteFromEntitiesDropped( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisBaseEntity entity) { + public void deleteFromEntitiesDropped(@Nonnull PolarisBaseEntity entity) { // delete it this.store.getSliceEntitiesDropped().delete(entity); this.store.getSliceEntitiesDroppedToPurge().delete(entity); @@ -181,20 +171,17 @@ public void deleteFromEntitiesDropped( /** * {@inheritDoc} * - * @param callCtx * @param entity entity record to delete */ @Override - public void deleteFromEntitiesChangeTracking( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisEntityCore entity) { + public void deleteFromEntitiesChangeTracking(@Nonnull PolarisEntityCore entity) { // delete it this.store.getSliceEntitiesChangeTracking().delete(this.store.buildEntitiesKey(entity)); } /** {@inheritDoc} */ @Override - public void deleteFromGrantRecords( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisGrantRecord grantRec) { + public void deleteFromGrantRecords(@Nonnull PolarisGrantRecord grantRec) { // delete it this.store.getSliceGrantRecords().delete(grantRec); @@ -204,7 +191,6 @@ public void deleteFromGrantRecords( /** {@inheritDoc} */ @Override public void deleteAllEntityGrantRecords( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisEntityCore entity, @Nonnull List grantsOnGrantee, @Nonnull List grantsOnSecurable) { @@ -222,22 +208,20 @@ public void deleteAllEntityGrantRecords( /** {@inheritDoc} */ @Override - public void deleteAll(@Nonnull PolarisCallContext callCtx) { + public void deleteAll() { // clear all slices this.store.deleteAll(); } /** {@inheritDoc} */ @Override - public @Nullable PolarisBaseEntity lookupEntity( - @Nonnull PolarisCallContext callCtx, long catalogId, long entityId) { + public @Nullable PolarisBaseEntity lookupEntity(long catalogId, long entityId) { return this.store.getSliceEntities().read(this.store.buildKeyComposite(catalogId, entityId)); } /** {@inheritDoc} */ @Override - public @Nonnull List lookupEntities( - @Nonnull PolarisCallContext callCtx, List entityIds) { + public @Nonnull List lookupEntities(List entityIds) { // allocate return list return entityIds.stream() .map( @@ -250,8 +234,7 @@ public void deleteAll(@Nonnull PolarisCallContext callCtx) { /** {@inheritDoc} */ @Override - public int lookupEntityVersion( - @Nonnull PolarisCallContext callCtx, long catalogId, long entityId) { + public int lookupEntityVersion(long catalogId, long entityId) { PolarisBaseEntity baseEntity = this.store .getSliceEntitiesChangeTracking() @@ -263,7 +246,7 @@ public int lookupEntityVersion( /** {@inheritDoc} */ @Override public @Nonnull List lookupEntityVersions( - @Nonnull PolarisCallContext callCtx, List entityIds) { + List entityIds) { // allocate return list return entityIds.stream() .map( @@ -284,7 +267,7 @@ public int lookupEntityVersion( @Override @Nullable public PolarisEntityActiveRecord lookupEntityActive( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisEntitiesActiveKey entityActiveKey) { + @Nonnull PolarisEntitiesActiveKey entityActiveKey) { // lookup the active entity slice PolarisBaseEntity entity = this.store @@ -312,34 +295,28 @@ public PolarisEntityActiveRecord lookupEntityActive( @Override @Nonnull public List lookupEntityActiveBatch( - @Nonnull PolarisCallContext callCtx, @Nonnull List entityActiveKeys) { // now build a list to quickly verify that nothing has changed return entityActiveKeys.stream() - .map(entityActiveKey -> this.lookupEntityActive(callCtx, entityActiveKey)) + .map(entityActiveKey -> this.lookupEntityActive(entityActiveKey)) .collect(Collectors.toList()); } /** {@inheritDoc} */ @Override public @Nonnull List listActiveEntities( - @Nonnull PolarisCallContext callCtx, - long catalogId, - long parentId, - @Nonnull PolarisEntityType entityType) { - return listActiveEntities(callCtx, catalogId, parentId, entityType, Predicates.alwaysTrue()); + long catalogId, long parentId, @Nonnull PolarisEntityType entityType) { + return listActiveEntities(catalogId, parentId, entityType, Predicates.alwaysTrue()); } @Override public @Nonnull List listActiveEntities( - @Nonnull PolarisCallContext callCtx, long catalogId, long parentId, @Nonnull PolarisEntityType entityType, @Nonnull Predicate entityFilter) { // full range scan under the parent for that type return listActiveEntities( - callCtx, catalogId, parentId, entityType, @@ -357,7 +334,6 @@ public List lookupEntityActiveBatch( @Override public @Nonnull List listActiveEntities( - @Nonnull PolarisCallContext callCtx, long catalogId, long parentId, @Nonnull PolarisEntityType entityType, @@ -378,10 +354,7 @@ public List lookupEntityActiveBatch( /** {@inheritDoc} */ @Override public boolean hasChildren( - @Nonnull PolarisCallContext callContext, - @Nullable PolarisEntityType entityType, - long catalogId, - long parentId) { + @Nullable PolarisEntityType entityType, long catalogId, long parentId) { // determine key prefix, add type if one is passed-in String prefixKey = entityType == null @@ -393,8 +366,7 @@ public boolean hasChildren( /** {@inheritDoc} */ @Override - public int lookupEntityGrantRecordsVersion( - @Nonnull PolarisCallContext callCtx, long catalogId, long entityId) { + public int lookupEntityGrantRecordsVersion(long catalogId, long entityId) { PolarisBaseEntity entity = this.store .getSliceEntitiesChangeTracking() @@ -407,7 +379,6 @@ public int lookupEntityGrantRecordsVersion( /** {@inheritDoc} */ @Override public @Nullable PolarisGrantRecord lookupGrantRecord( - @Nonnull PolarisCallContext callCtx, long securableCatalogId, long securableId, long granteeCatalogId, @@ -424,7 +395,7 @@ public int lookupEntityGrantRecordsVersion( /** {@inheritDoc} */ @Override public @Nonnull List loadAllGrantRecordsOnSecurable( - @Nonnull PolarisCallContext callCtx, long securableCatalogId, long securableId) { + long securableCatalogId, long securableId) { // now fetch all grants for this securable return this.store .getSliceGrantRecords() @@ -434,7 +405,7 @@ public int lookupEntityGrantRecordsVersion( /** {@inheritDoc} */ @Override public @Nonnull List loadAllGrantRecordsOnGrantee( - @Nonnull PolarisCallContext callCtx, long granteeCatalogId, long granteeId) { + long granteeCatalogId, long granteeId) { // now fetch all grants assigned to this grantee return this.store .getSliceGrantRecordsByGrantee() @@ -443,15 +414,14 @@ public int lookupEntityGrantRecordsVersion( /** {@inheritDoc} */ @Override - public @Nullable PolarisPrincipalSecrets loadPrincipalSecrets( - @Nonnull PolarisCallContext callCtx, @Nonnull String clientId) { + public @Nullable PolarisPrincipalSecrets loadPrincipalSecrets(@Nonnull String clientId) { return this.store.getSlicePrincipalSecrets().read(clientId); } /** {@inheritDoc} */ @Override public @Nonnull PolarisPrincipalSecrets generateNewPrincipalSecrets( - @Nonnull PolarisCallContext callCtx, @Nonnull String principalName, long principalId) { + @Nonnull String principalName, long principalId) { // ensure principal client id is unique PolarisPrincipalSecrets principalSecrets; PolarisPrincipalSecrets lookupPrincipalSecrets; @@ -474,34 +444,26 @@ public int lookupEntityGrantRecordsVersion( /** {@inheritDoc} */ @Override public @Nonnull PolarisPrincipalSecrets rotatePrincipalSecrets( - @Nonnull PolarisCallContext callCtx, - @Nonnull String clientId, - long principalId, - boolean reset, - @Nonnull String oldSecretHash) { + @Nonnull String clientId, long principalId, boolean reset, @Nonnull String oldSecretHash) { // load the existing secrets PolarisPrincipalSecrets principalSecrets = this.store.getSlicePrincipalSecrets().read(clientId); // should be found - callCtx - .getDiagServices() - .checkNotNull( - principalSecrets, - "cannot_find_secrets", - "client_id={} principalId={}", - clientId, - principalId); + diagnostics.checkNotNull( + principalSecrets, + "cannot_find_secrets", + "client_id={} principalId={}", + clientId, + principalId); // ensure principal id is matching - callCtx - .getDiagServices() - .check( - principalId == principalSecrets.getPrincipalId(), - "principal_id_mismatch", - "expectedId={} id={}", - principalId, - principalSecrets.getPrincipalId()); + diagnostics.check( + principalId == principalSecrets.getPrincipalId(), + "principal_id_mismatch", + "expectedId={} id={}", + principalId, + principalSecrets.getPrincipalId()); // rotate the secrets principalSecrets.rotateSecrets(oldSecretHash); @@ -518,30 +480,25 @@ public int lookupEntityGrantRecordsVersion( /** {@inheritDoc} */ @Override - public void deletePrincipalSecrets( - @Nonnull PolarisCallContext callCtx, @Nonnull String clientId, long principalId) { + public void deletePrincipalSecrets(@Nonnull String clientId, long principalId) { // load the existing secrets PolarisPrincipalSecrets principalSecrets = this.store.getSlicePrincipalSecrets().read(clientId); // should be found - callCtx - .getDiagServices() - .checkNotNull( - principalSecrets, - "cannot_find_secrets", - "client_id={} principalId={}", - clientId, - principalId); + diagnostics.checkNotNull( + principalSecrets, + "cannot_find_secrets", + "client_id={} principalId={}", + clientId, + principalId); // ensure principal id is matching - callCtx - .getDiagServices() - .check( - principalId == principalSecrets.getPrincipalId(), - "principal_id_mismatch", - "expectedId={} id={}", - principalId, - principalSecrets.getPrincipalId()); + diagnostics.check( + principalId == principalSecrets.getPrincipalId(), + "principal_id_mismatch", + "expectedId={} id={}", + principalId, + principalSecrets.getPrincipalId()); // delete these secrets this.store.getSlicePrincipalSecrets().delete(clientId); @@ -551,7 +508,6 @@ public void deletePrincipalSecrets( @Override public @Nullable PolarisStorageIntegration createStorageIntegration( - @Nonnull PolarisCallContext callCtx, long catalogId, long entityId, PolarisStorageConfigurationInfo polarisStorageConfigurationInfo) { @@ -563,9 +519,9 @@ PolarisStorageIntegration createStorageIntegration( @Override public @Nullable PolarisStorageIntegration loadPolarisStorageIntegration( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisBaseEntity entity) { + @Nonnull PolarisBaseEntity entity) { PolarisStorageConfigurationInfo storageConfig = - PolarisMetaStoreManagerImpl.readStorageConfiguration(callCtx, entity); + PolarisMetaStoreManagerImpl.readStorageConfiguration(diagnostics, entity); return storageIntegrationProvider.getStorageIntegrationForConfig(storageConfig); } diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisTreeMapStore.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisTreeMapStore.java index 544bcf0ff..e6710e87a 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisTreeMapStore.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/PolarisTreeMapStore.java @@ -25,7 +25,6 @@ import java.util.concurrent.atomic.AtomicLong; import java.util.function.Function; import java.util.function.Supplier; -import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.entity.PolarisBaseEntity; import org.apache.polaris.core.entity.PolarisEntityCore; @@ -189,7 +188,7 @@ public boolean isWrite() { private Transaction tr; // diagnostic services - private PolarisDiagnostics diagnosticServices; + private final PolarisDiagnostics diagnosticServices; // all entities private final Slice sliceEntities; @@ -404,18 +403,15 @@ private void ensureReadTr() { /** * Run inside a read/write transaction * - * @param callCtx call context to use * @param transactionCode transaction code * @return the result of the execution */ - public T runInTransaction( - @Nonnull PolarisCallContext callCtx, @Nonnull Supplier transactionCode) { + public T runInTransaction(@Nonnull Supplier transactionCode) { synchronized (lock) { // execute transaction try { // init diagnostic services - this.diagnosticServices = callCtx.getDiagServices(); this.startWriteTransaction(); return transactionCode.get(); } catch (Throwable e) { @@ -423,7 +419,6 @@ public T runInTransaction( throw e; } finally { this.tr = null; - this.diagnosticServices = null; } } } @@ -431,18 +426,15 @@ public T runInTransaction( /** * Run inside a read/write transaction * - * @param callCtx call context to use * @param transactionCode transaction code */ - public void runActionInTransaction( - @Nonnull PolarisCallContext callCtx, @Nonnull Runnable transactionCode) { + public void runActionInTransaction(@Nonnull Runnable transactionCode) { synchronized (lock) { // execute transaction try { // init diagnostic services - this.diagnosticServices = callCtx.getDiagServices(); this.startWriteTransaction(); transactionCode.run(); } catch (Throwable e) { @@ -450,7 +442,6 @@ public void runActionInTransaction( throw e; } finally { this.tr = null; - this.diagnosticServices = null; } } } @@ -458,23 +449,19 @@ public void runActionInTransaction( /** * Run inside a read only transaction * - * @param callCtx call context to use * @param transactionCode transaction code * @return the result of the execution */ - public T runInReadTransaction( - @Nonnull PolarisCallContext callCtx, @Nonnull Supplier transactionCode) { + public T runInReadTransaction(@Nonnull Supplier transactionCode) { synchronized (lock) { // execute transaction try { // init diagnostic services - this.diagnosticServices = callCtx.getDiagServices(); this.startReadTransaction(); return transactionCode.get(); } finally { this.tr = null; - this.diagnosticServices = null; } } } @@ -482,22 +469,18 @@ public T runInReadTransaction( /** * Run inside a read only transaction * - * @param callCtx call context to use * @param transactionCode transaction code */ - public void runActionInReadTransaction( - @Nonnull PolarisCallContext callCtx, @Nonnull Runnable transactionCode) { + public void runActionInReadTransaction(@Nonnull Runnable transactionCode) { synchronized (lock) { // execute transaction try { // init diagnostic services - this.diagnosticServices = callCtx.getDiagServices(); this.startReadTransaction(); transactionCode.run(); } finally { this.tr = null; - this.diagnosticServices = null; } } } diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/TransactionWorkspaceMetaStoreManager.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/TransactionWorkspaceMetaStoreManager.java index 44134751b..fe1ea5b70 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/TransactionWorkspaceMetaStoreManager.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/TransactionWorkspaceMetaStoreManager.java @@ -25,7 +25,7 @@ import java.util.List; import java.util.Map; import java.util.Set; -import org.apache.polaris.core.PolarisCallContext; +import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.entity.PolarisBaseEntity; import org.apache.polaris.core.entity.PolarisEntity; import org.apache.polaris.core.entity.PolarisEntityCore; @@ -61,9 +61,12 @@ public class TransactionWorkspaceMetaStoreManager implements PolarisMetaStoreMan // to serve reads within the same transaction while also storing the ordered list of // pendingUpdates that ultimately need to be applied in order within the real MetaStoreManager. private final List pendingUpdates = new ArrayList<>(); + private final PolarisDiagnostics diagnostics; - public TransactionWorkspaceMetaStoreManager(PolarisMetaStoreManager delegate) { + public TransactionWorkspaceMetaStoreManager( + PolarisMetaStoreManager delegate, PolarisDiagnostics diagnostics) { this.delegate = delegate; + this.diagnostics = diagnostics; } public List getPendingUpdates() { @@ -71,111 +74,99 @@ public List getPendingUpdates() { } @Override - public BaseResult bootstrapPolarisService(@Nonnull PolarisCallContext callCtx) { - callCtx - .getDiagServices() - .fail("illegal_method_in_transaction_workspace", "bootstrapPolarisService"); + public @Nonnull BaseResult bootstrapPolarisService(@Nonnull PolarisMetaStoreSession session) { + diagnostics.fail("illegal_method_in_transaction_workspace", "bootstrapPolarisService"); return null; } @Override - public BaseResult purge(@Nonnull PolarisCallContext callCtx) { - callCtx.getDiagServices().fail("illegal_method_in_transaction_workspace", "purge"); + public @Nonnull BaseResult purge(@Nonnull PolarisMetaStoreSession session) { + diagnostics.fail("illegal_method_in_transaction_workspace", "purge"); return null; } @Override - public PolarisMetaStoreManager.EntityResult readEntityByName( - @Nonnull PolarisCallContext callCtx, + public @Nonnull PolarisMetaStoreManager.EntityResult readEntityByName( + @Nonnull PolarisMetaStoreSession ms, @Nullable List catalogPath, @Nonnull PolarisEntityType entityType, @Nonnull PolarisEntitySubType entitySubType, @Nonnull String name) { - callCtx.getDiagServices().fail("illegal_method_in_transaction_workspace", "readEntityByName"); + diagnostics.fail("illegal_method_in_transaction_workspace", "readEntityByName"); return null; } @Override - public ListEntitiesResult listEntities( - @Nonnull PolarisCallContext callCtx, + public @Nonnull ListEntitiesResult listEntities( + @Nonnull PolarisMetaStoreSession ms, @Nullable List catalogPath, @Nonnull PolarisEntityType entityType, @Nonnull PolarisEntitySubType entitySubType) { - callCtx.getDiagServices().fail("illegal_method_in_transaction_workspace", "listEntities"); + diagnostics.fail("illegal_method_in_transaction_workspace", "listEntities"); return null; } @Override - public GenerateEntityIdResult generateNewEntityId(@Nonnull PolarisCallContext callCtx) { - callCtx - .getDiagServices() - .fail("illegal_method_in_transaction_workspace", "generateNewEntityId"); + public @Nonnull GenerateEntityIdResult generateNewEntityId(@Nonnull PolarisMetaStoreSession ms) { + diagnostics.fail("illegal_method_in_transaction_workspace", "generateNewEntityId"); return null; } @Override - public CreatePrincipalResult createPrincipal( - @Nonnull PolarisCallContext callCtx, @Nonnull PolarisBaseEntity principal) { - callCtx.getDiagServices().fail("illegal_method_in_transaction_workspace", "createPrincipal"); + public @Nonnull CreatePrincipalResult createPrincipal( + @Nonnull PolarisMetaStoreSession ms, @Nonnull PolarisBaseEntity principal) { + diagnostics.fail("illegal_method_in_transaction_workspace", "createPrincipal"); return null; } @Override - public PrincipalSecretsResult loadPrincipalSecrets( - @Nonnull PolarisCallContext callCtx, @Nonnull String clientId) { - callCtx - .getDiagServices() - .fail("illegal_method_in_transaction_workspace", "loadPrincipalSecrets"); + public @Nonnull PrincipalSecretsResult loadPrincipalSecrets( + @Nonnull PolarisMetaStoreSession ms, @Nonnull String clientId) { + diagnostics.fail("illegal_method_in_transaction_workspace", "loadPrincipalSecrets"); return null; } @Override - public PrincipalSecretsResult rotatePrincipalSecrets( - @Nonnull PolarisCallContext callCtx, + public @Nonnull PrincipalSecretsResult rotatePrincipalSecrets( + @Nonnull PolarisMetaStoreSession ms, @Nonnull String clientId, long principalId, boolean reset, @Nonnull String oldSecretHash) { - callCtx - .getDiagServices() - .fail("illegal_method_in_transaction_workspace", "rotatePrincipalSecrets"); + diagnostics.fail("illegal_method_in_transaction_workspace", "rotatePrincipalSecrets"); return null; } @Override - public CreateCatalogResult createCatalog( - @Nonnull PolarisCallContext callCtx, + public @Nonnull CreateCatalogResult createCatalog( + @Nonnull PolarisMetaStoreSession ms, @Nonnull PolarisBaseEntity catalog, @Nonnull List principalRoles) { - callCtx.getDiagServices().fail("illegal_method_in_transaction_workspace", "createCatalog"); + diagnostics.fail("illegal_method_in_transaction_workspace", "createCatalog"); return null; } @Override - public EntityResult createEntityIfNotExists( - @Nonnull PolarisCallContext callCtx, + public @Nonnull EntityResult createEntityIfNotExists( + @Nonnull PolarisMetaStoreSession ms, @Nullable List catalogPath, @Nonnull PolarisBaseEntity entity) { - callCtx - .getDiagServices() - .fail("illegal_method_in_transaction_workspace", "createEntityIfNotExists"); + diagnostics.fail("illegal_method_in_transaction_workspace", "createEntityIfNotExists"); return null; } @Override - public EntitiesResult createEntitiesIfNotExist( - @Nonnull PolarisCallContext callCtx, + public @Nonnull EntitiesResult createEntitiesIfNotExist( + @Nonnull PolarisMetaStoreSession ms, @Nullable List catalogPath, @Nonnull List entities) { - callCtx - .getDiagServices() - .fail("illegal_method_in_transaction_workspace", "createEntitiesIfNotExist"); + diagnostics.fail("illegal_method_in_transaction_workspace", "createEntitiesIfNotExist"); return null; } @Override - public EntityResult updateEntityPropertiesIfNotChanged( - @Nonnull PolarisCallContext callCtx, + public @Nonnull EntityResult updateEntityPropertiesIfNotChanged( + @Nonnull PolarisMetaStoreSession ms, @Nullable List catalogPath, @Nonnull PolarisBaseEntity entity) { pendingUpdates.add(new EntityWithPath(catalogPath, entity)); @@ -183,137 +174,125 @@ public EntityResult updateEntityPropertiesIfNotChanged( } @Override - public EntitiesResult updateEntitiesPropertiesIfNotChanged( - @Nonnull PolarisCallContext callCtx, @Nonnull List entities) { - callCtx - .getDiagServices() - .fail("illegal_method_in_transaction_workspace", "updateEntitiesPropertiesIfNotChanged"); + public @Nonnull EntitiesResult updateEntitiesPropertiesIfNotChanged( + @Nonnull PolarisMetaStoreSession ms, @Nonnull List entities) { + diagnostics.fail( + "illegal_method_in_transaction_workspace", "updateEntitiesPropertiesIfNotChanged"); return null; } @Override - public EntityResult renameEntity( - @Nonnull PolarisCallContext callCtx, + public @Nonnull EntityResult renameEntity( + @Nonnull PolarisMetaStoreSession ms, @Nullable List catalogPath, @Nonnull PolarisEntityCore entityToRename, @Nullable List newCatalogPath, @Nonnull PolarisEntity renamedEntity) { - callCtx.getDiagServices().fail("illegal_method_in_transaction_workspace", "renameEntity"); + diagnostics.fail("illegal_method_in_transaction_workspace", "renameEntity"); return null; } @Override - public DropEntityResult dropEntityIfExists( - @Nonnull PolarisCallContext callCtx, + public @Nonnull DropEntityResult dropEntityIfExists( + @Nonnull PolarisMetaStoreSession ms, @Nullable List catalogPath, @Nonnull PolarisEntityCore entityToDrop, @Nullable Map cleanupProperties, boolean cleanup) { - callCtx.getDiagServices().fail("illegal_method_in_transaction_workspace", "dropEntityIfExists"); + diagnostics.fail("illegal_method_in_transaction_workspace", "dropEntityIfExists"); return null; } @Override - public PrivilegeResult grantUsageOnRoleToGrantee( - @Nonnull PolarisCallContext callCtx, + public @Nonnull PrivilegeResult grantUsageOnRoleToGrantee( + @Nonnull PolarisMetaStoreSession ms, @Nullable PolarisEntityCore catalog, @Nonnull PolarisEntityCore role, @Nonnull PolarisEntityCore grantee) { - callCtx - .getDiagServices() - .fail("illegal_method_in_transaction_workspace", "grantUsageOnRoleToGrantee"); + diagnostics.fail("illegal_method_in_transaction_workspace", "grantUsageOnRoleToGrantee"); return null; } @Override - public PrivilegeResult revokeUsageOnRoleFromGrantee( - @Nonnull PolarisCallContext callCtx, + public @Nonnull PrivilegeResult revokeUsageOnRoleFromGrantee( + @Nonnull PolarisMetaStoreSession ms, @Nullable PolarisEntityCore catalog, @Nonnull PolarisEntityCore role, @Nonnull PolarisEntityCore grantee) { - callCtx - .getDiagServices() - .fail("illegal_method_in_transaction_workspace", "revokeUsageOnRoleFromGrantee"); + diagnostics.fail("illegal_method_in_transaction_workspace", "revokeUsageOnRoleFromGrantee"); return null; } @Override - public PrivilegeResult grantPrivilegeOnSecurableToRole( - @Nonnull PolarisCallContext callCtx, + public @Nonnull PrivilegeResult grantPrivilegeOnSecurableToRole( + @Nonnull PolarisMetaStoreSession ms, @Nonnull PolarisEntityCore grantee, @Nullable List catalogPath, @Nonnull PolarisEntityCore securable, @Nonnull PolarisPrivilege privilege) { - callCtx - .getDiagServices() - .fail("illegal_method_in_transaction_workspace", "grantPrivilegeOnSecurableToRole"); + diagnostics.fail("illegal_method_in_transaction_workspace", "grantPrivilegeOnSecurableToRole"); return null; } @Override - public PrivilegeResult revokePrivilegeOnSecurableFromRole( - @Nonnull PolarisCallContext callCtx, + public @Nonnull PrivilegeResult revokePrivilegeOnSecurableFromRole( + @Nonnull PolarisMetaStoreSession ms, @Nonnull PolarisEntityCore grantee, @Nullable List catalogPath, @Nonnull PolarisEntityCore securable, @Nonnull PolarisPrivilege privilege) { - callCtx - .getDiagServices() - .fail("illegal_method_in_transaction_workspace", "revokePrivilegeOnSecurableFromRole"); + diagnostics.fail( + "illegal_method_in_transaction_workspace", "revokePrivilegeOnSecurableFromRole"); return null; } @Override - public LoadGrantsResult loadGrantsOnSecurable( - @Nonnull PolarisCallContext callCtx, long securableCatalogId, long securableId) { - callCtx - .getDiagServices() - .fail("illegal_method_in_transaction_workspace", "loadGrantsOnSecurable"); + public @Nonnull LoadGrantsResult loadGrantsOnSecurable( + @Nonnull PolarisMetaStoreSession ms, long securableCatalogId, long securableId) { + diagnostics.fail("illegal_method_in_transaction_workspace", "loadGrantsOnSecurable"); return null; } @Override - public LoadGrantsResult loadGrantsToGrantee( - PolarisCallContext callCtx, long granteeCatalogId, long granteeId) { - callCtx - .getDiagServices() - .fail("illegal_method_in_transaction_workspace", "loadGrantsToGrantee"); + public @Nonnull LoadGrantsResult loadGrantsToGrantee( + @Nonnull PolarisMetaStoreSession ms, long granteeCatalogId, long granteeId) { + diagnostics.fail("illegal_method_in_transaction_workspace", "loadGrantsToGrantee"); return null; } @Override - public ChangeTrackingResult loadEntitiesChangeTracking( - @Nonnull PolarisCallContext callCtx, @Nonnull List entityIds) { - callCtx - .getDiagServices() - .fail("illegal_method_in_transaction_workspace", "loadEntitiesChangeTracking"); + public @Nonnull ChangeTrackingResult loadEntitiesChangeTracking( + @Nonnull PolarisMetaStoreSession ms, @Nonnull List entityIds) { + diagnostics.fail("illegal_method_in_transaction_workspace", "loadEntitiesChangeTracking"); return null; } @Override + @Nonnull public EntityResult loadEntity( - @Nonnull PolarisCallContext callCtx, long entityCatalogId, long entityId) { - callCtx.getDiagServices().fail("illegal_method_in_transaction_workspace", "loadEntity"); + @Nonnull PolarisMetaStoreSession ms, long entityCatalogId, long entityId) { + diagnostics.fail("illegal_method_in_transaction_workspace", "loadEntity"); return null; } @Override + @Nonnull public EntitiesResult loadTasks( - @Nonnull PolarisCallContext callCtx, String executorId, int limit) { - callCtx.getDiagServices().fail("illegal_method_in_transaction_workspace", "loadTasks"); + @Nonnull PolarisMetaStoreSession ms, String executorId, int limit) { + diagnostics.fail("illegal_method_in_transaction_workspace", "loadTasks"); return null; } @Override - public ScopedCredentialsResult getSubscopedCredsForEntity( - @Nonnull PolarisCallContext callCtx, + public @Nonnull ScopedCredentialsResult getSubscopedCredsForEntity( + @Nonnull PolarisMetaStoreSession metaStoreSession, long catalogId, long entityId, boolean allowListOperation, @Nonnull Set allowedReadLocations, @Nonnull Set allowedWriteLocations) { return delegate.getSubscopedCredsForEntity( - callCtx, + metaStoreSession, catalogId, entityId, allowListOperation, @@ -322,51 +301,43 @@ public ScopedCredentialsResult getSubscopedCredsForEntity( } @Override - public ValidateAccessResult validateAccessToLocations( - @Nonnull PolarisCallContext callCtx, + public @Nonnull ValidateAccessResult validateAccessToLocations( + @Nonnull PolarisMetaStoreSession metaStoreSession, long catalogId, long entityId, @Nonnull Set actions, @Nonnull Set locations) { - callCtx - .getDiagServices() - .fail("illegal_method_in_transaction_workspace", "validateAccessToLocations"); + diagnostics.fail("illegal_method_in_transaction_workspace", "validateAccessToLocations"); return null; } @Override - public CachedEntryResult loadCachedEntryById( - @Nonnull PolarisCallContext callCtx, long entityCatalogId, long entityId) { - callCtx - .getDiagServices() - .fail("illegal_method_in_transaction_workspace", "loadCachedEntryById"); + public @Nonnull CachedEntryResult loadCachedEntryById( + @Nonnull PolarisMetaStoreSession ms, long entityCatalogId, long entityId) { + diagnostics.fail("illegal_method_in_transaction_workspace", "loadCachedEntryById"); return null; } @Override - public CachedEntryResult loadCachedEntryByName( - @Nonnull PolarisCallContext callCtx, + public @Nonnull CachedEntryResult loadCachedEntryByName( + @Nonnull PolarisMetaStoreSession ms, long entityCatalogId, long parentId, @Nonnull PolarisEntityType entityType, @Nonnull String entityName) { - callCtx - .getDiagServices() - .fail("illegal_method_in_transaction_workspace", "loadCachedEntryByName"); + diagnostics.fail("illegal_method_in_transaction_workspace", "loadCachedEntryByName"); return null; } @Override - public CachedEntryResult refreshCachedEntity( - @Nonnull PolarisCallContext callCtx, + public @Nonnull CachedEntryResult refreshCachedEntity( + @Nonnull PolarisMetaStoreSession ms, int entityVersion, int entityGrantRecordsVersion, @Nonnull PolarisEntityType entityType, long entityCatalogId, long entityId) { - callCtx - .getDiagServices() - .fail("illegal_method_in_transaction_workspace", "refreshCachedEntity"); + diagnostics.fail("illegal_method_in_transaction_workspace", "refreshCachedEntity"); return null; } } diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/cache/EntityCache.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/cache/EntityCache.java index 05efba7fb..c7d96f56b 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/cache/EntityCache.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/cache/EntityCache.java @@ -27,10 +27,11 @@ import java.util.List; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.TimeUnit; -import org.apache.polaris.core.PolarisCallContext; +import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.entity.PolarisBaseEntity; import org.apache.polaris.core.entity.PolarisEntityType; import org.apache.polaris.core.entity.PolarisGrantRecord; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; import org.apache.polaris.core.persistence.cache.PolarisRemoteCache.CachedEntryResult; /** The entity cache, can be private or shared */ @@ -39,6 +40,8 @@ public class EntityCache { // cache mode private EntityCacheMode cacheMode; + private final PolarisDiagnostics diagServices; + // the meta store manager private final PolarisRemoteCache polarisRemoteCache; @@ -53,7 +56,9 @@ public class EntityCache { * * @param polarisRemoteCache the meta store manager implementation */ - public EntityCache(@Nonnull PolarisRemoteCache polarisRemoteCache) { + public EntityCache( + @Nonnull PolarisRemoteCache polarisRemoteCache, @Nonnull PolarisDiagnostics diagServices) { + this.diagServices = diagServices; // by name cache this.byName = new ConcurrentHashMap<>(); @@ -241,7 +246,7 @@ public void setCacheMode(EntityCacheMode cacheMode) { * Refresh the cache if needs be with a version of the entity/grant records matching the minimum * specified version. * - * @param callContext the Polaris call context + * @param session the metastore session * @param entityToValidate copy of the entity held by the caller to validate * @param entityMinVersion minimum expected version. Should be reloaded if found in a cache with a * version less than this one @@ -250,7 +255,7 @@ public void setCacheMode(EntityCacheMode cacheMode) { * @return the cache entry for the entity or null if the specified entity does not exist */ public @Nullable EntityCacheEntry getAndRefreshIfNeeded( - @Nonnull PolarisCallContext callContext, + @Nonnull PolarisMetaStoreSession session, @Nonnull PolarisBaseEntity entityToValidate, int entityMinVersion, int entityGrantRecordsMinVersion) { @@ -290,7 +295,7 @@ && isNewer(existingCacheEntry, existingCacheEntryByName)) { if (existingCacheEntry == null) { // try to load it refreshedCacheEntry = - this.polarisRemoteCache.loadCachedEntryById(callContext, entityCatalogId, entityId); + this.polarisRemoteCache.loadCachedEntryById(session, entityCatalogId, entityId); if (refreshedCacheEntry.isSuccess()) { entity = refreshedCacheEntry.getEntity(); grantRecords = refreshedCacheEntry.getEntityGrantRecords(); @@ -302,7 +307,7 @@ && isNewer(existingCacheEntry, existingCacheEntryByName)) { // refresh it refreshedCacheEntry = this.polarisRemoteCache.refreshCachedEntity( - callContext, + session, existingCacheEntry.getEntity().getEntityVersion(), existingCacheEntry.getEntity().getGrantRecordsVersion(), entityType, @@ -328,16 +333,14 @@ && isNewer(existingCacheEntry, existingCacheEntryByName)) { } // assert that entity, grant records and version are all set - callContext.getDiagServices().checkNotNull(entity, "unexpected_null_entity"); - callContext.getDiagServices().checkNotNull(grantRecords, "unexpected_null_grant_records"); - callContext - .getDiagServices() - .check(grantRecordsVersion > 0, "unexpected_null_grant_records_version"); + diagServices.checkNotNull(entity, "unexpected_null_entity"); + diagServices.checkNotNull(grantRecords, "unexpected_null_grant_records"); + diagServices.check(grantRecordsVersion > 0, "unexpected_null_grant_records_version"); // create new cache entry newCacheEntry = new EntityCacheEntry( - callContext.getDiagServices(), + diagServices, existingCacheEntry == null ? System.nanoTime() : existingCacheEntry.getCreatedOnNanoTimestamp(), @@ -359,14 +362,14 @@ && isNewer(existingCacheEntry, existingCacheEntryByName)) { /** * Get the specified entity by name and load it if it is not found. * - * @param callContext the Polaris call context + * @param session the metastore session * @param entityCatalogId id of the catalog where this entity resides or NULL_ID if top-level * @param entityId id of the entity to lookup * @return null if the entity does not exist or was dropped. Else return the entry for that * entity, either as found in the cache or loaded from the backend */ public @Nullable EntityCacheLookupResult getOrLoadEntityById( - @Nonnull PolarisCallContext callContext, long entityCatalogId, long entityId) { + @Nonnull PolarisMetaStoreSession session, long entityCatalogId, long entityId) { // if it exists, we are set EntityCacheEntry entry = this.getEntityById(entityId); @@ -379,7 +382,7 @@ && isNewer(existingCacheEntry, existingCacheEntryByName)) { // load it CachedEntryResult result = - polarisRemoteCache.loadCachedEntryById(callContext, entityCatalogId, entityId); + polarisRemoteCache.loadCachedEntryById(session, entityCatalogId, entityId); // not found, exit if (!result.isSuccess()) { @@ -387,13 +390,12 @@ && isNewer(existingCacheEntry, existingCacheEntryByName)) { } // if found, setup entry - callContext.getDiagServices().checkNotNull(result.getEntity(), "entity_should_loaded"); - callContext - .getDiagServices() - .checkNotNull(result.getEntityGrantRecords(), "entity_grant_records_should_loaded"); + diagServices.checkNotNull(result.getEntity(), "entity_should_loaded"); + diagServices.checkNotNull( + result.getEntityGrantRecords(), "entity_grant_records_should_loaded"); entry = new EntityCacheEntry( - callContext.getDiagServices(), + diagServices, System.nanoTime(), result.getEntity(), result.getEntityGrantRecords(), @@ -412,13 +414,13 @@ && isNewer(existingCacheEntry, existingCacheEntryByName)) { /** * Get the specified entity by name and load it if it is not found. * - * @param callContext the Polaris call context + * @param session the metastore session * @param entityNameKey name of the entity to load * @return null if the entity does not exist or was dropped. Else return the entry for that * entity, either as found in the cache or loaded from the backend */ public @Nullable EntityCacheLookupResult getOrLoadEntityByName( - @Nonnull PolarisCallContext callContext, @Nonnull EntityCacheByNameKey entityNameKey) { + @Nonnull PolarisMetaStoreSession session, @Nonnull EntityCacheByNameKey entityNameKey) { // if it exists, we are set EntityCacheEntry entry = this.getEntityByName(entityNameKey); @@ -432,7 +434,7 @@ && isNewer(existingCacheEntry, existingCacheEntryByName)) { // load it CachedEntryResult result = polarisRemoteCache.loadCachedEntryByName( - callContext, + session, entityNameKey.getCatalogId(), entityNameKey.getParentId(), entityNameKey.getType(), @@ -444,15 +446,14 @@ && isNewer(existingCacheEntry, existingCacheEntryByName)) { } // validate return - callContext.getDiagServices().checkNotNull(result.getEntity(), "entity_should_loaded"); - callContext - .getDiagServices() - .checkNotNull(result.getEntityGrantRecords(), "entity_grant_records_should_loaded"); + diagServices.checkNotNull(result.getEntity(), "entity_should_loaded"); + diagServices.checkNotNull( + result.getEntityGrantRecords(), "entity_grant_records_should_loaded"); // if found, setup entry entry = new EntityCacheEntry( - callContext.getDiagServices(), + diagServices, System.nanoTime(), result.getEntity(), result.getEntityGrantRecords(), diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/cache/PolarisRemoteCache.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/cache/PolarisRemoteCache.java index e45cc8b1d..1bd7f68a3 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/cache/PolarisRemoteCache.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/cache/PolarisRemoteCache.java @@ -23,13 +23,13 @@ import jakarta.annotation.Nonnull; import jakarta.annotation.Nullable; import java.util.List; -import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.entity.PolarisBaseEntity; import org.apache.polaris.core.entity.PolarisChangeTrackingVersions; import org.apache.polaris.core.entity.PolarisEntityId; import org.apache.polaris.core.entity.PolarisEntityType; import org.apache.polaris.core.entity.PolarisGrantRecord; import org.apache.polaris.core.persistence.BaseResult; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; /** * Interface to the remote entity cache. This allows the local cache to detect remote entity changes @@ -40,7 +40,7 @@ public interface PolarisRemoteCache { * Load change tracking information for a set of entities in one single shot and return for each * the version for the entity itself and the version associated to its grant records. * - * @param callCtx call context + * @param session the metastore session * @param entityIds list of catalog/entity pair ids for which we need to efficiently load the * version information, both entity version and grant records version. * @return a list of version tracking information. Order in that returned list is the same as the @@ -48,7 +48,7 @@ public interface PolarisRemoteCache { */ @Nonnull ChangeTrackingResult loadEntitiesChangeTracking( - @Nonnull PolarisCallContext callCtx, @Nonnull List entityIds); + @Nonnull PolarisMetaStoreSession session, @Nonnull List entityIds); /** * Load a cached entry, i.e. an entity definition and associated grant records, from the backend @@ -57,7 +57,7 @@ ChangeTrackingResult loadEntitiesChangeTracking( *

For entities that can be grantees, the associated grant records will include both the grant * records for this entity as a grantee and for this entity as a securable. * - * @param callCtx call context + * @param session the metastore session * @param entityCatalogId id of the catalog for that entity * @param entityId id of the entity * @return cached entry for this entity. Status will be ENTITY_NOT_FOUND if the entity was not @@ -65,17 +65,17 @@ ChangeTrackingResult loadEntitiesChangeTracking( */ @Nonnull CachedEntryResult loadCachedEntryById( - @Nonnull PolarisCallContext callCtx, long entityCatalogId, long entityId); + @Nonnull PolarisMetaStoreSession session, long entityCatalogId, long entityId); /** - * Load a cached entry, i.e. an entity definition and associated grant records, from the backend + * ¬ Load a cached entry, i.e. an entity definition and associated grant records, from the backend * store. The entity is identified by its name. Will return NULL if the entity does not exist, * i.e. has been purged or dropped. * *

For entities that can be grantees, the associated grant records will include both the grant * records for this entity as a grantee and for this entity as a securable. * - * @param callCtx call context + * @param session the metastore session * @param entityCatalogId id of the catalog for that entity * @param parentId the id of the parent of that entity * @param entityType the type of this entity @@ -85,7 +85,7 @@ CachedEntryResult loadCachedEntryById( */ @Nonnull CachedEntryResult loadCachedEntryByName( - @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisMetaStoreSession session, long entityCatalogId, long parentId, @Nonnull PolarisEntityType entityType, @@ -99,7 +99,7 @@ CachedEntryResult loadCachedEntryByName( *

For entities that can be grantees, the associated grant records will include both the grant * records for this entity as a grantee and for this entity as a securable. * - * @param callCtx call context + * @param session the metastore session * @param entityType type of the entity whose cached entry we are refreshing * @param entityCatalogId id of the catalog for that entity * @param entityId the id of the entity to load @@ -108,7 +108,7 @@ CachedEntryResult loadCachedEntryByName( */ @Nonnull CachedEntryResult refreshCachedEntity( - @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisMetaStoreSession session, int entityVersion, int entityGrantRecordsVersion, @Nonnull PolarisEntityType entityType, diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/resolver/PolarisResolutionManifest.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/resolver/PolarisResolutionManifest.java index 2ab0140a0..7e874c5dc 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/resolver/PolarisResolutionManifest.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/resolver/PolarisResolutionManifest.java @@ -30,13 +30,13 @@ import java.util.stream.Collectors; import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.auth.AuthenticatedPolarisPrincipal; -import org.apache.polaris.core.context.CallContext; import org.apache.polaris.core.entity.PolarisBaseEntity; 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.entity.PrincipalRoleEntity; import org.apache.polaris.core.persistence.PolarisEntityManager; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; import org.apache.polaris.core.persistence.PolarisResolvedPathWrapper; import org.apache.polaris.core.persistence.ResolvedPolarisEntity; import org.apache.polaris.core.persistence.cache.EntityCacheEntry; @@ -54,8 +54,8 @@ public class PolarisResolutionManifest implements PolarisResolutionManifestCatalogView { private static final Logger LOGGER = LoggerFactory.getLogger(PolarisResolutionManifest.class); + private final PolarisMetaStoreSession metaStoreSession; private final PolarisEntityManager entityManager; - private final CallContext callContext; private final SecurityContext securityContext; private final AuthenticatedPolarisPrincipal authenticatedPrincipal; private final String catalogName; @@ -81,15 +81,17 @@ public class PolarisResolutionManifest implements PolarisResolutionManifestCatal private ResolverStatus primaryResolverStatus = null; public PolarisResolutionManifest( - CallContext callContext, + PolarisMetaStoreSession metaStoreSession, + PolarisDiagnostics diagnostics, PolarisEntityManager entityManager, SecurityContext securityContext, String catalogName) { + this.metaStoreSession = metaStoreSession; + this.diagnostics = diagnostics; this.entityManager = entityManager; - this.callContext = callContext; this.catalogName = catalogName; - this.primaryResolver = entityManager.prepareResolver(callContext, securityContext, catalogName); - this.diagnostics = callContext.getPolarisCallContext().getDiagServices(); + this.primaryResolver = + entityManager.prepareResolver(metaStoreSession, securityContext, catalogName); this.diagnostics.checkNotNull(securityContext, "null_security_context_for_resolution_manifest"); this.securityContext = securityContext; diagnostics.check( @@ -202,7 +204,7 @@ public PolarisResolvedPathWrapper getPassthroughResolvedPath(Object key) { // Run a single-use Resolver for this path. Resolver passthroughResolver = - entityManager.prepareResolver(callContext, securityContext, catalogName); + entityManager.prepareResolver(metaStoreSession, securityContext, catalogName); passthroughResolver.addPath(requestedPath); ResolverStatus status = passthroughResolver.resolveAll(); diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/resolver/Resolver.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/resolver/Resolver.java index 814124823..7b71631a5 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/resolver/Resolver.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/persistence/resolver/Resolver.java @@ -30,7 +30,6 @@ import java.util.Map; import java.util.Set; import java.util.stream.Collectors; -import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.auth.AuthenticatedPolarisPrincipal; import org.apache.polaris.core.entity.PolarisBaseEntity; @@ -40,6 +39,7 @@ import org.apache.polaris.core.entity.PolarisEntityType; import org.apache.polaris.core.entity.PolarisGrantRecord; import org.apache.polaris.core.entity.PolarisPrivilege; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; import org.apache.polaris.core.persistence.cache.EntityCache; import org.apache.polaris.core.persistence.cache.EntityCacheByNameKey; import org.apache.polaris.core.persistence.cache.EntityCacheEntry; @@ -53,8 +53,7 @@ */ public class Resolver { - // we stash the Polaris call context here - private final @Nonnull PolarisCallContext polarisCallContext; + private final @Nonnull PolarisMetaStoreSession metaStoreSession; // the diagnostic services private final @Nonnull PolarisDiagnostics diagnostics; @@ -106,7 +105,8 @@ public class Resolver { /** * Constructor, effectively starts an entity resolver session * - * @param polarisCallContext the polaris call context + * @param metaStoreSession the meta store session + * @param diagnostics the diagnostics service * @param polarisRemoteCache meta store manager * @param securityContext The {@link AuthenticatedPolarisPrincipal} for the current request * @param cache shared entity cache @@ -119,21 +119,20 @@ public class Resolver { * service admin should use null for that parameter. */ public Resolver( - @Nonnull PolarisCallContext polarisCallContext, + @Nonnull PolarisMetaStoreSession metaStoreSession, + @Nonnull PolarisDiagnostics diagnostics, @Nonnull PolarisRemoteCache polarisRemoteCache, @Nonnull SecurityContext securityContext, @Nonnull EntityCache cache, @Nullable String referenceCatalogName) { - this.polarisCallContext = polarisCallContext; - this.diagnostics = polarisCallContext.getDiagServices(); + this.metaStoreSession = metaStoreSession; + this.diagnostics = diagnostics; this.polarisRemoteCache = polarisRemoteCache; this.cache = cache; this.securityContext = securityContext; this.referenceCatalogName = referenceCatalogName; // validate inputs - this.diagnostics.checkNotNull(polarisCallContext, "unexpected_null_polarisCallContext"); - this.diagnostics.checkNotNull(polarisRemoteCache, "unexpected_null_polarisRemoteCache"); this.diagnostics.checkNotNull(cache, "unexpected_null_cache"); this.diagnostics.checkNotNull(securityContext, "security_context_must_be_specified"); this.diagnostics.checkNotNull( @@ -526,7 +525,7 @@ private boolean bulkValidate(List toValidate) { // now get the current backend versions of all these entities ChangeTrackingResult changeTrackingResult = - this.polarisRemoteCache.loadEntitiesChangeTracking(this.polarisCallContext, entityIds); + this.polarisRemoteCache.loadEntitiesChangeTracking(metaStoreSession, entityIds); // refresh any entity which is not fresh. If an entity is missing, reload it Iterator entityIterator = toValidate.iterator(); @@ -556,7 +555,7 @@ private boolean bulkValidate(List toValidate) { // refresh that entity. If versions is null, it has been dropped refreshedCacheEntry = this.cache.getAndRefreshIfNeeded( - this.polarisCallContext, + metaStoreSession, entity, versions.getEntityVersion(), versions.getGrantRecordsVersion()); @@ -929,7 +928,7 @@ private EntityCacheEntry resolveByName( // get or load by name EntityCacheLookupResult lookupResult = this.cache.getOrLoadEntityByName( - this.polarisCallContext, + metaStoreSession, new EntityCacheByNameKey(catalogId, parentId, entityType, entityName)); // if not found @@ -967,7 +966,7 @@ private EntityCacheEntry resolveById( long entityId) { // get or load by name EntityCacheLookupResult lookupResult = - this.cache.getOrLoadEntityById(this.polarisCallContext, catalogId, entityId); + this.cache.getOrLoadEntityById(metaStoreSession, catalogId, entityId); // if not found, return null if (lookupResult == null) { diff --git a/polaris-core/src/main/java/org/apache/polaris/core/storage/InMemoryStorageIntegration.java b/polaris-core/src/main/java/org/apache/polaris/core/storage/InMemoryStorageIntegration.java index ae8726428..307650279 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/storage/InMemoryStorageIntegration.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/storage/InMemoryStorageIntegration.java @@ -22,32 +22,39 @@ import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.Optional; import java.util.Set; import java.util.function.Function; import java.util.stream.Collectors; -import org.apache.polaris.core.context.CallContext; +import org.apache.polaris.core.PolarisConfigurationStore; +import org.apache.polaris.core.context.RealmContext; /** * Base class for in-memory implementations of {@link PolarisStorageIntegration}. A basic - * implementation of {@link #validateAccessToLocations(PolarisStorageConfigurationInfo, Set, Set)} - * is provided that checks to see that the list of locations being accessed is among the list of - * {@link PolarisStorageConfigurationInfo#getAllowedLocations()}. Locations being accessed must be - * equal to or a subdirectory of at least one of the allowed locations. + * implementation of {@link PolarisStorageIntegration#validateAccessToLocations(RealmContext, + * PolarisStorageConfigurationInfo, Set, Set)} is provided that checks to see that the list of + * locations being accessed is among the list of {@link + * PolarisStorageConfigurationInfo#getAllowedLocations()}. Locations being accessed must be equal to + * or a subdirectory of at least one of the allowed locations. * * @param */ public abstract class InMemoryStorageIntegration extends PolarisStorageIntegration { - public InMemoryStorageIntegration(String identifierOrId) { + private final PolarisConfigurationStore configurationStore; + + public InMemoryStorageIntegration( + PolarisConfigurationStore configurationStore, String identifierOrId) { super(identifierOrId); + this.configurationStore = configurationStore; } /** * Check that the locations being accessed are all equal to or subdirectories of at least one of * the {@link PolarisStorageConfigurationInfo#getAllowedLocations}. * + * @param realmContext + * @param configurationStore * @param actions a set of operation actions to validate, like LIST/READ/DELETE/WRITE/ALL * @param locations a set of locations to get access to * @return a map of location to a validation result for each action passed in. In this @@ -56,6 +63,8 @@ public InMemoryStorageIntegration(String identifierOrId) { */ public static Map> validateSubpathsOfAllowedLocations( + @Nonnull RealmContext realmContext, + @Nonnull PolarisConfigurationStore configurationStore, @Nonnull PolarisStorageConfigurationInfo storageConfig, @Nonnull Set actions, @Nonnull Set locations) { @@ -77,13 +86,7 @@ public InMemoryStorageIntegration(String identifierOrId) { allowedLocationStrings.stream().map(StorageLocation::of).collect(Collectors.toList()); boolean allowWildcardLocation = - Optional.ofNullable(CallContext.getCurrentContext()) - .flatMap(c -> Optional.ofNullable(c.getPolarisCallContext())) - .map( - pc -> - pc.getConfigurationStore() - .getConfiguration(pc, "ALLOW_WILDCARD_LOCATION", false)) - .orElse(false); + configurationStore.getConfiguration(realmContext, "ALLOW_WILDCARD_LOCATION", false); if (allowWildcardLocation && allowedLocationStrings.contains("*")) { return locations.stream() @@ -126,9 +129,11 @@ public InMemoryStorageIntegration(String identifierOrId) { @Override @Nonnull public Map> validateAccessToLocations( + @Nonnull RealmContext realmContext, @Nonnull T storageConfig, @Nonnull Set actions, @Nonnull Set locations) { - return validateSubpathsOfAllowedLocations(storageConfig, actions, locations); + return validateSubpathsOfAllowedLocations( + realmContext, configurationStore, storageConfig, actions, locations); } } diff --git a/polaris-core/src/main/java/org/apache/polaris/core/storage/PolarisCredentialVendor.java b/polaris-core/src/main/java/org/apache/polaris/core/storage/PolarisCredentialVendor.java index 9f03082a9..02523056b 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/storage/PolarisCredentialVendor.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/storage/PolarisCredentialVendor.java @@ -25,8 +25,8 @@ import java.util.EnumMap; import java.util.Map; import java.util.Set; -import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.persistence.BaseResult; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; /** Manage credentials for storage locations. */ public interface PolarisCredentialVendor { @@ -34,7 +34,7 @@ public interface PolarisCredentialVendor { * Get a sub-scoped credentials for an entity against the provided allowed read and write * locations. * - * @param callCtx the polaris call context + * @param metaStoreSession the meta store session * @param catalogId the catalog id * @param entityId the entity id * @param allowListOperation whether to allow LIST operation on the allowedReadLocations and @@ -45,7 +45,7 @@ public interface PolarisCredentialVendor { */ @Nonnull ScopedCredentialsResult getSubscopedCredsForEntity( - @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisMetaStoreSession metaStoreSession, long catalogId, long entityId, boolean allowListOperation, @@ -55,7 +55,7 @@ ScopedCredentialsResult getSubscopedCredsForEntity( /** * Validate whether the entity has access to the locations with the provided target operations * - * @param callCtx the polaris call context + * @param metaStoreSession the meta store session * @param catalogId the catalog id * @param entityId the entity id * @param actions a set of operation actions: READ/WRITE/LIST/DELETE/ALL @@ -87,7 +87,7 @@ ScopedCredentialsResult getSubscopedCredsForEntity( */ @Nonnull ValidateAccessResult validateAccessToLocations( - @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisMetaStoreSession metaStoreSession, long catalogId, long entityId, @Nonnull Set actions, diff --git a/polaris-core/src/main/java/org/apache/polaris/core/storage/PolarisStorageConfigurationInfo.java b/polaris-core/src/main/java/org/apache/polaris/core/storage/PolarisStorageConfigurationInfo.java index 6b0638e83..a65e1e8eb 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/storage/PolarisStorageConfigurationInfo.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/storage/PolarisStorageConfigurationInfo.java @@ -37,9 +37,10 @@ import java.util.stream.Collectors; import java.util.stream.Stream; import org.apache.polaris.core.PolarisConfiguration; +import org.apache.polaris.core.PolarisConfigurationStore; import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.admin.model.Catalog; -import org.apache.polaris.core.context.CallContext; +import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.CatalogEntity; import org.apache.polaris.core.entity.PolarisEntity; import org.apache.polaris.core.entity.PolarisEntityConstants; @@ -136,7 +137,10 @@ public static PolarisStorageConfigurationInfo deserialize( } public static Optional forEntityPath( - PolarisDiagnostics diagnostics, List entityPath) { + RealmContext realmContext, + PolarisConfigurationStore configurationStore, + PolarisDiagnostics diagnostics, + List entityPath) { return findStorageInfoFromHierarchy(entityPath) .map( storageInfo -> @@ -162,13 +166,10 @@ public static Optional forEntityPath( .orElse(null); CatalogEntity catalog = CatalogEntity.of(entityPath.get(0)); boolean allowEscape = - CallContext.getCurrentContext() - .getPolarisCallContext() - .getConfigurationStore() - .getConfiguration( - CallContext.getCurrentContext().getPolarisCallContext(), - catalog, - PolarisConfiguration.ALLOW_UNSTRUCTURED_TABLE_LOCATION); + configurationStore.getConfiguration( + realmContext, + catalog, + PolarisConfiguration.ALLOW_UNSTRUCTURED_TABLE_LOCATION); if (!allowEscape && catalog.getCatalogType() != Catalog.TypeEnum.EXTERNAL && baseLocation != null) { diff --git a/polaris-core/src/main/java/org/apache/polaris/core/storage/PolarisStorageIntegration.java b/polaris-core/src/main/java/org/apache/polaris/core/storage/PolarisStorageIntegration.java index eec9a4209..9c50c6305 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/storage/PolarisStorageIntegration.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/storage/PolarisStorageIntegration.java @@ -24,6 +24,7 @@ import java.util.Objects; import java.util.Set; import org.apache.polaris.core.PolarisDiagnostics; +import org.apache.polaris.core.context.RealmContext; /** * Abstract of Polaris Storage Integration. It holds the reference to an object that having the @@ -46,6 +47,7 @@ public String getStorageIdentifierOrId() { /** * Subscope the creds against the allowed read and write locations. * + * @param realmContext the realm context * @param diagnostics the diagnostics service * @param storageConfig storage configuration * @param allowListOperation whether to allow LIST on all the provided allowed read/write @@ -55,6 +57,7 @@ public String getStorageIdentifierOrId() { * @return An enum map including the scoped credentials */ public abstract EnumMap getSubscopedCreds( + @Nonnull RealmContext realmContext, @Nonnull PolarisDiagnostics diagnostics, @Nonnull T storageConfig, boolean allowListOperation, @@ -64,6 +67,7 @@ public abstract EnumMap getSubscopedCreds( /** * Validate access for the provided operation actions and locations. * + * @param realmContext * @param actions a set of operation actions to validate, like LIST/READ/DELETE/WRITE/ALL * @param locations a set of locations to get access to * @return A Map of string, representing the result of validation, the key value is {@code @@ -95,12 +99,14 @@ public abstract EnumMap getSubscopedCreds( @Nonnull public abstract Map> validateAccessToLocations( + RealmContext realmContext, @Nonnull T storageConfig, @Nonnull Set actions, @Nonnull Set locations); /** - * Result of calling {@link #validateAccessToLocations(PolarisStorageConfigurationInfo, Set, Set)} + * Result of calling {@link PolarisStorageIntegration#validateAccessToLocations(RealmContext, + * PolarisStorageConfigurationInfo, Set, Set)} */ public static final class ValidationResult { private final boolean success; diff --git a/polaris-core/src/main/java/org/apache/polaris/core/storage/aws/AwsCredentialsStorageIntegration.java b/polaris-core/src/main/java/org/apache/polaris/core/storage/aws/AwsCredentialsStorageIntegration.java index 4832a0d16..00c1358aa 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/storage/aws/AwsCredentialsStorageIntegration.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/storage/aws/AwsCredentialsStorageIntegration.java @@ -19,7 +19,6 @@ package org.apache.polaris.core.storage.aws; import static org.apache.polaris.core.PolarisConfiguration.STORAGE_CREDENTIAL_DURATION_SECONDS; -import static org.apache.polaris.core.PolarisConfiguration.loadConfig; import jakarta.annotation.Nonnull; import java.net.URI; @@ -29,7 +28,9 @@ import java.util.Optional; import java.util.Set; import java.util.stream.Stream; +import org.apache.polaris.core.PolarisConfigurationStore; import org.apache.polaris.core.PolarisDiagnostics; +import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.storage.InMemoryStorageIntegration; import org.apache.polaris.core.storage.PolarisCredentialProperty; import org.apache.polaris.core.storage.StorageUtil; @@ -45,16 +46,21 @@ /** Credential vendor that supports generating */ public class AwsCredentialsStorageIntegration extends InMemoryStorageIntegration { + + private final PolarisConfigurationStore configurationStore; private final StsClient stsClient; - public AwsCredentialsStorageIntegration(StsClient stsClient) { - super(AwsCredentialsStorageIntegration.class.getName()); + public AwsCredentialsStorageIntegration( + PolarisConfigurationStore configurationStore, StsClient stsClient) { + super(configurationStore, AwsCredentialsStorageIntegration.class.getName()); + this.configurationStore = configurationStore; this.stsClient = stsClient; } /** {@inheritDoc} */ @Override public EnumMap getSubscopedCreds( + @Nonnull RealmContext realmContext, @Nonnull PolarisDiagnostics diagnostics, @Nonnull AwsStorageConfigurationInfo storageConfig, boolean allowListOperation, @@ -73,7 +79,9 @@ public EnumMap getSubscopedCreds( allowedReadLocations, allowedWriteLocations) .toJson()) - .durationSeconds(loadConfig(STORAGE_CREDENTIAL_DURATION_SECONDS)) + .durationSeconds( + configurationStore.getConfiguration( + realmContext, STORAGE_CREDENTIAL_DURATION_SECONDS)) .build()); EnumMap credentialMap = new EnumMap<>(PolarisCredentialProperty.class); diff --git a/polaris-core/src/main/java/org/apache/polaris/core/storage/azure/AzureCredentialsStorageIntegration.java b/polaris-core/src/main/java/org/apache/polaris/core/storage/azure/AzureCredentialsStorageIntegration.java index 4013cbd81..5da34596d 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/storage/azure/AzureCredentialsStorageIntegration.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/storage/azure/AzureCredentialsStorageIntegration.java @@ -46,7 +46,9 @@ import java.util.Objects; import java.util.Set; import org.apache.polaris.core.PolarisConfiguration; +import org.apache.polaris.core.PolarisConfigurationStore; import org.apache.polaris.core.PolarisDiagnostics; +import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.storage.InMemoryStorageIntegration; import org.apache.polaris.core.storage.PolarisCredentialProperty; import org.slf4j.Logger; @@ -60,10 +62,12 @@ public class AzureCredentialsStorageIntegration private static final Logger LOGGER = LoggerFactory.getLogger(AzureCredentialsStorageIntegration.class); - final DefaultAzureCredential defaultAzureCredential; + private final PolarisConfigurationStore configurationStore; + private final DefaultAzureCredential defaultAzureCredential; - public AzureCredentialsStorageIntegration() { - super(AzureCredentialsStorageIntegration.class.getName()); + public AzureCredentialsStorageIntegration(PolarisConfigurationStore configurationStore) { + super(configurationStore, AzureCredentialsStorageIntegration.class.getName()); + this.configurationStore = configurationStore; // The DefaultAzureCredential will by default load the environment variables for client id, // client secret, tenant id defaultAzureCredential = new DefaultAzureCredentialBuilder().build(); @@ -71,6 +75,7 @@ public AzureCredentialsStorageIntegration() { @Override public EnumMap getSubscopedCreds( + @Nonnull RealmContext realmContext, @Nonnull PolarisDiagnostics diagnostics, @Nonnull AzureStorageConfigurationInfo storageConfig, boolean allowListOperation, @@ -126,7 +131,8 @@ public EnumMap getSubscopedCreds( // clock skew between the client and server, OffsetDateTime startTime = start.truncatedTo(ChronoUnit.SECONDS).atOffset(ZoneOffset.UTC); int intendedDurationSeconds = - PolarisConfiguration.loadConfig(PolarisConfiguration.STORAGE_CREDENTIAL_DURATION_SECONDS); + configurationStore.getConfiguration( + realmContext, PolarisConfiguration.STORAGE_CREDENTIAL_DURATION_SECONDS); OffsetDateTime intendedEndTime = start.plusSeconds(intendedDurationSeconds).atOffset(ZoneOffset.UTC); OffsetDateTime maxAllowedEndTime = diff --git a/polaris-core/src/main/java/org/apache/polaris/core/storage/cache/StorageCredentialCache.java b/polaris-core/src/main/java/org/apache/polaris/core/storage/cache/StorageCredentialCache.java index 44a1b0fd4..871180160 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/storage/cache/StorageCredentialCache.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/storage/cache/StorageCredentialCache.java @@ -28,10 +28,12 @@ import java.util.concurrent.TimeUnit; import java.util.function.Function; import org.apache.iceberg.exceptions.UnprocessableEntityException; -import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.PolarisConfiguration; +import org.apache.polaris.core.PolarisConfigurationStore; +import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.entity.PolarisEntity; import org.apache.polaris.core.entity.PolarisEntityType; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; import org.apache.polaris.core.storage.PolarisCredentialVendor; import org.jetbrains.annotations.VisibleForTesting; import org.slf4j.Logger; @@ -44,9 +46,14 @@ public class StorageCredentialCache { private static final long CACHE_MAX_NUMBER_OF_ENTRIES = 10_000L; private final LoadingCache cache; + private final PolarisDiagnostics diagnostics; + private final PolarisConfigurationStore configurationStore; /** Initialize the creds cache */ - public StorageCredentialCache() { + public StorageCredentialCache( + PolarisDiagnostics diagnostics, PolarisConfigurationStore configurationStore) { + this.diagnostics = diagnostics; + this.configurationStore = configurationStore; cache = Caffeine.newBuilder() .maximumSize(CACHE_MAX_NUMBER_OF_ENTRIES) @@ -92,12 +99,15 @@ public long expireAfterRead( } /** How long credentials should remain in the cache. */ - private static long maxCacheDurationMs() { + private long maxCacheDurationMs() { + // FIXME no realm context available; in practice this means that these values are not + // overridable by realm-specific values var cacheDurationSeconds = - PolarisConfiguration.loadConfig( - PolarisConfiguration.STORAGE_CREDENTIAL_CACHE_DURATION_SECONDS); + configurationStore.getConfiguration( + null, PolarisConfiguration.STORAGE_CREDENTIAL_CACHE_DURATION_SECONDS); var credentialDurationSeconds = - PolarisConfiguration.loadConfig(PolarisConfiguration.STORAGE_CREDENTIAL_DURATION_SECONDS); + configurationStore.getConfiguration( + null, PolarisConfiguration.STORAGE_CREDENTIAL_DURATION_SECONDS); if (cacheDurationSeconds >= credentialDurationSeconds) { throw new IllegalArgumentException( String.format( @@ -113,7 +123,7 @@ private static long maxCacheDurationMs() { * Either get from the cache or generate a new entry for a scoped creds * * @param credentialVendor the credential vendor used to generate a new scoped creds if needed - * @param callCtx the call context + * @param metaStoreSession the meta store session * @param polarisEntity the polaris entity that is going to scoped creds * @param allowListOperation whether allow list action on the provided read and write locations * @param allowedReadLocations a set of allowed to read locations @@ -122,30 +132,25 @@ private static long maxCacheDurationMs() { */ public Map getOrGenerateSubScopeCreds( @Nonnull PolarisCredentialVendor credentialVendor, - @Nonnull PolarisCallContext callCtx, + @Nonnull PolarisMetaStoreSession metaStoreSession, @Nonnull PolarisEntity polarisEntity, boolean allowListOperation, @Nonnull Set allowedReadLocations, @Nonnull Set allowedWriteLocations) { if (!isTypeSupported(polarisEntity.getType())) { - callCtx - .getDiagServices() - .fail("entity_type_not_suppported_to_scope_creds", "type={}", polarisEntity.getType()); + diagnostics.fail( + "entity_type_not_suppported_to_scope_creds", "type={}", polarisEntity.getType()); } StorageCredentialCacheKey key = new StorageCredentialCacheKey( - polarisEntity, - allowListOperation, - allowedReadLocations, - allowedWriteLocations, - callCtx); + polarisEntity, allowListOperation, allowedReadLocations, allowedWriteLocations); LOGGER.atDebug().addKeyValue("key", key).log("subscopedCredsCache"); Function loader = k -> { LOGGER.atDebug().log("StorageCredentialCache::load"); PolarisCredentialVendor.ScopedCredentialsResult scopedCredentialsResult = credentialVendor.getSubscopedCredsForEntity( - k.getCallContext(), + metaStoreSession, k.getCatalogId(), k.getEntityId(), k.isAllowedListAction(), diff --git a/polaris-core/src/main/java/org/apache/polaris/core/storage/cache/StorageCredentialCacheKey.java b/polaris-core/src/main/java/org/apache/polaris/core/storage/cache/StorageCredentialCacheKey.java index 48d83ef16..3d0a77648 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/storage/cache/StorageCredentialCacheKey.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/storage/cache/StorageCredentialCacheKey.java @@ -18,11 +18,8 @@ */ package org.apache.polaris.core.storage.cache; -import jakarta.annotation.Nullable; import java.util.Objects; import java.util.Set; -import org.apache.polaris.core.PolarisCallContext; -import org.apache.polaris.core.context.CallContext; import org.apache.polaris.core.entity.PolarisEntity; import org.apache.polaris.core.entity.PolarisEntityConstants; @@ -44,18 +41,11 @@ public class StorageCredentialCacheKey { private final Set allowedWriteLocations; - /** - * The callContext is passed to be used to fetch subscoped creds, but is not used to hash/equals - * as part of the cache key. - */ - private @Nullable PolarisCallContext callContext; - public StorageCredentialCacheKey( PolarisEntity entity, boolean allowedListAction, Set allowedReadLocations, - Set allowedWriteLocations, - @Nullable PolarisCallContext callContext) { + Set allowedWriteLocations) { this.catalogId = entity.getCatalogId(); this.storageConfigSerializedStr = entity @@ -65,10 +55,6 @@ public StorageCredentialCacheKey( this.allowedListAction = allowedListAction; this.allowedReadLocations = allowedReadLocations; this.allowedWriteLocations = allowedWriteLocations; - this.callContext = callContext; - if (this.callContext == null) { - this.callContext = CallContext.getCurrentContext().getPolarisCallContext(); - } } public long getCatalogId() { @@ -95,10 +81,6 @@ public Set getAllowedWriteLocations() { return allowedWriteLocations; } - public @Nullable PolarisCallContext getCallContext() { - return callContext; - } - @Override public boolean equals(Object o) { if (this == o) return true; diff --git a/polaris-core/src/main/java/org/apache/polaris/core/storage/gcp/GcpCredentialsStorageIntegration.java b/polaris-core/src/main/java/org/apache/polaris/core/storage/gcp/GcpCredentialsStorageIntegration.java index 7d886e32c..fb88e4bec 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/storage/gcp/GcpCredentialsStorageIntegration.java +++ b/polaris-core/src/main/java/org/apache/polaris/core/storage/gcp/GcpCredentialsStorageIntegration.java @@ -38,7 +38,9 @@ import java.util.Objects; import java.util.Set; import java.util.stream.Stream; +import org.apache.polaris.core.PolarisConfigurationStore; import org.apache.polaris.core.PolarisDiagnostics; +import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.storage.InMemoryStorageIntegration; import org.apache.polaris.core.storage.PolarisCredentialProperty; import org.apache.polaris.core.storage.PolarisStorageIntegration; @@ -59,8 +61,10 @@ public class GcpCredentialsStorageIntegration private final HttpTransportFactory transportFactory; public GcpCredentialsStorageIntegration( - GoogleCredentials sourceCredentials, HttpTransportFactory transportFactory) { - super(GcpCredentialsStorageIntegration.class.getName()); + PolarisConfigurationStore configurationStore, + GoogleCredentials sourceCredentials, + HttpTransportFactory transportFactory) { + super(configurationStore, GcpCredentialsStorageIntegration.class.getName()); // Needed for when environment variable GOOGLE_APPLICATION_CREDENTIALS points to google service // account key json this.sourceCredentials = @@ -70,6 +74,7 @@ public GcpCredentialsStorageIntegration( @Override public EnumMap getSubscopedCreds( + @Nonnull RealmContext realmContext, @Nonnull PolarisDiagnostics diagnostics, @Nonnull GcpStorageConfigurationInfo storageConfig, boolean allowListOperation, diff --git a/polaris-core/src/test/java/org/apache/polaris/core/persistence/EntityCacheTest.java b/polaris-core/src/test/java/org/apache/polaris/core/persistence/EntityCacheTest.java index c47367769..4e918a42d 100644 --- a/polaris-core/src/test/java/org/apache/polaris/core/persistence/EntityCacheTest.java +++ b/polaris-core/src/test/java/org/apache/polaris/core/persistence/EntityCacheTest.java @@ -20,9 +20,10 @@ import static org.apache.polaris.core.persistence.PrincipalSecretsGenerator.RANDOM_SECRETS; +import java.time.Clock; import java.util.List; import java.util.stream.Collectors; -import org.apache.polaris.core.PolarisCallContext; +import org.apache.polaris.core.PolarisConfigurationStore; import org.apache.polaris.core.PolarisDefaultDiagServiceImpl; import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.entity.PolarisBaseEntity; @@ -44,15 +45,9 @@ public class EntityCacheTest { // diag services private final PolarisDiagnostics diagServices; - // the entity store, use treemap implementation - private final PolarisTreeMapStore store; - // to interact with the metastore private final PolarisMetaStoreSession metaStore; - // polaris call context - private final PolarisCallContext callCtx; - // utility to bootstrap the mata store private final PolarisTestMetaStoreManager tm; @@ -83,13 +78,16 @@ public class EntityCacheTest { */ public EntityCacheTest() { diagServices = new PolarisDefaultDiagServiceImpl(); - store = new PolarisTreeMapStore(diagServices); - metaStore = new PolarisTreeMapMetaStoreSessionImpl(store, Mockito.mock(), RANDOM_SECRETS); - callCtx = new PolarisCallContext(metaStore, diagServices); - metaStoreManager = new PolarisMetaStoreManagerImpl(); + // the entity store, use treemap implementation + PolarisTreeMapStore store = new PolarisTreeMapStore(diagServices); + metaStore = + new PolarisTreeMapMetaStoreSessionImpl(store, Mockito.mock(), RANDOM_SECRETS, diagServices); + metaStoreManager = + new PolarisMetaStoreManagerImpl( + () -> "test", diagServices, new PolarisConfigurationStore() {}, Clock.systemUTC()); // bootstrap the mata store with our test schema - tm = new PolarisTestMetaStoreManager(metaStoreManager, callCtx); + tm = new PolarisTestMetaStoreManager(metaStoreManager, metaStore, diagServices); tm.testCreateTestCatalog(); } @@ -97,7 +95,7 @@ public EntityCacheTest() { * @return new cache for the entity store */ EntityCache allocateNewCache() { - return new EntityCache(this.metaStoreManager); + return new EntityCache(this.metaStoreManager, diagServices); } @Test @@ -108,7 +106,7 @@ void testGetOrLoadEntityByName() { // should exist and no cache hit EntityCacheLookupResult lookup = cache.getOrLoadEntityByName( - this.callCtx, new EntityCacheByNameKey(PolarisEntityType.CATALOG, "test")); + metaStore, new EntityCacheByNameKey(PolarisEntityType.CATALOG, "test")); Assertions.assertThat(lookup).isNotNull(); Assertions.assertThat(lookup.isCacheHit()).isFalse(); Assertions.assertThat(lookup.getCacheEntry()).isNotNull(); @@ -121,12 +119,12 @@ void testGetOrLoadEntityByName() { // do it again, should be found in the cache lookup = cache.getOrLoadEntityByName( - this.callCtx, new EntityCacheByNameKey(PolarisEntityType.CATALOG, "test")); + metaStore, new EntityCacheByNameKey(PolarisEntityType.CATALOG, "test")); Assertions.assertThat(lookup).isNotNull(); Assertions.assertThat(lookup.isCacheHit()).isTrue(); // do it again by id, should be found in the cache - lookup = cache.getOrLoadEntityById(this.callCtx, catalog.getCatalogId(), catalog.getId()); + lookup = cache.getOrLoadEntityById(metaStore, catalog.getCatalogId(), catalog.getId()); Assertions.assertThat(lookup).isNotNull(); Assertions.assertThat(lookup.isCacheHit()).isTrue(); Assertions.assertThat(lookup.getCacheEntry()).isNotNull(); @@ -145,7 +143,7 @@ void testGetOrLoadEntityByName() { Assertions.assertThat(cacheEntry).isNull(); // try to find it in the cache by id. Should not be there, i.e. no cache hit - lookup = cache.getOrLoadEntityById(this.callCtx, N1.getCatalogId(), N1.getId()); + lookup = cache.getOrLoadEntityById(metaStore, N1.getCatalogId(), N1.getId()); Assertions.assertThat(lookup).isNotNull(); Assertions.assertThat(lookup.isCacheHit()).isFalse(); @@ -168,18 +166,17 @@ void testGetOrLoadEntityByName() { Assertions.assertThat(N1_entry.getGrantRecordsAsSecurable()).isNotNull(); // negative tests, load an entity which does not exist - lookup = cache.getOrLoadEntityById(this.callCtx, N1.getCatalogId(), 10000); + lookup = cache.getOrLoadEntityById(metaStore, N1.getCatalogId(), 10000); Assertions.assertThat(lookup).isNull(); lookup = cache.getOrLoadEntityByName( - this.callCtx, - new EntityCacheByNameKey(PolarisEntityType.CATALOG, "non_existant_catalog")); + metaStore, new EntityCacheByNameKey(PolarisEntityType.CATALOG, "non_existant_catalog")); Assertions.assertThat(lookup).isNull(); // lookup N2 to validate grants EntityCacheByNameKey N2_name = new EntityCacheByNameKey(catalog.getId(), N1.getId(), PolarisEntityType.NAMESPACE, "N2"); - lookup = cache.getOrLoadEntityByName(callCtx, N2_name); + lookup = cache.getOrLoadEntityByName(metaStore, N2_name); Assertions.assertThat(lookup).isNotNull(); EntityCacheEntry cacheEntry_N1 = lookup.getCacheEntry(); Assertions.assertThat(cacheEntry_N1).isNotNull(); @@ -190,7 +187,7 @@ void testGetOrLoadEntityByName() { EntityCacheByNameKey R1_name = new EntityCacheByNameKey( catalog.getId(), catalog.getId(), PolarisEntityType.CATALOG_ROLE, "R1"); - lookup = cache.getOrLoadEntityByName(callCtx, R1_name); + lookup = cache.getOrLoadEntityByName(metaStore, R1_name); Assertions.assertThat(lookup).isNotNull(); EntityCacheEntry cacheEntry_R1 = lookup.getCacheEntry(); Assertions.assertThat(cacheEntry_R1).isNotNull(); @@ -234,7 +231,7 @@ void testGetOrLoadEntityByName() { // lookup principal role PR1 EntityCacheByNameKey PR1_name = new EntityCacheByNameKey(PolarisEntityType.PRINCIPAL_ROLE, "PR1"); - lookup = cache.getOrLoadEntityByName(callCtx, PR1_name); + lookup = cache.getOrLoadEntityByName(metaStore, PR1_name); Assertions.assertThat(lookup).isNotNull(); EntityCacheEntry cacheEntry_PR1 = lookup.getCacheEntry(); Assertions.assertThat(cacheEntry_PR1).isNotNull(); @@ -275,7 +272,7 @@ void testRefresh() { // should exist and no cache hit EntityCacheLookupResult lookup = cache.getOrLoadEntityByName( - this.callCtx, new EntityCacheByNameKey(PolarisEntityType.CATALOG, "test")); + metaStore, new EntityCacheByNameKey(PolarisEntityType.CATALOG, "test")); Assertions.assertThat(lookup).isNotNull(); Assertions.assertThat(lookup.isCacheHit()).isFalse(); @@ -305,7 +302,7 @@ void testRefresh() { // now load that table in the cache cacheEntry = cache.getAndRefreshIfNeeded( - this.callCtx, T6v1, T6v1.getEntityVersion(), T6v1.getGrantRecordsVersion()); + metaStore, T6v1, T6v1.getEntityVersion(), T6v1.getGrantRecordsVersion()); Assertions.assertThat(cacheEntry).isNotNull(); Assertions.assertThat(cacheEntry.getEntity()).isNotNull(); Assertions.assertThat(cacheEntry.getGrantRecordsAsSecurable()).isNotNull(); @@ -326,7 +323,7 @@ void testRefresh() { // now refresh that entity. But because we don't change the versions, nothing should be reloaded cacheEntry = cache.getAndRefreshIfNeeded( - this.callCtx, T6v1, T6v1.getEntityVersion(), T6v1.getGrantRecordsVersion()); + metaStore, T6v1, T6v1.getEntityVersion(), T6v1.getGrantRecordsVersion()); Assertions.assertThat(cacheEntry).isNotNull(); Assertions.assertThat(cacheEntry.getEntity()).isNotNull(); Assertions.assertThat(cacheEntry.getGrantRecordsAsSecurable()).isNotNull(); @@ -338,7 +335,7 @@ void testRefresh() { // now refresh again, this time with the new versions. Should be reloaded cacheEntry = cache.getAndRefreshIfNeeded( - this.callCtx, T6v2, T6v2.getEntityVersion(), T6v2.getGrantRecordsVersion()); + metaStore, T6v2, T6v2.getEntityVersion(), T6v2.getGrantRecordsVersion()); Assertions.assertThat(cacheEntry).isNotNull(); Assertions.assertThat(cacheEntry.getEntity()).isNotNull(); Assertions.assertThat(cacheEntry.getGrantRecordsAsSecurable()).isNotNull(); @@ -367,7 +364,7 @@ void testRefresh() { // load that namespace cacheEntry = cache.getAndRefreshIfNeeded( - this.callCtx, N2, N2.getEntityVersion(), N2.getGrantRecordsVersion()); + metaStore, N2, N2.getEntityVersion(), N2.getGrantRecordsVersion()); // should have one single grant Assertions.assertThat(cacheEntry).isNotNull(); @@ -388,7 +385,7 @@ void testRefresh() { // the cache is outdated now lookup = cache.getOrLoadEntityByName( - this.callCtx, + metaStore, new EntityCacheByNameKey( catalog.getId(), N1.getId(), PolarisEntityType.NAMESPACE, "N2")); Assertions.assertThat(lookup).isNotNull(); @@ -403,7 +400,7 @@ void testRefresh() { // now refresh cacheEntry = cache.getAndRefreshIfNeeded( - this.callCtx, N2, N2v2.getEntityVersion(), N2v2.getGrantRecordsVersion()); + metaStore, N2, N2v2.getEntityVersion(), N2v2.getGrantRecordsVersion()); Assertions.assertThat(cacheEntry).isNotNull(); Assertions.assertThat(cacheEntry.getEntity()).isNotNull(); Assertions.assertThat(cacheEntry.getGrantRecordsAsSecurable()).isNotNull(); @@ -419,19 +416,19 @@ void testRenameAndCacheDestinationBeforeLoadingSource() { EntityCacheLookupResult lookup = cache.getOrLoadEntityByName( - this.callCtx, new EntityCacheByNameKey(PolarisEntityType.CATALOG, "test")); + metaStore, new EntityCacheByNameKey(PolarisEntityType.CATALOG, "test")); Assertions.assertThat(lookup).isNotNull(); Assertions.assertThat(lookup.getCacheEntry()).isNotNull(); PolarisBaseEntity catalog = lookup.getCacheEntry().getEntity(); PolarisBaseEntity N1 = this.tm.ensureExistsByName(List.of(catalog), PolarisEntityType.NAMESPACE, "N1"); - lookup = cache.getOrLoadEntityById(this.callCtx, N1.getCatalogId(), N1.getId()); + lookup = cache.getOrLoadEntityById(metaStore, N1.getCatalogId(), N1.getId()); Assertions.assertThat(lookup).isNotNull(); EntityCacheByNameKey T4_name = new EntityCacheByNameKey(N1.getCatalogId(), N1.getId(), PolarisEntityType.TABLE_LIKE, "T4"); - lookup = cache.getOrLoadEntityByName(callCtx, T4_name); + lookup = cache.getOrLoadEntityByName(metaStore, T4_name); Assertions.assertThat(lookup).isNotNull(); EntityCacheEntry cacheEntry_T4 = lookup.getCacheEntry(); Assertions.assertThat(cacheEntry_T4).isNotNull(); @@ -446,7 +443,7 @@ void testRenameAndCacheDestinationBeforeLoadingSource() { EntityCacheByNameKey T4_renamed = new EntityCacheByNameKey( N1.getCatalogId(), N1.getId(), PolarisEntityType.TABLE_LIKE, "T4_renamed"); - lookup = cache.getOrLoadEntityByName(callCtx, T4_renamed); + lookup = cache.getOrLoadEntityByName(metaStore, T4_renamed); Assertions.assertThat(lookup).isNotNull(); EntityCacheEntry cacheEntry_T4_renamed = lookup.getCacheEntry(); Assertions.assertThat(cacheEntry_T4_renamed).isNotNull(); @@ -454,23 +451,23 @@ void testRenameAndCacheDestinationBeforeLoadingSource() { // new entry if lookup by id EntityCacheLookupResult lookupResult = - cache.getOrLoadEntityById(callCtx, T4.getCatalogId(), T4.getId()); + cache.getOrLoadEntityById(metaStore, T4.getCatalogId(), T4.getId()); Assertions.assertThat(lookupResult).isNotNull(); Assertions.assertThat(lookupResult.getCacheEntry()).isNotNull(); Assertions.assertThat(lookupResult.getCacheEntry().getEntity().getName()) .isEqualTo("T4_renamed"); // old name is gone, replaced by new name - // Assertions.assertNull(cache.getOrLoadEntityByName(callCtx, T4_name)); + // Assertions.assertNull(cache.getOrLoadEntityByName(metaStore, T4_name)); // refreshing should return null since our current held T4 is outdated cache.getAndRefreshIfNeeded( - callCtx, + metaStore, T4, T4_renamed_entity.getEntityVersion(), T4_renamed_entity.getGrantRecordsVersion()); // now the loading by the old name should return null - Assertions.assertThat(cache.getOrLoadEntityByName(callCtx, T4_name)).isNull(); + Assertions.assertThat(cache.getOrLoadEntityByName(metaStore, T4_name)).isNull(); } } diff --git a/polaris-core/src/test/java/org/apache/polaris/core/persistence/PolarisTreeMapMetaStoreManagerTest.java b/polaris-core/src/test/java/org/apache/polaris/core/persistence/PolarisTreeMapMetaStoreManagerTest.java index e44b45577..2a0a96aea 100644 --- a/polaris-core/src/test/java/org/apache/polaris/core/persistence/PolarisTreeMapMetaStoreManagerTest.java +++ b/polaris-core/src/test/java/org/apache/polaris/core/persistence/PolarisTreeMapMetaStoreManagerTest.java @@ -21,7 +21,6 @@ import static org.apache.polaris.core.persistence.PrincipalSecretsGenerator.RANDOM_SECRETS; import java.time.ZoneId; -import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.PolarisConfigurationStore; import org.apache.polaris.core.PolarisDefaultDiagServiceImpl; import org.apache.polaris.core.PolarisDiagnostics; @@ -32,13 +31,15 @@ public class PolarisTreeMapMetaStoreManagerTest extends BasePolarisMetaStoreMana public PolarisTestMetaStoreManager createPolarisTestMetaStoreManager() { PolarisDiagnostics diagServices = new PolarisDefaultDiagServiceImpl(); PolarisTreeMapStore store = new PolarisTreeMapStore(diagServices); - PolarisCallContext callCtx = - new PolarisCallContext( - new PolarisTreeMapMetaStoreSessionImpl(store, Mockito.mock(), RANDOM_SECRETS), + PolarisMetaStoreSession session = + new PolarisTreeMapMetaStoreSessionImpl(store, Mockito.mock(), RANDOM_SECRETS, diagServices); + return new PolarisTestMetaStoreManager( + new PolarisMetaStoreManagerImpl( + () -> "test", diagServices, new PolarisConfigurationStore() {}, - timeSource.withZone(ZoneId.systemDefault())); - - return new PolarisTestMetaStoreManager(new PolarisMetaStoreManagerImpl(), callCtx); + timeSource.withZone(ZoneId.systemDefault())), + session, + diagServices); } } diff --git a/polaris-core/src/test/java/org/apache/polaris/core/persistence/ResolverTest.java b/polaris-core/src/test/java/org/apache/polaris/core/persistence/ResolverTest.java index 9bc6cabfe..d30559719 100644 --- a/polaris-core/src/test/java/org/apache/polaris/core/persistence/ResolverTest.java +++ b/polaris-core/src/test/java/org/apache/polaris/core/persistence/ResolverTest.java @@ -24,6 +24,7 @@ import jakarta.annotation.Nullable; import jakarta.ws.rs.core.SecurityContext; import java.security.Principal; +import java.time.Clock; import java.util.ArrayList; import java.util.Comparator; import java.util.Iterator; @@ -32,7 +33,7 @@ import java.util.Optional; import java.util.Set; import java.util.stream.Collectors; -import org.apache.polaris.core.PolarisCallContext; +import org.apache.polaris.core.PolarisConfigurationStore; import org.apache.polaris.core.PolarisDefaultDiagServiceImpl; import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.auth.AuthenticatedPolarisPrincipal; @@ -59,15 +60,9 @@ public class ResolverTest { // diag services private final PolarisDiagnostics diagServices; - // the entity store, use treemap implementation - private final PolarisTreeMapStore store; - // to interact with the metastore private final PolarisMetaStoreSession metaStore; - // polaris call context - private final PolarisCallContext callCtx; - // utility to bootstrap the mata store private final PolarisTestMetaStoreManager tm; @@ -104,13 +99,16 @@ public class ResolverTest { */ public ResolverTest() { diagServices = new PolarisDefaultDiagServiceImpl(); - store = new PolarisTreeMapStore(diagServices); - metaStore = new PolarisTreeMapMetaStoreSessionImpl(store, Mockito.mock(), RANDOM_SECRETS); - callCtx = new PolarisCallContext(metaStore, diagServices); - metaStoreManager = new PolarisMetaStoreManagerImpl(); + // the entity store, use treemap implementation + PolarisTreeMapStore store = new PolarisTreeMapStore(diagServices); + metaStore = + new PolarisTreeMapMetaStoreSessionImpl(store, Mockito.mock(), RANDOM_SECRETS, diagServices); + metaStoreManager = + new PolarisMetaStoreManagerImpl( + () -> "test", diagServices, new PolarisConfigurationStore() {}, Clock.systemUTC()); // bootstrap the mata store with our test schema - tm = new PolarisTestMetaStoreManager(metaStoreManager, callCtx); + tm = new PolarisTestMetaStoreManager(metaStoreManager, metaStore, diagServices); tm.testCreateTestCatalog(); // principal P1 @@ -444,7 +442,7 @@ private Resolver allocateResolver( // create a new cache if needs be if (cache == null) { - this.cache = new EntityCache(this.metaStoreManager); + this.cache = new EntityCache(this.metaStoreManager, diagServices); } boolean allRoles = principalRolesScope == null; Optional> roleEntities = @@ -455,7 +453,7 @@ private Resolver allocateResolver( .map( role -> metaStoreManager.readEntityByName( - callCtx, + this.metaStore, null, PolarisEntityType.PRINCIPAL_ROLE, PolarisEntitySubType.NULL_SUBTYPE, @@ -468,7 +466,8 @@ private Resolver allocateResolver( new AuthenticatedPolarisPrincipal( PrincipalEntity.of(P1), Optional.ofNullable(principalRolesScope).orElse(Set.of())); return new Resolver( - this.callCtx, + this.metaStore, + this.diagServices, metaStoreManager, new SecurityContext() { @Override @@ -751,7 +750,7 @@ private Resolver resolveDriver( // see if the principal exists PolarisMetaStoreManager.EntityResult result = this.metaStoreManager.readEntityByName( - this.callCtx, + this.metaStore, null, PolarisEntityType.PRINCIPAL, PolarisEntitySubType.NULL_SUBTYPE, @@ -954,7 +953,7 @@ private void ensureResolved( // reload the cached entry from the backend CachedEntryResult refCachedEntry = this.metaStoreManager.loadCachedEntryById( - this.callCtx, refEntity.getCatalogId(), refEntity.getId()); + this.metaStore, refEntity.getCatalogId(), refEntity.getId()); // should exist Assertions.assertThat(refCachedEntry).isNotNull(); diff --git a/polaris-core/src/test/java/org/apache/polaris/core/storage/InMemoryStorageIntegrationTest.java b/polaris-core/src/test/java/org/apache/polaris/core/storage/InMemoryStorageIntegrationTest.java index aa5317e99..413cad5ee 100644 --- a/polaris-core/src/test/java/org/apache/polaris/core/storage/InMemoryStorageIntegrationTest.java +++ b/polaris-core/src/test/java/org/apache/polaris/core/storage/InMemoryStorageIntegrationTest.java @@ -20,28 +20,28 @@ import jakarta.annotation.Nonnull; import jakarta.annotation.Nullable; -import java.time.Clock; import java.util.EnumMap; import java.util.List; import java.util.Map; import java.util.Set; -import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.PolarisConfigurationStore; -import org.apache.polaris.core.PolarisDefaultDiagServiceImpl; import org.apache.polaris.core.PolarisDiagnostics; -import org.apache.polaris.core.context.CallContext; +import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.storage.aws.AwsStorageConfigurationInfo; import org.assertj.core.api.Assertions; import org.junit.jupiter.api.Test; -import org.mockito.Mockito; class InMemoryStorageIntegrationTest { + private RealmContext realmContext = () -> "test"; + @Test public void testValidateAccessToLocations() { - MockInMemoryStorageIntegration storage = new MockInMemoryStorageIntegration(); + MockInMemoryStorageIntegration storage = + new MockInMemoryStorageIntegration(new PolarisConfigurationStore() {}); Map> result = storage.validateAccessToLocations( + realmContext, new AwsStorageConfigurationInfo( PolarisStorageConfigurationInfo.StorageType.S3, List.of( @@ -91,64 +91,60 @@ public void testAwsAccountIdParsing() { @Test public void testValidateAccessToLocationsWithWildcard() { - MockInMemoryStorageIntegration storage = new MockInMemoryStorageIntegration(); Map config = Map.of("ALLOW_WILDCARD_LOCATION", true); - PolarisCallContext polarisCallContext = - new PolarisCallContext( - Mockito.mock(), - new PolarisDefaultDiagServiceImpl(), - new PolarisConfigurationStore() { - @SuppressWarnings("unchecked") - @Override - public @Nullable T getConfiguration(PolarisCallContext ctx, String configName) { - return (T) config.get(configName); - } - }, - Clock.systemUTC()); - try (CallContext ignored = - CallContext.setCurrentContext(CallContext.of(() -> "realm", polarisCallContext))) { - Map> result = - storage.validateAccessToLocations( - new FileStorageConfigurationInfo(List.of("file://", "*")), - Set.of(PolarisStorageActions.READ), - Set.of( - "s3://bucket/path/to/warehouse/namespace/table", - "file:///etc/passwd", - "a/relative/subdirectory")); - Assertions.assertThat(result) - .hasSize(3) - .hasEntrySatisfying( - "s3://bucket/path/to/warehouse/namespace/table", - val -> - Assertions.assertThat(val) - .hasSize(1) - .containsKey(PolarisStorageActions.READ) - .extractingByKey(PolarisStorageActions.READ) - .returns(true, PolarisStorageIntegration.ValidationResult::isSuccess)) - .hasEntrySatisfying( - "file:///etc/passwd", - val -> - Assertions.assertThat(val) - .hasSize(1) - .containsKey(PolarisStorageActions.READ) - .extractingByKey(PolarisStorageActions.READ) - .returns(true, PolarisStorageIntegration.ValidationResult::isSuccess)) - .hasEntrySatisfying( - "a/relative/subdirectory", - val -> - Assertions.assertThat(val) - .hasSize(1) - .containsKey(PolarisStorageActions.READ) - .extractingByKey(PolarisStorageActions.READ) - .returns(true, PolarisStorageIntegration.ValidationResult::isSuccess)); - } + PolarisConfigurationStore configurationStore = + new PolarisConfigurationStore() { + @SuppressWarnings("unchecked") + @Override + public @Nullable T getConfiguration(RealmContext realmContext, String configName) { + return (T) config.get(configName); + } + }; + MockInMemoryStorageIntegration storage = new MockInMemoryStorageIntegration(configurationStore); + Map> result = + storage.validateAccessToLocations( + realmContext, + new FileStorageConfigurationInfo(List.of("file://", "*")), + Set.of(PolarisStorageActions.READ), + Set.of( + "s3://bucket/path/to/warehouse/namespace/table", + "file:///etc/passwd", + "a/relative/subdirectory")); + Assertions.assertThat(result) + .hasSize(3) + .hasEntrySatisfying( + "s3://bucket/path/to/warehouse/namespace/table", + val -> + Assertions.assertThat(val) + .hasSize(1) + .containsKey(PolarisStorageActions.READ) + .extractingByKey(PolarisStorageActions.READ) + .returns(true, PolarisStorageIntegration.ValidationResult::isSuccess)) + .hasEntrySatisfying( + "file:///etc/passwd", + val -> + Assertions.assertThat(val) + .hasSize(1) + .containsKey(PolarisStorageActions.READ) + .extractingByKey(PolarisStorageActions.READ) + .returns(true, PolarisStorageIntegration.ValidationResult::isSuccess)) + .hasEntrySatisfying( + "a/relative/subdirectory", + val -> + Assertions.assertThat(val) + .hasSize(1) + .containsKey(PolarisStorageActions.READ) + .extractingByKey(PolarisStorageActions.READ) + .returns(true, PolarisStorageIntegration.ValidationResult::isSuccess)); } @Test public void testValidateAccessToLocationsNoAllowedLocations() { - MockInMemoryStorageIntegration storage = new MockInMemoryStorageIntegration(); + MockInMemoryStorageIntegration storage = + new MockInMemoryStorageIntegration(new PolarisConfigurationStore() {}); Map> result = storage.validateAccessToLocations( + realmContext, new AwsStorageConfigurationInfo( PolarisStorageConfigurationInfo.StorageType.S3, List.of(), @@ -180,9 +176,11 @@ public void testValidateAccessToLocationsNoAllowedLocations() { @Test public void testValidateAccessToLocationsWithPrefixOfAllowedLocation() { - MockInMemoryStorageIntegration storage = new MockInMemoryStorageIntegration(); + MockInMemoryStorageIntegration storage = + new MockInMemoryStorageIntegration(new PolarisConfigurationStore() {}); Map> result = storage.validateAccessToLocations( + realmContext, new AwsStorageConfigurationInfo( PolarisStorageConfigurationInfo.StorageType.S3, List.of("s3://bucket/path/to/warehouse"), @@ -202,12 +200,13 @@ public void testValidateAccessToLocationsWithPrefixOfAllowedLocation() { private static final class MockInMemoryStorageIntegration extends InMemoryStorageIntegration { - public MockInMemoryStorageIntegration() { - super(MockInMemoryStorageIntegration.class.getName()); + public MockInMemoryStorageIntegration(PolarisConfigurationStore configurationStore) { + super(configurationStore, MockInMemoryStorageIntegration.class.getName()); } @Override public EnumMap getSubscopedCreds( + @Nonnull RealmContext realmContext, @Nonnull PolarisDiagnostics diagnostics, @Nonnull PolarisStorageConfigurationInfo storageConfig, boolean allowListOperation, diff --git a/polaris-core/src/test/java/org/apache/polaris/core/storage/cache/StorageCredentialCacheTest.java b/polaris-core/src/test/java/org/apache/polaris/core/storage/cache/StorageCredentialCacheTest.java index d19bc41c8..06fbccd4b 100644 --- a/polaris-core/src/test/java/org/apache/polaris/core/storage/cache/StorageCredentialCacheTest.java +++ b/polaris-core/src/test/java/org/apache/polaris/core/storage/cache/StorageCredentialCacheTest.java @@ -18,7 +18,10 @@ */ package org.apache.polaris.core.storage.cache; +import static org.apache.polaris.core.PolarisConfiguration.STORAGE_CREDENTIAL_CACHE_DURATION_SECONDS; +import static org.apache.polaris.core.PolarisConfiguration.STORAGE_CREDENTIAL_DURATION_SECONDS; import static org.apache.polaris.core.persistence.PrincipalSecretsGenerator.RANDOM_SECRETS; +import static org.mockito.Mockito.when; import com.google.common.collect.ImmutableMap; import jakarta.annotation.Nonnull; @@ -29,7 +32,7 @@ import java.util.List; import java.util.Map; import org.apache.iceberg.exceptions.UnprocessableEntityException; -import org.apache.polaris.core.PolarisCallContext; +import org.apache.polaris.core.PolarisConfigurationStore; import org.apache.polaris.core.PolarisDefaultDiagServiceImpl; import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.entity.PolarisBaseEntity; @@ -53,41 +56,40 @@ public class StorageCredentialCacheTest { - // polaris call context - private final PolarisCallContext callCtx; - - // the meta store manager + private final PolarisDiagnostics diagServices; + private final PolarisMetaStoreSession metaStoreSession; private final PolarisMetaStoreManager metaStoreManager; - - StorageCredentialCache storageCredentialCache; + private final StorageCredentialCache storageCredentialCache; public StorageCredentialCacheTest() { // diag services - PolarisDiagnostics diagServices = new PolarisDefaultDiagServiceImpl(); + diagServices = new PolarisDefaultDiagServiceImpl(); // the entity store, use treemap implementation PolarisTreeMapStore store = new PolarisTreeMapStore(diagServices); // to interact with the metastore - PolarisMetaStoreSession metaStore = - new PolarisTreeMapMetaStoreSessionImpl(store, Mockito.mock(), RANDOM_SECRETS); - callCtx = new PolarisCallContext(metaStore, diagServices); + metaStoreSession = + new PolarisTreeMapMetaStoreSessionImpl(store, Mockito.mock(), RANDOM_SECRETS, diagServices); metaStoreManager = Mockito.mock(PolarisMetaStoreManagerImpl.class); - storageCredentialCache = new StorageCredentialCache(); + PolarisConfigurationStore configurationStore = Mockito.mock(PolarisConfigurationStore.class); + when(configurationStore.getConfiguration(null, STORAGE_CREDENTIAL_CACHE_DURATION_SECONDS)) + .thenReturn(300); + when(configurationStore.getConfiguration(null, STORAGE_CREDENTIAL_DURATION_SECONDS)) + .thenReturn(600); + storageCredentialCache = new StorageCredentialCache(diagServices, configurationStore); } @Test public void testBadResult() { - storageCredentialCache = new StorageCredentialCache(); ScopedCredentialsResult badResult = new ScopedCredentialsResult( BaseResult.ReturnStatus.SUBSCOPE_CREDS_ERROR, "extra_error_info"); - Mockito.when( - metaStoreManager.getSubscopedCredsForEntity( - Mockito.any(), - Mockito.anyLong(), - Mockito.anyLong(), - Mockito.anyBoolean(), - Mockito.anySet(), - Mockito.anySet())) + when(metaStoreManager.getSubscopedCredsForEntity( + Mockito.any(), + Mockito.anyLong(), + Mockito.anyLong(), + Mockito.anyBoolean(), + Mockito.anySet(), + Mockito.anySet())) .thenReturn(badResult); PolarisEntity polarisEntity = new PolarisEntity( @@ -97,7 +99,7 @@ public void testBadResult() { () -> storageCredentialCache.getOrGenerateSubScopeCreds( metaStoreManager, - callCtx, + metaStoreSession, polarisEntity, true, new HashSet<>(Arrays.asList("s3://bucket1/path")), @@ -108,17 +110,15 @@ public void testBadResult() { @Test public void testCacheHit() { - storageCredentialCache = new StorageCredentialCache(); List mockedScopedCreds = getFakeScopedCreds(3, /* expireSoon= */ false); - Mockito.when( - metaStoreManager.getSubscopedCredsForEntity( - Mockito.any(), - Mockito.anyLong(), - Mockito.anyLong(), - Mockito.anyBoolean(), - Mockito.anySet(), - Mockito.anySet())) + when(metaStoreManager.getSubscopedCredsForEntity( + Mockito.any(), + Mockito.anyLong(), + Mockito.anyLong(), + Mockito.anyBoolean(), + Mockito.anySet(), + Mockito.anySet())) .thenReturn(mockedScopedCreds.get(0)) .thenReturn(mockedScopedCreds.get(1)) .thenReturn(mockedScopedCreds.get(1)); @@ -130,7 +130,7 @@ public void testCacheHit() { // add an item to the cache storageCredentialCache.getOrGenerateSubScopeCreds( metaStoreManager, - callCtx, + metaStoreSession, polarisEntity, true, new HashSet<>(Arrays.asList("s3://bucket1/path", "s3://bucket2/path")), @@ -140,7 +140,7 @@ public void testCacheHit() { // subscope for the same entity and same allowed locations, will hit the cache storageCredentialCache.getOrGenerateSubScopeCreds( metaStoreManager, - callCtx, + metaStoreSession, polarisEntity, true, new HashSet<>(Arrays.asList("s3://bucket1/path", "s3://bucket2/path")), @@ -150,17 +150,15 @@ public void testCacheHit() { @RepeatedTest(10) public void testCacheEvict() throws InterruptedException { - storageCredentialCache = new StorageCredentialCache(); List mockedScopedCreds = getFakeScopedCreds(3, /* expireSoon= */ true); - Mockito.when( - metaStoreManager.getSubscopedCredsForEntity( - Mockito.any(), - Mockito.anyLong(), - Mockito.anyLong(), - Mockito.anyBoolean(), - Mockito.anySet(), - Mockito.anySet())) + when(metaStoreManager.getSubscopedCredsForEntity( + Mockito.any(), + Mockito.anyLong(), + Mockito.anyLong(), + Mockito.anyBoolean(), + Mockito.anySet(), + Mockito.anySet())) .thenReturn(mockedScopedCreds.get(0)) .thenReturn(mockedScopedCreds.get(1)) .thenReturn(mockedScopedCreds.get(2)); @@ -173,13 +171,12 @@ public void testCacheEvict() throws InterruptedException { polarisEntity, true, new HashSet<>(Arrays.asList("s3://bucket1/path", "s3://bucket2/path")), - new HashSet<>(Arrays.asList("s3://bucket/path")), - callCtx); + new HashSet<>(Arrays.asList("s3://bucket/path"))); // the entry will be evicted immediately because the token is expired storageCredentialCache.getOrGenerateSubScopeCreds( metaStoreManager, - callCtx, + metaStoreSession, polarisEntity, true, new HashSet<>(Arrays.asList("s3://bucket1/path", "s3://bucket2/path")), @@ -188,7 +185,7 @@ public void testCacheEvict() throws InterruptedException { storageCredentialCache.getOrGenerateSubScopeCreds( metaStoreManager, - callCtx, + metaStoreSession, polarisEntity, true, new HashSet<>(Arrays.asList("s3://bucket1/path", "s3://bucket2/path")), @@ -197,7 +194,7 @@ public void testCacheEvict() throws InterruptedException { storageCredentialCache.getOrGenerateSubScopeCreds( metaStoreManager, - callCtx, + metaStoreSession, polarisEntity, true, new HashSet<>(Arrays.asList("s3://bucket1/path", "s3://bucket2/path")), @@ -207,17 +204,15 @@ public void testCacheEvict() throws InterruptedException { @Test public void testCacheGenerateNewEntries() { - storageCredentialCache = new StorageCredentialCache(); List mockedScopedCreds = getFakeScopedCreds(3, /* expireSoon= */ false); - Mockito.when( - metaStoreManager.getSubscopedCredsForEntity( - Mockito.any(), - Mockito.anyLong(), - Mockito.anyLong(), - Mockito.anyBoolean(), - Mockito.anySet(), - Mockito.anySet())) + when(metaStoreManager.getSubscopedCredsForEntity( + Mockito.any(), + Mockito.anyLong(), + Mockito.anyLong(), + Mockito.anyBoolean(), + Mockito.anySet(), + Mockito.anySet())) .thenReturn(mockedScopedCreds.get(0)) .thenReturn(mockedScopedCreds.get(1)) .thenReturn(mockedScopedCreds.get(2)); @@ -227,7 +222,7 @@ public void testCacheGenerateNewEntries() { for (PolarisEntity entity : entityList) { storageCredentialCache.getOrGenerateSubScopeCreds( metaStoreManager, - callCtx, + metaStoreSession, entity, true, new HashSet<>(Arrays.asList("s3://bucket1/path", "s3://bucket2/path")), @@ -241,10 +236,10 @@ public void testCacheGenerateNewEntries() { internalMap.put( PolarisEntityConstants.getStorageConfigInfoPropertyName(), "newStorageConfig"); entity.setInternalProperties( - PolarisObjectMapperUtil.serializeProperties(callCtx, internalMap)); + PolarisObjectMapperUtil.serializeProperties(diagServices, internalMap)); storageCredentialCache.getOrGenerateSubScopeCreds( metaStoreManager, - callCtx, + metaStoreSession, entity, /* allowedListAction= */ true, new HashSet<>(Arrays.asList("s3://bucket1/path", "s3://bucket2/path")), @@ -255,7 +250,7 @@ public void testCacheGenerateNewEntries() { for (PolarisEntity entity : entityList) { storageCredentialCache.getOrGenerateSubScopeCreds( metaStoreManager, - callCtx, + metaStoreSession, entity, /* allowedListAction= */ false, new HashSet<>(Arrays.asList("s3://bucket1/path", "s3://bucket2/path")), @@ -266,7 +261,7 @@ public void testCacheGenerateNewEntries() { for (PolarisEntity entity : entityList) { storageCredentialCache.getOrGenerateSubScopeCreds( metaStoreManager, - callCtx, + metaStoreSession, entity, /* allowedListAction= */ false, new HashSet<>(Arrays.asList("s3://bucket1/path", "s3://bucket2/path")), @@ -279,10 +274,10 @@ public void testCacheGenerateNewEntries() { internalMap.put( PolarisEntityConstants.getStorageConfigInfoPropertyName(), "newStorageConfig"); entity.setInternalProperties( - PolarisObjectMapperUtil.serializeProperties(callCtx, internalMap)); + PolarisObjectMapperUtil.serializeProperties(diagServices, internalMap)); storageCredentialCache.getOrGenerateSubScopeCreds( metaStoreManager, - callCtx, + metaStoreSession, entity, /* allowedListAction= */ false, new HashSet<>(Arrays.asList("s3://differentbucket/path", "s3://bucket2/path")), @@ -293,18 +288,16 @@ public void testCacheGenerateNewEntries() { @Test public void testCacheNotAffectedBy() { - storageCredentialCache = new StorageCredentialCache(); List mockedScopedCreds = getFakeScopedCreds(3, /* expireSoon= */ false); - Mockito.when( - metaStoreManager.getSubscopedCredsForEntity( - Mockito.any(), - Mockito.anyLong(), - Mockito.anyLong(), - Mockito.anyBoolean(), - Mockito.anySet(), - Mockito.anySet())) + when(metaStoreManager.getSubscopedCredsForEntity( + Mockito.any(), + Mockito.anyLong(), + Mockito.anyLong(), + Mockito.anyBoolean(), + Mockito.anySet(), + Mockito.anySet())) .thenReturn(mockedScopedCreds.get(0)) .thenReturn(mockedScopedCreds.get(1)) .thenReturn(mockedScopedCreds.get(2)); @@ -312,7 +305,7 @@ public void testCacheNotAffectedBy() { for (PolarisEntity entity : entityList) { storageCredentialCache.getOrGenerateSubScopeCreds( metaStoreManager, - callCtx, + metaStoreSession, entity, true, new HashSet<>(Arrays.asList("s3://bucket1/path", "s3://bucket2/path")), @@ -325,7 +318,7 @@ public void testCacheNotAffectedBy() { entity.setId(1234); storageCredentialCache.getOrGenerateSubScopeCreds( metaStoreManager, - callCtx, + metaStoreSession, entity, true, new HashSet<>(Arrays.asList("s3://bucket1/path", "s3://bucket2/path")), @@ -338,7 +331,7 @@ public void testCacheNotAffectedBy() { entity.setEntityVersion(5); storageCredentialCache.getOrGenerateSubScopeCreds( metaStoreManager, - callCtx, + metaStoreSession, entity, true, new HashSet<>(Arrays.asList("s3://bucket1/path", "s3://bucket2/path")), @@ -350,7 +343,7 @@ public void testCacheNotAffectedBy() { entity.setEntityVersion(5); storageCredentialCache.getOrGenerateSubScopeCreds( metaStoreManager, - callCtx, + metaStoreSession, entity, true, new HashSet<>(Arrays.asList("s3://bucket2/path", "s3://bucket1/path")), @@ -363,7 +356,7 @@ public void testCacheNotAffectedBy() { entity.setEntityVersion(5); storageCredentialCache.getOrGenerateSubScopeCreds( metaStoreManager, - callCtx, + metaStoreSession, entity, true, new HashSet<>(Arrays.asList("s3://bucket2/path", "s3://bucket1/path")), diff --git a/polaris-core/src/test/java/org/apache/polaris/service/storage/PolarisConfigurationStoreTest.java b/polaris-core/src/test/java/org/apache/polaris/service/storage/PolarisConfigurationStoreTest.java index a0da2c5d8..6e9379775 100644 --- a/polaris-core/src/test/java/org/apache/polaris/service/storage/PolarisConfigurationStoreTest.java +++ b/polaris-core/src/test/java/org/apache/polaris/service/storage/PolarisConfigurationStoreTest.java @@ -20,14 +20,17 @@ import jakarta.annotation.Nullable; import java.util.List; -import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.PolarisConfiguration; import org.apache.polaris.core.PolarisConfigurationStore; +import org.apache.polaris.core.context.RealmContext; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; /** Unit test for the default behaviors of the PolarisConfigurationStore interface. */ public class PolarisConfigurationStoreTest { + + private RealmContext realmContext = () -> "test"; + @Test public void testConfigsCanBeCastedFromString() { List> configs = @@ -45,7 +48,7 @@ public void testConfigsCanBeCastedFromString() { */ @SuppressWarnings("unchecked") @Override - public @Nullable T getConfiguration(PolarisCallContext ctx, String configName) { + public @Nullable T getConfiguration(RealmContext realmContext, String configName) { for (PolarisConfiguration c : configs) { if (c.key.equals(configName)) { return (T) String.valueOf(c.defaultValue); @@ -62,7 +65,7 @@ public void testConfigsCanBeCastedFromString() { // Ensure that we can fetch all the configs and that the value is what we expect, which // is the config's default value based on how we've implemented PolarisConfigurationStore above. for (PolarisConfiguration c : configs) { - Assertions.assertEquals(c.defaultValue, store.getConfiguration(null, c)); + Assertions.assertEquals(c.defaultValue, store.getConfiguration(realmContext, c)); } } @@ -76,13 +79,14 @@ public void testInvalidCastThrowsException() { new PolarisConfigurationStore() { @SuppressWarnings("unchecked") @Override - public T getConfiguration(PolarisCallContext ctx, String configName) { + public T getConfiguration(RealmContext realmContext, String configName) { return (T) "abc123"; } }; for (PolarisConfiguration c : configs) { - Assertions.assertThrows(NumberFormatException.class, () -> store.getConfiguration(null, c)); + Assertions.assertThrows( + NumberFormatException.class, () -> store.getConfiguration(realmContext, c)); } } diff --git a/polaris-core/src/test/java/org/apache/polaris/service/storage/aws/AwsCredentialsStorageIntegrationTest.java b/polaris-core/src/test/java/org/apache/polaris/service/storage/aws/AwsCredentialsStorageIntegrationTest.java index c552494e5..bf563da45 100644 --- a/polaris-core/src/test/java/org/apache/polaris/service/storage/aws/AwsCredentialsStorageIntegrationTest.java +++ b/polaris-core/src/test/java/org/apache/polaris/service/storage/aws/AwsCredentialsStorageIntegrationTest.java @@ -24,7 +24,9 @@ import java.util.EnumMap; import java.util.List; import java.util.Set; +import org.apache.polaris.core.PolarisConfigurationStore; import org.apache.polaris.core.PolarisDiagnostics; +import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.storage.PolarisCredentialProperty; import org.apache.polaris.core.storage.PolarisStorageConfigurationInfo; import org.apache.polaris.core.storage.aws.AwsCredentialsStorageIntegration; @@ -59,6 +61,8 @@ class AwsCredentialsStorageIntegrationTest { .build(); public static final String AWS_PARTITION = "aws"; + private static final RealmContext REALM_CONTEXT = () -> "realm"; + @Test public void testGetSubscopedCreds() { StsClient stsClient = Mockito.mock(StsClient.class); @@ -76,8 +80,9 @@ public void testGetSubscopedCreds() { }); String warehouseDir = "s3://bucket/path/to/warehouse"; EnumMap credentials = - new AwsCredentialsStorageIntegration(stsClient) + new AwsCredentialsStorageIntegration(new PolarisConfigurationStore() {}, stsClient) .getSubscopedCreds( + REALM_CONTEXT, Mockito.mock(PolarisDiagnostics.class), new AwsStorageConfigurationInfo( PolarisStorageConfigurationInfo.StorageType.S3, @@ -214,8 +219,9 @@ public void testGetSubscopedCredsInlinePolicy(String awsPartition) { return ASSUME_ROLE_RESPONSE; }); EnumMap credentials = - new AwsCredentialsStorageIntegration(stsClient) + new AwsCredentialsStorageIntegration(new PolarisConfigurationStore() {}, stsClient) .getSubscopedCreds( + REALM_CONTEXT, Mockito.mock(PolarisDiagnostics.class), new AwsStorageConfigurationInfo( storageType, @@ -308,16 +314,17 @@ public void testGetSubscopedCredsInlinePolicyWithoutList() { PolarisStorageConfigurationInfo.StorageType storageType = PolarisStorageConfigurationInfo.StorageType.S3; EnumMap credentials = - new AwsCredentialsStorageIntegration(stsClient) + new AwsCredentialsStorageIntegration(new PolarisConfigurationStore() {}, stsClient) .getSubscopedCreds( + REALM_CONTEXT, Mockito.mock(PolarisDiagnostics.class), new AwsStorageConfigurationInfo( PolarisStorageConfigurationInfo.StorageType.S3, List.of(s3Path(bucket, warehouseKeyPrefix)), roleARN, externalId, - "us-east-2"), - false, /* allowList = false*/ + "us-east-2"), /* allowList = false*/ + false, Set.of(s3Path(bucket, firstPath), s3Path(bucket, secondPath)), Set.of(s3Path(bucket, firstPath))); assertThat(credentials) @@ -400,16 +407,17 @@ public void testGetSubscopedCredsInlinePolicyWithoutWrites() { PolarisStorageConfigurationInfo.StorageType storageType = PolarisStorageConfigurationInfo.StorageType.S3; EnumMap credentials = - new AwsCredentialsStorageIntegration(stsClient) + new AwsCredentialsStorageIntegration(new PolarisConfigurationStore() {}, stsClient) .getSubscopedCreds( + REALM_CONTEXT, Mockito.mock(PolarisDiagnostics.class), new AwsStorageConfigurationInfo( storageType, List.of(s3Path(bucket, warehouseKeyPrefix)), roleARN, externalId, - "us-east-2"), - true, /* allowList = true */ + "us-east-2"), /* allowList = true */ + true, Set.of(s3Path(bucket, firstPath), s3Path(bucket, secondPath)), Set.of()); assertThat(credentials) @@ -462,16 +470,17 @@ public void testGetSubscopedCredsInlinePolicyWithEmptyReadAndWrite() { return ASSUME_ROLE_RESPONSE; }); EnumMap credentials = - new AwsCredentialsStorageIntegration(stsClient) + new AwsCredentialsStorageIntegration(new PolarisConfigurationStore() {}, stsClient) .getSubscopedCreds( + REALM_CONTEXT, Mockito.mock(PolarisDiagnostics.class), new AwsStorageConfigurationInfo( PolarisStorageConfigurationInfo.StorageType.S3, List.of(s3Path(bucket, warehouseKeyPrefix)), roleARN, externalId, - "us-east-2"), - true, /* allowList = true */ + "us-east-2"), /* allowList = true */ + true, Set.of(), Set.of()); assertThat(credentials) @@ -495,16 +504,17 @@ public void testClientRegion() { return ASSUME_ROLE_RESPONSE; }); EnumMap credentials = - new AwsCredentialsStorageIntegration(stsClient) + new AwsCredentialsStorageIntegration(new PolarisConfigurationStore() {}, stsClient) .getSubscopedCreds( + REALM_CONTEXT, Mockito.mock(PolarisDiagnostics.class), new AwsStorageConfigurationInfo( PolarisStorageConfigurationInfo.StorageType.S3, List.of(s3Path(bucket, warehouseKeyPrefix)), roleARN, externalId, - clientRegion), - true, /* allowList = true */ + clientRegion), /* allowList = true */ + true, Set.of(), Set.of()); assertThat(credentials) @@ -525,16 +535,17 @@ public void testNoClientRegion() { return ASSUME_ROLE_RESPONSE; }); EnumMap credentials = - new AwsCredentialsStorageIntegration(stsClient) + new AwsCredentialsStorageIntegration(new PolarisConfigurationStore() {}, stsClient) .getSubscopedCreds( + REALM_CONTEXT, Mockito.mock(PolarisDiagnostics.class), new AwsStorageConfigurationInfo( PolarisStorageConfigurationInfo.StorageType.S3, List.of(s3Path(bucket, warehouseKeyPrefix)), roleARN, externalId, - null), - true, /* allowList = true */ + null), /* allowList = true */ + true, Set.of(), Set.of()); assertThat(credentials).isNotEmpty().doesNotContainKey(PolarisCredentialProperty.CLIENT_REGION); diff --git a/polaris-core/src/test/java/org/apache/polaris/service/storage/azure/AzureCredentialStorageIntegrationTest.java b/polaris-core/src/test/java/org/apache/polaris/service/storage/azure/AzureCredentialStorageIntegrationTest.java index fd7ac9a88..8c366ee6a 100644 --- a/polaris-core/src/test/java/org/apache/polaris/service/storage/azure/AzureCredentialStorageIntegrationTest.java +++ b/polaris-core/src/test/java/org/apache/polaris/service/storage/azure/AzureCredentialStorageIntegrationTest.java @@ -47,7 +47,9 @@ import java.util.List; import java.util.Map; import java.util.stream.Stream; +import org.apache.polaris.core.PolarisConfigurationStore; import org.apache.polaris.core.PolarisDefaultDiagServiceImpl; +import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.storage.PolarisCredentialProperty; import org.apache.polaris.core.storage.azure.AzureCredentialsStorageIntegration; import org.apache.polaris.core.storage.azure.AzureStorageConfigurationInfo; @@ -66,6 +68,7 @@ public class AzureCredentialStorageIntegrationTest { private final String clientId = System.getenv("AZURE_CLIENT_ID"); private final String clientSecret = System.getenv("AZURE_CLIENT_SECRET"); private final String tenantId = System.getenv("AZURE_TENANT_ID"); + private final RealmContext realmContext = () -> "realm"; private void assumeEnvVariablesNotNull() { Assumptions.assumeThat( @@ -85,7 +88,10 @@ public void testNegativeCases() { Assertions.assertThatThrownBy( () -> subscopedCredsForOperations( - differentEndpointList, /* allowedWriteLoc= */ new ArrayList<>(), true)) + realmContext, + differentEndpointList, + /* allowedWriteLoc= */ new ArrayList<>(), + true)) .isInstanceOf(RuntimeException.class); List differentStorageAccts = @@ -95,7 +101,10 @@ public void testNegativeCases() { Assertions.assertThatThrownBy( () -> subscopedCredsForOperations( - differentStorageAccts, /* allowedWriteLoc= */ new ArrayList<>(), true)) + realmContext, + differentStorageAccts, + /* allowedWriteLoc= */ new ArrayList<>(), + true)) .isInstanceOf(RuntimeException.class); List differentContainers = Arrays.asList( @@ -105,7 +114,10 @@ public void testNegativeCases() { Assertions.assertThatThrownBy( () -> subscopedCredsForOperations( - differentContainers, /* allowedWriteLoc= */ new ArrayList<>(), true)) + realmContext, + differentContainers, + /* allowedWriteLoc= */ new ArrayList<>(), + true)) .isInstanceOf(RuntimeException.class); } @@ -122,7 +134,8 @@ public void testGetSubscopedTokenList(boolean allowListAction, String service) { service)); Map credsMap = subscopedCredsForOperations( - /* allowedReadLoc= */ allowedLoc, + /* allowedReadLoc= */ realmContext, + allowedLoc, /* allowedWriteLoc= */ new ArrayList<>(), allowListAction); Assertions.assertThat(credsMap).hasSize(2); @@ -193,7 +206,8 @@ public void testGetSubscopedTokenRead( service, allowedPrefix)); Map credsMap = subscopedCredsForOperations( - /* allowedReadLoc= */ allowedLoc, + /* allowedReadLoc= */ realmContext, + allowedLoc, /* allowedWriteLoc= */ new ArrayList<>(), /* allowListAction= */ false); @@ -263,7 +277,8 @@ public void testGetSubscopedTokenWrite( service, allowedPrefix)); Map credsMap = subscopedCredsForOperations( - /* allowedReadLoc= */ new ArrayList<>(), + /* allowedReadLoc= */ realmContext, + new ArrayList<>(), /* allowedWriteLoc= */ allowedLoc, /* allowListAction= */ false); String serviceEndpoint = @@ -339,16 +354,20 @@ public void testGetSubscopedTokenWrite( } private Map subscopedCredsForOperations( - List allowedReadLoc, List allowedWriteLoc, boolean allowListAction) { + RealmContext realmContext, + List allowedReadLoc, + List allowedWriteLoc, + boolean allowListAction) { List allowedLoc = new ArrayList<>(); allowedLoc.addAll(allowedReadLoc); allowedLoc.addAll(allowedWriteLoc); AzureStorageConfigurationInfo azureConfig = new AzureStorageConfigurationInfo(allowedLoc, tenantId); AzureCredentialsStorageIntegration azureCredsIntegration = - new AzureCredentialsStorageIntegration(); + new AzureCredentialsStorageIntegration(new PolarisConfigurationStore() {}); EnumMap credsMap = azureCredsIntegration.getSubscopedCreds( + realmContext, new PolarisDefaultDiagServiceImpl(), azureConfig, allowListAction, diff --git a/polaris-core/src/test/java/org/apache/polaris/service/storage/gcp/GcpCredentialsStorageIntegrationTest.java b/polaris-core/src/test/java/org/apache/polaris/service/storage/gcp/GcpCredentialsStorageIntegrationTest.java index 35d05c251..17bf389f1 100644 --- a/polaris-core/src/test/java/org/apache/polaris/service/storage/gcp/GcpCredentialsStorageIntegrationTest.java +++ b/polaris-core/src/test/java/org/apache/polaris/service/storage/gcp/GcpCredentialsStorageIntegrationTest.java @@ -48,7 +48,9 @@ import java.util.List; import java.util.Map; import java.util.Set; +import org.apache.polaris.core.PolarisConfigurationStore; import org.apache.polaris.core.PolarisDefaultDiagServiceImpl; +import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.storage.PolarisCredentialProperty; import org.apache.polaris.core.storage.gcp.GcpCredentialsStorageIntegration; import org.apache.polaris.core.storage.gcp.GcpStorageConfigurationInfo; @@ -71,7 +73,7 @@ public void testSubscope(boolean allowedListAction) throws Exception { .describedAs("Environment variable GOOGLE_APPLICATION_CREDENTIALS not exits") .isNotNull() .isNotEmpty(); - + RealmContext realmContext = () -> "realm"; List allowedRead = Arrays.asList( "gs://sfc-dev1-regtest/polaris-test/subscoped-test/read1/", @@ -80,7 +82,8 @@ public void testSubscope(boolean allowedListAction) throws Exception { Arrays.asList( "gs://sfc-dev1-regtest/polaris-test/subscoped-test/write1/", "gs://sfc-dev1-regtest/polaris-test/subscoped-test/write2/"); - Storage storageClient = setupStorageClient(allowedRead, allowedWrite, allowedListAction); + Storage storageClient = + setupStorageClient(realmContext, allowedRead, allowedWrite, allowedListAction); BlobInfo blobInfoGoodWrite = createStorageBlob("sfc-dev1-regtest", "polaris-test/subscoped-test/write1/", "file.txt"); BlobInfo blobInfoBad = @@ -122,7 +125,8 @@ public void testSubscope(boolean allowedListAction) throws Exception { Arrays.asList( "gs://sfc-dev1-regtest/polaris-test/subscoped-test/write2/", "gs://sfc-dev1-regtest/polaris-test/subscoped-test/write3/"); - Storage clientForDelete = setupStorageClient(List.of(), allowedWrite2, allowedListAction); + Storage clientForDelete = + setupStorageClient(realmContext, List.of(), allowedWrite2, allowedListAction); // can not delete because it is not in allowed write path for this client Assertions.assertThatThrownBy(() -> clientForDelete.delete(blobInfoGoodWrite.getBlobId())) @@ -134,10 +138,13 @@ public void testSubscope(boolean allowedListAction) throws Exception { } private Storage setupStorageClient( - List allowedReadLoc, List allowedWriteLoc, boolean allowListAction) + RealmContext realmContext, + List allowedReadLoc, + List allowedWriteLoc, + boolean allowListAction) throws IOException { Map credsMap = - subscopedCredsForOperations(allowedReadLoc, allowedWriteLoc, allowListAction); + subscopedCredsForOperations(realmContext, allowedReadLoc, allowedWriteLoc, allowListAction); return createStorageClient(credsMap); } @@ -160,7 +167,10 @@ private Storage createStorageClient(Map creds } private Map subscopedCredsForOperations( - List allowedReadLoc, List allowedWriteLoc, boolean allowListAction) + RealmContext realmContext, + List allowedReadLoc, + List allowedWriteLoc, + boolean allowListAction) throws IOException { List allowedLoc = new ArrayList<>(); allowedLoc.addAll(allowedReadLoc); @@ -168,10 +178,12 @@ private Map subscopedCredsForOperations( GcpStorageConfigurationInfo gcpConfig = new GcpStorageConfigurationInfo(allowedLoc); GcpCredentialsStorageIntegration gcpCredsIntegration = new GcpCredentialsStorageIntegration( + new PolarisConfigurationStore() {}, GoogleCredentials.getApplicationDefault(), ServiceOptions.getFromServiceLoader(HttpTransportFactory.class, NetHttpTransport::new)); EnumMap credsMap = gcpCredsIntegration.getSubscopedCreds( + realmContext, new PolarisDefaultDiagServiceImpl(), gcpConfig, allowListAction, @@ -184,6 +196,7 @@ private Map subscopedCredsForOperations( public void testGenerateAccessBoundary() throws IOException { GcpCredentialsStorageIntegration integration = new GcpCredentialsStorageIntegration( + new PolarisConfigurationStore() {}, GoogleCredentials.newBuilder() .setAccessToken( new AccessToken( @@ -213,6 +226,7 @@ public void testGenerateAccessBoundary() throws IOException { public void testGenerateAccessBoundaryWithMultipleBuckets() throws IOException { GcpCredentialsStorageIntegration integration = new GcpCredentialsStorageIntegration( + new PolarisConfigurationStore() {}, GoogleCredentials.newBuilder() .setAccessToken( new AccessToken( @@ -247,6 +261,7 @@ public void testGenerateAccessBoundaryWithMultipleBuckets() throws IOException { public void testGenerateAccessBoundaryWithoutList() throws IOException { GcpCredentialsStorageIntegration integration = new GcpCredentialsStorageIntegration( + new PolarisConfigurationStore() {}, GoogleCredentials.newBuilder() .setAccessToken( new AccessToken( @@ -278,6 +293,7 @@ public void testGenerateAccessBoundaryWithoutList() throws IOException { public void testGenerateAccessBoundaryWithoutWrites() throws IOException { GcpCredentialsStorageIntegration integration = new GcpCredentialsStorageIntegration( + new PolarisConfigurationStore() {}, GoogleCredentials.newBuilder() .setAccessToken( new AccessToken( diff --git a/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/BasePolarisMetaStoreManagerTest.java b/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/BasePolarisMetaStoreManagerTest.java index 1b5483afe..92382cec5 100644 --- a/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/BasePolarisMetaStoreManagerTest.java +++ b/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/BasePolarisMetaStoreManagerTest.java @@ -34,8 +34,6 @@ import java.util.function.Function; import java.util.stream.Collectors; import java.util.stream.Stream; -import org.apache.polaris.core.PolarisCallContext; -import org.apache.polaris.core.context.CallContext; import org.apache.polaris.core.entity.AsyncTaskType; import org.apache.polaris.core.entity.PolarisBaseEntity; import org.apache.polaris.core.entity.PolarisEntity; @@ -103,144 +101,125 @@ void testBrowse() { @Test void testCreateEntities() { PolarisMetaStoreManager metaStoreManager = polarisTestMetaStoreManager.polarisMetaStoreManager; - try (CallContext callCtx = - CallContext.of(() -> "testRealm", polarisTestMetaStoreManager.polarisCallContext)) { - if (CallContext.getCurrentContext() == null) { - CallContext.setCurrentContext(callCtx); - } - TaskEntity task1 = createTask("task1", 100L); - TaskEntity task2 = createTask("task2", 101L); - List createdEntities = - metaStoreManager - .createEntitiesIfNotExist( - polarisTestMetaStoreManager.polarisCallContext, null, List.of(task1, task2)) - .getEntities(); - - Assertions.assertThat(createdEntities) - .isNotNull() - .hasSize(2) - .extracting(PolarisEntity::toCore) - .containsExactly(PolarisEntity.toCore(task1), PolarisEntity.toCore(task2)); - - List listedEntities = - metaStoreManager - .listEntities( - polarisTestMetaStoreManager.polarisCallContext, - null, - PolarisEntityType.TASK, - PolarisEntitySubType.NULL_SUBTYPE) - .getEntities(); - Assertions.assertThat(listedEntities) - .isNotNull() - .hasSize(2) - .containsExactly( - new PolarisEntityActiveRecord( - task1.getCatalogId(), - task1.getId(), - task1.getParentId(), - task1.getName(), - task1.getTypeCode(), - task1.getSubTypeCode()), - new PolarisEntityActiveRecord( - task2.getCatalogId(), - task2.getId(), - task2.getParentId(), - task2.getName(), - task2.getTypeCode(), - task2.getSubTypeCode())); - } + TaskEntity task1 = createTask("task1", 100L); + TaskEntity task2 = createTask("task2", 101L); + List createdEntities = + metaStoreManager + .createEntitiesIfNotExist( + polarisTestMetaStoreManager.polarisMetaStoreSession, null, List.of(task1, task2)) + .getEntities(); + + Assertions.assertThat(createdEntities) + .isNotNull() + .hasSize(2) + .extracting(PolarisEntity::toCore) + .containsExactly(PolarisEntity.toCore(task1), PolarisEntity.toCore(task2)); + + List listedEntities = + metaStoreManager + .listEntities( + polarisTestMetaStoreManager.polarisMetaStoreSession, + null, + PolarisEntityType.TASK, + PolarisEntitySubType.NULL_SUBTYPE) + .getEntities(); + Assertions.assertThat(listedEntities) + .isNotNull() + .hasSize(2) + .containsExactly( + new PolarisEntityActiveRecord( + task1.getCatalogId(), + task1.getId(), + task1.getParentId(), + task1.getName(), + task1.getTypeCode(), + task1.getSubTypeCode()), + new PolarisEntityActiveRecord( + task2.getCatalogId(), + task2.getId(), + task2.getParentId(), + task2.getName(), + task2.getTypeCode(), + task2.getSubTypeCode())); } @Test void testCreateEntitiesAlreadyExisting() { PolarisMetaStoreManager metaStoreManager = polarisTestMetaStoreManager.polarisMetaStoreManager; - try (CallContext callCtx = - CallContext.of(() -> "testRealm", polarisTestMetaStoreManager.polarisCallContext)) { - if (CallContext.getCurrentContext() == null) { - CallContext.setCurrentContext(callCtx); - } - TaskEntity task1 = createTask("task1", 100L); - TaskEntity task2 = createTask("task2", 101L); - List createdEntities = - metaStoreManager - .createEntitiesIfNotExist( - polarisTestMetaStoreManager.polarisCallContext, null, List.of(task1, task2)) - .getEntities(); - - Assertions.assertThat(createdEntities) - .isNotNull() - .hasSize(2) - .extracting(PolarisEntity::toCore) - .containsExactly(PolarisEntity.toCore(task1), PolarisEntity.toCore(task2)); - - TaskEntity task3 = createTask("task3", 103L); - - // entities task1 and task2 already exist with the same identifier, so the full list is - // returned - createdEntities = - metaStoreManager - .createEntitiesIfNotExist( - polarisTestMetaStoreManager.polarisCallContext, - null, - List.of(task1, task2, task3)) - .getEntities(); - Assertions.assertThat(createdEntities) - .isNotNull() - .hasSize(3) - .extracting(PolarisEntity::toCore) - .containsExactly( - PolarisEntity.toCore(task1), - PolarisEntity.toCore(task2), - PolarisEntity.toCore(task3)); - } + TaskEntity task1 = createTask("task1", 100L); + TaskEntity task2 = createTask("task2", 101L); + List createdEntities = + metaStoreManager + .createEntitiesIfNotExist( + polarisTestMetaStoreManager.polarisMetaStoreSession, null, List.of(task1, task2)) + .getEntities(); + + Assertions.assertThat(createdEntities) + .isNotNull() + .hasSize(2) + .extracting(PolarisEntity::toCore) + .containsExactly(PolarisEntity.toCore(task1), PolarisEntity.toCore(task2)); + + TaskEntity task3 = createTask("task3", 103L); + + // entities task1 and task2 already exist with the same identifier, so the full list is + // returned + createdEntities = + metaStoreManager + .createEntitiesIfNotExist( + polarisTestMetaStoreManager.polarisMetaStoreSession, + null, + List.of(task1, task2, task3)) + .getEntities(); + Assertions.assertThat(createdEntities) + .isNotNull() + .hasSize(3) + .extracting(PolarisEntity::toCore) + .containsExactly( + PolarisEntity.toCore(task1), PolarisEntity.toCore(task2), PolarisEntity.toCore(task3)); } @Test void testCreateEntitiesWithConflict() { PolarisMetaStoreManager metaStoreManager = polarisTestMetaStoreManager.polarisMetaStoreManager; - try (CallContext callCtx = - CallContext.of(() -> "testRealm", polarisTestMetaStoreManager.polarisCallContext)) { - if (CallContext.getCurrentContext() == null) { - CallContext.setCurrentContext(callCtx); - } - TaskEntity task1 = createTask("task1", 100L); - TaskEntity task2 = createTask("task2", 101L); - TaskEntity task3 = createTask("task3", 103L); - List createdEntities = - metaStoreManager - .createEntitiesIfNotExist( - polarisTestMetaStoreManager.polarisCallContext, - null, - List.of(task1, task2, task3)) - .getEntities(); - - Assertions.assertThat(createdEntities) - .isNotNull() - .hasSize(3) - .extracting(PolarisEntity::toCore) - .containsExactly( - PolarisEntity.toCore(task1), - PolarisEntity.toCore(task2), - PolarisEntity.toCore(task3)); - - TaskEntity secondTask3 = createTask("task3", 104L); - - TaskEntity task4 = createTask("task4", 105L); - createdEntities = - metaStoreManager - .createEntitiesIfNotExist( - polarisTestMetaStoreManager.polarisCallContext, null, List.of(secondTask3, task4)) - .getEntities(); - Assertions.assertThat(createdEntities).isNull(); - } + TaskEntity task1 = createTask("task1", 100L); + TaskEntity task2 = createTask("task2", 101L); + TaskEntity task3 = createTask("task3", 103L); + List createdEntities = + metaStoreManager + .createEntitiesIfNotExist( + polarisTestMetaStoreManager.polarisMetaStoreSession, + null, + List.of(task1, task2, task3)) + .getEntities(); + + Assertions.assertThat(createdEntities) + .isNotNull() + .hasSize(3) + .extracting(PolarisEntity::toCore) + .containsExactly( + PolarisEntity.toCore(task1), PolarisEntity.toCore(task2), PolarisEntity.toCore(task3)); + + TaskEntity secondTask3 = createTask("task3", 104L); + + TaskEntity task4 = createTask("task4", 105L); + createdEntities = + metaStoreManager + .createEntitiesIfNotExist( + polarisTestMetaStoreManager.polarisMetaStoreSession, + null, + List.of(secondTask3, task4)) + .getEntities(); + Assertions.assertThat(createdEntities).isNull(); } - private static TaskEntity createTask(String taskName, long id) { + private TaskEntity createTask(String taskName, long id) { return new TaskEntity.Builder() .setName(taskName) - .withData("data") + .withData(polarisTestMetaStoreManager.polarisDiagnostics, "data") .setId(id) - .withTaskType(AsyncTaskType.MANIFEST_FILE_CLEANUP) + .withTaskType( + polarisTestMetaStoreManager.polarisDiagnostics, AsyncTaskType.MANIFEST_FILE_CLEANUP) .setCreateTimestamp(Instant.now().toEpochMilli()) .build(); } @@ -288,9 +267,11 @@ void testLoadTasks() { } String executorId = "testExecutor_abc"; PolarisMetaStoreManager metaStoreManager = polarisTestMetaStoreManager.polarisMetaStoreManager; - PolarisCallContext callCtx = polarisTestMetaStoreManager.polarisCallContext; + List taskList = - metaStoreManager.loadTasks(callCtx, executorId, 5).getEntities(); + metaStoreManager + .loadTasks(polarisTestMetaStoreManager.polarisMetaStoreSession, executorId, 5) + .getEntities(); Assertions.assertThat(taskList) .isNotNull() .isNotEmpty() @@ -301,7 +282,7 @@ void testLoadTasks() { .extracting( e -> PolarisObjectMapperUtil.deserializeProperties( - callCtx, e.getProperties())) + polarisTestMetaStoreManager.polarisDiagnostics, e.getProperties())) .asInstanceOf(InstanceOfAssertFactories.map(String.class, String.class)) .containsEntry("lastAttemptExecutorId", executorId) .containsEntry("attemptCount", "1")); @@ -310,7 +291,9 @@ void testLoadTasks() { // grab a second round of tasks. Assert that none of the original 5 are in the list List newTaskList = - metaStoreManager.loadTasks(callCtx, executorId, 5).getEntities(); + metaStoreManager + .loadTasks(polarisTestMetaStoreManager.polarisMetaStoreSession, executorId, 5) + .getEntities(); Assertions.assertThat(newTaskList) .isNotNull() .isNotEmpty() @@ -324,7 +307,9 @@ void testLoadTasks() { // only 10 tasks are unassigned. Requesting 20, we should only receive those 10 List lastTen = - metaStoreManager.loadTasks(callCtx, executorId, 20).getEntities(); + metaStoreManager + .loadTasks(polarisTestMetaStoreManager.polarisMetaStoreSession, executorId, 20) + .getEntities(); Assertions.assertThat(lastTen) .isNotNull() @@ -338,7 +323,9 @@ void testLoadTasks() { .collect(Collectors.toSet()); List emtpyList = - metaStoreManager.loadTasks(callCtx, executorId, 20).getEntities(); + metaStoreManager + .loadTasks(polarisTestMetaStoreManager.polarisMetaStoreSession, executorId, 20) + .getEntities(); Assertions.assertThat(emtpyList).isNotNull().isEmpty(); @@ -346,7 +333,9 @@ void testLoadTasks() { // all the tasks are unassigned. Fetch them all List allTasks = - metaStoreManager.loadTasks(callCtx, executorId, 20).getEntities(); + metaStoreManager + .loadTasks(polarisTestMetaStoreManager.polarisMetaStoreSession, executorId, 20) + .getEntities(); Assertions.assertThat(allTasks) .isNotNull() @@ -357,11 +346,19 @@ void testLoadTasks() { // drop all the tasks. Skip the clock forward and fetch. empty list expected allTasks.forEach( - entity -> metaStoreManager.dropEntityIfExists(callCtx, null, entity, Map.of(), false)); + entity -> + metaStoreManager.dropEntityIfExists( + polarisTestMetaStoreManager.polarisMetaStoreSession, + null, + entity, + Map.of(), + false)); timeSource.add(Duration.ofMinutes(10)); List finalList = - metaStoreManager.loadTasks(callCtx, executorId, 20).getEntities(); + metaStoreManager + .loadTasks(polarisTestMetaStoreManager.polarisMetaStoreSession, executorId, 20) + .getEntities(); Assertions.assertThat(finalList).isNotNull().isEmpty(); } @@ -373,7 +370,6 @@ void testLoadTasksInParallel() throws Exception { null, PolarisEntityType.TASK, PolarisEntitySubType.NULL_SUBTYPE, "task_" + i); } PolarisMetaStoreManager metaStoreManager = polarisTestMetaStoreManager.polarisMetaStoreManager; - PolarisCallContext callCtx = polarisTestMetaStoreManager.polarisCallContext; List>> futureList = new ArrayList<>(); List> responses; ExecutorService executorService = Executors.newCachedThreadPool(); @@ -390,7 +386,13 @@ void testLoadTasksInParallel() throws Exception { do { retry = false; try { - taskList = metaStoreManager.loadTasks(callCtx, executorId, 5).getEntities(); + taskList = + metaStoreManager + .loadTasks( + polarisTestMetaStoreManager.polarisMetaStoreSession, + executorId, + 5) + .getEntities(); taskList.stream().map(PolarisBaseEntity::getName).forEach(taskNames::add); } catch (RetryOnConcurrencyException e) { retry = true; diff --git a/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/PolarisTestMetaStoreManager.java b/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/PolarisTestMetaStoreManager.java index 11ee4a87f..027d651a8 100644 --- a/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/PolarisTestMetaStoreManager.java +++ b/polaris-core/src/testFixtures/java/org/apache/polaris/core/persistence/PolarisTestMetaStoreManager.java @@ -28,7 +28,7 @@ import java.util.Map; import org.apache.commons.lang3.tuple.ImmutablePair; import org.apache.commons.lang3.tuple.Pair; -import org.apache.polaris.core.PolarisCallContext; +import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.auth.PolarisGrantManager.LoadGrantsResult; import org.apache.polaris.core.entity.PolarisBaseEntity; import org.apache.polaris.core.entity.PolarisChangeTrackingVersions; @@ -49,10 +49,8 @@ /** Test the Polaris persistence layer */ public class PolarisTestMetaStoreManager { - // call context - final PolarisCallContext polarisCallContext; - - // call metastore manager + final PolarisMetaStoreSession polarisMetaStoreSession; + final PolarisDiagnostics polarisDiagnostics; final PolarisMetaStoreManager polarisMetaStoreManager; // the start time @@ -64,14 +62,17 @@ public class PolarisTestMetaStoreManager { // initialize the test public PolarisTestMetaStoreManager( - PolarisMetaStoreManager polarisMetaStoreManager, PolarisCallContext polarisCallContext) { - this.polarisCallContext = polarisCallContext; + PolarisMetaStoreManager polarisMetaStoreManager, + PolarisMetaStoreSession polarisMetaStoreSession, + PolarisDiagnostics polarisDiagnostics) { this.polarisMetaStoreManager = polarisMetaStoreManager; + this.polarisMetaStoreSession = polarisMetaStoreSession; + this.polarisDiagnostics = polarisDiagnostics; this.doRetry = false; // bootstrap the Polaris service - polarisMetaStoreManager.purge(polarisCallContext); - polarisMetaStoreManager.bootstrapPolarisService(polarisCallContext); + polarisMetaStoreManager.purge(polarisMetaStoreSession); + polarisMetaStoreManager.bootstrapPolarisService(polarisMetaStoreSession); } public void forceRetry() { @@ -113,7 +114,7 @@ private PolarisBaseEntity ensureExistsById( // make sure this entity was persisted PolarisBaseEntity entity = polarisMetaStoreManager - .loadEntity(this.polarisCallContext, catalogId, entityId) + .loadEntity(polarisMetaStoreSession, catalogId, entityId) .getEntity(); // assert all expected values @@ -137,7 +138,7 @@ private PolarisBaseEntity ensureExistsById( // we should find it PolarisMetaStoreManager.EntityResult result = polarisMetaStoreManager.readEntityByName( - this.polarisCallContext, catalogPath, expectedType, expectedSubType, expectedName); + polarisMetaStoreSession, catalogPath, expectedType, expectedSubType, expectedName); // should be success, nothing changed Assertions.assertThat(result).isNotNull(); @@ -154,7 +155,7 @@ private PolarisBaseEntity ensureExistsById( // we should not find it PolarisMetaStoreManager.EntityResult result = polarisMetaStoreManager.readEntityByName( - this.polarisCallContext, catalogPath, expectedType, expectedSubType, expectedName); + polarisMetaStoreSession, catalogPath, expectedType, expectedSubType, expectedName); // lookup must be success, nothing changed Assertions.assertThat(result).isNotNull(); @@ -244,12 +245,12 @@ void ensureGrantRecordExists( // re-load both entities, ensure not null securable = polarisMetaStoreManager - .loadEntity(this.polarisCallContext, securable.getCatalogId(), securable.getId()) + .loadEntity(polarisMetaStoreSession, securable.getCatalogId(), securable.getId()) .getEntity(); Assertions.assertThat(securable).isNotNull(); grantee = polarisMetaStoreManager - .loadEntity(this.polarisCallContext, grantee.getCatalogId(), grantee.getId()) + .loadEntity(polarisMetaStoreSession, grantee.getCatalogId(), grantee.getId()) .getEntity(); Assertions.assertThat(grantee).isNotNull(); @@ -259,7 +260,7 @@ void ensureGrantRecordExists( // load all grant records on that securable, should not fail LoadGrantsResult loadGrantsOnSecurable = polarisMetaStoreManager.loadGrantsOnSecurable( - this.polarisCallContext, securable.getCatalogId(), securable.getId()); + polarisMetaStoreSession, securable.getCatalogId(), securable.getId()); // ensure entities for these grant records have been properly loaded this.validateLoadedGrants(loadGrantsOnSecurable, false); @@ -269,7 +270,7 @@ void ensureGrantRecordExists( // load all grant records on that grantee, should not fail LoadGrantsResult loadGrantsOnGrantee = polarisMetaStoreManager.loadGrantsToGrantee( - this.polarisCallContext, grantee.getCatalogId(), grantee.getId()); + polarisMetaStoreSession, grantee.getCatalogId(), grantee.getId()); // ensure entities for these grant records have been properly loaded this.validateLoadedGrants(loadGrantsOnGrantee, true); @@ -302,7 +303,7 @@ private void validateLoadedGrants(LoadGrantsResult loadGrantRecords, boolean isG // load that entity PolarisBaseEntity entity = polarisMetaStoreManager - .loadEntity(this.polarisCallContext, catalogId, entityId) + .loadEntity(polarisMetaStoreSession, catalogId, entityId) .getEntity(); Assertions.assertThat(entity).isNotNull(); Assertions.assertThat(entities.get(entityId)).isEqualTo(entity); @@ -321,12 +322,12 @@ void ensureGrantRecordRemoved( // re-load both entities, ensure not null securable = polarisMetaStoreManager - .loadEntity(this.polarisCallContext, securable.getCatalogId(), securable.getId()) + .loadEntity(polarisMetaStoreSession, securable.getCatalogId(), securable.getId()) .getEntity(); Assertions.assertThat(securable).isNotNull(); grantee = polarisMetaStoreManager - .loadEntity(this.polarisCallContext, grantee.getCatalogId(), grantee.getId()) + .loadEntity(polarisMetaStoreSession, grantee.getCatalogId(), grantee.getId()) .getEntity(); Assertions.assertThat(grantee).isNotNull(); @@ -336,7 +337,7 @@ void ensureGrantRecordRemoved( // load all grant records on that securable, should not fail LoadGrantsResult loadGrantsOnSecurable = polarisMetaStoreManager.loadGrantsOnSecurable( - this.polarisCallContext, securable.getCatalogId(), securable.getId()); + polarisMetaStoreSession, securable.getCatalogId(), securable.getId()); // ensure entities for these grant records have been properly loaded this.validateLoadedGrants(loadGrantsOnSecurable, false); @@ -346,7 +347,7 @@ void ensureGrantRecordRemoved( // load all grant records on that grantee, should not fail LoadGrantsResult loadGrantsOnGrantee = polarisMetaStoreManager.loadGrantsToGrantee( - this.polarisCallContext, grantee.getCatalogId(), grantee.getId()); + polarisMetaStoreSession, grantee.getCatalogId(), grantee.getId()); this.validateLoadedGrants(loadGrantsOnGrantee, true); // check that the grant record has been removed @@ -359,17 +360,17 @@ PolarisBaseEntity createPrincipal(String name) { PolarisBaseEntity principalEntity = new PolarisBaseEntity( PolarisEntityConstants.getNullId(), - polarisMetaStoreManager.generateNewEntityId(this.polarisCallContext).getId(), + polarisMetaStoreManager.generateNewEntityId(polarisMetaStoreSession).getId(), PolarisEntityType.PRINCIPAL, PolarisEntitySubType.NULL_SUBTYPE, PolarisEntityConstants.getRootEntityId(), name); principalEntity.setInternalProperties( PolarisObjectMapperUtil.serializeProperties( - this.polarisCallContext, + polarisDiagnostics, Map.of(PolarisEntityConstants.PRINCIPAL_CREDENTIAL_ROTATION_REQUIRED_STATE, "true"))); PolarisMetaStoreManager.CreatePrincipalResult createPrincipalResult = - polarisMetaStoreManager.createPrincipal(this.polarisCallContext, principalEntity); + polarisMetaStoreManager.createPrincipal(polarisMetaStoreSession, principalEntity); Assertions.assertThat(createPrincipalResult).isNotNull(); // ensure well created @@ -397,7 +398,7 @@ PolarisBaseEntity createPrincipal(String name) { // ensure that the secrets have been properly saved and match PolarisPrincipalSecrets reloadSecrets = polarisMetaStoreManager - .loadPrincipalSecrets(this.polarisCallContext, clientId) + .loadPrincipalSecrets(polarisMetaStoreSession, clientId) .getPrincipalSecrets(); Assertions.assertThat(reloadSecrets).isNotNull(); Assertions.assertThat(reloadSecrets.getPrincipalId()).isEqualTo(secrets.getPrincipalId()); @@ -409,7 +410,7 @@ PolarisBaseEntity createPrincipal(String name) { Map internalProperties = PolarisObjectMapperUtil.deserializeProperties( - this.polarisCallContext, createPrincipalResult.getPrincipal().getInternalProperties()); + polarisDiagnostics, createPrincipalResult.getPrincipal().getInternalProperties()); Assertions.assertThat( internalProperties.get( PolarisEntityConstants.PRINCIPAL_CREDENTIAL_ROTATION_REQUIRED_STATE)) @@ -419,7 +420,7 @@ PolarisBaseEntity createPrincipal(String name) { if (this.doRetry) { // simulate that we retried PolarisMetaStoreManager.CreatePrincipalResult newCreatePrincipalResult = - polarisMetaStoreManager.createPrincipal(this.polarisCallContext, principalEntity); + polarisMetaStoreManager.createPrincipal(polarisMetaStoreSession, principalEntity); Assertions.assertThat(newCreatePrincipalResult).isNotNull(); // ensure same @@ -436,7 +437,7 @@ PolarisBaseEntity createPrincipal(String name) { secrets = polarisMetaStoreManager .rotatePrincipalSecrets( - this.polarisCallContext, + polarisMetaStoreSession, clientId, principalEntity.getId(), false, @@ -446,11 +447,11 @@ PolarisBaseEntity createPrincipal(String name) { PolarisBaseEntity reloadPrincipal = polarisMetaStoreManager - .loadEntity(this.polarisCallContext, 0L, createPrincipalResult.getPrincipal().getId()) + .loadEntity(polarisMetaStoreSession, 0L, createPrincipalResult.getPrincipal().getId()) .getEntity(); internalProperties = PolarisObjectMapperUtil.deserializeProperties( - this.polarisCallContext, reloadPrincipal.getInternalProperties()); + polarisDiagnostics, reloadPrincipal.getInternalProperties()); Assertions.assertThat( internalProperties.get( PolarisEntityConstants.PRINCIPAL_CREDENTIAL_ROTATION_REQUIRED_STATE)) @@ -458,13 +459,13 @@ PolarisBaseEntity createPrincipal(String name) { // rotate the secrets, twice! polarisMetaStoreManager.rotatePrincipalSecrets( - this.polarisCallContext, + polarisMetaStoreSession, clientId, principalEntity.getId(), false, secrets.getMainSecretHash()); polarisMetaStoreManager.rotatePrincipalSecrets( - this.polarisCallContext, + polarisMetaStoreSession, clientId, principalEntity.getId(), false, @@ -473,7 +474,7 @@ PolarisBaseEntity createPrincipal(String name) { // reload and check that now the main should be secondary reloadSecrets = polarisMetaStoreManager - .loadPrincipalSecrets(this.polarisCallContext, clientId) + .loadPrincipalSecrets(polarisMetaStoreSession, clientId) .getPrincipalSecrets(); Assertions.assertThat(reloadSecrets).isNotNull(); Assertions.assertThat(reloadSecrets.getPrincipalId()).isEqualTo(secrets.getPrincipalId()); @@ -485,14 +486,14 @@ PolarisBaseEntity createPrincipal(String name) { // reset - the previous main secret is no longer one of the secrets polarisMetaStoreManager.rotatePrincipalSecrets( - this.polarisCallContext, + polarisMetaStoreSession, clientId, principalEntity.getId(), true, reloadSecrets.getMainSecretHash()); reloadSecrets = polarisMetaStoreManager - .loadPrincipalSecrets(this.polarisCallContext, clientId) + .loadPrincipalSecrets(polarisMetaStoreSession, clientId) .getPrincipalSecrets(); Assertions.assertThat(reloadSecrets).isNotNull(); Assertions.assertThat(reloadSecrets.getPrincipalId()).isEqualTo(secrets.getPrincipalId()); @@ -503,11 +504,11 @@ PolarisBaseEntity createPrincipal(String name) { PolarisBaseEntity newPrincipal = polarisMetaStoreManager - .loadEntity(this.polarisCallContext, 0L, principalEntity.getId()) + .loadEntity(polarisMetaStoreSession, 0L, principalEntity.getId()) .getEntity(); internalProperties = PolarisObjectMapperUtil.deserializeProperties( - this.polarisCallContext, newPrincipal.getInternalProperties()); + polarisDiagnostics, newPrincipal.getInternalProperties()); Assertions.assertThat( internalProperties.get( PolarisEntityConstants.PRINCIPAL_CREDENTIAL_ROTATION_REQUIRED_STATE)) @@ -516,14 +517,14 @@ PolarisBaseEntity createPrincipal(String name) { // reset again. we should get new secrets and the CREDENTIAL_ROTATION_REQUIRED flag should be // gone polarisMetaStoreManager.rotatePrincipalSecrets( - this.polarisCallContext, + polarisMetaStoreSession, clientId, principalEntity.getId(), true, reloadSecrets.getMainSecretHash()); PolarisPrincipalSecrets postResetCredentials = polarisMetaStoreManager - .loadPrincipalSecrets(this.polarisCallContext, clientId) + .loadPrincipalSecrets(polarisMetaStoreSession, clientId) .getPrincipalSecrets(); Assertions.assertThat(reloadSecrets).isNotNull(); Assertions.assertThat(postResetCredentials.getPrincipalId()) @@ -537,11 +538,11 @@ PolarisBaseEntity createPrincipal(String name) { PolarisBaseEntity finalPrincipal = polarisMetaStoreManager - .loadEntity(this.polarisCallContext, 0L, principalEntity.getId()) + .loadEntity(polarisMetaStoreSession, 0L, principalEntity.getId()) .getEntity(); internalProperties = PolarisObjectMapperUtil.deserializeProperties( - this.polarisCallContext, finalPrincipal.getInternalProperties()); + polarisDiagnostics, finalPrincipal.getInternalProperties()); Assertions.assertThat( internalProperties.get( PolarisEntityConstants.PRINCIPAL_CREDENTIAL_ROTATION_REQUIRED_STATE)) @@ -562,7 +563,7 @@ public PolarisBaseEntity createEntity( entityType, entitySubType, name, - polarisMetaStoreManager.generateNewEntityId(this.polarisCallContext).getId()); + polarisMetaStoreManager.generateNewEntityId(polarisMetaStoreSession).getId()); } PolarisBaseEntity createEntity( @@ -584,7 +585,7 @@ PolarisBaseEntity createEntity( new PolarisBaseEntity(catalogId, entityId, entityType, entitySubType, parentId, name); PolarisBaseEntity entity = polarisMetaStoreManager - .createEntityIfNotExists(this.polarisCallContext, catalogPath, newEntity) + .createEntityIfNotExists(polarisMetaStoreSession, catalogPath, newEntity) .getEntity(); Assertions.assertThat(entity).isNotNull(); @@ -598,7 +599,7 @@ PolarisBaseEntity createEntity( if (this.doRetry) { PolarisBaseEntity retryEntity = polarisMetaStoreManager - .createEntityIfNotExists(this.polarisCallContext, catalogPath, newEntity) + .createEntityIfNotExists(polarisMetaStoreSession, catalogPath, newEntity) .getEntity(); Assertions.assertThat(retryEntity).isNotNull(); @@ -633,12 +634,12 @@ void dropEntity(List catalogPath, PolarisEntityCore entityToD // check if it exists PolarisBaseEntity entity = polarisMetaStoreManager - .loadEntity(this.polarisCallContext, entityToDrop.getCatalogId(), entityToDrop.getId()) + .loadEntity(polarisMetaStoreSession, entityToDrop.getCatalogId(), entityToDrop.getId()) .getEntity(); if (entity != null) { PolarisMetaStoreManager.EntityResult entityFound = polarisMetaStoreManager.readEntityByName( - this.polarisCallContext, + polarisMetaStoreSession, catalogPath, entity.getType(), entity.getSubType(), @@ -660,7 +661,7 @@ void dropEntity(List catalogPath, PolarisEntityCore entityToD List children = polarisMetaStoreManager .listEntities( - this.polarisCallContext, + polarisMetaStoreSession, path, PolarisEntityType.NAMESPACE, PolarisEntitySubType.NULL_SUBTYPE) @@ -670,7 +671,7 @@ void dropEntity(List catalogPath, PolarisEntityCore entityToD children = polarisMetaStoreManager .listEntities( - this.polarisCallContext, + polarisMetaStoreSession, path, PolarisEntityType.TABLE_LIKE, PolarisEntitySubType.ANY_SUBTYPE) @@ -680,7 +681,7 @@ void dropEntity(List catalogPath, PolarisEntityCore entityToD children = polarisMetaStoreManager .listEntities( - this.polarisCallContext, + polarisMetaStoreSession, path, PolarisEntityType.CATALOG_ROLE, PolarisEntitySubType.ANY_SUBTYPE) @@ -705,13 +706,13 @@ void dropEntity(List catalogPath, PolarisEntityCore entityToD new ArrayList<>( polarisMetaStoreManager .loadGrantsOnSecurable( - this.polarisCallContext, entity.getCatalogId(), entity.getId()) + polarisMetaStoreSession, entity.getCatalogId(), entity.getId()) .getEntities()); securableEntities = new ArrayList<>( polarisMetaStoreManager .loadGrantsToGrantee( - this.polarisCallContext, entity.getCatalogId(), entity.getId()) + polarisMetaStoreSession, entity.getCatalogId(), entity.getId()) .getEntities()); } else { granteeEntities = List.of(); @@ -723,7 +724,7 @@ void dropEntity(List catalogPath, PolarisEntityCore entityToD Map.of("taskId", String.valueOf(entity.getId()), "cleanupProperty", "cleanupValue"); PolarisMetaStoreManager.DropEntityResult dropResult = polarisMetaStoreManager.dropEntityIfExists( - this.polarisCallContext, catalogPath, entityToDrop, cleanupProperties, true); + polarisMetaStoreSession, catalogPath, entityToDrop, cleanupProperties, true); // should have been dropped if exists if (entityToDrop.cannotBeDroppedOrRenamed()) { @@ -741,23 +742,23 @@ void dropEntity(List catalogPath, PolarisEntityCore entityToD Assertions.assertThat(dropResult.getCleanupTaskId()).isNotNull(); PolarisBaseEntity cleanupTask = polarisMetaStoreManager - .loadEntity(this.polarisCallContext, 0L, dropResult.getCleanupTaskId()) + .loadEntity(polarisMetaStoreSession, 0L, dropResult.getCleanupTaskId()) .getEntity(); Assertions.assertThat(cleanupTask).isNotNull(); Assertions.assertThat(cleanupTask.getType()).isEqualTo(PolarisEntityType.TASK); Assertions.assertThat(cleanupTask.getInternalProperties()).isNotNull(); Map internalProperties = PolarisObjectMapperUtil.deserializeProperties( - polarisCallContext, cleanupTask.getInternalProperties()); + polarisDiagnostics, cleanupTask.getInternalProperties()); Assertions.assertThat(internalProperties).isEqualTo(cleanupProperties); Map properties = PolarisObjectMapperUtil.deserializeProperties( - polarisCallContext, cleanupTask.getProperties()); + polarisDiagnostics, cleanupTask.getProperties()); Assertions.assertThat(properties).isNotNull(); Assertions.assertThat(properties.get(PolarisTaskConstants.TASK_DATA)).isNotNull(); PolarisBaseEntity droppedEntity = PolarisObjectMapperUtil.deserialize( - polarisCallContext, + polarisDiagnostics, properties.get(PolarisTaskConstants.TASK_DATA), PolarisBaseEntity.class); Assertions.assertThat(droppedEntity).isNotNull(); @@ -770,7 +771,7 @@ void dropEntity(List catalogPath, PolarisEntityCore entityToD PolarisBaseEntity entityAfterDrop = polarisMetaStoreManager .loadEntity( - this.polarisCallContext, entityToDrop.getCatalogId(), entityToDrop.getId()) + polarisMetaStoreSession, entityToDrop.getCatalogId(), entityToDrop.getId()) .getEntity(); // ensure dropped @@ -780,7 +781,7 @@ void dropEntity(List catalogPath, PolarisEntityCore entityToD Assertions.assertThat(entity).isNotNull(); PolarisMetaStoreManager.EntityResult entityFound = polarisMetaStoreManager.readEntityByName( - this.polarisCallContext, + polarisMetaStoreSession, catalogPath, entity.getType(), entity.getSubType(), @@ -795,7 +796,7 @@ void dropEntity(List catalogPath, PolarisEntityCore entityToD for (PolarisBaseEntity connectedEntity : granteeEntities) { LoadGrantsResult grantResult = polarisMetaStoreManager.loadGrantsToGrantee( - this.polarisCallContext, connectedEntity.getCatalogId(), connectedEntity.getId()); + polarisMetaStoreSession, connectedEntity.getCatalogId(), connectedEntity.getId()); if (grantResult.isSuccess()) { long cnt = grantResult.getGrantRecords().stream() @@ -817,7 +818,7 @@ void dropEntity(List catalogPath, PolarisEntityCore entityToD for (PolarisBaseEntity connectedEntity : securableEntities) { LoadGrantsResult grantResult = polarisMetaStoreManager.loadGrantsOnSecurable( - this.polarisCallContext, connectedEntity.getCatalogId(), connectedEntity.getId()); + polarisMetaStoreSession, connectedEntity.getCatalogId(), connectedEntity.getId()); long cnt = grantResult.getGrantRecords().stream() .filter(gr -> gr.getGranteeId() == entityToDrop.getId()) @@ -835,7 +836,7 @@ void grantPrivilege( PolarisPrivilege priv) { // grant the privilege polarisMetaStoreManager.grantPrivilegeOnSecurableToRole( - this.polarisCallContext, role, catalogPath, securable, priv); + polarisMetaStoreSession, role, catalogPath, securable, priv); // now validate the privilege this.ensureGrantRecordExists(securable, role, priv); @@ -849,7 +850,7 @@ void revokePrivilege( PolarisPrivilege priv) { // grant the privilege polarisMetaStoreManager.revokePrivilegeOnSecurableFromRole( - this.polarisCallContext, role, catalogPath, securable, priv); + polarisMetaStoreSession, role, catalogPath, securable, priv); // now validate the privilege this.ensureGrantRecordRemoved(securable, role, priv); @@ -863,7 +864,7 @@ void grantToGrantee( PolarisPrivilege priv) { // grant the privilege polarisMetaStoreManager.grantUsageOnRoleToGrantee( - this.polarisCallContext, catalog, granted, grantee); + polarisMetaStoreSession, catalog, granted, grantee); // now validate the privilege this.ensureGrantRecordExists(granted, grantee, priv); @@ -877,7 +878,7 @@ void revokeToGrantee( PolarisPrivilege priv) { // revoked the privilege polarisMetaStoreManager.revokeUsageOnRoleFromGrantee( - this.polarisCallContext, catalog, granted, grantee); + polarisMetaStoreSession, catalog, granted, grantee); // now validate that the privilege is gone this.ensureGrantRecordRemoved(granted, grantee, priv); @@ -911,13 +912,13 @@ PolarisBaseEntity createTestCatalog(String catalogName) { PolarisBaseEntity catalog = new PolarisBaseEntity( PolarisEntityConstants.getNullId(), - polarisMetaStoreManager.generateNewEntityId(this.polarisCallContext).getId(), + polarisMetaStoreManager.generateNewEntityId(polarisMetaStoreSession).getId(), PolarisEntityType.CATALOG, PolarisEntitySubType.NULL_SUBTYPE, PolarisEntityConstants.getRootEntityId(), catalogName); PolarisMetaStoreManager.CreateCatalogResult catalogCreated = - polarisMetaStoreManager.createCatalog(this.polarisCallContext, catalog, List.of()); + polarisMetaStoreManager.createCatalog(polarisMetaStoreSession, catalog, List.of()); Assertions.assertThat(catalogCreated).isNotNull(); catalog = catalogCreated.getCatalog(); @@ -1014,7 +1015,7 @@ PolarisBaseEntity ensureExistsByName( // find by name, ensure we found it PolarisMetaStoreManager.EntityResult entityFound = polarisMetaStoreManager.readEntityByName( - this.polarisCallContext, catalogPath, entityType, entitySubType, name); + polarisMetaStoreSession, catalogPath, entityType, entitySubType, name); Assertions.assertThat(entityFound).isNotNull(); Assertions.assertThat(entityFound.isSuccess()).isTrue(); @@ -1088,13 +1089,13 @@ PolarisBaseEntity updateEntity( // lookup that entity, ensure it exists PolarisBaseEntity beforeUpdateEntity = polarisMetaStoreManager - .loadEntity(this.polarisCallContext, entity.getCatalogId(), entity.getId()) + .loadEntity(polarisMetaStoreSession, entity.getCatalogId(), entity.getId()) .getEntity(); // update that property PolarisBaseEntity updatedEntity = polarisMetaStoreManager - .updateEntityPropertiesIfNotChanged(this.polarisCallContext, catalogPath, entity) + .updateEntityPropertiesIfNotChanged(polarisMetaStoreSession, catalogPath, entity) .getEntity(); // if version mismatch, nothing should be updated @@ -1105,7 +1106,7 @@ PolarisBaseEntity updateEntity( // refresh catalog info entity = polarisMetaStoreManager - .loadEntity(this.polarisCallContext, entity.getCatalogId(), entity.getId()) + .loadEntity(polarisMetaStoreSession, entity.getCatalogId(), entity.getId()) .getEntity(); // ensure nothing has changed @@ -1159,7 +1160,7 @@ PolarisBaseEntity updateEntity( List versions = polarisMetaStoreManager .loadEntitiesChangeTracking( - this.polarisCallContext, List.of(new PolarisEntityId(catalogId, entity.getId()))) + polarisMetaStoreSession, List.of(new PolarisEntityId(catalogId, entity.getId()))) .getChangeTrackingVersions(); Assertions.assertThat(versions).hasSize(1); Assertions.assertThat(versions.get(0).getEntityVersion()) @@ -1191,7 +1192,7 @@ private void validateListReturn( // list the entities under the specified path List result = polarisMetaStoreManager - .listEntities(this.polarisCallContext, path, entityType, entitySubType) + .listEntities(polarisMetaStoreSession, path, entityType, entitySubType) .getEntities(); Assertions.assertThat(result).isNotNull(); @@ -1243,7 +1244,7 @@ private void validateCacheEntryLoad(CachedEntryResult cacheEntry) { PolarisEntity refEntity = PolarisEntity.of( this.polarisMetaStoreManager.loadEntity( - this.polarisCallContext, entity.getCatalogId(), entity.getId())); + polarisMetaStoreSession, entity.getCatalogId(), entity.getId())); Assertions.assertThat(refEntity).isNotNull(); // same entity @@ -1256,7 +1257,7 @@ private void validateCacheEntryLoad(CachedEntryResult cacheEntry) { if (refEntity.getType().isGrantee()) { LoadGrantsResult loadGrantResult = this.polarisMetaStoreManager.loadGrantsToGrantee( - this.polarisCallContext, refEntity.getCatalogId(), refEntity.getId()); + polarisMetaStoreSession, refEntity.getCatalogId(), refEntity.getId()); this.validateLoadedGrants(loadGrantResult, true); // same version @@ -1268,7 +1269,7 @@ private void validateCacheEntryLoad(CachedEntryResult cacheEntry) { LoadGrantsResult loadGrantResult = this.polarisMetaStoreManager.loadGrantsOnSecurable( - this.polarisCallContext, refEntity.getCatalogId(), refEntity.getId()); + polarisMetaStoreSession, refEntity.getCatalogId(), refEntity.getId()); this.validateLoadedGrants(loadGrantResult, false); // same version @@ -1300,7 +1301,7 @@ private void validateCacheEntryRefresh( // reload the entity PolarisBaseEntity refEntity = this.polarisMetaStoreManager - .loadEntity(this.polarisCallContext, catalogId, entityId) + .loadEntity(polarisMetaStoreSession, catalogId, entityId) .getEntity(); Assertions.assertThat(refEntity).isNotNull(); @@ -1308,9 +1309,9 @@ private void validateCacheEntryRefresh( LoadGrantsResult loadGrantResult = refEntity.getType().isGrantee() ? this.polarisMetaStoreManager.loadGrantsToGrantee( - this.polarisCallContext, catalogId, entityId) + polarisMetaStoreSession, catalogId, entityId) : this.polarisMetaStoreManager.loadGrantsOnSecurable( - this.polarisCallContext, catalogId, entityId); + polarisMetaStoreSession, catalogId, entityId); this.validateLoadedGrants(loadGrantResult, refEntity.getType().isGrantee()); Assertions.assertThat(cacheEntry.getGrantRecordsVersion()) .isEqualTo(loadGrantResult.getGrantsVersion()); @@ -1364,7 +1365,7 @@ private PolarisBaseEntity loadCacheEntryByName( // load cached entry CachedEntryResult cacheEntry = this.polarisMetaStoreManager.loadCachedEntryByName( - this.polarisCallContext, entityCatalogId, parentId, entityType, entityName); + polarisMetaStoreSession, entityCatalogId, parentId, entityType, entityName); // if null, validate that indeed the entry does not exist Assertions.assertThat(cacheEntry.isSuccess()).isEqualTo(expectExists); @@ -1410,7 +1411,7 @@ private PolarisBaseEntity loadCacheEntryById( // load cached entry CachedEntryResult cacheEntry = this.polarisMetaStoreManager.loadCachedEntryById( - this.polarisCallContext, entityCatalogId, entityId); + polarisMetaStoreSession, entityCatalogId, entityId); // if null, validate that indeed the entry does not exist Assertions.assertThat(cacheEntry.isSuccess()).isEqualTo(expectExists); @@ -1457,7 +1458,7 @@ private void refreshCacheEntry( // load cached entry CachedEntryResult cacheEntry = this.polarisMetaStoreManager.refreshCachedEntity( - this.polarisCallContext, + polarisMetaStoreSession, entityVersion, entityGrantRecordsVersion, entityType, @@ -1501,7 +1502,7 @@ void validateBootstrap() { List principals = polarisMetaStoreManager .listEntities( - this.polarisCallContext, + polarisMetaStoreSession, null, PolarisEntityType.PRINCIPAL, PolarisEntitySubType.NULL_SUBTYPE) @@ -1527,7 +1528,7 @@ void validateBootstrap() { List principalRoles = polarisMetaStoreManager .listEntities( - this.polarisCallContext, + polarisMetaStoreSession, null, PolarisEntityType.PRINCIPAL_ROLE, PolarisEntitySubType.NULL_SUBTYPE) @@ -1932,7 +1933,7 @@ void testDropEntities() { // catalog exists PolarisMetaStoreManager.EntityResult catalogFound = polarisMetaStoreManager.readEntityByName( - this.polarisCallContext, + polarisMetaStoreSession, null, PolarisEntityType.CATALOG, PolarisEntitySubType.NULL_SUBTYPE, @@ -1965,7 +1966,7 @@ void testDropEntities() { // catalog exists? catalogFound = polarisMetaStoreManager.readEntityByName( - this.polarisCallContext, + polarisMetaStoreSession, null, PolarisEntityType.CATALOG, PolarisEntitySubType.NULL_SUBTYPE, @@ -2053,7 +2054,7 @@ public void testPrivileges() { grantToGrantee(catalog, R1, PR9000, PolarisPrivilege.CATALOG_ROLE_USAGE); LoadGrantsResult loadGrantsResult = - polarisMetaStoreManager.loadGrantsToGrantee(this.polarisCallContext, 0L, PR9000.getId()); + polarisMetaStoreManager.loadGrantsToGrantee(polarisMetaStoreSession, 0L, PR9000.getId()); this.validateLoadedGrants(loadGrantsResult, true); Assertions.assertThat(loadGrantsResult.getGrantRecords()).hasSize(1); Assertions.assertThat(loadGrantsResult.getGrantRecords().get(0).getSecurableCatalogId()) @@ -2062,7 +2063,7 @@ public void testPrivileges() { .isEqualTo(R1.getId()); loadGrantsResult = - polarisMetaStoreManager.loadGrantsToGrantee(this.polarisCallContext, 0L, PR900.getId()); + polarisMetaStoreManager.loadGrantsToGrantee(polarisMetaStoreSession, 0L, PR900.getId()); Assertions.assertThat(loadGrantsResult).isNotNull(); Assertions.assertThat(loadGrantsResult.getGrantRecords()).hasSize(0); } @@ -2097,7 +2098,7 @@ void renameEntity( // check to see if we would have a name conflict PolarisMetaStoreManager.EntityResult newNameLookup = polarisMetaStoreManager.readEntityByName( - polarisCallContext, + polarisMetaStoreSession, newCatPath == null ? catPath : newCatPath, entity.getType(), PolarisEntitySubType.ANY_SUBTYPE, @@ -2106,7 +2107,7 @@ void renameEntity( // rename it PolarisBaseEntity renamedEntity = polarisMetaStoreManager - .renameEntity(polarisCallContext, catPath, entity, newCatPath, renamedEntityInput) + .renameEntity(polarisMetaStoreSession, catPath, entity, newCatPath, renamedEntityInput) .getEntity(); // ensure success @@ -2132,7 +2133,7 @@ void renameEntity( // ensure the old one is gone PolarisMetaStoreManager.EntityResult res = polarisMetaStoreManager.readEntityByName( - polarisCallContext, catPath, entity.getType(), entity.getSubType(), oldName); + polarisMetaStoreSession, catPath, entity.getType(), entity.getSubType(), oldName); // not found Assertions.assertThat(res.getReturnStatus()) diff --git a/quarkus/service/src/main/java/org/apache/polaris/service/quarkus/config/QuarkusProducers.java b/quarkus/service/src/main/java/org/apache/polaris/service/quarkus/config/QuarkusProducers.java index f7c892730..65aa48ede 100644 --- a/quarkus/service/src/main/java/org/apache/polaris/service/quarkus/config/QuarkusProducers.java +++ b/quarkus/service/src/main/java/org/apache/polaris/service/quarkus/config/QuarkusProducers.java @@ -33,22 +33,25 @@ import jakarta.ws.rs.core.Context; import java.time.Clock; import java.util.HashMap; -import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.PolarisConfigurationStore; import org.apache.polaris.core.PolarisDefaultDiagServiceImpl; import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.auth.AuthenticatedPolarisPrincipal; import org.apache.polaris.core.auth.PolarisAuthorizer; import org.apache.polaris.core.auth.PolarisAuthorizerImpl; -import org.apache.polaris.core.context.CallContext; import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.persistence.MetaStoreManagerFactory; +import org.apache.polaris.core.persistence.PolarisEntityManager; +import org.apache.polaris.core.persistence.PolarisMetaStoreManager; import org.apache.polaris.core.persistence.PolarisMetaStoreSession; +import org.apache.polaris.core.persistence.cache.EntityCache; import org.apache.polaris.core.storage.cache.StorageCredentialCache; import org.apache.polaris.service.auth.Authenticator; +import org.apache.polaris.service.auth.TokenBroker; import org.apache.polaris.service.auth.TokenBrokerFactory; import org.apache.polaris.service.catalog.api.IcebergRestOAuth2ApiService; import org.apache.polaris.service.catalog.io.FileIOFactory; +import org.apache.polaris.service.config.RealmEntityManagerFactory; import org.apache.polaris.service.context.RealmContextConfiguration; import org.apache.polaris.service.context.RealmContextResolver; import org.apache.polaris.service.persistence.InMemoryPolarisMetaStoreManagerFactory; @@ -75,8 +78,9 @@ public Clock clock() { @Produces @ApplicationScoped - public StorageCredentialCache storageCredentialCache() { - return new StorageCredentialCache(); + public StorageCredentialCache storageCredentialCache( + PolarisDiagnostics diagnostics, PolarisConfigurationStore configurationStore) { + return new StorageCredentialCache(diagnostics, configurationStore); } @Produces @@ -107,25 +111,44 @@ public RealmContext realmContext( @Produces @RequestScoped - public PolarisCallContext polarisCallContext( - RealmContext realmContext, - PolarisDiagnostics diagServices, - PolarisConfigurationStore configurationStore, - MetaStoreManagerFactory metaStoreManagerFactory, - Clock clock) { - PolarisMetaStoreSession metaStoreSession = - metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(); - return new PolarisCallContext(metaStoreSession, diagServices, configurationStore, clock); + public PolarisMetaStoreSession metaStoreSession( + MetaStoreManagerFactory metaStoreManagerFactory, RealmContext realmContext) { + return metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(); } @Produces @RequestScoped - public CallContext callContext(RealmContext realmContext, PolarisCallContext polarisCallContext) { - return CallContext.of(realmContext, polarisCallContext); + // TODO break into separate beans + public PolarisMetaStoreManager polarisMetaStoreManager( + MetaStoreManagerFactory metaStoreManagerFactory, RealmContext realmContext) { + return metaStoreManagerFactory.getOrCreateMetaStoreManager(realmContext); } - public void closeCallContext(@Disposes CallContext callContext) { - callContext.close(); + @Produces + @RequestScoped + public StorageCredentialCache storageCredentialCache( + MetaStoreManagerFactory metaStoreManagerFactory, RealmContext realmContext) { + return metaStoreManagerFactory.getOrCreateStorageCredentialCache(realmContext); + } + + @Produces + @RequestScoped + public EntityCache entityCache( + MetaStoreManagerFactory metaStoreManagerFactory, RealmContext realmContext) { + return metaStoreManagerFactory.getOrCreateEntityCache(realmContext); + } + + @Produces + @RequestScoped + public PolarisEntityManager polarisEntityManager( + RealmEntityManagerFactory realmEntityManagerFactory, RealmContext realmContext) { + return realmEntityManagerFactory.getOrCreateEntityManager(realmContext); + } + + @Produces + @RequestScoped + public TokenBroker tokenBroker(TokenBrokerFactory tokenBrokerFactory, RealmContext realmContext) { + return tokenBrokerFactory.apply(realmContext); } // Polaris service beans - selected from @Identifier-annotated beans diff --git a/quarkus/service/src/main/java/org/apache/polaris/service/quarkus/task/QuarkusTaskExecutorImpl.java b/quarkus/service/src/main/java/org/apache/polaris/service/quarkus/task/QuarkusTaskExecutorImpl.java index 8f38de648..d14a9c2da 100644 --- a/quarkus/service/src/main/java/org/apache/polaris/service/quarkus/task/QuarkusTaskExecutorImpl.java +++ b/quarkus/service/src/main/java/org/apache/polaris/service/quarkus/task/QuarkusTaskExecutorImpl.java @@ -26,8 +26,11 @@ import io.smallrye.common.annotation.Identifier; import jakarta.enterprise.context.ApplicationScoped; import jakarta.inject.Inject; +import java.time.Clock; import java.util.concurrent.ExecutorService; -import org.apache.polaris.core.context.CallContext; +import org.apache.polaris.core.PolarisConfigurationStore; +import org.apache.polaris.core.PolarisDiagnostics; +import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.persistence.MetaStoreManagerFactory; import org.apache.polaris.service.quarkus.tracing.QuarkusTracingFilter; import org.apache.polaris.service.task.TaskExecutorImpl; @@ -39,16 +42,25 @@ public class QuarkusTaskExecutorImpl extends TaskExecutorImpl { private final Tracer tracer; public QuarkusTaskExecutorImpl() { - this(null, null, null, null); + this(null, null, null, null, null, null, null); } @Inject public QuarkusTaskExecutorImpl( @Identifier("task-executor") ExecutorService executorService, MetaStoreManagerFactory metaStoreManagerFactory, + PolarisConfigurationStore configurationStore, + PolarisDiagnostics diagnostics, TaskFileIOSupplier fileIOSupplier, + Clock clock, Tracer tracer) { - super(executorService, metaStoreManagerFactory, fileIOSupplier); + super( + executorService, + metaStoreManagerFactory, + configurationStore, + diagnostics, + fileIOSupplier, + clock); this.tracer = tracer; } @@ -59,19 +71,18 @@ public void init() { } @Override - protected void handleTask(long taskEntityId, CallContext callContext, int attempt) { + protected void handleTask(long taskEntityId, RealmContext realmContext, int attempt) { Span span = tracer .spanBuilder("polaris.task") .setParent(Context.current()) .setAttribute( - QuarkusTracingFilter.REALM_ID_ATTRIBUTE, - callContext.getRealmContext().getRealmIdentifier()) + QuarkusTracingFilter.REALM_ID_ATTRIBUTE, realmContext.getRealmIdentifier()) .setAttribute("polaris.task.entity.id", taskEntityId) .setAttribute("polaris.task.attempt", attempt) .startSpan(); try (Scope ignored = span.makeCurrent()) { - super.handleTask(taskEntityId, callContext, attempt); + super.handleTask(taskEntityId, realmContext, attempt); } finally { span.end(); } diff --git a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/TestServices.java b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/TestServices.java index a31d4bb47..6bc0be0a5 100644 --- a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/TestServices.java +++ b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/TestServices.java @@ -27,14 +27,13 @@ import java.util.Date; import java.util.Map; import java.util.Set; -import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.auth.AuthenticatedPolarisPrincipal; import org.apache.polaris.core.auth.PolarisAuthorizer; -import org.apache.polaris.core.context.CallContext; import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.PolarisEntity; import org.apache.polaris.core.entity.PrincipalEntity; +import org.apache.polaris.core.persistence.PolarisEntityManager; import org.apache.polaris.core.persistence.PolarisMetaStoreManager; import org.apache.polaris.core.persistence.PolarisMetaStoreSession; import org.apache.polaris.service.admin.PolarisServiceImpl; @@ -70,10 +69,22 @@ public static TestServices inMemory(FileIOFactory ioFactory) { } public static TestServices inMemory(FileIOFactory ioFactory, Map config) { + + DefaultConfigurationStore configurationStore = new DefaultConfigurationStore(config); + PolarisDiagnostics polarisDiagnostics = Mockito.mock(PolarisDiagnostics.class); + + PolarisStorageIntegrationProviderImpl storageIntegrationProvider = + new PolarisStorageIntegrationProviderImpl( + Mockito::mock, + () -> GoogleCredentials.create(new AccessToken("abc", new Date())), + configurationStore); + InMemoryPolarisMetaStoreManagerFactory metaStoreManagerFactory = new InMemoryPolarisMetaStoreManagerFactory( - new PolarisStorageIntegrationProviderImpl( - Mockito::mock, () -> GoogleCredentials.create(new AccessToken("abc", new Date())))); + storageIntegrationProvider, + configurationStore, + polarisDiagnostics, + Clock.systemDefaultZone()); PolarisMetaStoreManager metaStoreManager = metaStoreManagerFactory.getOrCreateMetaStoreManager(testRealm); @@ -81,36 +92,40 @@ public static TestServices inMemory(FileIOFactory ioFactory, Map PolarisMetaStoreSession session = metaStoreManagerFactory.getOrCreateSessionSupplier(testRealm).get(); - PolarisCallContext context = - new PolarisCallContext( - session, - Mockito.mock(PolarisDiagnostics.class), - new DefaultConfigurationStore(config), - Clock.systemDefaultZone()); + RealmEntityManagerFactory realmEntityManagerFactory = + new RealmEntityManagerFactory(metaStoreManagerFactory, polarisDiagnostics) {}; - CallContext callContext = CallContext.of(testRealm, context); + PolarisEntityManager entityManager = + realmEntityManagerFactory.getOrCreateEntityManager(testRealm); - RealmEntityManagerFactory realmEntityManagerFactory = - new RealmEntityManagerFactory(metaStoreManagerFactory) {}; CallContextCatalogFactory callContextFactory = new PolarisCallContextCatalogFactory( - realmEntityManagerFactory, - metaStoreManagerFactory, + entityManager, + metaStoreManager, + session, + configurationStore, + polarisDiagnostics, Mockito.mock(TaskExecutor.class), ioFactory); + PolarisAuthorizer authorizer = Mockito.mock(PolarisAuthorizer.class); + IcebergRestCatalogApiService service = new IcebergCatalogAdapter( - callContext, + testRealm, callContextFactory, - realmEntityManagerFactory, - metaStoreManagerFactory, + entityManager, + metaStoreManager, + session, + configurationStore, + polarisDiagnostics, authorizer); + IcebergRestCatalogApi restApi = new IcebergRestCatalogApi(service); PolarisMetaStoreManager.CreatePrincipalResult createdPrincipal = metaStoreManager.createPrincipal( - context, + session, new PrincipalEntity.Builder() .setName("test-principal") .setCreateTimestamp(Instant.now().toEpochMilli()) @@ -147,9 +162,13 @@ public String getAuthenticationScheme() { PolarisCatalogsApi catalogsApi = new PolarisCatalogsApi( new PolarisServiceImpl( - realmEntityManagerFactory, metaStoreManagerFactory, authorizer, callContext)); + entityManager, + metaStoreManager, + session, + configurationStore, + authorizer, + polarisDiagnostics)); - CallContext.setCurrentContext(CallContext.of(testRealm, context)); return new TestServices(restApi, catalogsApi, testRealm, securityContext); } } diff --git a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/admin/PolarisAdminServiceAuthzTest.java b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/admin/PolarisAdminServiceAuthzTest.java index 5c055eec4..303746d5d 100644 --- a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/admin/PolarisAdminServiceAuthzTest.java +++ b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/admin/PolarisAdminServiceAuthzTest.java @@ -49,9 +49,12 @@ private PolarisAdminService newTestAdminService(Set activatedPrincipalRo final AuthenticatedPolarisPrincipal authenticatedPrincipal = new AuthenticatedPolarisPrincipal(principalEntity, activatedPrincipalRoles); return new PolarisAdminService( - callContext, + realmContext, entityManager, metaStoreManager, + metaStoreSession, + configurationStore, + diagServices, securityContext(authenticatedPrincipal, activatedPrincipalRoles), polarisAuthorizer); } diff --git a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/admin/PolarisAuthzTestBase.java b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/admin/PolarisAuthzTestBase.java index 93f5fa6fb..ed802c9fc 100644 --- a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/admin/PolarisAuthzTestBase.java +++ b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/admin/PolarisAuthzTestBase.java @@ -26,13 +26,11 @@ import io.quarkus.test.junit.QuarkusMock; import io.quarkus.test.junit.QuarkusTestProfile; import jakarta.annotation.Nonnull; -import jakarta.annotation.Nullable; -import jakarta.enterprise.context.ApplicationScoped; +import jakarta.enterprise.context.RequestScoped; import jakarta.enterprise.inject.Alternative; import jakarta.inject.Inject; import jakarta.ws.rs.core.SecurityContext; import java.io.IOException; -import java.time.Clock; import java.util.Date; import java.util.List; import java.util.Map; @@ -46,7 +44,6 @@ import org.apache.iceberg.catalog.TableIdentifier; import org.apache.iceberg.exceptions.ForbiddenException; import org.apache.iceberg.types.Types; -import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.PolarisConfiguration; import org.apache.polaris.core.PolarisConfigurationStore; import org.apache.polaris.core.PolarisDiagnostics; @@ -57,7 +54,6 @@ import org.apache.polaris.core.auth.AuthenticatedPolarisPrincipal; import org.apache.polaris.core.auth.PolarisAuthorizer; import org.apache.polaris.core.auth.PolarisAuthorizerImpl; -import org.apache.polaris.core.context.CallContext; import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.CatalogEntity; import org.apache.polaris.core.entity.CatalogRoleEntity; @@ -71,6 +67,7 @@ import org.apache.polaris.core.persistence.MetaStoreManagerFactory; import org.apache.polaris.core.persistence.PolarisEntityManager; import org.apache.polaris.core.persistence.PolarisMetaStoreManager; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; import org.apache.polaris.core.persistence.resolver.PolarisResolutionManifest; import org.apache.polaris.service.admin.PolarisAdminService; import org.apache.polaris.service.catalog.BasePolarisCatalog; @@ -78,7 +75,6 @@ import org.apache.polaris.service.catalog.io.FileIOFactory; import org.apache.polaris.service.config.DefaultConfigurationStore; import org.apache.polaris.service.config.RealmEntityManagerFactory; -import org.apache.polaris.service.context.CallContextCatalogFactory; import org.apache.polaris.service.context.PolarisCallContextCatalogFactory; import org.apache.polaris.service.quarkus.catalog.PolarisPassthroughResolutionView; import org.apache.polaris.service.storage.PolarisStorageIntegrationProviderImpl; @@ -99,6 +95,15 @@ public static class Profile implements QuarkusTestProfile { public Set> getEnabledAlternatives() { return Set.of(TestPolarisCallContextCatalogFactory.class); } + + @Override + public Map getConfigOverrides() { + return Map.of( + "polaris.config.defaults.ALLOW_SPECIFYING_FILE_IO_IMPL", + "true", + "polaris.config.defaults.ALLOW_EXTERNAL_METADATA_FILE_LOCATION", + "true"); + } } protected static final String CATALOG_NAME = "polaris-catalog"; @@ -159,61 +164,42 @@ public Set> getEnabledAlternatives() { @Inject protected MetaStoreManagerFactory managerFactory; @Inject protected RealmEntityManagerFactory realmEntityManagerFactory; - @Inject protected CallContextCatalogFactory callContextCatalogFactory; + @Inject protected PolarisConfigurationStore configurationStore; @Inject protected PolarisDiagnostics diagServices; - @Inject protected Clock clock; protected BasePolarisCatalog baseCatalog; protected PolarisAdminService adminService; protected PolarisEntityManager entityManager; protected PolarisMetaStoreManager metaStoreManager; + protected PolarisMetaStoreSession metaStoreSession; protected PolarisBaseEntity catalogEntity; protected PrincipalEntity principalEntity; - protected CallContext callContext; + protected RealmContext realmContext; protected AuthenticatedPolarisPrincipal authenticatedRoot; - private PolarisCallContext polarisContext; - @BeforeAll public static void setUpMocks() { PolarisStorageIntegrationProviderImpl mock = new PolarisStorageIntegrationProviderImpl( - Mockito::mock, () -> GoogleCredentials.create(new AccessToken("abc", new Date()))); + Mockito::mock, + () -> GoogleCredentials.create(new AccessToken("abc", new Date())), + null); QuarkusMock.installMockForType(mock, PolarisStorageIntegrationProviderImpl.class); } @BeforeEach public void before(TestInfo testInfo) { - RealmContext realmContext = testInfo::getDisplayName; + realmContext = testInfo::getDisplayName; metaStoreManager = managerFactory.getOrCreateMetaStoreManager(realmContext); - - Map configMap = - Map.of( - "ALLOW_SPECIFYING_FILE_IO_IMPL", true, "ALLOW_EXTERNAL_METADATA_FILE_LOCATION", true); - polarisContext = - new PolarisCallContext( - managerFactory.getOrCreateSessionSupplier(realmContext).get(), - diagServices, - new PolarisConfigurationStore() { - @Override - public @Nullable T getConfiguration(PolarisCallContext ctx, String configName) { - @SuppressWarnings("unchecked") - var r = (T) configMap.get(configName); - return r; - } - }, - clock); - this.entityManager = realmEntityManagerFactory.getOrCreateEntityManager(realmContext); - - callContext = CallContext.of(realmContext, polarisContext); - CallContext.setCurrentContext(callContext); + metaStoreSession = managerFactory.getOrCreateSessionSupplier(realmContext).get(); + entityManager = realmEntityManagerFactory.getOrCreateEntityManager(realmContext); PrincipalEntity rootEntity = new PrincipalEntity( PolarisEntity.of( metaStoreManager .readEntityByName( - polarisContext, + metaStoreSession, null, PolarisEntityType.PRINCIPAL, PolarisEntitySubType.NULL_SUBTYPE, @@ -224,9 +210,12 @@ public void before(TestInfo testInfo) { this.adminService = new PolarisAdminService( - callContext, + realmContext, entityManager, metaStoreManager, + metaStoreSession, + configurationStore, + diagServices, securityContext(authenticatedRoot, Set.of()), polarisAuthorizer); @@ -249,9 +238,7 @@ public void before(TestInfo testInfo) { PrincipalWithCredentials principal = adminService.createPrincipal(new PrincipalEntity.Builder().setName(PRINCIPAL_NAME).build()); - principalEntity = - rotateAndRefreshPrincipal( - metaStoreManager, PRINCIPAL_NAME, principal.getCredentials(), polarisContext); + principalEntity = rotateAndRefreshPrincipal(PRINCIPAL_NAME, principal.getCredentials()); // Pre-create the principal roles and catalog roles without any grants on securables, but // assign both principal roles to the principal, then CATALOG_ROLE1 to PRINCIPAL_ROLE1, @@ -334,7 +321,7 @@ public void after() { } } } finally { - metaStoreManager.purge(polarisContext); + metaStoreManager.purge(metaStoreSession); } } @@ -350,34 +337,27 @@ public void after() { protected @Nonnull Set loadPrincipalRolesNames(AuthenticatedPolarisPrincipal p) { return metaStoreManager - .loadGrantsToGrantee( - callContext.getPolarisCallContext(), 0L, p.getPrincipalEntity().getId()) + .loadGrantsToGrantee(metaStoreSession, 0L, p.getPrincipalEntity().getId()) .getGrantRecords() .stream() .filter(gr -> gr.getPrivilegeCode() == PolarisPrivilege.PRINCIPAL_ROLE_USAGE.getCode()) - .map( - gr -> - metaStoreManager.loadEntity( - callContext.getPolarisCallContext(), 0L, gr.getSecurableId())) + .map(gr -> metaStoreManager.loadEntity(metaStoreSession, 0L, gr.getSecurableId())) .map(PolarisMetaStoreManager.EntityResult::getEntity) .map(PolarisBaseEntity::getName) .collect(Collectors.toSet()); } protected @Nonnull PrincipalEntity rotateAndRefreshPrincipal( - PolarisMetaStoreManager metaStoreManager, - String principalName, - PrincipalWithCredentialsCredentials credentials, - PolarisCallContext polarisContext) { + String principalName, PrincipalWithCredentialsCredentials credentials) { PolarisMetaStoreManager.EntityResult lookupEntity = metaStoreManager.readEntityByName( - callContext.getPolarisCallContext(), + metaStoreSession, null, PolarisEntityType.PRINCIPAL, PolarisEntitySubType.NULL_SUBTYPE, principalName); metaStoreManager.rotatePrincipalSecrets( - callContext.getPolarisCallContext(), + metaStoreSession, credentials.getClientId(), lookupEntity.getEntity().getId(), false, @@ -387,7 +367,7 @@ public void after() { PolarisEntity.of( metaStoreManager .readEntityByName( - polarisContext, + metaStoreSession, null, PolarisEntityType.PRINCIPAL, PolarisEntitySubType.NULL_SUBTYPE, @@ -414,15 +394,17 @@ private void initBaseCatalog() { Mockito.when(securityContext.isUserInRole(Mockito.anyString())).thenReturn(true); PolarisPassthroughResolutionView passthroughView = new PolarisPassthroughResolutionView( - callContext, entityManager, securityContext, CATALOG_NAME); + entityManager, metaStoreSession, securityContext, CATALOG_NAME); this.baseCatalog = new BasePolarisCatalog( + realmContext, entityManager, metaStoreManager, - callContext, + metaStoreSession, + configurationStore, + diagServices, passthroughView, securityContext, - authenticatedRoot, Mockito.mock(), new DefaultFileIOFactory()); this.baseCatalog.initialize( @@ -432,26 +414,36 @@ private void initBaseCatalog() { } @Alternative - @ApplicationScoped + @RequestScoped public static class TestPolarisCallContextCatalogFactory extends PolarisCallContextCatalogFactory { public TestPolarisCallContextCatalogFactory() { - super(null, null, null, null); + super(null, null, null, null, null, null, null); } @Inject public TestPolarisCallContextCatalogFactory( - RealmEntityManagerFactory entityManagerFactory, - MetaStoreManagerFactory metaStoreManagerFactory, + PolarisEntityManager entityManager, + PolarisMetaStoreManager metaStoreManager, + PolarisMetaStoreSession metaStoreSession, + PolarisConfigurationStore configurationStore, + PolarisDiagnostics diagnostics, TaskExecutor taskExecutor, FileIOFactory fileIOFactory) { - super(entityManagerFactory, metaStoreManagerFactory, taskExecutor, fileIOFactory); + super( + entityManager, + metaStoreManager, + metaStoreSession, + configurationStore, + diagnostics, + taskExecutor, + fileIOFactory); } @Override public Catalog createCallContextCatalog( - CallContext context, + RealmContext realmContext, AuthenticatedPolarisPrincipal authenticatedPolarisPrincipal, SecurityContext securityContext, final PolarisResolutionManifest resolvedManifest) { @@ -459,7 +451,7 @@ public Catalog createCallContextCatalog( // to override the previous config. Catalog catalog = super.createCallContextCatalog( - context, authenticatedPolarisPrincipal, securityContext, resolvedManifest); + realmContext, authenticatedPolarisPrincipal, securityContext, resolvedManifest); catalog.initialize( CATALOG_NAME, ImmutableMap.of( diff --git a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/auth/JWTRSAKeyPairTest.java b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/auth/JWTRSAKeyPairTest.java index ba624e621..6627fbf95 100644 --- a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/auth/JWTRSAKeyPairTest.java +++ b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/auth/JWTRSAKeyPairTest.java @@ -28,21 +28,19 @@ import java.nio.file.Path; import java.security.interfaces.RSAPrivateKey; import java.security.interfaces.RSAPublicKey; -import java.util.HashMap; -import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.auth.PolarisSecretsManager.PrincipalSecretsResult; import org.apache.polaris.core.entity.PolarisBaseEntity; import org.apache.polaris.core.entity.PolarisEntitySubType; import org.apache.polaris.core.entity.PolarisEntityType; import org.apache.polaris.core.entity.PolarisPrincipalSecrets; import org.apache.polaris.core.persistence.PolarisMetaStoreManager; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; import org.apache.polaris.service.auth.JWTRSAKeyPair; import org.apache.polaris.service.auth.LocalRSAKeyProvider; import org.apache.polaris.service.auth.PemUtils; import org.apache.polaris.service.auth.TokenBroker; import org.apache.polaris.service.auth.TokenRequestValidator; import org.apache.polaris.service.auth.TokenResponse; -import org.apache.polaris.service.config.DefaultConfigurationStore; import org.junit.jupiter.api.Test; import org.mockito.Mockito; @@ -57,13 +55,12 @@ public void testSuccessfulTokenGeneration() throws Exception { final String clientId = "test-client-id"; final String scope = "PRINCIPAL_ROLE:TEST"; - DefaultConfigurationStore store = new DefaultConfigurationStore(new HashMap<>()); - PolarisCallContext polarisCallContext = new PolarisCallContext(null, null, store, null); PolarisMetaStoreManager metastoreManager = Mockito.mock(PolarisMetaStoreManager.class); String mainSecret = "client-secret"; PolarisPrincipalSecrets principalSecrets = new PolarisPrincipalSecrets(1L, clientId, mainSecret, "otherSecret"); - Mockito.when(metastoreManager.loadPrincipalSecrets(polarisCallContext, clientId)) + PolarisMetaStoreSession session = Mockito.mock(PolarisMetaStoreSession.class); + Mockito.when(metastoreManager.loadPrincipalSecrets(session, clientId)) .thenReturn(new PrincipalSecretsResult(principalSecrets)); PolarisBaseEntity principal = new PolarisBaseEntity( @@ -73,17 +70,13 @@ public void testSuccessfulTokenGeneration() throws Exception { PolarisEntitySubType.NULL_SUBTYPE, 0L, "principal"); - Mockito.when(metastoreManager.loadEntity(polarisCallContext, 0L, 1L)) + Mockito.when(metastoreManager.loadEntity(session, 0L, 1L)) .thenReturn(new PolarisMetaStoreManager.EntityResult(principal)); TokenBroker tokenBroker = - new JWTRSAKeyPair(metastoreManager, 420, publicFileLocation, privateFileLocation); + new JWTRSAKeyPair(metastoreManager, session, 420, publicFileLocation, privateFileLocation); TokenResponse token = tokenBroker.generateFromClientSecrets( - clientId, - mainSecret, - TokenRequestValidator.CLIENT_CREDENTIALS, - scope, - polarisCallContext); + clientId, mainSecret, TokenRequestValidator.CLIENT_CREDENTIALS, scope); assertThat(token).isNotNull(); assertThat(token.getExpiresIn()).isEqualTo(420); diff --git a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/auth/JWTSymmetricKeyGeneratorTest.java b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/auth/JWTSymmetricKeyGeneratorTest.java index aaa62643d..039e575c3 100644 --- a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/auth/JWTSymmetricKeyGeneratorTest.java +++ b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/auth/JWTSymmetricKeyGeneratorTest.java @@ -24,15 +24,12 @@ import com.auth0.jwt.JWTVerifier; import com.auth0.jwt.algorithms.Algorithm; import com.auth0.jwt.interfaces.DecodedJWT; -import java.util.Map; -import org.apache.polaris.core.PolarisCallContext; -import org.apache.polaris.core.context.CallContext; -import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.PolarisBaseEntity; import org.apache.polaris.core.entity.PolarisEntitySubType; import org.apache.polaris.core.entity.PolarisEntityType; import org.apache.polaris.core.entity.PolarisPrincipalSecrets; import org.apache.polaris.core.persistence.PolarisMetaStoreManager; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; import org.apache.polaris.service.auth.JWTSymmetricKeyBroker; import org.apache.polaris.service.auth.TokenBroker; import org.apache.polaris.service.auth.TokenRequestValidator; @@ -45,30 +42,13 @@ public class JWTSymmetricKeyGeneratorTest { /** Sanity test to verify that we can generate a token */ @Test public void testJWTSymmetricKeyGenerator() { - PolarisCallContext polarisCallContext = new PolarisCallContext(null, null, null, null); - CallContext.setCurrentContext( - new CallContext() { - @Override - public RealmContext getRealmContext() { - return () -> "realm"; - } - - @Override - public PolarisCallContext getPolarisCallContext() { - return polarisCallContext; - } - - @Override - public Map contextVariables() { - return Map.of(); - } - }); PolarisMetaStoreManager metastoreManager = Mockito.mock(PolarisMetaStoreManager.class); + PolarisMetaStoreSession metaStoreSession = Mockito.mock(PolarisMetaStoreSession.class); String mainSecret = "test_secret"; String clientId = "test_client_id"; PolarisPrincipalSecrets principalSecrets = new PolarisPrincipalSecrets(1L, clientId, mainSecret, "otherSecret"); - Mockito.when(metastoreManager.loadPrincipalSecrets(polarisCallContext, clientId)) + Mockito.when(metastoreManager.loadPrincipalSecrets(metaStoreSession, clientId)) .thenReturn(new PolarisMetaStoreManager.PrincipalSecretsResult(principalSecrets)); PolarisBaseEntity principal = new PolarisBaseEntity( @@ -78,16 +58,13 @@ public Map contextVariables() { PolarisEntitySubType.NULL_SUBTYPE, 0L, "principal"); - Mockito.when(metastoreManager.loadEntity(polarisCallContext, 0L, 1L)) + Mockito.when(metastoreManager.loadEntity(metaStoreSession, 0L, 1L)) .thenReturn(new PolarisMetaStoreManager.EntityResult(principal)); - TokenBroker generator = new JWTSymmetricKeyBroker(metastoreManager, 666, () -> "polaris"); + TokenBroker generator = + new JWTSymmetricKeyBroker(metastoreManager, metaStoreSession, 666, () -> "polaris"); TokenResponse token = generator.generateFromClientSecrets( - clientId, - mainSecret, - TokenRequestValidator.CLIENT_CREDENTIALS, - "PRINCIPAL_ROLE:TEST", - polarisCallContext); + clientId, mainSecret, TokenRequestValidator.CLIENT_CREDENTIALS, "PRINCIPAL_ROLE:TEST"); assertThat(token).isNotNull(); JWTVerifier verifier = JWT.require(Algorithm.HMAC256("polaris")).withIssuer("polaris").build(); diff --git a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/catalog/BasePolarisCatalogTest.java b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/catalog/BasePolarisCatalogTest.java index 316f115b9..20d0fd0bf 100644 --- a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/catalog/BasePolarisCatalogTest.java +++ b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/catalog/BasePolarisCatalogTest.java @@ -20,7 +20,11 @@ import static java.nio.charset.StandardCharsets.UTF_8; import static org.apache.iceberg.types.Types.NestedField.required; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.isA; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; import static org.mockito.Mockito.when; import com.google.common.collect.ImmutableMap; @@ -60,7 +64,6 @@ import org.apache.iceberg.inmemory.InMemoryFileIO; import org.apache.iceberg.io.FileIO; import org.apache.iceberg.types.Types; -import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.PolarisConfiguration; import org.apache.polaris.core.PolarisConfigurationStore; import org.apache.polaris.core.PolarisDiagnostics; @@ -69,7 +72,6 @@ import org.apache.polaris.core.auth.AuthenticatedPolarisPrincipal; import org.apache.polaris.core.auth.PolarisAuthorizerImpl; import org.apache.polaris.core.auth.PolarisSecretsManager.PrincipalSecretsResult; -import org.apache.polaris.core.context.CallContext; import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.CatalogEntity; import org.apache.polaris.core.entity.PolarisBaseEntity; @@ -93,6 +95,7 @@ import org.apache.polaris.service.catalog.BasePolarisCatalog; import org.apache.polaris.service.catalog.io.DefaultFileIOFactory; import org.apache.polaris.service.catalog.io.FileIOFactory; +import org.apache.polaris.service.config.RealmEntityManagerFactory; import org.apache.polaris.service.exception.IcebergExceptionMapper; import org.apache.polaris.service.quarkus.catalog.io.TestFileIOFactory; import org.apache.polaris.service.storage.PolarisStorageIntegrationProviderImpl; @@ -129,21 +132,22 @@ public class BasePolarisCatalogTest extends CatalogTests { public static final String SESSION_TOKEN = "session_token"; @Inject MetaStoreManagerFactory managerFactory; + @Inject RealmEntityManagerFactory entityManagerFactory; @Inject PolarisConfigurationStore configurationStore; @Inject PolarisStorageIntegrationProvider storageIntegrationProvider; @Inject PolarisDiagnostics diagServices; + @Inject Clock clock; private BasePolarisCatalog catalog; - private AwsStorageConfigInfo storageConfigModel; - private StsClient stsClient; - private String realmName; + private RealmContext realmContext; private PolarisMetaStoreManager metaStoreManager; - private PolarisCallContext polarisContext; + private PolarisMetaStoreSession metaStoreSession; private PolarisAdminService adminService; private PolarisEntityManager entityManager; private AuthenticatedPolarisPrincipal authenticatedRoot; private PolarisEntity catalogEntity; private SecurityContext securityContext; + private PolarisPassthroughResolutionView passthroughView; @BeforeAll public static void setUpMocks() { @@ -155,31 +159,21 @@ public static void setUpMocks() { @BeforeEach @SuppressWarnings("unchecked") public void before(TestInfo testInfo) { - realmName = + String realmName = "realm_%s_%s" .formatted( testInfo.getTestMethod().map(Method::getName).orElse("test"), System.nanoTime()); - RealmContext realmContext = () -> realmName; + realmContext = () -> realmName; metaStoreManager = managerFactory.getOrCreateMetaStoreManager(realmContext); - polarisContext = - new PolarisCallContext( - managerFactory.getOrCreateSessionSupplier(realmContext).get(), - diagServices, - configurationStore, - Clock.systemDefaultZone()); - entityManager = - new PolarisEntityManager( - metaStoreManager, new StorageCredentialCache(), new EntityCache(metaStoreManager)); - - CallContext callContext = CallContext.of(realmContext, polarisContext); - CallContext.setCurrentContext(callContext); + metaStoreSession = managerFactory.getOrCreateSessionSupplier(realmContext).get(); + entityManager = entityManagerFactory.getOrCreateEntityManager(realmContext); PrincipalEntity rootEntity = new PrincipalEntity( PolarisEntity.of( metaStoreManager .readEntityByName( - polarisContext, + metaStoreSession, null, PolarisEntityType.PRINCIPAL, PolarisEntitySubType.NULL_SUBTYPE, @@ -193,14 +187,17 @@ public void before(TestInfo testInfo) { when(securityContext.isUserInRole(isA(String.class))).thenReturn(true); adminService = new PolarisAdminService( - callContext, + realmContext, entityManager, metaStoreManager, + metaStoreSession, + configurationStore, + diagServices, securityContext, new PolarisAuthorizerImpl(new PolarisConfigurationStore() {})); String storageLocation = "s3://my-bucket/path/to/data"; - storageConfigModel = + AwsStorageConfigInfo storageConfigModel = AwsStorageConfigInfo.builder() .setRoleArn("arn:aws:iam::012345678901:role/jdoe") .setExternalId("externalId") @@ -221,25 +218,27 @@ public void before(TestInfo testInfo) { .setStorageConfigurationInfo(storageConfigModel, storageLocation) .build()); - PolarisPassthroughResolutionView passthroughView = + passthroughView = new PolarisPassthroughResolutionView( - callContext, entityManager, securityContext, CATALOG_NAME); + entityManager, metaStoreSession, securityContext, CATALOG_NAME); TaskExecutor taskExecutor = Mockito.mock(); this.catalog = new BasePolarisCatalog( + realmContext, entityManager, metaStoreManager, - callContext, + metaStoreSession, + configurationStore, + diagServices, passthroughView, securityContext, - authenticatedRoot, taskExecutor, new DefaultFileIOFactory()); this.catalog.initialize( CATALOG_NAME, ImmutableMap.of( CatalogProperties.FILE_IO_IMPL, "org.apache.iceberg.inmemory.InMemoryFileIO")); - stsClient = Mockito.mock(StsClient.class); + StsClient stsClient = Mockito.mock(StsClient.class); when(stsClient.assumeRole(isA(AssumeRoleRequest.class))) .thenReturn( AssumeRoleResponse.builder() @@ -251,7 +250,7 @@ public void before(TestInfo testInfo) { .build()) .build()); PolarisStorageIntegration storageIntegration = - new AwsCredentialsStorageIntegration(stsClient); + new AwsCredentialsStorageIntegration(configurationStore, stsClient); when(storageIntegrationProvider.getStorageIntegrationForConfig( isA(AwsStorageConfigurationInfo.class))) .thenReturn((PolarisStorageIntegration) storageIntegration); @@ -260,7 +259,7 @@ public void before(TestInfo testInfo) { @AfterEach public void after() throws IOException { catalog().close(); - metaStoreManager.purge(polarisContext); + metaStoreManager.purge(metaStoreSession); } @Override @@ -297,17 +296,17 @@ public PolarisMetaStoreManager getOrCreateMetaStoreManager(RealmContext realmCon @Override public Supplier getOrCreateSessionSupplier( RealmContext realmContext) { - return () -> polarisContext.getMetaStore(); + return () -> metaStoreSession; } @Override public StorageCredentialCache getOrCreateStorageCredentialCache(RealmContext realmContext) { - return new StorageCredentialCache(); + return new StorageCredentialCache(diagServices, configurationStore); } @Override public EntityCache getOrCreateEntityCache(RealmContext realmContext) { - return new EntityCache(metaStoreManager); + return new EntityCache(metaStoreManager, diagServices); } @Override @@ -476,7 +475,7 @@ public void testValidateNotificationInDisallowedLocation() { } @Test - public void testValidateNotificationFailToCreateFileIO() { + public void testValidateNotificationFailToCreateFileIO() throws IOException { Assumptions.assumeTrue( requiresNamespaceCreate(), "Only applicable if namespaces must be created before adding children"); @@ -491,7 +490,23 @@ public void testValidateNotificationFailToCreateFileIO() { // filename. final String tableLocation = "s3://externally-owned-bucket/validate_table/"; final String tableMetadataLocation = tableLocation + "metadata/"; - BasePolarisCatalog catalog = catalog(); + FileIOFactory fileIoFactory = spy(new DefaultFileIOFactory()); + BasePolarisCatalog catalog = + new BasePolarisCatalog( + realmContext, + entityManager, + metaStoreManager, + metaStoreSession, + configurationStore, + diagServices, + passthroughView, + securityContext, + Mockito.mock(), + fileIoFactory); + catalog.initialize( + CATALOG_NAME, + ImmutableMap.of( + CatalogProperties.FILE_IO_IMPL, "org.apache.iceberg.inmemory.InMemoryFileIO")); Namespace namespace = Namespace.of("parent", "child1"); TableIdentifier table = TableIdentifier.of(namespace, "table"); @@ -505,16 +520,14 @@ public void testValidateNotificationFailToCreateFileIO() { update.setTimestamp(230950845L); request.setPayload(update); - catalog.setFileIOFactory( - new FileIOFactory() { - @Override - public FileIO loadFileIO(String impl, Map properties) { - throw new ForbiddenException("Fake failure applying downscoped credentials"); - } - }); + doThrow(new ForbiddenException("Fake failure applying downscoped credentials")) + .when(fileIoFactory) + .loadFileIO(any(), any()); Assertions.assertThatThrownBy(() -> catalog.sendNotification(table, request)) .isInstanceOf(ForbiddenException.class) .hasMessageContaining("Fake failure applying downscoped credentials"); + + catalog.close(); } @Test @@ -616,7 +629,7 @@ public void testCreateNotificationCreateTableInExternalLocation() { final String anotherTableLocation = "s3://my-bucket/path/to/data/another_table/"; metaStoreManager.updateEntityPropertiesIfNotChanged( - polarisContext, + metaStoreSession, List.of(PolarisEntity.toCore(catalogEntity)), new CatalogEntity.Builder(CatalogEntity.of(catalogEntity)) .addProperty( @@ -673,7 +686,7 @@ public void testCreateNotificationCreateTableOutsideOfMetadataLocation() { final String anotherTableLocation = "s3://my-bucket/path/to/data/another_table"; metaStoreManager.updateEntityPropertiesIfNotChanged( - polarisContext, + metaStoreSession, List.of(PolarisEntity.toCore(catalogEntity)), new CatalogEntity.Builder(CatalogEntity.of(catalogEntity)) .addProperty( @@ -727,7 +740,7 @@ public void testUpdateNotificationCreateTableInExternalLocation() { final String anotherTableLocation = "s3://my-bucket/path/to/data/another_table/"; metaStoreManager.updateEntityPropertiesIfNotChanged( - polarisContext, + metaStoreSession, List.of(PolarisEntity.toCore(catalogEntity)), new CatalogEntity.Builder(CatalogEntity.of(catalogEntity)) .addProperty( @@ -796,26 +809,26 @@ public void testUpdateNotificationCreateTableWithLocalFilePrefix() { // The location of the metadata JSON file specified in the create will be forbidden. final String metadataLocation = "file:///etc/metadata.json/../passwd"; String catalogWithoutStorage = "catalogWithoutStorage"; - PolarisEntity catalogEntity = - adminService.createCatalog( - new CatalogEntity.Builder() - .setDefaultBaseLocation("file://") - .setName(catalogWithoutStorage) - .build()); + adminService.createCatalog( + new CatalogEntity.Builder() + .setDefaultBaseLocation("file://") + .setName(catalogWithoutStorage) + .build()); - CallContext callContext = CallContext.getCurrentContext(); PolarisPassthroughResolutionView passthroughView = new PolarisPassthroughResolutionView( - callContext, entityManager, securityContext, catalogWithoutStorage); + entityManager, metaStoreSession, securityContext, catalogWithoutStorage); TaskExecutor taskExecutor = Mockito.mock(); BasePolarisCatalog catalog = new BasePolarisCatalog( + realmContext, entityManager, metaStoreManager, - callContext, + metaStoreSession, + configurationStore, + diagServices, passthroughView, securityContext, - authenticatedRoot, taskExecutor, new DefaultFileIOFactory()); catalog.initialize( @@ -841,10 +854,8 @@ public void testUpdateNotificationCreateTableWithLocalFilePrefix() { metadataLocation, TableMetadataParser.toJson(createSampleTableMetadata(metadataLocation)).getBytes(UTF_8)); - PolarisCallContext polarisCallContext = callContext.getPolarisCallContext(); - if (!polarisCallContext - .getConfigurationStore() - .getConfiguration(polarisCallContext, PolarisConfiguration.SUPPORTED_CATALOG_STORAGE_TYPES) + if (!configurationStore + .getConfiguration(realmContext, PolarisConfiguration.SUPPORTED_CATALOG_STORAGE_TYPES) .contains("FILE")) { Assertions.assertThatThrownBy(() -> catalog.sendNotification(table, request)) .isInstanceOf(ForbiddenException.class) @@ -869,19 +880,20 @@ public void testUpdateNotificationCreateTableWithHttpPrefix() { .setName(catalogName) .build()); - CallContext callContext = CallContext.getCurrentContext(); PolarisPassthroughResolutionView passthroughView = new PolarisPassthroughResolutionView( - callContext, entityManager, securityContext, catalogName); + entityManager, metaStoreSession, securityContext, catalogName); TaskExecutor taskExecutor = Mockito.mock(); BasePolarisCatalog catalog = new BasePolarisCatalog( + realmContext, entityManager, metaStoreManager, - callContext, + metaStoreSession, + configurationStore, + diagServices, passthroughView, securityContext, - authenticatedRoot, taskExecutor, new DefaultFileIOFactory()); catalog.initialize( @@ -909,10 +921,8 @@ public void testUpdateNotificationCreateTableWithHttpPrefix() { metadataLocation, TableMetadataParser.toJson(createSampleTableMetadata(metadataLocation)).getBytes(UTF_8)); - PolarisCallContext polarisCallContext = callContext.getPolarisCallContext(); - if (!polarisCallContext - .getConfigurationStore() - .getConfiguration(polarisCallContext, PolarisConfiguration.SUPPORTED_CATALOG_STORAGE_TYPES) + if (!configurationStore + .getConfiguration(realmContext, PolarisConfiguration.SUPPORTED_CATALOG_STORAGE_TYPES) .contains("FILE")) { Assertions.assertThatThrownBy(() -> catalog.sendNotification(table, request)) .isInstanceOf(ForbiddenException.class) @@ -931,9 +941,8 @@ public void testUpdateNotificationCreateTableWithHttpPrefix() { httpsMetadataLocation, TableMetadataParser.toJson(createSampleTableMetadata(metadataLocation)).getBytes(UTF_8)); - if (!polarisCallContext - .getConfigurationStore() - .getConfiguration(polarisCallContext, PolarisConfiguration.SUPPORTED_CATALOG_STORAGE_TYPES) + if (!configurationStore + .getConfiguration(realmContext, PolarisConfiguration.SUPPORTED_CATALOG_STORAGE_TYPES) .contains("FILE")) { Assertions.assertThatThrownBy(() -> catalog.sendNotification(table, newRequest)) .isInstanceOf(ForbiddenException.class) @@ -1359,13 +1368,13 @@ public void testDropTableWithPurge() { .as("Table should not exist after drop") .rejects(TABLE); List tasks = - metaStoreManager.loadTasks(polarisContext, "testExecutor", 1).getEntities(); + metaStoreManager.loadTasks(metaStoreSession, "testExecutor", 1).getEntities(); Assertions.assertThat(tasks).hasSize(1); TaskEntity taskEntity = TaskEntity.of(tasks.get(0)); EnumMap credentials = metaStoreManager .getSubscopedCredsForEntity( - polarisContext, + metaStoreSession, 0, taskEntity.getId(), true, @@ -1380,10 +1389,8 @@ public void testDropTableWithPurge() { .containsEntry(PolarisCredentialProperty.AWS_TOKEN, SESSION_TOKEN); FileIO fileIO = new TaskFileIOSupplier( - createMockMetaStoreManagerFactory(), - new DefaultFileIOFactory(), - polarisContext.getConfigurationStore()) - .apply(taskEntity); + createMockMetaStoreManagerFactory(), new DefaultFileIOFactory(), configurationStore) + .apply(taskEntity, realmContext); Assertions.assertThat(fileIO).isNotNull().isInstanceOf(InMemoryFileIO.class); } @@ -1410,19 +1417,19 @@ public void testDropTableWithPurgeDisabled() { .addProperty(PolarisConfiguration.DROP_WITH_PURGE_ENABLED.catalogConfig(), "false") .setStorageConfigurationInfo(noPurgeStorageConfigModel, storageLocation) .build()); - RealmContext realmContext = () -> "realm"; - CallContext callContext = CallContext.of(realmContext, polarisContext); PolarisPassthroughResolutionView passthroughView = new PolarisPassthroughResolutionView( - callContext, entityManager, securityContext, noPurgeCatalogName); + entityManager, metaStoreSession, securityContext, noPurgeCatalogName); BasePolarisCatalog noPurgeCatalog = new BasePolarisCatalog( + realmContext, entityManager, metaStoreManager, - callContext, + metaStoreSession, + configurationStore, + diagServices, passthroughView, securityContext, - authenticatedRoot, Mockito.mock(), new DefaultFileIOFactory()); noPurgeCatalog.initialize( @@ -1492,22 +1499,21 @@ public void testRetriableException() { @Test public void testFileIOWrapper() { - RealmContext realmContext = () -> "realm"; - CallContext callContext = CallContext.of(realmContext, polarisContext); - CallContext.setCurrentContext(callContext); PolarisPassthroughResolutionView passthroughView = new PolarisPassthroughResolutionView( - callContext, entityManager, securityContext, CATALOG_NAME); + entityManager, metaStoreSession, securityContext, CATALOG_NAME); TestFileIOFactory measured = new TestFileIOFactory(); BasePolarisCatalog catalog = new BasePolarisCatalog( + realmContext, entityManager, metaStoreManager, - callContext, + metaStoreSession, + configurationStore, + diagServices, passthroughView, securityContext, - authenticatedRoot, Mockito.mock(), measured); catalog.initialize( @@ -1537,13 +1543,18 @@ public void testFileIOWrapper() { new TableCleanupTaskHandler( Mockito.mock(), createMockMetaStoreManagerFactory(), - (task) -> measured.loadFileIO("org.apache.iceberg.inmemory.InMemoryFileIO", Map.of())); + configurationStore, + diagServices, + (task, rc) -> + measured.loadFileIO("org.apache.iceberg.inmemory.InMemoryFileIO", Map.of()), + clock); handler.handleTask( TaskEntity.of( metaStoreManager - .loadTasks(polarisContext, "testExecutor", 1) + .loadTasks(metaStoreSession, "testExecutor", 1) .getEntities() - .getFirst())); + .getFirst()), + realmContext); Assertions.assertThat(measured.getNumDeletedFiles()).as("A table was deleted").isGreaterThan(0); } } diff --git a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/catalog/BasePolarisCatalogViewTest.java b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/catalog/BasePolarisCatalogViewTest.java index e1785c616..f896c63a0 100644 --- a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/catalog/BasePolarisCatalogViewTest.java +++ b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/catalog/BasePolarisCatalogViewTest.java @@ -29,13 +29,11 @@ import java.lang.reflect.Field; import java.lang.reflect.Method; import java.nio.file.Path; -import java.time.Clock; import java.util.List; import java.util.Set; import org.apache.iceberg.CatalogProperties; import org.apache.iceberg.catalog.Catalog; import org.apache.iceberg.view.ViewCatalogTests; -import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.PolarisConfiguration; import org.apache.polaris.core.PolarisConfigurationStore; import org.apache.polaris.core.PolarisDiagnostics; @@ -43,7 +41,6 @@ import org.apache.polaris.core.admin.model.StorageConfigInfo; import org.apache.polaris.core.auth.AuthenticatedPolarisPrincipal; import org.apache.polaris.core.auth.PolarisAuthorizerImpl; -import org.apache.polaris.core.context.CallContext; import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.CatalogEntity; import org.apache.polaris.core.entity.PolarisEntity; @@ -53,11 +50,11 @@ import org.apache.polaris.core.persistence.MetaStoreManagerFactory; import org.apache.polaris.core.persistence.PolarisEntityManager; import org.apache.polaris.core.persistence.PolarisMetaStoreManager; -import org.apache.polaris.core.persistence.cache.EntityCache; -import org.apache.polaris.core.storage.cache.StorageCredentialCache; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; import org.apache.polaris.service.admin.PolarisAdminService; import org.apache.polaris.service.catalog.BasePolarisCatalog; import org.apache.polaris.service.catalog.io.DefaultFileIOFactory; +import org.apache.polaris.service.config.RealmEntityManagerFactory; import org.apache.polaris.service.storage.PolarisStorageIntegrationProviderImpl; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeAll; @@ -71,14 +68,14 @@ public class BasePolarisCatalogViewTest extends ViewCatalogTests realmName; metaStoreManager = managerFactory.getOrCreateMetaStoreManager(realmContext); - polarisContext = - new PolarisCallContext( - managerFactory.getOrCreateSessionSupplier(realmContext).get(), - diagServices, - configurationStore, - Clock.systemDefaultZone()); - - PolarisEntityManager entityManager = - new PolarisEntityManager( - metaStoreManager, new StorageCredentialCache(), new EntityCache(metaStoreManager)); - - CallContext callContext = CallContext.of(realmContext, polarisContext); - CallContext.setCurrentContext(callContext); + metaStoreSession = managerFactory.getOrCreateSessionSupplier(realmContext).get(); PrincipalEntity rootEntity = new PrincipalEntity( PolarisEntity.of( metaStoreManager .readEntityByName( - polarisContext, + metaStoreSession, null, PolarisEntityType.PRINCIPAL, PolarisEntitySubType.NULL_SUBTYPE, @@ -132,14 +117,20 @@ public void before(TestInfo testInfo) { AuthenticatedPolarisPrincipal authenticatedRoot = new AuthenticatedPolarisPrincipal(rootEntity, Set.of()); + PolarisEntityManager entityManager = + entityManagerFactory.getOrCreateEntityManager(realmContext); + SecurityContext securityContext = Mockito.mock(SecurityContext.class); when(securityContext.getUserPrincipal()).thenReturn(authenticatedRoot); when(securityContext.isUserInRole(Mockito.anyString())).thenReturn(true); PolarisAdminService adminService = new PolarisAdminService( - callContext, + realmContext, entityManager, metaStoreManager, + metaStoreSession, + configurationStore, + diagServices, securityContext, new PolarisAuthorizerImpl(new PolarisConfigurationStore() {})); adminService.createCatalog( @@ -157,15 +148,17 @@ public void before(TestInfo testInfo) { PolarisPassthroughResolutionView passthroughView = new PolarisPassthroughResolutionView( - callContext, entityManager, securityContext, CATALOG_NAME); + entityManager, metaStoreSession, securityContext, CATALOG_NAME); this.catalog = new BasePolarisCatalog( + realmContext, entityManager, metaStoreManager, - callContext, + metaStoreSession, + configurationStore, + diagServices, passthroughView, securityContext, - authenticatedRoot, Mockito.mock(), new DefaultFileIOFactory()); this.catalog.initialize( @@ -177,7 +170,7 @@ public void before(TestInfo testInfo) { @AfterEach public void after() throws IOException { catalog().close(); - metaStoreManager.purge(polarisContext); + metaStoreManager.purge(metaStoreSession); } @Override diff --git a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/catalog/PolarisCatalogHandlerWrapperAuthzTest.java b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/catalog/PolarisCatalogHandlerWrapperAuthzTest.java index f56356e24..4cff22227 100644 --- a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/catalog/PolarisCatalogHandlerWrapperAuthzTest.java +++ b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/catalog/PolarisCatalogHandlerWrapperAuthzTest.java @@ -54,18 +54,15 @@ import org.apache.polaris.core.admin.model.PrincipalWithCredentialsCredentials; import org.apache.polaris.core.admin.model.StorageConfigInfo; import org.apache.polaris.core.auth.AuthenticatedPolarisPrincipal; -import org.apache.polaris.core.context.CallContext; import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.CatalogEntity; import org.apache.polaris.core.entity.CatalogRoleEntity; import org.apache.polaris.core.entity.PolarisPrivilege; import org.apache.polaris.core.entity.PrincipalEntity; -import org.apache.polaris.core.persistence.PolarisEntityManager; import org.apache.polaris.core.persistence.PolarisMetaStoreManager; import org.apache.polaris.core.persistence.resolver.PolarisResolutionManifest; import org.apache.polaris.service.catalog.PolarisCatalogHandlerWrapper; import org.apache.polaris.service.catalog.io.DefaultFileIOFactory; -import org.apache.polaris.service.config.RealmEntityManagerFactory; import org.apache.polaris.service.context.CallContextCatalogFactory; import org.apache.polaris.service.context.PolarisCallContextCatalogFactory; import org.apache.polaris.service.quarkus.admin.PolarisAuthzTestBase; @@ -93,7 +90,7 @@ private PolarisCatalogHandlerWrapper newWrapper() { } private PolarisCatalogHandlerWrapper newWrapper(Set activatedPrincipalRoles) { - return newWrapper(activatedPrincipalRoles, CATALOG_NAME, callContextCatalogFactory); + return newWrapper(activatedPrincipalRoles, CATALOG_NAME, newCatalogFactory()); } private PolarisCatalogHandlerWrapper newWrapper( @@ -101,7 +98,10 @@ private PolarisCatalogHandlerWrapper newWrapper( final AuthenticatedPolarisPrincipal authenticatedPrincipal = new AuthenticatedPolarisPrincipal(principalEntity, activatedPrincipalRoles); return new PolarisCatalogHandlerWrapper( - callContext, + realmContext, + metaStoreSession, + configurationStore, + diagServices, entityManager, metaStoreManager, securityContext(authenticatedPrincipal, activatedPrincipalRoles), @@ -110,6 +110,31 @@ private PolarisCatalogHandlerWrapper newWrapper( polarisAuthorizer); } + private PolarisCatalogHandlerWrapper newWrapper(SecurityContext securityContext) { + return new PolarisCatalogHandlerWrapper( + realmContext, + metaStoreSession, + configurationStore, + diagServices, + entityManager, + metaStoreManager, + securityContext, + newCatalogFactory(), + CATALOG_NAME, + polarisAuthorizer); + } + + private CallContextCatalogFactory newCatalogFactory() { + return new TestPolarisCallContextCatalogFactory( + entityManager, + metaStoreManager, + metaStoreSession, + configurationStore, + diagServices, + Mockito.mock(), + new DefaultFileIOFactory()); + } + /** * Tests each "sufficient" privilege individually using CATALOG_ROLE1 by granting at the * CATALOG_NAME level, revoking after each test, and also ensuring that the request fails after @@ -224,7 +249,7 @@ public void testInsufficientPermissionsPriorToSecretRotation() { String principalName = "all_the_powers"; PolarisMetaStoreManager.CreatePrincipalResult newPrincipal = metaStoreManager.createPrincipal( - callContext.getPolarisCallContext(), + metaStoreSession, new PrincipalEntity.Builder() .setName(principalName) .setCreateTimestamp(Instant.now().toEpochMilli()) @@ -238,14 +263,8 @@ public void testInsufficientPermissionsPriorToSecretRotation() { PrincipalEntity.of(newPrincipal.getPrincipal()), Set.of(PRINCIPAL_ROLE1, PRINCIPAL_ROLE2)); PolarisCatalogHandlerWrapper wrapper = - new PolarisCatalogHandlerWrapper( - callContext, - entityManager, - metaStoreManager, - securityContext(authenticatedPrincipal, Set.of(PRINCIPAL_ROLE1, PRINCIPAL_ROLE2)), - callContextCatalogFactory, - CATALOG_NAME, - polarisAuthorizer); + newWrapper( + securityContext(authenticatedPrincipal, Set.of(PRINCIPAL_ROLE1, PRINCIPAL_ROLE2))); // a variety of actions are all disallowed because the principal's credentials must be rotated doTestInsufficientPrivileges( @@ -263,21 +282,13 @@ public void testInsufficientPermissionsPriorToSecretRotation() { new PrincipalWithCredentialsCredentials( newPrincipal.getPrincipalSecrets().getPrincipalClientId(), newPrincipal.getPrincipalSecrets().getMainSecret()); - PrincipalEntity refreshPrincipal = - rotateAndRefreshPrincipal( - metaStoreManager, principalName, credentials, callContext.getPolarisCallContext()); + PrincipalEntity refreshPrincipal = rotateAndRefreshPrincipal(principalName, credentials); final AuthenticatedPolarisPrincipal authenticatedPrincipal1 = new AuthenticatedPolarisPrincipal( PrincipalEntity.of(refreshPrincipal), Set.of(PRINCIPAL_ROLE1, PRINCIPAL_ROLE2)); PolarisCatalogHandlerWrapper refreshedWrapper = - new PolarisCatalogHandlerWrapper( - callContext, - entityManager, - metaStoreManager, - securityContext(authenticatedPrincipal1, Set.of(PRINCIPAL_ROLE1, PRINCIPAL_ROLE2)), - callContextCatalogFactory, - CATALOG_NAME, - polarisAuthorizer); + newWrapper( + securityContext(authenticatedPrincipal1, Set.of(PRINCIPAL_ROLE1, PRINCIPAL_ROLE2))); doTestSufficientPrivilegeSets( List.of(Set.of(PolarisPrivilege.NAMESPACE_LIST)), @@ -1690,42 +1701,43 @@ public void testSendNotificationSufficientPrivileges() { PolarisCallContextCatalogFactory factory = new PolarisCallContextCatalogFactory( - new RealmEntityManagerFactory(null) { - @Override - public PolarisEntityManager getOrCreateEntityManager(RealmContext realmContext) { - return entityManager; - } - }, - managerFactory, + entityManager, + metaStoreManager, + metaStoreSession, + configurationStore, + diagServices, Mockito.mock(), new DefaultFileIOFactory()) { @Override public Catalog createCallContextCatalog( - CallContext context, + RealmContext realmContext, AuthenticatedPolarisPrincipal authenticatedPolarisPrincipal, SecurityContext securityContext, PolarisResolutionManifest resolvedManifest) { Catalog catalog = super.createCallContextCatalog( - context, authenticatedPolarisPrincipal, securityContext, resolvedManifest); + realmContext, authenticatedPolarisPrincipal, securityContext, resolvedManifest); String fileIoImpl = "org.apache.iceberg.inmemory.InMemoryFileIO"; catalog.initialize( externalCatalog, ImmutableMap.of(CatalogProperties.FILE_IO_IMPL, fileIoImpl)); - FileIO fileIO = CatalogUtil.loadFileIO(fileIoImpl, Map.of(), new Configuration()); - TableMetadata tableMetadata = - TableMetadata.buildFromEmpty() - .addSchema(SCHEMA, SCHEMA.highestFieldId()) - .setLocation( - String.format("%s/bucket/table/metadata/v1.metadata.json", storageLocation)) - .addPartitionSpec(PartitionSpec.unpartitioned()) - .addSortOrder(SortOrder.unsorted()) - .assignUUID() - .build(); - TableMetadataParser.overwrite( - tableMetadata, fileIO.newOutputFile(createPayload.getMetadataLocation())); - TableMetadataParser.overwrite( - tableMetadata, fileIO.newOutputFile(updatePayload.getMetadataLocation())); + try (FileIO fileIO = + CatalogUtil.loadFileIO(fileIoImpl, Map.of(), new Configuration())) { + TableMetadata tableMetadata = + TableMetadata.buildFromEmpty() + .addSchema(SCHEMA, SCHEMA.highestFieldId()) + .setLocation( + String.format( + "%s/bucket/table/metadata/v1.metadata.json", storageLocation)) + .addPartitionSpec(PartitionSpec.unpartitioned()) + .addSortOrder(SortOrder.unsorted()) + .assignUUID() + .build(); + TableMetadataParser.overwrite( + tableMetadata, fileIO.newOutputFile(createPayload.getMetadataLocation())); + TableMetadataParser.overwrite( + tableMetadata, fileIO.newOutputFile(updatePayload.getMetadataLocation())); + } return catalog; } }; diff --git a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/catalog/PolarisPassthroughResolutionView.java b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/catalog/PolarisPassthroughResolutionView.java index 6ff31fd35..db74ad3b4 100644 --- a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/catalog/PolarisPassthroughResolutionView.java +++ b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/catalog/PolarisPassthroughResolutionView.java @@ -23,10 +23,10 @@ import org.apache.iceberg.catalog.Namespace; import org.apache.iceberg.catalog.TableIdentifier; import org.apache.polaris.core.catalog.PolarisCatalogHelpers; -import org.apache.polaris.core.context.CallContext; import org.apache.polaris.core.entity.PolarisEntitySubType; import org.apache.polaris.core.entity.PolarisEntityType; import org.apache.polaris.core.persistence.PolarisEntityManager; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; import org.apache.polaris.core.persistence.PolarisResolvedPathWrapper; import org.apache.polaris.core.persistence.resolver.PolarisResolutionManifest; import org.apache.polaris.core.persistence.resolver.PolarisResolutionManifestCatalogView; @@ -41,17 +41,17 @@ */ public class PolarisPassthroughResolutionView implements PolarisResolutionManifestCatalogView { private final PolarisEntityManager entityManager; - private final CallContext callContext; + private final PolarisMetaStoreSession metaStoreSession; private final SecurityContext securityContext; private final String catalogName; public PolarisPassthroughResolutionView( - CallContext callContext, PolarisEntityManager entityManager, + PolarisMetaStoreSession metaStoreSession, SecurityContext securityContext, String catalogName) { this.entityManager = entityManager; - this.callContext = callContext; + this.metaStoreSession = metaStoreSession; this.securityContext = securityContext; this.catalogName = catalogName; } @@ -59,7 +59,7 @@ public PolarisPassthroughResolutionView( @Override public PolarisResolvedPathWrapper getResolvedReferenceCatalogEntity() { PolarisResolutionManifest manifest = - entityManager.prepareResolutionManifest(callContext, securityContext, catalogName); + entityManager.prepareResolutionManifest(metaStoreSession, securityContext, catalogName); manifest.resolveAll(); return manifest.getResolvedReferenceCatalogEntity(); } @@ -67,7 +67,7 @@ public PolarisResolvedPathWrapper getResolvedReferenceCatalogEntity() { @Override public PolarisResolvedPathWrapper getResolvedPath(Object key) { PolarisResolutionManifest manifest = - entityManager.prepareResolutionManifest(callContext, securityContext, catalogName); + entityManager.prepareResolutionManifest(metaStoreSession, securityContext, catalogName); if (key instanceof Namespace namespace) { manifest.addPath( @@ -85,7 +85,7 @@ public PolarisResolvedPathWrapper getResolvedPath(Object key) { @Override public PolarisResolvedPathWrapper getResolvedPath(Object key, PolarisEntitySubType subType) { PolarisResolutionManifest manifest = - entityManager.prepareResolutionManifest(callContext, securityContext, catalogName); + entityManager.prepareResolutionManifest(metaStoreSession, securityContext, catalogName); if (key instanceof TableIdentifier identifier) { manifest.addPath( @@ -106,7 +106,7 @@ public PolarisResolvedPathWrapper getResolvedPath(Object key, PolarisEntitySubTy @Override public PolarisResolvedPathWrapper getPassthroughResolvedPath(Object key) { PolarisResolutionManifest manifest = - entityManager.prepareResolutionManifest(callContext, securityContext, catalogName); + entityManager.prepareResolutionManifest(metaStoreSession, securityContext, catalogName); if (key instanceof Namespace namespace) { manifest.addPassthroughPath( @@ -124,7 +124,7 @@ public PolarisResolvedPathWrapper getPassthroughResolvedPath(Object key) { public PolarisResolvedPathWrapper getPassthroughResolvedPath( Object key, PolarisEntitySubType subType) { PolarisResolutionManifest manifest = - entityManager.prepareResolutionManifest(callContext, securityContext, catalogName); + entityManager.prepareResolutionManifest(metaStoreSession, securityContext, catalogName); if (key instanceof TableIdentifier identifier) { manifest.addPassthroughPath( diff --git a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/config/DefaultConfigurationStoreTest.java b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/config/DefaultConfigurationStoreTest.java index ac7c0116a..937dd3429 100644 --- a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/config/DefaultConfigurationStoreTest.java +++ b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/config/DefaultConfigurationStoreTest.java @@ -21,10 +21,8 @@ import static org.assertj.core.api.Assertions.assertThat; import java.util.Map; -import org.apache.polaris.core.PolarisCallContext; -import org.apache.polaris.core.PolarisDefaultDiagServiceImpl; +import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.service.config.DefaultConfigurationStore; -import org.apache.polaris.service.persistence.InMemoryPolarisMetaStoreManagerFactory; import org.junit.jupiter.api.Test; public class DefaultConfigurationStoreTest { @@ -33,21 +31,17 @@ public class DefaultConfigurationStoreTest { public void testGetConfiguration() { DefaultConfigurationStore defaultConfigurationStore = new DefaultConfigurationStore(Map.of("key1", 1, "key2", "value")); - InMemoryPolarisMetaStoreManagerFactory metastoreFactory = - new InMemoryPolarisMetaStoreManagerFactory(); - PolarisCallContext callCtx = - new PolarisCallContext( - metastoreFactory.getOrCreateSessionSupplier(() -> "realm1").get(), - new PolarisDefaultDiagServiceImpl()); - Object value = defaultConfigurationStore.getConfiguration(callCtx, "missingKeyWithoutDefault"); + RealmContext realmContext = () -> "test"; + Object value = + defaultConfigurationStore.getConfiguration(realmContext, "missingKeyWithoutDefault"); assertThat(value).isNull(); Object defaultValue = defaultConfigurationStore.getConfiguration( - callCtx, "missingKeyWithDefault", "defaultValue"); + realmContext, "missingKeyWithDefault", "defaultValue"); assertThat(defaultValue).isEqualTo("defaultValue"); - Integer keyOne = defaultConfigurationStore.getConfiguration(callCtx, "key1"); + Integer keyOne = defaultConfigurationStore.getConfiguration(realmContext, "key1"); assertThat(keyOne).isEqualTo(1); - String keyTwo = defaultConfigurationStore.getConfiguration(callCtx, "key2"); + String keyTwo = defaultConfigurationStore.getConfiguration(realmContext, "key2"); assertThat(keyTwo).isEqualTo("value"); } @@ -67,44 +61,33 @@ public void testGetRealmConfiguration() { Map.of("key1", realm1KeyOneValue), "realm2", Map.of("key1", realm2KeyOneValue, "key2", realm2KeyTwoValue))); - InMemoryPolarisMetaStoreManagerFactory metastoreFactory = - new InMemoryPolarisMetaStoreManagerFactory(); // check realm1 values - PolarisCallContext realm1Ctx = - new PolarisCallContext( - metastoreFactory.getOrCreateSessionSupplier(() -> "realm1").get(), - new PolarisDefaultDiagServiceImpl()); + RealmContext realmContext = () -> "realm1"; Object value = - defaultConfigurationStore.getConfiguration(realm1Ctx, "missingKeyWithoutDefault"); + defaultConfigurationStore.getConfiguration(realmContext, "missingKeyWithoutDefault"); assertThat(value).isNull(); Object defaultValue = defaultConfigurationStore.getConfiguration( - realm1Ctx, "missingKeyWithDefault", "defaultValue"); + realmContext, "missingKeyWithDefault", "defaultValue"); assertThat(defaultValue).isEqualTo("defaultValue"); - Integer keyOneRealm1 = defaultConfigurationStore.getConfiguration(realm1Ctx, "key1"); + Integer keyOneRealm1 = defaultConfigurationStore.getConfiguration(realmContext, "key1"); assertThat(keyOneRealm1).isEqualTo(realm1KeyOneValue); - String keyTwoRealm1 = defaultConfigurationStore.getConfiguration(realm1Ctx, "key2"); + String keyTwoRealm1 = defaultConfigurationStore.getConfiguration(realmContext, "key2"); assertThat(keyTwoRealm1).isEqualTo(defaultKeyTwoValue); // check realm2 values - PolarisCallContext realm2Ctx = - new PolarisCallContext( - metastoreFactory.getOrCreateSessionSupplier(() -> "realm2").get(), - new PolarisDefaultDiagServiceImpl()); - Integer keyOneRealm2 = defaultConfigurationStore.getConfiguration(realm2Ctx, "key1"); + realmContext = () -> "realm2"; + Integer keyOneRealm2 = defaultConfigurationStore.getConfiguration(realmContext, "key1"); assertThat(keyOneRealm2).isEqualTo(realm2KeyOneValue); - String keyTwoRealm2 = defaultConfigurationStore.getConfiguration(realm2Ctx, "key2"); + String keyTwoRealm2 = defaultConfigurationStore.getConfiguration(realmContext, "key2"); assertThat(keyTwoRealm2).isEqualTo(realm2KeyTwoValue); // realm3 has no realm-overrides, so just returns default values - PolarisCallContext realm3Ctx = - new PolarisCallContext( - metastoreFactory.getOrCreateSessionSupplier(() -> "realm3").get(), - new PolarisDefaultDiagServiceImpl()); - Integer keyOneRealm3 = defaultConfigurationStore.getConfiguration(realm3Ctx, "key1"); + realmContext = () -> "realm3"; + Integer keyOneRealm3 = defaultConfigurationStore.getConfiguration(realmContext, "key1"); assertThat(keyOneRealm3).isEqualTo(defaultKeyOneValue); - String keyTwoRealm3 = defaultConfigurationStore.getConfiguration(realm3Ctx, "key2"); + String keyTwoRealm3 = defaultConfigurationStore.getConfiguration(realmContext, "key2"); assertThat(keyTwoRealm3).isEqualTo(defaultKeyTwoValue); } } diff --git a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/task/ManifestFileCleanupTaskHandlerTest.java b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/task/ManifestFileCleanupTaskHandlerTest.java index 18196d476..e0932252b 100644 --- a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/task/ManifestFileCleanupTaskHandlerTest.java +++ b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/task/ManifestFileCleanupTaskHandlerTest.java @@ -20,7 +20,6 @@ import static java.nio.charset.StandardCharsets.UTF_8; import static org.assertj.core.api.Assertions.assertThat; -import static org.assertj.core.api.Assertions.assertThatPredicate; import io.quarkus.test.junit.QuarkusTest; import jakarta.inject.Inject; @@ -46,418 +45,363 @@ import org.apache.iceberg.io.FileIO; import org.apache.iceberg.io.OutputFile; import org.apache.iceberg.io.PositionOutputStream; -import org.apache.polaris.core.PolarisCallContext; -import org.apache.polaris.core.PolarisDefaultDiagServiceImpl; -import org.apache.polaris.core.context.CallContext; +import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.AsyncTaskType; import org.apache.polaris.core.entity.TaskEntity; -import org.apache.polaris.core.persistence.MetaStoreManagerFactory; import org.apache.polaris.service.task.ManifestFileCleanupTaskHandler; import org.apache.polaris.service.task.TaskUtils; import org.junit.jupiter.api.Test; @QuarkusTest class ManifestFileCleanupTaskHandlerTest { - @Inject MetaStoreManagerFactory metaStoreManagerFactory; + + @Inject PolarisDiagnostics diagnostics; private final RealmContext realmContext = () -> "realmName"; @Test public void testCleanupFileNotExists() throws IOException { - PolarisCallContext polarisCallContext = - new PolarisCallContext( - metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(), - new PolarisDefaultDiagServiceImpl()); - try (CallContext callCtx = CallContext.of(realmContext, polarisCallContext)) { - CallContext.setCurrentContext(callCtx); - FileIO fileIO = new InMemoryFileIO(); - TableIdentifier tableIdentifier = - TableIdentifier.of(Namespace.of("db1", "schema1"), "table1"); - ManifestFileCleanupTaskHandler handler = - new ManifestFileCleanupTaskHandler((task) -> fileIO, Executors.newSingleThreadExecutor()); - ManifestFile manifestFile = - TaskTestUtils.manifestFile( - fileIO, "manifest1.avro", 1L, "dataFile1.parquet", "dataFile2.parquet"); - fileIO.deleteFile(manifestFile.path()); - TaskEntity task = - new TaskEntity.Builder() - .withTaskType(AsyncTaskType.MANIFEST_FILE_CLEANUP) - .withData( - new ManifestFileCleanupTaskHandler.ManifestCleanupTask( - tableIdentifier, - Base64.encodeBase64String(ManifestFiles.encode(manifestFile)))) - .setName(UUID.randomUUID().toString()) - .build(); - assertThatPredicate(handler::canHandleTask).accepts(task); - assertThatPredicate(handler::handleTask).accepts(task); - } + FileIO fileIO = new InMemoryFileIO(); + TableIdentifier tableIdentifier = TableIdentifier.of(Namespace.of("db1", "schema1"), "table1"); + ManifestFileCleanupTaskHandler handler = + new ManifestFileCleanupTaskHandler( + (task, rc) -> fileIO, Executors.newSingleThreadExecutor(), diagnostics); + ManifestFile manifestFile = + TaskTestUtils.manifestFile( + fileIO, "manifest1.avro", 1L, "dataFile1.parquet", "dataFile2.parquet"); + fileIO.deleteFile(manifestFile.path()); + TaskEntity task = + new TaskEntity.Builder() + .withTaskType(diagnostics, AsyncTaskType.MANIFEST_FILE_CLEANUP) + .withData( + diagnostics, + new ManifestFileCleanupTaskHandler.ManifestCleanupTask( + tableIdentifier, Base64.encodeBase64String(ManifestFiles.encode(manifestFile)))) + .setName(UUID.randomUUID().toString()) + .build(); + assertThat(handler.canHandleTask(task)).isTrue(); + assertThat(handler.handleTask(task, realmContext)).isTrue(); } @Test public void testCleanupFileManifestExistsDataFilesDontExist() throws IOException { - PolarisCallContext polarisCallContext = - new PolarisCallContext( - metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(), - new PolarisDefaultDiagServiceImpl()); - try (CallContext callCtx = CallContext.of(realmContext, polarisCallContext)) { - CallContext.setCurrentContext(callCtx); - FileIO fileIO = new InMemoryFileIO(); - TableIdentifier tableIdentifier = - TableIdentifier.of(Namespace.of("db1", "schema1"), "table1"); - ManifestFileCleanupTaskHandler handler = - new ManifestFileCleanupTaskHandler((task) -> fileIO, Executors.newSingleThreadExecutor()); - ManifestFile manifestFile = - TaskTestUtils.manifestFile( - fileIO, "manifest1.avro", 100L, "dataFile1.parquet", "dataFile2.parquet"); - TaskEntity task = - new TaskEntity.Builder() - .withTaskType(AsyncTaskType.MANIFEST_FILE_CLEANUP) - .withData( - new ManifestFileCleanupTaskHandler.ManifestCleanupTask( - tableIdentifier, - Base64.encodeBase64String(ManifestFiles.encode(manifestFile)))) - .setName(UUID.randomUUID().toString()) - .build(); - assertThatPredicate(handler::canHandleTask).accepts(task); - assertThatPredicate(handler::handleTask).accepts(task); - } + FileIO fileIO = new InMemoryFileIO(); + TableIdentifier tableIdentifier = TableIdentifier.of(Namespace.of("db1", "schema1"), "table1"); + ManifestFileCleanupTaskHandler handler = + new ManifestFileCleanupTaskHandler( + (task, rc) -> fileIO, Executors.newSingleThreadExecutor(), diagnostics); + ManifestFile manifestFile = + TaskTestUtils.manifestFile( + fileIO, "manifest1.avro", 100L, "dataFile1.parquet", "dataFile2.parquet"); + TaskEntity task = + new TaskEntity.Builder() + .withTaskType(diagnostics, AsyncTaskType.MANIFEST_FILE_CLEANUP) + .withData( + diagnostics, + new ManifestFileCleanupTaskHandler.ManifestCleanupTask( + tableIdentifier, Base64.encodeBase64String(ManifestFiles.encode(manifestFile)))) + .setName(UUID.randomUUID().toString()) + .build(); + assertThat(handler.canHandleTask(task)).isTrue(); + assertThat(handler.handleTask(task, realmContext)).isTrue(); } @Test public void testCleanupFiles() throws IOException { - PolarisCallContext polarisCallContext = - new PolarisCallContext( - metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(), - new PolarisDefaultDiagServiceImpl()); - try (CallContext callCtx = CallContext.of(realmContext, polarisCallContext)) { - CallContext.setCurrentContext(callCtx); - FileIO fileIO = - new InMemoryFileIO() { - @Override - public void close() { - // no-op - } - }; - TableIdentifier tableIdentifier = - TableIdentifier.of(Namespace.of("db1", "schema1"), "table1"); - ManifestFileCleanupTaskHandler handler = - new ManifestFileCleanupTaskHandler((task) -> fileIO, Executors.newSingleThreadExecutor()); - String dataFile1Path = "dataFile1.parquet"; - OutputFile dataFile1 = fileIO.newOutputFile(dataFile1Path); - PositionOutputStream out1 = dataFile1.createOrOverwrite(); - out1.write("the data".getBytes(UTF_8)); - out1.close(); - String dataFile2Path = "dataFile2.parquet"; - OutputFile dataFile2 = fileIO.newOutputFile(dataFile2Path); - PositionOutputStream out2 = dataFile2.createOrOverwrite(); - out2.write("the data".getBytes(UTF_8)); - out2.close(); - ManifestFile manifestFile = - TaskTestUtils.manifestFile(fileIO, "manifest1.avro", 100L, dataFile1Path, dataFile2Path); - TaskEntity task = - new TaskEntity.Builder() - .withTaskType(AsyncTaskType.MANIFEST_FILE_CLEANUP) - .withData( - new ManifestFileCleanupTaskHandler.ManifestCleanupTask( - tableIdentifier, - Base64.encodeBase64String(ManifestFiles.encode(manifestFile)))) - .setName(UUID.randomUUID().toString()) - .build(); - assertThatPredicate(handler::canHandleTask).accepts(task); - assertThatPredicate(handler::handleTask).accepts(task); - assertThatPredicate((String f) -> TaskUtils.exists(f, fileIO)).rejects(dataFile1Path); - assertThatPredicate((String f) -> TaskUtils.exists(f, fileIO)).rejects(dataFile2Path); - } + FileIO fileIO = + new InMemoryFileIO() { + @Override + public void close() { + // no-op + } + }; + TableIdentifier tableIdentifier = TableIdentifier.of(Namespace.of("db1", "schema1"), "table1"); + ManifestFileCleanupTaskHandler handler = + new ManifestFileCleanupTaskHandler( + (task, rc) -> fileIO, Executors.newSingleThreadExecutor(), diagnostics); + String dataFile1Path = "dataFile1.parquet"; + OutputFile dataFile1 = fileIO.newOutputFile(dataFile1Path); + PositionOutputStream out1 = dataFile1.createOrOverwrite(); + out1.write("the data".getBytes(UTF_8)); + out1.close(); + String dataFile2Path = "dataFile2.parquet"; + OutputFile dataFile2 = fileIO.newOutputFile(dataFile2Path); + PositionOutputStream out2 = dataFile2.createOrOverwrite(); + out2.write("the data".getBytes(UTF_8)); + out2.close(); + ManifestFile manifestFile = + TaskTestUtils.manifestFile(fileIO, "manifest1.avro", 100L, dataFile1Path, dataFile2Path); + TaskEntity task = + new TaskEntity.Builder() + .withTaskType(diagnostics, AsyncTaskType.MANIFEST_FILE_CLEANUP) + .withData( + diagnostics, + new ManifestFileCleanupTaskHandler.ManifestCleanupTask( + tableIdentifier, Base64.encodeBase64String(ManifestFiles.encode(manifestFile)))) + .setName(UUID.randomUUID().toString()) + .build(); + assertThat(handler.canHandleTask(task)).isTrue(); + assertThat(handler.handleTask(task, realmContext)).isTrue(); + assertThat(TaskUtils.exists(dataFile1Path, fileIO)).isFalse(); + assertThat(TaskUtils.exists(dataFile2Path, fileIO)).isFalse(); } @Test public void testCleanupFilesWithRetries() throws IOException { - PolarisCallContext polarisCallContext = - new PolarisCallContext( - metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(), - new PolarisDefaultDiagServiceImpl()); - try (CallContext callCtx = CallContext.of(realmContext, polarisCallContext)) { - CallContext.setCurrentContext(callCtx); - Map retryCounter = new HashMap<>(); - FileIO fileIO = - new InMemoryFileIO() { - @Override - public void close() { - // no-op - } + Map retryCounter = new HashMap<>(); + FileIO fileIO = + new InMemoryFileIO() { + @Override + public void close() { + // no-op + } - @Override - public void deleteFile(String location) { - int attempts = - retryCounter - .computeIfAbsent(location, k -> new AtomicInteger(0)) - .incrementAndGet(); - if (attempts < 3) { - throw new RuntimeException("I'm failing to test retries"); - } else { - // succeed on the third attempt - super.deleteFile(location); - } + @Override + public void deleteFile(String location) { + int attempts = + retryCounter.computeIfAbsent(location, k -> new AtomicInteger(0)).incrementAndGet(); + if (attempts < 3) { + throw new RuntimeException("I'm failing to test retries"); + } else { + // succeed on the third attempt + super.deleteFile(location); } - }; + } + }; - TableIdentifier tableIdentifier = - TableIdentifier.of(Namespace.of("db1", "schema1"), "table1"); - ManifestFileCleanupTaskHandler handler = - new ManifestFileCleanupTaskHandler((task) -> fileIO, Executors.newSingleThreadExecutor()); - String dataFile1Path = "dataFile1.parquet"; - OutputFile dataFile1 = fileIO.newOutputFile(dataFile1Path); - PositionOutputStream out1 = dataFile1.createOrOverwrite(); - out1.write("the data".getBytes(UTF_8)); - out1.close(); - String dataFile2Path = "dataFile2.parquet"; - OutputFile dataFile2 = fileIO.newOutputFile(dataFile2Path); - PositionOutputStream out2 = dataFile2.createOrOverwrite(); - out2.write("the data".getBytes(UTF_8)); - out2.close(); - ManifestFile manifestFile = - TaskTestUtils.manifestFile(fileIO, "manifest1.avro", 100L, dataFile1Path, dataFile2Path); - TaskEntity task = - new TaskEntity.Builder() - .withTaskType(AsyncTaskType.MANIFEST_FILE_CLEANUP) - .withData( - new ManifestFileCleanupTaskHandler.ManifestCleanupTask( - tableIdentifier, - Base64.encodeBase64String(ManifestFiles.encode(manifestFile)))) - .setName(UUID.randomUUID().toString()) - .build(); - assertThatPredicate(handler::canHandleTask).accepts(task); - assertThatPredicate(handler::handleTask).accepts(task); - assertThatPredicate((String f) -> TaskUtils.exists(f, fileIO)).rejects(dataFile1Path); - assertThatPredicate((String f) -> TaskUtils.exists(f, fileIO)).rejects(dataFile2Path); - } + TableIdentifier tableIdentifier = TableIdentifier.of(Namespace.of("db1", "schema1"), "table1"); + ManifestFileCleanupTaskHandler handler = + new ManifestFileCleanupTaskHandler( + (task, rc) -> fileIO, Executors.newSingleThreadExecutor(), diagnostics); + String dataFile1Path = "dataFile1.parquet"; + OutputFile dataFile1 = fileIO.newOutputFile(dataFile1Path); + PositionOutputStream out1 = dataFile1.createOrOverwrite(); + out1.write("the data".getBytes(UTF_8)); + out1.close(); + String dataFile2Path = "dataFile2.parquet"; + OutputFile dataFile2 = fileIO.newOutputFile(dataFile2Path); + PositionOutputStream out2 = dataFile2.createOrOverwrite(); + out2.write("the data".getBytes(UTF_8)); + out2.close(); + ManifestFile manifestFile = + TaskTestUtils.manifestFile(fileIO, "manifest1.avro", 100L, dataFile1Path, dataFile2Path); + TaskEntity task = + new TaskEntity.Builder() + .withTaskType(diagnostics, AsyncTaskType.MANIFEST_FILE_CLEANUP) + .withData( + diagnostics, + new ManifestFileCleanupTaskHandler.ManifestCleanupTask( + tableIdentifier, Base64.encodeBase64String(ManifestFiles.encode(manifestFile)))) + .setName(UUID.randomUUID().toString()) + .build(); + assertThat(handler.canHandleTask(task)).isTrue(); + assertThat(handler.handleTask(task, realmContext)).isTrue(); + assertThat(TaskUtils.exists(dataFile1Path, fileIO)).isFalse(); + assertThat(TaskUtils.exists(dataFile2Path, fileIO)).isFalse(); } @Test public void testMetadataFileCleanup() throws IOException { - PolarisCallContext polarisCallContext = - new PolarisCallContext( - metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(), - new PolarisDefaultDiagServiceImpl()); - try (CallContext callCtx = CallContext.of(realmContext, polarisCallContext)) { - CallContext.setCurrentContext(callCtx); - FileIO fileIO = - new InMemoryFileIO() { - @Override - public void close() { - // no-op - } - }; - TableIdentifier tableIdentifier = - TableIdentifier.of(Namespace.of("db1", "schema1"), "table1"); - ManifestFileCleanupTaskHandler handler = - new ManifestFileCleanupTaskHandler((task) -> fileIO, Executors.newSingleThreadExecutor()); + FileIO fileIO = + new InMemoryFileIO() { + @Override + public void close() { + // no-op + } + }; + TableIdentifier tableIdentifier = TableIdentifier.of(Namespace.of("db1", "schema1"), "table1"); + ManifestFileCleanupTaskHandler handler = + new ManifestFileCleanupTaskHandler( + (task, rc) -> fileIO, Executors.newSingleThreadExecutor(), diagnostics); - long snapshotId1 = 100L; - ManifestFile manifestFile1 = - TaskTestUtils.manifestFile( - fileIO, "manifest1.avro", snapshotId1, "dataFile1.parquet", "dataFile2.parquet"); - ManifestFile manifestFile2 = - TaskTestUtils.manifestFile( - fileIO, "manifest2.avro", snapshotId1, "dataFile3.parquet", "dataFile4.parquet"); - Snapshot snapshot = - TaskTestUtils.newSnapshot( - fileIO, "manifestList.avro", 1, snapshotId1, 99L, manifestFile1, manifestFile2); - StatisticsFile statisticsFile1 = - TaskTestUtils.writeStatsFile( - snapshot.snapshotId(), - snapshot.sequenceNumber(), - "/metadata/" + UUID.randomUUID() + ".stats", - fileIO); - String firstMetadataFile = "v1-295495059.metadata.json"; - TableMetadata firstMetadata = - TaskTestUtils.writeTableMetadata( - fileIO, firstMetadataFile, List.of(statisticsFile1), snapshot); - assertThat(TaskUtils.exists(firstMetadataFile, fileIO)).isTrue(); + long snapshotId1 = 100L; + ManifestFile manifestFile1 = + TaskTestUtils.manifestFile( + fileIO, "manifest1.avro", snapshotId1, "dataFile1.parquet", "dataFile2.parquet"); + ManifestFile manifestFile2 = + TaskTestUtils.manifestFile( + fileIO, "manifest2.avro", snapshotId1, "dataFile3.parquet", "dataFile4.parquet"); + Snapshot snapshot = + TaskTestUtils.newSnapshot( + fileIO, "manifestList.avro", 1, snapshotId1, 99L, manifestFile1, manifestFile2); + StatisticsFile statisticsFile1 = + TaskTestUtils.writeStatsFile( + snapshot.snapshotId(), + snapshot.sequenceNumber(), + "/metadata/" + UUID.randomUUID() + ".stats", + fileIO); + String firstMetadataFile = "v1-295495059.metadata.json"; + TableMetadata firstMetadata = + TaskTestUtils.writeTableMetadata( + fileIO, firstMetadataFile, List.of(statisticsFile1), snapshot); + assertThat(TaskUtils.exists(firstMetadataFile, fileIO)).isTrue(); - ManifestFile manifestFile3 = - TaskTestUtils.manifestFile( - fileIO, "manifest3.avro", snapshot.snapshotId() + 1, "dataFile5.parquet"); - Snapshot snapshot2 = - TaskTestUtils.newSnapshot( - fileIO, - "manifestList2.avro", - snapshot.sequenceNumber() + 1, - snapshot.snapshotId() + 1, - snapshot.snapshotId(), - manifestFile1, - manifestFile3); // exclude manifest2 from the new snapshot - StatisticsFile statisticsFile2 = - TaskTestUtils.writeStatsFile( - snapshot2.snapshotId(), - snapshot2.sequenceNumber(), - "/metadata/" + UUID.randomUUID() + ".stats", - fileIO); - String secondMetadataFile = "v1-295495060.metadata.json"; - TableMetadata secondMetadata = - TaskTestUtils.writeTableMetadata( - fileIO, - secondMetadataFile, - firstMetadata, - firstMetadataFile, - List.of(statisticsFile2), - snapshot2); - assertThat(TaskUtils.exists(firstMetadataFile, fileIO)).isTrue(); - assertThat(TaskUtils.exists(secondMetadataFile, fileIO)).isTrue(); + ManifestFile manifestFile3 = + TaskTestUtils.manifestFile( + fileIO, "manifest3.avro", snapshot.snapshotId() + 1, "dataFile5.parquet"); + Snapshot snapshot2 = + TaskTestUtils.newSnapshot( + fileIO, + "manifestList2.avro", + snapshot.sequenceNumber() + 1, + snapshot.snapshotId() + 1, + snapshot.snapshotId(), + manifestFile1, + manifestFile3); // exclude manifest2 from the new snapshot + StatisticsFile statisticsFile2 = + TaskTestUtils.writeStatsFile( + snapshot2.snapshotId(), + snapshot2.sequenceNumber(), + "/metadata/" + UUID.randomUUID() + ".stats", + fileIO); + String secondMetadataFile = "v1-295495060.metadata.json"; + TableMetadata secondMetadata = + TaskTestUtils.writeTableMetadata( + fileIO, + secondMetadataFile, + firstMetadata, + firstMetadataFile, + List.of(statisticsFile2), + snapshot2); + assertThat(TaskUtils.exists(firstMetadataFile, fileIO)).isTrue(); + assertThat(TaskUtils.exists(secondMetadataFile, fileIO)).isTrue(); - List cleanupFiles = - Stream.concat( - secondMetadata.previousFiles().stream() - .map(TableMetadata.MetadataLogEntry::file) - .filter(file -> TaskUtils.exists(file, fileIO)), - secondMetadata.statisticsFiles().stream() - .map(StatisticsFile::path) - .filter(file -> TaskUtils.exists(file, fileIO))) - .toList(); + List cleanupFiles = + Stream.concat( + secondMetadata.previousFiles().stream() + .map(metadataLogEntry -> metadataLogEntry.file()) + .filter(file -> TaskUtils.exists(file, fileIO)), + secondMetadata.statisticsFiles().stream() + .map(statisticsFile -> statisticsFile.path()) + .filter(file -> TaskUtils.exists(file, fileIO))) + .toList(); - TaskEntity task = - new TaskEntity.Builder() - .withTaskType(AsyncTaskType.METADATA_FILE_BATCH_CLEANUP) - .withData( - new ManifestFileCleanupTaskHandler.ManifestCleanupTask( - tableIdentifier, cleanupFiles)) - .setName(UUID.randomUUID().toString()) - .build(); + TaskEntity task = + new TaskEntity.Builder() + .withTaskType(diagnostics, AsyncTaskType.METADATA_FILE_BATCH_CLEANUP) + .withData( + diagnostics, + new ManifestFileCleanupTaskHandler.ManifestCleanupTask( + tableIdentifier, cleanupFiles)) + .setName(UUID.randomUUID().toString()) + .build(); - assertThatPredicate(handler::canHandleTask).accepts(task); - assertThatPredicate(handler::handleTask).accepts(task); + assertThat(handler.canHandleTask(task)).isTrue(); + assertThat(handler.handleTask(task, realmContext)).isTrue(); - assertThatPredicate((String file) -> TaskUtils.exists(file, fileIO)) - .rejects(firstMetadataFile); - assertThatPredicate((String file) -> TaskUtils.exists(file, fileIO)) - .rejects(statisticsFile1.path()); - assertThatPredicate((String file) -> TaskUtils.exists(file, fileIO)) - .rejects(statisticsFile2.path()); - } + assertThat(TaskUtils.exists(firstMetadataFile, fileIO)).isFalse(); + assertThat(TaskUtils.exists(statisticsFile1.path(), fileIO)).isFalse(); + assertThat(TaskUtils.exists(statisticsFile2.path(), fileIO)).isFalse(); } @Test public void testMetadataFileCleanupIfFileNotExist() throws IOException { - PolarisCallContext polarisCallContext = - new PolarisCallContext( - metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(), - new PolarisDefaultDiagServiceImpl()); - try (CallContext callCtx = CallContext.of(realmContext, polarisCallContext)) { - CallContext.setCurrentContext(callCtx); - FileIO fileIO = new InMemoryFileIO(); - TableIdentifier tableIdentifier = - TableIdentifier.of(Namespace.of("db1", "schema1"), "table1"); - ManifestFileCleanupTaskHandler handler = - new ManifestFileCleanupTaskHandler((task) -> fileIO, Executors.newSingleThreadExecutor()); - long snapshotId = 100L; - ManifestFile manifestFile = - TaskTestUtils.manifestFile( - fileIO, "manifest1.avro", snapshotId, "dataFile1.parquet", "dataFile2.parquet"); - TestSnapshot snapshot = - TaskTestUtils.newSnapshot(fileIO, "manifestList.avro", 1, snapshotId, 99L, manifestFile); - String metadataFile = "v1-49494949.metadata.json"; - StatisticsFile statisticsFile = - TaskTestUtils.writeStatsFile( - snapshot.snapshotId(), - snapshot.sequenceNumber(), - "/metadata/" + UUID.randomUUID() + ".stats", - fileIO); - TaskTestUtils.writeTableMetadata(fileIO, metadataFile, List.of(statisticsFile), snapshot); + FileIO fileIO = new InMemoryFileIO(); + TableIdentifier tableIdentifier = TableIdentifier.of(Namespace.of("db1", "schema1"), "table1"); + ManifestFileCleanupTaskHandler handler = + new ManifestFileCleanupTaskHandler( + (task, rc) -> fileIO, Executors.newSingleThreadExecutor(), diagnostics); + long snapshotId = 100L; + ManifestFile manifestFile = + TaskTestUtils.manifestFile( + fileIO, "manifest1.avro", snapshotId, "dataFile1.parquet", "dataFile2.parquet"); + TestSnapshot snapshot = + TaskTestUtils.newSnapshot(fileIO, "manifestList.avro", 1, snapshotId, 99L, manifestFile); + String metadataFile = "v1-49494949.metadata.json"; + StatisticsFile statisticsFile = + TaskTestUtils.writeStatsFile( + snapshot.snapshotId(), + snapshot.sequenceNumber(), + "/metadata/" + UUID.randomUUID() + ".stats", + fileIO); + TaskTestUtils.writeTableMetadata(fileIO, metadataFile, List.of(statisticsFile), snapshot); - fileIO.deleteFile(statisticsFile.path()); - assertThat(TaskUtils.exists(statisticsFile.path(), fileIO)).isFalse(); + fileIO.deleteFile(statisticsFile.path()); + assertThat(TaskUtils.exists(statisticsFile.path(), fileIO)).isFalse(); - TaskEntity task = - new TaskEntity.Builder() - .withTaskType(AsyncTaskType.METADATA_FILE_BATCH_CLEANUP) - .withData( - new ManifestFileCleanupTaskHandler.ManifestCleanupTask( - tableIdentifier, List.of(statisticsFile.path()))) - .setName(UUID.randomUUID().toString()) - .build(); - assertThatPredicate(handler::canHandleTask).accepts(task); - assertThatPredicate(handler::handleTask).accepts(task); - } + TaskEntity task = + new TaskEntity.Builder() + .withTaskType(diagnostics, AsyncTaskType.METADATA_FILE_BATCH_CLEANUP) + .withData( + diagnostics, + new ManifestFileCleanupTaskHandler.ManifestCleanupTask( + tableIdentifier, List.of(statisticsFile.path()))) + .setName(UUID.randomUUID().toString()) + .build(); + assertThat(handler.canHandleTask(task)).isTrue(); + assertThat(handler.handleTask(task, realmContext)).isTrue(); } @Test public void testCleanupWithRetries() throws IOException { - PolarisCallContext polarisCallContext = - new PolarisCallContext( - metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(), - new PolarisDefaultDiagServiceImpl()); - try (CallContext callCtx = CallContext.of(realmContext, polarisCallContext)) { - CallContext.setCurrentContext(callCtx); - Map retryCounter = new HashMap<>(); - FileIO fileIO = - new InMemoryFileIO() { - @Override - public void close() { - // no-op - } + Map retryCounter = new HashMap<>(); + FileIO fileIO = + new InMemoryFileIO() { + @Override + public void close() { + // no-op + } - @Override - public void deleteFile(String location) { - int attempts = - retryCounter - .computeIfAbsent(location, k -> new AtomicInteger(0)) - .incrementAndGet(); - if (attempts < 3) { - throw new RuntimeException("Simulating failure to test retries"); - } else { - super.deleteFile(location); - } + @Override + public void deleteFile(String location) { + int attempts = + retryCounter.computeIfAbsent(location, k -> new AtomicInteger(0)).incrementAndGet(); + if (attempts < 3) { + throw new RuntimeException("Simulating failure to test retries"); + } else { + super.deleteFile(location); } - }; - TableIdentifier tableIdentifier = - TableIdentifier.of(Namespace.of("db1", "schema1"), "table1"); - ManifestFileCleanupTaskHandler handler = - new ManifestFileCleanupTaskHandler((task) -> fileIO, Executors.newSingleThreadExecutor()); - long snapshotId = 100L; - ManifestFile manifestFile = - TaskTestUtils.manifestFile( - fileIO, "manifest1.avro", snapshotId, "dataFile1.parquet", "dataFile2.parquet"); - TestSnapshot snapshot = - TaskTestUtils.newSnapshot(fileIO, "manifestList.avro", 1, snapshotId, 99L, manifestFile); - String metadataFile = "v1-49494949.metadata.json"; - StatisticsFile statisticsFile = - TaskTestUtils.writeStatsFile( - snapshot.snapshotId(), - snapshot.sequenceNumber(), - "/metadata/" + UUID.randomUUID() + ".stats", - fileIO); - TaskTestUtils.writeTableMetadata(fileIO, metadataFile, List.of(statisticsFile), snapshot); - assertThat(TaskUtils.exists(statisticsFile.path(), fileIO)).isTrue(); + } + }; + TableIdentifier tableIdentifier = TableIdentifier.of(Namespace.of("db1", "schema1"), "table1"); + ManifestFileCleanupTaskHandler handler = + new ManifestFileCleanupTaskHandler( + (task, rc) -> fileIO, Executors.newSingleThreadExecutor(), diagnostics); + long snapshotId = 100L; + ManifestFile manifestFile = + TaskTestUtils.manifestFile( + fileIO, "manifest1.avro", snapshotId, "dataFile1.parquet", "dataFile2.parquet"); + TestSnapshot snapshot = + TaskTestUtils.newSnapshot(fileIO, "manifestList.avro", 1, snapshotId, 99L, manifestFile); + String metadataFile = "v1-49494949.metadata.json"; + StatisticsFile statisticsFile = + TaskTestUtils.writeStatsFile( + snapshot.snapshotId(), + snapshot.sequenceNumber(), + "/metadata/" + UUID.randomUUID() + ".stats", + fileIO); + TaskTestUtils.writeTableMetadata(fileIO, metadataFile, List.of(statisticsFile), snapshot); + assertThat(TaskUtils.exists(statisticsFile.path(), fileIO)).isTrue(); - TaskEntity task = - new TaskEntity.Builder() - .withTaskType(AsyncTaskType.METADATA_FILE_BATCH_CLEANUP) - .withData( - new ManifestFileCleanupTaskHandler.ManifestCleanupTask( - tableIdentifier, List.of(statisticsFile.path()))) - .setName(UUID.randomUUID().toString()) - .build(); + TaskEntity task = + new TaskEntity.Builder() + .withTaskType(diagnostics, AsyncTaskType.METADATA_FILE_BATCH_CLEANUP) + .withData( + diagnostics, + new ManifestFileCleanupTaskHandler.ManifestCleanupTask( + tableIdentifier, List.of(statisticsFile.path()))) + .setName(UUID.randomUUID().toString()) + .build(); - try (ExecutorService executor = Executors.newSingleThreadExecutor()) { - CompletableFuture future; - future = - CompletableFuture.runAsync( - () -> { - assertThatPredicate(handler::canHandleTask).accepts(task); - handler.handleTask(task); // this will schedule the batch deletion - }, - executor); - // Wait for all async tasks to finish - future.join(); - } + try (ExecutorService executor = Executors.newSingleThreadExecutor()) { + CompletableFuture future; + future = + CompletableFuture.runAsync( + () -> { + assertThat(handler.canHandleTask(task)).isTrue(); + handler.handleTask(task, realmContext); // this will schedule the batch deletion + }, + executor); + // Wait for all async tasks to finish + future.join(); + } - // Check if the file was successfully deleted after retries - assertThat(TaskUtils.exists(statisticsFile.path(), fileIO)).isFalse(); + // Check if the file was successfully deleted after retries + assertThat(TaskUtils.exists(statisticsFile.path(), fileIO)).isFalse(); - // Ensure that retries happened as expected - assertThat(retryCounter.containsKey(statisticsFile.path())).isTrue(); - assertThat(retryCounter.get(statisticsFile.path()).get()).isEqualTo(3); - } + // Ensure that retries happened as expected + assertThat(retryCounter.containsKey(statisticsFile.path())).isTrue(); + assertThat(retryCounter.get(statisticsFile.path()).get()).isEqualTo(3); } } diff --git a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/task/TableCleanupTaskHandlerTest.java b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/task/TableCleanupTaskHandlerTest.java index 42efb95d2..fecb7ca8e 100644 --- a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/task/TableCleanupTaskHandlerTest.java +++ b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/task/TableCleanupTaskHandlerTest.java @@ -23,6 +23,7 @@ import io.quarkus.test.junit.QuarkusTest; import jakarta.inject.Inject; import java.io.IOException; +import java.time.Clock; import java.util.List; import java.util.UUID; import org.apache.commons.codec.binary.Base64; @@ -35,9 +36,8 @@ import org.apache.iceberg.catalog.TableIdentifier; import org.apache.iceberg.inmemory.InMemoryFileIO; import org.apache.iceberg.io.FileIO; -import org.apache.polaris.core.PolarisCallContext; -import org.apache.polaris.core.PolarisDefaultDiagServiceImpl; -import org.apache.polaris.core.context.CallContext; +import org.apache.polaris.core.PolarisConfigurationStore; +import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.AsyncTaskType; import org.apache.polaris.core.entity.PolarisBaseEntity; @@ -45,6 +45,7 @@ import org.apache.polaris.core.entity.TableLikeEntity; import org.apache.polaris.core.entity.TaskEntity; import org.apache.polaris.core.persistence.MetaStoreManagerFactory; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; import org.apache.polaris.service.task.ManifestFileCleanupTaskHandler; import org.apache.polaris.service.task.TableCleanupTaskHandler; import org.apache.polaris.service.task.TaskUtils; @@ -56,541 +57,557 @@ @QuarkusTest class TableCleanupTaskHandlerTest { @Inject MetaStoreManagerFactory metaStoreManagerFactory; + @Inject PolarisConfigurationStore configurationStore; + @Inject PolarisDiagnostics diagnostics; private final RealmContext realmContext = () -> "realmName"; @Test public void testTableCleanup() throws IOException { - PolarisCallContext polarisCallContext = - new PolarisCallContext( - metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(), - new PolarisDefaultDiagServiceImpl()); - try (CallContext callCtx = CallContext.of(realmContext, polarisCallContext)) { - CallContext.setCurrentContext(callCtx); - FileIO fileIO = new InMemoryFileIO(); - TableIdentifier tableIdentifier = - TableIdentifier.of(Namespace.of("db1", "schema1"), "table1"); - TableCleanupTaskHandler handler = - new TableCleanupTaskHandler(Mockito.mock(), metaStoreManagerFactory, (task) -> fileIO); - long snapshotId = 100L; - ManifestFile manifestFile = - TaskTestUtils.manifestFile( - fileIO, "manifest1.avro", snapshotId, "dataFile1.parquet", "dataFile2.parquet"); - TestSnapshot snapshot = - TaskTestUtils.newSnapshot(fileIO, "manifestList.avro", 1, snapshotId, 99L, manifestFile); - String metadataFile = "v1-49494949.metadata.json"; - StatisticsFile statisticsFile = - TaskTestUtils.writeStatsFile( - snapshot.snapshotId(), - snapshot.sequenceNumber(), - "/metadata/" + UUID.randomUUID() + ".stats", - fileIO); - TaskTestUtils.writeTableMetadata(fileIO, metadataFile, List.of(statisticsFile), snapshot); - - TaskEntity task = - new TaskEntity.Builder() - .setName("cleanup_" + tableIdentifier.toString()) - .withTaskType(AsyncTaskType.ENTITY_CLEANUP_SCHEDULER) - .withData( - new TableLikeEntity.Builder(tableIdentifier, metadataFile) - .setName("table1") - .setCatalogId(1) - .setCreateTimestamp(100) - .build()) - .build(); - Assertions.assertThatPredicate(handler::canHandleTask).accepts(task); - - CallContext.setCurrentContext(CallContext.of(realmContext, polarisCallContext)); - handler.handleTask(task); - - assertThat( - metaStoreManagerFactory - .getOrCreateMetaStoreManager(realmContext) - .loadTasks(polarisCallContext, "test", 2) - .getEntities()) - .hasSize(2) - .satisfiesExactlyInAnyOrder( - taskEntity -> - assertThat(taskEntity) - .returns(PolarisEntityType.TASK.getCode(), PolarisBaseEntity::getTypeCode) - .extracting(TaskEntity::of) - .returns(AsyncTaskType.MANIFEST_FILE_CLEANUP, TaskEntity::getTaskType) - .returns( - new ManifestFileCleanupTaskHandler.ManifestCleanupTask( - tableIdentifier, - Base64.encodeBase64String(ManifestFiles.encode(manifestFile))), - entity -> - entity.readData( - ManifestFileCleanupTaskHandler.ManifestCleanupTask.class)), - taskEntity -> - assertThat(taskEntity) - .returns(PolarisEntityType.TASK.getCode(), PolarisBaseEntity::getTypeCode) - .extracting(TaskEntity::of) - .returns(AsyncTaskType.METADATA_FILE_BATCH_CLEANUP, TaskEntity::getTaskType) - .returns( - new ManifestFileCleanupTaskHandler.ManifestCleanupTask( - tableIdentifier, List.of(statisticsFile.path())), - entity -> - entity.readData( - ManifestFileCleanupTaskHandler.ManifestCleanupTask.class))); - } + PolarisMetaStoreSession metaStoreSession = + metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(); + FileIO fileIO = new InMemoryFileIO(); + TableIdentifier tableIdentifier = TableIdentifier.of(Namespace.of("db1", "schema1"), "table1"); + TableCleanupTaskHandler handler = + new TableCleanupTaskHandler( + Mockito.mock(), + metaStoreManagerFactory, + configurationStore, + diagnostics, + (task, rc) -> fileIO, + Clock.systemUTC()); + long snapshotId = 100L; + ManifestFile manifestFile = + TaskTestUtils.manifestFile( + fileIO, "manifest1.avro", snapshotId, "dataFile1.parquet", "dataFile2.parquet"); + TestSnapshot snapshot = + TaskTestUtils.newSnapshot(fileIO, "manifestList.avro", 1, snapshotId, 99L, manifestFile); + String metadataFile = "v1-49494949.metadata.json"; + StatisticsFile statisticsFile = + TaskTestUtils.writeStatsFile( + snapshot.snapshotId(), + snapshot.sequenceNumber(), + "/metadata/" + UUID.randomUUID() + ".stats", + fileIO); + TaskTestUtils.writeTableMetadata(fileIO, metadataFile, List.of(statisticsFile), snapshot); + + TaskEntity task = + new TaskEntity.Builder() + .setName("cleanup_" + tableIdentifier) + .withTaskType(diagnostics, AsyncTaskType.ENTITY_CLEANUP_SCHEDULER) + .withData( + diagnostics, + new TableLikeEntity.Builder(tableIdentifier, metadataFile) + .setName("table1") + .setCatalogId(1) + .setCreateTimestamp(100) + .build()) + .build(); + Assertions.assertThatPredicate(handler::canHandleTask).accepts(task); + + handler.handleTask(task, realmContext); + + assertThat( + metaStoreManagerFactory + .getOrCreateMetaStoreManager(realmContext) + .loadTasks(metaStoreSession, "test", 2) + .getEntities()) + .hasSize(2) + .satisfiesExactlyInAnyOrder( + taskEntity -> + assertThat(taskEntity) + .returns(PolarisEntityType.TASK.getCode(), PolarisBaseEntity::getTypeCode) + .extracting(TaskEntity::of) + .returns( + AsyncTaskType.MANIFEST_FILE_CLEANUP, + taskEntity1 -> taskEntity1.getTaskType(diagnostics)) + .returns( + new ManifestFileCleanupTaskHandler.ManifestCleanupTask( + tableIdentifier, + Base64.encodeBase64String(ManifestFiles.encode(manifestFile))), + entity -> + entity.readData( + diagnostics, + ManifestFileCleanupTaskHandler.ManifestCleanupTask.class)), + taskEntity -> + assertThat(taskEntity) + .returns(PolarisEntityType.TASK.getCode(), PolarisBaseEntity::getTypeCode) + .extracting(TaskEntity::of) + .returns( + AsyncTaskType.METADATA_FILE_BATCH_CLEANUP, + taskEntity2 -> taskEntity2.getTaskType(diagnostics)) + .returns( + new ManifestFileCleanupTaskHandler.ManifestCleanupTask( + tableIdentifier, List.of(statisticsFile.path())), + entity -> + entity.readData( + diagnostics, + ManifestFileCleanupTaskHandler.ManifestCleanupTask.class))); } @Test public void testTableCleanupHandlesAlreadyDeletedMetadata() throws IOException { - PolarisCallContext polarisCallContext = - new PolarisCallContext( - metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(), - new PolarisDefaultDiagServiceImpl()); - try (CallContext callCtx = CallContext.of(realmContext, polarisCallContext)) { - CallContext.setCurrentContext(callCtx); - FileIO fileIO = - new InMemoryFileIO() { - @Override - public void close() { - // no-op - } - }; - TableIdentifier tableIdentifier = - TableIdentifier.of(Namespace.of("db1", "schema1"), "table1"); - TableCleanupTaskHandler handler = - new TableCleanupTaskHandler(Mockito.mock(), metaStoreManagerFactory, (task) -> fileIO); - long snapshotId = 100L; - ManifestFile manifestFile = - TaskTestUtils.manifestFile( - fileIO, "manifest1.avro", snapshotId, "dataFile1.parquet", "dataFile2.parquet"); - TestSnapshot snapshot = - TaskTestUtils.newSnapshot(fileIO, "manifestList.avro", 1, snapshotId, 99L, manifestFile); - String metadataFile = "v1-49494949.metadata.json"; - TaskTestUtils.writeTableMetadata(fileIO, metadataFile, snapshot); - - TableLikeEntity tableLikeEntity = - new TableLikeEntity.Builder(tableIdentifier, metadataFile) - .setName("table1") - .setCatalogId(1) - .setCreateTimestamp(100) - .build(); - TaskEntity task = - new TaskEntity.Builder() - .setName("cleanup_" + tableIdentifier.toString()) - .withTaskType(AsyncTaskType.ENTITY_CLEANUP_SCHEDULER) - .withData(tableLikeEntity) - .build(); - Assertions.assertThatPredicate(handler::canHandleTask).accepts(task); - - CallContext.setCurrentContext(CallContext.of(realmContext, polarisCallContext)); - - // handle the same task twice - // the first one should successfully delete the metadata - List results = List.of(handler.handleTask(task), handler.handleTask(task)); - assertThat(results).containsExactly(true, true); - - // both tasks successfully executed, but only one should queue subtasks - assertThat( - metaStoreManagerFactory - .getOrCreateMetaStoreManager(realmContext) - .loadTasks(polarisCallContext, "test", 5) - .getEntities()) - .hasSize(1); - } + PolarisMetaStoreSession metaStoreSession = + metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(); + FileIO fileIO = + new InMemoryFileIO() { + @Override + public void close() { + // no-op + } + }; + TableIdentifier tableIdentifier = TableIdentifier.of(Namespace.of("db1", "schema1"), "table1"); + TableCleanupTaskHandler handler = + new TableCleanupTaskHandler( + Mockito.mock(), + metaStoreManagerFactory, + configurationStore, + diagnostics, + (task, rc) -> fileIO, + Clock.systemUTC()); + long snapshotId = 100L; + ManifestFile manifestFile = + TaskTestUtils.manifestFile( + fileIO, "manifest1.avro", snapshotId, "dataFile1.parquet", "dataFile2.parquet"); + TestSnapshot snapshot = + TaskTestUtils.newSnapshot(fileIO, "manifestList.avro", 1, snapshotId, 99L, manifestFile); + String metadataFile = "v1-49494949.metadata.json"; + TaskTestUtils.writeTableMetadata(fileIO, metadataFile, snapshot); + + TableLikeEntity tableLikeEntity = + new TableLikeEntity.Builder(tableIdentifier, metadataFile) + .setName("table1") + .setCatalogId(1) + .setCreateTimestamp(100) + .build(); + TaskEntity task = + new TaskEntity.Builder() + .setName("cleanup_" + tableIdentifier) + .withTaskType(diagnostics, AsyncTaskType.ENTITY_CLEANUP_SCHEDULER) + .withData(diagnostics, tableLikeEntity) + .build(); + Assertions.assertThatPredicate(handler::canHandleTask).accepts(task); + + // handle the same task twice + // the first one should successfully delete the metadata + List results = + List.of(handler.handleTask(task, realmContext), handler.handleTask(task, realmContext)); + assertThat(results).containsExactly(true, true); + + // both tasks successfully executed, but only one should queue subtasks + assertThat( + metaStoreManagerFactory + .getOrCreateMetaStoreManager(realmContext) + .loadTasks(metaStoreSession, "test", 5) + .getEntities()) + .hasSize(1); } @Test public void testTableCleanupDuplicatesTasksIfFileStillExists() throws IOException { - PolarisCallContext polarisCallContext = - new PolarisCallContext( - metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(), - new PolarisDefaultDiagServiceImpl()); - try (CallContext callCtx = CallContext.of(realmContext, polarisCallContext)) { - CallContext.setCurrentContext(callCtx); - FileIO fileIO = - new InMemoryFileIO() { - @Override - public void deleteFile(String location) { - LoggerFactory.getLogger(TableCleanupTaskHandler.class) - .info( - "Not deleting file at location {} to simulate concurrent tasks runs", - location); - // don't do anything - } - - @Override - public void close() { - // no-op - } - }; - TableIdentifier tableIdentifier = - TableIdentifier.of(Namespace.of("db1", "schema1"), "table1"); - TableCleanupTaskHandler handler = - new TableCleanupTaskHandler(Mockito.mock(), metaStoreManagerFactory, (task) -> fileIO); - long snapshotId = 100L; - ManifestFile manifestFile = - TaskTestUtils.manifestFile( - fileIO, "manifest1.avro", snapshotId, "dataFile1.parquet", "dataFile2.parquet"); - TestSnapshot snapshot = - TaskTestUtils.newSnapshot(fileIO, "manifestList.avro", 1, snapshotId, 99L, manifestFile); - String metadataFile = "v1-49494949.metadata.json"; - TaskTestUtils.writeTableMetadata(fileIO, metadataFile, snapshot); - - TaskEntity task = - new TaskEntity.Builder() - .setName("cleanup_" + tableIdentifier.toString()) - .withTaskType(AsyncTaskType.ENTITY_CLEANUP_SCHEDULER) - .withData( - new TableLikeEntity.Builder(tableIdentifier, metadataFile) - .setName("table1") - .setCatalogId(1) - .setCreateTimestamp(100) - .build()) - .build(); - Assertions.assertThatPredicate(handler::canHandleTask).accepts(task); - - CallContext.setCurrentContext(CallContext.of(realmContext, polarisCallContext)); - - // handle the same task twice - // the first one should successfully delete the metadata - List results = List.of(handler.handleTask(task), handler.handleTask(task)); - assertThat(results).containsExactly(true, true); - - // both tasks successfully executed, but only one should queue subtasks - assertThat( - metaStoreManagerFactory - .getOrCreateMetaStoreManager(realmContext) - .loadTasks(polarisCallContext, "test", 5) - .getEntities()) - .hasSize(2) - .satisfiesExactly( - taskEntity -> - assertThat(taskEntity) - .returns(PolarisEntityType.TASK.getCode(), PolarisBaseEntity::getTypeCode) - .extracting(TaskEntity::of) - .returns(AsyncTaskType.MANIFEST_FILE_CLEANUP, TaskEntity::getTaskType) - .returns( - new ManifestFileCleanupTaskHandler.ManifestCleanupTask( - tableIdentifier, - Base64.encodeBase64String(ManifestFiles.encode(manifestFile))), - entity -> - entity.readData( - ManifestFileCleanupTaskHandler.ManifestCleanupTask.class)), - taskEntity -> - assertThat(taskEntity) - .returns(PolarisEntityType.TASK.getCode(), PolarisBaseEntity::getTypeCode) - .extracting(TaskEntity::of) - .returns(AsyncTaskType.MANIFEST_FILE_CLEANUP, TaskEntity::getTaskType) - .returns( - new ManifestFileCleanupTaskHandler.ManifestCleanupTask( - tableIdentifier, - Base64.encodeBase64String(ManifestFiles.encode(manifestFile))), - entity -> - entity.readData( - ManifestFileCleanupTaskHandler.ManifestCleanupTask.class))); - } + PolarisMetaStoreSession metaStoreSession = + metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(); + FileIO fileIO = + new InMemoryFileIO() { + @Override + public void deleteFile(String location) { + LoggerFactory.getLogger(TableCleanupTaskHandler.class) + .info( + "Not deleting file at location {} to simulate concurrent tasks runs", location); + // don't do anything + } + + @Override + public void close() { + // no-op + } + }; + TableIdentifier tableIdentifier = TableIdentifier.of(Namespace.of("db1", "schema1"), "table1"); + TableCleanupTaskHandler handler = + new TableCleanupTaskHandler( + Mockito.mock(), + metaStoreManagerFactory, + configurationStore, + diagnostics, + (task, rc) -> fileIO, + Clock.systemUTC()); + long snapshotId = 100L; + ManifestFile manifestFile = + TaskTestUtils.manifestFile( + fileIO, "manifest1.avro", snapshotId, "dataFile1.parquet", "dataFile2.parquet"); + TestSnapshot snapshot = + TaskTestUtils.newSnapshot(fileIO, "manifestList.avro", 1, snapshotId, 99L, manifestFile); + String metadataFile = "v1-49494949.metadata.json"; + TaskTestUtils.writeTableMetadata(fileIO, metadataFile, snapshot); + + TaskEntity task = + new TaskEntity.Builder() + .setName("cleanup_" + tableIdentifier) + .withTaskType(diagnostics, AsyncTaskType.ENTITY_CLEANUP_SCHEDULER) + .withData( + diagnostics, + new TableLikeEntity.Builder(tableIdentifier, metadataFile) + .setName("table1") + .setCatalogId(1) + .setCreateTimestamp(100) + .build()) + .build(); + Assertions.assertThatPredicate(handler::canHandleTask).accepts(task); + + // handle the same task twice + // the first one should successfully delete the metadata + List results = + List.of(handler.handleTask(task, realmContext), handler.handleTask(task, realmContext)); + assertThat(results).containsExactly(true, true); + + // both tasks successfully executed, but only one should queue subtasks + assertThat( + metaStoreManagerFactory + .getOrCreateMetaStoreManager(realmContext) + .loadTasks(metaStoreSession, "test", 5) + .getEntities()) + .hasSize(2) + .satisfiesExactly( + taskEntity -> + assertThat(taskEntity) + .returns(PolarisEntityType.TASK.getCode(), PolarisBaseEntity::getTypeCode) + .extracting(TaskEntity::of) + .returns( + AsyncTaskType.MANIFEST_FILE_CLEANUP, + taskEntity1 -> taskEntity1.getTaskType(diagnostics)) + .returns( + new ManifestFileCleanupTaskHandler.ManifestCleanupTask( + tableIdentifier, + Base64.encodeBase64String(ManifestFiles.encode(manifestFile))), + entity -> + entity.readData( + diagnostics, + ManifestFileCleanupTaskHandler.ManifestCleanupTask.class)), + taskEntity -> + assertThat(taskEntity) + .returns(PolarisEntityType.TASK.getCode(), PolarisBaseEntity::getTypeCode) + .extracting(TaskEntity::of) + .returns( + AsyncTaskType.MANIFEST_FILE_CLEANUP, + taskEntity2 -> taskEntity2.getTaskType(diagnostics)) + .returns( + new ManifestFileCleanupTaskHandler.ManifestCleanupTask( + tableIdentifier, + Base64.encodeBase64String(ManifestFiles.encode(manifestFile))), + entity -> + entity.readData( + diagnostics, + ManifestFileCleanupTaskHandler.ManifestCleanupTask.class))); } @Test public void testTableCleanupMultipleSnapshots() throws IOException { - PolarisCallContext polarisCallContext = - new PolarisCallContext( - metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(), - new PolarisDefaultDiagServiceImpl()); - try (CallContext callCtx = CallContext.of(realmContext, polarisCallContext)) { - CallContext.setCurrentContext(callCtx); - FileIO fileIO = new InMemoryFileIO(); - TableIdentifier tableIdentifier = - TableIdentifier.of(Namespace.of("db1", "schema1"), "table1"); - TableCleanupTaskHandler handler = - new TableCleanupTaskHandler(Mockito.mock(), metaStoreManagerFactory, (task) -> fileIO); - long snapshotId1 = 100L; - ManifestFile manifestFile1 = - TaskTestUtils.manifestFile( - fileIO, "manifest1.avro", snapshotId1, "dataFile1.parquet", "dataFile2.parquet"); - ManifestFile manifestFile2 = - TaskTestUtils.manifestFile( - fileIO, "manifest2.avro", snapshotId1, "dataFile3.parquet", "dataFile4.parquet"); - Snapshot snapshot = - TaskTestUtils.newSnapshot( - fileIO, "manifestList.avro", 1, snapshotId1, 99L, manifestFile1, manifestFile2); - ManifestFile manifestFile3 = - TaskTestUtils.manifestFile( - fileIO, "manifest3.avro", snapshot.snapshotId() + 1, "dataFile5.parquet"); - Snapshot snapshot2 = - TaskTestUtils.newSnapshot( - fileIO, - "manifestList2.avro", - snapshot.sequenceNumber() + 1, - snapshot.snapshotId() + 1, - snapshot.snapshotId(), - manifestFile1, - manifestFile3); // exclude manifest2 from the new snapshot - String metadataFile = "v1-295495059.metadata.json"; - StatisticsFile statisticsFile1 = - TaskTestUtils.writeStatsFile( - snapshot.snapshotId(), - snapshot.sequenceNumber(), - "/metadata/" + UUID.randomUUID() + ".stats", - fileIO); - StatisticsFile statisticsFile2 = - TaskTestUtils.writeStatsFile( - snapshot2.snapshotId(), - snapshot2.sequenceNumber(), - "/metadata/" + UUID.randomUUID() + ".stats", - fileIO); - TaskTestUtils.writeTableMetadata( - fileIO, metadataFile, List.of(statisticsFile1, statisticsFile2), snapshot, snapshot2); - - TaskEntity task = - new TaskEntity.Builder() - .withTaskType(AsyncTaskType.ENTITY_CLEANUP_SCHEDULER) - .withData( - new TableLikeEntity.Builder(tableIdentifier, metadataFile) - .setName("table1") - .setCatalogId(1) - .setCreateTimestamp(100) - .build()) - .build(); - Assertions.assertThatPredicate(handler::canHandleTask).accepts(task); - - CallContext.setCurrentContext(CallContext.of(realmContext, polarisCallContext)); - - handler.handleTask(task); - - List entities = - metaStoreManagerFactory - .getOrCreateMetaStoreManager(realmContext) - .loadTasks(polarisCallContext, "test", 5) - .getEntities(); - - List manifestCleanupTasks = - entities.stream() - .filter( - entity -> { - AsyncTaskType taskType = TaskEntity.of(entity).getTaskType(); - return taskType == AsyncTaskType.MANIFEST_FILE_CLEANUP; - }) - .toList(); - List metadataCleanupTasks = - entities.stream() - .filter( - entity -> { - AsyncTaskType taskType = TaskEntity.of(entity).getTaskType(); - return taskType == AsyncTaskType.METADATA_FILE_BATCH_CLEANUP; - }) - .toList(); - - assertThat(metadataCleanupTasks) - .hasSize(1) - .satisfiesExactlyInAnyOrder( - taskEntity -> - assertThat(taskEntity) - .returns(PolarisEntityType.TASK.getCode(), PolarisBaseEntity::getTypeCode) - .extracting(TaskEntity::of) - .returns( - new ManifestFileCleanupTaskHandler.ManifestCleanupTask( - tableIdentifier, - List.of(statisticsFile1.path(), statisticsFile2.path())), - entity -> - entity.readData( - ManifestFileCleanupTaskHandler.ManifestCleanupTask.class))); - - assertThat(manifestCleanupTasks) - // all three manifests should be present, even though one is excluded from the latest - // snapshot - .hasSize(3) - .satisfiesExactlyInAnyOrder( - taskEntity -> - assertThat(taskEntity) - .returns(PolarisEntityType.TASK.getCode(), PolarisBaseEntity::getTypeCode) - .extracting(TaskEntity::of) - .returns( - new ManifestFileCleanupTaskHandler.ManifestCleanupTask( - tableIdentifier, - Base64.encodeBase64String(ManifestFiles.encode(manifestFile1))), - entity -> - entity.readData( - ManifestFileCleanupTaskHandler.ManifestCleanupTask.class)), - taskEntity -> - assertThat(taskEntity) - .returns(PolarisEntityType.TASK.getCode(), PolarisBaseEntity::getTypeCode) - .extracting(TaskEntity::of) - .returns( - new ManifestFileCleanupTaskHandler.ManifestCleanupTask( - tableIdentifier, - Base64.encodeBase64String(ManifestFiles.encode(manifestFile2))), - entity -> - entity.readData( - ManifestFileCleanupTaskHandler.ManifestCleanupTask.class)), - taskEntity -> - assertThat(taskEntity) - .returns(PolarisEntityType.TASK.getCode(), PolarisBaseEntity::getTypeCode) - .extracting(TaskEntity::of) - .returns( - new ManifestFileCleanupTaskHandler.ManifestCleanupTask( - tableIdentifier, - Base64.encodeBase64String(ManifestFiles.encode(manifestFile3))), - entity -> - entity.readData( - ManifestFileCleanupTaskHandler.ManifestCleanupTask.class))); - } + PolarisMetaStoreSession metaStoreSession = + metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(); + FileIO fileIO = new InMemoryFileIO(); + TableIdentifier tableIdentifier = TableIdentifier.of(Namespace.of("db1", "schema1"), "table1"); + TableCleanupTaskHandler handler = + new TableCleanupTaskHandler( + Mockito.mock(), + metaStoreManagerFactory, + configurationStore, + diagnostics, + (task, rc) -> fileIO, + Clock.systemUTC()); + long snapshotId1 = 100L; + ManifestFile manifestFile1 = + TaskTestUtils.manifestFile( + fileIO, "manifest1.avro", snapshotId1, "dataFile1.parquet", "dataFile2.parquet"); + ManifestFile manifestFile2 = + TaskTestUtils.manifestFile( + fileIO, "manifest2.avro", snapshotId1, "dataFile3.parquet", "dataFile4.parquet"); + Snapshot snapshot = + TaskTestUtils.newSnapshot( + fileIO, "manifestList.avro", 1, snapshotId1, 99L, manifestFile1, manifestFile2); + ManifestFile manifestFile3 = + TaskTestUtils.manifestFile( + fileIO, "manifest3.avro", snapshot.snapshotId() + 1, "dataFile5.parquet"); + Snapshot snapshot2 = + TaskTestUtils.newSnapshot( + fileIO, + "manifestList2.avro", + snapshot.sequenceNumber() + 1, + snapshot.snapshotId() + 1, + snapshot.snapshotId(), + manifestFile1, + manifestFile3); // exclude manifest2 from the new snapshot + String metadataFile = "v1-295495059.metadata.json"; + StatisticsFile statisticsFile1 = + TaskTestUtils.writeStatsFile( + snapshot.snapshotId(), + snapshot.sequenceNumber(), + "/metadata/" + UUID.randomUUID() + ".stats", + fileIO); + StatisticsFile statisticsFile2 = + TaskTestUtils.writeStatsFile( + snapshot2.snapshotId(), + snapshot2.sequenceNumber(), + "/metadata/" + UUID.randomUUID() + ".stats", + fileIO); + TaskTestUtils.writeTableMetadata( + fileIO, metadataFile, List.of(statisticsFile1, statisticsFile2), snapshot, snapshot2); + + TaskEntity task = + new TaskEntity.Builder() + .withTaskType(diagnostics, AsyncTaskType.ENTITY_CLEANUP_SCHEDULER) + .withData( + diagnostics, + new TableLikeEntity.Builder(tableIdentifier, metadataFile) + .setName("table1") + .setCatalogId(1) + .setCreateTimestamp(100) + .build()) + .build(); + Assertions.assertThatPredicate(handler::canHandleTask).accepts(task); + + handler.handleTask(task, realmContext); + + List entities = + metaStoreManagerFactory + .getOrCreateMetaStoreManager(realmContext) + .loadTasks(metaStoreSession, "test", 5) + .getEntities(); + + List manifestCleanupTasks = + entities.stream() + .filter( + entity -> { + AsyncTaskType taskType = TaskEntity.of(entity).getTaskType(diagnostics); + return taskType == AsyncTaskType.MANIFEST_FILE_CLEANUP; + }) + .toList(); + List metadataCleanupTasks = + entities.stream() + .filter( + entity -> { + AsyncTaskType taskType = TaskEntity.of(entity).getTaskType(diagnostics); + return taskType == AsyncTaskType.METADATA_FILE_BATCH_CLEANUP; + }) + .toList(); + + assertThat(metadataCleanupTasks) + .hasSize(1) + .satisfiesExactlyInAnyOrder( + taskEntity -> + assertThat(taskEntity) + .returns(PolarisEntityType.TASK.getCode(), PolarisBaseEntity::getTypeCode) + .extracting(TaskEntity::of) + .returns( + new ManifestFileCleanupTaskHandler.ManifestCleanupTask( + tableIdentifier, + List.of(statisticsFile1.path(), statisticsFile2.path())), + entity -> + entity.readData( + diagnostics, + ManifestFileCleanupTaskHandler.ManifestCleanupTask.class))); + + assertThat(manifestCleanupTasks) + // all three manifests should be present, even though one is excluded from the latest + // snapshot + .hasSize(3) + .satisfiesExactlyInAnyOrder( + taskEntity -> + assertThat(taskEntity) + .returns(PolarisEntityType.TASK.getCode(), PolarisBaseEntity::getTypeCode) + .extracting(TaskEntity::of) + .returns( + new ManifestFileCleanupTaskHandler.ManifestCleanupTask( + tableIdentifier, + Base64.encodeBase64String(ManifestFiles.encode(manifestFile1))), + entity -> + entity.readData( + diagnostics, + ManifestFileCleanupTaskHandler.ManifestCleanupTask.class)), + taskEntity -> + assertThat(taskEntity) + .returns(PolarisEntityType.TASK.getCode(), PolarisBaseEntity::getTypeCode) + .extracting(TaskEntity::of) + .returns( + new ManifestFileCleanupTaskHandler.ManifestCleanupTask( + tableIdentifier, + Base64.encodeBase64String(ManifestFiles.encode(manifestFile2))), + entity -> + entity.readData( + diagnostics, + ManifestFileCleanupTaskHandler.ManifestCleanupTask.class)), + taskEntity -> + assertThat(taskEntity) + .returns(PolarisEntityType.TASK.getCode(), PolarisBaseEntity::getTypeCode) + .extracting(TaskEntity::of) + .returns( + new ManifestFileCleanupTaskHandler.ManifestCleanupTask( + tableIdentifier, + Base64.encodeBase64String(ManifestFiles.encode(manifestFile3))), + entity -> + entity.readData( + diagnostics, + ManifestFileCleanupTaskHandler.ManifestCleanupTask.class))); } @Test public void testTableCleanupMultipleMetadata() throws IOException { - PolarisCallContext polarisCallContext = - new PolarisCallContext( - metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(), - new PolarisDefaultDiagServiceImpl()); - try (CallContext callCtx = CallContext.of(realmContext, polarisCallContext)) { - CallContext.setCurrentContext(callCtx); - FileIO fileIO = new InMemoryFileIO(); - TableIdentifier tableIdentifier = - TableIdentifier.of(Namespace.of("db1", "schema1"), "table1"); - TableCleanupTaskHandler handler = - new TableCleanupTaskHandler(Mockito.mock(), metaStoreManagerFactory, (task) -> fileIO); - long snapshotId1 = 100L; - ManifestFile manifestFile1 = - TaskTestUtils.manifestFile( - fileIO, "manifest1.avro", snapshotId1, "dataFile1.parquet", "dataFile2.parquet"); - ManifestFile manifestFile2 = - TaskTestUtils.manifestFile( - fileIO, "manifest2.avro", snapshotId1, "dataFile3.parquet", "dataFile4.parquet"); - Snapshot snapshot = - TaskTestUtils.newSnapshot( - fileIO, "manifestList.avro", 1, snapshotId1, 99L, manifestFile1, manifestFile2); - StatisticsFile statisticsFile1 = - TaskTestUtils.writeStatsFile( - snapshot.snapshotId(), - snapshot.sequenceNumber(), - "/metadata/" + UUID.randomUUID() + ".stats", - fileIO); - String firstMetadataFile = "v1-295495059.metadata.json"; - TableMetadata firstMetadata = - TaskTestUtils.writeTableMetadata( - fileIO, firstMetadataFile, List.of(statisticsFile1), snapshot); - assertThat(TaskUtils.exists(firstMetadataFile, fileIO)).isTrue(); - - ManifestFile manifestFile3 = - TaskTestUtils.manifestFile( - fileIO, "manifest3.avro", snapshot.snapshotId() + 1, "dataFile5.parquet"); - Snapshot snapshot2 = - TaskTestUtils.newSnapshot( - fileIO, - "manifestList2.avro", - snapshot.sequenceNumber() + 1, - snapshot.snapshotId() + 1, - snapshot.snapshotId(), - manifestFile1, - manifestFile3); // exclude manifest2 from the new snapshot - StatisticsFile statisticsFile2 = - TaskTestUtils.writeStatsFile( - snapshot2.snapshotId(), - snapshot2.sequenceNumber(), - "/metadata/" + UUID.randomUUID() + ".stats", - fileIO); - String secondMetadataFile = "v1-295495060.metadata.json"; - TaskTestUtils.writeTableMetadata( - fileIO, - secondMetadataFile, - firstMetadata, - firstMetadataFile, - List.of(statisticsFile2), - snapshot2); - assertThat(TaskUtils.exists(firstMetadataFile, fileIO)).isTrue(); - assertThat(TaskUtils.exists(secondMetadataFile, fileIO)).isTrue(); - - TaskEntity task = - new TaskEntity.Builder() - .withTaskType(AsyncTaskType.ENTITY_CLEANUP_SCHEDULER) - .withData( - new TableLikeEntity.Builder(tableIdentifier, secondMetadataFile) - .setName("table1") - .setCatalogId(1) - .setCreateTimestamp(100) - .build()) - .build(); - - Assertions.assertThatPredicate(handler::canHandleTask).accepts(task); - - CallContext.setCurrentContext(CallContext.of(realmContext, polarisCallContext)); - - handler.handleTask(task); - - List entities = - metaStoreManagerFactory - .getOrCreateMetaStoreManager(realmContext) - .loadTasks(polarisCallContext, "test", 6) - .getEntities(); - - List manifestCleanupTasks = - entities.stream() - .filter( - entity -> { - AsyncTaskType taskType = TaskEntity.of(entity).getTaskType(); - return taskType == AsyncTaskType.MANIFEST_FILE_CLEANUP; - }) - .toList(); - List metadataCleanupTasks = - entities.stream() - .filter( - entity -> { - AsyncTaskType taskType = TaskEntity.of(entity).getTaskType(); - return taskType == AsyncTaskType.METADATA_FILE_BATCH_CLEANUP; - }) - .toList(); - - assertThat(metadataCleanupTasks) - .hasSize(1) - .satisfiesExactlyInAnyOrder( - taskEntity -> - assertThat(taskEntity) - .returns(PolarisEntityType.TASK.getCode(), PolarisBaseEntity::getTypeCode) - .extracting(TaskEntity::of) - .returns( - new ManifestFileCleanupTaskHandler.ManifestCleanupTask( - tableIdentifier, - List.of( - firstMetadataFile, - statisticsFile1.path(), - statisticsFile2.path())), - entity -> - entity.readData( - ManifestFileCleanupTaskHandler.ManifestCleanupTask.class))); - - assertThat(manifestCleanupTasks) - .hasSize(3) - .satisfiesExactlyInAnyOrder( - taskEntity -> - assertThat(taskEntity) - .returns(PolarisEntityType.TASK.getCode(), PolarisBaseEntity::getTypeCode) - .extracting(TaskEntity::of) - .returns( - new ManifestFileCleanupTaskHandler.ManifestCleanupTask( - tableIdentifier, - Base64.encodeBase64String(ManifestFiles.encode(manifestFile1))), - entity -> - entity.readData( - ManifestFileCleanupTaskHandler.ManifestCleanupTask.class)), - taskEntity -> - assertThat(taskEntity) - .returns(PolarisEntityType.TASK.getCode(), PolarisBaseEntity::getTypeCode) - .extracting(TaskEntity::of) - .returns( - new ManifestFileCleanupTaskHandler.ManifestCleanupTask( - tableIdentifier, - Base64.encodeBase64String(ManifestFiles.encode(manifestFile2))), - entity -> - entity.readData( - ManifestFileCleanupTaskHandler.ManifestCleanupTask.class)), - taskEntity -> - assertThat(taskEntity) - .returns(PolarisEntityType.TASK.getCode(), PolarisBaseEntity::getTypeCode) - .extracting(TaskEntity::of) - .returns( - new ManifestFileCleanupTaskHandler.ManifestCleanupTask( - tableIdentifier, - Base64.encodeBase64String(ManifestFiles.encode(manifestFile3))), - entity -> - entity.readData( - ManifestFileCleanupTaskHandler.ManifestCleanupTask.class))); - } + PolarisMetaStoreSession metaStoreSession = + metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(); + FileIO fileIO = new InMemoryFileIO(); + TableIdentifier tableIdentifier = TableIdentifier.of(Namespace.of("db1", "schema1"), "table1"); + TableCleanupTaskHandler handler = + new TableCleanupTaskHandler( + Mockito.mock(), + metaStoreManagerFactory, + configurationStore, + diagnostics, + (task, rc) -> fileIO, + Clock.systemUTC()); + long snapshotId1 = 100L; + ManifestFile manifestFile1 = + TaskTestUtils.manifestFile( + fileIO, "manifest1.avro", snapshotId1, "dataFile1.parquet", "dataFile2.parquet"); + ManifestFile manifestFile2 = + TaskTestUtils.manifestFile( + fileIO, "manifest2.avro", snapshotId1, "dataFile3.parquet", "dataFile4.parquet"); + Snapshot snapshot = + TaskTestUtils.newSnapshot( + fileIO, "manifestList.avro", 1, snapshotId1, 99L, manifestFile1, manifestFile2); + StatisticsFile statisticsFile1 = + TaskTestUtils.writeStatsFile( + snapshot.snapshotId(), + snapshot.sequenceNumber(), + "/metadata/" + UUID.randomUUID() + ".stats", + fileIO); + String firstMetadataFile = "v1-295495059.metadata.json"; + TableMetadata firstMetadata = + TaskTestUtils.writeTableMetadata( + fileIO, firstMetadataFile, List.of(statisticsFile1), snapshot); + assertThat(TaskUtils.exists(firstMetadataFile, fileIO)).isTrue(); + + ManifestFile manifestFile3 = + TaskTestUtils.manifestFile( + fileIO, "manifest3.avro", snapshot.snapshotId() + 1, "dataFile5.parquet"); + Snapshot snapshot2 = + TaskTestUtils.newSnapshot( + fileIO, + "manifestList2.avro", + snapshot.sequenceNumber() + 1, + snapshot.snapshotId() + 1, + snapshot.snapshotId(), + manifestFile1, + manifestFile3); // exclude manifest2 from the new snapshot + StatisticsFile statisticsFile2 = + TaskTestUtils.writeStatsFile( + snapshot2.snapshotId(), + snapshot2.sequenceNumber(), + "/metadata/" + UUID.randomUUID() + ".stats", + fileIO); + String secondMetadataFile = "v1-295495060.metadata.json"; + TaskTestUtils.writeTableMetadata( + fileIO, + secondMetadataFile, + firstMetadata, + firstMetadataFile, + List.of(statisticsFile2), + snapshot2); + assertThat(TaskUtils.exists(firstMetadataFile, fileIO)).isTrue(); + assertThat(TaskUtils.exists(secondMetadataFile, fileIO)).isTrue(); + + TaskEntity task = + new TaskEntity.Builder() + .withTaskType(diagnostics, AsyncTaskType.ENTITY_CLEANUP_SCHEDULER) + .withData( + diagnostics, + new TableLikeEntity.Builder(tableIdentifier, secondMetadataFile) + .setName("table1") + .setCatalogId(1) + .setCreateTimestamp(100) + .build()) + .build(); + + Assertions.assertThatPredicate(handler::canHandleTask).accepts(task); + + handler.handleTask(task, realmContext); + + List entities = + metaStoreManagerFactory + .getOrCreateMetaStoreManager(realmContext) + .loadTasks(metaStoreSession, "test", 6) + .getEntities(); + + List manifestCleanupTasks = + entities.stream() + .filter( + entity -> { + AsyncTaskType taskType = TaskEntity.of(entity).getTaskType(diagnostics); + return taskType == AsyncTaskType.MANIFEST_FILE_CLEANUP; + }) + .toList(); + List metadataCleanupTasks = + entities.stream() + .filter( + entity -> { + AsyncTaskType taskType = TaskEntity.of(entity).getTaskType(diagnostics); + return taskType == AsyncTaskType.METADATA_FILE_BATCH_CLEANUP; + }) + .toList(); + + assertThat(metadataCleanupTasks) + .hasSize(1) + .satisfiesExactlyInAnyOrder( + taskEntity -> + assertThat(taskEntity) + .returns(PolarisEntityType.TASK.getCode(), PolarisBaseEntity::getTypeCode) + .extracting(TaskEntity::of) + .returns( + new ManifestFileCleanupTaskHandler.ManifestCleanupTask( + tableIdentifier, + List.of( + firstMetadataFile, statisticsFile1.path(), statisticsFile2.path())), + entity -> + entity.readData( + diagnostics, + ManifestFileCleanupTaskHandler.ManifestCleanupTask.class))); + + assertThat(manifestCleanupTasks) + .hasSize(3) + .satisfiesExactlyInAnyOrder( + taskEntity -> + assertThat(taskEntity) + .returns(PolarisEntityType.TASK.getCode(), PolarisBaseEntity::getTypeCode) + .extracting(TaskEntity::of) + .returns( + new ManifestFileCleanupTaskHandler.ManifestCleanupTask( + tableIdentifier, + Base64.encodeBase64String(ManifestFiles.encode(manifestFile1))), + entity -> + entity.readData( + diagnostics, + ManifestFileCleanupTaskHandler.ManifestCleanupTask.class)), + taskEntity -> + assertThat(taskEntity) + .returns(PolarisEntityType.TASK.getCode(), PolarisBaseEntity::getTypeCode) + .extracting(TaskEntity::of) + .returns( + new ManifestFileCleanupTaskHandler.ManifestCleanupTask( + tableIdentifier, + Base64.encodeBase64String(ManifestFiles.encode(manifestFile2))), + entity -> + entity.readData( + diagnostics, + ManifestFileCleanupTaskHandler.ManifestCleanupTask.class)), + taskEntity -> + assertThat(taskEntity) + .returns(PolarisEntityType.TASK.getCode(), PolarisBaseEntity::getTypeCode) + .extracting(TaskEntity::of) + .returns( + new ManifestFileCleanupTaskHandler.ManifestCleanupTask( + tableIdentifier, + Base64.encodeBase64String(ManifestFiles.encode(manifestFile3))), + entity -> + entity.readData( + diagnostics, + ManifestFileCleanupTaskHandler.ManifestCleanupTask.class))); } } diff --git a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/test/PolarisIntegrationTestFixture.java b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/test/PolarisIntegrationTestFixture.java index 963c917c6..5f7e769ae 100644 --- a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/test/PolarisIntegrationTestFixture.java +++ b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/test/PolarisIntegrationTestFixture.java @@ -31,12 +31,10 @@ import java.net.URI; import java.util.List; import java.util.Map; -import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.admin.model.GrantPrincipalRoleRequest; import org.apache.polaris.core.admin.model.Principal; import org.apache.polaris.core.admin.model.PrincipalRole; import org.apache.polaris.core.admin.model.PrincipalWithCredentials; -import org.apache.polaris.core.context.CallContext; import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.PolarisEntityConstants; import org.apache.polaris.core.entity.PolarisEntitySubType; @@ -100,34 +98,24 @@ private PolarisPrincipalSecrets fetchAdminSecrets() { helper.metaStoreManagerFactory.bootstrapRealms(List.of(realm)); } - RealmContext realmContext = - helper.realmContextResolver.resolveRealmContext( - baseUri.toString(), "GET", "/", Map.of(REALM_PROPERTY_KEY, realm)); + RealmContext realmContext = () -> realm; PolarisMetaStoreSession metaStoreSession = helper.metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(); - PolarisCallContext polarisContext = - new PolarisCallContext( - metaStoreSession, helper.diagServices, helper.configurationStore, helper.clock); - try (CallContext ctx = CallContext.of(realmContext, polarisContext)) { - CallContext.setCurrentContext(ctx); - PolarisMetaStoreManager metaStoreManager = - helper.metaStoreManagerFactory.getOrCreateMetaStoreManager(ctx.getRealmContext()); - PolarisMetaStoreManager.EntityResult principal = - metaStoreManager.readEntityByName( - ctx.getPolarisCallContext(), - null, - PolarisEntityType.PRINCIPAL, - PolarisEntitySubType.NULL_SUBTYPE, - PolarisEntityConstants.getRootPrincipalName()); - - Map propertiesMap = readInternalProperties(principal); - return metaStoreManager - .loadPrincipalSecrets(ctx.getPolarisCallContext(), propertiesMap.get("client_id")) - .getPrincipalSecrets(); - } finally { - CallContext.unsetCurrentContext(); - } + PolarisMetaStoreManager metaStoreManager = + helper.metaStoreManagerFactory.getOrCreateMetaStoreManager(realmContext); + PolarisMetaStoreManager.EntityResult principal = + metaStoreManager.readEntityByName( + metaStoreSession, + null, + PolarisEntityType.PRINCIPAL, + PolarisEntitySubType.NULL_SUBTYPE, + PolarisEntityConstants.getRootPrincipalName()); + + Map propertiesMap = readInternalProperties(principal); + return metaStoreManager + .loadPrincipalSecrets(metaStoreSession, propertiesMap.get("client_id")) + .getPrincipalSecrets(); } private SnowmanCredentials createSnowmanCredentials(TestEnvironment testEnv) { diff --git a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/test/PolarisIntegrationTestHelper.java b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/test/PolarisIntegrationTestHelper.java index f7f5ed719..1ebb69b6f 100644 --- a/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/test/PolarisIntegrationTestHelper.java +++ b/quarkus/service/src/test/java/org/apache/polaris/service/quarkus/test/PolarisIntegrationTestHelper.java @@ -21,22 +21,14 @@ import com.fasterxml.jackson.databind.ObjectMapper; import jakarta.inject.Inject; import jakarta.inject.Singleton; -import java.time.Clock; -import org.apache.polaris.core.PolarisConfigurationStore; -import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.persistence.MetaStoreManagerFactory; -import org.apache.polaris.service.context.RealmContextResolver; import org.junit.jupiter.api.TestInfo; @Singleton public class PolarisIntegrationTestHelper { @Inject MetaStoreManagerFactory metaStoreManagerFactory; - @Inject RealmContextResolver realmContextResolver; @Inject ObjectMapper objectMapper; - @Inject PolarisDiagnostics diagServices; - @Inject PolarisConfigurationStore configurationStore; - @Inject Clock clock; public PolarisIntegrationTestFixture createFixture(TestEnvironment testEnv, TestInfo testInfo) { return new PolarisIntegrationTestFixture(this, testEnv, testInfo); diff --git a/service/common/src/main/java/org/apache/polaris/service/admin/PolarisAdminService.java b/service/common/src/main/java/org/apache/polaris/service/admin/PolarisAdminService.java index 6a90e2ab4..9d183033b 100644 --- a/service/common/src/main/java/org/apache/polaris/service/admin/PolarisAdminService.java +++ b/service/common/src/main/java/org/apache/polaris/service/admin/PolarisAdminService.java @@ -22,7 +22,6 @@ import jakarta.annotation.Nonnull; import jakarta.annotation.Nullable; -import jakarta.validation.constraints.NotNull; import jakarta.ws.rs.core.SecurityContext; import java.util.ArrayList; import java.util.Arrays; @@ -44,8 +43,8 @@ import org.apache.iceberg.exceptions.NoSuchViewException; import org.apache.iceberg.exceptions.NotFoundException; import org.apache.iceberg.exceptions.ValidationException; -import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.PolarisConfiguration; +import org.apache.polaris.core.PolarisConfigurationStore; import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.admin.model.CatalogGrant; import org.apache.polaris.core.admin.model.CatalogPrivilege; @@ -67,7 +66,7 @@ import org.apache.polaris.core.auth.PolarisAuthorizer; import org.apache.polaris.core.auth.PolarisGrantManager.LoadGrantsResult; import org.apache.polaris.core.catalog.PolarisCatalogHelpers; -import org.apache.polaris.core.context.CallContext; +import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.CatalogEntity; import org.apache.polaris.core.entity.CatalogRoleEntity; import org.apache.polaris.core.entity.NamespaceEntity; @@ -83,6 +82,7 @@ import org.apache.polaris.core.entity.TableLikeEntity; import org.apache.polaris.core.persistence.PolarisEntityManager; import org.apache.polaris.core.persistence.PolarisMetaStoreManager; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; import org.apache.polaris.core.persistence.PolarisResolvedPathWrapper; import org.apache.polaris.core.persistence.resolver.PolarisResolutionManifest; import org.apache.polaris.core.persistence.resolver.ResolverPath; @@ -107,7 +107,9 @@ public class PolarisAdminService { private static final Logger LOGGER = LoggerFactory.getLogger(PolarisAdminService.class); - private final CallContext callContext; + private final RealmContext realmContext; + private final PolarisMetaStoreSession metaStoreSession; + private final PolarisConfigurationStore configurationStore; private final PolarisEntityManager entityManager; private final SecurityContext securityContext; private final AuthenticatedPolarisPrincipal authenticatedPrincipal; @@ -118,16 +120,20 @@ public class PolarisAdminService { private PolarisResolutionManifest resolutionManifest = null; public PolarisAdminService( - @NotNull CallContext callContext, - @NotNull PolarisEntityManager entityManager, - @NotNull PolarisMetaStoreManager metaStoreManager, - @NotNull SecurityContext securityContext, - @NotNull PolarisAuthorizer authorizer) { - this.callContext = callContext; + RealmContext realmContext, + PolarisEntityManager entityManager, + PolarisMetaStoreManager metaStoreManager, + PolarisMetaStoreSession metaStoreSession, + PolarisConfigurationStore configurationStore, + PolarisDiagnostics diagServices, + SecurityContext securityContext, + PolarisAuthorizer authorizer) { + this.realmContext = realmContext; + this.metaStoreSession = metaStoreSession; + this.configurationStore = configurationStore; this.entityManager = entityManager; this.metaStoreManager = metaStoreManager; this.securityContext = securityContext; - PolarisDiagnostics diagServices = callContext.getPolarisCallContext().getDiagServices(); diagServices.checkNotNull(securityContext, "null_security_context"); diagServices.checkNotNull(securityContext.getUserPrincipal(), "null_security_context"); diagServices.check( @@ -140,10 +146,6 @@ public PolarisAdminService( this.authorizer = authorizer; } - private PolarisCallContext getCurrentPolarisContext() { - return callContext.getPolarisCallContext(); - } - private Optional findCatalogByName(String name) { return Optional.ofNullable(resolutionManifest.getResolvedReferenceCatalogEntity()) .map(path -> CatalogEntity.of(path.getRawLeafEntity())); @@ -169,11 +171,12 @@ private Optional findCatalogRoleByName(String catalogName, St private void authorizeBasicRootOperationOrThrow(PolarisAuthorizableOperation op) { resolutionManifest = entityManager.prepareResolutionManifest( - callContext, securityContext, null /* referenceCatalogName */); + metaStoreSession, securityContext, null /* referenceCatalogName */); resolutionManifest.resolveAll(); PolarisResolvedPathWrapper rootContainerWrapper = resolutionManifest.getResolvedRootContainerEntityAsPath(); authorizer.authorizeOrThrow( + realmContext, authenticatedPrincipal, resolutionManifest.getAllActivatedPrincipalRoleEntities(), op, @@ -195,7 +198,8 @@ private void authorizeBasicTopLevelEntityOperationOrThrow( PolarisEntityType entityType, @Nullable String referenceCatalogName) { resolutionManifest = - entityManager.prepareResolutionManifest(callContext, securityContext, referenceCatalogName); + entityManager.prepareResolutionManifest( + metaStoreSession, securityContext, referenceCatalogName); resolutionManifest.addTopLevelName(topLevelEntityName, entityType, false /* isOptional */); ResolverStatus status = resolutionManifest.resolveAll(); if (status.getStatus() == ResolverStatus.StatusEnum.ENTITY_COULD_NOT_BE_RESOLVED) { @@ -218,6 +222,7 @@ private void authorizeBasicTopLevelEntityOperationOrThrow( return; } authorizer.authorizeOrThrow( + realmContext, authenticatedPrincipal, resolutionManifest.getAllActivatedCatalogRoleAndPrincipalRoles(), op, @@ -228,7 +233,7 @@ private void authorizeBasicTopLevelEntityOperationOrThrow( private void authorizeBasicCatalogRoleOperationOrThrow( PolarisAuthorizableOperation op, String catalogName, String catalogRoleName) { resolutionManifest = - entityManager.prepareResolutionManifest(callContext, securityContext, catalogName); + entityManager.prepareResolutionManifest(metaStoreSession, securityContext, catalogName); resolutionManifest.addPath( new ResolverPath(List.of(catalogRoleName), PolarisEntityType.CATALOG_ROLE), catalogRoleName); @@ -238,6 +243,7 @@ private void authorizeBasicCatalogRoleOperationOrThrow( throw new NotFoundException("CatalogRole does not exist: %s", catalogRoleName); } authorizer.authorizeOrThrow( + realmContext, authenticatedPrincipal, resolutionManifest.getAllActivatedCatalogRoleAndPrincipalRoles(), op, @@ -248,7 +254,7 @@ private void authorizeBasicCatalogRoleOperationOrThrow( private void authorizeGrantOnRootContainerToPrincipalRoleOperationOrThrow( PolarisAuthorizableOperation op, String principalRoleName) { resolutionManifest = - entityManager.prepareResolutionManifest(callContext, securityContext, null); + entityManager.prepareResolutionManifest(metaStoreSession, securityContext, null); resolutionManifest.addTopLevelName( principalRoleName, PolarisEntityType.PRINCIPAL_ROLE, false /* isOptional */); ResolverStatus status = resolutionManifest.resolveAll(); @@ -268,6 +274,7 @@ private void authorizeGrantOnRootContainerToPrincipalRoleOperationOrThrow( principalRoleName, PolarisEntityType.PRINCIPAL_ROLE); authorizer.authorizeOrThrow( + realmContext, authenticatedPrincipal, resolutionManifest.getAllActivatedCatalogRoleAndPrincipalRoles(), op, @@ -281,7 +288,7 @@ private void authorizeGrantOnTopLevelEntityToPrincipalRoleOperationOrThrow( PolarisEntityType topLevelEntityType, String principalRoleName) { resolutionManifest = - entityManager.prepareResolutionManifest(callContext, securityContext, null); + entityManager.prepareResolutionManifest(metaStoreSession, securityContext, null); resolutionManifest.addTopLevelName( topLevelEntityName, topLevelEntityType, false /* isOptional */); resolutionManifest.addTopLevelName( @@ -304,6 +311,7 @@ private void authorizeGrantOnTopLevelEntityToPrincipalRoleOperationOrThrow( principalRoleName, PolarisEntityType.PRINCIPAL_ROLE); authorizer.authorizeOrThrow( + realmContext, authenticatedPrincipal, resolutionManifest.getAllActivatedCatalogRoleAndPrincipalRoles(), op, @@ -314,7 +322,7 @@ private void authorizeGrantOnTopLevelEntityToPrincipalRoleOperationOrThrow( private void authorizeGrantOnPrincipalRoleToPrincipalOperationOrThrow( PolarisAuthorizableOperation op, String principalRoleName, String principalName) { resolutionManifest = - entityManager.prepareResolutionManifest(callContext, securityContext, null); + entityManager.prepareResolutionManifest(metaStoreSession, securityContext, null); resolutionManifest.addTopLevelName( principalRoleName, PolarisEntityType.PRINCIPAL_ROLE, false /* isOptional */); resolutionManifest.addTopLevelName( @@ -334,6 +342,7 @@ private void authorizeGrantOnPrincipalRoleToPrincipalOperationOrThrow( resolutionManifest.getResolvedTopLevelEntity(principalName, PolarisEntityType.PRINCIPAL); authorizer.authorizeOrThrow( + realmContext, authenticatedPrincipal, resolutionManifest.getAllActivatedCatalogRoleAndPrincipalRoles(), op, @@ -347,7 +356,7 @@ private void authorizeGrantOnCatalogRoleToPrincipalRoleOperationOrThrow( String catalogRoleName, String principalRoleName) { resolutionManifest = - entityManager.prepareResolutionManifest(callContext, securityContext, catalogName); + entityManager.prepareResolutionManifest(metaStoreSession, securityContext, catalogName); resolutionManifest.addPath( new ResolverPath(List.of(catalogRoleName), PolarisEntityType.CATALOG_ROLE), catalogRoleName); @@ -372,6 +381,7 @@ private void authorizeGrantOnCatalogRoleToPrincipalRoleOperationOrThrow( resolutionManifest.getResolvedPath(catalogRoleName, true); authorizer.authorizeOrThrow( + realmContext, authenticatedPrincipal, resolutionManifest.getAllActivatedCatalogRoleAndPrincipalRoles(), op, @@ -382,7 +392,7 @@ private void authorizeGrantOnCatalogRoleToPrincipalRoleOperationOrThrow( private void authorizeGrantOnCatalogOperationOrThrow( PolarisAuthorizableOperation op, String catalogName, String catalogRoleName) { resolutionManifest = - entityManager.prepareResolutionManifest(callContext, securityContext, catalogName); + entityManager.prepareResolutionManifest(metaStoreSession, securityContext, catalogName); resolutionManifest.addTopLevelName( catalogName, PolarisEntityType.CATALOG, false /* isOptional */); resolutionManifest.addPath( @@ -401,6 +411,7 @@ private void authorizeGrantOnCatalogOperationOrThrow( PolarisResolvedPathWrapper catalogRoleWrapper = resolutionManifest.getResolvedPath(catalogRoleName, true); authorizer.authorizeOrThrow( + realmContext, authenticatedPrincipal, resolutionManifest.getAllActivatedCatalogRoleAndPrincipalRoles(), op, @@ -414,7 +425,7 @@ private void authorizeGrantOnNamespaceOperationOrThrow( Namespace namespace, String catalogRoleName) { resolutionManifest = - entityManager.prepareResolutionManifest(callContext, securityContext, catalogName); + entityManager.prepareResolutionManifest(metaStoreSession, securityContext, catalogName); resolutionManifest.addPath( new ResolverPath(Arrays.asList(namespace.levels()), PolarisEntityType.NAMESPACE), namespace); @@ -440,6 +451,7 @@ private void authorizeGrantOnNamespaceOperationOrThrow( resolutionManifest.getResolvedPath(catalogRoleName, true); authorizer.authorizeOrThrow( + realmContext, authenticatedPrincipal, resolutionManifest.getAllActivatedCatalogRoleAndPrincipalRoles(), op, @@ -454,7 +466,7 @@ private void authorizeGrantOnTableLikeOperationOrThrow( TableIdentifier identifier, String catalogRoleName) { resolutionManifest = - entityManager.prepareResolutionManifest(callContext, securityContext, catalogName); + entityManager.prepareResolutionManifest(metaStoreSession, securityContext, catalogName); resolutionManifest.addPath( new ResolverPath( PolarisCatalogHelpers.tableIdentifierToList(identifier), PolarisEntityType.TABLE_LIKE), @@ -484,6 +496,7 @@ private void authorizeGrantOnTableLikeOperationOrThrow( resolutionManifest.getResolvedPath(catalogRoleName, true); authorizer.authorizeOrThrow( + realmContext, authenticatedPrincipal, resolutionManifest.getAllActivatedCatalogRoleAndPrincipalRoles(), op, @@ -521,10 +534,8 @@ private String terminateWithSlash(String path) { */ private boolean catalogOverlapsWithExistingCatalog(CatalogEntity catalogEntity) { boolean allowOverlappingCatalogUrls = - getCurrentPolarisContext() - .getConfigurationStore() - .getConfiguration( - getCurrentPolarisContext(), PolarisConfiguration.ALLOW_OVERLAPPING_CATALOG_URLS); + configurationStore.getConfiguration( + realmContext, PolarisConfiguration.ALLOW_OVERLAPPING_CATALOG_URLS); if (allowOverlappingCatalogUrls) { return false; @@ -569,11 +580,11 @@ public PolarisEntity createCatalog(PolarisEntity entity) { PolarisEntity polarisEntity = new PolarisEntity.Builder(entity) - .setId(metaStoreManager.generateNewEntityId(getCurrentPolarisContext()).getId()) + .setId(metaStoreManager.generateNewEntityId(metaStoreSession).getId()) .setCreateTimestamp(System.currentTimeMillis()) .build(); PolarisMetaStoreManager.CreateCatalogResult catalogResult = - metaStoreManager.createCatalog(getCurrentPolarisContext(), polarisEntity, List.of()); + metaStoreManager.createCatalog(metaStoreSession, polarisEntity, List.of()); if (catalogResult.alreadyExists()) { throw new AlreadyExistsException( "Cannot create Catalog %s. Catalog already exists or resolution failed", @@ -590,14 +601,11 @@ public void deleteCatalog(String name) { findCatalogByName(name) .orElseThrow(() -> new NotFoundException("Catalog %s not found", name)); // TODO: Handle return value in case of concurrent modification - PolarisCallContext polarisCallContext = callContext.getPolarisCallContext(); boolean cleanup = - polarisCallContext - .getConfigurationStore() - .getConfiguration(polarisCallContext, PolarisConfiguration.CLEANUP_ON_CATALOG_DROP); + configurationStore.getConfiguration( + realmContext, PolarisConfiguration.CLEANUP_ON_CATALOG_DROP); PolarisMetaStoreManager.DropEntityResult dropEntityResult = - metaStoreManager.dropEntityIfExists( - getCurrentPolarisContext(), null, entity, Map.of(), cleanup); + metaStoreManager.dropEntityIfExists(metaStoreSession, null, entity, Map.of(), cleanup); // at least some handling of error if (!dropEntityResult.isSuccess()) { @@ -727,7 +735,7 @@ private void validateUpdateCatalogDiffOrThrow( CatalogEntity.of( PolarisEntity.of( metaStoreManager.updateEntityPropertiesIfNotChanged( - getCurrentPolarisContext(), null, updatedEntity)))) + metaStoreSession, null, updatedEntity)))) .orElseThrow( () -> new CommitFailedException( @@ -749,16 +757,13 @@ public List listCatalogs() { private List listCatalogsUnsafe() { return metaStoreManager .listEntities( - getCurrentPolarisContext(), - null, - PolarisEntityType.CATALOG, - PolarisEntitySubType.ANY_SUBTYPE) + metaStoreSession, null, PolarisEntityType.CATALOG, PolarisEntitySubType.ANY_SUBTYPE) .getEntities() .stream() .map( nameAndId -> PolarisEntity.of( - metaStoreManager.loadEntity(getCurrentPolarisContext(), 0, nameAndId.getId()))) + metaStoreManager.loadEntity(metaStoreSession, 0, nameAndId.getId()))) .toList(); } @@ -770,9 +775,9 @@ public PrincipalWithCredentials createPrincipal(PolarisEntity entity) { PolarisMetaStoreManager.CreatePrincipalResult principalResult = metaStoreManager.createPrincipal( - getCurrentPolarisContext(), + metaStoreSession, new PolarisEntity.Builder(entity) - .setId(metaStoreManager.generateNewEntityId(getCurrentPolarisContext()).getId()) + .setId(metaStoreManager.generateNewEntityId(metaStoreSession).getId()) .setCreateTimestamp(System.currentTimeMillis()) .build()); if (principalResult.alreadyExists()) { @@ -796,8 +801,7 @@ public void deletePrincipal(String name) { .orElseThrow(() -> new NotFoundException("Principal %s not found", name)); // TODO: Handle return value in case of concurrent modification PolarisMetaStoreManager.DropEntityResult dropEntityResult = - metaStoreManager.dropEntityIfExists( - getCurrentPolarisContext(), null, entity, Map.of(), false); + metaStoreManager.dropEntityIfExists(metaStoreSession, null, entity, Map.of(), false); // at least some handling of error if (!dropEntityResult.isSuccess()) { @@ -844,7 +848,7 @@ public void deletePrincipal(String name) { PrincipalEntity.of( PolarisEntity.of( metaStoreManager.updateEntityPropertiesIfNotChanged( - getCurrentPolarisContext(), null, updatedEntity)))) + metaStoreSession, null, updatedEntity)))) .orElseThrow( () -> new CommitFailedException( @@ -860,7 +864,7 @@ public void deletePrincipal(String name) { PolarisPrincipalSecrets currentSecrets = metaStoreManager - .loadPrincipalSecrets(getCurrentPolarisContext(), currentPrincipalEntity.getClientId()) + .loadPrincipalSecrets(metaStoreSession, currentPrincipalEntity.getClientId()) .getPrincipalSecrets(); if (currentSecrets == null) { throw new IllegalArgumentException( @@ -869,7 +873,7 @@ public void deletePrincipal(String name) { PolarisPrincipalSecrets newSecrets = metaStoreManager .rotatePrincipalSecrets( - getCurrentPolarisContext(), + metaStoreSession, currentPrincipalEntity.getClientId(), currentPrincipalEntity.getId(), shouldReset, @@ -883,8 +887,7 @@ public void deletePrincipal(String name) { } PolarisEntity newPrincipal = PolarisEntity.of( - metaStoreManager.loadEntity( - getCurrentPolarisContext(), 0L, currentPrincipalEntity.getId())); + metaStoreManager.loadEntity(metaStoreSession, 0L, currentPrincipalEntity.getId())); return new PrincipalWithCredentials( PrincipalEntity.of(newPrincipal).asPrincipal(), new PrincipalWithCredentialsCredentials( @@ -911,16 +914,13 @@ public List listPrincipals() { return metaStoreManager .listEntities( - getCurrentPolarisContext(), - null, - PolarisEntityType.PRINCIPAL, - PolarisEntitySubType.NULL_SUBTYPE) + metaStoreSession, null, PolarisEntityType.PRINCIPAL, PolarisEntitySubType.NULL_SUBTYPE) .getEntities() .stream() .map( nameAndId -> PolarisEntity.of( - metaStoreManager.loadEntity(getCurrentPolarisContext(), 0, nameAndId.getId()))) + metaStoreManager.loadEntity(metaStoreSession, 0, nameAndId.getId()))) .toList(); } @@ -933,10 +933,10 @@ public PolarisEntity createPrincipalRole(PolarisEntity entity) { PolarisEntity returnedEntity = PolarisEntity.of( metaStoreManager.createEntityIfNotExists( - getCurrentPolarisContext(), + metaStoreSession, null, new PolarisEntity.Builder(entity) - .setId(metaStoreManager.generateNewEntityId(getCurrentPolarisContext()).getId()) + .setId(metaStoreManager.generateNewEntityId(metaStoreSession).getId()) .setCreateTimestamp(System.currentTimeMillis()) .build())); if (returnedEntity == null) { @@ -957,7 +957,7 @@ public void deletePrincipalRole(String name) { // TODO: Handle return value in case of concurrent modification PolarisMetaStoreManager.DropEntityResult dropEntityResult = metaStoreManager.dropEntityIfExists( - getCurrentPolarisContext(), null, entity, Map.of(), true); // cleanup grants + metaStoreSession, null, entity, Map.of(), true); // cleanup grants // at least some handling of error if (!dropEntityResult.isSuccess()) { @@ -1005,7 +1005,7 @@ public void deletePrincipalRole(String name) { PrincipalRoleEntity.of( PolarisEntity.of( metaStoreManager.updateEntityPropertiesIfNotChanged( - getCurrentPolarisContext(), null, updatedEntity)))) + metaStoreSession, null, updatedEntity)))) .orElseThrow( () -> new CommitFailedException( @@ -1019,7 +1019,7 @@ public List listPrincipalRoles() { return metaStoreManager .listEntities( - getCurrentPolarisContext(), + metaStoreSession, null, PolarisEntityType.PRINCIPAL_ROLE, PolarisEntitySubType.NULL_SUBTYPE) @@ -1028,7 +1028,7 @@ public List listPrincipalRoles() { .map( nameAndId -> PolarisEntity.of( - metaStoreManager.loadEntity(getCurrentPolarisContext(), 0, nameAndId.getId()))) + metaStoreManager.loadEntity(metaStoreSession, 0, nameAndId.getId()))) .toList(); } @@ -1045,10 +1045,10 @@ public PolarisEntity createCatalogRole(String catalogName, PolarisEntity entity) PolarisEntity returnedEntity = PolarisEntity.of( metaStoreManager.createEntityIfNotExists( - getCurrentPolarisContext(), + metaStoreSession, PolarisEntity.toCoreList(List.of(catalogEntity)), new PolarisEntity.Builder(entity) - .setId(metaStoreManager.generateNewEntityId(getCurrentPolarisContext()).getId()) + .setId(metaStoreManager.generateNewEntityId(metaStoreSession).getId()) .setCatalogId(catalogEntity.getId()) .setParentId(catalogEntity.getId()) .setCreateTimestamp(System.currentTimeMillis()) @@ -1072,7 +1072,7 @@ public void deleteCatalogRole(String catalogName, String name) { // TODO: Handle return value in case of concurrent modification PolarisMetaStoreManager.DropEntityResult dropEntityResult = metaStoreManager.dropEntityIfExists( - getCurrentPolarisContext(), + metaStoreSession, PolarisEntity.toCoreList(resolvedCatalogRoleEntity.getRawParentPath()), resolvedCatalogRoleEntity.getRawLeafEntity(), Map.of(), @@ -1127,7 +1127,7 @@ public void deleteCatalogRole(String catalogName, String name) { CatalogRoleEntity.of( PolarisEntity.of( metaStoreManager.updateEntityPropertiesIfNotChanged( - getCurrentPolarisContext(), + metaStoreSession, PolarisEntity.toCoreList(List.of(catalogEntity)), updatedEntity)))) .orElseThrow( @@ -1146,7 +1146,7 @@ public List listCatalogRoles(String catalogName) { .orElseThrow(() -> new NotFoundException("Parent catalog %s not found", catalogName)); return metaStoreManager .listEntities( - getCurrentPolarisContext(), + metaStoreSession, PolarisEntity.toCoreList(List.of(catalogEntity)), PolarisEntityType.CATALOG_ROLE, PolarisEntitySubType.NULL_SUBTYPE) @@ -1156,7 +1156,7 @@ public List listCatalogRoles(String catalogName) { nameAndId -> PolarisEntity.of( metaStoreManager.loadEntity( - getCurrentPolarisContext(), catalogEntity.getId(), nameAndId.getId()))) + metaStoreSession, catalogEntity.getId(), nameAndId.getId()))) .toList(); } @@ -1173,8 +1173,7 @@ public boolean assignPrincipalRole(String principalName, String principalRoleNam () -> new NotFoundException("PrincipalRole %s not found", principalRoleName)); return metaStoreManager - .grantUsageOnRoleToGrantee( - getCurrentPolarisContext(), null, principalRoleEntity, principalEntity) + .grantUsageOnRoleToGrantee(metaStoreSession, null, principalRoleEntity, principalEntity) .isSuccess(); } @@ -1190,8 +1189,7 @@ public boolean revokePrincipalRole(String principalName, String principalRoleNam .orElseThrow( () -> new NotFoundException("PrincipalRole %s not found", principalRoleName)); return metaStoreManager - .revokeUsageOnRoleFromGrantee( - getCurrentPolarisContext(), null, principalRoleEntity, principalEntity) + .revokeUsageOnRoleFromGrantee(metaStoreSession, null, principalRoleEntity, principalEntity) .isSuccess(); } @@ -1205,7 +1203,7 @@ public List listPrincipalRolesAssigned(String principalName) { .orElseThrow(() -> new NotFoundException("Principal %s not found", principalName)); LoadGrantsResult grantList = metaStoreManager.loadGrantsToGrantee( - getCurrentPolarisContext(), principalEntity.getCatalogId(), principalEntity.getId()); + metaStoreSession, principalEntity.getCatalogId(), principalEntity.getId()); return buildEntitiesFromGrantResults(grantList, false, null); } @@ -1229,7 +1227,7 @@ public boolean assignCatalogRoleToPrincipalRole( return metaStoreManager .grantUsageOnRoleToGrantee( - getCurrentPolarisContext(), catalogEntity, catalogRoleEntity, principalRoleEntity) + metaStoreSession, catalogEntity, catalogRoleEntity, principalRoleEntity) .isSuccess(); } @@ -1252,7 +1250,7 @@ public boolean revokeCatalogRoleFromPrincipalRole( .orElseThrow(() -> new NotFoundException("CatalogRole %s not found", catalogRoleName)); return metaStoreManager .revokeUsageOnRoleFromGrantee( - getCurrentPolarisContext(), catalogEntity, catalogRoleEntity, principalRoleEntity) + metaStoreSession, catalogEntity, catalogRoleEntity, principalRoleEntity) .isSuccess(); } @@ -1269,9 +1267,7 @@ public List listAssigneePrincipalsForPrincipalRole(String princip () -> new NotFoundException("PrincipalRole %s not found", principalRoleName)); LoadGrantsResult grantList = metaStoreManager.loadGrantsOnSecurable( - getCurrentPolarisContext(), - principalRoleEntity.getCatalogId(), - principalRoleEntity.getId()); + metaStoreSession, principalRoleEntity.getCatalogId(), principalRoleEntity.getId()); return buildEntitiesFromGrantResults(grantList, true, null); } @@ -1322,9 +1318,7 @@ public List listCatalogRolesForPrincipalRole( () -> new NotFoundException("PrincipalRole %s not found", principalRoleName)); LoadGrantsResult grantList = metaStoreManager.loadGrantsToGrantee( - getCurrentPolarisContext(), - principalRoleEntity.getCatalogId(), - principalRoleEntity.getId()); + metaStoreSession, principalRoleEntity.getCatalogId(), principalRoleEntity.getId()); return buildEntitiesFromGrantResults( grantList, false, grantRec -> grantRec.getSecurableCatalogId() == catalogEntity.getId()); } @@ -1344,7 +1338,7 @@ public boolean grantPrivilegeOnRootContainerToPrincipalRole( return metaStoreManager .grantPrivilegeOnSecurableToRole( - getCurrentPolarisContext(), principalRoleEntity, null, rootContainerEntity, privilege) + metaStoreSession, principalRoleEntity, null, rootContainerEntity, privilege) .isSuccess(); } @@ -1364,7 +1358,7 @@ public boolean revokePrivilegeOnRootContainerFromPrincipalRole( return metaStoreManager .revokePrivilegeOnSecurableFromRole( - getCurrentPolarisContext(), principalRoleEntity, null, rootContainerEntity, privilege) + metaStoreSession, principalRoleEntity, null, rootContainerEntity, privilege) .isSuccess(); } @@ -1388,7 +1382,7 @@ public boolean grantPrivilegeOnCatalogToRole( return metaStoreManager .grantPrivilegeOnSecurableToRole( - getCurrentPolarisContext(), + metaStoreSession, catalogRoleEntity, PolarisEntity.toCoreList(List.of(catalogEntity)), catalogEntity, @@ -1412,7 +1406,7 @@ public boolean revokePrivilegeOnCatalogFromRole( return metaStoreManager .revokePrivilegeOnSecurableFromRole( - getCurrentPolarisContext(), + metaStoreSession, catalogRoleEntity, PolarisEntity.toCoreList(List.of(catalogEntity)), catalogEntity, @@ -1440,7 +1434,7 @@ public boolean grantPrivilegeOnNamespaceToRole( return metaStoreManager .grantPrivilegeOnSecurableToRole( - getCurrentPolarisContext(), + metaStoreSession, catalogRoleEntity, PolarisEntity.toCoreList(catalogPath), namespaceEntity, @@ -1468,7 +1462,7 @@ public boolean revokePrivilegeOnNamespaceFromRole( return metaStoreManager .revokePrivilegeOnSecurableFromRole( - getCurrentPolarisContext(), + metaStoreSession, catalogRoleEntity, PolarisEntity.toCoreList(catalogPath), namespaceEntity, @@ -1548,9 +1542,7 @@ public List listAssigneePrincipalRolesForCatalogRole( .orElseThrow(() -> new NotFoundException("CatalogRole %s not found", catalogRoleName)); LoadGrantsResult grantList = metaStoreManager.loadGrantsOnSecurable( - getCurrentPolarisContext(), - catalogRoleEntity.getCatalogId(), - catalogRoleEntity.getId()); + metaStoreSession, catalogRoleEntity.getCatalogId(), catalogRoleEntity.getId()); return buildEntitiesFromGrantResults(grantList, true, null); } @@ -1567,9 +1559,7 @@ public List listGrantsForCatalogRole(String catalogName, String c .orElseThrow(() -> new NotFoundException("CatalogRole %s not found", catalogRoleName)); LoadGrantsResult grantList = metaStoreManager.loadGrantsToGrantee( - getCurrentPolarisContext(), - catalogRoleEntity.getCatalogId(), - catalogRoleEntity.getId()); + metaStoreSession, catalogRoleEntity.getCatalogId(), catalogRoleEntity.getId()); List catalogGrants = new ArrayList<>(); List namespaceGrants = new ArrayList<>(); List tableGrants = new ArrayList<>(); @@ -1652,7 +1642,7 @@ public List listGrantsForCatalogRole(String catalogName, String c private @Nullable PolarisBaseEntity getOrLoadEntity( @Nullable Map entitiesMap, long catalogId, long id) { return (entitiesMap == null) - ? metaStoreManager.loadEntity(getCurrentPolarisContext(), catalogId, id).getEntity() + ? metaStoreManager.loadEntity(metaStoreSession, catalogId, id).getEntity() : entitiesMap.get(id); } @@ -1684,7 +1674,7 @@ private boolean grantPrivilegeOnTableLikeToRole( return metaStoreManager .grantPrivilegeOnSecurableToRole( - getCurrentPolarisContext(), + metaStoreSession, catalogRoleEntity, PolarisEntity.toCoreList(catalogPath), tableLikeEntity, @@ -1722,7 +1712,7 @@ private boolean revokePrivilegeOnTableLikeFromRole( return metaStoreManager .revokePrivilegeOnSecurableFromRole( - getCurrentPolarisContext(), + metaStoreSession, catalogRoleEntity, PolarisEntity.toCoreList(catalogPath), tableLikeEntity, diff --git a/service/common/src/main/java/org/apache/polaris/service/admin/PolarisServiceImpl.java b/service/common/src/main/java/org/apache/polaris/service/admin/PolarisServiceImpl.java index 15fb3854d..4f0a8dc7c 100644 --- a/service/common/src/main/java/org/apache/polaris/service/admin/PolarisServiceImpl.java +++ b/service/common/src/main/java/org/apache/polaris/service/admin/PolarisServiceImpl.java @@ -26,8 +26,9 @@ import org.apache.iceberg.catalog.Namespace; import org.apache.iceberg.catalog.TableIdentifier; import org.apache.iceberg.exceptions.NotAuthorizedException; -import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.PolarisConfiguration; +import org.apache.polaris.core.PolarisConfigurationStore; +import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.admin.model.AddGrantRequest; import org.apache.polaris.core.admin.model.Catalog; import org.apache.polaris.core.admin.model.CatalogGrant; @@ -58,20 +59,18 @@ import org.apache.polaris.core.admin.model.ViewGrant; import org.apache.polaris.core.auth.AuthenticatedPolarisPrincipal; import org.apache.polaris.core.auth.PolarisAuthorizer; -import org.apache.polaris.core.context.CallContext; import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.CatalogEntity; import org.apache.polaris.core.entity.CatalogRoleEntity; import org.apache.polaris.core.entity.PolarisPrivilege; import org.apache.polaris.core.entity.PrincipalEntity; import org.apache.polaris.core.entity.PrincipalRoleEntity; -import org.apache.polaris.core.persistence.MetaStoreManagerFactory; import org.apache.polaris.core.persistence.PolarisEntityManager; import org.apache.polaris.core.persistence.PolarisMetaStoreManager; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; import org.apache.polaris.service.admin.api.PolarisCatalogsApiService; import org.apache.polaris.service.admin.api.PolarisPrincipalRolesApiService; import org.apache.polaris.service.admin.api.PolarisPrincipalsApiService; -import org.apache.polaris.service.config.RealmEntityManagerFactory; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -82,23 +81,28 @@ public class PolarisServiceImpl PolarisPrincipalsApiService, PolarisPrincipalRolesApiService { private static final Logger LOGGER = LoggerFactory.getLogger(PolarisServiceImpl.class); - private final RealmEntityManagerFactory entityManagerFactory; + + private final PolarisEntityManager entityManager; + private final PolarisMetaStoreManager metaStoreManager; + private final PolarisMetaStoreSession session; + private final PolarisConfigurationStore configurationStore; private final PolarisAuthorizer polarisAuthorizer; - private final MetaStoreManagerFactory metaStoreManagerFactory; - private final CallContext callContext; + private final PolarisDiagnostics diagnostics; @Inject public PolarisServiceImpl( - RealmEntityManagerFactory entityManagerFactory, - MetaStoreManagerFactory metaStoreManagerFactory, + PolarisEntityManager entityManager, + PolarisMetaStoreManager metaStoreManager, + PolarisMetaStoreSession session, + PolarisConfigurationStore configurationStore, PolarisAuthorizer polarisAuthorizer, - CallContext callContext) { - this.entityManagerFactory = entityManagerFactory; - this.metaStoreManagerFactory = metaStoreManagerFactory; + PolarisDiagnostics diagnostics) { + this.entityManager = entityManager; + this.metaStoreManager = metaStoreManager; + this.session = session; + this.configurationStore = configurationStore; this.polarisAuthorizer = polarisAuthorizer; - this.callContext = callContext; - // FIXME: This is a hack to set the current context for downstream calls. - CallContext.setCurrentContext(callContext); + this.diagnostics = diagnostics; } private PolarisAdminService newAdminService( @@ -109,12 +113,15 @@ private PolarisAdminService newAdminService( throw new NotAuthorizedException("Failed to find authenticatedPrincipal in SecurityContext"); } - PolarisEntityManager entityManager = - entityManagerFactory.getOrCreateEntityManager(realmContext); - PolarisMetaStoreManager metaStoreManager = - metaStoreManagerFactory.getOrCreateMetaStoreManager(realmContext); return new PolarisAdminService( - callContext, entityManager, metaStoreManager, securityContext, polarisAuthorizer); + realmContext, + entityManager, + metaStoreManager, + session, + configurationStore, + diagnostics, + securityContext, + polarisAuthorizer); } /** From PolarisCatalogsApiService */ @@ -123,7 +130,7 @@ public Response createCatalog( CreateCatalogRequest request, RealmContext realmContext, SecurityContext securityContext) { PolarisAdminService adminService = newAdminService(realmContext, securityContext); Catalog catalog = request.getCatalog(); - validateStorageConfig(catalog.getStorageConfigInfo()); + validateStorageConfig(catalog.getStorageConfigInfo(), realmContext); Catalog newCatalog = new CatalogEntity(adminService.createCatalog(CatalogEntity.fromCatalog(catalog))) .asCatalog(); @@ -131,13 +138,11 @@ public Response createCatalog( return Response.status(Response.Status.CREATED).build(); } - private void validateStorageConfig(StorageConfigInfo storageConfigInfo) { - PolarisCallContext polarisCallContext = callContext.getPolarisCallContext(); + private void validateStorageConfig( + StorageConfigInfo storageConfigInfo, RealmContext realmContext) { List allowedStorageTypes = - polarisCallContext - .getConfigurationStore() - .getConfiguration( - polarisCallContext, PolarisConfiguration.SUPPORTED_CATALOG_STORAGE_TYPES); + configurationStore.getConfiguration( + realmContext, PolarisConfiguration.SUPPORTED_CATALOG_STORAGE_TYPES); if (!allowedStorageTypes.contains(storageConfigInfo.getStorageType().name())) { LOGGER .atWarn() @@ -174,7 +179,7 @@ public Response updateCatalog( SecurityContext securityContext) { PolarisAdminService adminService = newAdminService(realmContext, securityContext); if (updateRequest.getStorageConfigInfo() != null) { - validateStorageConfig(updateRequest.getStorageConfigInfo()); + validateStorageConfig(updateRequest.getStorageConfigInfo(), realmContext); } return Response.ok(adminService.updateCatalog(catalogName, updateRequest).asCatalog()).build(); } diff --git a/service/common/src/main/java/org/apache/polaris/service/auth/BasePolarisAuthenticator.java b/service/common/src/main/java/org/apache/polaris/service/auth/BasePolarisAuthenticator.java index fb713659b..2ad132c34 100644 --- a/service/common/src/main/java/org/apache/polaris/service/auth/BasePolarisAuthenticator.java +++ b/service/common/src/main/java/org/apache/polaris/service/auth/BasePolarisAuthenticator.java @@ -26,14 +26,13 @@ import org.apache.iceberg.exceptions.NotAuthorizedException; import org.apache.iceberg.exceptions.ServiceFailureException; import org.apache.polaris.core.auth.AuthenticatedPolarisPrincipal; -import org.apache.polaris.core.context.CallContext; import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.PolarisEntity; import org.apache.polaris.core.entity.PolarisEntitySubType; import org.apache.polaris.core.entity.PolarisEntityType; import org.apache.polaris.core.entity.PrincipalEntity; -import org.apache.polaris.core.persistence.MetaStoreManagerFactory; import org.apache.polaris.core.persistence.PolarisMetaStoreManager; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -50,29 +49,26 @@ public abstract class BasePolarisAuthenticator public static final String PRINCIPAL_ROLE_PREFIX = "PRINCIPAL_ROLE:"; private static final Logger LOGGER = LoggerFactory.getLogger(BasePolarisAuthenticator.class); - protected final MetaStoreManagerFactory metaStoreManagerFactory; - protected final CallContext callContext; + protected final PolarisMetaStoreManager metaStoreManager; + protected final PolarisMetaStoreSession metaStoreSession; protected BasePolarisAuthenticator( - MetaStoreManagerFactory metaStoreManagerFactory, CallContext callContext) { - this.metaStoreManagerFactory = metaStoreManagerFactory; - this.callContext = callContext; + PolarisMetaStoreManager metaStoreManager, PolarisMetaStoreSession metaStoreSession) { + this.metaStoreManager = metaStoreManager; + this.metaStoreSession = metaStoreSession; } protected Optional getPrincipal(DecodedToken tokenInfo) { LOGGER.debug("Resolving principal for tokenInfo client_id={}", tokenInfo.getClientId()); - PolarisMetaStoreManager metaStoreManager = - metaStoreManagerFactory.getOrCreateMetaStoreManager(callContext.getRealmContext()); PolarisEntity principal; try { principal = tokenInfo.getPrincipalId() > 0 ? PolarisEntity.of( - metaStoreManager.loadEntity( - callContext.getPolarisCallContext(), 0L, tokenInfo.getPrincipalId())) + metaStoreManager.loadEntity(metaStoreSession, 0L, tokenInfo.getPrincipalId())) : PolarisEntity.of( metaStoreManager.readEntityByName( - callContext.getPolarisCallContext(), + metaStoreSession, null, PolarisEntityType.PRINCIPAL, PolarisEntitySubType.NULL_SUBTYPE, @@ -109,7 +105,6 @@ protected Optional getPrincipal(DecodedToken toke AuthenticatedPolarisPrincipal authenticatedPrincipal = new AuthenticatedPolarisPrincipal(new PrincipalEntity(principal), activatedPrincipalRoles); LOGGER.debug("Populating authenticatedPrincipal into CallContext: {}", authenticatedPrincipal); - callContext.contextVariables().put(CallContext.AUTHENTICATED_PRINCIPAL, authenticatedPrincipal); return Optional.of(authenticatedPrincipal); } } diff --git a/service/common/src/main/java/org/apache/polaris/service/auth/DefaultActiveRolesProvider.java b/service/common/src/main/java/org/apache/polaris/service/auth/DefaultActiveRolesProvider.java index a9cea1f00..71f289831 100644 --- a/service/common/src/main/java/org/apache/polaris/service/auth/DefaultActiveRolesProvider.java +++ b/service/common/src/main/java/org/apache/polaris/service/auth/DefaultActiveRolesProvider.java @@ -25,15 +25,15 @@ import java.util.function.Predicate; import java.util.stream.Collectors; import org.apache.iceberg.exceptions.NotAuthorizedException; -import org.apache.polaris.core.PolarisCallContext; +import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.auth.AuthenticatedPolarisPrincipal; import org.apache.polaris.core.auth.PolarisGrantManager; -import org.apache.polaris.core.context.CallContext; import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.PolarisEntity; import org.apache.polaris.core.entity.PrincipalRoleEntity; import org.apache.polaris.core.persistence.MetaStoreManagerFactory; import org.apache.polaris.core.persistence.PolarisMetaStoreManager; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -49,6 +49,7 @@ public class DefaultActiveRolesProvider implements ActiveRolesProvider { @Inject RealmContext realmContext; @Inject MetaStoreManagerFactory metaStoreManagerFactory; + @Inject PolarisDiagnostics diagnostics; @Override public Set getActiveRoles(AuthenticatedPolarisPrincipal principal) { @@ -56,22 +57,23 @@ public Set getActiveRoles(AuthenticatedPolarisPrincipal principal) { loadActivePrincipalRoles( principal.getActivatedPrincipalRoleNames(), principal.getPrincipalEntity(), - metaStoreManagerFactory.getOrCreateMetaStoreManager(realmContext)); + metaStoreManagerFactory.getOrCreateMetaStoreManager(realmContext), + metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get()); return activeRoles.stream().map(PrincipalRoleEntity::getName).collect(Collectors.toSet()); } protected List loadActivePrincipalRoles( - Set tokenRoles, PolarisEntity principal, PolarisMetaStoreManager metaStoreManager) { - PolarisCallContext polarisContext = CallContext.getCurrentContext().getPolarisCallContext(); + Set tokenRoles, + PolarisEntity principal, + PolarisMetaStoreManager metaStoreManager, + PolarisMetaStoreSession metaStoreSession) { PolarisGrantManager.LoadGrantsResult principalGrantResults = - metaStoreManager.loadGrantsToGrantee(polarisContext, principal); - polarisContext - .getDiagServices() - .check( - principalGrantResults.isSuccess(), - "Failed to resolve principal roles for principal name={} id={}", - principal.getName(), - principal.getId()); + metaStoreManager.loadGrantsToGrantee(metaStoreSession, principal); + diagnostics.check( + principalGrantResults.isSuccess(), + "Failed to resolve principal roles for principal name={} id={}", + principal.getName(), + principal.getId()); if (!principalGrantResults.isSuccess()) { LOGGER.warn( "Failed to resolve principal roles for principal name={} id={}", @@ -87,7 +89,7 @@ protected List loadActivePrincipalRoles( .map( gr -> metaStoreManager.loadEntity( - polarisContext, gr.getSecurableCatalogId(), gr.getSecurableId())) + metaStoreSession, gr.getSecurableCatalogId(), gr.getSecurableId())) .filter(PolarisMetaStoreManager.EntityResult::isSuccess) .map(PolarisMetaStoreManager.EntityResult::getEntity) .map(PrincipalRoleEntity::of) diff --git a/service/common/src/main/java/org/apache/polaris/service/auth/DefaultOAuth2ApiService.java b/service/common/src/main/java/org/apache/polaris/service/auth/DefaultOAuth2ApiService.java index da9a5c1c7..f35dd370f 100644 --- a/service/common/src/main/java/org/apache/polaris/service/auth/DefaultOAuth2ApiService.java +++ b/service/common/src/main/java/org/apache/polaris/service/auth/DefaultOAuth2ApiService.java @@ -27,7 +27,6 @@ import jakarta.ws.rs.core.SecurityContext; import org.apache.commons.codec.binary.Base64; import org.apache.iceberg.rest.responses.OAuthTokenResponse; -import org.apache.polaris.core.context.CallContext; import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.service.catalog.api.IcebergRestOAuth2ApiService; import org.apache.polaris.service.types.TokenType; @@ -48,12 +47,10 @@ public class DefaultOAuth2ApiService implements IcebergRestOAuth2ApiService { private static final String BEARER = "bearer"; private final TokenBrokerFactory tokenBrokerFactory; - private final CallContext callContext; @Inject - public DefaultOAuth2ApiService(TokenBrokerFactory tokenBrokerFactory, CallContext callContext) { + public DefaultOAuth2ApiService(TokenBrokerFactory tokenBrokerFactory) { this.tokenBrokerFactory = tokenBrokerFactory; - this.callContext = callContext; } @Override @@ -107,18 +104,13 @@ public Response getToken( // secret and treat it as a new token request if (clientId != null && clientSecret != null) { yield tokenBroker.generateFromClientSecrets( - clientId, - clientSecret, - CLIENT_CREDENTIALS, - scope, - callContext.getPolarisCallContext()); + clientId, clientSecret, CLIENT_CREDENTIALS, scope); } else { yield tokenBroker.generateFromToken(subjectTokenType, subjectToken, grantType, scope); } } case null -> - tokenBroker.generateFromClientSecrets( - clientId, clientSecret, grantType, scope, callContext.getPolarisCallContext()); + tokenBroker.generateFromClientSecrets(clientId, clientSecret, grantType, scope); }; if (tokenResponse == null) { return OAuthUtils.getResponseFromError(OAuthTokenErrorResponse.Error.unsupported_grant_type); diff --git a/service/common/src/main/java/org/apache/polaris/service/auth/DefaultPolarisAuthenticator.java b/service/common/src/main/java/org/apache/polaris/service/auth/DefaultPolarisAuthenticator.java index 900a4ddd3..f932c8d47 100644 --- a/service/common/src/main/java/org/apache/polaris/service/auth/DefaultPolarisAuthenticator.java +++ b/service/common/src/main/java/org/apache/polaris/service/auth/DefaultPolarisAuthenticator.java @@ -23,14 +23,14 @@ import jakarta.inject.Inject; import java.util.Optional; import org.apache.polaris.core.auth.AuthenticatedPolarisPrincipal; -import org.apache.polaris.core.context.CallContext; -import org.apache.polaris.core.persistence.MetaStoreManagerFactory; +import org.apache.polaris.core.persistence.PolarisMetaStoreManager; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; @RequestScoped @Identifier("default") public class DefaultPolarisAuthenticator extends BasePolarisAuthenticator { - private final TokenBrokerFactory tokenBrokerFactory; + private final TokenBroker tokenBroker; public DefaultPolarisAuthenticator() { this(null, null, null); @@ -38,17 +38,16 @@ public DefaultPolarisAuthenticator() { @Inject public DefaultPolarisAuthenticator( - MetaStoreManagerFactory metaStoreManagerFactory, - TokenBrokerFactory tokenBrokerFactory, - CallContext callContext) { - super(metaStoreManagerFactory, callContext); - this.tokenBrokerFactory = tokenBrokerFactory; + PolarisMetaStoreManager metaStoreManager, + PolarisMetaStoreSession metaStoreSession, + TokenBroker tokenBroker) { + super(metaStoreManager, metaStoreSession); + this.tokenBroker = tokenBroker; } @Override public Optional authenticate(String credentials) { - TokenBroker handler = tokenBrokerFactory.apply(callContext.getRealmContext()); - DecodedToken decodedToken = handler.verify(credentials); + DecodedToken decodedToken = tokenBroker.verify(credentials); return getPrincipal(decodedToken); } } diff --git a/service/common/src/main/java/org/apache/polaris/service/auth/JWTBroker.java b/service/common/src/main/java/org/apache/polaris/service/auth/JWTBroker.java index f1945e3c4..cd4b24d13 100644 --- a/service/common/src/main/java/org/apache/polaris/service/auth/JWTBroker.java +++ b/service/common/src/main/java/org/apache/polaris/service/auth/JWTBroker.java @@ -29,11 +29,10 @@ import java.util.UUID; import org.apache.commons.lang3.StringUtils; import org.apache.iceberg.exceptions.NotAuthorizedException; -import org.apache.polaris.core.PolarisCallContext; -import org.apache.polaris.core.context.CallContext; import org.apache.polaris.core.entity.PolarisEntityType; import org.apache.polaris.core.entity.PrincipalEntity; import org.apache.polaris.core.persistence.PolarisMetaStoreManager; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; import org.apache.polaris.service.types.TokenType; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -49,10 +48,15 @@ public abstract class JWTBroker implements TokenBroker { private static final String CLAIM_KEY_SCOPE = "scope"; private final PolarisMetaStoreManager metaStoreManager; + private final PolarisMetaStoreSession metaStoreSession; private final int maxTokenGenerationInSeconds; - JWTBroker(PolarisMetaStoreManager metaStoreManager, int maxTokenGenerationInSeconds) { + JWTBroker( + PolarisMetaStoreManager metaStoreManager, + PolarisMetaStoreSession metaStoreSession, + int maxTokenGenerationInSeconds) { this.metaStoreManager = metaStoreManager; + this.metaStoreSession = metaStoreSession; this.maxTokenGenerationInSeconds = maxTokenGenerationInSeconds; } @@ -103,10 +107,7 @@ public TokenResponse generateFromToken( } DecodedToken decodedToken = verify(subjectToken); PolarisMetaStoreManager.EntityResult principalLookup = - metaStoreManager.loadEntity( - CallContext.getCurrentContext().getPolarisCallContext(), - 0L, - decodedToken.getPrincipalId()); + metaStoreManager.loadEntity(metaStoreSession, 0L, decodedToken.getPrincipalId()); if (!principalLookup.isSuccess() || principalLookup.getEntity().getType() != PolarisEntityType.PRINCIPAL) { return new TokenResponse(OAuthTokenErrorResponse.Error.unauthorized_client); @@ -120,11 +121,7 @@ public TokenResponse generateFromToken( @Override public TokenResponse generateFromClientSecrets( - String clientId, - String clientSecret, - String grantType, - String scope, - PolarisCallContext polarisCallContext) { + String clientId, String clientSecret, String grantType, String scope) { // Initial sanity checks TokenRequestValidator validator = new TokenRequestValidator(); Optional initialValidationResponse = @@ -134,8 +131,7 @@ public TokenResponse generateFromClientSecrets( } Optional principal = - TokenBroker.findPrincipalEntity( - metaStoreManager, clientId, clientSecret, polarisCallContext); + TokenBroker.findPrincipalEntity(metaStoreManager, metaStoreSession, clientId, clientSecret); if (principal.isEmpty()) { return new TokenResponse(OAuthTokenErrorResponse.Error.unauthorized_client); } diff --git a/service/common/src/main/java/org/apache/polaris/service/auth/JWTRSAKeyPair.java b/service/common/src/main/java/org/apache/polaris/service/auth/JWTRSAKeyPair.java index 18f270238..6c3dfc3cc 100644 --- a/service/common/src/main/java/org/apache/polaris/service/auth/JWTRSAKeyPair.java +++ b/service/common/src/main/java/org/apache/polaris/service/auth/JWTRSAKeyPair.java @@ -23,6 +23,7 @@ import java.security.interfaces.RSAPrivateKey; import java.security.interfaces.RSAPublicKey; import org.apache.polaris.core.persistence.PolarisMetaStoreManager; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; /** Generates a JWT using a Public/Private RSA Key */ public class JWTRSAKeyPair extends JWTBroker { @@ -31,10 +32,11 @@ public class JWTRSAKeyPair extends JWTBroker { public JWTRSAKeyPair( PolarisMetaStoreManager metaStoreManager, + PolarisMetaStoreSession metaStoreSession, int maxTokenGenerationInSeconds, Path publicKeyFile, Path privateKeyFile) { - super(metaStoreManager, maxTokenGenerationInSeconds); + super(metaStoreManager, metaStoreSession, maxTokenGenerationInSeconds); keyProvider = new LocalRSAKeyProvider(publicKeyFile, privateKeyFile); } diff --git a/service/common/src/main/java/org/apache/polaris/service/auth/JWTRSAKeyPairFactory.java b/service/common/src/main/java/org/apache/polaris/service/auth/JWTRSAKeyPairFactory.java index d9537b95d..89e43caba 100644 --- a/service/common/src/main/java/org/apache/polaris/service/auth/JWTRSAKeyPairFactory.java +++ b/service/common/src/main/java/org/apache/polaris/service/auth/JWTRSAKeyPairFactory.java @@ -56,6 +56,7 @@ public JWTRSAKeyPairFactory( public TokenBroker apply(RealmContext realmContext) { return new JWTRSAKeyPair( metaStoreManagerFactory.getOrCreateMetaStoreManager(realmContext), + metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(), (int) tokenBrokerConfiguration.maxTokenGeneration().toSeconds(), keyPairConfiguration.publicKeyFile(), keyPairConfiguration.privateKeyFile()); diff --git a/service/common/src/main/java/org/apache/polaris/service/auth/JWTSymmetricKeyBroker.java b/service/common/src/main/java/org/apache/polaris/service/auth/JWTSymmetricKeyBroker.java index 16c9e1551..239c75d8a 100644 --- a/service/common/src/main/java/org/apache/polaris/service/auth/JWTSymmetricKeyBroker.java +++ b/service/common/src/main/java/org/apache/polaris/service/auth/JWTSymmetricKeyBroker.java @@ -21,6 +21,7 @@ import com.auth0.jwt.algorithms.Algorithm; import java.util.function.Supplier; import org.apache.polaris.core.persistence.PolarisMetaStoreManager; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; /** Generates a JWT using a Symmetric Key. */ public class JWTSymmetricKeyBroker extends JWTBroker { @@ -28,9 +29,10 @@ public class JWTSymmetricKeyBroker extends JWTBroker { public JWTSymmetricKeyBroker( PolarisMetaStoreManager metaStoreManager, + PolarisMetaStoreSession metaStoreSession, int maxTokenGenerationInSeconds, Supplier secretSupplier) { - super(metaStoreManager, maxTokenGenerationInSeconds); + super(metaStoreManager, metaStoreSession, maxTokenGenerationInSeconds); this.secretSupplier = secretSupplier; } diff --git a/service/common/src/main/java/org/apache/polaris/service/auth/JWTSymmetricKeyFactory.java b/service/common/src/main/java/org/apache/polaris/service/auth/JWTSymmetricKeyFactory.java index d08754ef0..38cdb3663 100644 --- a/service/common/src/main/java/org/apache/polaris/service/auth/JWTSymmetricKeyFactory.java +++ b/service/common/src/main/java/org/apache/polaris/service/auth/JWTSymmetricKeyFactory.java @@ -61,6 +61,7 @@ public JWTSymmetricKeyFactory( public TokenBroker apply(RealmContext realmContext) { return new JWTSymmetricKeyBroker( metaStoreManagerFactory.getOrCreateMetaStoreManager(realmContext), + metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(), (int) maxTokenGeneration.toSeconds(), secretSupplier); } diff --git a/service/common/src/main/java/org/apache/polaris/service/auth/NoneTokenBrokerFactory.java b/service/common/src/main/java/org/apache/polaris/service/auth/NoneTokenBrokerFactory.java index 311175176..9f642a2b8 100644 --- a/service/common/src/main/java/org/apache/polaris/service/auth/NoneTokenBrokerFactory.java +++ b/service/common/src/main/java/org/apache/polaris/service/auth/NoneTokenBrokerFactory.java @@ -20,7 +20,6 @@ import io.smallrye.common.annotation.Identifier; import jakarta.enterprise.context.ApplicationScoped; -import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.service.types.TokenType; @@ -43,11 +42,7 @@ public boolean supportsRequestedTokenType(TokenType tokenType) { @Override public TokenResponse generateFromClientSecrets( - String clientId, - String clientSecret, - String grantType, - String scope, - PolarisCallContext polarisCallContext) { + String clientId, String clientSecret, String grantType, String scope) { return null; } diff --git a/service/common/src/main/java/org/apache/polaris/service/auth/TestInlineBearerTokenPolarisAuthenticator.java b/service/common/src/main/java/org/apache/polaris/service/auth/TestInlineBearerTokenPolarisAuthenticator.java index cd4b914aa..36e8df845 100644 --- a/service/common/src/main/java/org/apache/polaris/service/auth/TestInlineBearerTokenPolarisAuthenticator.java +++ b/service/common/src/main/java/org/apache/polaris/service/auth/TestInlineBearerTokenPolarisAuthenticator.java @@ -28,10 +28,9 @@ import java.util.Optional; import java.util.stream.Collectors; import org.apache.polaris.core.auth.AuthenticatedPolarisPrincipal; -import org.apache.polaris.core.context.CallContext; import org.apache.polaris.core.entity.PolarisPrincipalSecrets; -import org.apache.polaris.core.persistence.MetaStoreManagerFactory; import org.apache.polaris.core.persistence.PolarisMetaStoreManager; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -59,15 +58,13 @@ public TestInlineBearerTokenPolarisAuthenticator() { @Inject public TestInlineBearerTokenPolarisAuthenticator( - MetaStoreManagerFactory metaStoreManagerFactory, CallContext callContext) { - super(metaStoreManagerFactory, callContext); + PolarisMetaStoreManager metaStoreManager, PolarisMetaStoreSession metaStoreSession) { + super(metaStoreManager, metaStoreSession); } @Override public Optional authenticate(String credentials) { Map properties = extractPrincipal(credentials); - PolarisMetaStoreManager metaStoreManager = - metaStoreManagerFactory.getOrCreateMetaStoreManager(callContext.getRealmContext()); String principal = properties.get("principal"); LOGGER.info("Checking for existence of principal {} in map {}", principal, properties); @@ -82,9 +79,7 @@ public Optional authenticate(String credentials) } PolarisPrincipalSecrets secrets = - metaStoreManager - .loadPrincipalSecrets(callContext.getPolarisCallContext(), principal) - .getPrincipalSecrets(); + metaStoreManager.loadPrincipalSecrets(metaStoreSession, principal).getPrincipalSecrets(); if (secrets == null) { // For test scenarios, if we're allowing short-circuiting into the bearer flow, there may // not be a clientId/clientSecret, and instead we'll let the BasePolarisAuthenticator diff --git a/service/common/src/main/java/org/apache/polaris/service/auth/TestOAuth2ApiService.java b/service/common/src/main/java/org/apache/polaris/service/auth/TestOAuth2ApiService.java index 910defbfa..2f53c20e5 100644 --- a/service/common/src/main/java/org/apache/polaris/service/auth/TestOAuth2ApiService.java +++ b/service/common/src/main/java/org/apache/polaris/service/auth/TestOAuth2ApiService.java @@ -27,14 +27,12 @@ import java.util.Map; import java.util.Objects; import org.apache.iceberg.exceptions.NotAuthorizedException; -import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.auth.PolarisSecretsManager.PrincipalSecretsResult; -import org.apache.polaris.core.context.CallContext; import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.PolarisEntitySubType; import org.apache.polaris.core.entity.PolarisEntityType; -import org.apache.polaris.core.persistence.MetaStoreManagerFactory; import org.apache.polaris.core.persistence.PolarisMetaStoreManager; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; import org.apache.polaris.service.catalog.api.IcebergRestOAuth2ApiService; import org.apache.polaris.service.types.TokenType; import org.slf4j.Logger; @@ -45,8 +43,8 @@ public class TestOAuth2ApiService implements IcebergRestOAuth2ApiService { private static final Logger LOGGER = LoggerFactory.getLogger(TestOAuth2ApiService.class); - @Inject MetaStoreManagerFactory metaStoreManagerFactory; - @Inject CallContext callContext; + @Inject PolarisMetaStoreManager metaStoreManager; + @Inject PolarisMetaStoreSession metaStoreSession; @Override public Response getToken( @@ -63,7 +61,7 @@ public Response getToken( RealmContext realmContext, SecurityContext securityContext) { Map response = new HashMap<>(); - String principalName = getPrincipalName(clientId, realmContext); + String principalName = getPrincipalName(clientId); response.put( "access_token", "principal:" @@ -80,17 +78,14 @@ public Response getToken( return Response.ok(response).build(); } - private String getPrincipalName(String clientId, RealmContext realmContext) { - PolarisMetaStoreManager metaStoreManager = - metaStoreManagerFactory.getOrCreateMetaStoreManager(realmContext); - PolarisCallContext polarisCallContext = callContext.getPolarisCallContext(); + private String getPrincipalName(String clientId) { PrincipalSecretsResult secretsResult = - metaStoreManager.loadPrincipalSecrets(polarisCallContext, clientId); + metaStoreManager.loadPrincipalSecrets(metaStoreSession, clientId); if (secretsResult.isSuccess()) { LOGGER.debug("Found principal secrets for client id {}", clientId); PolarisMetaStoreManager.EntityResult principalResult = metaStoreManager.loadEntity( - polarisCallContext, 0L, secretsResult.getPrincipalSecrets().getPrincipalId()); + metaStoreSession, 0L, secretsResult.getPrincipalSecrets().getPrincipalId()); if (!principalResult.isSuccess()) { throw new NotAuthorizedException("Failed to load principal entity"); } @@ -100,7 +95,7 @@ private String getPrincipalName(String clientId, RealmContext realmContext) { "Unable to find principal secrets for client id {} - trying as principal name", clientId); PolarisMetaStoreManager.EntityResult principalResult = metaStoreManager.readEntityByName( - polarisCallContext, + metaStoreSession, null, PolarisEntityType.PRINCIPAL, PolarisEntitySubType.NULL_SUBTYPE, diff --git a/service/common/src/main/java/org/apache/polaris/service/auth/TokenBroker.java b/service/common/src/main/java/org/apache/polaris/service/auth/TokenBroker.java index 2ace9bb92..31b6c4009 100644 --- a/service/common/src/main/java/org/apache/polaris/service/auth/TokenBroker.java +++ b/service/common/src/main/java/org/apache/polaris/service/auth/TokenBroker.java @@ -20,11 +20,11 @@ import jakarta.annotation.Nonnull; import java.util.Optional; -import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.auth.PolarisSecretsManager.PrincipalSecretsResult; import org.apache.polaris.core.entity.PolarisEntityType; import org.apache.polaris.core.entity.PrincipalEntity; import org.apache.polaris.core.persistence.PolarisMetaStoreManager; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; import org.apache.polaris.service.types.TokenType; /** Generic token class intended to be extended by different token types */ @@ -35,11 +35,7 @@ public interface TokenBroker { boolean supportsRequestedTokenType(TokenType tokenType); TokenResponse generateFromClientSecrets( - final String clientId, - final String clientSecret, - final String grantType, - final String scope, - PolarisCallContext polarisCallContext); + final String clientId, final String clientSecret, final String grantType, final String scope); TokenResponse generateFromToken( TokenType tokenType, String subjectToken, final String grantType, final String scope); @@ -48,12 +44,12 @@ TokenResponse generateFromToken( static @Nonnull Optional findPrincipalEntity( PolarisMetaStoreManager metaStoreManager, + PolarisMetaStoreSession metaStoreSession, String clientId, - String clientSecret, - PolarisCallContext polarisCallContext) { + String clientSecret) { // Validate the principal is present and secrets match PrincipalSecretsResult principalSecrets = - metaStoreManager.loadPrincipalSecrets(polarisCallContext, clientId); + metaStoreManager.loadPrincipalSecrets(metaStoreSession, clientId); if (!principalSecrets.isSuccess()) { return Optional.empty(); } @@ -62,7 +58,7 @@ TokenResponse generateFromToken( } PolarisMetaStoreManager.EntityResult result = metaStoreManager.loadEntity( - polarisCallContext, 0L, principalSecrets.getPrincipalSecrets().getPrincipalId()); + metaStoreSession, 0L, principalSecrets.getPrincipalSecrets().getPrincipalId()); if (!result.isSuccess() || result.getEntity().getType() != PolarisEntityType.PRINCIPAL) { return Optional.empty(); } diff --git a/service/common/src/main/java/org/apache/polaris/service/catalog/BasePolarisCatalog.java b/service/common/src/main/java/org/apache/polaris/service/catalog/BasePolarisCatalog.java index 4b20d95f8..26458892b 100644 --- a/service/common/src/main/java/org/apache/polaris/service/catalog/BasePolarisCatalog.java +++ b/service/common/src/main/java/org/apache/polaris/service/catalog/BasePolarisCatalog.java @@ -24,7 +24,6 @@ import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableMap; import jakarta.annotation.Nonnull; -import jakarta.annotation.Nullable; import jakarta.ws.rs.core.SecurityContext; import java.io.Closeable; import java.io.IOException; @@ -74,12 +73,13 @@ import org.apache.iceberg.view.ViewMetadataParser; import org.apache.iceberg.view.ViewOperations; import org.apache.iceberg.view.ViewUtil; -import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.PolarisConfiguration; +import org.apache.polaris.core.PolarisConfigurationStore; +import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.admin.model.StorageConfigInfo; import org.apache.polaris.core.auth.AuthenticatedPolarisPrincipal; import org.apache.polaris.core.catalog.PolarisCatalogHelpers; -import org.apache.polaris.core.context.CallContext; +import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.CatalogEntity; import org.apache.polaris.core.entity.NamespaceEntity; import org.apache.polaris.core.entity.PolarisEntity; @@ -91,6 +91,7 @@ import org.apache.polaris.core.persistence.BaseResult; import org.apache.polaris.core.persistence.PolarisEntityManager; import org.apache.polaris.core.persistence.PolarisMetaStoreManager; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; import org.apache.polaris.core.persistence.PolarisResolvedPathWrapper; import org.apache.polaris.core.persistence.resolver.PolarisResolutionManifest; import org.apache.polaris.core.persistence.resolver.PolarisResolutionManifestCatalogView; @@ -149,52 +150,62 @@ public class BasePolarisCatalog extends BaseMetastoreViewCatalog }; private final PolarisEntityManager entityManager; - private final CallContext callContext; + private PolarisMetaStoreManager metaStoreManager; + private final PolarisMetaStoreSession metaStoreSession; + private final PolarisConfigurationStore configurationStore; + private final PolarisDiagnostics diagnostics; + private final RealmContext realmContext; private final PolarisResolutionManifestCatalogView resolvedEntityView; private final CatalogEntity catalogEntity; private final TaskExecutor taskExecutor; private final SecurityContext securityContext; private final AuthenticatedPolarisPrincipal authenticatedPrincipal; + private final FileIOFactory fileIOFactory; + private final long catalogId; + private final String catalogName; + private String ioImplClassName; private FileIO catalogFileIO; - private final String catalogName; - private long catalogId = -1; private String defaultBaseLocation; private CloseableGroup closeableGroup; private Map catalogProperties; private Map tableDefaultProperties; - private FileIOFactory fileIOFactory; - private PolarisMetaStoreManager metaStoreManager; /** + * @param realmContext the current RealmContext * @param entityManager provides handle to underlying PolarisMetaStoreManager with which to * perform mutations on entities. - * @param callContext the current CallContext * @param resolvedEntityView accessor to resolved entity paths that have been pre-vetted to ensure * this catalog instance only interacts with authorized resolved paths. * @param taskExecutor Executor we use to register cleanup task handlers */ public BasePolarisCatalog( + RealmContext realmContext, PolarisEntityManager entityManager, PolarisMetaStoreManager metaStoreManager, - CallContext callContext, + PolarisMetaStoreSession metaStoreSession, + PolarisConfigurationStore configurationStore, + PolarisDiagnostics diagnostics, PolarisResolutionManifestCatalogView resolvedEntityView, SecurityContext securityContext, - AuthenticatedPolarisPrincipal authenticatedPrincipal, TaskExecutor taskExecutor, FileIOFactory fileIOFactory) { + this.realmContext = realmContext; this.entityManager = entityManager; - this.callContext = callContext; + this.metaStoreManager = metaStoreManager; + this.metaStoreSession = metaStoreSession; + this.configurationStore = configurationStore; + this.diagnostics = diagnostics; this.resolvedEntityView = resolvedEntityView; + this.taskExecutor = taskExecutor; + this.fileIOFactory = fileIOFactory; this.catalogEntity = CatalogEntity.of(resolvedEntityView.getResolvedReferenceCatalogEntity().getRawLeafEntity()); this.securityContext = securityContext; - this.authenticatedPrincipal = authenticatedPrincipal; - this.taskExecutor = taskExecutor; + this.authenticatedPrincipal = + (AuthenticatedPolarisPrincipal) securityContext.getUserPrincipal(); this.catalogId = catalogEntity.getId(); this.catalogName = catalogEntity.getName(); - this.fileIOFactory = fileIOFactory; - this.metaStoreManager = metaStoreManager; } @Override @@ -256,7 +267,6 @@ public void initialize(String name, Map properties) { CatalogProperties.FILE_IO_IMPL); } } - callContext.closeables().addCloseable(this); this.closeableGroup = new CloseableGroup(); closeableGroup.addCloseable(metricsReporter()); closeableGroup.setSuppressCloseFailure(true); @@ -356,7 +366,7 @@ protected String defaultWarehouseLocation(TableIdentifier tableIdentifier) { "Namespace does not exist: %s", tableIdentifier.namespace()); } List namespacePath = resolvedNamespace.getRawFullPath(); - String namespaceLocation = resolveLocationForPath(callContext, namespacePath); + String namespaceLocation = resolveLocationForPath(namespacePath); return SLASH.join(namespaceLocation, tableIdentifier.name()); } } @@ -440,7 +450,7 @@ public boolean dropTable(TableIdentifier tableIdentifier, boolean purge) { "Scheduled cleanup task {} for table {}", dropEntityResult.getCleanupTaskId(), tableIdentifier); - taskExecutor.addTaskHandlerContext(dropEntityResult.getCleanupTaskId(), callContext); + taskExecutor.addTaskHandlerContext(dropEntityResult.getCleanupTaskId(), realmContext); } return true; @@ -497,18 +507,14 @@ private void createNamespaceInternal( NamespaceEntity entity = new NamespaceEntity.Builder(namespace) .setCatalogId(getCatalogId()) - .setId(getMetaStoreManager().generateNewEntityId(getCurrentPolarisContext()).getId()) + .setId(getMetaStoreManager().generateNewEntityId(metaStoreSession).getId()) .setParentId(resolvedParent.getRawLeafEntity().getId()) .setProperties(metadata) .setCreateTimestamp(System.currentTimeMillis()) .setBaseLocation(baseLocation) .build(); - if (!callContext - .getPolarisCallContext() - .getConfigurationStore() - .getConfiguration( - callContext.getPolarisCallContext(), - PolarisConfiguration.ALLOW_NAMESPACE_LOCATION_OVERLAP)) { + if (!configurationStore.getConfiguration( + realmContext, PolarisConfiguration.ALLOW_NAMESPACE_LOCATION_OVERLAP)) { LOGGER.debug("Validating no overlap for {} with sibling tables or namespaces", namespace); validateNoLocationOverlap( entity.getBaseLocation(), resolvedParent.getRawFullPath(), entity.getName()); @@ -519,7 +525,7 @@ private void createNamespaceInternal( PolarisEntity.of( getMetaStoreManager() .createEntityIfNotExists( - getCurrentPolarisContext(), + metaStoreSession, PolarisEntity.toCoreList(resolvedParent.getRawFullPath()), entity)); if (returnedEntity == null) { @@ -537,14 +543,13 @@ private String resolveNamespaceLocation(Namespace namespace, Map ? getResolvedParentNamespace(namespace).getRawFullPath() : List.of(resolvedEntityView.getResolvedReferenceCatalogEntity().getRawLeafEntity()); - String parentLocation = resolveLocationForPath(callContext, parentPath); + String parentLocation = resolveLocationForPath(parentPath); return parentLocation + "/" + namespace.level(namespace.length() - 1); } } - private static @Nonnull String resolveLocationForPath( - @Nonnull CallContext callContext, List parentPath) { + private @Nonnull String resolveLocationForPath(List parentPath) { // always take the first object. If it has the base-location, stop there AtomicBoolean foundBaseLocation = new AtomicBoolean(false); return parentPath.reversed().stream() @@ -557,24 +562,20 @@ private String resolveNamespaceLocation(Namespace namespace, Map .toList() .reversed() .stream() - .map(entity -> baseLocation(callContext, entity)) + .map(this::baseLocation) .map(BasePolarisCatalog::stripLeadingTrailingSlash) .collect(Collectors.joining("/")); } - private static @Nullable String baseLocation( - @Nonnull CallContext callContext, PolarisEntity entity) { + private @Nonnull String baseLocation(PolarisEntity entity) { if (entity.getType().equals(PolarisEntityType.CATALOG)) { CatalogEntity catEntity = CatalogEntity.of(entity); String catalogDefaultBaseLocation = catEntity.getDefaultBaseLocation(); - callContext - .getPolarisCallContext() - .getDiagServices() - .checkNotNull( - catalogDefaultBaseLocation, - "Tried to resolve location with catalog with null default base location", - "catalog = {}", - catEntity); + diagnostics.checkNotNull( + catalogDefaultBaseLocation, + "Tried to resolve location with catalog with null default base location", + "catalog = {}", + catEntity); return catalogDefaultBaseLocation; } else { String baseLocation = @@ -583,14 +584,11 @@ private String resolveNamespaceLocation(Namespace namespace, Map return baseLocation; } else { String entityName = entity.getName(); - callContext - .getPolarisCallContext() - .getDiagServices() - .checkNotNull( - entityName, - "Tried to resolve location with entity without base location or name", - "entity = {}", - entity); + diagnostics.checkNotNull( + entityName, + "Tried to resolve location with entity without base location or name", + "entity = {}", + entity); return entityName; } } @@ -633,18 +631,15 @@ public boolean dropNamespace(Namespace namespace) throws NamespaceNotEmptyExcept PolarisEntity leafEntity = resolvedEntities.getRawLeafEntity(); // drop if exists and is empty - PolarisCallContext polarisCallContext = callContext.getPolarisCallContext(); PolarisMetaStoreManager.DropEntityResult dropEntityResult = getMetaStoreManager() .dropEntityIfExists( - getCurrentPolarisContext(), + metaStoreSession, PolarisEntity.toCoreList(catalogPath), leafEntity, Map.of(), - polarisCallContext - .getConfigurationStore() - .getConfiguration( - polarisCallContext, PolarisConfiguration.CLEANUP_ON_NAMESPACE_DROP)); + configurationStore.getConfiguration( + realmContext, PolarisConfiguration.CLEANUP_ON_NAMESPACE_DROP)); if (!dropEntityResult.isSuccess() && dropEntityResult.failedBecauseNotEmpty()) { throw new NamespaceNotEmptyException("Namespace %s is not empty", namespace); @@ -669,12 +664,8 @@ public boolean setProperties(Namespace namespace, Map properties PolarisEntity updatedEntity = new PolarisEntity.Builder(entity).setProperties(newProperties).build(); - if (!callContext - .getPolarisCallContext() - .getConfigurationStore() - .getConfiguration( - callContext.getPolarisCallContext(), - PolarisConfiguration.ALLOW_NAMESPACE_LOCATION_OVERLAP)) { + if (!configurationStore.getConfiguration( + realmContext, PolarisConfiguration.ALLOW_NAMESPACE_LOCATION_OVERLAP)) { LOGGER.debug("Validating no overlap with sibling tables or namespaces"); validateNoLocationOverlap( NamespaceEntity.of(updatedEntity).getBaseLocation(), @@ -689,9 +680,7 @@ public boolean setProperties(Namespace namespace, Map properties Optional.ofNullable( getMetaStoreManager() .updateEntityPropertiesIfNotChanged( - getCurrentPolarisContext(), - PolarisEntity.toCoreList(parentPath), - updatedEntity) + metaStoreSession, PolarisEntity.toCoreList(parentPath), updatedEntity) .getEntity()) .map(PolarisEntity::new) .orElse(null); @@ -721,9 +710,7 @@ public boolean removeProperties(Namespace namespace, Set properties) Optional.ofNullable( getMetaStoreManager() .updateEntityPropertiesIfNotChanged( - getCurrentPolarisContext(), - PolarisEntity.toCoreList(parentPath), - updatedEntity) + metaStoreSession, PolarisEntity.toCoreList(parentPath), updatedEntity) .getEntity()) .map(PolarisEntity::new) .orElse(null); @@ -767,7 +754,7 @@ public List listNamespaces(Namespace namespace) throws NoSuchNamespac PolarisEntity.toNameAndIdList( getMetaStoreManager() .listEntities( - getCurrentPolarisContext(), + metaStoreSession, PolarisEntity.toCoreList(catalogPath), PolarisEntityType.NAMESPACE, PolarisEntitySubType.NULL_SUBTYPE) @@ -912,7 +899,7 @@ private Map refreshCredentials( .getCredentialCache() .getOrGenerateSubScopeCreds( getCredentialVendor(), - callContext.getPolarisCallContext(), + metaStoreSession, entity, allowList, tableLocations, @@ -966,15 +953,18 @@ private void validateLocationsForTableLike( PolarisResolvedPathWrapper resolvedStorageEntity) { Optional optStorageConfiguration = PolarisStorageConfigurationInfo.forEntityPath( - callContext.getPolarisCallContext().getDiagServices(), - resolvedStorageEntity.getRawFullPath()); + realmContext, configurationStore, diagnostics, resolvedStorageEntity.getRawFullPath()); optStorageConfiguration.ifPresentOrElse( storageConfigInfo -> { Map> validationResults = InMemoryStorageIntegration.validateSubpathsOfAllowedLocations( - storageConfigInfo, Set.of(PolarisStorageActions.ALL), locations); + realmContext, + configurationStore, + storageConfigInfo, + Set.of(PolarisStorageActions.ALL), + locations); validationResults .values() .forEach( @@ -1016,12 +1006,8 @@ private void validateLocationsForTableLike( }, () -> { List allowedStorageTypes = - callContext - .getPolarisCallContext() - .getConfigurationStore() - .getConfiguration( - callContext.getPolarisCallContext(), - PolarisConfiguration.SUPPORTED_CATALOG_STORAGE_TYPES); + configurationStore.getConfiguration( + realmContext, PolarisConfiguration.SUPPORTED_CATALOG_STORAGE_TYPES); if (!allowedStorageTypes.contains(StorageConfigInfo.StorageTypeEnum.FILE.name())) { List invalidLocations = locations.stream() @@ -1045,13 +1031,8 @@ private void validateNoLocationOverlap( TableIdentifier identifier, List resolvedNamespace, String location) { - if (callContext - .getPolarisCallContext() - .getConfigurationStore() - .getConfiguration( - callContext.getPolarisCallContext(), - catalog, - PolarisConfiguration.ALLOW_TABLE_LOCATION_OVERLAP)) { + if (configurationStore.getConfiguration( + realmContext, catalog, PolarisConfiguration.ALLOW_TABLE_LOCATION_OVERLAP)) { LOGGER.debug("Skipping location overlap validation for identifier '{}'", identifier); } else { // if (entity.getSubType().equals(PolarisEntitySubType.TABLE)) { // TODO - is this necessary for views? overlapping views do not expose subdirectories via the @@ -1073,7 +1054,7 @@ private void validateNoLocationOverlap( PolarisMetaStoreManager.ListEntitiesResult siblingNamespacesResult = getMetaStoreManager() .listEntities( - callContext.getPolarisCallContext(), + metaStoreSession, parentPath.stream().map(PolarisEntity::toCore).collect(Collectors.toList()), PolarisEntityType.NAMESPACE, PolarisEntitySubType.ANY_SUBTYPE); @@ -1096,7 +1077,7 @@ private void validateNoLocationOverlap( PolarisMetaStoreManager.ListEntitiesResult siblingTablesResult = getMetaStoreManager() .listEntities( - callContext.getPolarisCallContext(), + metaStoreSession, parentPath.stream() .map(PolarisEntity::toCore) .collect(Collectors.toList()), @@ -1130,7 +1111,11 @@ private void validateNoLocationOverlap( siblingTables.size() + siblingNamespaces.size()); PolarisResolutionManifest resolutionManifest = new PolarisResolutionManifest( - callContext, entityManager, securityContext, parentPath.getFirst().getName()); + metaStoreSession, + diagnostics, + entityManager, + securityContext, + parentPath.getFirst().getName()); siblingTables.forEach( tbl -> resolutionManifest.addPath( @@ -1368,8 +1353,7 @@ public void doCommit(TableMetadata base, TableMetadata metadata) { .setCatalogId(getCatalogId()) .setSubType(PolarisEntitySubType.TABLE) .setBaseLocation(metadata.location()) - .setId( - getMetaStoreManager().generateNewEntityId(getCurrentPolarisContext()).getId()) + .setId(getMetaStoreManager().generateNewEntityId(metaStoreSession).getId()) .build(); } else { existingLocation = entity.getMetadataLocation(); @@ -1413,17 +1397,12 @@ protected String tableName() { private void validateMetadataFileInTableDir( TableIdentifier identifier, TableMetadata metadata, CatalogEntity catalog) { - PolarisCallContext polarisCallContext = callContext.getPolarisCallContext(); boolean allowEscape = - polarisCallContext - .getConfigurationStore() - .getConfiguration( - polarisCallContext, PolarisConfiguration.ALLOW_EXTERNAL_TABLE_LOCATION); + configurationStore.getConfiguration( + realmContext, PolarisConfiguration.ALLOW_EXTERNAL_TABLE_LOCATION); if (!allowEscape - && !polarisCallContext - .getConfigurationStore() - .getConfiguration( - polarisCallContext, PolarisConfiguration.ALLOW_EXTERNAL_METADATA_FILE_LOCATION)) { + && !configurationStore.getConfiguration( + realmContext, PolarisConfiguration.ALLOW_EXTERNAL_METADATA_FILE_LOCATION)) { LOGGER.debug( "Validating base location {} for table {} in metadata file {}", metadata.location(), @@ -1569,8 +1548,7 @@ public void doCommit(ViewMetadata base, ViewMetadata metadata) { new TableLikeEntity.Builder(identifier, newLocation) .setCatalogId(getCatalogId()) .setSubType(PolarisEntitySubType.VIEW) - .setId( - getMetaStoreManager().generateNewEntityId(getCurrentPolarisContext()).getId()) + .setId(getMetaStoreManager().generateNewEntityId(metaStoreSession).getId()) .build(); } else { existingLocation = entity.getMetadataLocation(); @@ -1634,10 +1612,6 @@ private FileIO refreshIOWithCredentials( return fileIO; } - private PolarisCallContext getCurrentPolarisContext() { - return callContext.getPolarisCallContext(); - } - private PolarisMetaStoreManager getMetaStoreManager() { return metaStoreManager; } @@ -1646,11 +1620,6 @@ private PolarisCredentialVendor getCredentialVendor() { return metaStoreManager; } - @VisibleForTesting - public void setFileIOFactory(FileIOFactory newFactory) { - this.fileIOFactory = newFactory; - } - @VisibleForTesting long getCatalogId() { // TODO: Properly handle initialization @@ -1703,7 +1672,7 @@ private void renameTableLike( PolarisMetaStoreManager.EntityResult returnedEntityResult = getMetaStoreManager() .renameEntity( - getCurrentPolarisContext(), + metaStoreSession, PolarisEntity.toCoreList(catalogPath), leafEntity, PolarisEntity.toCoreList(newCatalogPath), @@ -1766,7 +1735,7 @@ private void renameTableLike( .log("Returned entity identifier doesn't match toEntity identifier"); getMetaStoreManager() .updateEntityPropertiesIfNotChanged( - getCurrentPolarisContext(), + metaStoreSession, PolarisEntity.toCoreList(newCatalogPath), new TableLikeEntity.Builder(returnedEntity).setTableIdentifier(to).build()); } @@ -1812,7 +1781,7 @@ private void createTableLike( PolarisEntity.of( getMetaStoreManager() .createEntityIfNotExists( - getCurrentPolarisContext(), PolarisEntity.toCoreList(catalogPath), entity)); + metaStoreSession, PolarisEntity.toCoreList(catalogPath), entity)); LOGGER.debug("Created TableLike entity {} with TableIdentifier {}", entity, identifier); if (returnedEntity == null) { // TODO: Error or retry? @@ -1837,7 +1806,7 @@ private void updateTableLike(TableIdentifier identifier, PolarisEntity entity) { Optional.ofNullable( getMetaStoreManager() .updateEntityPropertiesIfNotChanged( - getCurrentPolarisContext(), PolarisEntity.toCoreList(catalogPath), entity) + metaStoreSession, PolarisEntity.toCoreList(catalogPath), entity) .getEntity()) .map(PolarisEntity::new) .orElse(null); @@ -1866,13 +1835,8 @@ private void updateTableLike(TableIdentifier identifier, PolarisEntity entity) { // Check that purge is enabled, if it is set: if (catalogPath != null && !catalogPath.isEmpty() && purge) { boolean dropWithPurgeEnabled = - callContext - .getPolarisCallContext() - .getConfigurationStore() - .getConfiguration( - callContext.getPolarisCallContext(), - catalogEntity, - PolarisConfiguration.DROP_WITH_PURGE_ENABLED); + configurationStore.getConfiguration( + realmContext, catalogEntity, PolarisConfiguration.DROP_WITH_PURGE_ENABLED); if (!dropWithPurgeEnabled) { throw new ForbiddenException( String.format( @@ -1886,7 +1850,7 @@ private void updateTableLike(TableIdentifier identifier, PolarisEntity entity) { return getMetaStoreManager() .dropEntityIfExists( - getCurrentPolarisContext(), + metaStoreSession, PolarisEntity.toCoreList(catalogPath), leafEntity, storageProperties, @@ -1971,8 +1935,7 @@ private boolean sendNotificationForTableLike( new TableLikeEntity.Builder(tableIdentifier, newLocation) .setCatalogId(getCatalogId()) .setSubType(PolarisEntitySubType.TABLE) - .setId( - getMetaStoreManager().generateNewEntityId(getCurrentPolarisContext()).getId()) + .setId(getMetaStoreManager().generateNewEntityId(metaStoreSession).getId()) .setLastNotificationTimestamp(request.getPayload().getTimestamp()) .build(); } else { @@ -2060,7 +2023,7 @@ private List listTableLike(PolarisEntitySubType subType, Namesp PolarisEntity.toNameAndIdList( getMetaStoreManager() .listEntities( - getCurrentPolarisContext(), + metaStoreSession, PolarisEntity.toCoreList(catalogPath), PolarisEntityType.TABLE_LIKE, subType) @@ -2092,10 +2055,7 @@ private void blockedUserSpecifiedWriteLocation(Map properties) { /** Helper to retrieve dynamic context-based configuration that has a boolean value. */ private Boolean getBooleanContextConfiguration(String configKey, boolean defaultValue) { - return callContext - .getPolarisCallContext() - .getConfigurationStore() - .getConfiguration(callContext.getPolarisCallContext(), configKey, defaultValue); + return configurationStore.getConfiguration(realmContext, configKey, defaultValue); } /** diff --git a/service/common/src/main/java/org/apache/polaris/service/catalog/IcebergCatalogAdapter.java b/service/common/src/main/java/org/apache/polaris/service/catalog/IcebergCatalogAdapter.java index f29a7dc3c..337b1f12f 100644 --- a/service/common/src/main/java/org/apache/polaris/service/catalog/IcebergCatalogAdapter.java +++ b/service/common/src/main/java/org/apache/polaris/service/catalog/IcebergCatalogAdapter.java @@ -27,6 +27,7 @@ import jakarta.enterprise.context.RequestScoped; import jakarta.inject.Inject; import jakarta.ws.rs.core.Response; +import jakarta.ws.rs.core.Response.Status; import jakarta.ws.rs.core.SecurityContext; import java.net.URLEncoder; import java.nio.charset.Charset; @@ -34,6 +35,7 @@ import java.util.Map; import java.util.Optional; import java.util.Set; +import java.util.function.Function; import org.apache.iceberg.catalog.Catalog; import org.apache.iceberg.catalog.Namespace; import org.apache.iceberg.catalog.TableIdentifier; @@ -52,23 +54,26 @@ import org.apache.iceberg.rest.requests.ReportMetricsRequest; import org.apache.iceberg.rest.requests.UpdateNamespacePropertiesRequest; import org.apache.iceberg.rest.responses.ConfigResponse; +import org.apache.polaris.core.PolarisConfigurationStore; +import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.auth.AuthenticatedPolarisPrincipal; import org.apache.polaris.core.auth.PolarisAuthorizer; -import org.apache.polaris.core.context.CallContext; import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.PolarisEntity; -import org.apache.polaris.core.persistence.MetaStoreManagerFactory; import org.apache.polaris.core.persistence.PolarisEntityManager; +import org.apache.polaris.core.persistence.PolarisMetaStoreManager; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; import org.apache.polaris.core.persistence.cache.EntityCacheEntry; import org.apache.polaris.core.persistence.resolver.Resolver; import org.apache.polaris.core.persistence.resolver.ResolverStatus; import org.apache.polaris.service.catalog.api.IcebergRestCatalogApiService; import org.apache.polaris.service.catalog.api.IcebergRestConfigurationApiService; -import org.apache.polaris.service.config.RealmEntityManagerFactory; import org.apache.polaris.service.context.CallContextCatalogFactory; import org.apache.polaris.service.types.CommitTableRequest; import org.apache.polaris.service.types.CommitViewRequest; import org.apache.polaris.service.types.NotificationRequest; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * {@link IcebergRestCatalogApiService} implementation that delegates operations to {@link @@ -79,6 +84,8 @@ public class IcebergCatalogAdapter implements IcebergRestCatalogApiService, IcebergRestConfigurationApiService { + private static final Logger LOGGER = LoggerFactory.getLogger(IcebergCatalogAdapter.class); + private static final Set DEFAULT_ENDPOINTS = ImmutableSet.builder() .add(Endpoint.V1_LIST_NAMESPACES) @@ -111,43 +118,69 @@ public class IcebergCatalogAdapter .add(Endpoint.create("POST", ResourcePaths.V1_TRANSACTIONS_COMMIT)) .build(); - private final CallContext callContext; + private final RealmContext realmContext; private final CallContextCatalogFactory catalogFactory; - private final MetaStoreManagerFactory metaStoreManagerFactory; - private final RealmEntityManagerFactory entityManagerFactory; + private final PolarisMetaStoreManager metaStoreManager; + private final PolarisEntityManager entityManager; + private final PolarisMetaStoreSession session; + private final PolarisConfigurationStore configurationStore; + private final PolarisDiagnostics diagnostics; private final PolarisAuthorizer polarisAuthorizer; @Inject public IcebergCatalogAdapter( - CallContext callContext, + RealmContext realmContext, CallContextCatalogFactory catalogFactory, - RealmEntityManagerFactory entityManagerFactory, - MetaStoreManagerFactory metaStoreManagerFactory, + PolarisEntityManager entityManager, + PolarisMetaStoreManager metaStoreManager, + PolarisMetaStoreSession session, + PolarisConfigurationStore configurationStore, + PolarisDiagnostics diagnostics, PolarisAuthorizer polarisAuthorizer) { - this.callContext = callContext; + this.realmContext = realmContext; this.catalogFactory = catalogFactory; - this.entityManagerFactory = entityManagerFactory; - this.metaStoreManagerFactory = metaStoreManagerFactory; + this.entityManager = entityManager; + this.metaStoreManager = metaStoreManager; + this.session = session; + this.configurationStore = configurationStore; + this.diagnostics = diagnostics; this.polarisAuthorizer = polarisAuthorizer; - // FIXME: This is a hack to set the current context for downstream calls. - CallContext.setCurrentContext(callContext); + } + + /** + * Execute operations on a catalog wrapper and ensure we close the BaseCatalog afterward. This + * will typically ensure the underlying FileIO is closed. + */ + private Response withCatalog( + SecurityContext securityContext, + String catalogName, + Function action) { + try (PolarisCatalogHandlerWrapper wrapper = newHandlerWrapper(securityContext, catalogName)) { + return action.apply(wrapper); + } catch (RuntimeException e) { + LOGGER.error("Error while operating on catalog", e); + throw e; + } catch (Exception e) { + LOGGER.error("Error while operating on catalog", e); + throw new RuntimeException(e); + } } private PolarisCatalogHandlerWrapper newHandlerWrapper( - RealmContext realmContext, SecurityContext securityContext, String catalogName) { + SecurityContext securityContext, String catalogName) { AuthenticatedPolarisPrincipal authenticatedPrincipal = (AuthenticatedPolarisPrincipal) securityContext.getUserPrincipal(); if (authenticatedPrincipal == null) { throw new NotAuthorizedException("Failed to find authenticatedPrincipal in SecurityContext"); } - PolarisEntityManager entityManager = - entityManagerFactory.getOrCreateEntityManager(realmContext); - return new PolarisCatalogHandlerWrapper( - callContext, + realmContext, + session, + configurationStore, + diagnostics, entityManager, - metaStoreManagerFactory.getOrCreateMetaStoreManager(realmContext), + metaStoreManager, securityContext, catalogFactory, catalogName, @@ -160,10 +193,10 @@ public Response createNamespace( CreateNamespaceRequest createNamespaceRequest, RealmContext realmContext, SecurityContext securityContext) { - return Response.ok( - newHandlerWrapper(realmContext, securityContext, prefix) - .createNamespace(createNamespaceRequest)) - .build(); + return withCatalog( + securityContext, + prefix, + catalog -> Response.ok(catalog.createNamespace(createNamespaceRequest)).build()); } @Override @@ -176,19 +209,19 @@ public Response listNamespaces( SecurityContext securityContext) { Optional namespaceOptional = Optional.ofNullable(parent).map(IcebergCatalogAdapter::decodeNamespace); - return Response.ok( - newHandlerWrapper(realmContext, securityContext, prefix) - .listNamespaces(namespaceOptional.orElse(Namespace.of()))) - .build(); + return withCatalog( + securityContext, + prefix, + catalog -> + Response.ok(catalog.listNamespaces(namespaceOptional.orElse(Namespace.of()))).build()); } @Override public Response loadNamespaceMetadata( String prefix, String namespace, RealmContext realmContext, SecurityContext securityContext) { Namespace ns = decodeNamespace(namespace); - return Response.ok( - newHandlerWrapper(realmContext, securityContext, prefix).loadNamespaceMetadata(ns)) - .build(); + return withCatalog( + securityContext, prefix, catalog -> Response.ok(catalog.loadNamespaceMetadata(ns)).build()); } private static Namespace decodeNamespace(String namespace) { @@ -199,16 +232,26 @@ private static Namespace decodeNamespace(String namespace) { public Response namespaceExists( String prefix, String namespace, RealmContext realmContext, SecurityContext securityContext) { Namespace ns = decodeNamespace(namespace); - newHandlerWrapper(realmContext, securityContext, prefix).namespaceExists(ns); - return Response.status(Response.Status.NO_CONTENT).build(); + return withCatalog( + securityContext, + prefix, + catalog -> { + catalog.namespaceExists(ns); + return Response.status(Response.Status.NO_CONTENT).build(); + }); } @Override public Response dropNamespace( String prefix, String namespace, RealmContext realmContext, SecurityContext securityContext) { Namespace ns = decodeNamespace(namespace); - newHandlerWrapper(realmContext, securityContext, prefix).dropNamespace(ns); - return Response.status(Response.Status.NO_CONTENT).build(); + return withCatalog( + securityContext, + prefix, + catalog -> { + catalog.dropNamespace(ns); + return Response.status(Response.Status.NO_CONTENT).build(); + }); } @Override @@ -219,10 +262,12 @@ public Response updateProperties( RealmContext realmContext, SecurityContext securityContext) { Namespace ns = decodeNamespace(namespace); - return Response.ok( - newHandlerWrapper(realmContext, securityContext, prefix) - .updateNamespaceProperties(ns, updateNamespacePropertiesRequest)) - .build(); + return withCatalog( + securityContext, + prefix, + catalog -> + Response.ok(catalog.updateNamespaceProperties(ns, updateNamespacePropertiesRequest)) + .build()); } private EnumSet parseAccessDelegationModes(String accessDelegationMode) { @@ -246,29 +291,25 @@ public Response createTable( EnumSet delegationModes = parseAccessDelegationModes(accessDelegationMode); Namespace ns = decodeNamespace(namespace); - if (createTableRequest.stageCreate()) { - if (delegationModes.isEmpty()) { - return Response.ok( - newHandlerWrapper(realmContext, securityContext, prefix) - .createTableStaged(ns, createTableRequest)) - .build(); - } else { - return Response.ok( - newHandlerWrapper(realmContext, securityContext, prefix) - .createTableStagedWithWriteDelegation(ns, createTableRequest)) - .build(); - } - } else if (delegationModes.isEmpty()) { - return Response.ok( - newHandlerWrapper(realmContext, securityContext, prefix) - .createTableDirect(ns, createTableRequest)) - .build(); - } else { - return Response.ok( - newHandlerWrapper(realmContext, securityContext, prefix) - .createTableDirectWithWriteDelegation(ns, createTableRequest)) - .build(); - } + return withCatalog( + securityContext, + prefix, + catalog -> { + if (createTableRequest.stageCreate()) { + if (delegationModes.isEmpty()) { + return Response.ok(catalog.createTableStaged(ns, createTableRequest)).build(); + } else { + return Response.ok( + catalog.createTableStagedWithWriteDelegation(ns, createTableRequest)) + .build(); + } + } else if (delegationModes.isEmpty()) { + return Response.ok(catalog.createTableDirect(ns, createTableRequest)).build(); + } else { + return Response.ok(catalog.createTableDirectWithWriteDelegation(ns, createTableRequest)) + .build(); + } + }); } @Override @@ -280,8 +321,8 @@ public Response listTables( RealmContext realmContext, SecurityContext securityContext) { Namespace ns = decodeNamespace(namespace); - return Response.ok(newHandlerWrapper(realmContext, securityContext, prefix).listTables(ns)) - .build(); + return withCatalog( + securityContext, prefix, catalog -> Response.ok(catalog.listTables(ns)).build()); } @Override @@ -297,17 +338,17 @@ public Response loadTable( parseAccessDelegationModes(accessDelegationMode); Namespace ns = decodeNamespace(namespace); TableIdentifier tableIdentifier = TableIdentifier.of(ns, RESTUtil.decodeString(table)); - if (delegationModes.isEmpty()) { - return Response.ok( - newHandlerWrapper(realmContext, securityContext, prefix) - .loadTable(tableIdentifier, snapshots)) - .build(); - } else { - return Response.ok( - newHandlerWrapper(realmContext, securityContext, prefix) - .loadTableWithAccessDelegation(tableIdentifier, snapshots)) - .build(); - } + return withCatalog( + securityContext, + prefix, + catalog -> { + if (delegationModes.isEmpty()) { + return Response.ok(catalog.loadTable(tableIdentifier, snapshots)).build(); + } else { + return Response.ok(catalog.loadTableWithAccessDelegation(tableIdentifier, snapshots)) + .build(); + } + }); } @Override @@ -319,8 +360,13 @@ public Response tableExists( SecurityContext securityContext) { Namespace ns = decodeNamespace(namespace); TableIdentifier tableIdentifier = TableIdentifier.of(ns, RESTUtil.decodeString(table)); - newHandlerWrapper(realmContext, securityContext, prefix).tableExists(tableIdentifier); - return Response.status(Response.Status.NO_CONTENT).build(); + return withCatalog( + securityContext, + prefix, + catalog -> { + catalog.tableExists(tableIdentifier); + return Response.status(Status.NO_CONTENT).build(); + }); } @Override @@ -333,14 +379,17 @@ public Response dropTable( SecurityContext securityContext) { Namespace ns = decodeNamespace(namespace); TableIdentifier tableIdentifier = TableIdentifier.of(ns, RESTUtil.decodeString(table)); - - if (purgeRequested != null && purgeRequested) { - newHandlerWrapper(realmContext, securityContext, prefix).dropTableWithPurge(tableIdentifier); - } else { - newHandlerWrapper(realmContext, securityContext, prefix) - .dropTableWithoutPurge(tableIdentifier); - } - return Response.status(Response.Status.NO_CONTENT).build(); + return withCatalog( + securityContext, + prefix, + catalog -> { + if (purgeRequested != null && purgeRequested) { + catalog.dropTableWithPurge(tableIdentifier); + } else { + catalog.dropTableWithoutPurge(tableIdentifier); + } + return Response.status(Response.Status.NO_CONTENT).build(); + }); } @Override @@ -351,10 +400,10 @@ public Response registerTable( RealmContext realmContext, SecurityContext securityContext) { Namespace ns = decodeNamespace(namespace); - return Response.ok( - newHandlerWrapper(realmContext, securityContext, prefix) - .registerTable(ns, registerTableRequest)) - .build(); + return withCatalog( + securityContext, + prefix, + catalog -> Response.ok(catalog.registerTable(ns, registerTableRequest)).build()); } @Override @@ -363,8 +412,13 @@ public Response renameTable( RenameTableRequest renameTableRequest, RealmContext realmContext, SecurityContext securityContext) { - newHandlerWrapper(realmContext, securityContext, prefix).renameTable(renameTableRequest); - return Response.ok(javax.ws.rs.core.Response.Status.NO_CONTENT).build(); + return withCatalog( + securityContext, + prefix, + catalog -> { + catalog.renameTable(renameTableRequest); + return Response.ok(Response.Status.NO_CONTENT).build(); + }); } @Override @@ -377,18 +431,18 @@ public Response updateTable( SecurityContext securityContext) { Namespace ns = decodeNamespace(namespace); TableIdentifier tableIdentifier = TableIdentifier.of(ns, RESTUtil.decodeString(table)); - - if (PolarisCatalogHandlerWrapper.isCreate(commitTableRequest)) { - return Response.ok( - newHandlerWrapper(realmContext, securityContext, prefix) - .updateTableForStagedCreate(tableIdentifier, commitTableRequest)) - .build(); - } else { - return Response.ok( - newHandlerWrapper(realmContext, securityContext, prefix) - .updateTable(tableIdentifier, commitTableRequest)) - .build(); - } + return withCatalog( + securityContext, + prefix, + catalog -> { + if (PolarisCatalogHandlerWrapper.isCreate(commitTableRequest)) { + return Response.ok( + catalog.updateTableForStagedCreate(tableIdentifier, commitTableRequest)) + .build(); + } else { + return Response.ok(catalog.updateTable(tableIdentifier, commitTableRequest)).build(); + } + }); } @Override @@ -399,10 +453,10 @@ public Response createView( RealmContext realmContext, SecurityContext securityContext) { Namespace ns = decodeNamespace(namespace); - return Response.ok( - newHandlerWrapper(realmContext, securityContext, prefix) - .createView(ns, createViewRequest)) - .build(); + return withCatalog( + securityContext, + prefix, + catalog -> Response.ok(catalog.createView(ns, createViewRequest)).build()); } @Override @@ -414,8 +468,8 @@ public Response listViews( RealmContext realmContext, SecurityContext securityContext) { Namespace ns = decodeNamespace(namespace); - return Response.ok(newHandlerWrapper(realmContext, securityContext, prefix).listViews(ns)) - .build(); + return withCatalog( + securityContext, prefix, catalog -> Response.ok(catalog.listViews(ns)).build()); } @Override @@ -427,9 +481,8 @@ public Response loadView( SecurityContext securityContext) { Namespace ns = decodeNamespace(namespace); TableIdentifier tableIdentifier = TableIdentifier.of(ns, RESTUtil.decodeString(view)); - return Response.ok( - newHandlerWrapper(realmContext, securityContext, prefix).loadView(tableIdentifier)) - .build(); + return withCatalog( + securityContext, prefix, catalog -> Response.ok(catalog.loadView(tableIdentifier)).build()); } @Override @@ -441,8 +494,13 @@ public Response viewExists( SecurityContext securityContext) { Namespace ns = decodeNamespace(namespace); TableIdentifier tableIdentifier = TableIdentifier.of(ns, RESTUtil.decodeString(view)); - newHandlerWrapper(realmContext, securityContext, prefix).viewExists(tableIdentifier); - return Response.status(Response.Status.NO_CONTENT).build(); + return withCatalog( + securityContext, + prefix, + catalog -> { + catalog.viewExists(tableIdentifier); + return Response.status(Response.Status.NO_CONTENT).build(); + }); } @Override @@ -454,8 +512,13 @@ public Response dropView( SecurityContext securityContext) { Namespace ns = decodeNamespace(namespace); TableIdentifier tableIdentifier = TableIdentifier.of(ns, RESTUtil.decodeString(view)); - newHandlerWrapper(realmContext, securityContext, prefix).dropView(tableIdentifier); - return Response.status(Response.Status.NO_CONTENT).build(); + return withCatalog( + securityContext, + prefix, + catalog -> { + catalog.dropView(tableIdentifier); + return Response.status(Response.Status.NO_CONTENT).build(); + }); } @Override @@ -464,8 +527,13 @@ public Response renameView( RenameTableRequest renameTableRequest, RealmContext realmContext, SecurityContext securityContext) { - newHandlerWrapper(realmContext, securityContext, prefix).renameView(renameTableRequest); - return Response.status(Response.Status.NO_CONTENT).build(); + return withCatalog( + securityContext, + prefix, + catalog -> { + catalog.renameView(renameTableRequest); + return Response.status(Response.Status.NO_CONTENT).build(); + }); } @Override @@ -478,10 +546,10 @@ public Response replaceView( SecurityContext securityContext) { Namespace ns = decodeNamespace(namespace); TableIdentifier tableIdentifier = TableIdentifier.of(ns, RESTUtil.decodeString(view)); - return Response.ok( - newHandlerWrapper(realmContext, securityContext, prefix) - .replaceView(tableIdentifier, commitViewRequest)) - .build(); + return withCatalog( + securityContext, + prefix, + catalog -> Response.ok(catalog.replaceView(tableIdentifier, commitViewRequest)).build()); } @Override @@ -490,9 +558,13 @@ public Response commitTransaction( CommitTransactionRequest commitTransactionRequest, RealmContext realmContext, SecurityContext securityContext) { - newHandlerWrapper(realmContext, securityContext, prefix) - .commitTransaction(commitTransactionRequest); - return Response.status(Response.Status.NO_CONTENT).build(); + return withCatalog( + securityContext, + prefix, + catalog -> { + catalog.commitTransaction(commitTransactionRequest); + return Response.status(Response.Status.NO_CONTENT).build(); + }); } @Override @@ -516,9 +588,13 @@ public Response sendNotification( SecurityContext securityContext) { Namespace ns = decodeNamespace(namespace); TableIdentifier tableIdentifier = TableIdentifier.of(ns, RESTUtil.decodeString(table)); - newHandlerWrapper(realmContext, securityContext, prefix) - .sendNotification(tableIdentifier, notificationRequest); - return Response.status(Response.Status.NO_CONTENT).build(); + return withCatalog( + securityContext, + prefix, + catalog -> { + catalog.sendNotification(tableIdentifier, notificationRequest); + return Response.status(Response.Status.NO_CONTENT).build(); + }); } /** From IcebergRestConfigurationApiService. */ @@ -536,8 +612,6 @@ public Response getConfig( // the catalog being accessed. // TODO: Push this down into PolarisCatalogHandlerWrapper for authorizing "any" catalog // role in this catalog. - PolarisEntityManager entityManager = - entityManagerFactory.getOrCreateEntityManager(realmContext); AuthenticatedPolarisPrincipal authenticatedPrincipal = (AuthenticatedPolarisPrincipal) securityContext.getUserPrincipal(); if (authenticatedPrincipal == null) { @@ -546,7 +620,7 @@ public Response getConfig( if (warehouse == null) { throw new BadRequestException("Please specify a warehouse"); } - Resolver resolver = entityManager.prepareResolver(callContext, securityContext, warehouse); + Resolver resolver = entityManager.prepareResolver(session, securityContext, warehouse); ResolverStatus resolverStatus = resolver.resolveAll(); if (!resolverStatus.getStatus().equals(ResolverStatus.StatusEnum.SUCCESS)) { throw new NotFoundException("Unable to find warehouse %s", warehouse); diff --git a/service/common/src/main/java/org/apache/polaris/service/catalog/PolarisCatalogHandlerWrapper.java b/service/common/src/main/java/org/apache/polaris/service/catalog/PolarisCatalogHandlerWrapper.java index 823cb0db6..95c8c4856 100644 --- a/service/common/src/main/java/org/apache/polaris/service/catalog/PolarisCatalogHandlerWrapper.java +++ b/service/common/src/main/java/org/apache/polaris/service/catalog/PolarisCatalogHandlerWrapper.java @@ -32,7 +32,6 @@ import java.util.Map; import java.util.Optional; import java.util.Set; -import java.util.function.Supplier; import java.util.stream.Collectors; import org.apache.iceberg.BaseMetadataTable; import org.apache.iceberg.BaseTable; @@ -78,12 +77,13 @@ import org.apache.polaris.core.auth.PolarisAuthorizableOperation; import org.apache.polaris.core.auth.PolarisAuthorizer; import org.apache.polaris.core.catalog.PolarisCatalogHelpers; -import org.apache.polaris.core.context.CallContext; +import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.CatalogEntity; import org.apache.polaris.core.entity.PolarisEntitySubType; import org.apache.polaris.core.entity.PolarisEntityType; import org.apache.polaris.core.persistence.PolarisEntityManager; import org.apache.polaris.core.persistence.PolarisMetaStoreManager; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; import org.apache.polaris.core.persistence.PolarisResolvedPathWrapper; import org.apache.polaris.core.persistence.TransactionWorkspaceMetaStoreManager; import org.apache.polaris.core.persistence.resolver.PolarisResolutionManifest; @@ -110,10 +110,13 @@ * model objects used in this layer to still benefit from the shared implementation of * authorization-aware catalog protocols. */ -public class PolarisCatalogHandlerWrapper { +public class PolarisCatalogHandlerWrapper implements AutoCloseable { private static final Logger LOGGER = LoggerFactory.getLogger(PolarisCatalogHandlerWrapper.class); - private final CallContext callContext; + private final RealmContext realmContext; + private final PolarisMetaStoreSession session; + private final PolarisConfigurationStore configurationStore; + private final PolarisDiagnostics diagnostics; private final PolarisEntityManager entityManager; private final PolarisMetaStoreManager metaStoreManager; private final String catalogName; @@ -132,21 +135,26 @@ public class PolarisCatalogHandlerWrapper { private ViewCatalog viewCatalog = null; public PolarisCatalogHandlerWrapper( - CallContext callContext, + RealmContext realmContext, + PolarisMetaStoreSession session, + PolarisConfigurationStore configurationStore, + PolarisDiagnostics diagnostics, PolarisEntityManager entityManager, PolarisMetaStoreManager metaStoreManager, SecurityContext securityContext, CallContextCatalogFactory catalogFactory, String catalogName, PolarisAuthorizer authorizer) { - this.callContext = callContext; + this.realmContext = realmContext; + this.session = session; this.entityManager = entityManager; this.metaStoreManager = metaStoreManager; + this.diagnostics = diagnostics; + this.configurationStore = configurationStore; this.catalogName = catalogName; - PolarisDiagnostics diagServices = callContext.getPolarisCallContext().getDiagServices(); - diagServices.checkNotNull(securityContext, "null_security_context"); - diagServices.checkNotNull(securityContext.getUserPrincipal(), "null_user_principal"); - diagServices.check( + diagnostics.checkNotNull(securityContext, "null_security_context"); + diagnostics.checkNotNull(securityContext.getUserPrincipal(), "null_user_principal"); + diagnostics.check( securityContext.getUserPrincipal() instanceof AuthenticatedPolarisPrincipal, "invalid_principal_type", "Principal must be an AuthenticatedPolarisPrincipal"); @@ -178,10 +186,17 @@ public static boolean isCreate(UpdateTableRequest request) { return isCreate; } + @Override + public void close() throws IOException { + if (baseCatalog instanceof Closeable closeable) { + closeable.close(); + } + } + private void initializeCatalog() { this.baseCatalog = catalogFactory.createCallContextCatalog( - callContext, authenticatedPrincipal, securityContext, resolutionManifest); + realmContext, authenticatedPrincipal, securityContext, resolutionManifest); this.namespaceCatalog = (baseCatalog instanceof SupportsNamespaces) ? (SupportsNamespaces) baseCatalog : null; this.viewCatalog = (baseCatalog instanceof ViewCatalog) ? (ViewCatalog) baseCatalog : null; @@ -198,7 +213,7 @@ private void authorizeBasicNamespaceOperationOrThrow( List extraPassthroughNamespaces, List extraPassthroughTableLikes) { resolutionManifest = - entityManager.prepareResolutionManifest(callContext, securityContext, catalogName); + entityManager.prepareResolutionManifest(session, securityContext, catalogName); resolutionManifest.addPath( new ResolverPath(Arrays.asList(namespace.levels()), PolarisEntityType.NAMESPACE), namespace); @@ -227,6 +242,7 @@ private void authorizeBasicNamespaceOperationOrThrow( throw new NoSuchNamespaceException("Namespace does not exist: %s", namespace); } authorizer.authorizeOrThrow( + realmContext, authenticatedPrincipal, resolutionManifest.getAllActivatedCatalogRoleAndPrincipalRoles(), op, @@ -239,7 +255,7 @@ private void authorizeBasicNamespaceOperationOrThrow( private void authorizeCreateNamespaceUnderNamespaceOperationOrThrow( PolarisAuthorizableOperation op, Namespace namespace) { resolutionManifest = - entityManager.prepareResolutionManifest(callContext, securityContext, catalogName); + entityManager.prepareResolutionManifest(session, securityContext, catalogName); Namespace parentNamespace = PolarisCatalogHelpers.getParentNamespace(namespace); resolutionManifest.addPath( @@ -260,6 +276,7 @@ private void authorizeCreateNamespaceUnderNamespaceOperationOrThrow( throw new NoSuchNamespaceException("Namespace does not exist: %s", parentNamespace); } authorizer.authorizeOrThrow( + realmContext, authenticatedPrincipal, resolutionManifest.getAllActivatedCatalogRoleAndPrincipalRoles(), op, @@ -274,7 +291,7 @@ private void authorizeCreateTableLikeUnderNamespaceOperationOrThrow( Namespace namespace = identifier.namespace(); resolutionManifest = - entityManager.prepareResolutionManifest(callContext, securityContext, catalogName); + entityManager.prepareResolutionManifest(session, securityContext, catalogName); resolutionManifest.addPath( new ResolverPath(Arrays.asList(namespace.levels()), PolarisEntityType.NAMESPACE), namespace); @@ -297,6 +314,7 @@ private void authorizeCreateTableLikeUnderNamespaceOperationOrThrow( throw new NoSuchNamespaceException("Namespace does not exist: %s", namespace); } authorizer.authorizeOrThrow( + realmContext, authenticatedPrincipal, resolutionManifest.getAllActivatedCatalogRoleAndPrincipalRoles(), op, @@ -309,7 +327,7 @@ private void authorizeCreateTableLikeUnderNamespaceOperationOrThrow( private void authorizeBasicTableLikeOperationOrThrow( PolarisAuthorizableOperation op, PolarisEntitySubType subType, TableIdentifier identifier) { resolutionManifest = - entityManager.prepareResolutionManifest(callContext, securityContext, catalogName); + entityManager.prepareResolutionManifest(session, securityContext, catalogName); // The underlying Catalog is also allowed to fetch "fresh" versions of the target entity. resolutionManifest.addPassthroughPath( @@ -329,6 +347,7 @@ private void authorizeBasicTableLikeOperationOrThrow( } } authorizer.authorizeOrThrow( + realmContext, authenticatedPrincipal, resolutionManifest.getAllActivatedCatalogRoleAndPrincipalRoles(), op, @@ -343,7 +362,7 @@ private void authorizeCollectionOfTableLikeOperationOrThrow( final PolarisEntitySubType subType, List ids) { resolutionManifest = - entityManager.prepareResolutionManifest(callContext, securityContext, catalogName); + entityManager.prepareResolutionManifest(session, securityContext, catalogName); ids.forEach( identifier -> resolutionManifest.addPassthroughPath( @@ -382,6 +401,7 @@ private void authorizeCollectionOfTableLikeOperationOrThrow( "View does not exist: %s", identifier))) .toList(); authorizer.authorizeOrThrow( + realmContext, authenticatedPrincipal, resolutionManifest.getAllActivatedCatalogRoleAndPrincipalRoles(), op, @@ -397,7 +417,7 @@ private void authorizeRenameTableLikeOperationOrThrow( TableIdentifier src, TableIdentifier dst) { resolutionManifest = - entityManager.prepareResolutionManifest(callContext, securityContext, catalogName); + entityManager.prepareResolutionManifest(session, securityContext, catalogName); // Add src, dstParent, and dst(optional) resolutionManifest.addPath( new ResolverPath( @@ -442,6 +462,7 @@ private void authorizeRenameTableLikeOperationOrThrow( PolarisResolvedPathWrapper secondary = resolutionManifest.getResolvedPath(dst.namespace(), true); authorizer.authorizeOrThrow( + realmContext, authenticatedPrincipal, resolutionManifest.getAllActivatedCatalogRoleAndPrincipalRoles(), op, @@ -455,7 +476,7 @@ public ListNamespacesResponse listNamespaces(Namespace parent) { PolarisAuthorizableOperation op = PolarisAuthorizableOperation.LIST_NAMESPACES; authorizeBasicNamespaceOperationOrThrow(op, parent); - return doCatalogOperation(() -> CatalogHandlers.listNamespaces(namespaceCatalog, parent)); + return CatalogHandlers.listNamespaces(namespaceCatalog, parent); } public CreateNamespaceResponse createNamespace(CreateNamespaceRequest request) { @@ -478,20 +499,17 @@ public CreateNamespaceResponse createNamespace(CreateNamespaceRequest request) { // For CreateNamespace, we consider this a special case in that the creator is able to // retrieve the latest namespace metadata for the duration of the CreateNamespace // operation, even if the entityVersion and/or grantsVersion update in the interim. - return doCatalogOperation( - () -> { - namespaceCatalog.createNamespace(namespace, request.properties()); - return CreateNamespaceResponse.builder() - .withNamespace(namespace) - .setProperties( - resolutionManifest - .getPassthroughResolvedPath(namespace) - .getRawLeafEntity() - .getPropertiesAsMap()) - .build(); - }); + namespaceCatalog.createNamespace(namespace, request.properties()); + return CreateNamespaceResponse.builder() + .withNamespace(namespace) + .setProperties( + resolutionManifest + .getPassthroughResolvedPath(namespace) + .getRawLeafEntity() + .getPropertiesAsMap()) + .build(); } else { - return doCatalogOperation(() -> CatalogHandlers.createNamespace(namespaceCatalog, request)); + return CatalogHandlers.createNamespace(namespaceCatalog, request); } } @@ -500,37 +518,11 @@ private static boolean isExternal(CatalogEntity catalog) { catalog.getCatalogType()); } - private void doCatalogOperation(Runnable handler) { - doCatalogOperation( - () -> { - handler.run(); - return null; - }); - } - - /** - * Execute a catalog function and ensure we close the BaseCatalog afterward. This will typically - * ensure the underlying FileIO is closed - */ - private T doCatalogOperation(Supplier handler) { - try { - return handler.get(); - } finally { - if (baseCatalog instanceof Closeable closeable) { - try { - closeable.close(); - } catch (IOException e) { - LOGGER.error("Error while closing BaseCatalog", e); - } - } - } - } - public GetNamespaceResponse loadNamespaceMetadata(Namespace namespace) { PolarisAuthorizableOperation op = PolarisAuthorizableOperation.LOAD_NAMESPACE_METADATA; authorizeBasicNamespaceOperationOrThrow(op, namespace); - return doCatalogOperation(() -> CatalogHandlers.loadNamespace(namespaceCatalog, namespace)); + return CatalogHandlers.loadNamespace(namespaceCatalog, namespace); } public void namespaceExists(Namespace namespace) { @@ -545,14 +537,14 @@ public void namespaceExists(Namespace namespace) { authorizeBasicNamespaceOperationOrThrow(op, namespace); // TODO: Just skip CatalogHandlers for this one maybe - doCatalogOperation(() -> CatalogHandlers.loadNamespace(namespaceCatalog, namespace)); + CatalogHandlers.loadNamespace(namespaceCatalog, namespace); } public void dropNamespace(Namespace namespace) { PolarisAuthorizableOperation op = PolarisAuthorizableOperation.DROP_NAMESPACE; authorizeBasicNamespaceOperationOrThrow(op, namespace); - doCatalogOperation(() -> CatalogHandlers.dropNamespace(namespaceCatalog, namespace)); + CatalogHandlers.dropNamespace(namespaceCatalog, namespace); } public UpdateNamespacePropertiesResponse updateNamespaceProperties( @@ -560,15 +552,14 @@ public UpdateNamespacePropertiesResponse updateNamespaceProperties( PolarisAuthorizableOperation op = PolarisAuthorizableOperation.UPDATE_NAMESPACE_PROPERTIES; authorizeBasicNamespaceOperationOrThrow(op, namespace); - return doCatalogOperation( - () -> CatalogHandlers.updateNamespaceProperties(namespaceCatalog, namespace, request)); + return CatalogHandlers.updateNamespaceProperties(namespaceCatalog, namespace, request); } public ListTablesResponse listTables(Namespace namespace) { PolarisAuthorizableOperation op = PolarisAuthorizableOperation.LIST_TABLES; authorizeBasicNamespaceOperationOrThrow(op, namespace); - return doCatalogOperation(() -> CatalogHandlers.listTables(baseCatalog, namespace)); + return CatalogHandlers.listTables(baseCatalog, namespace); } public LoadTableResponse createTableDirect(Namespace namespace, CreateTableRequest request) { @@ -585,7 +576,7 @@ public LoadTableResponse createTableDirect(Namespace namespace, CreateTableReque if (isExternal(catalog)) { throw new BadRequestException("Cannot create table on external catalogs."); } - return doCatalogOperation(() -> CatalogHandlers.createTable(baseCatalog, namespace, request)); + return CatalogHandlers.createTable(baseCatalog, namespace, request); } public LoadTableResponse createTableDirectWithWriteDelegation( @@ -604,55 +595,52 @@ public LoadTableResponse createTableDirectWithWriteDelegation( if (isExternal(catalog)) { throw new BadRequestException("Cannot create table on external catalogs."); } - return doCatalogOperation( - () -> { - request.validate(); - - TableIdentifier tableIdentifier = TableIdentifier.of(namespace, request.name()); - if (baseCatalog.tableExists(tableIdentifier)) { - throw new AlreadyExistsException("Table already exists: %s", tableIdentifier); - } - - Map properties = Maps.newHashMap(); - properties.put("created-at", OffsetDateTime.now(ZoneOffset.UTC).toString()); - properties.putAll(request.properties()); - - Table table = - baseCatalog - .buildTable(tableIdentifier, request.schema()) - .withLocation(request.location()) - .withPartitionSpec(request.spec()) - .withSortOrder(request.writeOrder()) - .withProperties(properties) - .create(); - - if (table instanceof BaseTable baseTable) { - TableMetadata tableMetadata = baseTable.operations().current(); - LoadTableResponse.Builder responseBuilder = - LoadTableResponse.builder().withTableMetadata(tableMetadata); - if (baseCatalog instanceof SupportsCredentialDelegation credentialDelegation) { - LOGGER - .atDebug() - .addKeyValue("tableIdentifier", tableIdentifier) - .addKeyValue("tableLocation", tableMetadata.location()) - .log("Fetching client credentials for table"); - responseBuilder.addAllConfig( - credentialDelegation.getCredentialConfig( - tableIdentifier, - tableMetadata, - Set.of( - PolarisStorageActions.READ, - PolarisStorageActions.WRITE, - PolarisStorageActions.LIST))); - } - return responseBuilder.build(); - } else if (table instanceof BaseMetadataTable) { - // metadata tables are loaded on the client side, return NoSuchTableException for now - throw new NoSuchTableException("Table does not exist: %s", tableIdentifier.toString()); - } - - throw new IllegalStateException("Cannot wrap catalog that does not produce BaseTable"); - }); + request.validate(); + + TableIdentifier tableIdentifier = TableIdentifier.of(namespace, request.name()); + if (baseCatalog.tableExists(tableIdentifier)) { + throw new AlreadyExistsException("Table already exists: %s", tableIdentifier); + } + + Map properties = Maps.newHashMap(); + properties.put("created-at", OffsetDateTime.now(ZoneOffset.UTC).toString()); + properties.putAll(request.properties()); + + Table table = + baseCatalog + .buildTable(tableIdentifier, request.schema()) + .withLocation(request.location()) + .withPartitionSpec(request.spec()) + .withSortOrder(request.writeOrder()) + .withProperties(properties) + .create(); + + if (table instanceof BaseTable baseTable) { + TableMetadata tableMetadata = baseTable.operations().current(); + LoadTableResponse.Builder responseBuilder = + LoadTableResponse.builder().withTableMetadata(tableMetadata); + if (baseCatalog instanceof SupportsCredentialDelegation credentialDelegation) { + LOGGER + .atDebug() + .addKeyValue("tableIdentifier", tableIdentifier) + .addKeyValue("tableLocation", tableMetadata.location()) + .log("Fetching client credentials for table"); + responseBuilder.addAllConfig( + credentialDelegation.getCredentialConfig( + tableIdentifier, + tableMetadata, + Set.of( + PolarisStorageActions.READ, + PolarisStorageActions.WRITE, + PolarisStorageActions.LIST))); + } + return responseBuilder.build(); + } else if (table instanceof BaseMetadataTable) { + // metadata tables are loaded on the client side, return NoSuchTableException for now + throw new NoSuchTableException("Table does not exist: %s", tableIdentifier.toString()); + } + + throw new IllegalStateException("Cannot wrap catalog that does not produce BaseTable"); } private TableMetadata stageTableCreateHelper(Namespace namespace, CreateTableRequest request) { @@ -713,11 +701,8 @@ public LoadTableResponse createTableStaged(Namespace namespace, CreateTableReque if (isExternal(catalog)) { throw new BadRequestException("Cannot create table on external catalogs."); } - return doCatalogOperation( - () -> { - TableMetadata metadata = stageTableCreateHelper(namespace, request); - return LoadTableResponse.builder().withTableMetadata(metadata).build(); - }); + TableMetadata metadata = stageTableCreateHelper(namespace, request); + return LoadTableResponse.builder().withTableMetadata(metadata).build(); } public LoadTableResponse createTableStagedWithWriteDelegation( @@ -736,26 +721,23 @@ public LoadTableResponse createTableStagedWithWriteDelegation( if (isExternal(catalog)) { throw new BadRequestException("Cannot create table on external catalogs."); } - return doCatalogOperation( - () -> { - TableIdentifier ident = TableIdentifier.of(namespace, request.name()); - TableMetadata metadata = stageTableCreateHelper(namespace, request); - - LoadTableResponse.Builder responseBuilder = - LoadTableResponse.builder().withTableMetadata(metadata); - - if (baseCatalog instanceof SupportsCredentialDelegation credentialDelegation) { - LOGGER - .atDebug() - .addKeyValue("tableIdentifier", ident) - .addKeyValue("tableLocation", metadata.location()) - .log("Fetching client credentials for table"); - responseBuilder.addAllConfig( - credentialDelegation.getCredentialConfig( - ident, metadata, Set.of(PolarisStorageActions.ALL))); - } - return responseBuilder.build(); - }); + TableIdentifier ident = TableIdentifier.of(namespace, request.name()); + TableMetadata metadata = stageTableCreateHelper(namespace, request); + + LoadTableResponse.Builder responseBuilder = + LoadTableResponse.builder().withTableMetadata(metadata); + + if (baseCatalog instanceof SupportsCredentialDelegation credentialDelegation) { + LOGGER + .atDebug() + .addKeyValue("tableIdentifier", ident) + .addKeyValue("tableLocation", metadata.location()) + .log("Fetching client credentials for table"); + responseBuilder.addAllConfig( + credentialDelegation.getCredentialConfig( + ident, metadata, Set.of(PolarisStorageActions.ALL))); + } + return responseBuilder.build(); } public LoadTableResponse registerTable(Namespace namespace, RegisterTableRequest request) { @@ -763,7 +745,7 @@ public LoadTableResponse registerTable(Namespace namespace, RegisterTableRequest authorizeCreateTableLikeUnderNamespaceOperationOrThrow( op, TableIdentifier.of(namespace, request.name())); - return doCatalogOperation(() -> CatalogHandlers.registerTable(baseCatalog, namespace, request)); + return CatalogHandlers.registerTable(baseCatalog, namespace, request); } public boolean sendNotification(TableIdentifier identifier, NotificationRequest request) { @@ -808,7 +790,7 @@ public LoadTableResponse loadTable(TableIdentifier tableIdentifier, String snaps PolarisAuthorizableOperation op = PolarisAuthorizableOperation.LOAD_TABLE; authorizeBasicTableLikeOperationOrThrow(op, PolarisEntitySubType.TABLE, tableIdentifier); - return doCatalogOperation(() -> CatalogHandlers.loadTable(baseCatalog, tableIdentifier)); + return CatalogHandlers.loadTable(baseCatalog, tableIdentifier); } public LoadTableResponse loadTableWithAccessDelegation( @@ -835,18 +817,13 @@ public LoadTableResponse loadTableWithAccessDelegation( } PolarisResolvedPathWrapper catalogPath = resolutionManifest.getResolvedReferenceCatalogEntity(); - callContext - .getPolarisCallContext() - .getDiagServices() - .checkNotNull(catalogPath, "No catalog available for loadTable request"); + diagnostics.checkNotNull(catalogPath, "No catalog available for loadTable request"); CatalogEntity catalogEntity = CatalogEntity.of(catalogPath.getRawLeafEntity()); - PolarisConfigurationStore configurationStore = - callContext.getPolarisCallContext().getConfigurationStore(); if (catalogEntity .getCatalogType() .equals(org.apache.polaris.core.admin.model.Catalog.TypeEnum.EXTERNAL) && !configurationStore.getConfiguration( - callContext.getPolarisCallContext(), + realmContext, catalogEntity, PolarisConfiguration.ALLOW_EXTERNAL_CATALOG_CREDENTIAL_VENDING)) { throw new ForbiddenException( @@ -857,32 +834,29 @@ public LoadTableResponse loadTableWithAccessDelegation( // TODO: Find a way for the configuration or caller to better express whether to fail or omit // when data-access is specified but access delegation grants are not found. - return doCatalogOperation( - () -> { - Table table = baseCatalog.loadTable(tableIdentifier); - - if (table instanceof BaseTable baseTable) { - TableMetadata tableMetadata = baseTable.operations().current(); - LoadTableResponse.Builder responseBuilder = - LoadTableResponse.builder().withTableMetadata(tableMetadata); - if (baseCatalog instanceof SupportsCredentialDelegation credentialDelegation) { - LOGGER - .atDebug() - .addKeyValue("tableIdentifier", tableIdentifier) - .addKeyValue("tableLocation", tableMetadata.location()) - .log("Fetching client credentials for table"); - responseBuilder.addAllConfig( - credentialDelegation.getCredentialConfig( - tableIdentifier, tableMetadata, actionsRequested)); - } - return responseBuilder.build(); - } else if (table instanceof BaseMetadataTable) { - // metadata tables are loaded on the client side, return NoSuchTableException for now - throw new NoSuchTableException("Table does not exist: %s", tableIdentifier.toString()); - } - - throw new IllegalStateException("Cannot wrap catalog that does not produce BaseTable"); - }); + Table table = baseCatalog.loadTable(tableIdentifier); + + if (table instanceof BaseTable baseTable) { + TableMetadata tableMetadata = baseTable.operations().current(); + LoadTableResponse.Builder responseBuilder = + LoadTableResponse.builder().withTableMetadata(tableMetadata); + if (baseCatalog instanceof SupportsCredentialDelegation credentialDelegation) { + LOGGER + .atDebug() + .addKeyValue("tableIdentifier", tableIdentifier) + .addKeyValue("tableLocation", tableMetadata.location()) + .log("Fetching client credentials for table"); + responseBuilder.addAllConfig( + credentialDelegation.getCredentialConfig( + tableIdentifier, tableMetadata, actionsRequested)); + } + return responseBuilder.build(); + } else if (table instanceof BaseMetadataTable) { + // metadata tables are loaded on the client side, return NoSuchTableException for now + throw new NoSuchTableException("Table does not exist: %s", tableIdentifier.toString()); + } + + throw new IllegalStateException("Cannot wrap catalog that does not produce BaseTable"); } private UpdateTableRequest applyUpdateFilters(UpdateTableRequest request) { @@ -923,9 +897,7 @@ public LoadTableResponse updateTable( if (isExternal(catalog)) { throw new BadRequestException("Cannot update table on external catalogs."); } - return doCatalogOperation( - () -> - CatalogHandlers.updateTable(baseCatalog, tableIdentifier, applyUpdateFilters(request))); + return CatalogHandlers.updateTable(baseCatalog, tableIdentifier, applyUpdateFilters(request)); } public LoadTableResponse updateTableForStagedCreate( @@ -942,16 +914,14 @@ public LoadTableResponse updateTableForStagedCreate( if (isExternal(catalog)) { throw new BadRequestException("Cannot update table on external catalogs."); } - return doCatalogOperation( - () -> - CatalogHandlers.updateTable(baseCatalog, tableIdentifier, applyUpdateFilters(request))); + return CatalogHandlers.updateTable(baseCatalog, tableIdentifier, applyUpdateFilters(request)); } public void dropTableWithoutPurge(TableIdentifier tableIdentifier) { PolarisAuthorizableOperation op = PolarisAuthorizableOperation.DROP_TABLE_WITHOUT_PURGE; authorizeBasicTableLikeOperationOrThrow(op, PolarisEntitySubType.TABLE, tableIdentifier); - doCatalogOperation(() -> CatalogHandlers.dropTable(baseCatalog, tableIdentifier)); + CatalogHandlers.dropTable(baseCatalog, tableIdentifier); } public void dropTableWithPurge(TableIdentifier tableIdentifier) { @@ -967,7 +937,7 @@ public void dropTableWithPurge(TableIdentifier tableIdentifier) { if (isExternal(catalog)) { throw new BadRequestException("Cannot drop table on external catalogs."); } - doCatalogOperation(() -> CatalogHandlers.purgeTable(baseCatalog, tableIdentifier)); + CatalogHandlers.purgeTable(baseCatalog, tableIdentifier); } public void tableExists(TableIdentifier tableIdentifier) { @@ -975,7 +945,7 @@ public void tableExists(TableIdentifier tableIdentifier) { authorizeBasicTableLikeOperationOrThrow(op, PolarisEntitySubType.TABLE, tableIdentifier); // TODO: Just skip CatalogHandlers for this one maybe - doCatalogOperation(() -> CatalogHandlers.loadTable(baseCatalog, tableIdentifier)); + CatalogHandlers.loadTable(baseCatalog, tableIdentifier); } public void renameTable(RenameTableRequest request) { @@ -992,7 +962,7 @@ public void renameTable(RenameTableRequest request) { if (isExternal(catalog)) { throw new BadRequestException("Cannot rename table on external catalogs."); } - doCatalogOperation(() -> CatalogHandlers.renameTable(baseCatalog, request)); + CatalogHandlers.renameTable(baseCatalog, request); } public void commitTransaction(CommitTransactionRequest commitTransactionRequest) { @@ -1027,7 +997,7 @@ public void commitTransaction(CommitTransactionRequest commitTransactionRequest) // only go into an in-memory collection that we can commit as a single atomic unit after all // validations. TransactionWorkspaceMetaStoreManager transactionMetaStoreManager = - new TransactionWorkspaceMetaStoreManager(metaStoreManager); + new TransactionWorkspaceMetaStoreManager(metaStoreManager, diagnostics); ((BasePolarisCatalog) baseCatalog).setMetaStoreManager(transactionMetaStoreManager); commitTransactionRequest.tableChanges().stream() @@ -1063,12 +1033,9 @@ public void commitTransaction(CommitTransactionRequest commitTransactionRequest) if (!currentMetadata .location() .equals(((MetadataUpdate.SetLocation) singleUpdate).location()) - && !callContext - .getPolarisCallContext() - .getConfigurationStore() - .getConfiguration( - callContext.getPolarisCallContext(), - PolarisConfiguration.ALLOW_NAMESPACE_LOCATION_OVERLAP)) { + && !configurationStore.getConfiguration( + realmContext, + PolarisConfiguration.ALLOW_NAMESPACE_LOCATION_OVERLAP)) { throw new BadRequestException( "Unsupported operation: commitTransaction containing SetLocation" + " for table '%s' and new location '%s'", @@ -1092,8 +1059,7 @@ public void commitTransaction(CommitTransactionRequest commitTransactionRequest) List pendingUpdates = transactionMetaStoreManager.getPendingUpdates(); PolarisMetaStoreManager.EntitiesResult result = - metaStoreManager.updateEntitiesPropertiesIfNotChanged( - callContext.getPolarisCallContext(), pendingUpdates); + metaStoreManager.updateEntitiesPropertiesIfNotChanged(session, pendingUpdates); if (!result.isSuccess()) { // TODO: Retries and server-side cleanup on failure throw new CommitFailedException( @@ -1106,7 +1072,7 @@ public ListTablesResponse listViews(Namespace namespace) { PolarisAuthorizableOperation op = PolarisAuthorizableOperation.LIST_VIEWS; authorizeBasicNamespaceOperationOrThrow(op, namespace); - return doCatalogOperation(() -> CatalogHandlers.listViews(viewCatalog, namespace)); + return CatalogHandlers.listViews(viewCatalog, namespace); } public LoadViewResponse createView(Namespace namespace, CreateViewRequest request) { @@ -1123,14 +1089,14 @@ public LoadViewResponse createView(Namespace namespace, CreateViewRequest reques if (isExternal(catalog)) { throw new BadRequestException("Cannot create view on external catalogs."); } - return doCatalogOperation(() -> CatalogHandlers.createView(viewCatalog, namespace, request)); + return CatalogHandlers.createView(viewCatalog, namespace, request); } public LoadViewResponse loadView(TableIdentifier viewIdentifier) { PolarisAuthorizableOperation op = PolarisAuthorizableOperation.LOAD_VIEW; authorizeBasicTableLikeOperationOrThrow(op, PolarisEntitySubType.VIEW, viewIdentifier); - return doCatalogOperation(() -> CatalogHandlers.loadView(viewCatalog, viewIdentifier)); + return CatalogHandlers.loadView(viewCatalog, viewIdentifier); } public LoadViewResponse replaceView(TableIdentifier viewIdentifier, UpdateTableRequest request) { @@ -1146,15 +1112,14 @@ public LoadViewResponse replaceView(TableIdentifier viewIdentifier, UpdateTableR if (isExternal(catalog)) { throw new BadRequestException("Cannot replace view on external catalogs."); } - return doCatalogOperation( - () -> CatalogHandlers.updateView(viewCatalog, viewIdentifier, applyUpdateFilters(request))); + return CatalogHandlers.updateView(viewCatalog, viewIdentifier, applyUpdateFilters(request)); } public void dropView(TableIdentifier viewIdentifier) { PolarisAuthorizableOperation op = PolarisAuthorizableOperation.DROP_VIEW; authorizeBasicTableLikeOperationOrThrow(op, PolarisEntitySubType.VIEW, viewIdentifier); - doCatalogOperation(() -> CatalogHandlers.dropView(viewCatalog, viewIdentifier)); + CatalogHandlers.dropView(viewCatalog, viewIdentifier); } public void viewExists(TableIdentifier viewIdentifier) { @@ -1162,7 +1127,7 @@ public void viewExists(TableIdentifier viewIdentifier) { authorizeBasicTableLikeOperationOrThrow(op, PolarisEntitySubType.VIEW, viewIdentifier); // TODO: Just skip CatalogHandlers for this one maybe - doCatalogOperation(() -> CatalogHandlers.loadView(viewCatalog, viewIdentifier)); + CatalogHandlers.loadView(viewCatalog, viewIdentifier); } public void renameView(RenameTableRequest request) { @@ -1179,6 +1144,6 @@ public void renameView(RenameTableRequest request) { if (isExternal(catalog)) { throw new BadRequestException("Cannot rename view on external catalogs."); } - doCatalogOperation(() -> CatalogHandlers.renameView(viewCatalog, request)); + CatalogHandlers.renameView(viewCatalog, request); } } diff --git a/service/common/src/main/java/org/apache/polaris/service/config/DefaultConfigurationStore.java b/service/common/src/main/java/org/apache/polaris/service/config/DefaultConfigurationStore.java index 3db2fd998..17a0472a4 100644 --- a/service/common/src/main/java/org/apache/polaris/service/config/DefaultConfigurationStore.java +++ b/service/common/src/main/java/org/apache/polaris/service/config/DefaultConfigurationStore.java @@ -19,14 +19,12 @@ package org.apache.polaris.service.config; import com.fasterxml.jackson.databind.ObjectMapper; -import jakarta.annotation.Nonnull; import jakarta.annotation.Nullable; import jakarta.enterprise.context.ApplicationScoped; import jakarta.inject.Inject; import java.util.Map; -import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.PolarisConfigurationStore; -import org.apache.polaris.core.context.CallContext; +import org.apache.polaris.core.context.RealmContext; @ApplicationScoped public class DefaultConfigurationStore implements PolarisConfigurationStore { @@ -55,14 +53,16 @@ public DefaultConfigurationStore( } @Override - public @Nullable T getConfiguration(@Nonnull PolarisCallContext ctx, String configName) { - String realm = CallContext.getCurrentContext().getRealmContext().getRealmIdentifier(); + public @Nullable T getConfiguration(@Nullable RealmContext realmContext, String configName) { + Object rawValue = defaults.get(configName); + if (realmContext != null) { + rawValue = + realmOverrides + .getOrDefault(realmContext.getRealmIdentifier(), Map.of()) + .getOrDefault(configName, rawValue); + } @SuppressWarnings("unchecked") - T confgValue = - (T) - realmOverrides - .getOrDefault(realm, Map.of()) - .getOrDefault(configName, defaults.get(configName)); - return confgValue; + T value = (T) rawValue; + return value; } } diff --git a/service/common/src/main/java/org/apache/polaris/service/config/RealmEntityManagerFactory.java b/service/common/src/main/java/org/apache/polaris/service/config/RealmEntityManagerFactory.java index 03799bd40..4936dcf34 100644 --- a/service/common/src/main/java/org/apache/polaris/service/config/RealmEntityManagerFactory.java +++ b/service/common/src/main/java/org/apache/polaris/service/config/RealmEntityManagerFactory.java @@ -22,6 +22,7 @@ import jakarta.inject.Inject; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; +import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.persistence.MetaStoreManagerFactory; import org.apache.polaris.core.persistence.PolarisEntityManager; @@ -35,13 +36,16 @@ public class RealmEntityManagerFactory { private static final Logger LOGGER = LoggerFactory.getLogger(RealmEntityManagerFactory.class); private final MetaStoreManagerFactory metaStoreManagerFactory; + private final PolarisDiagnostics diagnostics; // Key: realmIdentifier private final Map cachedEntityManagers = new ConcurrentHashMap<>(); @Inject - public RealmEntityManagerFactory(MetaStoreManagerFactory metaStoreManagerFactory) { + public RealmEntityManagerFactory( + MetaStoreManagerFactory metaStoreManagerFactory, PolarisDiagnostics diagnostics) { this.metaStoreManagerFactory = metaStoreManagerFactory; + this.diagnostics = diagnostics; } public PolarisEntityManager getOrCreateEntityManager(RealmContext context) { @@ -56,7 +60,8 @@ public PolarisEntityManager getOrCreateEntityManager(RealmContext context) { return new PolarisEntityManager( metaStoreManagerFactory.getOrCreateMetaStoreManager(context), metaStoreManagerFactory.getOrCreateStorageCredentialCache(context), - metaStoreManagerFactory.getOrCreateEntityCache(context)); + metaStoreManagerFactory.getOrCreateEntityCache(context), + diagnostics); }); } } diff --git a/service/common/src/main/java/org/apache/polaris/service/context/CallContextCatalogFactory.java b/service/common/src/main/java/org/apache/polaris/service/context/CallContextCatalogFactory.java index 7c812d6d3..b300aa322 100644 --- a/service/common/src/main/java/org/apache/polaris/service/context/CallContextCatalogFactory.java +++ b/service/common/src/main/java/org/apache/polaris/service/context/CallContextCatalogFactory.java @@ -21,12 +21,12 @@ import jakarta.ws.rs.core.SecurityContext; import org.apache.iceberg.catalog.Catalog; import org.apache.polaris.core.auth.AuthenticatedPolarisPrincipal; -import org.apache.polaris.core.context.CallContext; +import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.persistence.resolver.PolarisResolutionManifest; public interface CallContextCatalogFactory { Catalog createCallContextCatalog( - CallContext context, + RealmContext realmContext, AuthenticatedPolarisPrincipal authenticatedPrincipal, SecurityContext securityContext, PolarisResolutionManifest resolvedManifest); diff --git a/service/common/src/main/java/org/apache/polaris/service/context/CallContextResolver.java b/service/common/src/main/java/org/apache/polaris/service/context/CallContextResolver.java deleted file mode 100644 index 16b8c180c..000000000 --- a/service/common/src/main/java/org/apache/polaris/service/context/CallContextResolver.java +++ /dev/null @@ -1,29 +0,0 @@ -/* - * 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.context; - -import java.util.Map; -import org.apache.polaris.core.context.CallContext; -import org.apache.polaris.core.context.RealmContext; - -/** Uses the resolved RealmContext to further resolve elements of the CallContext. */ -public interface CallContextResolver { - CallContext resolveCallContext( - RealmContext realmContext, String method, String path, Map headers); -} diff --git a/service/common/src/main/java/org/apache/polaris/service/context/DefaultCallContextResolver.java b/service/common/src/main/java/org/apache/polaris/service/context/DefaultCallContextResolver.java deleted file mode 100644 index ff3949724..000000000 --- a/service/common/src/main/java/org/apache/polaris/service/context/DefaultCallContextResolver.java +++ /dev/null @@ -1,69 +0,0 @@ -/* - * 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.context; - -import io.smallrye.common.annotation.Identifier; -import jakarta.enterprise.context.ApplicationScoped; -import jakarta.inject.Inject; -import java.time.Clock; -import java.util.Map; -import org.apache.polaris.core.PolarisCallContext; -import org.apache.polaris.core.PolarisConfigurationStore; -import org.apache.polaris.core.PolarisDiagnostics; -import org.apache.polaris.core.context.CallContext; -import org.apache.polaris.core.context.RealmContext; -import org.apache.polaris.core.persistence.MetaStoreManagerFactory; -import org.apache.polaris.core.persistence.PolarisMetaStoreSession; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -/** - * For local/dev testing, this resolver simply expects a custom bearer-token format that is a - * semicolon-separated list of colon-separated key/value pairs that constitute the realm properties. - * - *

Example: principal:data-engineer;password:test;realm:acct123 - */ -@ApplicationScoped -@Identifier("default") -public class DefaultCallContextResolver implements CallContextResolver { - private static final Logger LOGGER = LoggerFactory.getLogger(DefaultCallContextResolver.class); - - @Inject MetaStoreManagerFactory metaStoreManagerFactory; - @Inject PolarisConfigurationStore configurationStore; - @Inject PolarisDiagnostics diagnostics; - @Inject Clock clock; - - @Override - public CallContext resolveCallContext( - final RealmContext realmContext, String method, String path, Map headers) { - LOGGER - .atDebug() - .addKeyValue("realmContext", realmContext.getRealmIdentifier()) - .addKeyValue("method", method) - .addKeyValue("path", path) - .addKeyValue("headers", headers) - .log("Resolving CallContext"); - - PolarisMetaStoreSession metaStoreSession = - metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(); - PolarisCallContext polarisContext = - new PolarisCallContext(metaStoreSession, diagnostics, configurationStore, clock); - return CallContext.of(realmContext, polarisContext); - } -} diff --git a/service/common/src/main/java/org/apache/polaris/service/context/PolarisCallContextCatalogFactory.java b/service/common/src/main/java/org/apache/polaris/service/context/PolarisCallContextCatalogFactory.java index 66e50e406..85773a3bc 100644 --- a/service/common/src/main/java/org/apache/polaris/service/context/PolarisCallContextCatalogFactory.java +++ b/service/common/src/main/java/org/apache/polaris/service/context/PolarisCallContextCatalogFactory.java @@ -18,7 +18,7 @@ */ package org.apache.polaris.service.context; -import jakarta.enterprise.context.ApplicationScoped; +import jakarta.enterprise.context.RequestScoped; import jakarta.inject.Inject; import jakarta.ws.rs.core.SecurityContext; import java.nio.file.Paths; @@ -27,21 +27,23 @@ import java.util.Objects; import org.apache.iceberg.CatalogProperties; import org.apache.iceberg.catalog.Catalog; +import org.apache.polaris.core.PolarisConfigurationStore; +import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.auth.AuthenticatedPolarisPrincipal; -import org.apache.polaris.core.context.CallContext; +import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.CatalogEntity; import org.apache.polaris.core.entity.PolarisBaseEntity; -import org.apache.polaris.core.persistence.MetaStoreManagerFactory; import org.apache.polaris.core.persistence.PolarisEntityManager; +import org.apache.polaris.core.persistence.PolarisMetaStoreManager; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; import org.apache.polaris.core.persistence.resolver.PolarisResolutionManifest; import org.apache.polaris.service.catalog.BasePolarisCatalog; import org.apache.polaris.service.catalog.io.FileIOFactory; -import org.apache.polaris.service.config.RealmEntityManagerFactory; import org.apache.polaris.service.task.TaskExecutor; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -@ApplicationScoped +@RequestScoped public class PolarisCallContextCatalogFactory implements CallContextCatalogFactory { private static final Logger LOGGER = LoggerFactory.getLogger(PolarisCallContextCatalogFactory.class); @@ -49,26 +51,35 @@ public class PolarisCallContextCatalogFactory implements CallContextCatalogFacto private static final String WAREHOUSE_LOCATION_BASEDIR = "/tmp/iceberg_rest_server_warehouse_data/"; - private final RealmEntityManagerFactory entityManagerFactory; + private final PolarisEntityManager entityManager; + private final PolarisMetaStoreManager metaStoreManager; + private final PolarisMetaStoreSession metaStoreSession; + private final PolarisConfigurationStore configurationStore; + private final PolarisDiagnostics diagnostics; private final TaskExecutor taskExecutor; private final FileIOFactory fileIOFactory; - private final MetaStoreManagerFactory metaStoreManagerFactory; @Inject public PolarisCallContextCatalogFactory( - RealmEntityManagerFactory entityManagerFactory, - MetaStoreManagerFactory metaStoreManagerFactory, + PolarisEntityManager entityManager, + PolarisMetaStoreManager metaStoreManager, + PolarisMetaStoreSession metaStoreSession, + PolarisConfigurationStore configurationStore, + PolarisDiagnostics diagnostics, TaskExecutor taskExecutor, FileIOFactory fileIOFactory) { - this.entityManagerFactory = entityManagerFactory; - this.metaStoreManagerFactory = metaStoreManagerFactory; + this.entityManager = entityManager; + this.metaStoreManager = metaStoreManager; + this.metaStoreSession = metaStoreSession; + this.configurationStore = configurationStore; + this.diagnostics = diagnostics; this.taskExecutor = taskExecutor; this.fileIOFactory = fileIOFactory; } @Override public Catalog createCallContextCatalog( - CallContext context, + RealmContext realmContext, AuthenticatedPolarisPrincipal authenticatedPrincipal, SecurityContext securityContext, final PolarisResolutionManifest resolvedManifest) { @@ -76,26 +87,23 @@ public Catalog createCallContextCatalog( resolvedManifest.getResolvedReferenceCatalogEntity().getRawLeafEntity(); String catalogName = baseCatalogEntity.getName(); - String realm = context.getRealmContext().getRealmIdentifier(); + String realm = realmContext.getRealmIdentifier(); String catalogKey = realm + "/" + catalogName; LOGGER.info("Initializing new BasePolarisCatalog for key: {}", catalogKey); - PolarisEntityManager entityManager = - entityManagerFactory.getOrCreateEntityManager(context.getRealmContext()); - BasePolarisCatalog catalogInstance = new BasePolarisCatalog( + realmContext, entityManager, - metaStoreManagerFactory.getOrCreateMetaStoreManager(context.getRealmContext()), - context, + metaStoreManager, + metaStoreSession, + configurationStore, + diagnostics, resolvedManifest, securityContext, - authenticatedPrincipal, taskExecutor, fileIOFactory); - context.contextVariables().put(CallContext.REQUEST_PATH_CATALOG_INSTANCE_KEY, catalogInstance); - CatalogEntity catalog = CatalogEntity.of(baseCatalogEntity); Map catalogProperties = new HashMap<>(catalog.getPropertiesAsMap()); String defaultBaseLocation = catalog.getDefaultBaseLocation(); diff --git a/service/common/src/main/java/org/apache/polaris/service/persistence/InMemoryPolarisMetaStoreManagerFactory.java b/service/common/src/main/java/org/apache/polaris/service/persistence/InMemoryPolarisMetaStoreManagerFactory.java index 3fed45e0d..0c3cc6ec0 100644 --- a/service/common/src/main/java/org/apache/polaris/service/persistence/InMemoryPolarisMetaStoreManagerFactory.java +++ b/service/common/src/main/java/org/apache/polaris/service/persistence/InMemoryPolarisMetaStoreManagerFactory.java @@ -22,11 +22,13 @@ import jakarta.annotation.Nonnull; import jakarta.enterprise.context.ApplicationScoped; import jakarta.inject.Inject; +import java.time.Clock; import java.util.Collections; -import java.util.HashSet; import java.util.Map; import java.util.Set; +import java.util.concurrent.CopyOnWriteArraySet; import java.util.function.Supplier; +import org.apache.polaris.core.PolarisConfigurationStore; import org.apache.polaris.core.PolarisDiagnostics; import org.apache.polaris.core.auth.PolarisSecretsManager.PrincipalSecretsResult; import org.apache.polaris.core.context.RealmContext; @@ -44,15 +46,19 @@ public class InMemoryPolarisMetaStoreManagerFactory extends LocalPolarisMetaStoreManagerFactory { private final PolarisStorageIntegrationProvider storageIntegration; - private final Set bootstrappedRealms = new HashSet<>(); + private final Set bootstrappedRealms = new CopyOnWriteArraySet<>(); public InMemoryPolarisMetaStoreManagerFactory() { - this(null); + this(null, null, null, null); } @Inject public InMemoryPolarisMetaStoreManagerFactory( - PolarisStorageIntegrationProvider storageIntegration) { + PolarisStorageIntegrationProvider storageIntegration, + PolarisConfigurationStore configurationStore, + PolarisDiagnostics diagnostics, + Clock clock) { + super(configurationStore, diagnostics, clock); this.storageIntegration = storageIntegration; } @@ -67,9 +73,11 @@ protected PolarisTreeMapStore createBackingStore(@Nonnull PolarisDiagnostics dia @Override protected PolarisMetaStoreSession createMetaStoreSession( - @Nonnull PolarisTreeMapStore store, @Nonnull RealmContext realmContext) { + @Nonnull PolarisTreeMapStore store, + @Nonnull RealmContext realmContext, + @Nonnull PolarisDiagnostics diagnostics) { return new PolarisTreeMapMetaStoreSessionImpl( - store, storageIntegration, secretsGenerator(realmContext)); + store, storageIntegration, secretsGenerator(realmContext), diagnostics); } @Override diff --git a/service/common/src/main/java/org/apache/polaris/service/storage/PolarisStorageIntegrationProviderImpl.java b/service/common/src/main/java/org/apache/polaris/service/storage/PolarisStorageIntegrationProviderImpl.java index f61c67620..cd7759819 100644 --- a/service/common/src/main/java/org/apache/polaris/service/storage/PolarisStorageIntegrationProviderImpl.java +++ b/service/common/src/main/java/org/apache/polaris/service/storage/PolarisStorageIntegrationProviderImpl.java @@ -30,7 +30,9 @@ import java.util.Map; import java.util.Set; import java.util.function.Supplier; +import org.apache.polaris.core.PolarisConfigurationStore; import org.apache.polaris.core.PolarisDiagnostics; +import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.storage.PolarisCredentialProperty; import org.apache.polaris.core.storage.PolarisStorageActions; import org.apache.polaris.core.storage.PolarisStorageConfigurationInfo; @@ -46,16 +48,24 @@ public class PolarisStorageIntegrationProviderImpl implements PolarisStorageInte private final Supplier stsClientSupplier; private final Supplier gcpCredsProvider; + private final PolarisConfigurationStore configurationStore; @Inject - public PolarisStorageIntegrationProviderImpl(StorageConfiguration storageConfiguration) { - this(storageConfiguration.stsClientSupplier(), storageConfiguration.gcpCredentialsSupplier()); + public PolarisStorageIntegrationProviderImpl( + StorageConfiguration storageConfiguration, PolarisConfigurationStore configurationStore) { + this( + storageConfiguration.stsClientSupplier(), + storageConfiguration.gcpCredentialsSupplier(), + configurationStore); } public PolarisStorageIntegrationProviderImpl( - Supplier stsClientSupplier, Supplier gcpCredsProvider) { + Supplier stsClientSupplier, + Supplier gcpCredsProvider, + PolarisConfigurationStore configurationStore) { this.stsClientSupplier = stsClientSupplier; this.gcpCredsProvider = gcpCredsProvider; + this.configurationStore = configurationStore; } @Override @@ -71,25 +81,28 @@ public PolarisStorageIntegrationProviderImpl( case S3: storageIntegration = (PolarisStorageIntegration) - new AwsCredentialsStorageIntegration(stsClientSupplier.get()); + new AwsCredentialsStorageIntegration(configurationStore, stsClientSupplier.get()); break; case GCS: storageIntegration = (PolarisStorageIntegration) new GcpCredentialsStorageIntegration( + configurationStore, gcpCredsProvider.get(), ServiceOptions.getFromServiceLoader( HttpTransportFactory.class, NetHttpTransport::new)); break; case AZURE: storageIntegration = - (PolarisStorageIntegration) new AzureCredentialsStorageIntegration(); + (PolarisStorageIntegration) + new AzureCredentialsStorageIntegration(configurationStore); break; case FILE: storageIntegration = new PolarisStorageIntegration<>("file") { @Override public EnumMap getSubscopedCreds( + @Nonnull RealmContext realmContext, @Nonnull PolarisDiagnostics diagnostics, @Nonnull T storageConfig, boolean allowListOperation, @@ -101,6 +114,7 @@ public EnumMap getSubscopedCreds( @Override public @Nonnull Map> validateAccessToLocations( + @Nonnull RealmContext realmContext, @Nonnull T storageConfig, @Nonnull Set actions, @Nonnull Set locations) { diff --git a/service/common/src/main/java/org/apache/polaris/service/task/ManifestFileCleanupTaskHandler.java b/service/common/src/main/java/org/apache/polaris/service/task/ManifestFileCleanupTaskHandler.java index 772eb833d..f37400e6c 100644 --- a/service/common/src/main/java/org/apache/polaris/service/task/ManifestFileCleanupTaskHandler.java +++ b/service/common/src/main/java/org/apache/polaris/service/task/ManifestFileCleanupTaskHandler.java @@ -27,7 +27,7 @@ import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.TimeUnit; -import java.util.function.Function; +import java.util.function.BiFunction; import java.util.stream.StreamSupport; import org.apache.commons.codec.binary.Base64; import org.apache.iceberg.DataFile; @@ -36,6 +36,8 @@ import org.apache.iceberg.ManifestReader; import org.apache.iceberg.catalog.TableIdentifier; import org.apache.iceberg.io.FileIO; +import org.apache.polaris.core.PolarisDiagnostics; +import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.AsyncTaskType; import org.apache.polaris.core.entity.TaskEntity; import org.slf4j.Logger; @@ -54,36 +56,40 @@ public class ManifestFileCleanupTaskHandler implements TaskHandler { public static final int FILE_DELETION_RETRY_MILLIS = 100; private static final Logger LOGGER = LoggerFactory.getLogger(ManifestFileCleanupTaskHandler.class); - private final Function fileIOSupplier; + private final BiFunction fileIOSupplier; private final ExecutorService executorService; + private final PolarisDiagnostics diagnostics; public ManifestFileCleanupTaskHandler( - Function fileIOSupplier, ExecutorService executorService) { + BiFunction fileIOSupplier, + ExecutorService executorService, + PolarisDiagnostics diagnostics) { this.fileIOSupplier = fileIOSupplier; this.executorService = executorService; + this.diagnostics = diagnostics; } @Override public boolean canHandleTask(TaskEntity task) { - return task.getTaskType() == AsyncTaskType.MANIFEST_FILE_CLEANUP - || task.getTaskType() == AsyncTaskType.METADATA_FILE_BATCH_CLEANUP; + return task.getTaskType(diagnostics) == AsyncTaskType.MANIFEST_FILE_CLEANUP + || task.getTaskType(diagnostics) == AsyncTaskType.METADATA_FILE_BATCH_CLEANUP; } @Override - public boolean handleTask(TaskEntity task) { - ManifestCleanupTask cleanupTask = task.readData(ManifestCleanupTask.class); + public boolean handleTask(TaskEntity task, RealmContext realmContext) { + ManifestCleanupTask cleanupTask = task.readData(diagnostics, ManifestCleanupTask.class); TableIdentifier tableId = cleanupTask.getTableId(); - try (FileIO authorizedFileIO = fileIOSupplier.apply(task)) { - if (task.getTaskType() == AsyncTaskType.MANIFEST_FILE_CLEANUP) { + try (FileIO authorizedFileIO = fileIOSupplier.apply(task, realmContext)) { + if (task.getTaskType(diagnostics) == AsyncTaskType.MANIFEST_FILE_CLEANUP) { ManifestFile manifestFile = decodeManifestData(cleanupTask.getManifestFileData()); return cleanUpManifestFile(manifestFile, authorizedFileIO, tableId); - } else if (task.getTaskType() == AsyncTaskType.METADATA_FILE_BATCH_CLEANUP) { + } else if (task.getTaskType(diagnostics) == AsyncTaskType.METADATA_FILE_BATCH_CLEANUP) { return cleanUpMetadataFiles(cleanupTask.getMetadataFiles(), authorizedFileIO, tableId); } else { LOGGER .atWarn() .addKeyValue("tableId", tableId) - .log("Unknown task type {}", task.getTaskType()); + .log("Unknown task type {}", task.getTaskType(diagnostics)); return false; } } diff --git a/service/common/src/main/java/org/apache/polaris/service/task/TableCleanupTaskHandler.java b/service/common/src/main/java/org/apache/polaris/service/task/TableCleanupTaskHandler.java index cc722bbb5..e675121e9 100644 --- a/service/common/src/main/java/org/apache/polaris/service/task/TableCleanupTaskHandler.java +++ b/service/common/src/main/java/org/apache/polaris/service/task/TableCleanupTaskHandler.java @@ -18,9 +18,11 @@ */ package org.apache.polaris.service.task; +import java.time.Clock; import java.util.ArrayList; import java.util.List; import java.util.UUID; +import java.util.function.BiFunction; import java.util.function.Function; import java.util.stream.Collectors; import java.util.stream.Stream; @@ -29,8 +31,9 @@ import org.apache.iceberg.TableMetadata; import org.apache.iceberg.TableMetadataParser; import org.apache.iceberg.io.FileIO; -import org.apache.polaris.core.PolarisCallContext; -import org.apache.polaris.core.context.CallContext; +import org.apache.polaris.core.PolarisConfigurationStore; +import org.apache.polaris.core.PolarisDiagnostics; +import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.AsyncTaskType; import org.apache.polaris.core.entity.PolarisBaseEntity; import org.apache.polaris.core.entity.PolarisEntity; @@ -39,6 +42,7 @@ import org.apache.polaris.core.entity.TaskEntity; import org.apache.polaris.core.persistence.MetaStoreManagerFactory; import org.apache.polaris.core.persistence.PolarisMetaStoreManager; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -49,38 +53,50 @@ */ public class TableCleanupTaskHandler implements TaskHandler { private static final Logger LOGGER = LoggerFactory.getLogger(TableCleanupTaskHandler.class); + private static final String BATCH_SIZE_CONFIG_KEY = "TABLE_METADATA_CLEANUP_BATCH_SIZE"; + private final TaskExecutor taskExecutor; private final MetaStoreManagerFactory metaStoreManagerFactory; - private final Function fileIOSupplier; - private static final String BATCH_SIZE_CONFIG_KEY = "TABLE_METADATA_CLEANUP_BATCH_SIZE"; + private final PolarisConfigurationStore configurationStore; + private final PolarisDiagnostics diagnostics; + private final BiFunction fileIOSupplier; + private final Clock clock; public TableCleanupTaskHandler( TaskExecutor taskExecutor, MetaStoreManagerFactory metaStoreManagerFactory, - Function fileIOSupplier) { + PolarisConfigurationStore configurationStore, + PolarisDiagnostics diagnostics, + BiFunction fileIOSupplier, + Clock clock) { this.taskExecutor = taskExecutor; this.metaStoreManagerFactory = metaStoreManagerFactory; + this.configurationStore = configurationStore; + this.diagnostics = diagnostics; this.fileIOSupplier = fileIOSupplier; + this.clock = clock; } @Override public boolean canHandleTask(TaskEntity task) { - return task.getTaskType() == AsyncTaskType.ENTITY_CLEANUP_SCHEDULER && taskEntityIsTable(task); + return task.getTaskType(diagnostics) == AsyncTaskType.ENTITY_CLEANUP_SCHEDULER + && taskEntityIsTable(task); } private boolean taskEntityIsTable(TaskEntity task) { - PolarisEntity entity = PolarisEntity.of((task.readData(PolarisBaseEntity.class))); + PolarisEntity entity = PolarisEntity.of((task.readData(diagnostics, PolarisBaseEntity.class))); return entity.getType().equals(PolarisEntityType.TABLE_LIKE); } @Override - public boolean handleTask(TaskEntity cleanupTask) { - PolarisBaseEntity entity = cleanupTask.readData(PolarisBaseEntity.class); + public boolean handleTask(TaskEntity cleanupTask, RealmContext realmContext) { + PolarisBaseEntity entity = cleanupTask.readData(diagnostics, PolarisBaseEntity.class); PolarisMetaStoreManager metaStoreManager = - metaStoreManagerFactory.getOrCreateMetaStoreManager( - CallContext.getCurrentContext().getRealmContext()); + metaStoreManagerFactory.getOrCreateMetaStoreManager(realmContext); + PolarisMetaStoreSession metaStoreSession = + metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(); + TableLikeEntity tableEntity = TableLikeEntity.of(entity); - PolarisCallContext polarisCallContext = CallContext.getCurrentContext().getPolarisCallContext(); LOGGER .atInfo() .addKeyValue("tableIdentifier", tableEntity.getTableIdentifier()) @@ -90,7 +106,7 @@ public boolean handleTask(TaskEntity cleanupTask) { // It's likely the cleanupTask has already been completed, but wasn't dropped successfully. // Log a // warning and move on - try (FileIO fileIO = fileIOSupplier.apply(cleanupTask)) { + try (FileIO fileIO = fileIOSupplier.apply(cleanupTask, realmContext)) { if (!TaskUtils.exists(tableEntity.getMetadataLocation(), fileIO)) { LOGGER .atWarn() @@ -110,24 +126,27 @@ public boolean handleTask(TaskEntity cleanupTask) { fileIO, tableEntity, metaStoreManager, - polarisCallContext); + metaStoreSession, + clock); // TODO: handle partition statistics files Stream metadataFileCleanupTasks = getMetadataTaskStream( + realmContext, cleanupTask, tableMetadata, - fileIO, tableEntity, metaStoreManager, - polarisCallContext); + metaStoreSession, + configurationStore, + clock); List taskEntities = Stream.concat(manifestCleanupTasks, metadataFileCleanupTasks).toList(); List createdTasks = metaStoreManager - .createEntitiesIfNotExist(polarisCallContext, null, taskEntities) + .createEntitiesIfNotExist(metaStoreSession, null, taskEntities) .getEntities(); if (createdTasks != null) { LOGGER @@ -138,7 +157,7 @@ public boolean handleTask(TaskEntity cleanupTask) { .log( "Successfully queued tasks to delete manifests, previous metadata, and statistics files - deleting table metadata file"); for (PolarisBaseEntity createdTask : createdTasks) { - taskExecutor.addTaskHandlerContext(createdTask.getId(), CallContext.getCurrentContext()); + taskExecutor.addTaskHandlerContext(createdTask.getId(), realmContext); } fileIO.deleteFile(tableEntity.getMetadataLocation()); @@ -155,7 +174,8 @@ private Stream getManifestTaskStream( FileIO fileIO, TableLikeEntity tableEntity, PolarisMetaStoreManager metaStoreManager, - PolarisCallContext polarisCallContext) { + PolarisMetaStoreSession metaStoreSession, + Clock clock) { // read the manifest list for each snapshot. dedupe the manifest files and schedule a // cleanupTask // for each manifest file and its data files to be deleted @@ -187,13 +207,14 @@ private Stream getManifestTaskStream( .log("Queueing task to delete manifest file"); return new TaskEntity.Builder() .setName(taskName) - .setId(metaStoreManager.generateNewEntityId(polarisCallContext).getId()) - .setCreateTimestamp(polarisCallContext.getClock().millis()) - .withTaskType(AsyncTaskType.MANIFEST_FILE_CLEANUP) + .setId(metaStoreManager.generateNewEntityId(metaStoreSession).getId()) + .setCreateTimestamp(clock.millis()) + .withTaskType(diagnostics, AsyncTaskType.MANIFEST_FILE_CLEANUP) .withData( + diagnostics, new ManifestFileCleanupTaskHandler.ManifestCleanupTask( tableEntity.getTableIdentifier(), TaskUtils.encodeManifestFile(mf))) - .setId(metaStoreManager.generateNewEntityId(polarisCallContext).getId()) + .setId(metaStoreManager.generateNewEntityId(metaStoreSession).getId()) // copy the internal properties, which will have storage info .setInternalProperties(cleanupTask.getInternalPropertiesAsMap()) .build(); @@ -201,16 +222,15 @@ private Stream getManifestTaskStream( } private Stream getMetadataTaskStream( + RealmContext realmContext, TaskEntity cleanupTask, TableMetadata tableMetadata, - FileIO fileIO, TableLikeEntity tableEntity, PolarisMetaStoreManager metaStoreManager, - PolarisCallContext polarisCallContext) { - int batchSize = - polarisCallContext - .getConfigurationStore() - .getConfiguration(polarisCallContext, BATCH_SIZE_CONFIG_KEY, 10); + PolarisMetaStoreSession metaStoreSession, + PolarisConfigurationStore configurationStore, + Clock clock) { + int batchSize = configurationStore.getConfiguration(realmContext, BATCH_SIZE_CONFIG_KEY, 10); return getMetadataFileBatches(tableMetadata, batchSize).stream() .map( metadataBatch -> { @@ -229,10 +249,11 @@ private Stream getMetadataTaskStream( "Queueing task to delete metadata files (prev metadata and statistics files)"); return new TaskEntity.Builder() .setName(taskName) - .setId(metaStoreManager.generateNewEntityId(polarisCallContext).getId()) - .setCreateTimestamp(polarisCallContext.getClock().millis()) - .withTaskType(AsyncTaskType.METADATA_FILE_BATCH_CLEANUP) + .setId(metaStoreManager.generateNewEntityId(metaStoreSession).getId()) + .setCreateTimestamp(clock.millis()) + .withTaskType(diagnostics, AsyncTaskType.METADATA_FILE_BATCH_CLEANUP) .withData( + diagnostics, new ManifestFileCleanupTaskHandler.ManifestCleanupTask( tableEntity.getTableIdentifier(), metadataBatch)) .setInternalProperties(cleanupTask.getInternalPropertiesAsMap()) diff --git a/service/common/src/main/java/org/apache/polaris/service/task/TaskExecutor.java b/service/common/src/main/java/org/apache/polaris/service/task/TaskExecutor.java index 016518e6f..bcf2f6169 100644 --- a/service/common/src/main/java/org/apache/polaris/service/task/TaskExecutor.java +++ b/service/common/src/main/java/org/apache/polaris/service/task/TaskExecutor.java @@ -18,12 +18,12 @@ */ package org.apache.polaris.service.task; -import org.apache.polaris.core.context.CallContext; +import org.apache.polaris.core.context.RealmContext; /** * Execute a task asynchronously with a provided context. The context must be cloned so that callers * can close their own context and closables */ public interface TaskExecutor { - void addTaskHandlerContext(long taskEntityId, CallContext callContext); + void addTaskHandlerContext(long taskEntityId, RealmContext realmContext); } diff --git a/service/common/src/main/java/org/apache/polaris/service/task/TaskExecutorImpl.java b/service/common/src/main/java/org/apache/polaris/service/task/TaskExecutorImpl.java index cc084a4d2..ff245dcf4 100644 --- a/service/common/src/main/java/org/apache/polaris/service/task/TaskExecutorImpl.java +++ b/service/common/src/main/java/org/apache/polaris/service/task/TaskExecutorImpl.java @@ -19,6 +19,7 @@ package org.apache.polaris.service.task; import jakarta.annotation.Nonnull; +import java.time.Clock; import java.util.List; import java.util.Map; import java.util.Optional; @@ -27,19 +28,22 @@ import java.util.concurrent.Executor; import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; -import org.apache.polaris.core.context.CallContext; +import org.apache.polaris.core.PolarisConfigurationStore; +import org.apache.polaris.core.PolarisDiagnostics; +import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.PolarisBaseEntity; import org.apache.polaris.core.entity.PolarisEntity; import org.apache.polaris.core.entity.PolarisEntityType; import org.apache.polaris.core.entity.TaskEntity; import org.apache.polaris.core.persistence.MetaStoreManagerFactory; import org.apache.polaris.core.persistence.PolarisMetaStoreManager; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** * Given a list of registered {@link TaskHandler}s, execute tasks asynchronously with the provided - * {@link CallContext}. + * {@link RealmContext}. */ public class TaskExecutorImpl implements TaskExecutor { private static final Logger LOGGER = LoggerFactory.getLogger(TaskExecutorImpl.class); @@ -47,23 +51,34 @@ public class TaskExecutorImpl implements TaskExecutor { private final Executor executor; private final MetaStoreManagerFactory metaStoreManagerFactory; + private final PolarisConfigurationStore configurationStore; + private final PolarisDiagnostics diagnostics; private final TaskFileIOSupplier fileIOSupplier; + private final Clock clock; private final List taskHandlers = new CopyOnWriteArrayList<>(); public TaskExecutorImpl( Executor executor, MetaStoreManagerFactory metaStoreManagerFactory, - TaskFileIOSupplier fileIOSupplier) { + PolarisConfigurationStore configurationStore, + PolarisDiagnostics diagnostics, + TaskFileIOSupplier fileIOSupplier, + Clock clock) { this.executor = executor; this.metaStoreManagerFactory = metaStoreManagerFactory; + this.configurationStore = configurationStore; + this.diagnostics = diagnostics; this.fileIOSupplier = fileIOSupplier; + this.clock = clock; } public void init() { - addTaskHandler(new TableCleanupTaskHandler(this, metaStoreManagerFactory, fileIOSupplier)); + addTaskHandler( + new TableCleanupTaskHandler( + this, metaStoreManagerFactory, configurationStore, diagnostics, fileIOSupplier, clock)); addTaskHandler( new ManifestFileCleanupTaskHandler( - fileIOSupplier, Executors.newVirtualThreadPerTaskExecutor())); + fileIOSupplier, Executors.newVirtualThreadPerTaskExecutor(), diagnostics)); } /** @@ -76,45 +91,39 @@ public void addTaskHandler(TaskHandler taskHandler) { } /** - * Register a {@link CallContext} for a specific task id. That task will be loaded and executed - * asynchronously with a clone of the provided {@link CallContext}. + * Register a {@link RealmContext} for a specific task id. That task will be loaded and executed + * asynchronously with a copy of the provided {@link RealmContext} (because the realm context is a + * request-scoped component). */ @Override - public void addTaskHandlerContext(long taskEntityId, CallContext callContext) { - // Unfortunately CallContext is a request-scoped bean and must be cloned now, - // because its usage inside the TaskExecutor thread pool will outlive its - // lifespan, so the original CallContext will eventually be closed while - // the task is still running. - // Note: PolarisCallContext has request-scoped beans as well, and must be cloned. - // FIXME replace with context propagation? - CallContext clone = CallContext.copyOf(callContext); - tryHandleTask(taskEntityId, clone, null, 1); + public void addTaskHandlerContext(long taskEntityId, RealmContext realmContext) { + tryHandleTask(taskEntityId, RealmContext.copyOf(realmContext), null, 1); } private @Nonnull CompletableFuture tryHandleTask( - long taskEntityId, CallContext callContext, Throwable e, int attempt) { + long taskEntityId, RealmContext realmContext, Throwable e, int attempt) { if (attempt > 3) { return CompletableFuture.failedFuture(e); } return CompletableFuture.runAsync( - () -> handleTask(taskEntityId, callContext, attempt), executor) + () -> handleTask(taskEntityId, realmContext, attempt), executor) .exceptionallyComposeAsync( (t) -> { LOGGER.warn("Failed to handle task entity id {}", taskEntityId, t); - return tryHandleTask(taskEntityId, callContext, t, attempt + 1); + return tryHandleTask(taskEntityId, realmContext, t, attempt + 1); }, CompletableFuture.delayedExecutor( TASK_RETRY_DELAY * (long) attempt, TimeUnit.MILLISECONDS, executor)); } - protected void handleTask(long taskEntityId, CallContext ctx, int attempt) { - // set the call context INSIDE the async task - CallContext.setCurrentContext(ctx); + protected void handleTask(long taskEntityId, RealmContext realmContext, int attempt) { LOGGER.info("Handling task entity id {}", taskEntityId); PolarisMetaStoreManager metaStoreManager = - metaStoreManagerFactory.getOrCreateMetaStoreManager(ctx.getRealmContext()); + metaStoreManagerFactory.getOrCreateMetaStoreManager(realmContext); + PolarisMetaStoreSession metaStoreSession = + metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(); PolarisBaseEntity taskEntity = - metaStoreManager.loadEntity(ctx.getPolarisCallContext(), 0L, taskEntityId).getEntity(); + metaStoreManager.loadEntity(metaStoreSession, 0L, taskEntityId).getEntity(); if (!PolarisEntityType.TASK.equals(taskEntity.getType())) { throw new IllegalArgumentException("Provided taskId must be a task entity type"); } @@ -125,12 +134,12 @@ protected void handleTask(long taskEntityId, CallContext ctx, int attempt) { LOGGER .atWarn() .addKeyValue("taskEntityId", taskEntityId) - .addKeyValue("taskType", task.getTaskType()) + .addKeyValue("taskType", task.getTaskType(diagnostics)) .log("Unable to find handler for task type"); return; } TaskHandler handler = handlerOpt.get(); - boolean success = handler.handleTask(task); + boolean success = handler.handleTask(task, realmContext); if (success) { LOGGER .atInfo() @@ -138,7 +147,7 @@ protected void handleTask(long taskEntityId, CallContext ctx, int attempt) { .addKeyValue("handlerClass", handler.getClass()) .log("Task successfully handled"); metaStoreManager.dropEntityIfExists( - ctx.getPolarisCallContext(), null, PolarisEntity.toCore(taskEntity), Map.of(), false); + metaStoreSession, null, PolarisEntity.toCore(taskEntity), Map.of(), false); } else { LOGGER .atWarn() diff --git a/service/common/src/main/java/org/apache/polaris/service/task/TaskFileIOSupplier.java b/service/common/src/main/java/org/apache/polaris/service/task/TaskFileIOSupplier.java index e926e2d60..b11a6d8db 100644 --- a/service/common/src/main/java/org/apache/polaris/service/task/TaskFileIOSupplier.java +++ b/service/common/src/main/java/org/apache/polaris/service/task/TaskFileIOSupplier.java @@ -23,20 +23,21 @@ import java.util.HashMap; import java.util.Map; import java.util.Set; -import java.util.function.Function; +import java.util.function.BiFunction; import org.apache.iceberg.CatalogProperties; import org.apache.iceberg.io.FileIO; import org.apache.polaris.core.PolarisConfiguration; import org.apache.polaris.core.PolarisConfigurationStore; -import org.apache.polaris.core.context.CallContext; +import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.PolarisTaskConstants; import org.apache.polaris.core.entity.TaskEntity; import org.apache.polaris.core.persistence.MetaStoreManagerFactory; import org.apache.polaris.core.persistence.PolarisMetaStoreManager; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; import org.apache.polaris.service.catalog.io.FileIOFactory; @ApplicationScoped -public class TaskFileIOSupplier implements Function { +public class TaskFileIOSupplier implements BiFunction { private final MetaStoreManagerFactory metaStoreManagerFactory; private final FileIOFactory fileIOFactory; private final PolarisConfigurationStore configurationStore; @@ -52,27 +53,28 @@ public TaskFileIOSupplier( } @Override - public FileIO apply(TaskEntity task) { + public FileIO apply(TaskEntity task, RealmContext realmContext) { Map internalProperties = task.getInternalPropertiesAsMap(); String location = internalProperties.get(PolarisTaskConstants.STORAGE_LOCATION); PolarisMetaStoreManager metaStoreManager = - metaStoreManagerFactory.getOrCreateMetaStoreManager( - CallContext.getCurrentContext().getRealmContext()); + metaStoreManagerFactory.getOrCreateMetaStoreManager(realmContext); + PolarisMetaStoreSession metaStoreSession = + metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(); Map properties = new HashMap<>(internalProperties); Boolean skipCredentialSubscopingIndirection = configurationStore.getConfiguration( - null, + realmContext, PolarisConfiguration.SKIP_CREDENTIAL_SUBSCOPING_INDIRECTION.key, PolarisConfiguration.SKIP_CREDENTIAL_SUBSCOPING_INDIRECTION.defaultValue); if (!skipCredentialSubscopingIndirection) { properties.putAll( metaStoreManagerFactory - .getOrCreateStorageCredentialCache(CallContext.getCurrentContext().getRealmContext()) + .getOrCreateStorageCredentialCache(realmContext) .getOrGenerateSubScopeCreds( metaStoreManager, - CallContext.getCurrentContext().getPolarisCallContext(), + metaStoreSession, task, true, Set.of(location), diff --git a/service/common/src/main/java/org/apache/polaris/service/task/TaskHandler.java b/service/common/src/main/java/org/apache/polaris/service/task/TaskHandler.java index f903ddf80..f0d331494 100644 --- a/service/common/src/main/java/org/apache/polaris/service/task/TaskHandler.java +++ b/service/common/src/main/java/org/apache/polaris/service/task/TaskHandler.java @@ -18,10 +18,11 @@ */ package org.apache.polaris.service.task; +import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.entity.TaskEntity; public interface TaskHandler { boolean canHandleTask(TaskEntity task); - boolean handleTask(TaskEntity task); + boolean handleTask(TaskEntity task, RealmContext realmContext); } diff --git a/service/common/src/test/java/org/apache/polaris/service/auth/BasePolarisAuthenticatorTest.java b/service/common/src/test/java/org/apache/polaris/service/auth/BasePolarisAuthenticatorTest.java index b2511f7ca..7f4f031a8 100644 --- a/service/common/src/test/java/org/apache/polaris/service/auth/BasePolarisAuthenticatorTest.java +++ b/service/common/src/test/java/org/apache/polaris/service/auth/BasePolarisAuthenticatorTest.java @@ -23,13 +23,10 @@ import java.util.Optional; import org.apache.iceberg.exceptions.NotAuthorizedException; import org.apache.iceberg.exceptions.ServiceFailureException; -import org.apache.polaris.core.PolarisCallContext; import org.apache.polaris.core.auth.AuthenticatedPolarisPrincipal; -import org.apache.polaris.core.context.CallContext; -import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.core.persistence.BaseResult; -import org.apache.polaris.core.persistence.MetaStoreManagerFactory; import org.apache.polaris.core.persistence.PolarisMetaStoreManager; +import org.apache.polaris.core.persistence.PolarisMetaStoreSession; import org.assertj.core.api.Assertions; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -39,20 +36,14 @@ public class BasePolarisAuthenticatorTest { private BasePolarisAuthenticator authenticator; private PolarisMetaStoreManager metaStoreManager; - private PolarisCallContext polarisCallContext; - private CallContext callContext; + private PolarisMetaStoreSession metaStoreSession; @BeforeEach public void setUp() { - RealmContext realmContext = () -> "test"; - polarisCallContext = Mockito.mock(PolarisCallContext.class); - callContext = CallContext.of(realmContext, polarisCallContext); metaStoreManager = Mockito.mock(PolarisMetaStoreManager.class); - MetaStoreManagerFactory metaStoreManagerFactory = Mockito.mock(MetaStoreManagerFactory.class); - when(metaStoreManagerFactory.getOrCreateMetaStoreManager(realmContext)) - .thenReturn(metaStoreManager); + metaStoreSession = Mockito.mock(PolarisMetaStoreSession.class); authenticator = - new BasePolarisAuthenticator(metaStoreManagerFactory, callContext) { + new BasePolarisAuthenticator(metaStoreManager, metaStoreSession) { @Override public Optional authenticate(String credentials) { return Optional.empty(); @@ -65,7 +56,7 @@ public void testFetchPrincipalThrowsServiceExceptionOnMetastoreException() { DecodedToken token = Mockito.mock(DecodedToken.class); long principalId = 100L; when(token.getPrincipalId()).thenReturn(principalId); - when(metaStoreManager.loadEntity(polarisCallContext, 0L, principalId)) + when(metaStoreManager.loadEntity(metaStoreSession, 0L, principalId)) .thenThrow(new RuntimeException("Metastore exception")); Assertions.assertThatThrownBy(() -> authenticator.getPrincipal(token)) @@ -79,7 +70,7 @@ public void testFetchPrincipalThrowsNotAuthorizedWhenNotFound() { long principalId = 100L; when(token.getPrincipalId()).thenReturn(principalId); when(token.getClientId()).thenReturn("abc"); - when(metaStoreManager.loadEntity(polarisCallContext, 0L, principalId)) + when(metaStoreManager.loadEntity(metaStoreSession, 0L, principalId)) .thenReturn( new PolarisMetaStoreManager.EntityResult(BaseResult.ReturnStatus.ENTITY_NOT_FOUND, ""));