From 6958ba5c621180958945f71d5a8ed878b1c2376a Mon Sep 17 00:00:00 2001 From: Pedro Matias Date: Tue, 30 Jun 2026 05:41:30 +0100 Subject: [PATCH] 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 --- arrow-format/FlightSql.proto | 5 ++ .../ArrowFlightJdbcFlightStreamResultSet.java | 2 +- ...owFlightJdbcVectorSchemaRootResultSet.java | 2 +- .../driver/jdbc/ArrowFlightMetaImpl.java | 23 +++++-- .../client/ArrowFlightSqlClientHandler.java | 19 ++++++ .../ArrowFlightPreparedStatementTest.java | 48 ++++++++++++++ .../jdbc/ArrowFlightStatementExecuteTest.java | 66 +++++++++++++++++++ .../jdbc/utils/MockFlightSqlProducer.java | 42 ++++++++++++ .../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 566230c2a6..b1dc57b33b 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 2885f7895b..376e5b11e7 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 @@ 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); 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 49334951de..58a97e13bd 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 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); 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 64529b50c8..0d85b5eddb 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 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 = 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; } 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 f0ea284239..719cc38a2b 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 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 { 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 0369c3a162..078837adf3 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 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"; 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 632cb0ba56..6acce9c2a6 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 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 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 45c2a96404..6627d91ab6 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.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 parameterSchemas = new HashMap<>(); private final Map>> expectedParameterValues = new HashMap<>(); + private final Map isUpdateMap = new HashMap<>(); private final Map 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> 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); 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 623f9311e8..0af09faee1 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 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();