Tighten statement protocol regression coverage

Author: xborder

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

@@ -26,6 +26,7 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 
 import java.sql.Connection;
 import java.sql.ResultSet;
@@ -37,6 +38,7 @@
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 import org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer;
+import org.apache.arrow.flight.sql.FlightSqlUtils;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.RootAllocator;
 import org.apache.arrow.util.AutoCloseables;

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

@@ -174,6 +174,11 @@ public void testStatementExecuteUsesPreparedProtocolForQuery() throws SQLExcepti
             .getCommandTypeCounter()
             .getOrDefault(MockFlightSqlProducer.COMMAND_STATEMENT_QUERY, 0),
         is(0));
+    assertThat(
+        PRODUCER
+            .getActionTypeCounter()
+            .getOrDefault(FlightSqlUtils.FLIGHT_SQL_CLOSE_PREPARED_STATEMENT.getType(), 0),
+        is(1));
   }
 
   @Test
@@ -216,6 +221,16 @@ public void testStatementExecuteUsesPreparedProtocolForUpdate() throws SQLExcept
             .getCommandTypeCounter()
             .getOrDefault(MockFlightSqlProducer.COMMAND_PREPARED_STATEMENT_UPDATE, 0),
         is(1));
+    assertThat(
+        PRODUCER
+            .getCommandTypeCounter()
+            .getOrDefault(MockFlightSqlProducer.COMMAND_STATEMENT_UPDATE, 0),
+        is(0));
+    assertThat(
+        PRODUCER
+            .getActionTypeCounter()
+            .getOrDefault(FlightSqlUtils.FLIGHT_SQL_CLOSE_PREPARED_STATEMENT.getType(), 0),
+        is(1));
   }
 
   @Test
@@ -273,6 +288,60 @@ public void testStatementExecuteUpdateThenExecuteQueryUsesStatementProtocol()
         is(1));
   }
 
+  @Test
+  public void testStatementExecuteQueryThenExecuteUsesPreparedProtocol() throws SQLException {
+    try (Statement statement = connection.createStatement()) {
+      try (ResultSet resultSet = statement.executeQuery(SELECT_QUERY)) {
+        assertTrue(resultSet.next());
+      }
+      assertThat(statement.execute(SELECT_QUERY), is(true));
+      try (ResultSet resultSet = statement.getResultSet()) {
+        assertTrue(resultSet.next());
+      }
+    }
+
+    assertThat(
+        PRODUCER
+            .getCommandTypeCounter()
+            .getOrDefault(MockFlightSqlProducer.COMMAND_STATEMENT_QUERY, 0),
+        is(1));
+    assertThat(
+        PRODUCER
+            .getActionTypeCounter()
+            .getOrDefault(FlightSqlUtils.FLIGHT_SQL_CREATE_PREPARED_STATEMENT.getType(), 0),
+        is(1));
+    assertThat(
+        PRODUCER
+            .getCommandTypeCounter()
+            .getOrDefault(MockFlightSqlProducer.COMMAND_PREPARED_STATEMENT_QUERY, 0),
+        is(1));
+  }
+
+  @Test
+  public void testStatementExecuteUpdateThenExecuteUsesPreparedProtocol() throws SQLException {
+    try (Statement statement = connection.createStatement()) {
+      assertThat(statement.executeUpdate(UPDATE_QUERY), is(1));
+      assertThat(statement.execute(UPDATE_QUERY), is(false));
+      assertThat(statement.getUpdateCount(), is(1));
+    }
+
+    assertThat(
+        PRODUCER
+            .getCommandTypeCounter()
+            .getOrDefault(MockFlightSqlProducer.COMMAND_STATEMENT_UPDATE, 0),
+        is(1));
+    assertThat(
+        PRODUCER
+            .getActionTypeCounter()
+            .getOrDefault(FlightSqlUtils.FLIGHT_SQL_CREATE_PREPARED_STATEMENT.getType(), 0),
+        is(1));
+    assertThat(
+        PRODUCER
+            .getCommandTypeCounter()
+            .getOrDefault(MockFlightSqlProducer.COMMAND_PREPARED_STATEMENT_UPDATE, 0),
+        is(1));
+  }
+
   @Test
   public void testPreparedStatementExecuteQueryUsesPreparedProtocol() throws SQLException {
     try (PreparedStatement statement = connection.prepareStatement(SELECT_QUERY);
@@ -339,6 +408,11 @@ public void testPreparedStatementExecuteUpdateUsesPreparedProtocol() throws SQLE
             .getCommandTypeCounter()
             .getOrDefault(MockFlightSqlProducer.COMMAND_PREPARED_STATEMENT_UPDATE, 0),
         is(1));
+    assertThat(
+        PRODUCER
+            .getCommandTypeCounter()
+            .getOrDefault(MockFlightSqlProducer.COMMAND_STATEMENT_UPDATE, 0),
+        is(0));
   }
 
   @Test
@@ -358,6 +432,11 @@ public void testPreparedStatementExecuteUsesPreparedProtocolForUpdate() throws S
             .getCommandTypeCounter()
             .getOrDefault(MockFlightSqlProducer.COMMAND_PREPARED_STATEMENT_UPDATE, 0),
         is(1));
+    assertThat(
+        PRODUCER
+            .getCommandTypeCounter()
+            .getOrDefault(MockFlightSqlProducer.COMMAND_STATEMENT_UPDATE, 0),
+        is(0));
   }
 
   @Test

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

@@ -407,8 +407,14 @@ public Runnable acceptPutStatement(
       final FlightStream flightStream,
       final StreamListener<PutResult> streamListener) {
     incrementCommandTypeCounter(COMMAND_STATEMENT_UPDATE);
+    return acceptUpdate(commandStatementUpdate.getQuery(), flightStream, streamListener);
+  }
+
+  private Runnable acceptUpdate(
+      final String query,
+      final FlightStream flightStream,
+      final StreamListener<PutResult> streamListener) {
     return () -> {
-      final String query = commandStatementUpdate.getQuery();
       final BiConsumer<FlightStream, StreamListener<PutResult>> resultProvider =
           Preconditions.checkNotNull(
               updateResultProviders.get(query),
@@ -491,11 +497,7 @@ public Runnable acceptPutPreparedStatementUpdate(
       return () -> {};
     }
 
-    return acceptPutStatement(
-        CommandStatementUpdate.newBuilder().setQuery(query).build(),
-        callContext,
-        flightStream,
-        streamListener);
+    return acceptUpdate(query, flightStream, streamListener);
   }
 
   @Override