(polaris) branch main updated: Follow-up to #4795: use transaction connection for SELECT in inheritable policy check (#4810)

[di...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

dimas-b 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 f89897481b Follow-up to #4795: use transaction connection for SELECT in inheritable policy check (#4810)
f89897481b is described below

commit f89897481ba9d3df956e2fb6a6614f63d3647a74
Author: vignesh_A <16...@users.noreply.github.com>
AuthorDate: Thu Jun 25 21:28:02 2026 +0530

    Follow-up to #4795: use transaction connection for SELECT in inheritable policy check (#4810)
    
    * Make policy mapping SELECT use transaction connection for inheritable policies (follow-up to #4795)
    
    - Added connection-aware executeSelect / executeSelectOverStream in DatasourceOperations
    - Added private fetch + public overload of loadPoliciesOnTargetByType that accepts Connection
    - Updated handleInheritablePolicy to pass the tx connection for the SELECT
    
    This addresses the review comment that the SELECT should also participate in the same transaction as the insert/update.
    
    * Address review: add @NonNull to Connection parameter in connection-aware executeSelect* and fetch methods
    
    * Extract private helper for SELECT, make overload package-private, add unit tests
    
    - Extract duplicated prepare/bind/execute logic into private executeSelectOverStreamWithConnection
    - Change connection-aware loadPoliciesOnTargetByType to package-private
    - Add unit tests for the new Connection-aware SELECT paths
    
    Follow-up to #4795
---
 .../relational/jdbc/DatasourceOperations.java      | 64 ++++++++++++++++++----
 .../relational/jdbc/JdbcBasePersistenceImpl.java   | 44 ++++++++++++++-
 .../relational/jdbc/DatasourceOperationsTest.java  | 32 +++++++++++
 3 files changed, 127 insertions(+), 13 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 cea7f58b95..372ccfe1ca 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
@@ -173,22 +173,62 @@ public class DatasourceOperations {
       throws SQLException {
     withRetries(
         () -> {
-          logQuery(query);
-          try (Connection connection = borrowConnection();
-              PreparedStatement statement = connection.prepareStatement(query.sql())) {
-            List<Object> params = query.parameters();
-            for (int i = 0; i < params.size(); i++) {
-              statement.setObject(i + 1, params.get(i));
-            }
-            try (ResultSet resultSet = statement.executeQuery()) {
-              ResultSetIterator<T> iterator = new ResultSetIterator<>(resultSet, converterInstance);
-              consumer.accept(iterator.toStream());
-              return null;
-            }
+          try (Connection connection = borrowConnection()) {
+            executeSelectOverStreamWithConnection(query, converterInstance, consumer, connection);
+            return null;
           }
         });
   }
 
+  /** Connection-aware version for use inside runWithinTransaction. */
+  public <T> void executeSelectOverStream(
+      @NonNull Connection connection,
+      @NonNull PreparedQuery query,
+      @NonNull Converter<T> converterInstance,
+      @NonNull Consumer<Stream<T>> consumer)
+      throws SQLException {
+    withRetries(
+        () -> {
+          executeSelectOverStreamWithConnection(query, converterInstance, consumer, connection);
+          return null;
+        });
+  }
+
+  /**
+   * Internal implementation that executes the SELECT on the provided connection. Does not manage
+   * connection lifecycle or retries.
+   */
+  private <T> void executeSelectOverStreamWithConnection(
+      @NonNull PreparedQuery query,
+      @NonNull Converter<T> converterInstance,
+      @NonNull Consumer<Stream<T>> consumer,
+      @NonNull Connection connection)
+      throws SQLException {
+    logQuery(query);
+    try (PreparedStatement statement = connection.prepareStatement(query.sql())) {
+      List<Object> params = query.parameters();
+      for (int i = 0; i < params.size(); i++) {
+        statement.setObject(i + 1, params.get(i));
+      }
+      try (ResultSet resultSet = statement.executeQuery()) {
+        ResultSetIterator<T> iterator = new ResultSetIterator<>(resultSet, converterInstance);
+        consumer.accept(iterator.toStream());
+      }
+    }
+  }
+
+  /** Connection-aware version for use inside runWithinTransaction. */
+  public <T> List<T> executeSelect(
+      @NonNull Connection connection,
+      @NonNull PreparedQuery query,
+      @NonNull Converter<T> converterInstance)
+      throws SQLException {
+    ArrayList<T> results = new ArrayList<>();
+    executeSelectOverStream(
+        connection, query, converterInstance, stream -> stream.forEach(results::add));
+    return results;
+  }
+
   /**
    * Executes the UPDATE or INSERT Query
    *
diff --git a/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcBasePersistenceImpl.java b/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcBasePersistenceImpl.java
index b9bcc45141..fcc59c8e67 100644
--- a/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcBasePersistenceImpl.java
+++ b/persistence/relational-jdbc/src/main/java/org/apache/polaris/persistence/relational/jdbc/JdbcBasePersistenceImpl.java
@@ -1084,7 +1084,11 @@ public class JdbcBasePersistenceImpl
       throws SQLException {
     List<PolarisPolicyMappingRecord> existingRecords =
         loadPoliciesOnTargetByType(
-            callCtx, record.getTargetCatalogId(), record.getTargetId(), record.getPolicyTypeCode());
+            callCtx,
+            record.getTargetCatalogId(),
+            record.getTargetId(),
+            record.getPolicyTypeCode(),
+            connection);
     if (existingRecords.size() > 1) {
       throw new PolicyMappingAlreadyExistsException(existingRecords.getFirst());
     } else if (existingRecords.size() == 1) {
@@ -1231,6 +1235,44 @@ public class JdbcBasePersistenceImpl
             ModelPolicyMappingRecord.ALL_COLUMNS, ModelPolicyMappingRecord.TABLE_NAME, params));
   }
 
+  /**
+   * Connection-aware version for use inside runWithinTransaction (for inheritable policy check).
+   */
+  private List<PolarisPolicyMappingRecord> fetchPolicyMappingRecords(
+      QueryGenerator.PreparedQuery query, @NonNull Connection connection) {
+    try {
+      var results =
+          datasourceOperations.executeSelect(connection, query, new ModelPolicyMappingRecord());
+      return results == null ? Collections.emptyList() : results;
+    } catch (SQLException e) {
+      throw new RuntimeException(
+          String.format("Failed to retrieve policy mapping records %s", e.getMessage()), e);
+    }
+  }
+
+  /** Connection-aware overload for use inside transaction. */
+  @NonNull List<PolarisPolicyMappingRecord> loadPoliciesOnTargetByType(
+      @NonNull PolarisCallContext callCtx,
+      long targetCatalogId,
+      long targetId,
+      int policyTypeCode,
+      @NonNull Connection connection) {
+    Map<String, Object> params =
+        Map.of(
+            "target_catalog_id",
+            targetCatalogId,
+            "target_id",
+            targetId,
+            "policy_type_code",
+            policyTypeCode,
+            "realm_id",
+            realmId);
+    return fetchPolicyMappingRecords(
+        QueryGenerator.generateSelectQuery(
+            ModelPolicyMappingRecord.ALL_COLUMNS, ModelPolicyMappingRecord.TABLE_NAME, params),
+        connection);
+  }
+
   @NonNull
   @Override
   public List<PolarisPolicyMappingRecord> loadAllPoliciesOnTarget(
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 200579cabe..29bbd6c49b 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
@@ -144,6 +144,38 @@ public class DatasourceOperationsTest {
         SQLException.class, () -> datasourceOperations.executeSelect(query, new ModelEntity(1)));
   }
 
+  @Test
+  void testExecuteSelect_withConnection_usesProvidedConnection() throws Exception {
+    QueryGenerator.PreparedQuery query =
+        new QueryGenerator.PreparedQuery("SELECT * FROM users", List.of());
+    when(mockConnection.prepareStatement(query.sql())).thenReturn(mockPreparedStatement);
+    when(mockPreparedStatement.executeQuery()).thenThrow(new SQLException("demo", "42P07"));
+
+    // The connection-aware overload must use the caller-provided Connection
+    // (no internal borrow should occur for the SELECT itself).
+    assertThrows(
+        SQLException.class,
+        () -> datasourceOperations.executeSelect(mockConnection, query, new ModelEntity(1)));
+
+    verify(mockConnection).prepareStatement(query.sql());
+  }
+
+  @Test
+  void testExecuteSelectOverStream_withConnection_usesProvidedConnection() throws Exception {
+    QueryGenerator.PreparedQuery query =
+        new QueryGenerator.PreparedQuery("SELECT * FROM users", List.of());
+    when(mockConnection.prepareStatement(query.sql())).thenReturn(mockPreparedStatement);
+    when(mockPreparedStatement.executeQuery()).thenThrow(new SQLException("stream demo", "42P07"));
+
+    assertThrows(
+        SQLException.class,
+        () ->
+            datasourceOperations.executeSelectOverStream(
+                mockConnection, query, new ModelEntity(1), stream -> stream.forEach(x -> {})));
+
+    verify(mockConnection).prepareStatement(query.sql());
+  }
+
   @Test
   void testRunWithinTransaction_commit() throws Exception {
     // reset to ignore constructor interaction