GH-1063: Add is_update field to ActionCreatePreparedStatementResult (#1064)

## What's Changed

A new field, `optional bool is_update = 4;`, was added to `message
ActionCreatePreparedStatementResult`. When this field is sent by the
server, its value indicates whether the proper network flow to execute
the query that the driver should follow uses
`CommandPreparedStatementQuery` or `CommandPreparedStatementUpdate`.


For outdated servers that don't send the field, the driver maintains its
current behavior of using `CommandPreparedStatementQuery` when the
`dataset_schema` is not empty, thus ensuring the backward compatibility
of the new driver with old servers.

This change was created with AI assistance (Augment Code and Claude
code). All lines were manually reviewed by a human. The output is not
copyrightable subject matter.

- Closes #1063

---------

Co-authored-by: David Li <li.davidm96@gmail.com>

Author: ennuite

arrow-format/FlightSql.proto

@@ -1550,6 +1550,11 @@ message ActionCreatePreparedStatementResult {
   // If the query provided contained parameters, parameter_schema contains the
   // schema of the expected parameters.  It should be an IPC-encapsulated Schema, as described in Schema.fbs.
   bytes parameter_schema = 3;
+
+  // When set to true, the query should be executed with CommandPreparedStatementUpdate,
+  // when set to false, the query should be executed with CommandPreparedStatementQuery.
+  // If not set, the client can choose how to execute the query.
+  optional bool is_update = 4;
 }
 
 /*

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

@@ -106,7 +106,7 @@ static ArrowFlightJdbcFlightStreamResultSet fromFlightInfo(
     final TimeZone timeZone = TimeZone.getDefault();
     final QueryState state = new QueryState();
 
-    final Meta.Signature signature = ArrowFlightMetaImpl.newSignature(null, null, null);
+    final Meta.Signature signature = ArrowFlightMetaImpl.newSignature(null, null, null, null);
 
     final AvaticaResultSetMetaData resultSetMetaData =
         new AvaticaResultSetMetaData(null, null, signature);

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

@@ -73,7 +73,7 @@ public static ArrowFlightJdbcVectorSchemaRootResultSet fromVectorSchemaRoot(
     final TimeZone timeZone = TimeZone.getDefault();
     final QueryState state = new QueryState();
 
-    final Meta.Signature signature = ArrowFlightMetaImpl.newSignature(null, null, null);
+    final Meta.Signature signature = ArrowFlightMetaImpl.newSignature(null, null, null, null);
 
     final AvaticaResultSetMetaData resultSetMetaData =
         new AvaticaResultSetMetaData(null, null, signature);

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

@@ -53,7 +53,8 @@ public ArrowFlightMetaImpl(final AvaticaConnection connection) {
   }
 
   /** Construct a signature. */
-  static Signature newSignature(final String sql, Schema resultSetSchema, Schema parameterSchema) {
+  static Signature newSignature(
+      final String sql, Schema resultSetSchema, Schema parameterSchema, Boolean isUpdate) {
     List<ColumnMetaData> columnMetaData =
         resultSetSchema == null
             ? new ArrayList<>()
@@ -62,10 +63,17 @@ static Signature newSignature(final String sql, Schema resultSetSchema, Schema p
         parameterSchema == null
             ? new ArrayList<>()
             : ConvertUtils.convertArrowFieldsToAvaticaParameters(parameterSchema.getFields());
-    StatementType statementType =
-        resultSetSchema == null || resultSetSchema.getFields().isEmpty()
-            ? StatementType.IS_DML
-            : StatementType.SELECT;
+    // If the server provided the is_update field, use it to determine the statement type
+    StatementType statementType;
+    if (isUpdate != null) {
+      statementType = isUpdate ? StatementType.IS_DML : StatementType.SELECT;
+    } else {
+      // Fall back to the legacy logic: check if the result set schema is empty
+      statementType =
+          resultSetSchema == null || resultSetSchema.getFields().isEmpty()
+              ? StatementType.IS_DML
+              : StatementType.SELECT;
+    }
     return new Signature(
         columnMetaData,
         sql,
@@ -178,7 +186,10 @@ private PreparedStatement prepareForHandle(final String query, StatementHandle h
         ((ArrowFlightConnection) connection).getClientHandler().prepare(query);
     handle.signature =
         newSignature(
-            query, preparedStatement.getDataSetSchema(), preparedStatement.getParameterSchema());
+            query,
+            preparedStatement.getDataSetSchema(),
+            preparedStatement.getParameterSchema(),
+            preparedStatement.isUpdate());
     statementHandlePreparedStatementMap.put(new StatementHandleKey(handle), preparedStatement);
     return preparedStatement;
   }

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

@@ -388,6 +388,14 @@ public interface PreparedStatement extends AutoCloseable {
      */
     Schema getParameterSchema();
 
+    /**
+     * Gets whether this {@link PreparedStatement} is an update statement.
+     *
+     * @return {@code true} if this is an update statement, {@code false} if it's a query, or {@code
+     *     null} if the server did not provide this information.
+     */
+    @Nullable Boolean isUpdate();
+
     void setParameters(VectorSchemaRoot parameters);
 
     @Override
@@ -456,6 +464,12 @@ public long executeUpdate() {
 
       @Override
       public StatementType getType() {
+        // If the server provided the is_update field, use it to determine the statement type
+        final Boolean isUpdate = preparedStatement.isUpdate();
+        if (isUpdate != null) {
+          return isUpdate ? StatementType.UPDATE : StatementType.SELECT;
+        }
+        // Fall back to the legacy logic: check if the result set schema is empty
         final Schema schema = preparedStatement.getResultSetSchema();
         return schema.getFields().isEmpty() ? StatementType.UPDATE : StatementType.SELECT;
       }
@@ -475,6 +489,11 @@ public void setParameters(VectorSchemaRoot parameters) {
         preparedStatement.setParameters(parameters);
       }
 
+      @Override
+      public Boolean isUpdate() {
+        return preparedStatement.isUpdate();
+      }
+
       @Override
       public void close() {
         try {

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

@@ -98,6 +98,40 @@ public void testSimpleQueryNoParameterBindingWithExecute() throws SQLException {
     }
   }
 
+  @Test
+  public void testSimpleQueryNoParameterBindingWithExecuteV2() throws SQLException {
+    final String query = "SELECT * FROM TEST_V2";
+    final Schema schema =
+        new Schema(Collections.singletonList(Field.nullable("", Types.MinorType.INT.getType())));
+    PRODUCER.addSelectQuery(
+        query,
+        schema,
+        Collections.singletonList(
+            listener -> {
+              try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
+                  final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
+                root.allocateNew();
+                ((IntVector) root.getVector(0)).setSafe(0, 123);
+                root.setRowCount(1);
+                listener.start(root);
+                listener.putNext();
+              } finally {
+                listener.completed();
+              }
+            }),
+        false);
+    try (final PreparedStatement preparedStatement = connection.prepareStatement(query)) {
+      boolean isResultSet = preparedStatement.execute();
+      assertTrue(isResultSet);
+      final ResultSet resultSet = preparedStatement.getResultSet();
+      assertTrue(resultSet.next());
+      assertEquals(123, resultSet.getInt(1));
+      assertFalse(resultSet.next());
+      assertFalse(preparedStatement.getMoreResults());
+      assertEquals(-1, preparedStatement.getUpdateCount());
+    }
+  }
+
   @Test
   public void testQueryWithParameterBinding() throws SQLException {
     final String query = "Fake query with parameters";
@@ -203,6 +237,20 @@ public void testUpdateQueryWithExecute() throws SQLException {
     }
   }
 
+  @Test
+  public void testUpdateQueryWithExecuteV2() throws SQLException {
+    String query = "Fake update with execute V2";
+    PRODUCER.addUpdateQuery(query, /*updatedRows*/ 99, true);
+    try (final PreparedStatement stmt = connection.prepareStatement(query)) {
+      boolean isResultSet = stmt.execute();
+      assertFalse(isResultSet);
+      int updated = stmt.getUpdateCount();
+      assertEquals(99, updated);
+      assertFalse(stmt.getMoreResults());
+      assertEquals(-1, stmt.getUpdateCount());
+    }
+  }
+
   @Test
   public void testUpdateQueryWithParameters() throws SQLException {
     String query = "Fake update with parameters";

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

@@ -62,6 +62,9 @@ public class ArrowFlightStatementExecuteTest {
   private static final String SAMPLE_LARGE_UPDATE_QUERY =
       "UPDATE this_large_table SET this_large_field = that_large_field FROM this_large_test WHERE this_large_condition";
   private static final long SAMPLE_LARGE_UPDATE_COUNT = Long.MAX_VALUE;
+  private static final String SAMPLE_QUERY_CMD_V2 = "SELECT * FROM this_test_v2";
+  private static final String SAMPLE_LARGE_UPDATE_QUERY_V2 =
+      "UPDATE this_large_table_v2 SET this_large_field = that_large_field FROM this_large_test WHERE this_large_condition";
   private static final MockFlightSqlProducer PRODUCER = new MockFlightSqlProducer();
 
   @RegisterExtension
@@ -96,6 +99,31 @@ public static void setUpBeforeClass() {
             }));
     PRODUCER.addUpdateQuery(SAMPLE_UPDATE_QUERY, SAMPLE_UPDATE_COUNT);
     PRODUCER.addUpdateQuery(SAMPLE_LARGE_UPDATE_QUERY, SAMPLE_LARGE_UPDATE_COUNT);
+
+    // V2 queries with is_update field set
+    PRODUCER.addSelectQuery(
+        SAMPLE_QUERY_CMD_V2,
+        SAMPLE_QUERY_SCHEMA,
+        Collections.singletonList(
+            listener -> {
+              try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
+                  final VectorSchemaRoot root =
+                      VectorSchemaRoot.create(SAMPLE_QUERY_SCHEMA, allocator)) {
+                final UInt1Vector vector = (UInt1Vector) root.getVector(VECTOR_NAME);
+                IntStream.range(0, SAMPLE_QUERY_ROWS)
+                    .forEach(index -> vector.setSafe(index, index));
+                vector.setValueCount(SAMPLE_QUERY_ROWS);
+                root.setRowCount(SAMPLE_QUERY_ROWS);
+                listener.start(root);
+                listener.putNext();
+              } catch (final Throwable throwable) {
+                listener.error(throwable);
+              } finally {
+                listener.completed();
+              }
+            }),
+        false);
+    PRODUCER.addUpdateQuery(SAMPLE_LARGE_UPDATE_QUERY_V2, SAMPLE_LARGE_UPDATE_COUNT, true);
   }
 
   @BeforeEach
@@ -168,4 +196,42 @@ public void testUpdateCountShouldStartOnZero() throws SQLException {
         is(allOf(equalTo(statement.getLargeUpdateCount()), equalTo(0L))));
     assertThat(statement.getResultSet(), is(nullValue()));
   }
+
+  @Test
+  public void testExecuteShouldRunSelectQueryV2() throws SQLException {
+    assertThat(statement.execute(SAMPLE_QUERY_CMD_V2), is(true));
+    final Set<Byte> numbers =
+        IntStream.range(0, SAMPLE_QUERY_ROWS)
+            .boxed()
+            .map(Integer::byteValue)
+            .collect(Collectors.toCollection(HashSet::new));
+    try (final ResultSet resultSet = statement.getResultSet()) {
+      final int columnCount = resultSet.getMetaData().getColumnCount();
+      assertThat(columnCount, is(1));
+      int rowCount = 0;
+      for (; resultSet.next(); rowCount++) {
+        assertThat(numbers.remove(resultSet.getByte(1)), is(true));
+      }
+      assertThat(rowCount, is(equalTo(SAMPLE_QUERY_ROWS)));
+    }
+    assertThat(numbers, is(Collections.emptySet()));
+    assertThat(
+        (long) statement.getUpdateCount(),
+        is(allOf(equalTo(statement.getLargeUpdateCount()), equalTo(-1L))));
+  }
+
+  @Test
+  public void testExecuteShouldRunUpdateQueryForLargeUpdateV2() throws SQLException {
+    assertThat(statement.execute(SAMPLE_LARGE_UPDATE_QUERY_V2), is(false)); // UPDATE query.
+    final long updateCountSmall = statement.getUpdateCount();
+    final long updateCountLarge = statement.getLargeUpdateCount();
+    assertThat(updateCountLarge, is(equalTo(SAMPLE_LARGE_UPDATE_COUNT)));
+    assertThat(
+        updateCountSmall,
+        is(
+            allOf(
+                equalTo((long) AvaticaUtils.toSaturatedInt(updateCountLarge)),
+                not(equalTo(updateCountLarge)))));
+    assertThat(statement.getResultSet(), is(nullValue()));
+  }
 }

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

@@ -87,6 +87,7 @@
 import org.apache.arrow.vector.types.pojo.Schema;
 import org.apache.arrow.vector.util.JsonStringArrayList;
 import org.apache.calcite.avatica.Meta.StatementType;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 /** An ad-hoc {@link FlightSqlProducer} for tests. */
 public final class MockFlightSqlProducer implements FlightSqlProducer {
@@ -101,6 +102,7 @@ public final class MockFlightSqlProducer implements FlightSqlProducer {
   private final SqlInfoBuilder sqlInfoBuilder = new SqlInfoBuilder();
   private final Map<String, Schema> parameterSchemas = new HashMap<>();
   private final Map<String, List<List<Object>>> expectedParameterValues = new HashMap<>();
+  private final Map<String, Boolean> isUpdateMap = new HashMap<>();
 
   private final Map<String, Integer> actionTypeCounter = new HashMap<>();
 
@@ -176,6 +178,40 @@ public void addUpdateQuery(final String sqlCommand, final long updatedRows) {
         });
   }
 
+  /**
+   * Registers a new {@link StatementType#SELECT} SQL query, optionally setting the is_update field.
+   *
+   * @param sqlCommand the SQL command under which to register the new query.
+   * @param schema the schema to use for the query result.
+   * @param resultProviders the result provider for this query.
+   * @param isUpdate value to report for the is_update field, or {@code null} to leave it unset.
+   */
+  public void addSelectQuery(
+      final String sqlCommand,
+      final Schema schema,
+      final List<Consumer<ServerStreamListener>> resultProviders,
+      final @Nullable Boolean isUpdate) {
+    addSelectQuery(sqlCommand, schema, resultProviders);
+    if (isUpdate != null) {
+      isUpdateMap.put(sqlCommand, isUpdate);
+    }
+  }
+
+  /**
+   * Registers a new {@link StatementType#UPDATE} SQL query, optionally setting the is_update field.
+   *
+   * @param sqlCommand the SQL command.
+   * @param updatedRows the number of rows affected.
+   * @param isUpdate value to report for the is_update field, or {@code null} to leave it unset.
+   */
+  public void addUpdateQuery(
+      final String sqlCommand, final long updatedRows, final @Nullable Boolean isUpdate) {
+    addUpdateQuery(sqlCommand, updatedRows);
+    if (isUpdate != null) {
+      isUpdateMap.put(sqlCommand, isUpdate);
+    }
+  }
+
   /**
    * Adds a catalog query to the results.
    *
@@ -247,6 +283,12 @@ public void createPreparedStatement(
         resultBuilder.setParameterSchema(ByteString.copyFrom(outputStream.toByteArray()));
       }
 
+      // Set is_update field if present
+      final Boolean isUpdate = isUpdateMap.get(query);
+      if (isUpdate != null) {
+        resultBuilder.setIsUpdate(isUpdate);
+      }
+
       listener.onNext(new Result(pack(resultBuilder.build()).toByteArray()));
     } catch (final Throwable t) {
       listener.onError(t);

flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java

@@ -1284,6 +1284,19 @@ public Schema getParameterSchema() {
       return parameterSchema;
     }
 
+    /**
+     * Returns whether the server indicated this prepared statement is an update query.
+     *
+     * @return true if the server indicated this is an update query, false if the server indicated
+     *     this is a select query, or null if the server did not provide this information.
+     */
+    public Boolean isUpdate() {
+      if (preparedStatementResult.hasIsUpdate()) {
+        return preparedStatementResult.getIsUpdate();
+      }
+      return null;
+    }
+
     /** Get the schema of the result set (should be identical to {@link #getResultSetSchema()}). */
     public SchemaResult fetchSchema(CallOption... options) {
       checkOpen();