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,29 @@ public static DatabaseType fromDisplayName(String displayName) {
     };
   }
 
-  public String getInitScriptResource() {
-    return String.format("%s/schema-v2.sql", this.getDisplayName());
+  /**
+   * Open an InputStream that contains data from an init script. This stream should be closed by the
+   * caller.
+   */
+  public InputStream openInitScriptResource(@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 null -> schemaSuffix = "schema-v2.sql";
+        case 1 -> schemaSuffix = "schema-v1.sql";
+        case 2 -> schemaSuffix = "schema-v2.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;
@@ -75,46 +76,51 @@ DatabaseType getDatabaseType() {
   }
 
   /**
-   * Execute SQL script.
+   * Execute SQL script and close the associated input stream
    *
-   * @param scriptFilePath : Path of SQL script.
+   * @param scriptInputStream : Input stream containing the SQL script.
    * @throws SQLException : Exception while executing the script.
    */
-  public void executeScript(String scriptFilePath) throws SQLException {
-    ClassLoader classLoader = DatasourceOperations.class.getClassLoader();
-    runWithinTransaction(
-        connection -> {
-          try (Statement statement = connection.createStatement()) {
-            BufferedReader reader =
-                new BufferedReader(
-                    new InputStreamReader(
-                        Objects.requireNonNull(classLoader.getResourceAsStream(scriptFilePath)),
-                        UTF_8));
-            StringBuilder sqlBuffer = new StringBuilder();
-            String line;
-            while ((line = reader.readLine()) != null) {
-              line = line.trim();
-              if (!line.isEmpty() && !line.startsWith("--")) { // Ignore empty lines and comments
-                sqlBuffer.append(line).append("\n");
-                if (line.endsWith(";")) { // Execute statement when semicolon is found
-                  String sql = sqlBuffer.toString().trim();
-                  try {
-                    // since SQL is directly read from the file, there is close to 0 possibility
-                    // of this being injected plus this run via an Admin tool, if attacker can
-                    // fiddle with this that means lot of other things are already compromised.
-                    statement.execute(sql);
-                  } catch (SQLException e) {
-                    throw new RuntimeException(e);
+  public void executeScript(InputStream scriptInputStream) throws SQLException {
+    try {
+      runWithinTransaction(
+          connection -> {
+            try (Statement statement = connection.createStatement();
+                BufferedReader reader =
+                    new BufferedReader(
+                        new InputStreamReader(Objects.requireNonNull(scriptInputStream), UTF_8))) {
+              StringBuilder sqlBuffer = new StringBuilder();
+              String line;
+              while ((line = reader.readLine()) != null) {
+                line = line.trim();
+                if (!line.isEmpty() && !line.startsWith("--")) { // Ignore empty lines and comments
+                  sqlBuffer.append(line).append("\n");
+                  if (line.endsWith(";")) { // Execute statement when semicolon is found
+                    String sql = sqlBuffer.toString().trim();
+                    try {
+                      // since SQL is directly read from the file, there is close to 0 possibility
+                      // of this being injected plus this run via an Admin tool, if attacker can
+                      // fiddle with this that means lot of other things are already compromised.
+                      statement.execute(sql);
+                    } catch (SQLException e) {
+                      throw new RuntimeException(e);
+                    }
+                    sqlBuffer.setLength(0); // Clear the buffer for the next statement
                   }
-                  sqlBuffer.setLength(0); // Clear the buffer for the next statement
                 }
               }
+              return true;
+            } catch (IOException e) {
+              throw new RuntimeException(e);
             }
-            return true;
-          } catch (IOException e) {
-            throw new RuntimeException(e);
-          }
-        });
+          });
+    } finally {
+      try {
+        scriptInputStream.close();
+      } catch (IOException e) {
+        LOGGER.error("Failed to close input stream: {}", e.getMessage());
+      }
+    }
   }
 
   /**

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()
+                  .openInitScriptResource(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-v2.sql", DatabaseType.H2.getDisplayName()));
+      ClassLoader classLoader = DatasourceOperations.class.getClassLoader();
+      InputStream scriptStream =
+          classLoader.getResourceAsStream(
+              String.format("%s/schema-v2.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,40 @@
+/*
+ * 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 {
+  @Nullable
+  Integer schemaVersion();
+
+  @Nullable
+  String schemaFile();
+
+  @Value.Check
+  default void validate() {
+    if (schemaVersion() != null && schemaFile() != null) {
+      throw new IllegalStateException("Only one of schemaVersion or schemaFile can be set.");
+    }
+  }
+}