diff --git a/polaris-core/src/main/java/io/polaris/core/persistence/LocalPolarisMetaStoreManagerFactory.java b/polaris-core/src/main/java/io/polaris/core/persistence/LocalPolarisMetaStoreManagerFactory.java index faaf360e8..9288725cd 100644 --- a/polaris-core/src/main/java/io/polaris/core/persistence/LocalPolarisMetaStoreManagerFactory.java +++ b/polaris-core/src/main/java/io/polaris/core/persistence/LocalPolarisMetaStoreManagerFactory.java @@ -18,7 +18,6 @@ */ package io.polaris.core.persistence; -import com.fasterxml.jackson.annotation.JsonProperty; import io.polaris.core.PolarisCallContext; import io.polaris.core.PolarisDefaultDiagServiceImpl; import io.polaris.core.PolarisDiagnostics; @@ -34,12 +33,11 @@ import io.polaris.core.storage.cache.StorageCredentialCache; import java.util.Collections; import java.util.HashMap; -import java.util.HashSet; import java.util.List; import java.util.Map; -import java.util.Set; import java.util.function.Supplier; import org.jetbrains.annotations.NotNull; +import org.jetbrains.annotations.TestOnly; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -50,11 +48,6 @@ */ public abstract class LocalPolarisMetaStoreManagerFactory implements MetaStoreManagerFactory { - @JsonProperty("enable-bootstrap") - protected boolean bootstrapEnabled = false; - - final Set bootstrappedRealms = new HashSet<>(); - final Map metaStoreManagerMap = new HashMap<>(); final Map storageCredentialCacheMap = new HashMap<>(); final Map backingStoreMap = new HashMap<>(); @@ -114,18 +107,12 @@ public void purgeRealms(List realms) { backingStoreMap.remove(realm); sessionSupplierMap.remove(realm); metaStoreManagerMap.remove(realm); - bootstrappedRealms.remove(realm); } } @Override public synchronized PolarisMetaStoreManager getOrCreateMetaStoreManager( RealmContext realmContext) { - String realmId = realmContext.getRealmIdentifier(); - if (bootstrapEnabled && !bootstrappedRealms.contains(realmId)) { - bootstrapRealmAndPrintCredentials(realmId); - } - if (!metaStoreManagerMap.containsKey(realmContext.getRealmIdentifier())) { initializeForRealm(realmContext); checkPolarisServiceBootstrappedForRealm( @@ -137,11 +124,6 @@ public synchronized PolarisMetaStoreManager getOrCreateMetaStoreManager( @Override public synchronized Supplier getOrCreateSessionSupplier( RealmContext realmContext) { - String realmId = realmContext.getRealmIdentifier(); - if (bootstrapEnabled && !bootstrappedRealms.contains(realmId)) { - bootstrapRealmAndPrintCredentials(realmId); - } - if (!sessionSupplierMap.containsKey(realmContext.getRealmIdentifier())) { initializeForRealm(realmContext); checkPolarisServiceBootstrappedForRealm( @@ -228,11 +210,10 @@ public void setStorageIntegrationProvider(PolarisStorageIntegrationProvider stor return rotatedSecrets; } - private void bootstrapRealmAndPrintCredentials(String realmId) { + @TestOnly + public void bootstrapRealmAndPrintCredentials(String realmId) { Map results = this.bootstrapRealms(Collections.singletonList(realmId)); - bootstrappedRealms.add(realmId); - PolarisMetaStoreManager.PrincipalSecretsResult principalSecrets = results.get(realmId); String msg = diff --git a/polaris-service/src/main/java/io/polaris/service/persistence/InMemoryPolarisMetaStoreManagerFactory.java b/polaris-service/src/main/java/io/polaris/service/persistence/InMemoryPolarisMetaStoreManagerFactory.java index cec9ef528..a41f4861a 100644 --- a/polaris-service/src/main/java/io/polaris/service/persistence/InMemoryPolarisMetaStoreManagerFactory.java +++ b/polaris-service/src/main/java/io/polaris/service/persistence/InMemoryPolarisMetaStoreManagerFactory.java @@ -30,12 +30,6 @@ @JsonTypeName("in-memory") public class InMemoryPolarisMetaStoreManagerFactory extends LocalPolarisMetaStoreManagerFactory { - - public InMemoryPolarisMetaStoreManagerFactory() { - // Some tests use in-memory factory and require bootstrap. - bootstrapEnabled = true; - } - @Override protected PolarisTreeMapStore createBackingStore(@NotNull PolarisDiagnostics diagnostics) { return new PolarisTreeMapStore(diagnostics); diff --git a/polaris-service/src/test/java/io/polaris/service/admin/PolarisAuthzTestBase.java b/polaris-service/src/test/java/io/polaris/service/admin/PolarisAuthzTestBase.java index 56d8fcc59..a546f0541 100644 --- a/polaris-service/src/test/java/io/polaris/service/admin/PolarisAuthzTestBase.java +++ b/polaris-service/src/test/java/io/polaris/service/admin/PolarisAuthzTestBase.java @@ -144,13 +144,14 @@ public abstract class PolarisAuthzTestBase { @BeforeEach @SuppressWarnings("unchecked") public void before() { + RealmContext realmContext = () -> "realm"; PolarisDiagnostics diagServices = new PolarisDefaultDiagServiceImpl(); PolarisTreeMapStore backingStore = new PolarisTreeMapStore(diagServices); InMemoryPolarisMetaStoreManagerFactory managerFactory = new InMemoryPolarisMetaStoreManagerFactory(); + managerFactory.bootstrapRealmAndPrintCredentials(realmContext.getRealmIdentifier()); managerFactory.setStorageIntegrationProvider( new PolarisStorageIntegrationProviderImpl(Mockito::mock)); - RealmContext realmContext = () -> "realm"; PolarisMetaStoreManager metaStoreManager = managerFactory.getOrCreateMetaStoreManager(realmContext); diff --git a/polaris-service/src/test/java/io/polaris/service/catalog/BasePolarisCatalogTest.java b/polaris-service/src/test/java/io/polaris/service/catalog/BasePolarisCatalogTest.java index 56a15928f..5038b0065 100644 --- a/polaris-service/src/test/java/io/polaris/service/catalog/BasePolarisCatalogTest.java +++ b/polaris-service/src/test/java/io/polaris/service/catalog/BasePolarisCatalogTest.java @@ -133,6 +133,7 @@ public void before() { PolarisStorageIntegrationProvider storageIntegrationProvider = Mockito.mock(); InMemoryPolarisMetaStoreManagerFactory managerFactory = new InMemoryPolarisMetaStoreManagerFactory(); + managerFactory.bootstrapRealmAndPrintCredentials(realmContext.getRealmIdentifier()); managerFactory.setStorageIntegrationProvider(storageIntegrationProvider); metaStoreManager = managerFactory.getOrCreateMetaStoreManager(realmContext); Map configMap = new HashMap<>(); diff --git a/polaris-service/src/test/java/io/polaris/service/catalog/BasePolarisCatalogViewTest.java b/polaris-service/src/test/java/io/polaris/service/catalog/BasePolarisCatalogViewTest.java index 967420384..f1940546a 100644 --- a/polaris-service/src/test/java/io/polaris/service/catalog/BasePolarisCatalogViewTest.java +++ b/polaris-service/src/test/java/io/polaris/service/catalog/BasePolarisCatalogViewTest.java @@ -64,6 +64,7 @@ public void before() { RealmContext realmContext = () -> "realm"; InMemoryPolarisMetaStoreManagerFactory managerFactory = new InMemoryPolarisMetaStoreManagerFactory(); + managerFactory.bootstrapRealmAndPrintCredentials(realmContext.getRealmIdentifier()); managerFactory.setStorageIntegrationProvider( new PolarisStorageIntegrationProviderImpl(Mockito::mock)); PolarisMetaStoreManager metaStoreManager = diff --git a/polaris-service/src/test/java/io/polaris/service/task/ManifestFileCleanupTaskHandlerTest.java b/polaris-service/src/test/java/io/polaris/service/task/ManifestFileCleanupTaskHandlerTest.java index c79c34942..176721882 100644 --- a/polaris-service/src/test/java/io/polaris/service/task/ManifestFileCleanupTaskHandlerTest.java +++ b/polaris-service/src/test/java/io/polaris/service/task/ManifestFileCleanupTaskHandlerTest.java @@ -43,15 +43,22 @@ import org.apache.iceberg.io.FileIO; import org.apache.iceberg.io.OutputFile; import org.apache.iceberg.io.PositionOutputStream; +import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; class ManifestFileCleanupTaskHandlerTest { + private InMemoryPolarisMetaStoreManagerFactory metaStoreManagerFactory; + RealmContext realmContext; + + @BeforeEach + void setUp() { + metaStoreManagerFactory = new InMemoryPolarisMetaStoreManagerFactory(); + realmContext = () -> "realmName"; + metaStoreManagerFactory.bootstrapRealmAndPrintCredentials(realmContext.getRealmIdentifier()); + } @Test public void testCleanupFileNotExists() throws IOException { - InMemoryPolarisMetaStoreManagerFactory metaStoreManagerFactory = - new InMemoryPolarisMetaStoreManagerFactory(); - RealmContext realmContext = () -> "realmName"; PolarisCallContext polarisCallContext = new PolarisCallContext( metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(), @@ -83,9 +90,6 @@ public void testCleanupFileNotExists() throws IOException { @Test public void testCleanupFileManifestExistsDataFilesDontExist() throws IOException { - InMemoryPolarisMetaStoreManagerFactory metaStoreManagerFactory = - new InMemoryPolarisMetaStoreManagerFactory(); - RealmContext realmContext = () -> "realmName"; PolarisCallContext polarisCallContext = new PolarisCallContext( metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(), @@ -116,9 +120,6 @@ public void testCleanupFileManifestExistsDataFilesDontExist() throws IOException @Test public void testCleanupFiles() throws IOException { - InMemoryPolarisMetaStoreManagerFactory metaStoreManagerFactory = - new InMemoryPolarisMetaStoreManagerFactory(); - RealmContext realmContext = () -> "realmName"; PolarisCallContext polarisCallContext = new PolarisCallContext( metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(), @@ -166,9 +167,6 @@ public void close() { @Test public void testCleanupFilesWithRetries() throws IOException { - InMemoryPolarisMetaStoreManagerFactory metaStoreManagerFactory = - new InMemoryPolarisMetaStoreManagerFactory(); - RealmContext realmContext = () -> "realmName"; PolarisCallContext polarisCallContext = new PolarisCallContext( metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(), diff --git a/polaris-service/src/test/java/io/polaris/service/task/TableCleanupTaskHandlerTest.java b/polaris-service/src/test/java/io/polaris/service/task/TableCleanupTaskHandlerTest.java index d491cfc77..3435eb9f8 100644 --- a/polaris-service/src/test/java/io/polaris/service/task/TableCleanupTaskHandlerTest.java +++ b/polaris-service/src/test/java/io/polaris/service/task/TableCleanupTaskHandlerTest.java @@ -41,17 +41,30 @@ import org.apache.iceberg.inmemory.InMemoryFileIO; import org.apache.iceberg.io.FileIO; import org.assertj.core.api.Assertions; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.mockito.Mockito; import org.slf4j.LoggerFactory; class TableCleanupTaskHandlerTest { + private InMemoryPolarisMetaStoreManagerFactory metaStoreManagerFactory; + private RealmContext realmContext; + + @BeforeEach + void setUp() { + metaStoreManagerFactory = new InMemoryPolarisMetaStoreManagerFactory(); + realmContext = () -> "realmName"; + metaStoreManagerFactory.bootstrapRealmAndPrintCredentials(realmContext.getRealmIdentifier()); + } + + @AfterEach + void tearDown() { + metaStoreManagerFactory.purgeRealms(List.of(realmContext.getRealmIdentifier())); + } @Test public void testTableCleanup() throws IOException { - InMemoryPolarisMetaStoreManagerFactory metaStoreManagerFactory = - new InMemoryPolarisMetaStoreManagerFactory(); - RealmContext realmContext = () -> "realmName"; PolarisCallContext polarisCallContext = new PolarisCallContext( metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(), @@ -112,9 +125,6 @@ public void testTableCleanup() throws IOException { @Test public void testTableCleanupHandlesAlreadyDeletedMetadata() throws IOException { - InMemoryPolarisMetaStoreManagerFactory metaStoreManagerFactory = - new InMemoryPolarisMetaStoreManagerFactory(); - RealmContext realmContext = () -> "realmName"; PolarisCallContext polarisCallContext = new PolarisCallContext( metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(), @@ -174,9 +184,6 @@ public void close() { @Test public void testTableCleanupDuplicatesTasksIfFileStillExists() throws IOException { - InMemoryPolarisMetaStoreManagerFactory metaStoreManagerFactory = - new InMemoryPolarisMetaStoreManagerFactory(); - RealmContext realmContext = () -> "realmName"; PolarisCallContext polarisCallContext = new PolarisCallContext( metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(), @@ -269,9 +276,6 @@ public void close() { @Test public void testTableCleanupMultipleSnapshots() throws IOException { - InMemoryPolarisMetaStoreManagerFactory metaStoreManagerFactory = - new InMemoryPolarisMetaStoreManagerFactory(); - RealmContext realmContext = () -> "realmName"; PolarisCallContext polarisCallContext = new PolarisCallContext( metaStoreManagerFactory.getOrCreateSessionSupplier(realmContext).get(), diff --git a/polaris-service/src/test/java/io/polaris/service/test/PolarisConnectionExtension.java b/polaris-service/src/test/java/io/polaris/service/test/PolarisConnectionExtension.java index 4686e4c13..4a4c64384 100644 --- a/polaris-service/src/test/java/io/polaris/service/test/PolarisConnectionExtension.java +++ b/polaris-service/src/test/java/io/polaris/service/test/PolarisConnectionExtension.java @@ -50,11 +50,13 @@ import java.util.Arrays; import java.util.Comparator; import java.util.EnumMap; +import java.util.List; import java.util.Map; import java.util.Optional; import java.util.Set; import org.jetbrains.annotations.NotNull; import org.jetbrains.annotations.Nullable; +import org.junit.jupiter.api.extension.AfterAllCallback; import org.junit.jupiter.api.extension.BeforeAllCallback; import org.junit.jupiter.api.extension.ExtensionContext; import org.junit.jupiter.api.extension.ExtensionContext.Namespace; @@ -69,7 +71,8 @@ import software.amazon.awssdk.services.sts.model.AssumeRoleResponse; import software.amazon.awssdk.services.sts.model.Credentials; -public class PolarisConnectionExtension implements BeforeAllCallback, ParameterResolver { +public class PolarisConnectionExtension + implements BeforeAllCallback, AfterAllCallback, ParameterResolver { public static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); private MetaStoreManagerFactory metaStoreManagerFactory; @@ -99,6 +102,9 @@ public void beforeAll(ExtensionContext extensionContext) throws Exception { metaStoreManagerFactory = config.getMetaStoreManagerFactory(); if (metaStoreManagerFactory instanceof LocalPolarisMetaStoreManagerFactory msmf) { + ((LocalPolarisMetaStoreManagerFactory) metaStoreManagerFactory) + .bootstrapRealmAndPrintCredentials(realm); + StsClient mockSts = Mockito.mock(StsClient.class); Mockito.when(mockSts.assumeRole(Mockito.isA(AssumeRoleRequest.class))) .thenReturn( @@ -177,6 +183,13 @@ public EnumMap getSubscopedCreds( } } + @Override + public void afterAll(ExtensionContext context) { + if (metaStoreManagerFactory instanceof LocalPolarisMetaStoreManagerFactory) { + metaStoreManagerFactory.purgeRealms(List.of(realm)); + } + } + public static String getTestRealm(Class testClassName) { return testClassName.getName().replace('.', '_'); } @@ -228,7 +241,11 @@ static PolarisPrincipalSecrets getAdminSecrets() { public boolean supportsParameter( ParameterContext parameterContext, ExtensionContext extensionContext) throws ParameterResolutionException { - return parameterContext.getParameter().getType().equals(PolarisToken.class) + return parameterContext + .getParameter() + .getType() + .equals(PolarisConnectionExtension.PolarisToken.class) + || parameterContext.getParameter().getType().equals(MetaStoreManagerFactory.class) || parameterContext.getParameter().getType().equals(PolarisPrincipalSecrets.class); } diff --git a/polaris-service/src/test/resources/polaris-server-integrationtest.yml b/polaris-service/src/test/resources/polaris-server-integrationtest.yml index 0c7247de9..bdac81f3c 100644 --- a/polaris-service/src/test/resources/polaris-server-integrationtest.yml +++ b/polaris-service/src/test/resources/polaris-server-integrationtest.yml @@ -84,7 +84,6 @@ metaStoreManager: # url: http://sdp-devvm-mcollado:8080 type: eclipse-link # uncomment to use eclipse-link as metastore persistence-unit: polaris - enable-bootstrap: true oauth2: type: default