diff --git a/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/DatabaseType.java b/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/DatabaseType.java index b0f5dcd8c7..f731eb5508 100644 --- a/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/DatabaseType.java +++ b/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); + } } } diff --git a/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/DatasourceOperations.java b/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/DatasourceOperations.java index 7017f1919d..43d57ead59 100644 --- a/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/DatasourceOperations.java +++ b/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()); + } + } } /** diff --git a/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcMetaStoreManagerFactory.java b/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcMetaStoreManagerFactory.java index 118fcadd4f..a69410a94f 100644 --- a/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcMetaStoreManagerFactory.java +++ b/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 bootstrapRealms( Iterable 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 bootstrapRealms( + BootstrapOptions bootstrapOptions) { Map 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 purgeRealms(Iterable 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 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 { diff --git a/persistence/relational-jdbc/src/test/java/org/apache/polaris/persistence/relational/jdbc/AtomicMetastoreManagerWithJdbcBasePersistenceImplTest.java b/persistence/relational-jdbc/src/test/java/org/apache/polaris/persistence/relational/jdbc/AtomicMetastoreManagerWithJdbcBasePersistenceImplTest.java index a824a17a56..bbb57e68ba 100644 --- a/persistence/relational-jdbc/src/test/java/org/apache/polaris/persistence/relational/jdbc/AtomicMetastoreManagerWithJdbcBasePersistenceImplTest.java +++ b/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( diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/MetaStoreManagerFactory.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/MetaStoreManagerFactory.java index cb2523891f..5b5dadaa0f 100644 --- a/polaris-core/src/main/java/org/apache/polaris/core/persistence/MetaStoreManagerFactory.java +++ b/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 bootstrapRealms( Iterable realms, RootCredentialsSet rootCredentialsSet); + default Map bootstrapRealms(BootstrapOptions bootstrapOptions) { + return bootstrapRealms(bootstrapOptions.realms(), bootstrapOptions.rootCredentialsSet()); + } + /** Purge all metadata for the realms provided */ Map purgeRealms(Iterable realms); } diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/bootstrap/BootstrapOptions.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/bootstrap/BootstrapOptions.java new file mode 100644 index 0000000000..a9bd907788 --- /dev/null +++ b/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 realms(); + + RootCredentialsSet rootCredentialsSet(); + + SchemaOptions schemaOptions(); +} diff --git a/polaris-core/src/main/java/org/apache/polaris/core/persistence/bootstrap/SchemaOptions.java b/polaris-core/src/main/java/org/apache/polaris/core/persistence/bootstrap/SchemaOptions.java new file mode 100644 index 0000000000..f0779cc8e1 --- /dev/null +++ b/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."); + } + } +} diff --git a/runtime/admin/src/main/java/org/apache/polaris/admintool/BootstrapCommand.java b/runtime/admin/src/main/java/org/apache/polaris/admintool/BootstrapCommand.java index 8ddb350954..2d75b4c17d 100644 --- a/runtime/admin/src/main/java/org/apache/polaris/admintool/BootstrapCommand.java +++ b/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,22 @@ static class FileInputOptions { description = "A file containing root principal credentials to bootstrap.") Path file; } + + static class SchemaInputOptions { + @CommandLine.Option( + names = {"-v", "--schema-version"}, + paramLabel = "", + description = "The version of the schema to load in [1, 2, LATEST].") + Integer schemaVersion; + + @CommandLine.Option( + names = {"--schema-file"}, + paramLabel = "", + description = + "A schema file to bootstrap from. If unset, the bundled files will be used.", + defaultValue = "") + String schemaFile; + } } @Override @@ -103,9 +126,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 results = - metaStoreManagerFactory.bootstrapRealms(realms, rootCredentialsSet); + metaStoreManagerFactory.bootstrapRealms(bootstrapOptions); // Log any errors: boolean success = true; diff --git a/runtime/admin/src/test/java/org/apache/polaris/admintool/BootstrapCommandTestBase.java b/runtime/admin/src/test/java/org/apache/polaris/admintool/BootstrapCommandTestBase.java index 6ebbe0167f..1bb2543052 100644 --- a/runtime/admin/src/test/java/org/apache/polaris/admintool/BootstrapCommandTestBase.java +++ b/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= [-r=]... [-c=]... [-p]) " - + "and -f= are mutually exclusive (specify only one)"); + "(-r= [-r=]... [-c=]... [-p]) and -f= " + + "and (-v= | [--schema-file=]) are mutually exclusive " + + "(specify only one)"); } @Test