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);

Reply via email to