Explictly bootstrap from the tests

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<StoreType>
     implements MetaStoreManagerFactory {
-  @JsonProperty("enable-bootstrap")
-  protected boolean bootstrapEnabled = false;
-
-  final Set<String> bootstrappedRealms = new HashSet<>();
-
   final Map<String, PolarisMetaStoreManager> metaStoreManagerMap = new HashMap<>();
   final Map<String, StorageCredentialCache> storageCredentialCacheMap = new HashMap<>();
   final Map<String, StoreType> backingStoreMap = new HashMap<>();
@@ -114,18 +107,12 @@ public void purgeRealms(List<String> 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<PolarisMetaStoreSession> 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<String, PolarisMetaStoreManager.PrincipalSecretsResult> results =
         this.bootstrapRealms(Collections.singletonList(realmId));
-    bootstrappedRealms.add(realmId);
-
     PolarisMetaStoreManager.PrincipalSecretsResult principalSecrets = results.get(realmId);
 
     String msg =

polaris-service/src/main/java/io/polaris/service/persistence/InMemoryPolarisMetaStoreManagerFactory.java

@@ -30,12 +30,6 @@
 @JsonTypeName("in-memory")
 public class InMemoryPolarisMetaStoreManagerFactory
     extends LocalPolarisMetaStoreManagerFactory<PolarisTreeMapStore> {
-
-  public InMemoryPolarisMetaStoreManagerFactory() {
-    // Some tests use in-memory factory and require bootstrap.
-    bootstrapEnabled = true;
-  }
-
   @Override
   protected PolarisTreeMapStore createBackingStore(@NotNull PolarisDiagnostics diagnostics) {
     return new PolarisTreeMapStore(diagnostics);

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);
 

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<String, Object> configMap = new HashMap<>();

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 =

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(),

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(),

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<PolarisCredentialProperty, String> 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);
   }
 

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