JDBC: Fix Bootstrap with schema options

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

@@ -18,12 +18,8 @@
  */
 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"),
@@ -52,26 +48,16 @@ public static DatabaseType fromDisplayName(String displayName) {
    * 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-v3.sql";
-        case 1 -> schemaSuffix = "schema-v1.sql";
-        case 2 -> schemaSuffix = "schema-v2.sql";
-        case 3 -> schemaSuffix = "schema-v3.sql";
-        default ->
-            throw new IllegalArgumentException(
-                "Unknown schema version " + schemaOptions.schemaVersion());
-      }
-      ClassLoader classLoader = DatasourceOperations.class.getClassLoader();
-      return classLoader.getResourceAsStream(this.getDisplayName() + "/" + schemaSuffix);
+  public InputStream openInitScriptResource(int schemaVersion) {
+    // Preconditions check is simpler and more direct than a switch default
+    if (schemaVersion <= 0 || schemaVersion > 3) {
+      throw new IllegalArgumentException("Unknown or invalid schema version " + schemaVersion);
     }
+
+    final String resourceName =
+        String.format("%s/schema-v%d.sql", this.getDisplayName(), schemaVersion);
+
+    ClassLoader classLoader = DatasourceOperations.class.getClassLoader();
+    return classLoader.getResourceAsStream(resourceName);
   }
 }

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

@@ -52,9 +52,13 @@ public class DatasourceOperations {
 
   private static final Logger LOGGER = LoggerFactory.getLogger(DatasourceOperations.class);
 
+  // PG STATUS CODES
   private static final String CONSTRAINT_VIOLATION_SQL_CODE = "23505";
   private static final String RELATION_DOES_NOT_EXIST = "42P01";
 
+  // H2 STATUS CODES
+  private static final String H2_RELATION_DOES_NOT_EXIST = "90079";
+
   // POSTGRES RETRYABLE EXCEPTIONS
   private static final String SERIALIZATION_FAILURE_SQL_CODE = "40001";
 
@@ -396,7 +400,9 @@ public boolean isConstraintViolation(SQLException e) {
   }
 
   public boolean isRelationDoesNotExist(SQLException e) {
-    return RELATION_DOES_NOT_EXIST.equals(e.getSQLState());
+    return (RELATION_DOES_NOT_EXIST.equals(e.getSQLState())
+            && databaseType == DatabaseType.POSTGRES)
+        || (H2_RELATION_DOES_NOT_EXIST.equals(e.getSQLState()) && databaseType == DatabaseType.H2);
   }
 
   private Connection borrowConnection() throws SQLException {

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

@@ -748,14 +748,28 @@ static int loadSchemaVersion(
       }
       return schemaVersion.getFirst().getValue();
     } catch (SQLException e) {
-      LOGGER.error("Failed to load schema version due to {}", e.getMessage(), e);
       if (fallbackOnDoesNotExist && datasourceOperations.isRelationDoesNotExist(e)) {
         return SchemaVersion.MINIMUM.getValue();
       }
+      LOGGER.error("Failed to load schema version due to {}", e.getMessage(), e);
       throw new IllegalStateException("Failed to retrieve schema version", e);
     }
   }
 
+  static boolean entityTableExists(DatasourceOperations datasourceOperations) {
+    PreparedQuery query = QueryGenerator.generateEntityTableExistQuery();
+    try {
+      List<PolarisBaseEntity> entities =
+          datasourceOperations.executeSelect(query, new ModelEntity());
+      return entities != null && !entities.isEmpty();
+    } catch (SQLException e) {
+      if (datasourceOperations.isRelationDoesNotExist(e)) {
+        return false;
+      }
+      throw new IllegalStateException("Failed to check if Entities table exists", e);
+    }
+  }
+
   /** {@inheritDoc} */
   @Override
   public <T extends PolarisEntity & LocationBasedEntity>

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

@@ -0,0 +1,91 @@
+/*
+ * 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.persistence.relational.jdbc;
+
+import java.util.Optional;
+import org.apache.polaris.core.persistence.bootstrap.BootstrapOptions;
+import org.apache.polaris.core.persistence.bootstrap.SchemaOptions;
+
+public class JdbcBootstrapUtils {
+
+  private JdbcBootstrapUtils() {}
+
+  /**
+   * Determines the correct schema version to use for bootstrapping a realm.
+   *
+   * @param currentSchemaVersion The current version of the database schema.
+   * @param requiredSchemaVersion The requested schema version (-1 for auto-detection).
+   * @param hasAlreadyBootstrappedRealms Flag indicating if any realms already exist.
+   * @return The calculated bootstrap schema version.
+   * @throws IllegalStateException if the combination of parameters represents an invalid state.
+   */
+  public static int getRealmBootstrapSchemaVersion(
+      int currentSchemaVersion, int requiredSchemaVersion, boolean hasAlreadyBootstrappedRealms) {
+
+    // If versions already match, no change is needed.
+    if (currentSchemaVersion == requiredSchemaVersion) {
+      return requiredSchemaVersion;
+    }
+
+    // Handle fresh installations where no schema version is recorded (version 0).
+    if (currentSchemaVersion == 0) {
+      if (hasAlreadyBootstrappedRealms) {
+        // System was bootstrapped with v1 before schema versioning was introduced.
+        if (requiredSchemaVersion == -1 || requiredSchemaVersion == 1) {
+          return 1;
+        }
+      } else {
+        // A truly fresh start. Default to v3 for auto-detection, otherwise use the specified
+        // version.
+        return requiredSchemaVersion == -1 ? 3 : requiredSchemaVersion;
+      }
+    }
+
+    // Handle auto-detection on an existing installation (current version > 0).
+    if (requiredSchemaVersion == -1) {
+      // Use the current version if realms already exist; otherwise, use v3 for the new realm.
+      return hasAlreadyBootstrappedRealms ? currentSchemaVersion : 3;
+    }
+
+    // Any other combination is an unhandled or invalid migration path.
+    throw new IllegalStateException(
+        String.format(
+            "Cannot determine bootstrap schema version. Current: %d, Required: %d, Bootstrapped: %b",
+            currentSchemaVersion, requiredSchemaVersion, hasAlreadyBootstrappedRealms));
+  }
+
+  /**
+   * Extracts the requested schema version from the provided BootstrapOptions.
+   *
+   * @param bootstrapOptions: The bootstrap options containing schema information from which to
+   *     extract the version.
+   * @return The requested schema version, or -1 if not specified.
+   */
+  public static int getRequestedSchemaVersion(BootstrapOptions bootstrapOptions) {
+    SchemaOptions schemaOptions = bootstrapOptions.schemaOptions();
+    if (schemaOptions != null) {
+      Optional<Integer> version = schemaOptions.schemaVersion();
+      if (version.isPresent()) {
+        return version.get();
+      }
+    }
+    return -1;
+  }
+}

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

@@ -154,12 +154,27 @@ public synchronized Map<String, PrincipalSecretsResult> bootstrapRealms(
       RealmContext realmContext = () -> realm;
       if (!metaStoreManagerMap.containsKey(realm)) {
         DatasourceOperations datasourceOperations = getDatasourceOperations();
+        int currentSchemaVersion =
+            JdbcBasePersistenceImpl.loadSchemaVersion(
+                datasourceOperations,
+                configurationStore.getConfiguration(
+                    realmContext, BehaviorChangeConfiguration.SCHEMA_VERSION_FALL_BACK_ON_DNE));
+        int requestedSchemaVersion = JdbcBootstrapUtils.getRequestedSchemaVersion(bootstrapOptions);
+        int effectiveSchemaVersion =
+            JdbcBootstrapUtils.getRealmBootstrapSchemaVersion(
+                currentSchemaVersion,
+                requestedSchemaVersion,
+                JdbcBasePersistenceImpl.entityTableExists(datasourceOperations));
+        LOGGER.info(
+            "Effective schema version: {} for bootstrapping realm: {}",
+            effectiveSchemaVersion,
+            realm);
         try {
           // Run the set-up script to create the tables.
           datasourceOperations.executeScript(
               datasourceOperations
                   .getDatabaseType()
-                  .openInitScriptResource(bootstrapOptions.schemaOptions()));
+                  .openInitScriptResource(effectiveSchemaVersion));
         } catch (SQLException e) {
           throw new RuntimeException(
               String.format("Error executing sql script: %s", e.getMessage()), e);

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

@@ -256,6 +256,14 @@ static PreparedQuery generateVersionQuery() {
     return new PreparedQuery("SELECT version_value FROM POLARIS_SCHEMA.VERSION", List.of());
   }
 
+  @VisibleForTesting
+  static PreparedQuery generateEntityTableExistQuery() {
+    return new PreparedQuery(
+        String.format(
+            "SELECT * FROM %s LIMIT 1", getFullyQualifiedTableName(ModelEntity.TABLE_NAME)),
+        List.of());
+  }
+
   /**
    * Generate a SELECT query to find any entities that have a given realm & parent and that may
    * overlap with a given location. The check is performed without consideration for the scheme, so

persistence/relational-jdbc/src/main/resources/h2/schema-v3.sql

@@ -30,7 +30,7 @@ CREATE TABLE IF NOT EXISTS version (
 
 MERGE INTO version (version_key, version_value)
     KEY (version_key)
-    VALUES ('version', 2);
+    VALUES ('version', 3);
 
 -- H2 supports COMMENT, but some modes may ignore it
 COMMENT ON TABLE version IS 'the version of the JDBC schema in use';