This is an automated email from the ASF dual-hosted git repository.
lidavidm pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-java.git
The following commit(s) were added to refs/heads/main by this push:
new 301d3d8f8 GH-1063: Add is_update field to
ActionCreatePreparedStatementResult (#1064)
301d3d8f8 is described below
commit 301d3d8f889d3e4d1f45fc6e1707460ae5baaeb4
Author: Pedro Matias <[email protected]>
AuthorDate: Tue Jun 30 05:41:30 2026 +0100
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 <[email protected]>
---
arrow-format/FlightSql.proto | 5 ++
.../jdbc/ArrowFlightJdbcFlightStreamResultSet.java | 2 +-
.../ArrowFlightJdbcVectorSchemaRootResultSet.java | 2 +-
.../arrow/driver/jdbc/ArrowFlightMetaImpl.java | 23 ++++++--
.../jdbc/client/ArrowFlightSqlClientHandler.java | 19 +++++++
.../jdbc/ArrowFlightPreparedStatementTest.java | 48 ++++++++++++++++
.../jdbc/ArrowFlightStatementExecuteTest.java | 66 ++++++++++++++++++++++
.../driver/jdbc/utils/MockFlightSqlProducer.java | 42 ++++++++++++++
.../apache/arrow/flight/sql/FlightSqlClient.java | 13 +++++
9 files changed, 212 insertions(+), 8 deletions(-)
diff --git a/arrow-format/FlightSql.proto b/arrow-format/FlightSql.proto
index 566230c2a..b1dc57b33 100644
--- a/arrow-format/FlightSql.proto
+++ b/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;
}
/*
diff --git
a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSet.java
b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSet.java
index 2885f7895..376e5b11e 100644
---
a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSet.java
+++
b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSet.java
@@ -106,7 +106,7 @@ public final class ArrowFlightJdbcFlightStreamResultSet
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);
diff --git
a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcVectorSchemaRootResultSet.java
b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcVectorSchemaRootResultSet.java
index 5d02d6e84..ad6670a00 100644
---
a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcVectorSchemaRootResultSet.java
+++
b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcVectorSchemaRootResultSet.java
@@ -73,7 +73,7 @@ public class ArrowFlightJdbcVectorSchemaRootResultSet extends
AvaticaResultSet {
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);
diff --git
a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java
b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java
index 64529b50c..0d85b5edd 100644
---
a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java
+++
b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java
@@ -53,7 +53,8 @@ public class ArrowFlightMetaImpl extends MetaImpl {
}
/** 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 @@ public class ArrowFlightMetaImpl extends MetaImpl {
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 @@ public class ArrowFlightMetaImpl extends MetaImpl {
((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;
}
diff --git
a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java
b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java
index f0ea28423..719cc38a2 100644
---
a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java
+++
b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/client/ArrowFlightSqlClientHandler.java
@@ -388,6 +388,14 @@ public final class ArrowFlightSqlClientHandler implements
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 final class ArrowFlightSqlClientHandler implements
AutoCloseable {
@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 final class ArrowFlightSqlClientHandler implements
AutoCloseable {
preparedStatement.setParameters(parameters);
}
+ @Override
+ public Boolean isUpdate() {
+ return preparedStatement.isUpdate();
+ }
+
@Override
public void close() {
try {
diff --git
a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java
b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java
index 0369c3a16..078837adf 100644
---
a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java
+++
b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java
@@ -98,6 +98,40 @@ public class ArrowFlightPreparedStatementTest {
}
}
+ @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 class ArrowFlightPreparedStatementTest {
}
}
+ @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";
diff --git
a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementExecuteTest.java
b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementExecuteTest.java
index 632cb0ba5..6acce9c2a 100644
---
a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementExecuteTest.java
+++
b/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 class ArrowFlightStatementExecuteTest {
}));
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 class ArrowFlightStatementExecuteTest {
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()));
+ }
}
diff --git
a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java
b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java
index 45c2a9640..6627d91ab 100644
---
a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/MockFlightSqlProducer.java
+++
b/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.ipc.message.MessageSerializer;
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 final class MockFlightSqlProducer implements
FlightSqlProducer {
});
}
+ /**
+ * 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 final class MockFlightSqlProducer implements
FlightSqlProducer {
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);
diff --git
a/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java
b/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java
index 623f9311e..0af09faee 100644
---
a/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java
+++
b/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java
@@ -1284,6 +1284,19 @@ public class FlightSqlClient implements AutoCloseable {
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();
