test: Cover pstmt.execute() with DML and SELECT

Author: ennuite

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

@@ -20,6 +20,8 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.jupiter.api.Assertions.assertAll;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import java.nio.charset.StandardCharsets;
 import java.sql.Connection;
@@ -83,6 +85,19 @@ public void testSimpleQueryNoParameterBinding() throws SQLException {
     }
   }
 
+  @Test
+  public void testSimpleQueryNoParameterBindingWithExecute() throws SQLException {
+    final String query = CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD;
+    try (final PreparedStatement preparedStatement = connection.prepareStatement(query)) {
+      boolean isResultSet = preparedStatement.execute();
+      assertTrue(isResultSet);
+      final ResultSet resultSet = preparedStatement.getResultSet();
+      CoreMockedSqlProducers.assertLegacyRegularSqlResultSet(resultSet);
+      assertFalse(preparedStatement.getMoreResults());
+      assertEquals(-1, preparedStatement.getUpdateCount());
+    }
+  }
+
   @Test
   public void testQueryWithParameterBinding() throws SQLException {
     final String query = "Fake query with parameters";
@@ -174,6 +189,20 @@ public void testUpdateQuery() throws SQLException {
     }
   }
 
+  @Test
+  public void testUpdateQueryWithExecute() throws SQLException {
+    String query = "Fake update with execute";
+    PRODUCER.addUpdateQuery(query, /*updatedRows*/ 42);
+    try (final PreparedStatement stmt = connection.prepareStatement(query)) {
+      boolean isResultSet = stmt.execute();
+      assertFalse(isResultSet);
+      int updated = stmt.getUpdateCount();
+      assertEquals(42, updated);
+      assertFalse(stmt.getMoreResults());
+      assertEquals(-1, stmt.getUpdateCount());
+    }
+  }
+
   @Test
   public void testUpdateQueryWithParameters() throws SQLException {
     String query = "Fake update with parameters";