Fix remaining JDBC lifecycle validation gaps

Author: xborder

flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java

@@ -313,12 +313,16 @@ public PreparedStatement prepareStatement(
       final int resultSetHoldability)
       throws SQLException {
     checkOpen();
-    return ArrowFlightPreparedStatement.builder(this)
-        .withQuery(sql)
-        .withGeneratedHandle()
-        .withResultSetType(resultSetType)
-        .withResultSetConcurrency(resultSetConcurrency)
-        .withResultSetHoldability(resultSetHoldability)
-        .build();
+    try {
+      return ArrowFlightPreparedStatement.builder(this)
+          .withQuery(sql)
+          .withGeneratedHandle()
+          .withResultSetType(resultSetType)
+          .withResultSetConcurrency(resultSetConcurrency)
+          .withResultSetHoldability(resultSetHoldability)
+          .build();
+    } catch (final RuntimeException e) {
+      throw HELPER.createException("while preparing SQL:" + sql, e);
+    }
   }
 }

flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java

@@ -23,6 +23,7 @@
 import java.util.Collections;
 import java.util.List;
 import org.apache.arrow.driver.jdbc.utils.ConvertUtils;
+import org.apache.arrow.util.Preconditions;
 import org.apache.arrow.vector.types.pojo.Schema;
 import org.apache.calcite.avatica.AvaticaConnection;
 import org.apache.calcite.avatica.AvaticaParameter;
@@ -98,6 +99,7 @@ public Frame fetch(
   @Override
   public StatementHandle prepare(
       final ConnectionHandle connectionHandle, final String query, final long maxRowCount) {
+    validateConnectionId(connectionHandle.id);
     try {
       // This is the Avatica entry point used by Connection.prepareStatement(String).
       ArrowFlightPreparedStatement stmt =
@@ -192,6 +194,7 @@ void setDefaultConnectionProperties() {
   }
 
   private AvaticaStatement getStatement(final StatementHandle statementHandle) {
+    validateConnectionId(statementHandle.connectionId);
     final AvaticaStatement statement = connection.statementMap.get(statementHandle.id);
     if (statement instanceof ArrowFlightStatement
         || statement instanceof ArrowFlightPreparedStatement) {
@@ -209,6 +212,11 @@ private ArrowFlightPreparedStatement getPreparedStatement(final StatementHandle
         "Statement operation is not supported for handle: " + statementHandle);
   }
 
+  private void validateConnectionId(final String connectionId) {
+    Preconditions.checkArgument(
+        connection.id.equals(connectionId), "Connection IDs are not consistent");
+  }
+
   static Signature buildDefaultSignature() {
     return buildSignature(null, StatementType.SELECT);
   }

flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java

@@ -102,17 +102,32 @@ Schema getDataSetSchema() {
 
   @Override
   protected void close_() {
+    RuntimeException closeFailure = null;
     try {
       super.close_();
-    } finally {
-      ((ArrowFlightConnection) connection).unregisterStatementOwner(this);
+    } catch (final RuntimeException e) {
+      closeFailure = e;
+    }
+    try {
+      closePreparedResources();
+    } catch (final RuntimeException e) {
+      if (closeFailure == null) {
+        closeFailure = e;
+      } else {
+        closeFailure.addSuppressed(e);
+      }
+    }
+    ((ArrowFlightConnection) connection).unregisterStatementOwner(this);
+    if (closeFailure != null) {
+      throw closeFailure;
     }
   }
 
   void closePreparedResources() {
-    if (preparedStatement != null) {
-      preparedStatement.close();
-      preparedStatement = null;
+    final ArrowFlightSqlClientHandler.PreparedStatement statement = preparedStatement;
+    preparedStatement = null;
+    if (statement != null) {
+      statement.close();
     }
   }
 

flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java

@@ -26,6 +26,8 @@
 import static org.junit.jupiter.api.Assertions.assertSame;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
 
 import java.nio.charset.StandardCharsets;
 import java.sql.Connection;
@@ -52,6 +54,8 @@
 import org.apache.arrow.vector.types.pojo.FieldType;
 import org.apache.arrow.vector.types.pojo.Schema;
 import org.apache.arrow.vector.util.Text;
+import org.apache.calcite.avatica.AvaticaResultSet;
+import org.apache.calcite.avatica.AvaticaStatement;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.BeforeEach;
@@ -177,6 +181,82 @@ public void testPrepareFailureIsReportedAsSQLException() {
     assertThat(exception.getCause(), instanceOf(FlightRuntimeException.class));
   }
 
+  @Test
+  public void testNullQueryIsReportedAsSQLException() {
+    final SQLException exception =
+        assertThrows(SQLException.class, () -> connection.prepareStatement(null));
+
+    assertThat(exception.getCause(), instanceOf(NullPointerException.class));
+  }
+
+  @Test
+  public void testCloseReleasesPreparedResourcesWhenResultSetCloseFails() throws Exception {
+    try (final ArrowFlightConnection localConnection =
+        (ArrowFlightConnection) FLIGHT_SERVER_TEST_EXTENSION.getConnection(false)) {
+      final ArrowFlightPreparedStatement preparedStatement =
+          (ArrowFlightPreparedStatement)
+              localConnection.prepareStatement(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD);
+      final AvaticaResultSet resultSet = mock(AvaticaResultSet.class);
+      final RuntimeException resultSetCloseFailure =
+          new RuntimeException("Result set close failed");
+      doThrow(resultSetCloseFailure).when(resultSet).close();
+      final java.lang.reflect.Field openResultSetField =
+          AvaticaStatement.class.getDeclaredField("openResultSet");
+      openResultSetField.setAccessible(true);
+      openResultSetField.set(preparedStatement, resultSet);
+
+      final SQLException exception = assertThrows(SQLException.class, preparedStatement::close);
+
+      assertSame(resultSetCloseFailure, exception.getCause());
+      assertTrue(preparedStatement.isClosed());
+      assertEquals(
+          1,
+          PRODUCER
+              .getActionTypeCounter()
+              .getOrDefault(FlightSqlUtils.FLIGHT_SQL_CLOSE_PREPARED_STATEMENT.getType(), 0));
+    }
+  }
+
+  @Test
+  public void testMetaRejectsForeignStatementHandle() throws SQLException {
+    try (final ArrowFlightConnection firstConnection =
+            (ArrowFlightConnection) FLIGHT_SERVER_TEST_EXTENSION.getConnection(false);
+        final ArrowFlightConnection secondConnection =
+            (ArrowFlightConnection) FLIGHT_SERVER_TEST_EXTENSION.getConnection(false);
+        final PreparedStatement firstStatement =
+            firstConnection.prepareStatement(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD);
+        final PreparedStatement secondStatement =
+            secondConnection.prepareStatement(
+                CoreMockedSqlProducers.UUID_PREPARED_SELECT_SQL_CMD)) {
+      final ArrowFlightPreparedStatement firstPreparedStatement =
+          (ArrowFlightPreparedStatement) firstStatement;
+      final ArrowFlightPreparedStatement secondPreparedStatement =
+          (ArrowFlightPreparedStatement) secondStatement;
+      assertEquals(firstPreparedStatement.handle.id, secondPreparedStatement.handle.id);
+
+      assertThrows(
+          IllegalArgumentException.class,
+          () -> firstConnection.getMeta().closeStatement(secondPreparedStatement.handle));
+      assertFalse(firstPreparedStatement.isClosed());
+    }
+  }
+
+  @Test
+  public void testMetaRejectsForeignConnectionHandle() throws SQLException {
+    try (final ArrowFlightConnection firstConnection =
+            (ArrowFlightConnection) FLIGHT_SERVER_TEST_EXTENSION.getConnection(false);
+        final ArrowFlightConnection secondConnection =
+            (ArrowFlightConnection) FLIGHT_SERVER_TEST_EXTENSION.getConnection(false)) {
+      assertThrows(
+          IllegalArgumentException.class,
+          () ->
+              firstConnection
+                  .getMeta()
+                  .prepare(
+                      secondConnection.handle, CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD, -1));
+    }
+  }
+
   @Test
   public void testQueryWithParameterBinding() throws SQLException {
     final String query = "Fake query with parameters";