Add options to the bootstrap command to specify a schema file

persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/DatabaseType.java

@@ -18,7 +18,12 @@
  */
 package org.apache.polaris.persistence.relational.jdbc;
 
+import jakarta.annotation.Nonnull;
+import java.io.FileInputStream;
+import java.io.IOException;
+import java.io.InputStream;
 import java.util.Locale;
+import org.apache.polaris.core.persistence.bootstrap.SchemaOptions;
 
 public enum DatabaseType {
   POSTGRES("postgres"),
@@ -43,7 +48,24 @@ public static DatabaseType fromDisplayName(String displayName) {
     };
   }
 
-  public String getInitScriptResource() {
-    return String.format("%s/schema-v1.sql", this.getDisplayName());
+  public InputStream getInitScriptResource(@Nonnull SchemaOptions schemaOptions) {
+    if (schemaOptions.schemaFile() != null) {
+      try {
+        return new FileInputStream(schemaOptions.schemaFile());
+      } catch (IOException e) {
+        throw new IllegalArgumentException("Unable to load file " + schemaOptions.schemaFile(), e);
+      }
+    } else {
+      final String schemaSuffix;
+      switch (schemaOptions.schemaVersion()) {
+        case SchemaOptions.LATEST -> schemaSuffix = "schema-v1.sql";
+        case "1" -> schemaSuffix = "schema-v1.sql";
+        default ->
+            throw new IllegalArgumentException(
+                "Unknown schema version " + schemaOptions.schemaVersion());
+      }
+      ClassLoader classLoader = DatasourceOperations.class.getClassLoader();
+      return classLoader.getResourceAsStream(this.getDisplayName() + "/" + schemaSuffix);
+    }
   }
 }

persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/DatasourceOperations.java

@@ -24,6 +24,7 @@
 import jakarta.annotation.Nonnull;
 import java.io.BufferedReader;
 import java.io.IOException;
+import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.sql.Connection;
 import java.sql.PreparedStatement;
@@ -80,16 +81,13 @@ DatabaseType getDatabaseType() {
    * @param scriptFilePath : Path of SQL script.
    * @throws SQLException : Exception while executing the script.
    */
-  public void executeScript(String scriptFilePath) throws SQLException {
-    ClassLoader classLoader = DatasourceOperations.class.getClassLoader();
+  public void executeScript(InputStream scriptInputStream) throws SQLException {
     runWithinTransaction(
         connection -> {
           try (Statement statement = connection.createStatement()) {
             BufferedReader reader =
                 new BufferedReader(
-                    new InputStreamReader(
-                        Objects.requireNonNull(classLoader.getResourceAsStream(scriptFilePath)),
-                        UTF_8));
+                    new InputStreamReader(Objects.requireNonNull(scriptInputStream), UTF_8));
             StringBuilder sqlBuffer = new StringBuilder();
             String line;
             while ((line = reader.readLine()) != null) {

persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcMetaStoreManagerFactory.java

@@ -43,7 +43,11 @@
 import org.apache.polaris.core.persistence.MetaStoreManagerFactory;
 import org.apache.polaris.core.persistence.PolarisMetaStoreManager;
 import org.apache.polaris.core.persistence.PrincipalSecretsGenerator;
+import org.apache.polaris.core.persistence.bootstrap.BootstrapOptions;
+import org.apache.polaris.core.persistence.bootstrap.ImmutableBootstrapOptions;
+import org.apache.polaris.core.persistence.bootstrap.ImmutableSchemaOptions;
 import org.apache.polaris.core.persistence.bootstrap.RootCredentialsSet;
+import org.apache.polaris.core.persistence.bootstrap.SchemaOptions;
 import org.apache.polaris.core.persistence.cache.EntityCache;
 import org.apache.polaris.core.persistence.cache.InMemoryEntityCache;
 import org.apache.polaris.core.persistence.dao.entity.BaseResult;
@@ -93,13 +97,14 @@ protected PolarisMetaStoreManager createNewMetaStoreManager() {
   }
 
   private void initializeForRealm(
-      RealmContext realmContext, RootCredentialsSet rootCredentialsSet, boolean isBootstrap) {
-    DatasourceOperations databaseOperations = getDatasourceOperations(isBootstrap);
+      DatasourceOperations datasourceOperations,
+      RealmContext realmContext,
+      RootCredentialsSet rootCredentialsSet) {
     sessionSupplierMap.put(
         realmContext.getRealmIdentifier(),
         () ->
             new JdbcBasePersistenceImpl(
-                databaseOperations,
+                datasourceOperations,
                 secretsGenerator(realmContext, rootCredentialsSet),
                 storageIntegrationProvider,
                 realmContext.getRealmIdentifier()));
@@ -108,35 +113,52 @@ private void initializeForRealm(
     metaStoreManagerMap.put(realmContext.getRealmIdentifier(), metaStoreManager);
   }
 
-  private DatasourceOperations getDatasourceOperations(boolean isBootstrap) {
+  public DatasourceOperations getDatasourceOperations() {
     DatasourceOperations databaseOperations;
     try {
       databaseOperations = new DatasourceOperations(dataSource.get(), relationalJdbcConfiguration);
     } catch (SQLException sqlException) {
       throw new RuntimeException(sqlException);
     }
-    if (isBootstrap) {
-      try {
-        // Run the set-up script to create the tables.
-        databaseOperations.executeScript(
-            databaseOperations.getDatabaseType().getInitScriptResource());
-      } catch (SQLException e) {
-        throw new RuntimeException(
-            String.format("Error executing sql script: %s", e.getMessage()), e);
-      }
-    }
     return databaseOperations;
   }
 
   @Override
   public synchronized Map<String, PrincipalSecretsResult> bootstrapRealms(
       Iterable<String> realms, RootCredentialsSet rootCredentialsSet) {
+    SchemaOptions schemaOptions = ImmutableSchemaOptions.builder().build();
+
+    BootstrapOptions bootstrapOptions =
+        ImmutableBootstrapOptions.builder()
+            .realms(realms)
+            .rootCredentialsSet(rootCredentialsSet)
+            .schemaOptions(schemaOptions)
+            .build();
+
+    return bootstrapRealms(bootstrapOptions);
+  }
+
+  @Override
+  public synchronized Map<String, PrincipalSecretsResult> bootstrapRealms(
+      BootstrapOptions bootstrapOptions) {
     Map<String, PrincipalSecretsResult> results = new HashMap<>();
 
-    for (String realm : realms) {
+    for (String realm : bootstrapOptions.realms()) {
       RealmContext realmContext = () -> realm;
       if (!metaStoreManagerMap.containsKey(realm)) {
-        initializeForRealm(realmContext, rootCredentialsSet, true);
+        DatasourceOperations datasourceOperations = getDatasourceOperations();
+        try {
+          // Run the set-up script to create the tables.
+          datasourceOperations.executeScript(
+              datasourceOperations
+                  .getDatabaseType()
+                  .getInitScriptResource(bootstrapOptions.schemaOptions()));
+        } catch (SQLException e) {
+          throw new RuntimeException(
+              String.format("Error executing sql script: %s", e.getMessage()), e);
+        }
+        initializeForRealm(
+            datasourceOperations, realmContext, bootstrapOptions.rootCredentialsSet());
         PrincipalSecretsResult secretsResult =
             bootstrapServiceAndCreatePolarisPrincipalForRealm(
                 realmContext, metaStoreManagerMap.get(realm));
@@ -172,7 +194,8 @@ public Map<String, BaseResult> purgeRealms(Iterable<String> realms) {
   public synchronized PolarisMetaStoreManager getOrCreateMetaStoreManager(
       RealmContext realmContext) {
     if (!metaStoreManagerMap.containsKey(realmContext.getRealmIdentifier())) {
-      initializeForRealm(realmContext, null, false);
+      DatasourceOperations datasourceOperations = getDatasourceOperations();
+      initializeForRealm(datasourceOperations, realmContext, null);
       checkPolarisServiceBootstrappedForRealm(
           realmContext, metaStoreManagerMap.get(realmContext.getRealmIdentifier()));
     }
@@ -183,7 +206,8 @@ public synchronized PolarisMetaStoreManager getOrCreateMetaStoreManager(
   public synchronized Supplier<BasePersistence> getOrCreateSessionSupplier(
       RealmContext realmContext) {
     if (!sessionSupplierMap.containsKey(realmContext.getRealmIdentifier())) {
-      initializeForRealm(realmContext, null, false);
+      DatasourceOperations datasourceOperations = getDatasourceOperations();
+      initializeForRealm(datasourceOperations, realmContext, null);
       checkPolarisServiceBootstrappedForRealm(
           realmContext, metaStoreManagerMap.get(realmContext.getRealmIdentifier()));
     } else {

persistence/relational-jdbc/src/test/java/org/apache/polaris/persistence/relational/jdbc/AtomicMetastoreManagerWithJdbcBasePersistenceImplTest.java

@@ -20,6 +20,7 @@
 
 import static org.apache.polaris.core.persistence.PrincipalSecretsGenerator.RANDOM_SECRETS;
 
+import java.io.InputStream;
 import java.sql.SQLException;
 import java.time.ZoneId;
 import java.util.Optional;
@@ -49,8 +50,11 @@ protected PolarisTestMetaStoreManager createPolarisTestMetaStoreManager() {
     try {
       datasourceOperations =
           new DatasourceOperations(createH2DataSource(), new H2JdbcConfiguration());
-      datasourceOperations.executeScript(
-          String.format("%s/schema-v1.sql", DatabaseType.H2.getDisplayName()));
+      ClassLoader classLoader = DatasourceOperations.class.getClassLoader();
+      InputStream scriptStream =
+          classLoader.getResourceAsStream(
+              String.format("%s/schema-v1.sql", DatabaseType.H2.getDisplayName()));
+      datasourceOperations.executeScript(scriptStream);
     } catch (SQLException e) {
       throw new RuntimeException(
           String.format(

polaris-core/src/main/java/org/apache/polaris/core/persistence/MetaStoreManagerFactory.java

@@ -21,6 +21,7 @@
 import java.util.Map;
 import java.util.function.Supplier;
 import org.apache.polaris.core.context.RealmContext;
+import org.apache.polaris.core.persistence.bootstrap.BootstrapOptions;
 import org.apache.polaris.core.persistence.bootstrap.RootCredentialsSet;
 import org.apache.polaris.core.persistence.cache.EntityCache;
 import org.apache.polaris.core.persistence.dao.entity.BaseResult;
@@ -41,6 +42,10 @@ public interface MetaStoreManagerFactory {
   Map<String, PrincipalSecretsResult> bootstrapRealms(
       Iterable<String> realms, RootCredentialsSet rootCredentialsSet);
 
+  default Map<String, PrincipalSecretsResult> bootstrapRealms(BootstrapOptions bootstrapOptions) {
+    return bootstrapRealms(bootstrapOptions.realms(), bootstrapOptions.rootCredentialsSet());
+  }
+
   /** Purge all metadata for the realms provided */
   Map<String, BaseResult> purgeRealms(Iterable<String> realms);
 }

polaris-core/src/main/java/org/apache/polaris/core/persistence/bootstrap/BootstrapOptions.java

@@ -0,0 +1,31 @@
+/*
+ * 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.persistence.bootstrap;
+
+import org.apache.polaris.immutables.PolarisImmutable;
+
+@PolarisImmutable
+public interface BootstrapOptions {
+  Iterable<String> realms();
+
+  RootCredentialsSet rootCredentialsSet();
+
+  SchemaOptions schemaOptions();
+}

polaris-core/src/main/java/org/apache/polaris/core/persistence/bootstrap/SchemaOptions.java

@@ -0,0 +1,37 @@
+/*
+ * 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.persistence.bootstrap;
+
+import jakarta.annotation.Nullable;
+import org.apache.polaris.immutables.PolarisImmutable;
+import org.immutables.value.Value;
+
+@PolarisImmutable
+public interface SchemaOptions {
+  public static final String LATEST = "LATEST";
+
+  @Value.Default
+  default String schemaVersion() {
+    return LATEST;
+  }
+
+  @Nullable
+  String schemaFile();
+}

runtime/admin/src/main/java/org/apache/polaris/admintool/BootstrapCommand.java

@@ -21,7 +21,11 @@
 import java.nio.file.Path;
 import java.util.List;
 import java.util.Map;
+import org.apache.polaris.core.persistence.bootstrap.BootstrapOptions;
+import org.apache.polaris.core.persistence.bootstrap.ImmutableBootstrapOptions;
+import org.apache.polaris.core.persistence.bootstrap.ImmutableSchemaOptions;
 import org.apache.polaris.core.persistence.bootstrap.RootCredentialsSet;
+import org.apache.polaris.core.persistence.bootstrap.SchemaOptions;
 import org.apache.polaris.core.persistence.dao.entity.PrincipalSecretsResult;
 import picocli.CommandLine;
 
@@ -42,6 +46,9 @@ static class InputOptions {
     @CommandLine.ArgGroup(multiplicity = "1")
     FileInputOptions fileOptions;
 
+    @CommandLine.ArgGroup(multiplicity = "1")
+    SchemaInputOptions schemaInputOptions;
+
     static class StandardInputOptions {
 
       @CommandLine.Option(
@@ -71,6 +78,23 @@ static class FileInputOptions {
           description = "A file containing root principal credentials to bootstrap.")
       Path file;
     }
+
+    static class SchemaInputOptions {
+      @CommandLine.Option(
+          names = {"-v", "--schema-version"},
+          paramLabel = "<schema version>",
+          description = "The version of the schema to load in [1, 2, LATEST].",
+          defaultValue = SchemaOptions.LATEST)
+      String schemaVersion;
+
+      @CommandLine.Option(
+          names = {"--schema-file"},
+          paramLabel = "<schema file>",
+          description =
+              "A schema file to bootstrap from. If unset, the bundled files will be used.",
+          defaultValue = "")
+      String schemaFile;
+    }
   }
 
   @Override
@@ -103,9 +127,27 @@ public Integer call() {
         }
       }
 
+      final SchemaOptions schemaOptions;
+      if (inputOptions.schemaInputOptions != null) {
+        schemaOptions =
+            ImmutableSchemaOptions.builder()
+                .schemaFile(inputOptions.schemaInputOptions.schemaFile)
+                .schemaVersion(inputOptions.schemaInputOptions.schemaVersion)
+                .build();
+      } else {
+        schemaOptions = ImmutableSchemaOptions.builder().build();
+      }
+
+      BootstrapOptions bootstrapOptions =
+          ImmutableBootstrapOptions.builder()
+              .realms(realms)
+              .rootCredentialsSet(rootCredentialsSet)
+              .schemaOptions(schemaOptions)
+              .build();
+
       // Execute the bootstrap
       Map<String, PrincipalSecretsResult> results =
-          metaStoreManagerFactory.bootstrapRealms(realms, rootCredentialsSet);
+          metaStoreManagerFactory.bootstrapRealms(bootstrapOptions);
 
       // Log any errors:
       boolean success = true;

runtime/admin/src/test/java/org/apache/polaris/admintool/BootstrapCommandTestBase.java

@@ -91,8 +91,9 @@ public void testBootstrapInvalidCredentials(LaunchResult result) {
   public void testBootstrapInvalidArguments(LaunchResult result) {
     assertThat(result.getErrorOutput())
         .contains(
-            "Error: (-r=<realm> [-r=<realm>]... [-c=<realm,clientId,clientSecret>]... [-p]) "
-                + "and -f=<file> are mutually exclusive (specify only one)");
+            "(-r=<realm> [-r=<realm>]... [-c=<realm,clientId,clientSecret>]... [-p]) and -f=<file> "
+                + "and ([-v=<schema version>] | [--schema-file=<schema file>]) are mutually exclusive "
+                + "(specify only one)");
   }
 
   @Test