This is an automated email from the ASF dual-hosted git repository.
singhpk234 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/polaris.git
The following commit(s) were added to refs/heads/main by this push:
new 2a2bcde1f JDBC: Refactor DatabaseOps (#1843)
2a2bcde1f is described below
commit 2a2bcde1ff335ef7281ab23d46a6a6a5e8493d82
Author: Prashant Singh <[email protected]>
AuthorDate: Thu Jun 12 15:54:42 2025 -0700
JDBC: Refactor DatabaseOps (#1843)
* removes the databaseType computation from JDBCMetastoreManagerFactory to
DbOperations
* wraps the bootstrap in a transaction !
* refactor Production Readiness checks for Postgres
---
.../relational/jdbc/DatasourceOperations.java | 74 +++++++++++-----------
.../jdbc/JdbcMetaStoreManagerFactory.java | 29 +++------
.../RelationalJdbcProductionReadinessChecks.java | 13 ++--
...toreManagerWithJdbcBasePersistenceImplTest.java | 5 +-
.../relational/jdbc/DatasourceOperationsTest.java | 14 +++-
5 files changed, 71 insertions(+), 64 deletions(-)
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 84168727c..7017f1919 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
@@ -60,59 +60,61 @@ public class DatasourceOperations {
private final Random random = new Random();
public DatasourceOperations(
- DataSource datasource,
- DatabaseType databaseType,
- RelationalJdbcConfiguration relationalJdbcConfiguration) {
+ DataSource datasource, RelationalJdbcConfiguration
relationalJdbcConfiguration)
+ throws SQLException {
this.datasource = datasource;
- this.databaseType = databaseType;
this.relationalJdbcConfiguration = relationalJdbcConfiguration;
+ try (Connection connection = this.datasource.getConnection()) {
+ String productName = connection.getMetaData().getDatabaseProductName();
+ this.databaseType = DatabaseType.fromDisplayName(productName);
+ }
}
- public DatabaseType getDatabaseType() {
+ DatabaseType getDatabaseType() {
return databaseType;
}
/**
- * Execute SQL script
+ * Execute SQL script.
*
* @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();
- try (Connection connection = borrowConnection();
- Statement statement = connection.createStatement()) {
- boolean autoCommit = connection.getAutoCommit();
- connection.setAutoCommit(true);
- try {
- 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 {
- statement.executeUpdate(sql);
- } catch (SQLException e) {
- throw new RuntimeException(e);
+ 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);
+ }
+ 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);
}
- }
- } finally {
- connection.setAutoCommit(autoCommit);
- }
- } catch (IOException e) {
- throw new RuntimeException(e);
- }
+ });
}
/**
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 4b369953d..118fcadd4 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
@@ -23,7 +23,6 @@ import jakarta.annotation.Nullable;
import jakarta.enterprise.context.ApplicationScoped;
import jakarta.enterprise.inject.Instance;
import jakarta.inject.Inject;
-import java.sql.Connection;
import java.sql.SQLException;
import java.util.HashMap;
import java.util.Map;
@@ -95,13 +94,7 @@ public class JdbcMetaStoreManagerFactory implements
MetaStoreManagerFactory {
private void initializeForRealm(
RealmContext realmContext, RootCredentialsSet rootCredentialsSet,
boolean isBootstrap) {
- DatabaseType databaseType;
- try {
- databaseType = getDatabaseType();
- } catch (SQLException e) {
- throw new RuntimeException(e);
- }
- DatasourceOperations databaseOperations =
getDatasourceOperations(isBootstrap, databaseType);
+ DatasourceOperations databaseOperations =
getDatasourceOperations(isBootstrap);
sessionSupplierMap.put(
realmContext.getRealmIdentifier(),
() ->
@@ -115,20 +108,18 @@ public class JdbcMetaStoreManagerFactory implements
MetaStoreManagerFactory {
metaStoreManagerMap.put(realmContext.getRealmIdentifier(),
metaStoreManager);
}
- protected DatabaseType getDatabaseType() throws SQLException {
- try (Connection connection = dataSource.get().getConnection()) {
- String productName = connection.getMetaData().getDatabaseProductName();
- return DatabaseType.fromDisplayName(productName);
+ private DatasourceOperations getDatasourceOperations(boolean isBootstrap) {
+ DatasourceOperations databaseOperations;
+ try {
+ databaseOperations = new DatasourceOperations(dataSource.get(),
relationalJdbcConfiguration);
+ } catch (SQLException sqlException) {
+ throw new RuntimeException(sqlException);
}
- }
-
- private DatasourceOperations getDatasourceOperations(
- boolean isBootstrap, DatabaseType databaseType) {
- DatasourceOperations databaseOperations =
- new DatasourceOperations(dataSource.get(), databaseType,
relationalJdbcConfiguration);
if (isBootstrap) {
try {
- databaseOperations.executeScript(databaseType.getInitScriptResource());
+ // 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);
diff --git
a/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/RelationalJdbcProductionReadinessChecks.java
b/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/RelationalJdbcProductionReadinessChecks.java
index 51d209abd..d365b9e7f 100644
---
a/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/RelationalJdbcProductionReadinessChecks.java
+++
b/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/RelationalJdbcProductionReadinessChecks.java
@@ -20,8 +20,10 @@
package org.apache.polaris.persistence.relational.jdbc;
import jakarta.enterprise.context.ApplicationScoped;
+import jakarta.enterprise.inject.Instance;
import jakarta.enterprise.inject.Produces;
import java.sql.SQLException;
+import javax.sql.DataSource;
import org.apache.polaris.core.config.ProductionReadinessCheck;
import org.apache.polaris.core.persistence.MetaStoreManagerFactory;
@@ -29,15 +31,18 @@ import
org.apache.polaris.core.persistence.MetaStoreManagerFactory;
public class RelationalJdbcProductionReadinessChecks {
@Produces
public ProductionReadinessCheck checkRelationalJdbc(
- MetaStoreManagerFactory metaStoreManagerFactory) {
+ MetaStoreManagerFactory metaStoreManagerFactory,
+ Instance<DataSource> dataSource,
+ RelationalJdbcConfiguration relationalJdbcConfiguration) {
// This check should only be applicable when persistence uses
RelationalJdbc.
- if (!(metaStoreManagerFactory
- instanceof JdbcMetaStoreManagerFactory jdbcMetaStoreManagerFactory)) {
+ if (!(metaStoreManagerFactory instanceof JdbcMetaStoreManagerFactory)) {
return ProductionReadinessCheck.OK;
}
try {
- if
(jdbcMetaStoreManagerFactory.getDatabaseType().equals(DatabaseType.H2)) {
+ DatasourceOperations datasourceOperations =
+ new DatasourceOperations(dataSource.get(),
relationalJdbcConfiguration);
+ if (datasourceOperations.getDatabaseType().equals(DatabaseType.H2)) {
return ProductionReadinessCheck.of(
ProductionReadinessCheck.Error.of(
"The current persistence (jdbc:h2) is intended for tests
only.",
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 3018dc0e0..fbcd3a958 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
@@ -45,9 +45,10 @@ public class
AtomicMetastoreManagerWithJdbcBasePersistenceImplTest
@Override
protected PolarisTestMetaStoreManager createPolarisTestMetaStoreManager() {
PolarisDiagnostics diagServices = new PolarisDefaultDiagServiceImpl();
- DatasourceOperations datasourceOperations =
- new DatasourceOperations(createH2DataSource(), DatabaseType.H2, new
H2JdbcConfiguration());
+ DatasourceOperations datasourceOperations;
try {
+ datasourceOperations =
+ new DatasourceOperations(createH2DataSource(), new
H2JdbcConfiguration());
datasourceOperations.executeScript(
String.format("%s/schema-v1.sql", DatabaseType.H2.getDisplayName()));
} catch (SQLException e) {
diff --git
a/persistence/relational-jdbc/src/test/java/org/apache/polaris/persistence/relational/jdbc/DatasourceOperationsTest.java
b/persistence/relational-jdbc/src/test/java/org/apache/polaris/persistence/relational/jdbc/DatasourceOperationsTest.java
index bde721c3f..ae6f569cc 100644
---
a/persistence/relational-jdbc/src/test/java/org/apache/polaris/persistence/relational/jdbc/DatasourceOperationsTest.java
+++
b/persistence/relational-jdbc/src/test/java/org/apache/polaris/persistence/relational/jdbc/DatasourceOperationsTest.java
@@ -23,11 +23,13 @@ import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.Mockito.atLeast;
import static org.mockito.Mockito.atMost;
+import static org.mockito.Mockito.reset;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
import java.sql.Connection;
+import java.sql.DatabaseMetaData;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.time.Instant;
@@ -52,14 +54,18 @@ public class DatasourceOperationsTest {
@Mock private RelationalJdbcConfiguration relationalJdbcConfiguration;
+ @Mock private DatabaseMetaData mockDatabaseMetaData;
+
@Mock Operation<String> mockOperation;
private DatasourceOperations datasourceOperations;
@BeforeEach
- void setUp() {
- datasourceOperations =
- new DatasourceOperations(mockDataSource, DatabaseType.H2,
relationalJdbcConfiguration);
+ void setUp() throws SQLException {
+ when(mockDataSource.getConnection()).thenReturn(mockConnection);
+ when(mockConnection.getMetaData()).thenReturn(mockDatabaseMetaData);
+ when(mockDatabaseMetaData.getDatabaseProductName()).thenReturn("h2");
+ datasourceOperations = new DatasourceOperations(mockDataSource,
relationalJdbcConfiguration);
}
@Test
@@ -100,6 +106,8 @@ public class DatasourceOperationsTest {
@Test
void testRunWithinTransaction_commit() throws Exception {
+ // reset to ignore constructor interaction
+ reset(mockConnection);
when(mockDataSource.getConnection()).thenReturn(mockConnection);
DatasourceOperations.TransactionCallback callback = connection -> true;
when(mockConnection.getAutoCommit()).thenReturn(true);