diff --git a/dropwizard/service/src/main/java/org/apache/polaris/service/dropwizard/auth/QuarkusOAuthFilter.java b/dropwizard/service/src/main/java/org/apache/polaris/service/dropwizard/auth/QuarkusOAuthFilter.java index 50c56640b..e75ca3956 100644 --- a/dropwizard/service/src/main/java/org/apache/polaris/service/dropwizard/auth/QuarkusOAuthFilter.java +++ b/dropwizard/service/src/main/java/org/apache/polaris/service/dropwizard/auth/QuarkusOAuthFilter.java @@ -27,7 +27,7 @@ import java.security.Principal; import java.util.Optional; 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.service.auth.Authenticator; import org.checkerframework.checker.nullness.qual.Nullable; import org.jboss.resteasy.reactive.server.ServerRequestFilter; @@ -50,7 +50,7 @@ public class QuarkusOAuthFilter { public static final String OAUTH_ACCESS_TOKEN_PARAM = "access_token"; @Inject Authenticator authenticator; - @Inject CallContext callContext; + @Inject RealmContext realmContext; @ServerRequestFilter(priority = Priorities.AUTHENTICATION) public void authenticate(ContainerRequestContext requestContext) { @@ -71,7 +71,7 @@ public void authenticate(ContainerRequestContext requestContext) { if (!authenticate(requestContext, credentials, SecurityContext.BASIC_AUTH)) { throw new NotAuthorizedException( "Credentials are required to access this resource.", - String.format(CHALLENGE_FORMAT, callContext.getRealmContext().getRealmIdentifier())); + String.format(CHALLENGE_FORMAT, realmContext.getRealmIdentifier())); } } @@ -115,8 +115,6 @@ protected boolean authenticate( return false; } - CallContext.setCurrentContext(callContext); - Optional principal = authenticator.authenticate(credentials); if (principal.isEmpty()) { return false; diff --git a/dropwizard/service/src/main/java/org/apache/polaris/service/dropwizard/config/QuarkusProducers.java b/dropwizard/service/src/main/java/org/apache/polaris/service/dropwizard/config/QuarkusProducers.java index 6ffb10599..39eed70bb 100644 --- a/dropwizard/service/src/main/java/org/apache/polaris/service/dropwizard/config/QuarkusProducers.java +++ b/dropwizard/service/src/main/java/org/apache/polaris/service/dropwizard/config/QuarkusProducers.java @@ -37,7 +37,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.persistence.MetaStoreManagerFactory; import org.apache.polaris.core.persistence.PolarisEntityManager; @@ -107,16 +106,6 @@ public RealmContext realmContext( .collect(HashMap::new, (m, e) -> m.put(e.getKey(), e.getValue()), HashMap::putAll)); } - @Produces - @RequestScoped - public CallContext callContext(RealmContext realmContext) { - return CallContext.of(realmContext); - } - - public void closeCallContext(@Disposes CallContext callContext) { - callContext.close(); - } - @Produces @RequestScoped public PolarisMetaStoreSession metaStoreSession( diff --git a/dropwizard/service/src/main/java/org/apache/polaris/service/dropwizard/task/QuarkusTaskExecutorImpl.java b/dropwizard/service/src/main/java/org/apache/polaris/service/dropwizard/task/QuarkusTaskExecutorImpl.java index 2aaf1964c..5fe868a53 100644 --- a/dropwizard/service/src/main/java/org/apache/polaris/service/dropwizard/task/QuarkusTaskExecutorImpl.java +++ b/dropwizard/service/src/main/java/org/apache/polaris/service/dropwizard/task/QuarkusTaskExecutorImpl.java @@ -30,7 +30,7 @@ import java.util.concurrent.ExecutorService; 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.service.dropwizard.tracing.QuarkusTracingFilter; import org.apache.polaris.service.task.TaskExecutorImpl; @@ -71,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/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/admin/PolarisAdminServiceAuthzTest.java b/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/admin/PolarisAdminServiceAuthzTest.java index 8852c3929..467246905 100644 --- a/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/admin/PolarisAdminServiceAuthzTest.java +++ b/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/admin/PolarisAdminServiceAuthzTest.java @@ -49,6 +49,7 @@ private PolarisAdminService newTestAdminService(Set activatedPrincipalRo final AuthenticatedPolarisPrincipal authenticatedPrincipal = new AuthenticatedPolarisPrincipal(principalEntity, activatedPrincipalRoles); return new PolarisAdminService( + realmContext, entityManager, metaStoreManager, metaStoreSession, diff --git a/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/admin/PolarisAuthzTestBase.java b/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/admin/PolarisAuthzTestBase.java index c197bfded..362e6ab66 100644 --- a/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/admin/PolarisAuthzTestBase.java +++ b/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/admin/PolarisAuthzTestBase.java @@ -52,7 +52,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; @@ -163,7 +162,6 @@ public Map getConfigOverrides() { @Inject protected MetaStoreManagerFactory managerFactory; @Inject protected RealmEntityManagerFactory realmEntityManagerFactory; - // @Inject protected CallContextCatalogFactory callContextCatalogFactory; @Inject protected PolarisConfigurationStore configurationStore; @Inject protected PolarisDiagnostics diagServices; @@ -174,7 +172,7 @@ public Map getConfigOverrides() { protected PolarisMetaStoreSession metaStoreSession; protected PolarisBaseEntity catalogEntity; protected PrincipalEntity principalEntity; - protected CallContext callContext; + protected RealmContext realmContext; protected AuthenticatedPolarisPrincipal authenticatedRoot; @BeforeAll @@ -189,14 +187,11 @@ public static void setUpMocks() { @BeforeEach public void before(TestInfo testInfo) { - RealmContext realmContext = testInfo::getDisplayName; + realmContext = testInfo::getDisplayName; metaStoreManager = managerFactory.getOrCreateMetaStoreManager(realmContext); metaStoreSession = managerFactory.getOrCreateSessionSupplier(realmContext).get(); entityManager = realmEntityManagerFactory.getOrCreateEntityManager(realmContext); - callContext = CallContext.of(realmContext); - CallContext.setCurrentContext(callContext); - PrincipalEntity rootEntity = new PrincipalEntity( PolarisEntity.of( @@ -213,6 +208,7 @@ public void before(TestInfo testInfo) { this.adminService = new PolarisAdminService( + realmContext, entityManager, metaStoreManager, metaStoreSession, @@ -373,12 +369,12 @@ private void initBaseCatalog() { entityManager, metaStoreSession, authenticatedRoot, CATALOG_NAME); this.baseCatalog = new BasePolarisCatalog( + realmContext, entityManager, metaStoreManager, metaStoreSession, configurationStore, diagServices, - callContext, passthroughView, authenticatedRoot, Mockito.mock(), @@ -419,13 +415,14 @@ public TestPolarisCallContextCatalogFactory( @Override public Catalog createCallContextCatalog( - CallContext context, + RealmContext realmContext, AuthenticatedPolarisPrincipal authenticatedPolarisPrincipal, final PolarisResolutionManifest resolvedManifest) { // This depends on the BasePolarisCatalog allowing calling initialize multiple times // to override the previous config. Catalog catalog = - super.createCallContextCatalog(context, authenticatedPolarisPrincipal, resolvedManifest); + super.createCallContextCatalog( + realmContext, authenticatedPolarisPrincipal, resolvedManifest); catalog.initialize( CATALOG_NAME, ImmutableMap.of( diff --git a/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/auth/JWTSymmetricKeyGeneratorTest.java b/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/auth/JWTSymmetricKeyGeneratorTest.java index c47e0df11..58ddb373d 100644 --- a/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/auth/JWTSymmetricKeyGeneratorTest.java +++ b/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/auth/JWTSymmetricKeyGeneratorTest.java @@ -24,7 +24,6 @@ import com.auth0.jwt.JWTVerifier; import com.auth0.jwt.algorithms.Algorithm; import com.auth0.jwt.interfaces.DecodedJWT; -import org.apache.polaris.core.context.CallContext; import org.apache.polaris.core.entity.PolarisBaseEntity; import org.apache.polaris.core.entity.PolarisEntitySubType; import org.apache.polaris.core.entity.PolarisEntityType; @@ -43,7 +42,6 @@ public class JWTSymmetricKeyGeneratorTest { /** Sanity test to verify that we can generate a token */ @Test public void testJWTSymmetricKeyGenerator() { - CallContext.setCurrentContext(CallContext.of(() -> "realm")); PolarisMetaStoreManager metastoreManager = Mockito.mock(PolarisMetaStoreManager.class); PolarisMetaStoreSession metaStoreSession = Mockito.mock(PolarisMetaStoreSession.class); String mainSecret = "test_secret"; diff --git a/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/catalog/BasePolarisCatalogTest.java b/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/catalog/BasePolarisCatalogTest.java index abf6c5037..2694aa9da 100644 --- a/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/catalog/BasePolarisCatalogTest.java +++ b/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/catalog/BasePolarisCatalogTest.java @@ -71,7 +71,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; @@ -139,7 +138,6 @@ public class BasePolarisCatalogTest extends CatalogTests { @Inject Clock clock; private BasePolarisCatalog catalog; - private CallContext callContext; private RealmContext realmContext; private PolarisMetaStoreManager metaStoreManager; private PolarisMetaStoreSession metaStoreSession; @@ -168,9 +166,6 @@ public void before(TestInfo testInfo) { metaStoreSession = managerFactory.getOrCreateSessionSupplier(realmContext).get(); entityManager = entityManagerFactory.getOrCreateEntityManager(realmContext); - callContext = CallContext.of(realmContext); - CallContext.setCurrentContext(callContext); - PrincipalEntity rootEntity = new PrincipalEntity( PolarisEntity.of( @@ -187,6 +182,7 @@ public void before(TestInfo testInfo) { adminService = new PolarisAdminService( + realmContext, entityManager, metaStoreManager, metaStoreSession, @@ -221,12 +217,12 @@ public void before(TestInfo testInfo) { TaskExecutor taskExecutor = Mockito.mock(); this.catalog = new BasePolarisCatalog( + realmContext, entityManager, metaStoreManager, metaStoreSession, configurationStore, diagServices, - callContext, passthroughView, authenticatedRoot, taskExecutor, @@ -490,12 +486,12 @@ public void testValidateNotificationFailToCreateFileIO() throws IOException { FileIOFactory fileIoFactory = spy(new DefaultFileIOFactory()); BasePolarisCatalog catalog = new BasePolarisCatalog( + realmContext, entityManager, metaStoreManager, metaStoreSession, configurationStore, diagServices, - callContext, passthroughView, authenticatedRoot, Mockito.mock(), @@ -806,26 +802,24 @@ 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( entityManager, metaStoreSession, authenticatedRoot, catalogWithoutStorage); TaskExecutor taskExecutor = Mockito.mock(); BasePolarisCatalog catalog = new BasePolarisCatalog( + realmContext, entityManager, metaStoreManager, metaStoreSession, configurationStore, diagServices, - callContext, passthroughView, authenticatedRoot, taskExecutor, @@ -854,7 +848,7 @@ public void testUpdateNotificationCreateTableWithLocalFilePrefix() { TableMetadataParser.toJson(createSampleTableMetadata(metadataLocation)).getBytes(UTF_8)); if (!configurationStore - .getConfiguration(PolarisConfiguration.SUPPORTED_CATALOG_STORAGE_TYPES) + .getConfiguration(realmContext, PolarisConfiguration.SUPPORTED_CATALOG_STORAGE_TYPES) .contains("FILE")) { Assertions.assertThatThrownBy(() -> catalog.sendNotification(table, request)) .isInstanceOf(ForbiddenException.class) @@ -879,19 +873,18 @@ public void testUpdateNotificationCreateTableWithHttpPrefix() { .setName(catalogName) .build()); - CallContext callContext = CallContext.getCurrentContext(); PolarisPassthroughResolutionView passthroughView = new PolarisPassthroughResolutionView( entityManager, metaStoreSession, authenticatedRoot, catalogName); TaskExecutor taskExecutor = Mockito.mock(); BasePolarisCatalog catalog = new BasePolarisCatalog( + realmContext, entityManager, metaStoreManager, metaStoreSession, configurationStore, diagServices, - callContext, passthroughView, authenticatedRoot, taskExecutor, @@ -922,7 +915,7 @@ public void testUpdateNotificationCreateTableWithHttpPrefix() { TableMetadataParser.toJson(createSampleTableMetadata(metadataLocation)).getBytes(UTF_8)); if (!configurationStore - .getConfiguration(PolarisConfiguration.SUPPORTED_CATALOG_STORAGE_TYPES) + .getConfiguration(realmContext, PolarisConfiguration.SUPPORTED_CATALOG_STORAGE_TYPES) .contains("FILE")) { Assertions.assertThatThrownBy(() -> catalog.sendNotification(table, request)) .isInstanceOf(ForbiddenException.class) @@ -942,7 +935,7 @@ public void testUpdateNotificationCreateTableWithHttpPrefix() { TableMetadataParser.toJson(createSampleTableMetadata(metadataLocation)).getBytes(UTF_8)); if (!configurationStore - .getConfiguration(PolarisConfiguration.SUPPORTED_CATALOG_STORAGE_TYPES) + .getConfiguration(realmContext, PolarisConfiguration.SUPPORTED_CATALOG_STORAGE_TYPES) .contains("FILE")) { Assertions.assertThatThrownBy(() -> catalog.sendNotification(table, newRequest)) .isInstanceOf(ForbiddenException.class) @@ -1417,19 +1410,17 @@ public void testDropTableWithPurgeDisabled() { .addProperty(PolarisConfiguration.DROP_WITH_PURGE_ENABLED.catalogConfig(), "false") .setStorageConfigurationInfo(noPurgeStorageConfigModel, storageLocation) .build()); - realmContext = () -> "realm"; - CallContext callContext = CallContext.of(realmContext); PolarisPassthroughResolutionView passthroughView = new PolarisPassthroughResolutionView( entityManager, metaStoreSession, authenticatedRoot, noPurgeCatalogName); BasePolarisCatalog noPurgeCatalog = new BasePolarisCatalog( + realmContext, entityManager, metaStoreManager, metaStoreSession, configurationStore, diagServices, - callContext, passthroughView, authenticatedRoot, Mockito.mock(), @@ -1501,8 +1492,6 @@ public void testRetriableException() { @Test public void testFileIOWrapper() { - CallContext callContext = CallContext.of(realmContext); - CallContext.setCurrentContext(callContext); PolarisPassthroughResolutionView passthroughView = new PolarisPassthroughResolutionView( entityManager, metaStoreSession, authenticatedRoot, CATALOG_NAME); @@ -1510,12 +1499,12 @@ public void testFileIOWrapper() { TestFileIOFactory measured = new TestFileIOFactory(); BasePolarisCatalog catalog = new BasePolarisCatalog( + realmContext, entityManager, metaStoreManager, metaStoreSession, configurationStore, diagServices, - callContext, passthroughView, authenticatedRoot, Mockito.mock(), diff --git a/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/catalog/BasePolarisCatalogViewTest.java b/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/catalog/BasePolarisCatalogViewTest.java index d8fe7f96c..aaa9ed6e5 100644 --- a/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/catalog/BasePolarisCatalogViewTest.java +++ b/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/catalog/BasePolarisCatalogViewTest.java @@ -38,7 +38,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; @@ -101,9 +100,6 @@ public void before(TestInfo testInfo) { metaStoreManager = managerFactory.getOrCreateMetaStoreManager(realmContext); metaStoreSession = managerFactory.getOrCreateSessionSupplier(realmContext).get(); - CallContext callContext = CallContext.of(realmContext); - CallContext.setCurrentContext(callContext); - PrincipalEntity rootEntity = new PrincipalEntity( PolarisEntity.of( @@ -123,6 +119,7 @@ public void before(TestInfo testInfo) { PolarisAdminService adminService = new PolarisAdminService( + realmContext, entityManager, metaStoreManager, metaStoreSession, @@ -147,12 +144,12 @@ public void before(TestInfo testInfo) { entityManager, metaStoreSession, authenticatedRoot, CATALOG_NAME); this.catalog = new BasePolarisCatalog( + realmContext, entityManager, metaStoreManager, metaStoreSession, configurationStore, diagServices, - callContext, passthroughView, authenticatedRoot, Mockito.mock(), diff --git a/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/catalog/PolarisCatalogHandlerWrapperAuthzTest.java b/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/catalog/PolarisCatalogHandlerWrapperAuthzTest.java index d1e7983a6..ab96dc751 100644 --- a/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/catalog/PolarisCatalogHandlerWrapperAuthzTest.java +++ b/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/catalog/PolarisCatalogHandlerWrapperAuthzTest.java @@ -53,7 +53,7 @@ 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; @@ -97,7 +97,7 @@ private PolarisCatalogHandlerWrapper newWrapper( final AuthenticatedPolarisPrincipal authenticatedPrincipal = new AuthenticatedPolarisPrincipal(principalEntity, activatedPrincipalRoles); return new PolarisCatalogHandlerWrapper( - callContext, + realmContext, metaStoreSession, configurationStore, diagServices, @@ -112,7 +112,7 @@ private PolarisCatalogHandlerWrapper newWrapper( private PolarisCatalogHandlerWrapper newWrapper( AuthenticatedPolarisPrincipal authenticatedPrincipal) { return new PolarisCatalogHandlerWrapper( - callContext, + realmContext, metaStoreSession, configurationStore, diagServices, @@ -1704,12 +1704,12 @@ public void testSendNotificationSufficientPrivileges() { new DefaultFileIOFactory()) { @Override public Catalog createCallContextCatalog( - CallContext context, + RealmContext realmContext, AuthenticatedPolarisPrincipal authenticatedPolarisPrincipal, PolarisResolutionManifest resolvedManifest) { Catalog catalog = super.createCallContextCatalog( - context, authenticatedPolarisPrincipal, resolvedManifest); + realmContext, authenticatedPolarisPrincipal, resolvedManifest); String fileIoImpl = "org.apache.iceberg.inmemory.InMemoryFileIO"; catalog.initialize( externalCatalog, ImmutableMap.of(CatalogProperties.FILE_IO_IMPL, fileIoImpl)); diff --git a/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/config/DefaultConfigurationStoreTest.java b/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/config/DefaultConfigurationStoreTest.java index a605e5f67..dc16f2d32 100644 --- a/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/config/DefaultConfigurationStoreTest.java +++ b/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/config/DefaultConfigurationStoreTest.java @@ -21,7 +21,7 @@ import static org.assertj.core.api.Assertions.assertThat; import java.util.Map; -import org.apache.polaris.core.context.CallContext; +import org.apache.polaris.core.context.RealmContext; import org.apache.polaris.service.config.DefaultConfigurationStore; import org.junit.jupiter.api.Test; @@ -31,14 +31,17 @@ public class DefaultConfigurationStoreTest { public void testGetConfiguration() { DefaultConfigurationStore defaultConfigurationStore = new DefaultConfigurationStore(Map.of("key1", 1, "key2", "value")); - Object value = defaultConfigurationStore.getConfiguration("missingKeyWithoutDefault"); + RealmContext realmContext = () -> "test"; + Object value = + defaultConfigurationStore.getConfiguration(realmContext, "missingKeyWithoutDefault"); assertThat(value).isNull(); Object defaultValue = - defaultConfigurationStore.getConfiguration("missingKeyWithDefault", "defaultValue"); + defaultConfigurationStore.getConfiguration( + realmContext, "missingKeyWithDefault", "defaultValue"); assertThat(defaultValue).isEqualTo("defaultValue"); - Integer keyOne = defaultConfigurationStore.getConfiguration("key1"); + Integer keyOne = defaultConfigurationStore.getConfiguration(realmContext, "key1"); assertThat(keyOne).isEqualTo(1); - String keyTwo = defaultConfigurationStore.getConfiguration("key2"); + String keyTwo = defaultConfigurationStore.getConfiguration(realmContext, "key2"); assertThat(keyTwo).isEqualTo("value"); } @@ -60,29 +63,31 @@ public void testGetRealmConfiguration() { Map.of("key1", realm2KeyOneValue, "key2", realm2KeyTwoValue))); // check realm1 values - CallContext.setCurrentContext(CallContext.of(() -> "realm1")); - Object value = defaultConfigurationStore.getConfiguration("missingKeyWithoutDefault"); + RealmContext realmContext = () -> "realm1"; + Object value = + defaultConfigurationStore.getConfiguration(realmContext, "missingKeyWithoutDefault"); assertThat(value).isNull(); Object defaultValue = - defaultConfigurationStore.getConfiguration("missingKeyWithDefault", "defaultValue"); + defaultConfigurationStore.getConfiguration( + realmContext, "missingKeyWithDefault", "defaultValue"); assertThat(defaultValue).isEqualTo("defaultValue"); - Integer keyOneRealm1 = defaultConfigurationStore.getConfiguration("key1"); + Integer keyOneRealm1 = defaultConfigurationStore.getConfiguration(realmContext, "key1"); assertThat(keyOneRealm1).isEqualTo(realm1KeyOneValue); - String keyTwoRealm1 = defaultConfigurationStore.getConfiguration("key2"); + String keyTwoRealm1 = defaultConfigurationStore.getConfiguration(realmContext, "key2"); assertThat(keyTwoRealm1).isEqualTo(defaultKeyTwoValue); // check realm2 values - CallContext.setCurrentContext(CallContext.of(() -> "realm2")); - Integer keyOneRealm2 = defaultConfigurationStore.getConfiguration("key1"); + realmContext = () -> "realm2"; + Integer keyOneRealm2 = defaultConfigurationStore.getConfiguration(realmContext, "key1"); assertThat(keyOneRealm2).isEqualTo(realm2KeyOneValue); - String keyTwoRealm2 = defaultConfigurationStore.getConfiguration("key2"); + String keyTwoRealm2 = defaultConfigurationStore.getConfiguration(realmContext, "key2"); assertThat(keyTwoRealm2).isEqualTo(realm2KeyTwoValue); // realm3 has no realm-overrides, so just returns default values - CallContext.setCurrentContext(CallContext.of(() -> "realm3")); - Integer keyOneRealm3 = defaultConfigurationStore.getConfiguration("key1"); + realmContext = () -> "realm3"; + Integer keyOneRealm3 = defaultConfigurationStore.getConfiguration(realmContext, "key1"); assertThat(keyOneRealm3).isEqualTo(defaultKeyOneValue); - String keyTwoRealm3 = defaultConfigurationStore.getConfiguration("key2"); + String keyTwoRealm3 = defaultConfigurationStore.getConfiguration(realmContext, "key2"); assertThat(keyTwoRealm3).isEqualTo(defaultKeyTwoValue); } } diff --git a/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/task/ManifestFileCleanupTaskHandlerTest.java b/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/task/ManifestFileCleanupTaskHandlerTest.java index cc4997e1d..fd836f29b 100644 --- a/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/task/ManifestFileCleanupTaskHandlerTest.java +++ b/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/task/ManifestFileCleanupTaskHandlerTest.java @@ -46,7 +46,6 @@ import org.apache.iceberg.io.OutputFile; import org.apache.iceberg.io.PositionOutputStream; 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.entity.AsyncTaskType; import org.apache.polaris.core.entity.TaskEntity; @@ -63,382 +62,346 @@ class ManifestFileCleanupTaskHandlerTest { @Test public void testCleanupFileNotExists() throws IOException { - try (CallContext callCtx = CallContext.of(realmContext)) { - CallContext.setCurrentContext(callCtx); - 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(); - } + 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 { - try (CallContext callCtx = CallContext.of(realmContext)) { - CallContext.setCurrentContext(callCtx); - 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(); - } + 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 { - try (CallContext callCtx = CallContext.of(realmContext)) { - 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, 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(); - } + 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 { - try (CallContext callCtx = CallContext.of(realmContext)) { - 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, 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(); - } + 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 { - try (CallContext callCtx = CallContext.of(realmContext)) { - 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, rc) -> fileIO, Executors.newSingleThreadExecutor(), diagnostics); + 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(metadataLogEntry -> metadataLogEntry.file()) - .filter(file -> TaskUtils.exists(file, fileIO)), - secondMetadata.statisticsFiles().stream() - .map(statisticsFile -> 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(diagnostics, AsyncTaskType.METADATA_FILE_BATCH_CLEANUP) - .withData( - diagnostics, - 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(); - assertThat(handler.canHandleTask(task)).isTrue(); - assertThat(handler.handleTask(task, realmContext)).isTrue(); + assertThat(handler.canHandleTask(task)).isTrue(); + assertThat(handler.handleTask(task, realmContext)).isTrue(); - assertThat(TaskUtils.exists(firstMetadataFile, fileIO)).isFalse(); - assertThat(TaskUtils.exists(statisticsFile1.path(), fileIO)).isFalse(); - assertThat(TaskUtils.exists(statisticsFile2.path(), fileIO)).isFalse(); - } + 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 { - try (CallContext callCtx = CallContext.of(realmContext)) { - CallContext.setCurrentContext(callCtx); - 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 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(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(); - } + 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 { - try (CallContext callCtx = CallContext.of(realmContext)) { - 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, 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(); + } + }; + 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(diagnostics, AsyncTaskType.METADATA_FILE_BATCH_CLEANUP) - .withData( - diagnostics, - 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( - () -> { - 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(); - } + 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/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/task/TableCleanupTaskHandlerTest.java b/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/task/TableCleanupTaskHandlerTest.java index eac80573d..8ff6f143e 100644 --- a/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/task/TableCleanupTaskHandlerTest.java +++ b/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/task/TableCleanupTaskHandlerTest.java @@ -38,7 +38,6 @@ import org.apache.iceberg.io.FileIO; 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.entity.AsyncTaskType; import org.apache.polaris.core.entity.PolarisBaseEntity; @@ -67,580 +66,548 @@ class TableCleanupTaskHandlerTest { public void testTableCleanup() throws IOException { PolarisMetaStoreSession metaStoreSession = metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(); - try (CallContext callCtx = CallContext.of(realmContext)) { - CallContext.setCurrentContext(callCtx); - 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.toString()) - .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); - - CallContext.setCurrentContext(CallContext.of(realmContext)); - 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))); - } + 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 { PolarisMetaStoreSession metaStoreSession = metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(); - try (CallContext callCtx = CallContext.of(realmContext)) { - 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, - 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.toString()) - .withTaskType(diagnostics, AsyncTaskType.ENTITY_CLEANUP_SCHEDULER) - .withData(diagnostics, tableLikeEntity) - .build(); - Assertions.assertThatPredicate(handler::canHandleTask).accepts(task); - - CallContext.setCurrentContext(CallContext.of(realmContext)); - - // 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); - } + 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 { PolarisMetaStoreSession metaStoreSession = metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(); - try (CallContext callCtx = CallContext.of(realmContext)) { - 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, - 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.toString()) - .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); - - CallContext.setCurrentContext(CallContext.of(realmContext)); - - // 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))); - } + 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 { PolarisMetaStoreSession metaStoreSession = metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(); - try (CallContext callCtx = CallContext.of(realmContext)) { - CallContext.setCurrentContext(callCtx); - 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); - - CallContext.setCurrentContext(CallContext.of(realmContext)); - - 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))); - } + 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 { PolarisMetaStoreSession metaStoreSession = metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(); - try (CallContext callCtx = CallContext.of(realmContext)) { - CallContext.setCurrentContext(callCtx); - 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); - - CallContext.setCurrentContext(CallContext.of(realmContext)); - - 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))); - } + 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/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/test/PolarisIntegrationTestFixture.java b/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/test/PolarisIntegrationTestFixture.java index 88e00c4c2..ae9c9e427 100644 --- a/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/test/PolarisIntegrationTestFixture.java +++ b/dropwizard/service/src/test/java/org/apache/polaris/service/dropwizard/test/PolarisIntegrationTestFixture.java @@ -35,7 +35,6 @@ 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; @@ -103,25 +102,20 @@ private PolarisPrincipalSecrets fetchAdminSecrets() { PolarisMetaStoreSession metaStoreSession = helper.metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(); - try (CallContext ctx = CallContext.of(realmContext)) { - CallContext.setCurrentContext(ctx); - PolarisMetaStoreManager metaStoreManager = - helper.metaStoreManagerFactory.getOrCreateMetaStoreManager(ctx.getRealmContext()); - 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(); - } 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/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 57e7d457c..0b9a3b38f 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 @@ -30,6 +30,7 @@ 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; @@ -56,11 +57,13 @@ public class PolarisEclipseLinkMetaStoreManagerTest extends BasePolarisMetaStore protected PolarisTestMetaStoreManager createPolarisTestMetaStoreManager() { PolarisDiagnostics diagServices = new PolarisDefaultDiagServiceImpl(); PolarisEclipseLinkStore store = new PolarisEclipseLinkStore(diagServices); + RealmContext realmContext = () -> "realm"; PolarisMetaStoreSession session = new PolarisEclipseLinkMetaStoreSessionImpl( - store, Mockito.mock(), () -> "realm", null, "polaris", RANDOM_SECRETS, diagServices); + store, Mockito.mock(), realmContext, null, "polaris", RANDOM_SECRETS, diagServices); return new PolarisTestMetaStoreManager( new PolarisMetaStoreManagerImpl( + realmContext, diagServices, new PolarisConfigurationStore() {}, timeSource.withZone(ZoneId.systemDefault())), 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 f642f4b67..e577bb82c 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,11 +34,12 @@ public interface PolarisConfigurationStore { /** * Retrieve the current value for a configuration key. May be null if not set. * + * @param the type of the configuration value + * @param realmContext * @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(String configName) { + default @Nullable T getConfiguration(RealmContext realmContext, String configName) { return null; } @@ -45,14 +47,16 @@ public interface PolarisConfigurationStore { * Retrieve the current value for a configuration key. If not set, return the non-null default * value. * + * @param the type of the configuration value + * @param realmContext * @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(String configName, @Nonnull T defaultValue) { + default @Nonnull T getConfiguration( + RealmContext realmContext, String configName, @Nonnull T defaultValue) { Preconditions.checkNotNull(defaultValue, "Cannot pass null as a default value"); - T configValue = getConfiguration(configName); + T configValue = getConfiguration(realmContext, configName); return configValue != null ? configValue : defaultValue; } @@ -83,12 +87,14 @@ public interface PolarisConfigurationStore { /** * Retrieve the current value for a configuration. * + * @param the type of the configuration value + * @param realmContext * @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(PolarisConfiguration config) { - T result = getConfiguration(config.key, config.defaultValue); + default @Nonnull T getConfiguration( + RealmContext realmContext, PolarisConfiguration config) { + T result = getConfiguration(realmContext, config.key, config.defaultValue); return tryCast(config, result); } @@ -96,18 +102,21 @@ public interface PolarisConfigurationStore { * Retrieve the current value for a configuration, overriding with a catalog config if it is * present. * + * @param the type of the configuration value + * @param realmContext * @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( - @Nonnull CatalogEntity catalogEntity, PolarisConfiguration config) { + 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(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 c04e2460c..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,6 +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.RealmContext; import org.apache.polaris.core.entity.PolarisBaseEntity; import org.apache.polaris.core.entity.PolarisEntityConstants; import org.apache.polaris.core.entity.PolarisEntityCore; @@ -486,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, @@ -501,6 +504,7 @@ public void authorizeOrThrow( @Override public void authorizeOrThrow( + @Nonnull RealmContext realmContext, @Nonnull AuthenticatedPolarisPrincipal authenticatedPrincipal, @Nonnull Set activatedEntities, @Nonnull PolarisAuthorizableOperation authzOp, @@ -508,6 +512,7 @@ public void authorizeOrThrow( @Nullable List secondaries) { boolean enforceCredentialRotationRequiredState = featureConfig.getConfiguration( + realmContext, PolarisConfiguration.ENFORCE_PRINCIPAL_CREDENTIAL_ROTATION_REQUIRED_CHECKING); if (enforceCredentialRotationRequiredState && authenticatedPrincipal 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 105508608..000000000 --- a/polaris-core/src/main/java/org/apache/polaris/core/context/CallContext.java +++ /dev/null @@ -1,135 +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.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 AuthenticatedPolarisPrincipal getAuthenticatedPrincipal() { - return (AuthenticatedPolarisPrincipal) - CallContext.getCurrentContext().contextVariables().get(CallContext.AUTHENTICATED_PRINCIPAL); - } - - static void unsetCurrentContext() { - CURRENT_CONTEXT.remove(); - } - - static CallContext of(final RealmContext realmContext) { - Map map = new HashMap<>(); - return new CallContext() { - @Override - public RealmContext getRealmContext() { - return realmContext; - } - - @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; - 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 Map contextVariables() { - return contextVariables; - } - }; - } - - RealmContext getRealmContext(); - - 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/persistence/LocalPolarisMetaStoreManagerFactory.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/LocalPolarisMetaStoreManagerFactory.java index 30898affc..6b1f466d7 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 @@ -27,7 +27,6 @@ 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; @@ -90,7 +89,7 @@ private void initializeForRealm(RealmContext realmContext) { () -> createMetaStoreSession(backingStore, realmContext, diagnostics)); PolarisMetaStoreManager metaStoreManager = - new PolarisMetaStoreManagerImpl(diagnostics, configurationStore, clock); + new PolarisMetaStoreManagerImpl(realmContext, diagnostics, configurationStore, clock); metaStoreManagerMap.put(realmContext.getRealmIdentifier(), metaStoreManager); } @@ -189,7 +188,6 @@ private PrincipalSecretsResult bootstrapServiceAndCreatePolarisPrincipalForRealm // CallContext hasn't even been resolved yet. PolarisMetaStoreSession metaStoreSession = sessionSupplierMap.get(realmContext.getRealmIdentifier()).get(); - CallContext.setCurrentContext(CallContext.of(realmContext)); PolarisMetaStoreManager.EntityResult preliminaryRootPrincipalLookup = metaStoreManager.readEntityByName( @@ -245,7 +243,6 @@ private void checkPolarisServiceBootstrappedForRealm( PolarisMetaStoreSession metaStoreSession = sessionSupplierMap.get(realmContext.getRealmIdentifier()).get(); - CallContext.setCurrentContext(CallContext.of(realmContext)); PolarisMetaStoreManager.EntityResult rootPrincipalLookup = metaStoreManager.readEntityByName( 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 f52593fa9..0ee5e63c7 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 @@ -38,6 +38,7 @@ import java.util.stream.Collectors; 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; @@ -75,12 +76,17 @@ 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( - PolarisDiagnostics diagnostics, PolarisConfigurationStore configurationStore, Clock clock) { + RealmContext realmContext, + PolarisDiagnostics diagnostics, + PolarisConfigurationStore configurationStore, + Clock clock) { + this.realmContext = realmContext; this.diagnostics = diagnostics; this.configurationStore = configurationStore; this.clock = clock; @@ -1809,6 +1815,7 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( PolarisObjectMapperUtil.parseTaskState(entity); long taskAgeTimeout = configurationStore.getConfiguration( + realmContext, PolarisTaskConstants.TASK_TIMEOUT_MILLIS_CONFIG, PolarisTaskConstants.TASK_TIMEOUT_MILLIS); return taskState == null @@ -1927,7 +1934,7 @@ private PolarisEntityResolver resolveSecurableToRoleGrant( readStorageConfiguration(diagnostics, reloadedEntity.getEntity()); Map validateLocationAccess = storageIntegration - .validateAccessToLocations(storageConfigurationInfo, actions, locations) + .validateAccessToLocations(realmContext, storageConfigurationInfo, actions, locations) .entrySet() .stream() .collect( 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 9410d9ae3..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 @@ -26,13 +26,15 @@ import java.util.function.Function; import java.util.stream.Collectors; 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 */ @@ -51,6 +53,7 @@ public InMemoryStorageIntegration( * 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 @@ -60,6 +63,7 @@ public InMemoryStorageIntegration( */ public static Map> validateSubpathsOfAllowedLocations( + @Nonnull RealmContext realmContext, @Nonnull PolarisConfigurationStore configurationStore, @Nonnull PolarisStorageConfigurationInfo storageConfig, @Nonnull Set actions, @@ -82,7 +86,7 @@ public InMemoryStorageIntegration( allowedLocationStrings.stream().map(StorageLocation::of).collect(Collectors.toList()); boolean allowWildcardLocation = - configurationStore.getConfiguration("ALLOW_WILDCARD_LOCATION", false); + configurationStore.getConfiguration(realmContext, "ALLOW_WILDCARD_LOCATION", false); if (allowWildcardLocation && allowedLocationStrings.contains("*")) { return locations.stream() @@ -125,10 +129,11 @@ public InMemoryStorageIntegration( @Override @Nonnull public Map> validateAccessToLocations( + @Nonnull RealmContext realmContext, @Nonnull T storageConfig, @Nonnull Set actions, @Nonnull Set locations) { return validateSubpathsOfAllowedLocations( - configurationStore, storageConfig, actions, locations); + realmContext, configurationStore, storageConfig, actions, locations); } } 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 39c0e51fb..2fdc816e9 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 @@ -40,6 +40,7 @@ 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.RealmContext; import org.apache.polaris.core.entity.CatalogEntity; import org.apache.polaris.core.entity.PolarisEntity; import org.apache.polaris.core.entity.PolarisEntityConstants; @@ -136,6 +137,7 @@ public static PolarisStorageConfigurationInfo deserialize( } public static Optional forEntityPath( + RealmContext realmContext, PolarisConfigurationStore configurationStore, PolarisDiagnostics diagnostics, List entityPath) { @@ -165,7 +167,9 @@ public static Optional forEntityPath( CatalogEntity catalog = CatalogEntity.of(entityPath.get(0)); boolean allowEscape = configurationStore.getConfiguration( - catalog, PolarisConfiguration.ALLOW_UNSTRUCTURED_TABLE_LOCATION); + 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..c9358ba8d 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 @@ -64,6 +65,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 +97,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/test/java/org/apache/polaris/core/persistence/EntityCacheTest.java b/polaris-core/src/test/java/org/apache/polaris/core/persistence/EntityCacheTest.java index 380d3e916..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 @@ -84,7 +84,7 @@ public EntityCacheTest() { new PolarisTreeMapMetaStoreSessionImpl(store, Mockito.mock(), RANDOM_SECRETS, diagServices); metaStoreManager = new PolarisMetaStoreManagerImpl( - diagServices, new PolarisConfigurationStore() {}, Clock.systemUTC()); + () -> "test", diagServices, new PolarisConfigurationStore() {}, Clock.systemUTC()); // bootstrap the mata store with our test schema tm = new PolarisTestMetaStoreManager(metaStoreManager, metaStore, diagServices); 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 e7b7fd530..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 @@ -35,6 +35,7 @@ public PolarisTestMetaStoreManager createPolarisTestMetaStoreManager() { new PolarisTreeMapMetaStoreSessionImpl(store, Mockito.mock(), RANDOM_SECRETS, diagServices); return new PolarisTestMetaStoreManager( new PolarisMetaStoreManagerImpl( + () -> "test", diagServices, new PolarisConfigurationStore() {}, timeSource.withZone(ZoneId.systemDefault())), 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 76337de17..c82a65e85 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 @@ -98,7 +98,7 @@ public ResolverTest() { new PolarisTreeMapMetaStoreSessionImpl(store, Mockito.mock(), RANDOM_SECRETS, diagServices); metaStoreManager = new PolarisMetaStoreManagerImpl( - diagServices, new PolarisConfigurationStore() {}, Clock.systemUTC()); + () -> "test", diagServices, new PolarisConfigurationStore() {}, Clock.systemUTC()); // bootstrap the mata store with our test schema tm = new PolarisTestMetaStoreManager(metaStoreManager, metaStore, diagServices); 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 f2c016236..c63694dc2 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 @@ -26,19 +26,22 @@ import java.util.Set; 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.storage.aws.AwsStorageConfigurationInfo; import org.assertj.core.api.Assertions; import org.junit.jupiter.api.Test; class InMemoryStorageIntegrationTest { + private RealmContext realmContext = () -> "test"; + @Test public void testValidateAccessToLocations() { MockInMemoryStorageIntegration storage = new MockInMemoryStorageIntegration(new PolarisConfigurationStore() {}); Map> result = storage.validateAccessToLocations( + realmContext, new AwsStorageConfigurationInfo( PolarisStorageConfigurationInfo.StorageType.S3, List.of( @@ -78,47 +81,46 @@ public void testValidateAccessToLocationsWithWildcard() { new PolarisConfigurationStore() { @SuppressWarnings("unchecked") @Override - public @Nullable T getConfiguration(String configName) { + public @Nullable T getConfiguration(RealmContext realmContext, String configName) { return (T) config.get(configName); } }; MockInMemoryStorageIntegration storage = new MockInMemoryStorageIntegration(configurationStore); - try (CallContext ignored = CallContext.setCurrentContext(CallContext.of(() -> "realm"))) { - 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)); - } + 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 @@ -127,6 +129,7 @@ public void testValidateAccessToLocationsNoAllowedLocations() { new MockInMemoryStorageIntegration(new PolarisConfigurationStore() {}); Map> result = storage.validateAccessToLocations( + realmContext, new AwsStorageConfigurationInfo( PolarisStorageConfigurationInfo.StorageType.S3, List.of(), @@ -162,6 +165,7 @@ public void testValidateAccessToLocationsWithPrefixOfAllowedLocation() { new MockInMemoryStorageIntegration(new PolarisConfigurationStore() {}); Map> result = storage.validateAccessToLocations( + realmContext, new AwsStorageConfigurationInfo( PolarisStorageConfigurationInfo.StorageType.S3, List.of("s3://bucket/path/to/warehouse"), 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 57fa67a0f..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 @@ -22,11 +22,15 @@ import java.util.List; 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 = @@ -44,7 +48,7 @@ public void testConfigsCanBeCastedFromString() { */ @SuppressWarnings("unchecked") @Override - public @Nullable T getConfiguration(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); @@ -61,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(c)); + Assertions.assertEquals(c.defaultValue, store.getConfiguration(realmContext, c)); } } @@ -75,13 +79,14 @@ public void testInvalidCastThrowsException() { new PolarisConfigurationStore() { @SuppressWarnings("unchecked") @Override - public T getConfiguration(String configName) { + public T getConfiguration(RealmContext realmContext, String configName) { return (T) "abc123"; } }; for (PolarisConfiguration c : configs) { - Assertions.assertThrows(NumberFormatException.class, () -> store.getConfiguration(c)); + Assertions.assertThrows( + NumberFormatException.class, () -> store.getConfiguration(realmContext, c)); } } 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 23b05a7f9..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,7 +34,6 @@ import java.util.function.Function; import java.util.stream.Collectors; import java.util.stream.Stream; -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; @@ -102,135 +101,116 @@ void testBrowse() { @Test void testCreateEntities() { PolarisMetaStoreManager metaStoreManager = polarisTestMetaStoreManager.polarisMetaStoreManager; - try (CallContext callCtx = CallContext.of(() -> "testRealm")) { - if (CallContext.getCurrentContext() == null) { - CallContext.setCurrentContext(callCtx); - } - 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())); - } + 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")) { - if (CallContext.getCurrentContext() == null) { - CallContext.setCurrentContext(callCtx); - } - 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)); - } + 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")) { - 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.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(); - } + 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 TaskEntity createTask(String taskName, long id) { 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 ed8a88a51..34853ff40 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 @@ -64,6 +64,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.RealmContext; import org.apache.polaris.core.entity.CatalogEntity; import org.apache.polaris.core.entity.CatalogRoleEntity; import org.apache.polaris.core.entity.NamespaceEntity; @@ -104,6 +105,7 @@ public class PolarisAdminService { private static final Logger LOGGER = LoggerFactory.getLogger(PolarisAdminService.class); + private final RealmContext realmContext; private final PolarisMetaStoreSession metaStoreSession; private final PolarisConfigurationStore configurationStore; private final PolarisEntityManager entityManager; @@ -115,12 +117,14 @@ public class PolarisAdminService { private PolarisResolutionManifest resolutionManifest = null; public PolarisAdminService( + RealmContext realmContext, PolarisEntityManager entityManager, PolarisMetaStoreManager metaStoreManager, PolarisMetaStoreSession metaStoreSession, PolarisConfigurationStore configurationStore, AuthenticatedPolarisPrincipal authenticatedPrincipal, PolarisAuthorizer authorizer) { + this.realmContext = realmContext; this.metaStoreSession = metaStoreSession; this.configurationStore = configurationStore; this.entityManager = entityManager; @@ -159,6 +163,7 @@ private void authorizeBasicRootOperationOrThrow(PolarisAuthorizableOperation op) PolarisResolvedPathWrapper rootContainerWrapper = resolutionManifest.getResolvedRootContainerEntityAsPath(); authorizer.authorizeOrThrow( + realmContext, authenticatedPrincipal, resolutionManifest.getAllActivatedPrincipalRoleEntities(), op, @@ -204,6 +209,7 @@ private void authorizeBasicTopLevelEntityOperationOrThrow( return; } authorizer.authorizeOrThrow( + realmContext, authenticatedPrincipal, resolutionManifest.getAllActivatedCatalogRoleAndPrincipalRoles(), op, @@ -225,6 +231,7 @@ private void authorizeBasicCatalogRoleOperationOrThrow( throw new NotFoundException("CatalogRole does not exist: %s", catalogRoleName); } authorizer.authorizeOrThrow( + realmContext, authenticatedPrincipal, resolutionManifest.getAllActivatedCatalogRoleAndPrincipalRoles(), op, @@ -255,6 +262,7 @@ private void authorizeGrantOnRootContainerToPrincipalRoleOperationOrThrow( principalRoleName, PolarisEntityType.PRINCIPAL_ROLE); authorizer.authorizeOrThrow( + realmContext, authenticatedPrincipal, resolutionManifest.getAllActivatedCatalogRoleAndPrincipalRoles(), op, @@ -291,6 +299,7 @@ private void authorizeGrantOnTopLevelEntityToPrincipalRoleOperationOrThrow( principalRoleName, PolarisEntityType.PRINCIPAL_ROLE); authorizer.authorizeOrThrow( + realmContext, authenticatedPrincipal, resolutionManifest.getAllActivatedCatalogRoleAndPrincipalRoles(), op, @@ -321,6 +330,7 @@ private void authorizeGrantOnPrincipalRoleToPrincipalOperationOrThrow( resolutionManifest.getResolvedTopLevelEntity(principalName, PolarisEntityType.PRINCIPAL); authorizer.authorizeOrThrow( + realmContext, authenticatedPrincipal, resolutionManifest.getAllActivatedCatalogRoleAndPrincipalRoles(), op, @@ -360,6 +370,7 @@ private void authorizeGrantOnCatalogRoleToPrincipalRoleOperationOrThrow( resolutionManifest.getResolvedPath(catalogRoleName, true); authorizer.authorizeOrThrow( + realmContext, authenticatedPrincipal, resolutionManifest.getAllActivatedCatalogRoleAndPrincipalRoles(), op, @@ -390,6 +401,7 @@ private void authorizeGrantOnCatalogOperationOrThrow( PolarisResolvedPathWrapper catalogRoleWrapper = resolutionManifest.getResolvedPath(catalogRoleName, true); authorizer.authorizeOrThrow( + realmContext, authenticatedPrincipal, resolutionManifest.getAllActivatedCatalogRoleAndPrincipalRoles(), op, @@ -430,6 +442,7 @@ private void authorizeGrantOnNamespaceOperationOrThrow( resolutionManifest.getResolvedPath(catalogRoleName, true); authorizer.authorizeOrThrow( + realmContext, authenticatedPrincipal, resolutionManifest.getAllActivatedCatalogRoleAndPrincipalRoles(), op, @@ -475,6 +488,7 @@ private void authorizeGrantOnTableLikeOperationOrThrow( resolutionManifest.getResolvedPath(catalogRoleName, true); authorizer.authorizeOrThrow( + realmContext, authenticatedPrincipal, resolutionManifest.getAllActivatedCatalogRoleAndPrincipalRoles(), op, @@ -512,7 +526,8 @@ private String terminateWithSlash(String path) { */ private boolean catalogOverlapsWithExistingCatalog(CatalogEntity catalogEntity) { boolean allowOverlappingCatalogUrls = - configurationStore.getConfiguration(PolarisConfiguration.ALLOW_OVERLAPPING_CATALOG_URLS); + configurationStore.getConfiguration( + realmContext, PolarisConfiguration.ALLOW_OVERLAPPING_CATALOG_URLS); if (allowOverlappingCatalogUrls) { return false; @@ -579,7 +594,8 @@ public void deleteCatalog(String name) { .orElseThrow(() -> new NotFoundException("Catalog %s not found", name)); // TODO: Handle return value in case of concurrent modification boolean cleanup = - configurationStore.getConfiguration(PolarisConfiguration.CLEANUP_ON_CATALOG_DROP); + configurationStore.getConfiguration( + realmContext, PolarisConfiguration.CLEANUP_ON_CATALOG_DROP); PolarisMetaStoreManager.DropEntityResult dropEntityResult = metaStoreManager.dropEntityIfExists(metaStoreSession, null, entity, Map.of(), cleanup); 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 69a15be5b..8ed3cbba0 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 @@ -101,7 +101,8 @@ public PolarisServiceImpl( this.polarisAuthorizer = polarisAuthorizer; } - private PolarisAdminService newAdminService(SecurityContext securityContext) { + private PolarisAdminService newAdminService( + RealmContext realmContext, SecurityContext securityContext) { AuthenticatedPolarisPrincipal authenticatedPrincipal = (AuthenticatedPolarisPrincipal) securityContext.getUserPrincipal(); if (authenticatedPrincipal == null) { @@ -109,6 +110,7 @@ private PolarisAdminService newAdminService(SecurityContext securityContext) { } return new PolarisAdminService( + realmContext, entityManager, metaStoreManager, session, @@ -121,9 +123,9 @@ private PolarisAdminService newAdminService(SecurityContext securityContext) { @Override public Response createCatalog( CreateCatalogRequest request, RealmContext realmContext, SecurityContext securityContext) { - PolarisAdminService adminService = newAdminService(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,9 +133,11 @@ public Response createCatalog( return Response.status(Response.Status.CREATED).build(); } - private void validateStorageConfig(StorageConfigInfo storageConfigInfo) { + private void validateStorageConfig( + StorageConfigInfo storageConfigInfo, RealmContext realmContext) { List allowedStorageTypes = - configurationStore.getConfiguration(PolarisConfiguration.SUPPORTED_CATALOG_STORAGE_TYPES); + configurationStore.getConfiguration( + realmContext, PolarisConfiguration.SUPPORTED_CATALOG_STORAGE_TYPES); if (!allowedStorageTypes.contains(storageConfigInfo.getStorageType().name())) { LOGGER .atWarn() @@ -148,7 +152,7 @@ private void validateStorageConfig(StorageConfigInfo storageConfigInfo) { @Override public Response deleteCatalog( String catalogName, RealmContext realmContext, SecurityContext securityContext) { - PolarisAdminService adminService = newAdminService(securityContext); + PolarisAdminService adminService = newAdminService(realmContext, securityContext); adminService.deleteCatalog(catalogName); return Response.status(Response.Status.NO_CONTENT).build(); } @@ -157,7 +161,7 @@ public Response deleteCatalog( @Override public Response getCatalog( String catalogName, RealmContext realmContext, SecurityContext securityContext) { - PolarisAdminService adminService = newAdminService(securityContext); + PolarisAdminService adminService = newAdminService(realmContext, securityContext); return Response.ok(adminService.getCatalog(catalogName).asCatalog()).build(); } @@ -168,9 +172,9 @@ public Response updateCatalog( UpdateCatalogRequest updateRequest, RealmContext realmContext, SecurityContext securityContext) { - PolarisAdminService adminService = newAdminService(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(); } @@ -178,7 +182,7 @@ public Response updateCatalog( /** From PolarisCatalogsApiService */ @Override public Response listCatalogs(RealmContext realmContext, SecurityContext securityContext) { - PolarisAdminService adminService = newAdminService(securityContext); + PolarisAdminService adminService = newAdminService(realmContext, securityContext); List catalogList = adminService.listCatalogs().stream() .map(CatalogEntity::new) @@ -193,7 +197,7 @@ public Response listCatalogs(RealmContext realmContext, SecurityContext security @Override public Response createPrincipal( CreatePrincipalRequest request, RealmContext realmContext, SecurityContext securityContext) { - PolarisAdminService adminService = newAdminService(securityContext); + PolarisAdminService adminService = newAdminService(realmContext, securityContext); PrincipalEntity principal = PrincipalEntity.fromPrincipal(request.getPrincipal()); if (Boolean.TRUE.equals(request.getCredentialRotationRequired())) { principal = @@ -208,7 +212,7 @@ public Response createPrincipal( @Override public Response deletePrincipal( String principalName, RealmContext realmContext, SecurityContext securityContext) { - PolarisAdminService adminService = newAdminService(securityContext); + PolarisAdminService adminService = newAdminService(realmContext, securityContext); adminService.deletePrincipal(principalName); return Response.status(Response.Status.NO_CONTENT).build(); } @@ -217,7 +221,7 @@ public Response deletePrincipal( @Override public Response getPrincipal( String principalName, RealmContext realmContext, SecurityContext securityContext) { - PolarisAdminService adminService = newAdminService(securityContext); + PolarisAdminService adminService = newAdminService(realmContext, securityContext); return Response.ok(adminService.getPrincipal(principalName).asPrincipal()).build(); } @@ -228,7 +232,7 @@ public Response updatePrincipal( UpdatePrincipalRequest updateRequest, RealmContext realmContext, SecurityContext securityContext) { - PolarisAdminService adminService = newAdminService(securityContext); + PolarisAdminService adminService = newAdminService(realmContext, securityContext); return Response.ok(adminService.updatePrincipal(principalName, updateRequest).asPrincipal()) .build(); } @@ -237,14 +241,14 @@ public Response updatePrincipal( @Override public Response rotateCredentials( String principalName, RealmContext realmContext, SecurityContext securityContext) { - PolarisAdminService adminService = newAdminService(securityContext); + PolarisAdminService adminService = newAdminService(realmContext, securityContext); return Response.ok(adminService.rotateCredentials(principalName)).build(); } /** From PolarisPrincipalsApiService */ @Override public Response listPrincipals(RealmContext realmContext, SecurityContext securityContext) { - PolarisAdminService adminService = newAdminService(securityContext); + PolarisAdminService adminService = newAdminService(realmContext, securityContext); List principalList = adminService.listPrincipals().stream() .map(PrincipalEntity::new) @@ -261,7 +265,7 @@ public Response createPrincipalRole( CreatePrincipalRoleRequest request, RealmContext realmContext, SecurityContext securityContext) { - PolarisAdminService adminService = newAdminService(securityContext); + PolarisAdminService adminService = newAdminService(realmContext, securityContext); PrincipalRole newPrincipalRole = new PrincipalRoleEntity( adminService.createPrincipalRole( @@ -275,7 +279,7 @@ public Response createPrincipalRole( @Override public Response deletePrincipalRole( String principalRoleName, RealmContext realmContext, SecurityContext securityContext) { - PolarisAdminService adminService = newAdminService(securityContext); + PolarisAdminService adminService = newAdminService(realmContext, securityContext); adminService.deletePrincipalRole(principalRoleName); return Response.status(Response.Status.NO_CONTENT).build(); } @@ -284,7 +288,7 @@ public Response deletePrincipalRole( @Override public Response getPrincipalRole( String principalRoleName, RealmContext realmContext, SecurityContext securityContext) { - PolarisAdminService adminService = newAdminService(securityContext); + PolarisAdminService adminService = newAdminService(realmContext, securityContext); return Response.ok(adminService.getPrincipalRole(principalRoleName).asPrincipalRole()).build(); } @@ -295,7 +299,7 @@ public Response updatePrincipalRole( UpdatePrincipalRoleRequest updateRequest, RealmContext realmContext, SecurityContext securityContext) { - PolarisAdminService adminService = newAdminService(securityContext); + PolarisAdminService adminService = newAdminService(realmContext, securityContext); return Response.ok( adminService.updatePrincipalRole(principalRoleName, updateRequest).asPrincipalRole()) .build(); @@ -304,7 +308,7 @@ public Response updatePrincipalRole( /** From PolarisPrincipalRolesApiService */ @Override public Response listPrincipalRoles(RealmContext realmContext, SecurityContext securityContext) { - PolarisAdminService adminService = newAdminService(securityContext); + PolarisAdminService adminService = newAdminService(realmContext, securityContext); List principalRoleList = adminService.listPrincipalRoles().stream() .map(PrincipalRoleEntity::new) @@ -322,7 +326,7 @@ public Response createCatalogRole( CreateCatalogRoleRequest request, RealmContext realmContext, SecurityContext securityContext) { - PolarisAdminService adminService = newAdminService(securityContext); + PolarisAdminService adminService = newAdminService(realmContext, securityContext); CatalogRole newCatalogRole = new CatalogRoleEntity( adminService.createCatalogRole( @@ -339,7 +343,7 @@ public Response deleteCatalogRole( String catalogRoleName, RealmContext realmContext, SecurityContext securityContext) { - PolarisAdminService adminService = newAdminService(securityContext); + PolarisAdminService adminService = newAdminService(realmContext, securityContext); adminService.deleteCatalogRole(catalogName, catalogRoleName); return Response.status(Response.Status.NO_CONTENT).build(); } @@ -351,7 +355,7 @@ public Response getCatalogRole( String catalogRoleName, RealmContext realmContext, SecurityContext securityContext) { - PolarisAdminService adminService = newAdminService(securityContext); + PolarisAdminService adminService = newAdminService(realmContext, securityContext); return Response.ok(adminService.getCatalogRole(catalogName, catalogRoleName).asCatalogRole()) .build(); } @@ -364,7 +368,7 @@ public Response updateCatalogRole( UpdateCatalogRoleRequest updateRequest, RealmContext realmContext, SecurityContext securityContext) { - PolarisAdminService adminService = newAdminService(securityContext); + PolarisAdminService adminService = newAdminService(realmContext, securityContext); return Response.ok( adminService .updateCatalogRole(catalogName, catalogRoleName, updateRequest) @@ -376,7 +380,7 @@ public Response updateCatalogRole( @Override public Response listCatalogRoles( String catalogName, RealmContext realmContext, SecurityContext securityContext) { - PolarisAdminService adminService = newAdminService(securityContext); + PolarisAdminService adminService = newAdminService(realmContext, securityContext); List catalogRoleList = adminService.listCatalogRoles(catalogName).stream() .map(CatalogRoleEntity::new) @@ -398,7 +402,7 @@ public Response assignPrincipalRole( "Assigning principalRole {} to principal {}", request.getPrincipalRole().getName(), principalName); - PolarisAdminService adminService = newAdminService(securityContext); + PolarisAdminService adminService = newAdminService(realmContext, securityContext); adminService.assignPrincipalRole(principalName, request.getPrincipalRole().getName()); return Response.status(Response.Status.CREATED).build(); } @@ -411,7 +415,7 @@ public Response revokePrincipalRole( RealmContext realmContext, SecurityContext securityContext) { LOGGER.info("Revoking principalRole {} from principal {}", principalRoleName, principalName); - PolarisAdminService adminService = newAdminService(securityContext); + PolarisAdminService adminService = newAdminService(realmContext, securityContext); adminService.revokePrincipalRole(principalName, principalRoleName); return Response.status(Response.Status.NO_CONTENT).build(); } @@ -420,7 +424,7 @@ public Response revokePrincipalRole( @Override public Response listPrincipalRolesAssigned( String principalName, RealmContext realmContext, SecurityContext securityContext) { - PolarisAdminService adminService = newAdminService(securityContext); + PolarisAdminService adminService = newAdminService(realmContext, securityContext); List principalRoleList = adminService.listPrincipalRolesAssigned(principalName).stream() .map(PrincipalRoleEntity::new) @@ -444,7 +448,7 @@ public Response assignCatalogRoleToPrincipalRole( request.getCatalogRole().getName(), catalogName, principalRoleName); - PolarisAdminService adminService = newAdminService(securityContext); + PolarisAdminService adminService = newAdminService(realmContext, securityContext); adminService.assignCatalogRoleToPrincipalRole( principalRoleName, catalogName, request.getCatalogRole().getName()); return Response.status(Response.Status.CREATED).build(); @@ -463,7 +467,7 @@ public Response revokeCatalogRoleFromPrincipalRole( catalogRoleName, catalogName, principalRoleName); - PolarisAdminService adminService = newAdminService(securityContext); + PolarisAdminService adminService = newAdminService(realmContext, securityContext); adminService.revokeCatalogRoleFromPrincipalRole( principalRoleName, catalogName, catalogRoleName); return Response.status(Response.Status.NO_CONTENT).build(); @@ -473,7 +477,7 @@ public Response revokeCatalogRoleFromPrincipalRole( @Override public Response listAssigneePrincipalsForPrincipalRole( String principalRoleName, RealmContext realmContext, SecurityContext securityContext) { - PolarisAdminService adminService = newAdminService(securityContext); + PolarisAdminService adminService = newAdminService(realmContext, securityContext); List principalList = adminService.listAssigneePrincipalsForPrincipalRole(principalRoleName).stream() .map(PrincipalEntity::new) @@ -491,7 +495,7 @@ public Response listCatalogRolesForPrincipalRole( String catalogName, RealmContext realmContext, SecurityContext securityContext) { - PolarisAdminService adminService = newAdminService(securityContext); + PolarisAdminService adminService = newAdminService(realmContext, securityContext); List catalogRoleList = adminService.listCatalogRolesForPrincipalRole(principalRoleName, catalogName).stream() .map(CatalogRoleEntity::new) @@ -515,7 +519,7 @@ public Response addGrantToCatalogRole( grantRequest, catalogRoleName, catalogName); - PolarisAdminService adminService = newAdminService(securityContext); + PolarisAdminService adminService = newAdminService(realmContext, securityContext); switch (grantRequest.getGrant()) { // The per-securable-type Privilege enums must be exact String match for a subset of all // PolarisPrivilege values. @@ -591,7 +595,7 @@ public Response revokeGrantFromCatalogRole( return Response.status(501).build(); // not implemented } - PolarisAdminService adminService = newAdminService(securityContext); + PolarisAdminService adminService = newAdminService(realmContext, securityContext); switch (grantRequest.getGrant()) { // The per-securable-type Privilege enums must be exact String match for a subset of all // PolarisPrivilege values. @@ -655,7 +659,7 @@ public Response listAssigneePrincipalRolesForCatalogRole( String catalogRoleName, RealmContext realmContext, SecurityContext securityContext) { - PolarisAdminService adminService = newAdminService(securityContext); + PolarisAdminService adminService = newAdminService(realmContext, securityContext); List principalRoleList = adminService.listAssigneePrincipalRolesForCatalogRole(catalogName, catalogRoleName).stream() .map(PrincipalRoleEntity::new) @@ -673,7 +677,7 @@ public Response listGrantsForCatalogRole( String catalogRoleName, RealmContext realmContext, SecurityContext securityContext) { - PolarisAdminService adminService = newAdminService(securityContext); + PolarisAdminService adminService = newAdminService(realmContext, securityContext); List grantList = adminService.listGrantsForCatalogRole(catalogName, catalogRoleName); GrantResources grantResources = new GrantResources(grantList); 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 48ce3a2d5..7c966d63f 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 @@ -25,7 +25,6 @@ import org.apache.commons.lang3.exception.ExceptionUtils; import org.apache.iceberg.exceptions.NotAuthorizedException; 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; @@ -105,9 +104,6 @@ protected Optional getPrincipal(DecodedToken toke AuthenticatedPolarisPrincipal authenticatedPrincipal = new AuthenticatedPolarisPrincipal(new PrincipalEntity(principal), activatedPrincipalRoles); LOGGER.debug("Populating authenticatedPrincipal into CallContext: {}", authenticatedPrincipal); - CallContext.getCurrentContext() - .contextVariables() - .put(CallContext.AUTHENTICATED_PRINCIPAL, authenticatedPrincipal); return Optional.of(authenticatedPrincipal); } } 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 b6a83e31a..1059257fd 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 @@ -29,7 +29,6 @@ import org.apache.hadoop.hdfs.web.oauth2.OAuth2Constants; import org.apache.iceberg.rest.auth.OAuth2Properties; 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; @@ -46,12 +45,10 @@ public class DefaultOAuth2ApiService implements IcebergRestOAuth2ApiService { private static final Logger LOGGER = LoggerFactory.getLogger(DefaultOAuth2ApiService.class); 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 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 7fb1ea7d2..0e3bb315f 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 @@ -78,7 +78,7 @@ 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; @@ -153,7 +153,7 @@ public class BasePolarisCatalog extends BaseMetastoreViewCatalog private final PolarisMetaStoreSession metaStoreSession; private final PolarisConfigurationStore configurationStore; private final PolarisDiagnostics diagnostics; - private final CallContext callContext; + private final RealmContext realmContext; private final PolarisResolutionManifestCatalogView resolvedEntityView; private final CatalogEntity catalogEntity; private final TaskExecutor taskExecutor; @@ -170,30 +170,30 @@ public class BasePolarisCatalog extends BaseMetastoreViewCatalog private Map tableDefaultProperties; /** + * @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, PolarisMetaStoreSession metaStoreSession, PolarisConfigurationStore configurationStore, PolarisDiagnostics diagnostics, - CallContext callContext, PolarisResolutionManifestCatalogView resolvedEntityView, AuthenticatedPolarisPrincipal authenticatedPrincipal, TaskExecutor taskExecutor, FileIOFactory fileIOFactory) { + this.realmContext = realmContext; this.entityManager = entityManager; this.metaStoreManager = metaStoreManager; this.metaStoreSession = metaStoreSession; this.configurationStore = configurationStore; this.diagnostics = diagnostics; - this.callContext = callContext; this.resolvedEntityView = resolvedEntityView; this.taskExecutor = taskExecutor; this.fileIOFactory = fileIOFactory; @@ -263,7 +263,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); @@ -447,7 +446,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; @@ -511,7 +510,7 @@ private void createNamespaceInternal( .setBaseLocation(baseLocation) .build(); if (!configurationStore.getConfiguration( - PolarisConfiguration.ALLOW_NAMESPACE_LOCATION_OVERLAP)) { + 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()); @@ -622,7 +621,7 @@ public boolean dropNamespace(Namespace namespace) throws NamespaceNotEmptyExcept leafEntity, Map.of(), configurationStore.getConfiguration( - PolarisConfiguration.CLEANUP_ON_NAMESPACE_DROP)); + realmContext, PolarisConfiguration.CLEANUP_ON_NAMESPACE_DROP)); if (!dropEntityResult.isSuccess() && dropEntityResult.failedBecauseNotEmpty()) { throw new NamespaceNotEmptyException("Namespace %s is not empty", namespace); @@ -648,7 +647,7 @@ public boolean setProperties(Namespace namespace, Map properties new PolarisEntity.Builder(entity).setProperties(newProperties).build(); if (!configurationStore.getConfiguration( - PolarisConfiguration.ALLOW_NAMESPACE_LOCATION_OVERLAP)) { + realmContext, PolarisConfiguration.ALLOW_NAMESPACE_LOCATION_OVERLAP)) { LOGGER.debug("Validating no overlap with sibling tables or namespaces"); validateNoLocationOverlap( NamespaceEntity.of(updatedEntity).getBaseLocation(), @@ -936,13 +935,14 @@ private void validateLocationsForTableLike( PolarisResolvedPathWrapper resolvedStorageEntity) { Optional optStorageConfiguration = PolarisStorageConfigurationInfo.forEntityPath( - configurationStore, diagnostics, resolvedStorageEntity.getRawFullPath()); + realmContext, configurationStore, diagnostics, resolvedStorageEntity.getRawFullPath()); optStorageConfiguration.ifPresentOrElse( storageConfigInfo -> { Map> validationResults = InMemoryStorageIntegration.validateSubpathsOfAllowedLocations( + realmContext, configurationStore, storageConfigInfo, Set.of(PolarisStorageActions.ALL), @@ -989,7 +989,7 @@ private void validateLocationsForTableLike( () -> { List allowedStorageTypes = configurationStore.getConfiguration( - PolarisConfiguration.SUPPORTED_CATALOG_STORAGE_TYPES); + realmContext, PolarisConfiguration.SUPPORTED_CATALOG_STORAGE_TYPES); if (!allowedStorageTypes.contains(StorageConfigInfo.StorageTypeEnum.FILE.name())) { List invalidLocations = locations.stream() @@ -1014,7 +1014,7 @@ private void validateNoLocationOverlap( List resolvedNamespace, String location) { if (configurationStore.getConfiguration( - catalog, PolarisConfiguration.ALLOW_TABLE_LOCATION_OVERLAP)) { + 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 @@ -1380,10 +1380,11 @@ protected String tableName() { private void validateMetadataFileInTableDir( TableIdentifier identifier, TableMetadata metadata, CatalogEntity catalog) { boolean allowEscape = - configurationStore.getConfiguration(PolarisConfiguration.ALLOW_EXTERNAL_TABLE_LOCATION); + configurationStore.getConfiguration( + realmContext, PolarisConfiguration.ALLOW_EXTERNAL_TABLE_LOCATION); if (!allowEscape && !configurationStore.getConfiguration( - PolarisConfiguration.ALLOW_EXTERNAL_METADATA_FILE_LOCATION)) { + realmContext, PolarisConfiguration.ALLOW_EXTERNAL_METADATA_FILE_LOCATION)) { LOGGER.debug( "Validating base location {} for table {} in metadata file {}", metadata.location(), @@ -1817,7 +1818,7 @@ private void updateTableLike(TableIdentifier identifier, PolarisEntity entity) { if (catalogPath != null && !catalogPath.isEmpty() && purge) { boolean dropWithPurgeEnabled = configurationStore.getConfiguration( - catalogEntity, PolarisConfiguration.DROP_WITH_PURGE_ENABLED); + realmContext, catalogEntity, PolarisConfiguration.DROP_WITH_PURGE_ENABLED); if (!dropWithPurgeEnabled) { throw new ForbiddenException( String.format( @@ -2036,7 +2037,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 configurationStore.getConfiguration(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 bc55ea217..c9db7be90 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; @@ -56,7 +58,6 @@ 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.PolarisEntityManager; @@ -71,6 +72,8 @@ 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 @@ -81,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) @@ -113,7 +118,7 @@ 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 PolarisMetaStoreManager metaStoreManager; private final PolarisEntityManager entityManager; @@ -124,7 +129,7 @@ public class IcebergCatalogAdapter @Inject public IcebergCatalogAdapter( - CallContext callContext, + RealmContext realmContext, CallContextCatalogFactory catalogFactory, PolarisEntityManager entityManager, PolarisMetaStoreManager metaStoreManager, @@ -132,7 +137,7 @@ public IcebergCatalogAdapter( PolarisConfigurationStore configurationStore, PolarisDiagnostics diagnostics, PolarisAuthorizer polarisAuthorizer) { - this.callContext = callContext; + this.realmContext = realmContext; this.catalogFactory = catalogFactory; this.entityManager = entityManager; this.metaStoreManager = metaStoreManager; @@ -140,8 +145,25 @@ public IcebergCatalogAdapter( 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( @@ -153,7 +175,7 @@ private PolarisCatalogHandlerWrapper newHandlerWrapper( } return new PolarisCatalogHandlerWrapper( - callContext, + realmContext, session, configurationStore, diagnostics, @@ -171,9 +193,10 @@ public Response createNamespace( CreateNamespaceRequest createNamespaceRequest, RealmContext realmContext, SecurityContext securityContext) { - return Response.ok( - newHandlerWrapper(securityContext, prefix).createNamespace(createNamespaceRequest)) - .build(); + return withCatalog( + securityContext, + prefix, + catalog -> Response.ok(catalog.createNamespace(createNamespaceRequest)).build()); } @Override @@ -186,18 +209,19 @@ public Response listNamespaces( SecurityContext securityContext) { Optional namespaceOptional = Optional.ofNullable(parent).map(IcebergCatalogAdapter::decodeNamespace); - return Response.ok( - newHandlerWrapper(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(securityContext, prefix).loadNamespaceMetadata(ns)) - .build(); + return withCatalog( + securityContext, prefix, catalog -> Response.ok(catalog.loadNamespaceMetadata(ns)).build()); } private static Namespace decodeNamespace(String namespace) { @@ -208,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(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(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 @@ -228,10 +262,12 @@ public Response updateProperties( RealmContext realmContext, SecurityContext securityContext) { Namespace ns = decodeNamespace(namespace); - return Response.ok( - newHandlerWrapper(securityContext, prefix) - .updateNamespaceProperties(ns, updateNamespacePropertiesRequest)) - .build(); + return withCatalog( + securityContext, + prefix, + catalog -> + Response.ok(catalog.updateNamespaceProperties(ns, updateNamespacePropertiesRequest)) + .build()); } private EnumSet parseAccessDelegationModes(String accessDelegationMode) { @@ -255,28 +291,25 @@ public Response createTable( EnumSet delegationModes = parseAccessDelegationModes(accessDelegationMode); Namespace ns = decodeNamespace(namespace); - if (createTableRequest.stageCreate()) { - if (delegationModes.isEmpty()) { - return Response.ok( - newHandlerWrapper(securityContext, prefix) - .createTableStaged(ns, createTableRequest)) - .build(); - } else { - return Response.ok( - newHandlerWrapper(securityContext, prefix) - .createTableStagedWithWriteDelegation(ns, createTableRequest)) - .build(); - } - } else if (delegationModes.isEmpty()) { - return Response.ok( - newHandlerWrapper(securityContext, prefix).createTableDirect(ns, createTableRequest)) - .build(); - } else { - return Response.ok( - newHandlerWrapper(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 @@ -288,7 +321,8 @@ public Response listTables( RealmContext realmContext, SecurityContext securityContext) { Namespace ns = decodeNamespace(namespace); - return Response.ok(newHandlerWrapper(securityContext, prefix).listTables(ns)).build(); + return withCatalog( + securityContext, prefix, catalog -> Response.ok(catalog.listTables(ns)).build()); } @Override @@ -304,16 +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(securityContext, prefix).loadTable(tableIdentifier, snapshots)) - .build(); - } else { - return Response.ok( - newHandlerWrapper(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 @@ -325,8 +360,13 @@ public Response tableExists( SecurityContext securityContext) { Namespace ns = decodeNamespace(namespace); TableIdentifier tableIdentifier = TableIdentifier.of(ns, RESTUtil.decodeString(table)); - newHandlerWrapper(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 @@ -339,13 +379,17 @@ public Response dropTable( SecurityContext securityContext) { Namespace ns = decodeNamespace(namespace); TableIdentifier tableIdentifier = TableIdentifier.of(ns, RESTUtil.decodeString(table)); - - if (purgeRequested != null && purgeRequested) { - newHandlerWrapper(securityContext, prefix).dropTableWithPurge(tableIdentifier); - } else { - newHandlerWrapper(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 @@ -356,9 +400,10 @@ public Response registerTable( RealmContext realmContext, SecurityContext securityContext) { Namespace ns = decodeNamespace(namespace); - return Response.ok( - newHandlerWrapper(securityContext, prefix).registerTable(ns, registerTableRequest)) - .build(); + return withCatalog( + securityContext, + prefix, + catalog -> Response.ok(catalog.registerTable(ns, registerTableRequest)).build()); } @Override @@ -367,8 +412,13 @@ public Response renameTable( RenameTableRequest renameTableRequest, RealmContext realmContext, SecurityContext securityContext) { - newHandlerWrapper(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 @@ -381,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(securityContext, prefix) - .updateTableForStagedCreate(tableIdentifier, commitTableRequest)) - .build(); - } else { - return Response.ok( - newHandlerWrapper(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 @@ -403,8 +453,10 @@ public Response createView( RealmContext realmContext, SecurityContext securityContext) { Namespace ns = decodeNamespace(namespace); - return Response.ok(newHandlerWrapper(securityContext, prefix).createView(ns, createViewRequest)) - .build(); + return withCatalog( + securityContext, + prefix, + catalog -> Response.ok(catalog.createView(ns, createViewRequest)).build()); } @Override @@ -416,7 +468,8 @@ public Response listViews( RealmContext realmContext, SecurityContext securityContext) { Namespace ns = decodeNamespace(namespace); - return Response.ok(newHandlerWrapper(securityContext, prefix).listViews(ns)).build(); + return withCatalog( + securityContext, prefix, catalog -> Response.ok(catalog.listViews(ns)).build()); } @Override @@ -428,8 +481,8 @@ public Response loadView( SecurityContext securityContext) { Namespace ns = decodeNamespace(namespace); TableIdentifier tableIdentifier = TableIdentifier.of(ns, RESTUtil.decodeString(view)); - return Response.ok(newHandlerWrapper(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(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(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(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(securityContext, prefix) - .replaceView(tableIdentifier, commitViewRequest)) - .build(); + return withCatalog( + securityContext, + prefix, + catalog -> Response.ok(catalog.replaceView(tableIdentifier, commitViewRequest)).build()); } @Override @@ -490,8 +558,13 @@ public Response commitTransaction( CommitTransactionRequest commitTransactionRequest, RealmContext realmContext, SecurityContext securityContext) { - newHandlerWrapper(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 @@ -515,9 +588,13 @@ public Response sendNotification( SecurityContext securityContext) { Namespace ns = decodeNamespace(namespace); TableIdentifier tableIdentifier = TableIdentifier.of(ns, RESTUtil.decodeString(table)); - newHandlerWrapper(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. */ 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 5ab20ae2e..a85dacf48 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 @@ -31,7 +31,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; @@ -77,7 +76,7 @@ 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; @@ -110,10 +109,10 @@ * 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; @@ -134,7 +133,7 @@ public class PolarisCatalogHandlerWrapper { private ViewCatalog viewCatalog = null; public PolarisCatalogHandlerWrapper( - CallContext callContext, + RealmContext realmContext, PolarisMetaStoreSession session, PolarisConfigurationStore configurationStore, PolarisDiagnostics diagnostics, @@ -144,7 +143,7 @@ public PolarisCatalogHandlerWrapper( CallContextCatalogFactory catalogFactory, String catalogName, PolarisAuthorizer authorizer) { - this.callContext = callContext; + this.realmContext = realmContext; this.session = session; this.entityManager = entityManager; this.metaStoreManager = metaStoreManager; @@ -177,10 +176,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, resolutionManifest); + realmContext, authenticatedPrincipal, resolutionManifest); this.namespaceCatalog = (baseCatalog instanceof SupportsNamespaces) ? (SupportsNamespaces) baseCatalog : null; this.viewCatalog = (baseCatalog instanceof ViewCatalog) ? (ViewCatalog) baseCatalog : null; @@ -226,6 +232,7 @@ private void authorizeBasicNamespaceOperationOrThrow( throw new NoSuchNamespaceException("Namespace does not exist: %s", namespace); } authorizer.authorizeOrThrow( + realmContext, authenticatedPrincipal, resolutionManifest.getAllActivatedCatalogRoleAndPrincipalRoles(), op, @@ -259,6 +266,7 @@ private void authorizeCreateNamespaceUnderNamespaceOperationOrThrow( throw new NoSuchNamespaceException("Namespace does not exist: %s", parentNamespace); } authorizer.authorizeOrThrow( + realmContext, authenticatedPrincipal, resolutionManifest.getAllActivatedCatalogRoleAndPrincipalRoles(), op, @@ -296,6 +304,7 @@ private void authorizeCreateTableLikeUnderNamespaceOperationOrThrow( throw new NoSuchNamespaceException("Namespace does not exist: %s", namespace); } authorizer.authorizeOrThrow( + realmContext, authenticatedPrincipal, resolutionManifest.getAllActivatedCatalogRoleAndPrincipalRoles(), op, @@ -328,6 +337,7 @@ private void authorizeBasicTableLikeOperationOrThrow( } } authorizer.authorizeOrThrow( + realmContext, authenticatedPrincipal, resolutionManifest.getAllActivatedCatalogRoleAndPrincipalRoles(), op, @@ -381,6 +391,7 @@ private void authorizeCollectionOfTableLikeOperationOrThrow( "View does not exist: %s", identifier))) .toList(); authorizer.authorizeOrThrow( + realmContext, authenticatedPrincipal, resolutionManifest.getAllActivatedCatalogRoleAndPrincipalRoles(), op, @@ -441,6 +452,7 @@ private void authorizeRenameTableLikeOperationOrThrow( PolarisResolvedPathWrapper secondary = resolutionManifest.getResolvedPath(dst.namespace(), true); authorizer.authorizeOrThrow( + realmContext, authenticatedPrincipal, resolutionManifest.getAllActivatedCatalogRoleAndPrincipalRoles(), op, @@ -454,7 +466,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) { @@ -477,20 +489,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); } } @@ -499,37 +508,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) { @@ -544,14 +527,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( @@ -559,15 +542,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) { @@ -584,7 +566,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( @@ -603,55 +585,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) { @@ -712,11 +691,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( @@ -735,26 +711,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) { @@ -762,7 +735,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) { @@ -807,7 +780,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( @@ -840,7 +813,9 @@ public LoadTableResponse loadTableWithAccessDelegation( .getCatalogType() .equals(org.apache.polaris.core.admin.model.Catalog.TypeEnum.EXTERNAL) && !configurationStore.getConfiguration( - catalogEntity, PolarisConfiguration.ALLOW_EXTERNAL_CATALOG_CREDENTIAL_VENDING)) { + realmContext, + catalogEntity, + PolarisConfiguration.ALLOW_EXTERNAL_CATALOG_CREDENTIAL_VENDING)) { throw new ForbiddenException( "Access Delegation is not enabled for this catalog. Please consult applicable " + "documentation for the catalog config property '%s' to enable this feature", @@ -849,32 +824,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) { @@ -915,9 +887,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( @@ -934,16 +904,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) { @@ -959,7 +927,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) { @@ -967,7 +935,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) { @@ -984,7 +952,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) { @@ -1056,6 +1024,7 @@ public void commitTransaction(CommitTransactionRequest commitTransactionRequest) .location() .equals(((MetadataUpdate.SetLocation) singleUpdate).location()) && !configurationStore.getConfiguration( + realmContext, PolarisConfiguration.ALLOW_NAMESPACE_LOCATION_OVERLAP)) { throw new BadRequestException( "Unsupported operation: commitTransaction containing SetLocation" @@ -1093,7 +1062,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) { @@ -1110,14 +1079,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) { @@ -1133,15 +1102,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) { @@ -1149,7 +1117,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) { @@ -1166,6 +1134,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 bf42753d5..1e0834e11 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 @@ -24,7 +24,7 @@ import jakarta.inject.Inject; import java.util.Map; 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 { @@ -53,8 +53,8 @@ public DefaultConfigurationStore( } @Override - public @Nullable T getConfiguration(String configName) { - String realm = CallContext.getCurrentContext().getRealmContext().getRealmIdentifier(); + public @Nullable T getConfiguration(RealmContext realmContext, String configName) { + String realm = realmContext.getRealmIdentifier(); @SuppressWarnings("unchecked") T confgValue = (T) 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 24551a2d0..3c233f9f5 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 @@ -20,12 +20,12 @@ 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, 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 71b21d0f3..000000000 --- a/service/common/src/main/java/org/apache/polaris/service/context/CallContextResolver.java +++ /dev/null @@ -1,33 +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 queryParams, - 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 9bef1431c..000000000 --- a/service/common/src/main/java/org/apache/polaris/service/context/DefaultCallContextResolver.java +++ /dev/null @@ -1,57 +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 java.util.Map; -import org.apache.polaris.core.context.CallContext; -import org.apache.polaris.core.context.RealmContext; -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); - - @Override - public CallContext resolveCallContext( - final RealmContext realmContext, - String method, - String path, - Map queryParams, - Map headers) { - LOGGER - .atDebug() - .addKeyValue("realmContext", realmContext.getRealmIdentifier()) - .addKeyValue("method", method) - .addKeyValue("path", path) - .addKeyValue("queryParams", queryParams) - .addKeyValue("headers", headers) - .log("Resolving CallContext"); - return CallContext.of(realmContext); - } -} 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 048cd1dfc..a2869fb7b 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 @@ -29,7 +29,7 @@ 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.PolarisEntityManager; @@ -78,32 +78,30 @@ public PolarisCallContextCatalogFactory( @Override public Catalog createCallContextCatalog( - CallContext context, + RealmContext realmContext, AuthenticatedPolarisPrincipal authenticatedPrincipal, final PolarisResolutionManifest resolvedManifest) { PolarisBaseEntity baseCatalogEntity = 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); BasePolarisCatalog catalogInstance = new BasePolarisCatalog( + realmContext, entityManager, metaStoreManager, metaStoreSession, configurationStore, diagnostics, - context, resolvedManifest, 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/storage/PolarisStorageIntegrationProviderImpl.java b/service/common/src/main/java/org/apache/polaris/service/storage/PolarisStorageIntegrationProviderImpl.java index f47a23e63..32b63034b 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 @@ -32,6 +32,7 @@ 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; @@ -112,6 +113,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/TableCleanupTaskHandler.java b/service/common/src/main/java/org/apache/polaris/service/task/TableCleanupTaskHandler.java index c9383c492..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 @@ -33,7 +33,6 @@ import org.apache.iceberg.io.FileIO; 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.entity.AsyncTaskType; import org.apache.polaris.core.entity.PolarisBaseEntity; @@ -133,6 +132,7 @@ public boolean handleTask(TaskEntity cleanupTask, RealmContext realmContext) { // TODO: handle partition statistics files Stream metadataFileCleanupTasks = getMetadataTaskStream( + realmContext, cleanupTask, tableMetadata, tableEntity, @@ -157,7 +157,7 @@ public boolean handleTask(TaskEntity cleanupTask, RealmContext realmContext) { .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()); @@ -222,6 +222,7 @@ private Stream getManifestTaskStream( } private Stream getMetadataTaskStream( + RealmContext realmContext, TaskEntity cleanupTask, TableMetadata tableMetadata, TableLikeEntity tableEntity, @@ -229,7 +230,7 @@ private Stream getMetadataTaskStream( PolarisMetaStoreSession metaStoreSession, PolarisConfigurationStore configurationStore, Clock clock) { - int batchSize = configurationStore.getConfiguration(BATCH_SIZE_CONFIG_KEY, 10); + int batchSize = configurationStore.getConfiguration(realmContext, BATCH_SIZE_CONFIG_KEY, 10); return getMetadataFileBatches(tableMetadata, batchSize).stream() .map( metadataBatch -> { 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 94659245c..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 @@ -30,7 +30,7 @@ import java.util.concurrent.TimeUnit; 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.entity.PolarisBaseEntity; import org.apache.polaris.core.entity.PolarisEntity; import org.apache.polaris.core.entity.PolarisEntityType; @@ -43,7 +43,7 @@ /** * 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); @@ -91,45 +91,37 @@ 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(ctx.getRealmContext()).get(); + metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(); PolarisBaseEntity taskEntity = metaStoreManager.loadEntity(metaStoreSession, 0L, taskEntityId).getEntity(); if (!PolarisEntityType.TASK.equals(taskEntity.getType())) { @@ -147,7 +139,7 @@ protected void handleTask(long taskEntityId, CallContext ctx, int attempt) { return; } TaskHandler handler = handlerOpt.get(); - boolean success = handler.handleTask(task, ctx.getRealmContext()); + boolean success = handler.handleTask(task, realmContext); if (success) { LOGGER .atInfo() 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 55c4db080..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 @@ -28,7 +28,6 @@ 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; @@ -65,13 +64,14 @@ public FileIO apply(TaskEntity task, RealmContext realmContext) { Boolean skipCredentialSubscopingIndirection = configurationStore.getConfiguration( + 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, metaStoreSession,