Add a command to explicitly purge the metastore (#59)

polaris-core/src/main/java/io/polaris/core/persistence/LocalPolarisMetaStoreManagerFactory.java

@@ -89,6 +89,22 @@ public synchronized Map<String, PolarisMetaStoreManager.PrincipalSecretsResult>
     return results;
   }
 
+  @Override
+  public void purgeRealms(List<String> realms) {
+    for (String realm : realms) {
+      PolarisMetaStoreManager metaStoreManager = getOrCreateMetaStoreManager(() -> realm);
+      PolarisMetaStoreSession session = getOrCreateSessionSupplier(() -> realm).get();
+
+      PolarisCallContext callContext = new PolarisCallContext(session, diagServices);
+      metaStoreManager.purge(callContext);
+
+      storageCredentialCacheMap.remove(realm);
+      backingStoreMap.remove(realm);
+      sessionSupplierMap.remove(realm);
+      metaStoreManagerMap.remove(realm);
+    }
+  }
+
   @Override
   public synchronized PolarisMetaStoreManager getOrCreateMetaStoreManager(
       RealmContext realmContext) {
@@ -147,6 +163,21 @@ public void setStorageIntegrationProvider(PolarisStorageIntegrationProvider stor
             sessionSupplierMap.get(realmContext.getRealmIdentifier()).get(), diagServices);
     CallContext.setCurrentContext(CallContext.of(realmContext, polarisContext));
 
+    PolarisMetaStoreManager.EntityResult preliminaryRootPrincipalLookup =
+        metaStoreManager.readEntityByName(
+            polarisContext,
+            null,
+            PolarisEntityType.PRINCIPAL,
+            PolarisEntitySubType.NULL_SUBTYPE,
+            PolarisEntityConstants.getRootPrincipalName());
+    if (preliminaryRootPrincipalLookup.isSuccess()) {
+      String overrideMessage =
+          "It appears this metastore manager has already been bootstrapped. "
+              + "To continue bootstrapping, please first purge the metastore with the `purge` command.";
+      logger.error("\n\n {} \n\n", overrideMessage);
+      throw new IllegalArgumentException(overrideMessage);
+    }
+
     metaStoreManager.bootstrapPolarisService(polarisContext);
 
     PolarisMetaStoreManager.EntityResult rootPrincipalLookup =

polaris-core/src/main/java/io/polaris/core/persistence/MetaStoreManagerFactory.java

@@ -43,4 +43,7 @@ public interface MetaStoreManagerFactory extends Discoverable {
   void setMetricRegistry(PolarisMetricRegistry metricRegistry);
 
   Map<String, PolarisMetaStoreManager.PrincipalSecretsResult> bootstrapRealms(List<String> realms);
+
+  /** Purge all metadata for the realms provided */
+  void purgeRealms(List<String> realms);
 }

polaris-core/src/main/java/io/polaris/core/persistence/PolarisMetaStoreManager.java

@@ -176,18 +176,28 @@ public boolean alreadyExists() {
   }
 
   /**
-   * Bootstrap the Polaris service, will remove ALL existing persisted entities, then will create
-   * the root catalog, root principal and associated service admin role.
+   * Bootstrap the Polaris service, creating the root catalog, root principal, and associated
+   * service admin role. Will fail if the service has already been bootstrapped.
+   *
+   * @param callCtx call context
+   * @return the result of the bootstrap attempt
+   */
+  @NotNull
+  BaseResult bootstrapPolarisService(@NotNull PolarisCallContext callCtx);
+
+  /**
+   * Purge all metadata associated with the Polaris service, resetting the metastore to the state it
+   * was in prior to bootstrapping.
    *
    * <p>*************************** WARNING ************************
    *
-   * <p>This will destroy whatever Polaris metadata exists in this account
+   * <p>This will destroy whatever Polaris metadata exists in the metastore
    *
    * @param callCtx call context
    * @return always success or unexpected error
    */
   @NotNull
-  BaseResult bootstrapPolarisService(@NotNull PolarisCallContext callCtx);
+  BaseResult purge(@NotNull PolarisCallContext callCtx);
 
   /** the return for an entity lookup call */
   class EntityResult extends BaseResult {

polaris-core/src/main/java/io/polaris/core/persistence/PolarisMetaStoreManagerImpl.java

@@ -52,13 +52,16 @@
 import java.util.stream.Collectors;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
  * Default implementation of the Polaris Meta Store Manager. Uses the underlying meta store to store
  * and retrieve all Polaris metadata
  */
 @SuppressFBWarnings("NP_NULL_ON_SOME_PATH_FROM_RETURN_VALUE")
 public class PolarisMetaStoreManagerImpl implements PolarisMetaStoreManager {
+  private static final Logger LOGGER = LoggerFactory.getLogger(PolarisMetaStoreManagerImpl.class);
 
   /** mapper, allows to serialize/deserialize properties to/from JSON */
   private static final ObjectMapper MAPPER = new ObjectMapper();
@@ -634,9 +637,6 @@ private void revokeGrantRecord(
   private void bootstrapPolarisService(
       @NotNull PolarisCallContext callCtx, @NotNull PolarisMetaStoreSession ms) {
 
-    // cleanup everything, start from a blank slate
-    ms.deleteAll(callCtx);
-
     // Create a root container entity that can represent the securable for any top-level grants.
     PolarisBaseEntity rootContainer =
         new PolarisBaseEntity(
@@ -706,6 +706,20 @@ private void bootstrapPolarisService(
     return new BaseResult(ReturnStatus.SUCCESS);
   }
 
+  @Override
+  public @NotNull BaseResult purge(@NotNull PolarisCallContext callCtx) {
+    // get meta store we should be using
+    PolarisMetaStoreSession ms = callCtx.getMetaStore();
+
+    // run operation in a read/write transaction
+    LOGGER.warn("Deleting all metadata in the metastore...");
+    ms.runActionInTransaction(callCtx, () -> ms.deleteAll(callCtx));
+    LOGGER.warn("Finished deleting all metadata in the metastore");
+
+    // all good
+    return new BaseResult(ReturnStatus.SUCCESS);
+  }
+
   /**
    * See {@link #readEntityByName(PolarisCallContext, List, PolarisEntityType, PolarisEntitySubType,
    * String)}

polaris-core/src/testFixtures/java/io/polaris/core/persistence/PolarisTestMetaStoreManager.java

@@ -65,6 +65,7 @@ public PolarisTestMetaStoreManager(
     this.doRetry = false;
 
     // bootstrap the Polaris service
+    polarisMetaStoreManager.purge(polarisCallContext);
     polarisMetaStoreManager.bootstrapPolarisService(polarisCallContext);
   }
 

polaris-service/src/main/java/io/polaris/service/BootstrapRealmsCommand.java

@@ -19,18 +19,23 @@
 import io.dropwizard.core.setup.Bootstrap;
 import io.polaris.core.PolarisConfigurationStore;
 import io.polaris.core.persistence.MetaStoreManagerFactory;
+import io.polaris.core.persistence.PolarisMetaStoreManager;
 import io.polaris.service.config.ConfigurationStoreAware;
 import io.polaris.service.config.PolarisApplicationConfig;
 import io.polaris.service.config.RealmEntityManagerFactory;
 import io.polaris.service.context.CallContextResolver;
+import java.util.Map;
 import net.sourceforge.argparse4j.inf.Namespace;
+import org.slf4j.Logger;
 
 /**
  * Command for bootstrapping root level service principals for each realm. This command will invoke
  * a default implementation which generates random user id and secret. These credentials will be
  * printed out to the log and standard output (stdout).
  */
 public class BootstrapRealmsCommand extends ConfiguredCommand<PolarisApplicationConfig> {
+  private Logger LOGGER = org.slf4j.LoggerFactory.getLogger(BootstrapRealmsCommand.class);
+
   public BootstrapRealmsCommand() {
     super("bootstrap", "bootstraps principal credentials for all realms and prints them to log");
   }
@@ -55,6 +60,27 @@ protected void run(
       csa.setConfigurationStore(configurationStore);
     }
 
-    metaStoreManagerFactory.bootstrapRealms(configuration.getDefaultRealms());
+    // Execute the bootstrap
+    Map<String, PolarisMetaStoreManager.PrincipalSecretsResult> results =
+        metaStoreManagerFactory.bootstrapRealms(configuration.getDefaultRealms());
+
+    // Log any errors:
+    boolean success = true;
+    for (Map.Entry<String, PolarisMetaStoreManager.PrincipalSecretsResult> result :
+        results.entrySet()) {
+      if (!result.getValue().isSuccess()) {
+        LOGGER.error(
+            "Bootstrapping `{}` failed: {}",
+            result.getKey(),
+            result.getValue().getReturnStatus().toString());
+        success = false;
+      }
+    }
+
+    if (success) {
+      LOGGER.info("Bootstrap completed successfully.");
+    } else {
+      LOGGER.error("Bootstrap encountered errors during operation.");
+    }
   }
 }

polaris-service/src/main/java/io/polaris/service/PolarisApplication.java

@@ -136,6 +136,7 @@ public void initialize(Bootstrap<PolarisApplicationConfig> bootstrap) {
     bootstrap.setConfigurationSourceProvider(provider);
 
     bootstrap.addCommand(new BootstrapRealmsCommand());
+    bootstrap.addCommand(new PurgeRealmsCommand());
   }
 
   @Override
@@ -286,6 +287,8 @@ public void run(PolarisApplicationConfig configuration, Environment environment)
     if (metaStoreManagerFactory instanceof InMemoryPolarisMetaStoreManagerFactory) {
       metaStoreManagerFactory.getOrCreateMetaStoreManager(configuration::getDefaultRealm);
     }
+
+    LOGGER.info("Server started successfully.");
   }
 
   private static OpenTelemetry setupTracing() {

polaris-service/src/main/java/io/polaris/service/PurgeRealmsCommand.java

@@ -0,0 +1,61 @@
+/*
+ * Copyright (c) 2024 Snowflake Computing Inc.
+ *
+ * Licensed 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 io.polaris.service;
+
+import io.dropwizard.core.cli.ConfiguredCommand;
+import io.dropwizard.core.setup.Bootstrap;
+import io.polaris.core.PolarisConfigurationStore;
+import io.polaris.core.persistence.MetaStoreManagerFactory;
+import io.polaris.service.config.ConfigurationStoreAware;
+import io.polaris.service.config.PolarisApplicationConfig;
+import io.polaris.service.config.RealmEntityManagerFactory;
+import io.polaris.service.context.CallContextResolver;
+import net.sourceforge.argparse4j.inf.Namespace;
+import org.slf4j.Logger;
+
+/** Command for purging all metadata associated with a realm */
+public class PurgeRealmsCommand extends ConfiguredCommand<PolarisApplicationConfig> {
+  private Logger LOGGER = org.slf4j.LoggerFactory.getLogger(PurgeRealmsCommand.class);
+
+  public PurgeRealmsCommand() {
+    super("purge", "purge principal credentials for all realms and prints them to log");
+  }
+
+  @Override
+  protected void run(
+      Bootstrap<PolarisApplicationConfig> bootstrap,
+      Namespace namespace,
+      PolarisApplicationConfig configuration)
+      throws Exception {
+    MetaStoreManagerFactory metaStoreManagerFactory = configuration.getMetaStoreManagerFactory();
+
+    PolarisConfigurationStore configurationStore = configuration.getConfigurationStore();
+    if (metaStoreManagerFactory instanceof ConfigurationStoreAware) {
+      ((ConfigurationStoreAware) metaStoreManagerFactory).setConfigurationStore(configurationStore);
+    }
+    RealmEntityManagerFactory entityManagerFactory =
+        new RealmEntityManagerFactory(metaStoreManagerFactory);
+    CallContextResolver callContextResolver = configuration.getCallContextResolver();
+    callContextResolver.setEntityManagerFactory(entityManagerFactory);
+    if (callContextResolver instanceof ConfigurationStoreAware csa) {
+      csa.setConfigurationStore(configurationStore);
+    }
+
+    metaStoreManagerFactory.purgeRealms(configuration.getDefaultRealms());
+
+    LOGGER.info("Purge completed successfully.");
+  }
+}

polaris-service/src/test/java/io/polaris/service/catalog/BasePolarisCatalogTest.java

@@ -1113,6 +1113,11 @@ public void setMetricRegistry(PolarisMetricRegistry metricRegistry) {}
                     throw new NotImplementedException("Bootstrapping realms is not supported");
                   }
 
+                  @Override
+                  public void purgeRealms(List<String> realms) {
+                    throw new NotImplementedException("Purging realms is not supported");
+                  }
+
                   @Override
                   public void setStorageIntegrationProvider(
                       PolarisStorageIntegrationProvider storageIntegrationProvider) {}