From 2761b9f3ab52aeefa78ce03c7a4f80eb6ceb9506 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Fri, 13 Mar 2026 23:54:27 +0000 Subject: [PATCH 01/25] move orchestration out of ArrowFlightMetaImpl --- .../driver/jdbc/ArrowFlightConnection.java | 28 ++ .../driver/jdbc/ArrowFlightJdbcFactory.java | 25 +- .../ArrowFlightJdbcFlightStreamResultSet.java | 2 +- ...owFlightJdbcVectorSchemaRootResultSet.java | 2 +- .../driver/jdbc/ArrowFlightMetaImpl.java | 240 +++++++----------- .../jdbc/ArrowFlightPreparedStatement.java | 129 ++++++++-- .../driver/jdbc/ArrowFlightStatement.java | 20 +- 7 files changed, 250 insertions(+), 196 deletions(-) diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java index 623c2b81be..166c1157a1 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java @@ -19,6 +19,8 @@ import static org.apache.arrow.driver.jdbc.utils.ArrowFlightConnectionConfigImpl.ArrowFlightConnectionProperty.replaceSemiColons; import io.netty.util.concurrent.DefaultThreadFactory; +import java.sql.PreparedStatement; +import java.sql.ResultSet; import java.sql.SQLException; import java.util.ArrayList; import java.util.HashMap; @@ -257,4 +259,30 @@ BufferAllocator getBufferAllocator() { public ArrowFlightMetaImpl getMeta() { return (ArrowFlightMetaImpl) this.meta; } + + @Override + public PreparedStatement prepareStatement(final String sql) throws SQLException { + checkOpen(); + return prepareStatement(sql, ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY); + } + + @Override + public PreparedStatement prepareStatement( + final String sql, final int resultSetType, final int resultSetConcurrency) + throws SQLException { + checkOpen(); + return prepareStatement(sql, resultSetType, resultSetConcurrency, getHoldability()); + } + + @Override + public PreparedStatement prepareStatement( + final String sql, + final int resultSetType, + final int resultSetConcurrency, + final int resultSetHoldability) + throws SQLException { + checkOpen(); + return getMeta() + .createPreparedStatement(sql, resultSetType, resultSetConcurrency, resultSetHoldability); + } } diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFactory.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFactory.java index e1ccfc820f..202b491f5c 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFactory.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFactory.java @@ -20,7 +20,6 @@ import java.sql.SQLException; import java.util.Properties; import java.util.TimeZone; -import org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler; import org.apache.arrow.memory.RootAllocator; import org.apache.calcite.avatica.AvaticaConnection; import org.apache.calcite.avatica.AvaticaFactory; @@ -79,20 +78,20 @@ public ArrowFlightPreparedStatement newPreparedStatement( final Meta.Signature signature, final int resultType, final int resultSetConcurrency, - final int resultSetHoldability) - throws SQLException { + final int resultSetHoldability) { final ArrowFlightConnection flightConnection = (ArrowFlightConnection) connection; - ArrowFlightSqlClientHandler.PreparedStatement preparedStatement = - flightConnection.getMeta().getPreparedStatement(statementHandle); + final AvaticaStatement existingStatement = + flightConnection.statementMap.get(statementHandle.id); + if (existingStatement instanceof ArrowFlightPreparedStatement) { + return (ArrowFlightPreparedStatement) existingStatement; + } + if (existingStatement != null) { + throw new IllegalStateException( + "Unexpected statement type found for prepared statement handle: " + statementHandle); + } - return ArrowFlightPreparedStatement.newPreparedStatement( - flightConnection, - preparedStatement, - statementHandle, - signature, - resultType, - resultSetConcurrency, - resultSetHoldability); + throw new IllegalStateException( + "PreparedStatement was not pre-created for handle: " + statementHandle); } @Override 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 376e5b11e7..f230c95340 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, null); + final Meta.Signature signature = ArrowFlightMetaImpl.buildDefaultSignature(); 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 ad6670a001..d1d2ddc23a 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, null); + final Meta.Signature signature = ArrowFlightMetaImpl.buildDefaultSignature(); 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 0d85b5eddb..6415455db2 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 @@ -17,20 +17,17 @@ package org.apache.arrow.driver.jdbc; import java.sql.Connection; +import java.sql.ResultSet; import java.sql.SQLException; import java.sql.SQLTimeoutException; import java.util.ArrayList; import java.util.Collections; import java.util.List; -import java.util.Map; -import java.util.concurrent.ConcurrentHashMap; -import org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler.PreparedStatement; -import org.apache.arrow.driver.jdbc.utils.AvaticaParameterBinder; import org.apache.arrow.driver.jdbc.utils.ConvertUtils; -import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.calcite.avatica.AvaticaConnection; import org.apache.calcite.avatica.AvaticaParameter; +import org.apache.calcite.avatica.AvaticaStatement; import org.apache.calcite.avatica.ColumnMetaData; import org.apache.calcite.avatica.MetaImpl; import org.apache.calcite.avatica.NoSuchStatementException; @@ -39,8 +36,6 @@ /** Metadata handler for Arrow Flight. */ public class ArrowFlightMetaImpl extends MetaImpl { - private final Map statementHandlePreparedStatementMap; - /** * Constructs a {@link MetaImpl} object specific for Arrow Flight. * @@ -48,50 +43,14 @@ public class ArrowFlightMetaImpl extends MetaImpl { */ public ArrowFlightMetaImpl(final AvaticaConnection connection) { super(connection); - this.statementHandlePreparedStatementMap = new ConcurrentHashMap<>(); setDefaultConnectionProperties(); } - /** Construct a signature. */ - static Signature newSignature( - final String sql, Schema resultSetSchema, Schema parameterSchema, Boolean isUpdate) { - List columnMetaData = - resultSetSchema == null - ? new ArrayList<>() - : ConvertUtils.convertArrowFieldsToColumnMetaDataList(resultSetSchema.getFields()); - List parameters = - parameterSchema == null - ? new ArrayList<>() - : ConvertUtils.convertArrowFieldsToAvaticaParameters(parameterSchema.getFields()); - // 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, - parameters, - Collections.emptyMap(), - null, // unnecessary, as SQL requests use ArrowFlightJdbcCursor - statementType); - } - @Override public void closeStatement(final StatementHandle statementHandle) { - PreparedStatement preparedStatement = - statementHandlePreparedStatementMap.remove(new StatementHandleKey(statementHandle)); - // Testing if the prepared statement was created because the statement can be - // not created until - // this moment - if (preparedStatement != null) { - preparedStatement.close(); + AvaticaStatement statement = connection.statementMap.get(statementHandle.id); + if (statement instanceof ArrowFlightPreparedStatement) { + ((ArrowFlightPreparedStatement) statement).closePreparedResources(); } } @@ -105,36 +64,8 @@ public ExecuteResult execute( final StatementHandle statementHandle, final List typedValues, final long maxRowCount) { - Preconditions.checkArgument( - connection.id.equals(statementHandle.connectionId), "Connection IDs are not consistent"); - PreparedStatement preparedStatement = getPreparedStatement(statementHandle); - - if (preparedStatement == null) { - throw new IllegalStateException("Prepared statement not found: " + statementHandle); - } - - new AvaticaParameterBinder( - preparedStatement, ((ArrowFlightConnection) connection).getBufferAllocator()) - .bind(typedValues); - - if (statementHandle.signature == null - || statementHandle.signature.statementType == StatementType.IS_DML) { - // Update query - long updatedCount = preparedStatement.executeUpdate(); - return new ExecuteResult( - Collections.singletonList( - MetaResultSet.count(statementHandle.connectionId, statementHandle.id, updatedCount))); - } else { - // TODO Why is maxRowCount ignored? - return new ExecuteResult( - Collections.singletonList( - MetaResultSet.create( - statementHandle.connectionId, - statementHandle.id, - true, - statementHandle.signature, - null))); - } + return getPreparedStatementInstance(statementHandle) + .executeWithTypedValues(statementHandle, typedValues, maxRowCount); } @Override @@ -149,24 +80,8 @@ public ExecuteResult execute( public ExecuteBatchResult executeBatch( final StatementHandle statementHandle, final List> parameterValuesList) throws IllegalStateException { - Preconditions.checkArgument( - connection.id.equals(statementHandle.connectionId), "Connection IDs are not consistent"); - PreparedStatement preparedStatement = getPreparedStatement(statementHandle); - - if (preparedStatement == null) { - throw new IllegalStateException("Prepared statement not found: " + statementHandle); - } - - final AvaticaParameterBinder binder = - new AvaticaParameterBinder( - preparedStatement, ((ArrowFlightConnection) connection).getBufferAllocator()); - for (int i = 0; i < parameterValuesList.size(); i++) { - binder.bind(parameterValuesList.get(i), i); - } - - // Update query - long[] updatedCounts = {preparedStatement.executeUpdate()}; - return new ExecuteBatchResult(updatedCounts); + return getPreparedStatementInstance(statementHandle) + .executeBatchWithTypedValues(statementHandle, parameterValuesList); } @Override @@ -181,25 +96,36 @@ public Frame fetch( String.format("%s does not use frames.", this), AvaticaConnection.HELPER.unsupported()); } - private PreparedStatement prepareForHandle(final String query, StatementHandle handle) { - final PreparedStatement preparedStatement = - ((ArrowFlightConnection) connection).getClientHandler().prepare(query); - handle.signature = - newSignature( - query, - preparedStatement.getDataSetSchema(), - preparedStatement.getParameterSchema(), - preparedStatement.isUpdate()); - statementHandlePreparedStatementMap.put(new StatementHandleKey(handle), preparedStatement); - return preparedStatement; + ArrowFlightPreparedStatement createPreparedStatement( + final String query, + final int resultSetType, + final int resultSetConcurrency, + final int resultSetHoldability) + throws SQLException { + final StatementHandle handle = super.createStatement(connection.handle); + return ArrowFlightPreparedStatement.createPrepared( + (ArrowFlightConnection) connection, + handle, + // null, + query, + resultSetType, + resultSetConcurrency, + resultSetHoldability); } @Override public StatementHandle prepare( final ConnectionHandle connectionHandle, final String query, final long maxRowCount) { - final StatementHandle handle = super.createStatement(connectionHandle); - prepareForHandle(query, handle); - return handle; + try { + return createPreparedStatement( + query, + ResultSet.TYPE_FORWARD_ONLY, + ResultSet.CONCUR_READ_ONLY, + connection.getHoldability()) + .handle; + } catch (SQLException e) { + throw new RuntimeException(e); + } } @Override @@ -222,19 +148,20 @@ public ExecuteResult prepareAndExecute( final PrepareCallback callback) throws NoSuchStatementException { try { - PreparedStatement preparedStatement = prepareForHandle(query, handle); - final StatementType statementType = preparedStatement.getType(); - - final long updateCount = - statementType.equals(StatementType.UPDATE) ? preparedStatement.executeUpdate() : -1; - synchronized (callback.getMonitor()) { - callback.clear(); - callback.assign(handle.signature, null, updateCount); + final AvaticaStatement statement = connection.statementMap.get(handle.id); + if (!(statement instanceof ArrowFlightStatement) + && !(statement instanceof ArrowFlightPreparedStatement)) { + throw new IllegalStateException("Prepared statement not found: " + handle); } - callback.execute(); - final MetaResultSet metaResultSet = - MetaResultSet.create(handle.connectionId, handle.id, false, handle.signature, null); - return new ExecuteResult(Collections.singletonList(metaResultSet)); + final ArrowFlightPreparedStatement preparedStatement = + ArrowFlightPreparedStatement.createPrepared( + (ArrowFlightConnection) connection, + handle, + query, + statement.getResultSetType(), + statement.getResultSetConcurrency(), + statement.getResultSetHoldability()); + return preparedStatement.prepareAndExecute(callback); } catch (SQLTimeoutException e) { // So far AvaticaStatement(executeInternal) only handles NoSuchStatement and // Runtime @@ -291,45 +218,48 @@ void setDefaultConnectionProperties() { .setTransactionIsolation(Connection.TRANSACTION_NONE); } - PreparedStatement getPreparedStatement(StatementHandle statementHandle) { - return statementHandlePreparedStatementMap.get(new StatementHandleKey(statementHandle)); + private ArrowFlightPreparedStatement getPreparedStatementInstance( + StatementHandle statementHandle) { + AvaticaStatement statement = connection.statementMap.get(statementHandle.id); + if (!(statement instanceof ArrowFlightPreparedStatement)) { + throw new IllegalStateException("Prepared statement not found: " + statementHandle); + } + return (ArrowFlightPreparedStatement) statement; } - // Helper used to look up prepared statement instances later. Avatica doesn't - // give us the - // signature in - // an UPDATE code path so we can't directly use StatementHandle as a map key. - private static final class StatementHandleKey { - public final String connectionId; - public final int id; - - StatementHandleKey(StatementHandle statementHandle) { - this.connectionId = statementHandle.connectionId; - this.id = statementHandle.id; + ArrowFlightPreparedStatement getPreparedStatementInstanceOrNull(StatementHandle statementHandle) { + AvaticaStatement statement = connection.statementMap.get(statementHandle.id); + if (statement instanceof ArrowFlightPreparedStatement) { + return (ArrowFlightPreparedStatement) statement; } + return null; + } - @Override - public boolean equals(Object o) { - if (this == o) { - return true; - } - if (o == null || getClass() != o.getClass()) { - return false; - } - - StatementHandleKey that = (StatementHandleKey) o; - - if (id != that.id) { - return false; - } - return connectionId.equals(that.connectionId); - } + public static Signature buildDefaultSignature() { + return buildSignature(null, null, null); + } - @Override - public int hashCode() { - int result = connectionId.hashCode(); - result = 31 * result + id; - return result; - } + /** Builds an Avatica signature from Arrow result and parameter schemas. */ + public static Signature buildSignature( + final String sql, final Schema resultSetSchema, final Schema parameterSchema) { + List columnMetaData = + resultSetSchema == null + ? new ArrayList<>() + : ConvertUtils.convertArrowFieldsToColumnMetaDataList(resultSetSchema.getFields()); + List parameters = + parameterSchema == null + ? new ArrayList<>() + : ConvertUtils.convertArrowFieldsToAvaticaParameters(parameterSchema.getFields()); + StatementType statementType = + resultSetSchema == null || resultSetSchema.getFields().isEmpty() + ? StatementType.IS_DML + : StatementType.SELECT; + return new Signature( + columnMetaData, + sql, + parameters, + Collections.emptyMap(), + null, // unnecessary, as SQL requests use ArrowFlightJdbcCursor + statementType); } } diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java index d7af6902f4..7856d38f07 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java @@ -16,24 +16,32 @@ */ package org.apache.arrow.driver.jdbc; -import java.sql.PreparedStatement; import java.sql.SQLException; +import java.util.Collections; +import java.util.List; import org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler; +import org.apache.arrow.driver.jdbc.utils.AvaticaParameterBinder; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.util.Preconditions; +import org.apache.arrow.vector.types.pojo.Schema; import org.apache.calcite.avatica.AvaticaPreparedStatement; +import org.apache.calcite.avatica.Meta.ExecuteBatchResult; +import org.apache.calcite.avatica.Meta.ExecuteResult; +import org.apache.calcite.avatica.Meta.MetaResultSet; +import org.apache.calcite.avatica.Meta.PrepareCallback; import org.apache.calcite.avatica.Meta.Signature; import org.apache.calcite.avatica.Meta.StatementHandle; +import org.apache.calcite.avatica.Meta.StatementType; +import org.apache.calcite.avatica.remote.TypedValue; -/** Arrow Flight JBCS's implementation {@link PreparedStatement}. */ +/** Arrow Flight JDBC's implementation {@link java.sql.PreparedStatement}. */ public class ArrowFlightPreparedStatement extends AvaticaPreparedStatement implements ArrowFlightInfoStatement { - private final ArrowFlightSqlClientHandler.PreparedStatement preparedStatement; + private ArrowFlightSqlClientHandler.PreparedStatement preparedStatement; private ArrowFlightPreparedStatement( final ArrowFlightConnection connection, - final ArrowFlightSqlClientHandler.PreparedStatement preparedStatement, final StatementHandle handle, final Signature signature, final int resultSetType, @@ -41,26 +49,34 @@ private ArrowFlightPreparedStatement( final int resultSetHoldability) throws SQLException { super(connection, handle, signature, resultSetType, resultSetConcurrency, resultSetHoldability); - this.preparedStatement = Preconditions.checkNotNull(preparedStatement); } - static ArrowFlightPreparedStatement newPreparedStatement( + static ArrowFlightPreparedStatement createPrepared( final ArrowFlightConnection connection, - final ArrowFlightSqlClientHandler.PreparedStatement preparedStmt, final StatementHandle statementHandle, - final Signature signature, + final String query, final int resultSetType, final int resultSetConcurrency, final int resultSetHoldability) throws SQLException { - return new ArrowFlightPreparedStatement( - connection, - preparedStmt, - statementHandle, - signature, - resultSetType, - resultSetConcurrency, - resultSetHoldability); + final ArrowFlightSqlClientHandler.PreparedStatement preparedStatement = + connection.getClientHandler().prepare(query); + final Signature signature = + ArrowFlightMetaImpl.buildSignature( + query, preparedStatement.getDataSetSchema(), preparedStatement.getParameterSchema()); + statementHandle.signature = signature; + + final ArrowFlightPreparedStatement statement = + new ArrowFlightPreparedStatement( + connection, + statementHandle, + signature, + resultSetType, + resultSetConcurrency, + resultSetHoldability); + statement.preparedStatement = Preconditions.checkNotNull(preparedStatement); + statement.setSignature(signature); + return statement; } @Override @@ -68,14 +84,93 @@ public ArrowFlightConnection getConnection() throws SQLException { return (ArrowFlightConnection) super.getConnection(); } + ExecuteResult prepareAndExecute(final PrepareCallback callback) throws SQLException { + ensurePrepared(); + final StatementType statementType = preparedStatement.getType(); + final long updateCount = + statementType.equals(StatementType.UPDATE) ? preparedStatement.executeUpdate() : -1; + synchronized (callback.getMonitor()) { + callback.clear(); + callback.assign(handle.signature, null, updateCount); + } + callback.execute(); + final MetaResultSet metaResultSet = + MetaResultSet.create(handle.connectionId, handle.id, false, handle.signature, null); + return new ExecuteResult(Collections.singletonList(metaResultSet)); + } + + Schema getDataSetSchema() { + ensurePrepared(); + return preparedStatement.getDataSetSchema(); + } + @Override public synchronized void close() throws SQLException { - this.preparedStatement.close(); super.close(); } + void closePreparedResources() { + if (preparedStatement != null) { + preparedStatement.close(); + preparedStatement = null; + } + } + + ExecuteResult executeWithTypedValues( + final StatementHandle statementHandle, + final List typedValues, + final long maxRowCount) { + ensurePrepared(); + Preconditions.checkArgument( + connection.id.equals(statementHandle.connectionId), "Connection IDs are not consistent"); + new AvaticaParameterBinder( + preparedStatement, ((ArrowFlightConnection) connection).getBufferAllocator()) + .bind(typedValues); + + if (statementHandle.signature == null + || statementHandle.signature.statementType == StatementType.IS_DML) { + long updatedCount = preparedStatement.executeUpdate(); + return new ExecuteResult( + Collections.singletonList( + MetaResultSet.count(statementHandle.connectionId, statementHandle.id, updatedCount))); + } + + // TODO Why is maxRowCount ignored? + return new ExecuteResult( + Collections.singletonList( + MetaResultSet.create( + statementHandle.connectionId, + statementHandle.id, + true, + statementHandle.signature, + null))); + } + + ExecuteBatchResult executeBatchWithTypedValues( + final StatementHandle statementHandle, final List> parameterValuesList) { + ensurePrepared(); + Preconditions.checkArgument( + connection.id.equals(statementHandle.connectionId), "Connection IDs are not consistent"); + final AvaticaParameterBinder binder = + new AvaticaParameterBinder( + preparedStatement, ((ArrowFlightConnection) connection).getBufferAllocator()); + for (int i = 0; i < parameterValuesList.size(); i++) { + binder.bind(parameterValuesList.get(i), i); + } + + long[] updatedCounts = {preparedStatement.executeUpdate()}; + return new ExecuteBatchResult(updatedCounts); + } + @Override public FlightInfo executeFlightInfoQuery() throws SQLException { + ensurePrepared(); return preparedStatement.executeQuery(); } + + private void ensurePrepared() { + if (preparedStatement == null) { + throw new IllegalStateException("PreparedStatement is already closed."); + } + } } diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightStatement.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightStatement.java index ff3d060c50..66524c7dbb 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightStatement.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightStatement.java @@ -17,7 +17,6 @@ package org.apache.arrow.driver.jdbc; import java.sql.SQLException; -import org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler.PreparedStatement; import org.apache.arrow.driver.jdbc.utils.ConvertUtils; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.vector.types.pojo.Schema; @@ -44,19 +43,22 @@ public ArrowFlightConnection getConnection() throws SQLException { @Override public FlightInfo executeFlightInfoQuery() throws SQLException { - final PreparedStatement preparedStatement = - getConnection().getMeta().getPreparedStatement(handle); + final ArrowFlightPreparedStatement preparedStatement = + getConnection().getMeta().getPreparedStatementInstanceOrNull(handle); final Meta.Signature signature = getSignature(); if (signature == null) { return null; } - final Schema resultSetSchema = preparedStatement.getDataSetSchema(); - signature.columns.clear(); - signature.columns.addAll( - ConvertUtils.convertArrowFieldsToColumnMetaDataList(resultSetSchema.getFields())); - setSignature(signature); + if (preparedStatement != null) { + final Schema resultSetSchema = preparedStatement.getDataSetSchema(); + signature.columns.clear(); + signature.columns.addAll( + ConvertUtils.convertArrowFieldsToColumnMetaDataList(resultSetSchema.getFields())); + setSignature(signature); + return preparedStatement.executeFlightInfoQuery(); + } - return preparedStatement.executeQuery(); + throw new IllegalStateException("Prepared statement query not found: " + handle); } } From 1e2299c34db363a983498d2307bebd47a6160137 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Sat, 14 Mar 2026 23:13:01 +0000 Subject: [PATCH 02/25] add builder pattern --- .../driver/jdbc/ArrowFlightMetaImpl.java | 30 +++-- .../jdbc/ArrowFlightPreparedStatement.java | 108 +++++++++++++----- .../ArrowFlightPreparedStatementTest.java | 20 ++++ .../jdbc/ArrowFlightStatementExecuteTest.java | 17 +++ 4 files changed, 133 insertions(+), 42 deletions(-) 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 6415455db2..e0a8727630 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 @@ -102,15 +102,13 @@ ArrowFlightPreparedStatement createPreparedStatement( final int resultSetConcurrency, final int resultSetHoldability) throws SQLException { - final StatementHandle handle = super.createStatement(connection.handle); - return ArrowFlightPreparedStatement.createPrepared( - (ArrowFlightConnection) connection, - handle, - // null, - query, - resultSetType, - resultSetConcurrency, - resultSetHoldability); + return ArrowFlightPreparedStatement.builder((ArrowFlightConnection) connection) + .withQuery(query) + .withGeneratedHandle() + .withResultSetType(resultSetType) + .withResultSetConcurrency(resultSetConcurrency) + .withResultSetHoldability(resultSetHoldability) + .build(); } @Override @@ -153,14 +151,14 @@ public ExecuteResult prepareAndExecute( && !(statement instanceof ArrowFlightPreparedStatement)) { throw new IllegalStateException("Prepared statement not found: " + handle); } + if (statement instanceof ArrowFlightPreparedStatement) { + ((ArrowFlightPreparedStatement) statement).closePreparedResources(); + } final ArrowFlightPreparedStatement preparedStatement = - ArrowFlightPreparedStatement.createPrepared( - (ArrowFlightConnection) connection, - handle, - query, - statement.getResultSetType(), - statement.getResultSetConcurrency(), - statement.getResultSetHoldability()); + ArrowFlightPreparedStatement.builder((ArrowFlightConnection) connection) + .withQuery(query) + .withExistingStatement(statement) + .build(); return preparedStatement.prepareAndExecute(callback); } catch (SQLTimeoutException e) { // So far AvaticaStatement(executeInternal) only handles NoSuchStatement and diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java index 7856d38f07..1c6f0cdb21 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java @@ -25,6 +25,7 @@ import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.calcite.avatica.AvaticaPreparedStatement; +import org.apache.calcite.avatica.AvaticaStatement; import org.apache.calcite.avatica.Meta.ExecuteBatchResult; import org.apache.calcite.avatica.Meta.ExecuteResult; import org.apache.calcite.avatica.Meta.MetaResultSet; @@ -44,39 +45,19 @@ private ArrowFlightPreparedStatement( final ArrowFlightConnection connection, final StatementHandle handle, final Signature signature, + final ArrowFlightSqlClientHandler.PreparedStatement preparedStatement, final int resultSetType, final int resultSetConcurrency, final int resultSetHoldability) throws SQLException { super(connection, handle, signature, resultSetType, resultSetConcurrency, resultSetHoldability); + this.preparedStatement = Preconditions.checkNotNull(preparedStatement); + this.handle.signature = signature; + setSignature(signature); } - static ArrowFlightPreparedStatement createPrepared( - final ArrowFlightConnection connection, - final StatementHandle statementHandle, - final String query, - final int resultSetType, - final int resultSetConcurrency, - final int resultSetHoldability) - throws SQLException { - final ArrowFlightSqlClientHandler.PreparedStatement preparedStatement = - connection.getClientHandler().prepare(query); - final Signature signature = - ArrowFlightMetaImpl.buildSignature( - query, preparedStatement.getDataSetSchema(), preparedStatement.getParameterSchema()); - statementHandle.signature = signature; - - final ArrowFlightPreparedStatement statement = - new ArrowFlightPreparedStatement( - connection, - statementHandle, - signature, - resultSetType, - resultSetConcurrency, - resultSetHoldability); - statement.preparedStatement = Preconditions.checkNotNull(preparedStatement); - statement.setSignature(signature); - return statement; + static Builder builder(final ArrowFlightConnection connection) { + return new Builder(connection); } @Override @@ -173,4 +154,79 @@ private void ensurePrepared() { throw new IllegalStateException("PreparedStatement is already closed."); } } + + static final class Builder { + private final ArrowFlightConnection connection; + private StatementHandle handle; + private String query; + private Integer resultSetType; + private Integer resultSetConcurrency; + private Integer resultSetHoldability; + private boolean generateHandle; + + private Builder(final ArrowFlightConnection connection) { + this.connection = Preconditions.checkNotNull(connection); + } + + Builder withQuery(final String query) { + this.query = Preconditions.checkNotNull(query); + return this; + } + + Builder withGeneratedHandle() { + this.generateHandle = true; + this.handle = null; + return this; + } + + Builder withExistingStatement(final AvaticaStatement statement) throws SQLException { + Preconditions.checkNotNull(statement); + this.generateHandle = false; + this.handle = Preconditions.checkNotNull(statement.handle); + this.resultSetType = statement.getResultSetType(); + this.resultSetConcurrency = statement.getResultSetConcurrency(); + this.resultSetHoldability = statement.getResultSetHoldability(); + return this; + } + + Builder withResultSetType(final int resultSetType) { + this.resultSetType = resultSetType; + return this; + } + + Builder withResultSetConcurrency(final int resultSetConcurrency) { + this.resultSetConcurrency = resultSetConcurrency; + return this; + } + + Builder withResultSetHoldability(final int resultSetHoldability) { + this.resultSetHoldability = resultSetHoldability; + return this; + } + + ArrowFlightPreparedStatement build() throws SQLException { + Preconditions.checkNotNull(query); + Preconditions.checkNotNull(resultSetType); + Preconditions.checkNotNull(resultSetConcurrency); + Preconditions.checkNotNull(resultSetHoldability); + if (!generateHandle && handle == null) { + throw new IllegalStateException("PreparedStatement builder requires a handle."); + } + + final ArrowFlightSqlClientHandler.PreparedStatement preparedStatement = + connection.getClientHandler().prepare(query); + final Signature signature = + ArrowFlightMetaImpl.buildSignature( + query, preparedStatement.getDataSetSchema(), preparedStatement.getParameterSchema()); + + return new ArrowFlightPreparedStatement( + connection, + generateHandle ? null : handle, + signature, + preparedStatement, + resultSetType, + resultSetConcurrency, + resultSetHoldability); + } + } } 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 73ece46403..bbdb8b89e7 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 @@ -21,6 +21,8 @@ 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.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertTrue; import java.nio.charset.StandardCharsets; @@ -132,6 +134,24 @@ public void testSimpleQueryNoParameterBindingWithExecuteV2() throws SQLException assertFalse(resultSet.next()); assertFalse(preparedStatement.getMoreResults()); assertEquals(-1, preparedStatement.getUpdateCount()); + @Test + public void testPrepareStatementRegistersCreatedStatementByGeneratedHandle() throws SQLException { + final String query = CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD; + final ArrowFlightConnection flightConnection = (ArrowFlightConnection) connection; + + try (final PreparedStatement preparedStatement = connection.prepareStatement(query)) { + final ArrowFlightPreparedStatement arrowPreparedStatement = + (ArrowFlightPreparedStatement) preparedStatement; + + assertNotNull( + flightConnection + .getMeta() + .getPreparedStatementInstanceOrNull(arrowPreparedStatement.handle)); + assertSame( + arrowPreparedStatement, + flightConnection + .getMeta() + .getPreparedStatementInstanceOrNull(arrowPreparedStatement.handle)); } } 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 6acce9c2a6..e7fc93a4d2 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 @@ -22,6 +22,8 @@ import static org.hamcrest.CoreMatchers.not; import static org.hamcrest.CoreMatchers.nullValue; import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertSame; import java.sql.Connection; import java.sql.ResultSet; @@ -165,6 +167,21 @@ public void testExecuteShouldRunSelectQuery() throws SQLException { is(allOf(equalTo(statement.getLargeUpdateCount()), equalTo(-1L)))); } + @Test + public void testExecuteReplacesStatementMapEntryWithPreparedStatement() throws SQLException { + final ArrowFlightStatement arrowStatement = (ArrowFlightStatement) statement; + final ArrowFlightConnection arrowConnection = (ArrowFlightConnection) connection; + + assertThat(statement.execute(SAMPLE_QUERY_CMD), is(true)); + + final ArrowFlightPreparedStatement preparedStatement = + arrowConnection.getMeta().getPreparedStatementInstanceOrNull(arrowStatement.handle); + + assertNotNull(preparedStatement); + assertSame(preparedStatement, arrowConnection.statementMap.get(arrowStatement.handle.id)); + assertThat(preparedStatement.handle.id, is(equalTo(arrowStatement.handle.id))); + } + @Test public void testExecuteShouldRunUpdateQueryForSmallUpdate() throws SQLException { assertThat(statement.execute(SAMPLE_UPDATE_QUERY), is(false)); // Means this is an UPDATE query. From 0bbf948cd81a7d3fecfdab94aa9e3b3cb8c571b1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Fri, 20 Mar 2026 15:26:11 +0000 Subject: [PATCH 03/25] Optimize statement queries 3 (#3) * reduced number of requests for Statement * move meta orchestration * detect prepared statement --- .../driver/jdbc/ArrowFlightConnection.java | 9 +- .../driver/jdbc/ArrowFlightInfoStatement.java | 36 -- .../ArrowFlightJdbcFlightStreamResultSet.java | 2 +- .../driver/jdbc/ArrowFlightMetaImpl.java | 93 ++--- .../driver/jdbc/ArrowFlightMetaStatement.java | 60 +++ .../jdbc/ArrowFlightPreparedStatement.java | 36 +- .../driver/jdbc/ArrowFlightStatement.java | 139 ++++++- .../client/ArrowFlightSqlClientHandler.java | 10 + .../example/ArrowFlightJdbcSampleApp.java | 120 ++++++ .../ArrowFlightPreparedStatementTest.java | 9 +- .../jdbc/ArrowFlightStatementExecuteTest.java | 20 +- ...ArrowFlightStatementExecuteUpdateTest.java | 10 + .../ArrowFlightStatementProtocolTest.java | 376 ++++++++++++++++++ .../jdbc/utils/MockFlightSqlProducer.java | 22 + 14 files changed, 825 insertions(+), 117 deletions(-) delete mode 100644 flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightInfoStatement.java create mode 100644 flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaStatement.java create mode 100644 flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/example/ArrowFlightJdbcSampleApp.java create mode 100644 flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementProtocolTest.java diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java index 166c1157a1..107cfa0c2f 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java @@ -282,7 +282,12 @@ public PreparedStatement prepareStatement( final int resultSetHoldability) throws SQLException { checkOpen(); - return getMeta() - .createPreparedStatement(sql, resultSetType, resultSetConcurrency, resultSetHoldability); + return ArrowFlightPreparedStatement.builder(this) + .withQuery(sql) + .withGeneratedHandle() + .withResultSetType(resultSetType) + .withResultSetConcurrency(resultSetConcurrency) + .withResultSetHoldability(resultSetHoldability) + .build(); } } diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightInfoStatement.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightInfoStatement.java deleted file mode 100644 index 37ee93722a..0000000000 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightInfoStatement.java +++ /dev/null @@ -1,36 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.apache.arrow.driver.jdbc; - -import java.sql.SQLException; -import java.sql.Statement; -import org.apache.arrow.flight.FlightInfo; - -/** A {@link Statement} that deals with {@link FlightInfo}. */ -public interface ArrowFlightInfoStatement extends Statement { - - @Override - ArrowFlightConnection getConnection() throws SQLException; - - /** - * Executes the query this {@link Statement} is holding. - * - * @return the {@link FlightInfo} for the results. - * @throws SQLException on error. - */ - FlightInfo executeFlightInfoQuery() throws SQLException; -} 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 f230c95340..d383d239d1 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 @@ -67,7 +67,7 @@ public final class ArrowFlightJdbcFlightStreamResultSet throws SQLException { super(statement, state, signature, resultSetMetaData, timeZone, firstFrame); this.connection = (ArrowFlightConnection) statement.connection; - this.flightInfo = ((ArrowFlightInfoStatement) statement).executeFlightInfoQuery(); + this.flightInfo = ((ArrowFlightMetaStatement) statement).executeFlightInfoQuery(); } /** Private constructor for fromFlightInfo. */ 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 e0a8727630..4da182ca63 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 @@ -48,10 +48,7 @@ public ArrowFlightMetaImpl(final AvaticaConnection connection) { @Override public void closeStatement(final StatementHandle statementHandle) { - AvaticaStatement statement = connection.statementMap.get(statementHandle.id); - if (statement instanceof ArrowFlightPreparedStatement) { - ((ArrowFlightPreparedStatement) statement).closePreparedResources(); - } + getMetaStatement(statementHandle).closeStatement(); } @Override @@ -64,8 +61,7 @@ public ExecuteResult execute( final StatementHandle statementHandle, final List typedValues, final long maxRowCount) { - return getPreparedStatementInstance(statementHandle) - .executeWithTypedValues(statementHandle, typedValues, maxRowCount); + return getMetaStatement(statementHandle).execute(statementHandle, typedValues, maxRowCount); } @Override @@ -80,8 +76,7 @@ public ExecuteResult execute( public ExecuteBatchResult executeBatch( final StatementHandle statementHandle, final List> parameterValuesList) throws IllegalStateException { - return getPreparedStatementInstance(statementHandle) - .executeBatchWithTypedValues(statementHandle, parameterValuesList); + return getMetaStatement(statementHandle).executeBatch(statementHandle, parameterValuesList); } @Override @@ -96,31 +91,16 @@ public Frame fetch( String.format("%s does not use frames.", this), AvaticaConnection.HELPER.unsupported()); } - ArrowFlightPreparedStatement createPreparedStatement( - final String query, - final int resultSetType, - final int resultSetConcurrency, - final int resultSetHoldability) - throws SQLException { - return ArrowFlightPreparedStatement.builder((ArrowFlightConnection) connection) - .withQuery(query) - .withGeneratedHandle() - .withResultSetType(resultSetType) - .withResultSetConcurrency(resultSetConcurrency) - .withResultSetHoldability(resultSetHoldability) - .build(); - } - @Override public StatementHandle prepare( final ConnectionHandle connectionHandle, final String query, final long maxRowCount) { try { - return createPreparedStatement( - query, - ResultSet.TYPE_FORWARD_ONLY, - ResultSet.CONCUR_READ_ONLY, - connection.getHoldability()) - .handle; + // This is the Avatica entry point used by Connection.prepareStatement(String). + ArrowFlightPreparedStatement stmt = + (ArrowFlightPreparedStatement) + connection.prepareStatement( + query, ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY); + return stmt.handle; } catch (SQLException e) { throw new RuntimeException(e); } @@ -133,6 +113,7 @@ public ExecuteResult prepareAndExecute( final long maxRowCount, final PrepareCallback prepareCallback) throws NoSuchStatementException { + // This is the Avatica entry point used by Statement.execute(String). return prepareAndExecute( statementHandle, query, maxRowCount, -1 /* Not used */, prepareCallback); } @@ -146,20 +127,9 @@ public ExecuteResult prepareAndExecute( final PrepareCallback callback) throws NoSuchStatementException { try { - final AvaticaStatement statement = connection.statementMap.get(handle.id); - if (!(statement instanceof ArrowFlightStatement) - && !(statement instanceof ArrowFlightPreparedStatement)) { - throw new IllegalStateException("Prepared statement not found: " + handle); - } - if (statement instanceof ArrowFlightPreparedStatement) { - ((ArrowFlightPreparedStatement) statement).closePreparedResources(); - } - final ArrowFlightPreparedStatement preparedStatement = - ArrowFlightPreparedStatement.builder((ArrowFlightConnection) connection) - .withQuery(query) - .withExistingStatement(statement) - .build(); - return preparedStatement.prepareAndExecute(callback); + // This is the Avatica entry point used by Statement.execute(String). + return getMetaStatement(handle) + .prepareAndExecute(query, maxRowCount, maxRowsInFirstFrame, callback); } catch (SQLTimeoutException e) { // So far AvaticaStatement(executeInternal) only handles NoSuchStatement and // Runtime @@ -216,30 +186,37 @@ void setDefaultConnectionProperties() { .setTransactionIsolation(Connection.TRANSACTION_NONE); } - private ArrowFlightPreparedStatement getPreparedStatementInstance( - StatementHandle statementHandle) { + private ArrowFlightMetaStatement getMetaStatement(StatementHandle statementHandle) { AvaticaStatement statement = connection.statementMap.get(statementHandle.id); - if (!(statement instanceof ArrowFlightPreparedStatement)) { - throw new IllegalStateException("Prepared statement not found: " + statementHandle); + if (statement instanceof ArrowFlightMetaStatement) { + return (ArrowFlightMetaStatement) statement; } - return (ArrowFlightPreparedStatement) statement; + throw new IllegalStateException("Statement not found: " + statementHandle); } - ArrowFlightPreparedStatement getPreparedStatementInstanceOrNull(StatementHandle statementHandle) { - AvaticaStatement statement = connection.statementMap.get(statementHandle.id); - if (statement instanceof ArrowFlightPreparedStatement) { - return (ArrowFlightPreparedStatement) statement; - } - return null; + public static Signature buildDefaultSignature() { + return buildSignature(null, StatementType.SELECT); } - public static Signature buildDefaultSignature() { - return buildSignature(null, null, null); + public static Signature buildSignature(final String sql, final StatementType type) { + return buildSignature(sql, null, null, type); } /** Builds an Avatica signature from Arrow result and parameter schemas. */ public static Signature buildSignature( final String sql, final Schema resultSetSchema, final Schema parameterSchema) { + StatementType statementType = + resultSetSchema == null || resultSetSchema.getFields().isEmpty() + ? StatementType.IS_DML + : StatementType.SELECT; + return buildSignature(sql, resultSetSchema, parameterSchema, statementType); + } + + private static Signature buildSignature( + final String sql, + final Schema resultSetSchema, + final Schema parameterSchema, + final StatementType statementType) { List columnMetaData = resultSetSchema == null ? new ArrayList<>() @@ -248,10 +225,6 @@ public static Signature buildSignature( parameterSchema == null ? new ArrayList<>() : ConvertUtils.convertArrowFieldsToAvaticaParameters(parameterSchema.getFields()); - StatementType statementType = - resultSetSchema == null || resultSetSchema.getFields().isEmpty() - ? StatementType.IS_DML - : StatementType.SELECT; return new Signature( columnMetaData, sql, diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaStatement.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaStatement.java new file mode 100644 index 0000000000..415af19e8f --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaStatement.java @@ -0,0 +1,60 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.driver.jdbc; + +import java.sql.SQLException; +import java.sql.Statement; +import java.util.List; +import org.apache.arrow.flight.FlightInfo; +import org.apache.calcite.avatica.Meta.ExecuteBatchResult; +import org.apache.calcite.avatica.Meta.ExecuteResult; +import org.apache.calcite.avatica.Meta.PrepareCallback; +import org.apache.calcite.avatica.Meta.StatementHandle; +import org.apache.calcite.avatica.remote.TypedValue; + +/** Statement capabilities used by {@link ArrowFlightMetaImpl}. */ +interface ArrowFlightMetaStatement extends Statement { + + @Override + ArrowFlightConnection getConnection() throws SQLException; + + FlightInfo executeFlightInfoQuery() throws SQLException; + + /** + * Avatica routes {@link Statement#execute(String)} through Meta.prepareAndExecute(...), so plain + * statements still need this hook even when they support direct executeQuery/executeUpdate paths. + */ + ExecuteResult prepareAndExecute( + String query, long maxRowCount, int maxRowsInFirstFrame, PrepareCallback callback) + throws SQLException; + + default ExecuteResult execute( + final StatementHandle statementHandle, + final List typedValues, + final long maxRowCount) { + throw new IllegalStateException( + "Statement operation is not supported for handle: " + statementHandle); + } + + default ExecuteBatchResult executeBatch( + final StatementHandle statementHandle, final List> parameterValuesList) { + throw new IllegalStateException( + "Statement operation is not supported for handle: " + statementHandle); + } + + default void closeStatement() {} +} diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java index 1c6f0cdb21..bd7ebbe0e4 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java @@ -37,7 +37,7 @@ /** Arrow Flight JDBC's implementation {@link java.sql.PreparedStatement}. */ public class ArrowFlightPreparedStatement extends AvaticaPreparedStatement - implements ArrowFlightInfoStatement { + implements ArrowFlightMetaStatement { private ArrowFlightSqlClientHandler.PreparedStatement preparedStatement; @@ -80,6 +80,21 @@ ExecuteResult prepareAndExecute(final PrepareCallback callback) throws SQLExcept return new ExecuteResult(Collections.singletonList(metaResultSet)); } + @Override + public ExecuteResult prepareAndExecute( + final String query, + final long maxRowCount, + final int maxRowsInFirstFrame, + final PrepareCallback callback) + throws SQLException { + + return ArrowFlightPreparedStatement.builder(getConnection()) + .withQuery(query) + .withExistingStatement(this) + .build() + .prepareAndExecute(callback); + } + Schema getDataSetSchema() { ensurePrepared(); return preparedStatement.getDataSetSchema(); @@ -143,6 +158,25 @@ ExecuteBatchResult executeBatchWithTypedValues( return new ExecuteBatchResult(updatedCounts); } + @Override + public ExecuteResult execute( + final StatementHandle statementHandle, + final List typedValues, + final long maxRowCount) { + return executeWithTypedValues(statementHandle, typedValues, maxRowCount); + } + + @Override + public ExecuteBatchResult executeBatch( + final StatementHandle statementHandle, final List> parameterValuesList) { + return executeBatchWithTypedValues(statementHandle, parameterValuesList); + } + + @Override + public void closeStatement() { + closePreparedResources(); + } + @Override public FlightInfo executeFlightInfoQuery() throws SQLException { ensurePrepared(); diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightStatement.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightStatement.java index 66524c7dbb..5bbdb8ff87 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightStatement.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightStatement.java @@ -16,16 +16,24 @@ */ package org.apache.arrow.driver.jdbc; +import java.sql.ResultSet; import java.sql.SQLException; import org.apache.arrow.driver.jdbc.utils.ConvertUtils; import org.apache.arrow.flight.FlightInfo; +import org.apache.arrow.flight.FlightRuntimeException; +import org.apache.arrow.flight.FlightStatusCode; import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.calcite.avatica.AvaticaConnection; +import org.apache.calcite.avatica.AvaticaResultSet; import org.apache.calcite.avatica.AvaticaStatement; import org.apache.calcite.avatica.Meta; +import org.apache.calcite.avatica.Meta.ExecuteResult; +import org.apache.calcite.avatica.Meta.PrepareCallback; import org.apache.calcite.avatica.Meta.StatementHandle; +import org.apache.calcite.avatica.Meta.StatementType; /** A SQL statement for querying data from an Arrow Flight server. */ -public class ArrowFlightStatement extends AvaticaStatement implements ArrowFlightInfoStatement { +public class ArrowFlightStatement extends AvaticaStatement implements ArrowFlightMetaStatement { ArrowFlightStatement( final ArrowFlightConnection connection, @@ -41,24 +49,141 @@ public ArrowFlightConnection getConnection() throws SQLException { return (ArrowFlightConnection) super.getConnection(); } + @Override + public ExecuteResult prepareAndExecute( + final String query, + final long maxRowCount, + final int maxRowsInFirstFrame, + final PrepareCallback callback) + throws SQLException { + // Keep Avatica Statement.execute(String) behavior: Avatica calls Meta.prepareAndExecute, + // which resolves to this statement hook. + this.closeStatement(); + + return ArrowFlightPreparedStatement.builder(getConnection()) + .withQuery(query) + .withExistingStatement(this) + .build() + .prepareAndExecute(callback); + } + + @Override + public ResultSet executeQuery(final String sql) throws SQLException { + checkOpen(); + updateCount = -1; + switchToDirectStatementMode(); + try { + final Meta.Signature signature = + ArrowFlightMetaImpl.buildSignature(sql, StatementType.SELECT); + setSignature(signature); + return executeQueryInternal(signature, false); + } catch (Exception exception) { + throw wrapStatementExecutionException(sql, exception); + } + } + + @Override + public long executeLargeUpdate(final String sql) throws SQLException { + checkOpen(); + clearOpenResultSet(); + updateCount = -1; + switchToDirectStatementMode(); + + try { + final long updatedCount = getConnection().getClientHandler().executeUpdate(sql); + setSignature(ArrowFlightMetaImpl.buildSignature(sql, StatementType.IS_DML)); + updateCount = updatedCount; + return updatedCount; + } catch (Exception exception) { + throw wrapStatementExecutionException(sql, exception); + } + } + @Override public FlightInfo executeFlightInfoQuery() throws SQLException { - final ArrowFlightPreparedStatement preparedStatement = - getConnection().getMeta().getPreparedStatementInstanceOrNull(handle); + final ArrowFlightConnection connection = getConnection(); final Meta.Signature signature = getSignature(); if (signature == null) { return null; } - if (preparedStatement != null) { - final Schema resultSetSchema = preparedStatement.getDataSetSchema(); + // A Statement handle can point to either this direct statement instance or a prepared + // statement instance created by Avatica Statement.execute(String) through + // Meta.prepareAndExecute. + final AvaticaStatement currentStatement = connection.statementMap.get(handle.id); + if (currentStatement instanceof ArrowFlightPreparedStatement) { + // Prepared path: reuse the current statement implementation associated with the handle. + final FlightInfo flightInfo = + ((ArrowFlightPreparedStatement) currentStatement).executeFlightInfoQuery(); + updateSignatureColumnsFromFlightInfo(signature, flightInfo); + return flightInfo; + } + + // Direct Statement.executeQuery(String) / executeUpdate(String) path. + final FlightInfo flightInfo = connection.getClientHandler().getInfo(signature.sql); + updateSignatureColumnsFromFlightInfo(signature, flightInfo); + return flightInfo; + } + + private void updateSignatureColumnsFromFlightInfo( + final Meta.Signature signature, final FlightInfo flightInfo) { + final Schema resultSetSchema = flightInfo.getSchemaOptional().orElse(null); + if (resultSetSchema != null) { signature.columns.clear(); signature.columns.addAll( ConvertUtils.convertArrowFieldsToColumnMetaDataList(resultSetSchema.getFields())); setSignature(signature); - return preparedStatement.executeFlightInfoQuery(); } + } + + private SQLException wrapStatementExecutionException(final String sql, final Exception exception) + throws SQLException { + if (!(exception instanceof SQLException)) { + return AvaticaConnection.HELPER.createException( + "Error while executing SQL \"" + sql + "\": " + exception.getMessage(), exception); + } + final SQLException sqlException = (SQLException) exception; + final String prefix = "Error while executing SQL \"" + sql + "\""; + final String message = sqlException.getMessage(); + if (message != null && message.startsWith(prefix)) { + return sqlException; + } + final Throwable cause = sqlException.getCause(); + if (cause instanceof FlightRuntimeException) { + final FlightStatusCode statusCode = ((FlightRuntimeException) cause).status().code(); + if (statusCode == FlightStatusCode.UNAVAILABLE) { + return sqlException; + } + } + return AvaticaConnection.HELPER.createException(prefix + ": " + message, sqlException); + } - throw new IllegalStateException("Prepared statement query not found: " + handle); + private void clearOpenResultSet() throws SQLException { + synchronized (this) { + if (openResultSet != null) { + final AvaticaResultSet resultSet = openResultSet; + openResultSet = null; + try { + resultSet.close(); + } catch (Exception exception) { + throw AvaticaConnection.HELPER.createException( + "Error while closing previous result set", exception); + } + } + } + } + + private void switchToDirectStatementMode() throws SQLException { + final ArrowFlightConnection connection = getConnection(); + final AvaticaStatement existingStatement = connection.statementMap.get(handle.id); + if (existingStatement == this) { + return; + } + if (existingStatement instanceof ArrowFlightPreparedStatement) { + // Release resources from previously attached statement implementation before switching back + // to direct statement mode for executeQuery/executeUpdate. + ((ArrowFlightPreparedStatement) existingStatement).closeStatement(); + } + connection.statementMap.put(handle.id, this); } } 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 719cc38a2b..09b8dfecee 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 @@ -267,6 +267,16 @@ public FlightInfo getInfo(final String query) { return sqlClient.execute(query, getOptions()); } + /** + * Executes an update query directly, without creating a prepared statement first. + * + * @param query The update query. + * @return the number of rows affected. + */ + public long executeUpdate(final String query) { + return sqlClient.executeUpdate(query, getOptions()); + } + @Override public void close() throws SQLException { if (catalog.isPresent()) { diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/example/ArrowFlightJdbcSampleApp.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/example/ArrowFlightJdbcSampleApp.java new file mode 100644 index 0000000000..10ce0bd285 --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/example/ArrowFlightJdbcSampleApp.java @@ -0,0 +1,120 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.driver.jdbc.example; + +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.Properties; + +/** + * Minimal sample app for using the Arrow Flight SQL JDBC driver. + * + *

Defaults are configured for a local Dremio instance: + * + *

    + *
  • host: {@code localhost} + *
  • port: {@code 32010} + *
  • user: {@code dremio} + *
  • password: {@code dremio123} + *
+ * + *

Arguments are optional and positional: + * + * 

+ *   [host] [port] [user] [password] [selectSql] [updateSql]
+ *
+ * + *

If {@code updateSql} is omitted, only {@code Statement.executeQuery(...)} is executed. + */ +public final class ArrowFlightJdbcSampleApp { + private static final String DEFAULT_HOST = "localhost"; + private static final int DEFAULT_PORT = 32010; + private static final String DEFAULT_USER = "dremio"; + private static final String DEFAULT_PASSWORD = "dremio123"; + private static final String DEFAULT_SELECT_SQL = "SELECT 1 AS sample_value"; + + private ArrowFlightJdbcSampleApp() {} + + public static void main(final String[] args) throws Exception { + final String host = getArg(args, 0, DEFAULT_HOST); + final int port = Integer.parseInt(getArg(args, 1, Integer.toString(DEFAULT_PORT))); + final String user = getArg(args, 2, DEFAULT_USER); + final String password = getArg(args, 3, DEFAULT_PASSWORD); + final String selectSql = getArg(args, 4, DEFAULT_SELECT_SQL); + final String updateSql = getArg(args, 5, ""); + + final String url = String.format("jdbc:arrow-flight-sql://%s:%d", host, port); + final Properties properties = new Properties(); + properties.setProperty("user", user); + properties.setProperty("password", password); + properties.setProperty("useEncryption", "false"); + + System.out.println("Connecting to " + url); + try (Connection connection = DriverManager.getConnection(url, properties); + Statement statement = connection.createStatement()) { + runSelect(statement, selectSql); + + if (updateSql.isEmpty()) { + System.out.println( + "Skipping Statement.executeUpdate(...) because no updateSql argument was provided."); + } else { + runUpdate(statement, updateSql); + } + } + } + + private static void runSelect(final Statement statement, final String selectSql) + throws SQLException { + System.out.println("Running Statement.executeQuery: " + selectSql); + try (ResultSet resultSet = statement.executeQuery(selectSql)) { + final ResultSetMetaData metadata = resultSet.getMetaData(); + final int columnCount = metadata.getColumnCount(); + int rowCount = 0; + while (resultSet.next()) { + rowCount++; + final StringBuilder rowBuilder = new StringBuilder(); + for (int i = 1; i <= columnCount; i++) { + if (i > 1) { + rowBuilder.append(", "); + } + rowBuilder.append(metadata.getColumnLabel(i)).append('=').append(resultSet.getObject(i)); + } + System.out.println("row " + rowCount + ": " + rowBuilder); + } + System.out.println("Statement.executeQuery returned " + rowCount + " row(s)"); + } + } + + private static void runUpdate(final Statement statement, final String updateSql) + throws SQLException { + System.out.println("Running Statement.executeUpdate: " + updateSql); + final int updateCount = statement.executeUpdate(updateSql); + System.out.println("Statement.executeUpdate affected " + updateCount + " row(s)"); + } + + private static String getArg(final String[] args, final int index, final String defaultValue) { + if (index >= args.length) { + return defaultValue; + } + final String arg = args[index]; + return arg == null || arg.isEmpty() ? defaultValue : arg; + } +} 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 bbdb8b89e7..e8bfd4025b 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 @@ -143,15 +143,10 @@ public void testPrepareStatementRegistersCreatedStatementByGeneratedHandle() thr final ArrowFlightPreparedStatement arrowPreparedStatement = (ArrowFlightPreparedStatement) preparedStatement; - assertNotNull( - flightConnection - .getMeta() - .getPreparedStatementInstanceOrNull(arrowPreparedStatement.handle)); + assertNotNull(flightConnection.statementMap.get(arrowPreparedStatement.handle.id)); assertSame( arrowPreparedStatement, - flightConnection - .getMeta() - .getPreparedStatementInstanceOrNull(arrowPreparedStatement.handle)); + flightConnection.statementMap.get(arrowPreparedStatement.handle.id)); } } 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 e7fc93a4d2..f6f05c2bd0 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 @@ -18,6 +18,7 @@ import static org.hamcrest.CoreMatchers.allOf; import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.CoreMatchers.instanceOf; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.CoreMatchers.not; import static org.hamcrest.CoreMatchers.nullValue; @@ -174,12 +175,25 @@ public void testExecuteReplacesStatementMapEntryWithPreparedStatement() throws S assertThat(statement.execute(SAMPLE_QUERY_CMD), is(true)); - final ArrowFlightPreparedStatement preparedStatement = - arrowConnection.getMeta().getPreparedStatementInstanceOrNull(arrowStatement.handle); + final Object preparedStatement = arrowConnection.statementMap.get(arrowStatement.handle.id); assertNotNull(preparedStatement); assertSame(preparedStatement, arrowConnection.statementMap.get(arrowStatement.handle.id)); - assertThat(preparedStatement.handle.id, is(equalTo(arrowStatement.handle.id))); + assertThat(preparedStatement, instanceOf(ArrowFlightPreparedStatement.class)); + } + + @Test + public void testExecuteQueryRestoresStatementMapEntryWithStatement() throws SQLException { + final ArrowFlightStatement arrowStatement = (ArrowFlightStatement) statement; + final ArrowFlightConnection arrowConnection = (ArrowFlightConnection) connection; + + assertThat(statement.execute(SAMPLE_QUERY_CMD), is(true)); + + try (ResultSet resultSet = statement.executeQuery(SAMPLE_QUERY_CMD)) { + assertThat(resultSet.next(), is(true)); + } + + assertSame(arrowStatement, arrowConnection.statementMap.get(arrowStatement.handle.id)); } @Test diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementExecuteUpdateTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementExecuteUpdateTest.java index f7c31c590c..05e85227f0 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementExecuteUpdateTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementExecuteUpdateTest.java @@ -22,6 +22,7 @@ import static org.hamcrest.CoreMatchers.instanceOf; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.CoreMatchers.not; +import static org.hamcrest.CoreMatchers.startsWith; import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.jupiter.api.Assertions.assertThrows; @@ -224,4 +225,13 @@ public void testShouldFailToPrepareStatementForBadStatement() { } assertThat(count, is(1)); } + + @Test + public void testExecuteLargeUpdateShouldWrapBadStatement() { + final String badQuery = "BAD INVALID UPDATE"; + final SQLException exception = + assertThrows(SQLException.class, () -> statement.executeLargeUpdate(badQuery)); + assertThat( + exception.getMessage(), startsWith(format("Error while executing SQL \"%s\"", badQuery))); + } } diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementProtocolTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementProtocolTest.java new file mode 100644 index 0000000000..c5c35a9173 --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementProtocolTest.java @@ -0,0 +1,376 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.driver.jdbc; + +import static org.hamcrest.CoreMatchers.is; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.google.protobuf.Message; +import java.sql.Connection; +import java.sql.DatabaseMetaData; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.Collections; +import java.util.function.Consumer; +import org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer; +import org.apache.arrow.flight.FlightProducer.ServerStreamListener; +import org.apache.arrow.flight.sql.FlightSqlProducer.Schemas; +import org.apache.arrow.flight.sql.FlightSqlUtils; +import org.apache.arrow.flight.sql.impl.FlightSql.CommandGetDbSchemas; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.util.AutoCloseables; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.Types.MinorType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.arrow.vector.util.Text; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; + +public class ArrowFlightStatementProtocolTest { + private static final String SELECT_QUERY = "SELECT * FROM PROTOCOL_SELECT"; + private static final String UPDATE_QUERY = "UPDATE PROTOCOL_UPDATE"; + private static final Schema QUERY_SCHEMA = + new Schema(Collections.singletonList(Field.nullable("id", MinorType.INT.getType()))); + + private static final MockFlightSqlProducer PRODUCER = new MockFlightSqlProducer(); + + @RegisterExtension + public static final FlightServerTestExtension FLIGHT_SERVER_TEST_EXTENSION = + FlightServerTestExtension.createStandardTestExtension(PRODUCER); + + private Connection connection; + + @BeforeAll + public static void setUpBeforeClass() { + PRODUCER.addSelectQuery( + SELECT_QUERY, + QUERY_SCHEMA, + Collections.singletonList( + listener -> { + try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + final VectorSchemaRoot root = VectorSchemaRoot.create(QUERY_SCHEMA, allocator)) { + IntVector vector = (IntVector) root.getVector("id"); + vector.setSafe(0, 1); + root.setRowCount(1); + listener.start(root); + listener.putNext(); + } catch (final Throwable throwable) { + listener.error(throwable); + } finally { + listener.completed(); + } + })); + PRODUCER.addUpdateQuery(UPDATE_QUERY, 1); + + final Message commandGetDbSchemas = CommandGetDbSchemas.getDefaultInstance(); + final Consumer commandGetSchemasResultProducer = + listener -> { + try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + final VectorSchemaRoot root = + VectorSchemaRoot.create(Schemas.GET_SCHEMAS_SCHEMA, allocator)) { + final VarCharVector catalogName = (VarCharVector) root.getVector("catalog_name"); + final VarCharVector schemaName = (VarCharVector) root.getVector("db_schema_name"); + catalogName.setSafe(0, new Text("catalog_name #0")); + schemaName.setSafe(0, new Text("db_schema_name #0")); + root.setRowCount(1); + listener.start(root); + listener.putNext(); + } catch (final Throwable throwable) { + listener.error(throwable); + } finally { + listener.completed(); + } + }; + PRODUCER.addCatalogQuery(commandGetDbSchemas, commandGetSchemasResultProducer); + } + + @BeforeEach + public void setUp() throws SQLException { + PRODUCER.clearActionTypeCounter(); + PRODUCER.clearCommandTypeCounter(); + connection = FLIGHT_SERVER_TEST_EXTENSION.getConnection(false); + } + + @AfterEach + public void tearDown() throws Exception { + AutoCloseables.close(connection); + } + + @AfterAll + public static void tearDownAfterClass() throws Exception { + AutoCloseables.close(PRODUCER); + } + + @Test + public void testStatementExecuteQueryUsesStatementProtocol() throws SQLException { + try (Statement statement = connection.createStatement(); + ResultSet resultSet = statement.executeQuery(SELECT_QUERY)) { + assertTrue(resultSet.next()); + } + + assertThat( + PRODUCER + .getActionTypeCounter() + .getOrDefault(FlightSqlUtils.FLIGHT_SQL_CREATE_PREPARED_STATEMENT.getType(), 0), + is(0)); + assertThat( + PRODUCER + .getCommandTypeCounter() + .getOrDefault(MockFlightSqlProducer.COMMAND_STATEMENT_QUERY, 0), + is(1)); + assertThat( + PRODUCER + .getCommandTypeCounter() + .getOrDefault(MockFlightSqlProducer.COMMAND_PREPARED_STATEMENT_QUERY, 0), + is(0)); + } + + @Test + public void testStatementExecuteUsesPreparedProtocolForQuery() throws SQLException { + try (Statement statement = connection.createStatement()) { + assertThat(statement.execute(SELECT_QUERY), is(true)); + try (ResultSet resultSet = statement.getResultSet()) { + assertTrue(resultSet.next()); + } + } + + 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)); + assertThat( + PRODUCER + .getCommandTypeCounter() + .getOrDefault(MockFlightSqlProducer.COMMAND_STATEMENT_QUERY, 0), + is(0)); + } + + @Test + public void testStatementExecuteUpdateUsesStatementProtocol() throws SQLException { + try (Statement statement = connection.createStatement()) { + assertThat(statement.executeUpdate(UPDATE_QUERY), is(1)); + } + + assertThat( + PRODUCER + .getActionTypeCounter() + .getOrDefault(FlightSqlUtils.FLIGHT_SQL_CREATE_PREPARED_STATEMENT.getType(), 0), + is(0)); + assertThat( + PRODUCER + .getCommandTypeCounter() + .getOrDefault(MockFlightSqlProducer.COMMAND_STATEMENT_UPDATE, 0), + is(1)); + assertThat( + PRODUCER + .getCommandTypeCounter() + .getOrDefault(MockFlightSqlProducer.COMMAND_PREPARED_STATEMENT_UPDATE, 0), + is(0)); + } + + @Test + public void testStatementExecuteUsesPreparedProtocolForUpdate() throws SQLException { + try (Statement statement = connection.createStatement()) { + assertThat(statement.execute(UPDATE_QUERY), is(false)); + assertThat(statement.getUpdateCount(), 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 testStatementExecuteThenExecuteUpdateUsesStatementProtocol() throws SQLException { + try (Statement statement = connection.createStatement()) { + assertThat(statement.execute(SELECT_QUERY), is(true)); + try (ResultSet resultSet = statement.getResultSet()) { + assertTrue(resultSet.next()); + } + assertThat(statement.executeUpdate(UPDATE_QUERY), 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)); + assertThat( + PRODUCER + .getCommandTypeCounter() + .getOrDefault(MockFlightSqlProducer.COMMAND_STATEMENT_UPDATE, 0), + is(1)); + } + + @Test + public void testStatementExecuteUpdateThenExecuteQueryUsesStatementProtocol() + throws SQLException { + try (Statement statement = connection.createStatement()) { + assertThat(statement.execute(UPDATE_QUERY), is(false)); + assertThat(statement.getUpdateCount(), is(1)); + try (ResultSet resultSet = statement.executeQuery(SELECT_QUERY)) { + assertTrue(resultSet.next()); + } + } + + 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)); + assertThat( + PRODUCER + .getCommandTypeCounter() + .getOrDefault(MockFlightSqlProducer.COMMAND_STATEMENT_QUERY, 0), + is(1)); + } + + @Test + public void testPreparedStatementExecuteQueryUsesPreparedProtocol() throws SQLException { + try (PreparedStatement statement = connection.prepareStatement(SELECT_QUERY); + ResultSet resultSet = statement.executeQuery()) { + assertTrue(resultSet.next()); + } + + 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)); + assertThat( + PRODUCER + .getCommandTypeCounter() + .getOrDefault(MockFlightSqlProducer.COMMAND_STATEMENT_QUERY, 0), + is(0)); + } + + @Test + public void testPreparedStatementExecuteUsesPreparedProtocolForQuery() throws SQLException { + try (PreparedStatement statement = connection.prepareStatement(SELECT_QUERY)) { + assertThat(statement.execute(), is(true)); + try (ResultSet resultSet = statement.getResultSet()) { + assertTrue(resultSet.next()); + } + } + + 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)); + assertThat( + PRODUCER + .getCommandTypeCounter() + .getOrDefault(MockFlightSqlProducer.COMMAND_STATEMENT_QUERY, 0), + is(0)); + } + + @Test + public void testPreparedStatementExecuteUpdateUsesPreparedProtocol() throws SQLException { + try (PreparedStatement statement = connection.prepareStatement(UPDATE_QUERY)) { + assertThat(statement.executeUpdate(), 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 testPreparedStatementExecuteUsesPreparedProtocolForUpdate() throws SQLException { + try (PreparedStatement statement = connection.prepareStatement(UPDATE_QUERY)) { + assertThat(statement.execute(), is(false)); + assertThat(statement.getUpdateCount(), 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 testMetadataGetSchemasUsesJdbcApi() throws SQLException { + final DatabaseMetaData metaData = connection.getMetaData(); + try (ResultSet resultSet = metaData.getSchemas()) { + assertTrue(resultSet.next()); + } + + assertThat( + PRODUCER + .getActionTypeCounter() + .getOrDefault(FlightSqlUtils.FLIGHT_SQL_CREATE_PREPARED_STATEMENT.getType(), 0), + is(0)); + } +} 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 6627d91ab6..dfa536f230 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 @@ -105,6 +105,12 @@ public final class MockFlightSqlProducer implements FlightSqlProducer { private final Map isUpdateMap = new HashMap<>(); private final Map actionTypeCounter = new HashMap<>(); + private final Map commandTypeCounter = new HashMap<>(); + + public static final String COMMAND_STATEMENT_QUERY = "statement_query"; + public static final String COMMAND_STATEMENT_UPDATE = "statement_update"; + public static final String COMMAND_PREPARED_STATEMENT_QUERY = "prepared_statement_query"; + public static final String COMMAND_PREPARED_STATEMENT_UPDATE = "prepared_statement_update"; private static FlightInfo getFlightInfoExportedAndImportedKeys( final Message message, final FlightDescriptor descriptor) { @@ -311,6 +317,7 @@ public FlightInfo getFlightInfoStatement( final CommandStatementQuery commandStatementQuery, final CallContext callContext, final FlightDescriptor flightDescriptor) { + incrementCommandTypeCounter(COMMAND_STATEMENT_QUERY); final String query = commandStatementQuery.getQuery(); final Entry> queryInfo = Preconditions.checkNotNull( @@ -331,6 +338,7 @@ public FlightInfo getFlightInfoPreparedStatement( final CommandPreparedStatementQuery commandPreparedStatementQuery, final CallContext callContext, final FlightDescriptor flightDescriptor) { + incrementCommandTypeCounter(COMMAND_PREPARED_STATEMENT_QUERY); final ByteString preparedStatementHandle = commandPreparedStatementQuery.getPreparedStatementHandle(); @@ -398,6 +406,7 @@ public Runnable acceptPutStatement( final CallContext callContext, final FlightStream flightStream, final StreamListener streamListener) { + incrementCommandTypeCounter(COMMAND_STATEMENT_UPDATE); return () -> { final String query = commandStatementUpdate.getQuery(); final BiConsumer> resultProvider = @@ -471,6 +480,7 @@ public Runnable acceptPutPreparedStatementUpdate( final CallContext callContext, final FlightStream flightStream, final StreamListener streamListener) { + incrementCommandTypeCounter(COMMAND_PREPARED_STATEMENT_UPDATE); final ByteString handle = commandPreparedStatementUpdate.getPreparedStatementHandle(); final String query = Preconditions.checkNotNull( @@ -693,10 +703,22 @@ public void clearActionTypeCounter() { actionTypeCounter.clear(); } + public void clearCommandTypeCounter() { + commandTypeCounter.clear(); + } + public Map getActionTypeCounter() { return actionTypeCounter; } + public Map getCommandTypeCounter() { + return commandTypeCounter; + } + + private void incrementCommandTypeCounter(String commandType) { + commandTypeCounter.put(commandType, commandTypeCounter.getOrDefault(commandType, 0) + 1); + } + private void getStreamCatalogFunctions( final Message ticket, final ServerStreamListener serverStreamListener) { Preconditions.checkNotNull( From 97f9c05bc44d733ca99e4324996a8f27eb4a8551 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Fri, 10 Apr 2026 16:23:41 +0100 Subject: [PATCH 04/25] remove unrelated JDBC sample app --- .../example/ArrowFlightJdbcSampleApp.java | 120 ------------------ 1 file changed, 120 deletions(-) delete mode 100644 flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/example/ArrowFlightJdbcSampleApp.java diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/example/ArrowFlightJdbcSampleApp.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/example/ArrowFlightJdbcSampleApp.java deleted file mode 100644 index 10ce0bd285..0000000000 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/example/ArrowFlightJdbcSampleApp.java +++ /dev/null @@ -1,120 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.apache.arrow.driver.jdbc.example; - -import java.sql.Connection; -import java.sql.DriverManager; -import java.sql.ResultSet; -import java.sql.ResultSetMetaData; -import java.sql.SQLException; -import java.sql.Statement; -import java.util.Properties; - -/** - * Minimal sample app for using the Arrow Flight SQL JDBC driver. - * - *

Defaults are configured for a local Dremio instance: - * - *

    - *
  • host: {@code localhost} - *
  • port: {@code 32010} - *
  • user: {@code dremio} - *
  • password: {@code dremio123} - *
- * - *

Arguments are optional and positional: - * - * 

- *   [host] [port] [user] [password] [selectSql] [updateSql]
- *
- * - *

If {@code updateSql} is omitted, only {@code Statement.executeQuery(...)} is executed. - */ -public final class ArrowFlightJdbcSampleApp { - private static final String DEFAULT_HOST = "localhost"; - private static final int DEFAULT_PORT = 32010; - private static final String DEFAULT_USER = "dremio"; - private static final String DEFAULT_PASSWORD = "dremio123"; - private static final String DEFAULT_SELECT_SQL = "SELECT 1 AS sample_value"; - - private ArrowFlightJdbcSampleApp() {} - - public static void main(final String[] args) throws Exception { - final String host = getArg(args, 0, DEFAULT_HOST); - final int port = Integer.parseInt(getArg(args, 1, Integer.toString(DEFAULT_PORT))); - final String user = getArg(args, 2, DEFAULT_USER); - final String password = getArg(args, 3, DEFAULT_PASSWORD); - final String selectSql = getArg(args, 4, DEFAULT_SELECT_SQL); - final String updateSql = getArg(args, 5, ""); - - final String url = String.format("jdbc:arrow-flight-sql://%s:%d", host, port); - final Properties properties = new Properties(); - properties.setProperty("user", user); - properties.setProperty("password", password); - properties.setProperty("useEncryption", "false"); - - System.out.println("Connecting to " + url); - try (Connection connection = DriverManager.getConnection(url, properties); - Statement statement = connection.createStatement()) { - runSelect(statement, selectSql); - - if (updateSql.isEmpty()) { - System.out.println( - "Skipping Statement.executeUpdate(...) because no updateSql argument was provided."); - } else { - runUpdate(statement, updateSql); - } - } - } - - private static void runSelect(final Statement statement, final String selectSql) - throws SQLException { - System.out.println("Running Statement.executeQuery: " + selectSql); - try (ResultSet resultSet = statement.executeQuery(selectSql)) { - final ResultSetMetaData metadata = resultSet.getMetaData(); - final int columnCount = metadata.getColumnCount(); - int rowCount = 0; - while (resultSet.next()) { - rowCount++; - final StringBuilder rowBuilder = new StringBuilder(); - for (int i = 1; i <= columnCount; i++) { - if (i > 1) { - rowBuilder.append(", "); - } - rowBuilder.append(metadata.getColumnLabel(i)).append('=').append(resultSet.getObject(i)); - } - System.out.println("row " + rowCount + ": " + rowBuilder); - } - System.out.println("Statement.executeQuery returned " + rowCount + " row(s)"); - } - } - - private static void runUpdate(final Statement statement, final String updateSql) - throws SQLException { - System.out.println("Running Statement.executeUpdate: " + updateSql); - final int updateCount = statement.executeUpdate(updateSql); - System.out.println("Statement.executeUpdate affected " + updateCount + " row(s)"); - } - - private static String getArg(final String[] args, final int index, final String defaultValue) { - if (index >= args.length) { - return defaultValue; - } - final String arg = args[index]; - return arg == null || arg.isEmpty() ? defaultValue : arg; - } -} From 8d5dbbaed9db8873551c444f1996223ab7d4549a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Mon, 13 Apr 2026 14:27:36 +0100 Subject: [PATCH 05/25] relax JDBC cancellation test assertion --- .../org/apache/arrow/driver/jdbc/ResultSetTest.java | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java index 3a5a39be3d..7d9e38da92 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java @@ -302,14 +302,14 @@ public void testShouldCancelQueryUponCancelAfterQueryingResultSet() throws SQLEx @Test public void testShouldInterruptFlightStreamsIfQueryIsCancelledMidQuerying() throws SQLException, InterruptedException { + final String query = CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD; try (final Statement statement = connection.createStatement()) { final CountDownLatch latch = new CountDownLatch(1); final Set exceptions = synchronizedSet(new HashSet<>(1)); final Thread thread = new Thread( () -> { - try (final ResultSet resultSet = - statement.executeQuery(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD)) { + try (final ResultSet resultSet = statement.executeQuery(query)) { final int cachedColumnCount = resultSet.getMetaData().getColumnCount(); Thread.sleep(300); while (resultSet.next()) { @@ -332,7 +332,14 @@ public void testShouldInterruptFlightStreamsIfQueryIsCancelledMidQuerying() .reduce(StringBuilder::append) .orElseThrow(IllegalArgumentException::new) .toString(), - is("Statement canceled")); + anyOf( + is("Statement canceled"), + allOf( + containsString(format("Error while executing SQL \"%s\"", query)), + anyOf( + containsString("Query canceled"), + containsString("CANCELLED"), + containsString("Cancelling"))))); } } From d63674291e9893a1bb68e40c5d3b37ad5e7b1bbf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Fri, 17 Apr 2026 10:59:59 +0100 Subject: [PATCH 06/25] fix concurrent close --- .../arrow/driver/jdbc/utils/FlightEndpointDataQueue.java | 3 ++- .../arrow/driver/jdbc/utils/FlightEndpointDataQueueTest.java | 5 ++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/FlightEndpointDataQueue.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/FlightEndpointDataQueue.java index 7c1c569fda..67784890bb 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/FlightEndpointDataQueue.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/FlightEndpointDataQueue.java @@ -97,7 +97,8 @@ interface EndpointStreamSupplier { private CloseableEndpointStreamPair next(final EndpointStreamSupplier endpointStreamSupplier) throws SQLException { - checkOpen(); + // No checkOpen here: a concurrent close() is an end-of-stream signal for readers, + // not an error. Cancelled in-flight streams still surface via future.get() below. while (!futures.isEmpty()) { final Future future = endpointStreamSupplier.get(); futures.remove(future); diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/FlightEndpointDataQueueTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/FlightEndpointDataQueueTest.java index f40610a032..a5d9afbd16 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/FlightEndpointDataQueueTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/FlightEndpointDataQueueTest.java @@ -48,10 +48,9 @@ public void testNextShouldRetrieveNullIfEmpty() throws Exception { } @Test - public void testNextShouldThrowExceptionUponClose() throws Exception { + public void testNextShouldReturnNullUponClose() throws Exception { queue.close(); - ThrowableAssertionUtils.simpleAssertThrowableClass( - IllegalStateException.class, () -> queue.next()); + assertThat(queue.next(), is(nullValue())); } @Test From d0e54bc24fc0c2fd111ed6e164781fbdc63b5d08 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Fri, 17 Apr 2026 11:02:14 +0100 Subject: [PATCH 07/25] Revert "relax JDBC cancellation test assertion" This reverts commit b57550fa15c95a4011686f5d1addbba2186611d9. --- .../org/apache/arrow/driver/jdbc/ResultSetTest.java | 13 +++---------- 1 file changed, 3 insertions(+), 10 deletions(-) diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java index 7d9e38da92..3a5a39be3d 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java @@ -302,14 +302,14 @@ public void testShouldCancelQueryUponCancelAfterQueryingResultSet() throws SQLEx @Test public void testShouldInterruptFlightStreamsIfQueryIsCancelledMidQuerying() throws SQLException, InterruptedException { - final String query = CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD; try (final Statement statement = connection.createStatement()) { final CountDownLatch latch = new CountDownLatch(1); final Set exceptions = synchronizedSet(new HashSet<>(1)); final Thread thread = new Thread( () -> { - try (final ResultSet resultSet = statement.executeQuery(query)) { + try (final ResultSet resultSet = + statement.executeQuery(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD)) { final int cachedColumnCount = resultSet.getMetaData().getColumnCount(); Thread.sleep(300); while (resultSet.next()) { @@ -332,14 +332,7 @@ public void testShouldInterruptFlightStreamsIfQueryIsCancelledMidQuerying() .reduce(StringBuilder::append) .orElseThrow(IllegalArgumentException::new) .toString(), - anyOf( - is("Statement canceled"), - allOf( - containsString(format("Error while executing SQL \"%s\"", query)), - anyOf( - containsString("Query canceled"), - containsString("CANCELLED"), - containsString("Cancelling"))))); + is("Statement canceled")); } } From a2f054d6530bb2809a8a227c54209c126ffb2da0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Fri, 17 Apr 2026 16:22:56 +0100 Subject: [PATCH 08/25] fixing missing concurrent close --- .../jdbc/ArrowFlightJdbcFlightStreamResultSet.java | 8 ++++++++ .../driver/jdbc/utils/FlightEndpointDataQueue.java | 12 +++++++++++- .../jdbc/utils/FlightEndpointDataQueueTest.java | 8 +++++--- 3 files changed, 24 insertions(+), 4 deletions(-) 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 d383d239d1..e0c500e074 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 @@ -32,6 +32,7 @@ import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.calcite.avatica.AvaticaConnection; import org.apache.calcite.avatica.AvaticaResultSet; import org.apache.calcite.avatica.AvaticaResultSetMetaData; import org.apache.calcite.avatica.AvaticaStatement; @@ -209,6 +210,13 @@ public boolean next() throws SQLException { continue; } + // No more data. If the queue was closed concurrently (e.g. statement.cancel() + // racing with the reader past super.next()), surface as "Statement canceled" + // to match Avatica's cancellation semantics. + if (flightEndpointDataQueue.isClosed()) { + throw AvaticaConnection.HELPER.createException("Statement canceled"); + } + if (statement != null && statement.isCloseOnCompletion()) { statement.close(); } diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/FlightEndpointDataQueue.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/FlightEndpointDataQueue.java index 67784890bb..5ad2e17f80 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/FlightEndpointDataQueue.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/FlightEndpointDataQueue.java @@ -177,7 +177,17 @@ public void enqueue(final Collection endpointReques /** Adds given {@link FlightStream} to the queue. */ public synchronized void enqueue(final CloseableEndpointStreamPair endpointRequest) { checkNotNull(endpointRequest); - checkOpen(); + if (isClosed()) { + // Concurrent close() raced with enqueue (e.g. statement.cancel() while the reader + // was mid-loop). Treat as end-of-stream: close the endpoint and return silently. + // The reader detects cancellation via a subsequent isClosed() check. + try { + endpointRequest.close(); + } catch (final Exception e) { + LOGGER.error("Failed to close endpoint after queue was closed.", e); + } + return; + } endpointsToClose.add(endpointRequest); futures.add( completionService.submit( diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/FlightEndpointDataQueueTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/FlightEndpointDataQueueTest.java index a5d9afbd16..8dd39796e8 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/FlightEndpointDataQueueTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/FlightEndpointDataQueueTest.java @@ -21,6 +21,7 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; import java.util.concurrent.CompletionService; import org.apache.arrow.driver.jdbc.client.CloseableEndpointStreamPair; @@ -54,10 +55,11 @@ public void testNextShouldReturnNullUponClose() throws Exception { } @Test - public void testEnqueueShouldThrowExceptionUponClose() throws Exception { + public void testEnqueueShouldCloseEndpointSilentlyAfterClose() throws Exception { queue.close(); - ThrowableAssertionUtils.simpleAssertThrowableClass( - IllegalStateException.class, () -> queue.enqueue(mock(CloseableEndpointStreamPair.class))); + final CloseableEndpointStreamPair endpoint = mock(CloseableEndpointStreamPair.class); + assertDoesNotThrow(() -> queue.enqueue(endpoint)); + verify(endpoint).close(); } @Test From b594b90abd3cb3b202ed5e89381ba54608a1ecbf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Sun, 19 Apr 2026 12:25:18 +0100 Subject: [PATCH 09/25] throw Statement cancelled on queued streams --- .../ArrowFlightJdbcFlightStreamResultSet.java | 28 +++++++++++++------ 1 file changed, 20 insertions(+), 8 deletions(-) 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 e0c500e074..c39a6c6b6b 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 @@ -27,7 +27,9 @@ import org.apache.arrow.driver.jdbc.client.CloseableEndpointStreamPair; import org.apache.arrow.driver.jdbc.utils.FlightEndpointDataQueue; import org.apache.arrow.driver.jdbc.utils.VectorSchemaRootTransformer; +import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.FlightInfo; +import org.apache.arrow.flight.FlightRuntimeException; import org.apache.arrow.flight.FlightStream; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.vector.VectorSchemaRoot; @@ -193,18 +195,28 @@ public boolean next() throws SQLException { return true; } - if (currentEndpointData != null) { - currentEndpointData.getStream().getRoot().clear(); - if (currentEndpointData.getStream().next()) { - populateDataForCurrentFlightStream(); - continue; + try { + if (currentEndpointData != null) { + currentEndpointData.getStream().getRoot().clear(); + if (currentEndpointData.getStream().next()) { + populateDataForCurrentFlightStream(); + continue; + } + + flightEndpointDataQueue.enqueue(currentEndpointData); } - flightEndpointDataQueue.enqueue(currentEndpointData); + currentEndpointData = getNextEndpointStream(false); + } catch (final FlightRuntimeException e) { + // A concurrent statement.cancel() (or close) cancels in-flight FlightStreams, + // which surface here as CANCELLED. Normalize to Avatica's "Statement canceled". + if (flightEndpointDataQueue.isClosed() + && e.status().code() == CallStatus.CANCELLED.code()) { + throw AvaticaConnection.HELPER.createException("Statement canceled"); + } + throw e; } - currentEndpointData = getNextEndpointStream(false); - if (currentEndpointData != null) { populateDataForCurrentFlightStream(); continue; From dd5a5340766081c15abc3aff7039e6923b48b952 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Tue, 22 Sep 2026 23:18:04 +0100 Subject: [PATCH 10/25] Preserve Flight info statement API --- .../driver/jdbc/ArrowFlightInfoStatement.java | 36 +++++++++++++++++++ .../driver/jdbc/ArrowFlightMetaStatement.java | 3 +- .../ArrowFlightPreparedStatementTest.java | 11 ++++++ .../jdbc/ArrowFlightStatementExecuteTest.java | 7 ++++ 4 files changed, 56 insertions(+), 1 deletion(-) create mode 100644 flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightInfoStatement.java diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightInfoStatement.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightInfoStatement.java new file mode 100644 index 0000000000..37ee93722a --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightInfoStatement.java @@ -0,0 +1,36 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.driver.jdbc; + +import java.sql.SQLException; +import java.sql.Statement; +import org.apache.arrow.flight.FlightInfo; + +/** A {@link Statement} that deals with {@link FlightInfo}. */ +public interface ArrowFlightInfoStatement extends Statement { + + @Override + ArrowFlightConnection getConnection() throws SQLException; + + /** + * Executes the query this {@link Statement} is holding. + * + * @return the {@link FlightInfo} for the results. + * @throws SQLException on error. + */ + FlightInfo executeFlightInfoQuery() throws SQLException; +} diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaStatement.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaStatement.java index 415af19e8f..474fa03681 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaStatement.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaStatement.java @@ -27,11 +27,12 @@ import org.apache.calcite.avatica.remote.TypedValue; /** Statement capabilities used by {@link ArrowFlightMetaImpl}. */ -interface ArrowFlightMetaStatement extends Statement { +interface ArrowFlightMetaStatement extends ArrowFlightInfoStatement { @Override ArrowFlightConnection getConnection() throws SQLException; + @Override FlightInfo executeFlightInfoQuery() throws SQLException; /** 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 e8bfd4025b..3fbfedcbc6 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 @@ -81,6 +81,17 @@ public void before() { PRODUCER.clearActionTypeCounter(); } + @Test + public void testPreparedStatementKeepsFlightInfoStatementApi() throws SQLException { + try (final PreparedStatement preparedStatement = + connection.prepareStatement(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD)) { + final ArrowFlightInfoStatement flightInfoStatement = + (ArrowFlightInfoStatement) preparedStatement; + + assertSame(preparedStatement, flightInfoStatement); + } + } + @Test public void testSimpleQueryNoParameterBinding() throws SQLException { final String query = CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD; 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 f6f05c2bd0..2a70a0bb21 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 @@ -145,6 +145,13 @@ public static void tearDownAfterClass() throws Exception { AutoCloseables.close(PRODUCER); } + @Test + public void testStatementKeepsFlightInfoStatementApi() { + final ArrowFlightInfoStatement flightInfoStatement = (ArrowFlightInfoStatement) statement; + + assertSame(statement, flightInfoStatement); + } + @Test public void testExecuteShouldRunSelectQuery() throws SQLException { assertThat(statement.execute(SAMPLE_QUERY_CMD), is(true)); // Means this is a SELECT query. From e99e4b7de959a8767e4d81c4280a86db37a5d6a3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Thu, 24 Sep 2026 09:31:49 +0100 Subject: [PATCH 11/25] preserve server statement classification --- .../driver/jdbc/ArrowFlightMetaImpl.java | 22 +++++++++++++++---- .../jdbc/ArrowFlightPreparedStatement.java | 5 ++++- 2 files changed, 22 insertions(+), 5 deletions(-) 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 4da182ca63..9e6afd929e 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 @@ -205,10 +205,24 @@ public static Signature buildSignature(final String sql, final StatementType typ /** Builds an Avatica signature from Arrow result and parameter schemas. */ public static Signature buildSignature( final String sql, final Schema resultSetSchema, final Schema parameterSchema) { - StatementType statementType = - resultSetSchema == null || resultSetSchema.getFields().isEmpty() - ? StatementType.IS_DML - : StatementType.SELECT; + return buildSignature(sql, resultSetSchema, parameterSchema, (Boolean) null); + } + + /** Builds an Avatica signature, honoring the server-provided statement classification. */ + public static Signature buildSignature( + final String sql, + final Schema resultSetSchema, + final Schema parameterSchema, + final Boolean isUpdate) { + final StatementType statementType; + if (isUpdate != null) { + statementType = isUpdate ? StatementType.IS_DML : StatementType.SELECT; + } else { + statementType = + resultSetSchema == null || resultSetSchema.getFields().isEmpty() + ? StatementType.IS_DML + : StatementType.SELECT; + } return buildSignature(sql, resultSetSchema, parameterSchema, statementType); } diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java index bd7ebbe0e4..a3ae1829c9 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java @@ -251,7 +251,10 @@ ArrowFlightPreparedStatement build() throws SQLException { connection.getClientHandler().prepare(query); final Signature signature = ArrowFlightMetaImpl.buildSignature( - query, preparedStatement.getDataSetSchema(), preparedStatement.getParameterSchema()); + query, + preparedStatement.getDataSetSchema(), + preparedStatement.getParameterSchema(), + preparedStatement.isUpdate()); return new ArrowFlightPreparedStatement( connection, From fb4ff48a4880d7aa7d3ac00bf19a1a0cf412a19b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Wed, 23 Sep 2026 00:03:24 +0100 Subject: [PATCH 12/25] Close public statements with their connection --- .../driver/jdbc/ArrowFlightConnection.java | 34 +++++++++++++++++-- .../jdbc/ArrowFlightPreparedStatement.java | 27 ++++++++++----- .../driver/jdbc/ArrowFlightStatement.java | 10 ++++++ .../jdbc/ArrowFlightStatementExecuteTest.java | 22 ++++++++++++ 4 files changed, 82 insertions(+), 11 deletions(-) diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java index 107cfa0c2f..d7318b27d5 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java @@ -23,9 +23,13 @@ import java.sql.ResultSet; import java.sql.SQLException; import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; +import java.util.HashSet; +import java.util.IdentityHashMap; import java.util.Map; import java.util.Properties; +import java.util.Set; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler; @@ -37,6 +41,7 @@ import org.apache.arrow.util.Preconditions; import org.apache.calcite.avatica.AvaticaConnection; import org.apache.calcite.avatica.AvaticaFactory; +import org.apache.calcite.avatica.AvaticaStatement; import org.apache.calcite.avatica.DriverVersion; /** Connection to the Arrow Flight server. */ @@ -48,6 +53,8 @@ public final class ArrowFlightConnection extends AvaticaConnection { private ExecutorService executorService; private int metadataResultSetCount; private Map metadataResultSetMap = new HashMap<>(); + private final Set statementOwners = + Collections.newSetFromMap(new IdentityHashMap<>()); /** * Creates a new {@link ArrowFlightConnection}. @@ -144,12 +151,13 @@ private static ArrowFlightSqlClientHandler createNewClientHandler( void reset() throws SQLException { // Clean up any open Statements try { - AutoCloseables.close(statementMap.values()); + AutoCloseables.close(getStatementsToClose()); } catch (final Exception e) { throw AvaticaConnection.HELPER.createException(e.getMessage(), e); } statementMap.clear(); + statementOwners.clear(); // Reset Holdability this.setHoldability(this.metaData.getResultSetHoldability()); @@ -205,6 +213,28 @@ void onResultSetClose(Integer id) { metadataResultSetMap.remove(id); } + synchronized void registerStatementOwner(final AvaticaStatement statement) { + statementOwners.add(statement); + } + + synchronized void unregisterStatementOwner(final AvaticaStatement statement) { + statementOwners.remove(statement); + } + + private synchronized ArrayList getStatementsToClose() { + final ArrayList statements = new ArrayList<>(statementOwners); + final Set ownedHandles = new HashSet<>(); + for (AvaticaStatement statement : statementOwners) { + ownedHandles.add(statement.handle.id); + } + for (Map.Entry entry : statementMap.entrySet()) { + if (!ownedHandles.contains(entry.getKey())) { + statements.add(entry.getValue()); + } + } + return statements; + } + @Override public Properties getClientInfo() { final Properties copy = new Properties(); @@ -223,7 +253,7 @@ public void close() throws SQLException { topLevelException = e; } // copies of the collections are used to avoid concurrent modification problems - ArrayList closeables = new ArrayList<>(statementMap.values()); + ArrayList closeables = getStatementsToClose(); closeables.addAll(new ArrayList<>(metadataResultSetMap.values())); closeables.add(clientHandler); closeables.addAll(allocator.getChildAllocators()); diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java index a3ae1829c9..2c64b8e7ab 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java @@ -102,7 +102,11 @@ Schema getDataSetSchema() { @Override public synchronized void close() throws SQLException { - super.close(); + try { + super.close(); + } finally { + ((ArrowFlightConnection) connection).unregisterStatementOwner(this); + } } void closePreparedResources() { @@ -256,14 +260,19 @@ ArrowFlightPreparedStatement build() throws SQLException { preparedStatement.getParameterSchema(), preparedStatement.isUpdate()); - return new ArrowFlightPreparedStatement( - connection, - generateHandle ? null : handle, - signature, - preparedStatement, - resultSetType, - resultSetConcurrency, - resultSetHoldability); + final ArrowFlightPreparedStatement statement = + new ArrowFlightPreparedStatement( + connection, + generateHandle ? null : handle, + signature, + preparedStatement, + resultSetType, + resultSetConcurrency, + resultSetHoldability); + if (generateHandle) { + connection.registerStatementOwner(statement); + } + return statement; } } } diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightStatement.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightStatement.java index 5bbdb8ff87..a6122c7d72 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightStatement.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightStatement.java @@ -42,6 +42,16 @@ public class ArrowFlightStatement extends AvaticaStatement implements ArrowFligh final int resultSetConcurrency, final int resultSetHoldability) { super(connection, handle, resultSetType, resultSetConcurrency, resultSetHoldability); + connection.registerStatementOwner(this); + } + + @Override + public synchronized void close() throws SQLException { + try { + super.close(); + } finally { + ((ArrowFlightConnection) connection).unregisterStatementOwner(this); + } } @Override 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 2a70a0bb21..75f8efc9f7 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 @@ -189,6 +189,28 @@ public void testExecuteReplacesStatementMapEntryWithPreparedStatement() throws S assertThat(preparedStatement, instanceOf(ArrowFlightPreparedStatement.class)); } + @Test + public void testConnectionCloseClosesStatementAndResultSetAfterExecute() throws SQLException { + PRODUCER.clearActionTypeCounter(); + final ArrowFlightStatement arrowStatement = (ArrowFlightStatement) statement; + final ArrowFlightConnection arrowConnection = (ArrowFlightConnection) connection; + + assertThat(statement.execute(SAMPLE_QUERY_CMD), is(true)); + final ResultSet resultSet = statement.getResultSet(); + assertNotNull(resultSet); + + connection.close(); + + assertThat(statement.isClosed(), is(true)); + assertThat(resultSet.isClosed(), is(true)); + assertThat(arrowConnection.statementMap.containsKey(arrowStatement.handle.id), is(false)); + assertThat( + PRODUCER + .getActionTypeCounter() + .getOrDefault(FlightSqlUtils.FLIGHT_SQL_CLOSE_PREPARED_STATEMENT.getType(), 0), + is(1)); + } + @Test public void testExecuteQueryRestoresStatementMapEntryWithStatement() throws SQLException { final ArrowFlightStatement arrowStatement = (ArrowFlightStatement) statement; From a9bb7fa791379fd65bbe206dc6f4de4b1c377f4b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Wed, 23 Sep 2026 00:22:16 +0100 Subject: [PATCH 13/25] Map prepare failures to SQLException --- .../arrow/driver/jdbc/ArrowFlightMetaImpl.java | 11 ++++------- .../jdbc/client/ArrowFlightSqlClientHandler.java | 11 ++++++++--- .../jdbc/ArrowFlightPreparedStatementTest.java | 16 ++++++++++++++++ .../jdbc/ArrowFlightStatementExecuteTest.java | 10 ++++++++++ 4 files changed, 38 insertions(+), 10 deletions(-) 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 9e6afd929e..6c13f388c9 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 @@ -19,7 +19,6 @@ import java.sql.Connection; import java.sql.ResultSet; import java.sql.SQLException; -import java.sql.SQLTimeoutException; import java.util.ArrayList; import java.util.Collections; import java.util.List; @@ -130,13 +129,11 @@ public ExecuteResult prepareAndExecute( // This is the Avatica entry point used by Statement.execute(String). return getMetaStatement(handle) .prepareAndExecute(query, maxRowCount, maxRowsInFirstFrame, callback); - } catch (SQLTimeoutException e) { - // So far AvaticaStatement(executeInternal) only handles NoSuchStatement and - // Runtime - // Exceptions. - throw new RuntimeException(e); } catch (SQLException e) { - throw new NoSuchStatementException(handle); + // AvaticaStatement.executeInternal handles RuntimeException and preserves it as the cause of + // the public SQLException. NoSuchStatementException would incorrectly trigger a statement + // reset/retry and discard the actual server failure. + throw new RuntimeException(e.getMessage(), e); } } 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 09b8dfecee..5f49907bea 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 @@ -457,10 +457,15 @@ public void setCatalog(final String catalog) throws SQLException { * * @param query the SQL query. * @return a new prepared statement. + * @throws SQLException if the server rejects the prepare request. */ - public PreparedStatement prepare(final String query) { - final FlightSqlClient.PreparedStatement preparedStatement = - sqlClient.prepare(query, getOptions()); + public PreparedStatement prepare(final String query) throws SQLException { + final FlightSqlClient.PreparedStatement preparedStatement; + try { + preparedStatement = sqlClient.prepare(query, getOptions()); + } catch (final FlightRuntimeException e) { + throw new SQLException(e.getMessage(), e); + } return new PreparedStatement() { @Override public FlightInfo executeQuery() throws SQLException { 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 3fbfedcbc6..2ca485729d 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 @@ -17,12 +17,14 @@ package org.apache.arrow.driver.jdbc; import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.CoreMatchers.instanceOf; 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.assertNotNull; import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import java.nio.charset.StandardCharsets; @@ -37,6 +39,7 @@ import java.util.List; import org.apache.arrow.driver.jdbc.utils.CoreMockedSqlProducers; import org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer; +import org.apache.arrow.flight.FlightRuntimeException; import org.apache.arrow.flight.sql.FlightSqlUtils; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; @@ -145,6 +148,9 @@ public void testSimpleQueryNoParameterBindingWithExecuteV2() throws SQLException assertFalse(resultSet.next()); assertFalse(preparedStatement.getMoreResults()); assertEquals(-1, preparedStatement.getUpdateCount()); + } + } + @Test public void testPrepareStatementRegistersCreatedStatementByGeneratedHandle() throws SQLException { final String query = CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD; @@ -161,6 +167,16 @@ public void testPrepareStatementRegistersCreatedStatementByGeneratedHandle() thr } } + @Test + public void testPrepareFailureIsReportedAsSQLException() { + final SQLException exception = + assertThrows( + SQLException.class, + () -> connection.prepareStatement("SELECT * FROM unregistered_table")); + + assertThat(exception.getCause(), instanceOf(FlightRuntimeException.class)); + } + @Test public void testQueryWithParameterBinding() throws SQLException { final String query = "Fake query 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 75f8efc9f7..30b938733f 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 @@ -17,6 +17,7 @@ package org.apache.arrow.driver.jdbc; import static org.hamcrest.CoreMatchers.allOf; +import static org.hamcrest.CoreMatchers.containsString; import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.instanceOf; import static org.hamcrest.CoreMatchers.is; @@ -189,6 +190,15 @@ public void testExecuteReplacesStatementMapEntryWithPreparedStatement() throws S assertThat(preparedStatement, instanceOf(ArrowFlightPreparedStatement.class)); } + @Test + public void testExecutePrepareFailurePreservesServerError() { + final SQLException exception = + assertThrows( + SQLException.class, () -> statement.execute("SELECT * FROM unregistered_table")); + + assertThat(exception.getMessage(), containsString("Query not found")); + } + @Test public void testConnectionCloseClosesStatementAndResultSetAfterExecute() throws SQLException { PRODUCER.clearActionTypeCounter(); From 59c62987b13f4f626256d10e9796e03d0da9470b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Thu, 24 Sep 2026 09:39:38 +0100 Subject: [PATCH 14/25] Tighten statement protocol regression coverage --- .../jdbc/ArrowFlightStatementExecuteTest.java | 2 + .../ArrowFlightStatementProtocolTest.java | 79 +++++++++++++++++++ .../jdbc/utils/MockFlightSqlProducer.java | 14 ++-- 3 files changed, 89 insertions(+), 6 deletions(-) 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 30b938733f..dd5e4ecbea 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 @@ -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; diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementProtocolTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementProtocolTest.java index c5c35a9173..1ac31d3528 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementProtocolTest.java +++ b/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 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 dfa536f230..88355acdfb 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 @@ -407,8 +407,14 @@ public Runnable acceptPutStatement( final FlightStream flightStream, final StreamListener streamListener) { incrementCommandTypeCounter(COMMAND_STATEMENT_UPDATE); + return acceptUpdate(commandStatementUpdate.getQuery(), flightStream, streamListener); + } + + private Runnable acceptUpdate( + final String query, + final FlightStream flightStream, + final StreamListener streamListener) { return () -> { - final String query = commandStatementUpdate.getQuery(); final BiConsumer> 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 From 1f153dfb60882b35404b690fa683ee3c95d80b75 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Thu, 24 Sep 2026 09:41:02 +0100 Subject: [PATCH 15/25] Keep signature builders internal --- .../org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) 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 6c13f388c9..3bcd7da077 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 @@ -191,22 +191,22 @@ private ArrowFlightMetaStatement getMetaStatement(StatementHandle statementHandl throw new IllegalStateException("Statement not found: " + statementHandle); } - public static Signature buildDefaultSignature() { + static Signature buildDefaultSignature() { return buildSignature(null, StatementType.SELECT); } - public static Signature buildSignature(final String sql, final StatementType type) { + static Signature buildSignature(final String sql, final StatementType type) { return buildSignature(sql, null, null, type); } /** Builds an Avatica signature from Arrow result and parameter schemas. */ - public static Signature buildSignature( + static Signature buildSignature( final String sql, final Schema resultSetSchema, final Schema parameterSchema) { return buildSignature(sql, resultSetSchema, parameterSchema, (Boolean) null); } /** Builds an Avatica signature, honoring the server-provided statement classification. */ - public static Signature buildSignature( + static Signature buildSignature( final String sql, final Schema resultSetSchema, final Schema parameterSchema, From c184107077b0518da7e358709a929262d21d12fb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Thu, 24 Sep 2026 10:06:49 +0100 Subject: [PATCH 16/25] Clean up implementation comments --- .../org/apache/arrow/driver/jdbc/ArrowFlightConnection.java | 1 + .../java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java | 2 -- .../apache/arrow/driver/jdbc/utils/FlightEndpointDataQueue.java | 2 -- 3 files changed, 1 insertion(+), 4 deletions(-) diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java index d7318b27d5..5f0b04ec41 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java @@ -52,6 +52,7 @@ public final class ArrowFlightConnection extends AvaticaConnection { private final ArrowFlightConnectionConfigImpl config; private ExecutorService executorService; private int metadataResultSetCount; + // Track result sets and statements owned by this connection so close() can release them. private Map metadataResultSetMap = new HashMap<>(); private final Set statementOwners = Collections.newSetFromMap(new IdentityHashMap<>()); 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 3bcd7da077..f2b97b6044 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 @@ -112,7 +112,6 @@ public ExecuteResult prepareAndExecute( final long maxRowCount, final PrepareCallback prepareCallback) throws NoSuchStatementException { - // This is the Avatica entry point used by Statement.execute(String). return prepareAndExecute( statementHandle, query, maxRowCount, -1 /* Not used */, prepareCallback); } @@ -126,7 +125,6 @@ public ExecuteResult prepareAndExecute( final PrepareCallback callback) throws NoSuchStatementException { try { - // This is the Avatica entry point used by Statement.execute(String). return getMetaStatement(handle) .prepareAndExecute(query, maxRowCount, maxRowsInFirstFrame, callback); } catch (SQLException e) { diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/FlightEndpointDataQueue.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/FlightEndpointDataQueue.java index 5ad2e17f80..c6f485ee8b 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/FlightEndpointDataQueue.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/FlightEndpointDataQueue.java @@ -97,8 +97,6 @@ interface EndpointStreamSupplier { private CloseableEndpointStreamPair next(final EndpointStreamSupplier endpointStreamSupplier) throws SQLException { - // No checkOpen here: a concurrent close() is an end-of-stream signal for readers, - // not an error. Cancelled in-flight streams still surface via future.get() below. while (!futures.isEmpty()) { final Future future = endpointStreamSupplier.get(); futures.remove(future); From fee4892e79e9c9c54e0bc58f40cc31cc0a21dc8d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Thu, 24 Sep 2026 10:55:17 +0100 Subject: [PATCH 17/25] Preserve JDBC exception contracts Restore source-compatible exception declarations and map prepared-statement transition failures to SQLException. --- .../driver/jdbc/ArrowFlightJdbcFactory.java | 3 +- .../jdbc/ArrowFlightPreparedStatement.java | 9 +++- .../driver/jdbc/ArrowFlightStatement.java | 4 +- .../client/ArrowFlightSqlClientHandler.java | 11 ++--- .../jdbc/ArrowFlightStatementExecuteTest.java | 45 +++++++++++++++++++ 5 files changed, 59 insertions(+), 13 deletions(-) diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFactory.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFactory.java index 202b491f5c..10c361866e 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFactory.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFactory.java @@ -78,7 +78,8 @@ public ArrowFlightPreparedStatement newPreparedStatement( final Meta.Signature signature, final int resultType, final int resultSetConcurrency, - final int resultSetHoldability) { + final int resultSetHoldability) + throws SQLException { final ArrowFlightConnection flightConnection = (ArrowFlightConnection) connection; final AvaticaStatement existingStatement = flightConnection.statementMap.get(statementHandle.id); diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java index 2c64b8e7ab..2f6d0fe44d 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java @@ -22,6 +22,7 @@ import org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler; import org.apache.arrow.driver.jdbc.utils.AvaticaParameterBinder; import org.apache.arrow.flight.FlightInfo; +import org.apache.arrow.flight.FlightRuntimeException; import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.calcite.avatica.AvaticaPreparedStatement; @@ -251,8 +252,12 @@ ArrowFlightPreparedStatement build() throws SQLException { throw new IllegalStateException("PreparedStatement builder requires a handle."); } - final ArrowFlightSqlClientHandler.PreparedStatement preparedStatement = - connection.getClientHandler().prepare(query); + final ArrowFlightSqlClientHandler.PreparedStatement preparedStatement; + try { + preparedStatement = connection.getClientHandler().prepare(query); + } catch (final FlightRuntimeException e) { + throw new SQLException(e.getMessage(), e); + } final Signature signature = ArrowFlightMetaImpl.buildSignature( query, diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightStatement.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightStatement.java index a6122c7d72..104abdb9b4 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightStatement.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightStatement.java @@ -81,8 +81,8 @@ public ExecuteResult prepareAndExecute( public ResultSet executeQuery(final String sql) throws SQLException { checkOpen(); updateCount = -1; - switchToDirectStatementMode(); try { + switchToDirectStatementMode(); final Meta.Signature signature = ArrowFlightMetaImpl.buildSignature(sql, StatementType.SELECT); setSignature(signature); @@ -97,9 +97,9 @@ public long executeLargeUpdate(final String sql) throws SQLException { checkOpen(); clearOpenResultSet(); updateCount = -1; - switchToDirectStatementMode(); try { + switchToDirectStatementMode(); final long updatedCount = getConnection().getClientHandler().executeUpdate(sql); setSignature(ArrowFlightMetaImpl.buildSignature(sql, StatementType.IS_DML)); updateCount = updatedCount; 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 5f49907bea..09b8dfecee 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 @@ -457,15 +457,10 @@ public void setCatalog(final String catalog) throws SQLException { * * @param query the SQL query. * @return a new prepared statement. - * @throws SQLException if the server rejects the prepare request. */ - public PreparedStatement prepare(final String query) throws SQLException { - final FlightSqlClient.PreparedStatement preparedStatement; - try { - preparedStatement = sqlClient.prepare(query, getOptions()); - } catch (final FlightRuntimeException e) { - throw new SQLException(e.getMessage(), e); - } + public PreparedStatement prepare(final String query) { + final FlightSqlClient.PreparedStatement preparedStatement = + sqlClient.prepare(query, getOptions()); return new PreparedStatement() { @Override public FlightInfo executeQuery() throws SQLException { 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 dd5e4ecbea..abecd5ea83 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 @@ -27,6 +27,8 @@ 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 static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; import java.sql.Connection; import java.sql.ResultSet; @@ -38,6 +40,8 @@ import java.util.stream.Collectors; import java.util.stream.IntStream; import org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer; +import org.apache.arrow.flight.CallStatus; +import org.apache.arrow.flight.FlightRuntimeException; import org.apache.arrow.flight.sql.FlightSqlUtils; import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.memory.RootAllocator; @@ -237,6 +241,47 @@ public void testExecuteQueryRestoresStatementMapEntryWithStatement() throws SQLE assertSame(arrowStatement, arrowConnection.statementMap.get(arrowStatement.handle.id)); } + @Test + public void testExecuteQueryMapsPreparedStatementCloseFailureToSQLException() { + final ArrowFlightStatement arrowStatement = (ArrowFlightStatement) statement; + final ArrowFlightConnection arrowConnection = (ArrowFlightConnection) connection; + final ArrowFlightPreparedStatement preparedStatement = mock(ArrowFlightPreparedStatement.class); + final FlightRuntimeException closeFailure = + CallStatus.INTERNAL.withDescription("Prepared statement close failed").toRuntimeException(); + doThrow(closeFailure).when(preparedStatement).closeStatement(); + arrowConnection.statementMap.put(arrowStatement.handle.id, preparedStatement); + + try { + final SQLException exception = + assertThrows(SQLException.class, () -> statement.executeQuery(SAMPLE_QUERY_CMD)); + + assertSame(closeFailure, exception.getCause()); + } finally { + arrowConnection.statementMap.put(arrowStatement.handle.id, arrowStatement); + } + } + + @Test + public void testExecuteLargeUpdateMapsPreparedStatementCloseFailureToSQLException() { + final ArrowFlightStatement arrowStatement = (ArrowFlightStatement) statement; + final ArrowFlightConnection arrowConnection = (ArrowFlightConnection) connection; + final ArrowFlightPreparedStatement preparedStatement = mock(ArrowFlightPreparedStatement.class); + final FlightRuntimeException closeFailure = + CallStatus.INTERNAL.withDescription("Prepared statement close failed").toRuntimeException(); + doThrow(closeFailure).when(preparedStatement).closeStatement(); + arrowConnection.statementMap.put(arrowStatement.handle.id, preparedStatement); + + try { + final SQLException exception = + assertThrows( + SQLException.class, () -> statement.executeLargeUpdate(SAMPLE_LARGE_UPDATE_QUERY)); + + assertSame(closeFailure, exception.getCause()); + } finally { + arrowConnection.statementMap.put(arrowStatement.handle.id, arrowStatement); + } + } + @Test public void testExecuteShouldRunUpdateQueryForSmallUpdate() throws SQLException { assertThat(statement.execute(SAMPLE_UPDATE_QUERY), is(false)); // Means this is an UPDATE query. From 73c08e36d36f3cfdb9f210d738727eeeecbbd009 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Thu, 24 Sep 2026 11:28:51 +0100 Subject: [PATCH 18/25] Keep statement orchestration internal Dispatch Avatica operations through package-private hooks and guard the public statement API surface. --- .../ArrowFlightJdbcFlightStreamResultSet.java | 2 +- .../driver/jdbc/ArrowFlightMetaImpl.java | 38 ++++++-- .../driver/jdbc/ArrowFlightMetaStatement.java | 61 ------------- .../jdbc/ArrowFlightPreparedStatement.java | 28 +----- .../driver/jdbc/ArrowFlightStatement.java | 11 +-- .../jdbc/ArrowFlightStatementApiTest.java | 88 +++++++++++++++++++ .../jdbc/ArrowFlightStatementExecuteTest.java | 4 +- 7 files changed, 128 insertions(+), 104 deletions(-) delete mode 100644 flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaStatement.java create mode 100644 flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementApiTest.java 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 c39a6c6b6b..74c53104db 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 @@ -70,7 +70,7 @@ public final class ArrowFlightJdbcFlightStreamResultSet throws SQLException { super(statement, state, signature, resultSetMetaData, timeZone, firstFrame); this.connection = (ArrowFlightConnection) statement.connection; - this.flightInfo = ((ArrowFlightMetaStatement) statement).executeFlightInfoQuery(); + this.flightInfo = ((ArrowFlightInfoStatement) statement).executeFlightInfoQuery(); } /** Private constructor for fromFlightInfo. */ 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 f2b97b6044..0bd7789524 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 @@ -47,7 +47,10 @@ public ArrowFlightMetaImpl(final AvaticaConnection connection) { @Override public void closeStatement(final StatementHandle statementHandle) { - getMetaStatement(statementHandle).closeStatement(); + final AvaticaStatement statement = getStatement(statementHandle); + if (statement instanceof ArrowFlightPreparedStatement) { + ((ArrowFlightPreparedStatement) statement).closePreparedResources(); + } } @Override @@ -60,7 +63,8 @@ public ExecuteResult execute( final StatementHandle statementHandle, final List typedValues, final long maxRowCount) { - return getMetaStatement(statementHandle).execute(statementHandle, typedValues, maxRowCount); + return getPreparedStatement(statementHandle) + .executeWithTypedValues(statementHandle, typedValues, maxRowCount); } @Override @@ -75,7 +79,8 @@ public ExecuteResult execute( public ExecuteBatchResult executeBatch( final StatementHandle statementHandle, final List> parameterValuesList) throws IllegalStateException { - return getMetaStatement(statementHandle).executeBatch(statementHandle, parameterValuesList); + return getPreparedStatement(statementHandle) + .executeBatchWithTypedValues(statementHandle, parameterValuesList); } @Override @@ -125,8 +130,13 @@ public ExecuteResult prepareAndExecute( final PrepareCallback callback) throws NoSuchStatementException { try { - return getMetaStatement(handle) - .prepareAndExecute(query, maxRowCount, maxRowsInFirstFrame, callback); + final AvaticaStatement statement = getStatement(handle); + if (statement instanceof ArrowFlightStatement) { + return ((ArrowFlightStatement) statement) + .prepareAndExecuteInternal(query, maxRowCount, maxRowsInFirstFrame, callback); + } + return ((ArrowFlightPreparedStatement) statement) + .prepareAndExecuteInternal(query, maxRowCount, maxRowsInFirstFrame, callback); } catch (SQLException e) { // AvaticaStatement.executeInternal handles RuntimeException and preserves it as the cause of // the public SQLException. NoSuchStatementException would incorrectly trigger a statement @@ -181,14 +191,24 @@ void setDefaultConnectionProperties() { .setTransactionIsolation(Connection.TRANSACTION_NONE); } - private ArrowFlightMetaStatement getMetaStatement(StatementHandle statementHandle) { - AvaticaStatement statement = connection.statementMap.get(statementHandle.id); - if (statement instanceof ArrowFlightMetaStatement) { - return (ArrowFlightMetaStatement) statement; + private AvaticaStatement getStatement(final StatementHandle statementHandle) { + final AvaticaStatement statement = connection.statementMap.get(statementHandle.id); + if (statement instanceof ArrowFlightStatement + || statement instanceof ArrowFlightPreparedStatement) { + return statement; } throw new IllegalStateException("Statement not found: " + statementHandle); } + private ArrowFlightPreparedStatement getPreparedStatement(final StatementHandle statementHandle) { + final AvaticaStatement statement = getStatement(statementHandle); + if (statement instanceof ArrowFlightPreparedStatement) { + return (ArrowFlightPreparedStatement) statement; + } + throw new IllegalStateException( + "Statement operation is not supported for handle: " + statementHandle); + } + static Signature buildDefaultSignature() { return buildSignature(null, StatementType.SELECT); } diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaStatement.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaStatement.java deleted file mode 100644 index 474fa03681..0000000000 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaStatement.java +++ /dev/null @@ -1,61 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.apache.arrow.driver.jdbc; - -import java.sql.SQLException; -import java.sql.Statement; -import java.util.List; -import org.apache.arrow.flight.FlightInfo; -import org.apache.calcite.avatica.Meta.ExecuteBatchResult; -import org.apache.calcite.avatica.Meta.ExecuteResult; -import org.apache.calcite.avatica.Meta.PrepareCallback; -import org.apache.calcite.avatica.Meta.StatementHandle; -import org.apache.calcite.avatica.remote.TypedValue; - -/** Statement capabilities used by {@link ArrowFlightMetaImpl}. */ -interface ArrowFlightMetaStatement extends ArrowFlightInfoStatement { - - @Override - ArrowFlightConnection getConnection() throws SQLException; - - @Override - FlightInfo executeFlightInfoQuery() throws SQLException; - - /** - * Avatica routes {@link Statement#execute(String)} through Meta.prepareAndExecute(...), so plain - * statements still need this hook even when they support direct executeQuery/executeUpdate paths. - */ - ExecuteResult prepareAndExecute( - String query, long maxRowCount, int maxRowsInFirstFrame, PrepareCallback callback) - throws SQLException; - - default ExecuteResult execute( - final StatementHandle statementHandle, - final List typedValues, - final long maxRowCount) { - throw new IllegalStateException( - "Statement operation is not supported for handle: " + statementHandle); - } - - default ExecuteBatchResult executeBatch( - final StatementHandle statementHandle, final List> parameterValuesList) { - throw new IllegalStateException( - "Statement operation is not supported for handle: " + statementHandle); - } - - default void closeStatement() {} -} diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java index 2f6d0fe44d..0da2fe4a48 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java @@ -38,7 +38,7 @@ /** Arrow Flight JDBC's implementation {@link java.sql.PreparedStatement}. */ public class ArrowFlightPreparedStatement extends AvaticaPreparedStatement - implements ArrowFlightMetaStatement { + implements ArrowFlightInfoStatement { private ArrowFlightSqlClientHandler.PreparedStatement preparedStatement; @@ -66,7 +66,7 @@ public ArrowFlightConnection getConnection() throws SQLException { return (ArrowFlightConnection) super.getConnection(); } - ExecuteResult prepareAndExecute(final PrepareCallback callback) throws SQLException { + ExecuteResult prepareAndExecuteInternal(final PrepareCallback callback) throws SQLException { ensurePrepared(); final StatementType statementType = preparedStatement.getType(); final long updateCount = @@ -81,8 +81,7 @@ ExecuteResult prepareAndExecute(final PrepareCallback callback) throws SQLExcept return new ExecuteResult(Collections.singletonList(metaResultSet)); } - @Override - public ExecuteResult prepareAndExecute( + ExecuteResult prepareAndExecuteInternal( final String query, final long maxRowCount, final int maxRowsInFirstFrame, @@ -93,7 +92,7 @@ public ExecuteResult prepareAndExecute( .withQuery(query) .withExistingStatement(this) .build() - .prepareAndExecute(callback); + .prepareAndExecuteInternal(callback); } Schema getDataSetSchema() { @@ -163,25 +162,6 @@ ExecuteBatchResult executeBatchWithTypedValues( return new ExecuteBatchResult(updatedCounts); } - @Override - public ExecuteResult execute( - final StatementHandle statementHandle, - final List typedValues, - final long maxRowCount) { - return executeWithTypedValues(statementHandle, typedValues, maxRowCount); - } - - @Override - public ExecuteBatchResult executeBatch( - final StatementHandle statementHandle, final List> parameterValuesList) { - return executeBatchWithTypedValues(statementHandle, parameterValuesList); - } - - @Override - public void closeStatement() { - closePreparedResources(); - } - @Override public FlightInfo executeFlightInfoQuery() throws SQLException { ensurePrepared(); diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightStatement.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightStatement.java index 104abdb9b4..adb887dc5a 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightStatement.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightStatement.java @@ -33,7 +33,7 @@ import org.apache.calcite.avatica.Meta.StatementType; /** A SQL statement for querying data from an Arrow Flight server. */ -public class ArrowFlightStatement extends AvaticaStatement implements ArrowFlightMetaStatement { +public class ArrowFlightStatement extends AvaticaStatement implements ArrowFlightInfoStatement { ArrowFlightStatement( final ArrowFlightConnection connection, @@ -59,8 +59,7 @@ public ArrowFlightConnection getConnection() throws SQLException { return (ArrowFlightConnection) super.getConnection(); } - @Override - public ExecuteResult prepareAndExecute( + ExecuteResult prepareAndExecuteInternal( final String query, final long maxRowCount, final int maxRowsInFirstFrame, @@ -68,13 +67,11 @@ public ExecuteResult prepareAndExecute( throws SQLException { // Keep Avatica Statement.execute(String) behavior: Avatica calls Meta.prepareAndExecute, // which resolves to this statement hook. - this.closeStatement(); - return ArrowFlightPreparedStatement.builder(getConnection()) .withQuery(query) .withExistingStatement(this) .build() - .prepareAndExecute(callback); + .prepareAndExecuteInternal(callback); } @Override @@ -192,7 +189,7 @@ private void switchToDirectStatementMode() throws SQLException { if (existingStatement instanceof ArrowFlightPreparedStatement) { // Release resources from previously attached statement implementation before switching back // to direct statement mode for executeQuery/executeUpdate. - ((ArrowFlightPreparedStatement) existingStatement).closeStatement(); + ((ArrowFlightPreparedStatement) existingStatement).closePreparedResources(); } connection.statementMap.put(handle.id, this); } diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementApiTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementApiTest.java new file mode 100644 index 0000000000..49c6c3de10 --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementApiTest.java @@ -0,0 +1,88 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.driver.jdbc; + +import static org.junit.jupiter.api.Assertions.assertThrows; + +import java.util.List; +import org.apache.calcite.avatica.Meta.PrepareCallback; +import org.apache.calcite.avatica.Meta.StatementHandle; +import org.junit.jupiter.api.Test; + +public class ArrowFlightStatementApiTest { + + @Test + public void testMetaOrchestrationMethodsAreNotPublicStatementApis() { + assertNotPublic( + ArrowFlightStatement.class, + "prepareAndExecute", + String.class, + long.class, + int.class, + PrepareCallback.class); + assertNotPublic( + ArrowFlightPreparedStatement.class, + "prepareAndExecute", + String.class, + long.class, + int.class, + PrepareCallback.class); + assertNotPublic( + ArrowFlightPreparedStatement.class, + "execute", + StatementHandle.class, + List.class, + long.class); + assertNotPublic( + ArrowFlightPreparedStatement.class, "executeBatch", StatementHandle.class, List.class); + assertNotPublic(ArrowFlightPreparedStatement.class, "closeStatement"); + assertNotPublic( + ArrowFlightStatement.class, + "prepareAndExecuteInternal", + String.class, + long.class, + int.class, + PrepareCallback.class); + assertNotPublic( + ArrowFlightPreparedStatement.class, + "prepareAndExecuteInternal", + String.class, + long.class, + int.class, + PrepareCallback.class); + assertNotPublic( + ArrowFlightPreparedStatement.class, + "executeWithTypedValues", + StatementHandle.class, + List.class, + long.class); + assertNotPublic( + ArrowFlightPreparedStatement.class, + "executeBatchWithTypedValues", + StatementHandle.class, + List.class); + assertNotPublic(ArrowFlightPreparedStatement.class, "closePreparedResources"); + } + + private static void assertNotPublic( + final Class statementClass, final String methodName, final Class... parameterTypes) { + assertThrows( + NoSuchMethodException.class, + () -> statementClass.getMethod(methodName, parameterTypes), + methodName + " must remain an internal orchestration method"); + } +} 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 abecd5ea83..6b4fcf0676 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 @@ -248,7 +248,7 @@ public void testExecuteQueryMapsPreparedStatementCloseFailureToSQLException() { final ArrowFlightPreparedStatement preparedStatement = mock(ArrowFlightPreparedStatement.class); final FlightRuntimeException closeFailure = CallStatus.INTERNAL.withDescription("Prepared statement close failed").toRuntimeException(); - doThrow(closeFailure).when(preparedStatement).closeStatement(); + doThrow(closeFailure).when(preparedStatement).closePreparedResources(); arrowConnection.statementMap.put(arrowStatement.handle.id, preparedStatement); try { @@ -268,7 +268,7 @@ public void testExecuteLargeUpdateMapsPreparedStatementCloseFailureToSQLExceptio final ArrowFlightPreparedStatement preparedStatement = mock(ArrowFlightPreparedStatement.class); final FlightRuntimeException closeFailure = CallStatus.INTERNAL.withDescription("Prepared statement close failed").toRuntimeException(); - doThrow(closeFailure).when(preparedStatement).closeStatement(); + doThrow(closeFailure).when(preparedStatement).closePreparedResources(); arrowConnection.statementMap.put(arrowStatement.handle.id, preparedStatement); try { From c229eff81a9bdb7ea517ea8e43f35a3137ee2d7e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Thu, 24 Sep 2026 11:45:09 +0100 Subject: [PATCH 19/25] Restore queue post-close contract Keep FlightEndpointDataQueue fail-fast while normalizing cancellation races inside the JDBC result-set path. --- .../ArrowFlightJdbcFlightStreamResultSet.java | 69 ++++++++++++++++--- .../jdbc/utils/FlightEndpointDataQueue.java | 13 +--- ...owFlightJdbcFlightStreamResultSetTest.java | 66 ++++++++++++++++++ .../utils/FlightEndpointDataQueueTest.java | 13 ++-- 4 files changed, 134 insertions(+), 27 deletions(-) create mode 100644 flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSetTest.java 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 74c53104db..a9537eecc1 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 @@ -21,6 +21,8 @@ import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.SQLException; +import java.util.Collection; +import java.util.Iterator; import java.util.Optional; import java.util.TimeZone; import java.util.concurrent.TimeUnit; @@ -146,7 +148,8 @@ protected AvaticaResultSet execute() throws SQLException { private void populateData() throws SQLException { loadNewQueue(); - flightEndpointDataQueue.enqueue(connection.getClientHandler().getStreams(flightInfo)); + enqueueEndpointData( + flightEndpointDataQueue, connection.getClientHandler().getStreams(flightInfo)); loadNewFlightStream(); // Ownership of the root will be passed onto the cursor. @@ -203,7 +206,7 @@ public boolean next() throws SQLException { continue; } - flightEndpointDataQueue.enqueue(currentEndpointData); + enqueueEndpointData(currentEndpointData); } currentEndpointData = getNextEndpointStream(false); @@ -279,13 +282,61 @@ public synchronized void close() { private CloseableEndpointStreamPair getNextEndpointStream(final boolean canTimeout) throws SQLException { - if (canTimeout) { - final int statementTimeout = statement != null ? statement.getQueryTimeout() : 0; - return statementTimeout != 0 - ? flightEndpointDataQueue.next(statementTimeout, TimeUnit.SECONDS) - : flightEndpointDataQueue.next(); - } else { - return flightEndpointDataQueue.next(); + try { + if (canTimeout) { + final int statementTimeout = statement != null ? statement.getQueryTimeout() : 0; + return statementTimeout != 0 + ? flightEndpointDataQueue.next(statementTimeout, TimeUnit.SECONDS) + : flightEndpointDataQueue.next(); + } else { + return flightEndpointDataQueue.next(); + } + } catch (final IllegalStateException e) { + throw normalizeClosedQueueException(flightEndpointDataQueue, e); + } + } + + static void enqueueEndpointData( + final FlightEndpointDataQueue queue, + final Collection endpointRequests) + throws SQLException { + final Iterator endpointIterator = endpointRequests.iterator(); + while (endpointIterator.hasNext()) { + final CloseableEndpointStreamPair endpointRequest = endpointIterator.next(); + try { + queue.enqueue(endpointRequest); + } catch (final IllegalStateException e) { + final SQLException exception = normalizeClosedQueueException(queue, e); + closeRejectedEndpoint(endpointRequest, exception); + endpointIterator.forEachRemaining(endpoint -> closeRejectedEndpoint(endpoint, exception)); + throw exception; + } + } + } + + private void enqueueEndpointData(final CloseableEndpointStreamPair endpointRequest) + throws SQLException { + try { + flightEndpointDataQueue.enqueue(endpointRequest); + } catch (final IllegalStateException e) { + throw normalizeClosedQueueException(flightEndpointDataQueue, e); + } + } + + private static void closeRejectedEndpoint( + final CloseableEndpointStreamPair endpoint, final SQLException exception) { + try { + endpoint.close(); + } catch (final Exception closeException) { + exception.addSuppressed(closeException); + } + } + + private static SQLException normalizeClosedQueueException( + final FlightEndpointDataQueue queue, final IllegalStateException exception) { + if (queue.isClosed()) { + return AvaticaConnection.HELPER.createException("Statement canceled"); } + throw exception; } } diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/FlightEndpointDataQueue.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/FlightEndpointDataQueue.java index c6f485ee8b..7c1c569fda 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/FlightEndpointDataQueue.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/FlightEndpointDataQueue.java @@ -97,6 +97,7 @@ interface EndpointStreamSupplier { private CloseableEndpointStreamPair next(final EndpointStreamSupplier endpointStreamSupplier) throws SQLException { + checkOpen(); while (!futures.isEmpty()) { final Future future = endpointStreamSupplier.get(); futures.remove(future); @@ -175,17 +176,7 @@ public void enqueue(final Collection endpointReques /** Adds given {@link FlightStream} to the queue. */ public synchronized void enqueue(final CloseableEndpointStreamPair endpointRequest) { checkNotNull(endpointRequest); - if (isClosed()) { - // Concurrent close() raced with enqueue (e.g. statement.cancel() while the reader - // was mid-loop). Treat as end-of-stream: close the endpoint and return silently. - // The reader detects cancellation via a subsequent isClosed() check. - try { - endpointRequest.close(); - } catch (final Exception e) { - LOGGER.error("Failed to close endpoint after queue was closed.", e); - } - return; - } + checkOpen(); endpointsToClose.add(endpointRequest); futures.add( completionService.submit( diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSetTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSetTest.java new file mode 100644 index 0000000000..b57b2f92b6 --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSetTest.java @@ -0,0 +1,66 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.driver.jdbc; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.sql.SQLException; +import java.util.List; +import org.apache.arrow.driver.jdbc.client.CloseableEndpointStreamPair; +import org.apache.arrow.driver.jdbc.utils.FlightEndpointDataQueue; +import org.junit.jupiter.api.Test; + +public class ArrowFlightJdbcFlightStreamResultSetTest { + + @Test + public void testClosedQueueDuringCollectionHandoffClosesUnacceptedEndpoints() throws Exception { + final FlightEndpointDataQueue queue = mock(FlightEndpointDataQueue.class); + final CloseableEndpointStreamPair acceptedEndpoint = mock(CloseableEndpointStreamPair.class); + final CloseableEndpointStreamPair rejectedEndpoint = mock(CloseableEndpointStreamPair.class); + final CloseableEndpointStreamPair unacceptedEndpoint = mock(CloseableEndpointStreamPair.class); + final Exception closeFailure = new Exception("Failed to close rejected endpoint"); + doThrow(new IllegalStateException("FlightEndpointDataQueue closed")) + .when(queue) + .enqueue(rejectedEndpoint); + when(queue.isClosed()).thenReturn(true); + doThrow(closeFailure).when(rejectedEndpoint).close(); + + final SQLException exception = + assertThrows( + SQLException.class, + () -> + ArrowFlightJdbcFlightStreamResultSet.enqueueEndpointData( + queue, List.of(acceptedEndpoint, rejectedEndpoint, unacceptedEndpoint))); + + assertEquals("Statement canceled", exception.getMessage()); + assertEquals(1, exception.getSuppressed().length); + assertSame(closeFailure, exception.getSuppressed()[0]); + verify(acceptedEndpoint, never()).close(); + verify(rejectedEndpoint).close(); + verify(unacceptedEndpoint).close(); + verify(queue).enqueue(acceptedEndpoint); + verify(queue).enqueue(rejectedEndpoint); + verify(queue, never()).enqueue(unacceptedEndpoint); + } +} diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/FlightEndpointDataQueueTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/FlightEndpointDataQueueTest.java index 8dd39796e8..f40610a032 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/FlightEndpointDataQueueTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/FlightEndpointDataQueueTest.java @@ -21,7 +21,6 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.verify; import java.util.concurrent.CompletionService; import org.apache.arrow.driver.jdbc.client.CloseableEndpointStreamPair; @@ -49,17 +48,17 @@ public void testNextShouldRetrieveNullIfEmpty() throws Exception { } @Test - public void testNextShouldReturnNullUponClose() throws Exception { + public void testNextShouldThrowExceptionUponClose() throws Exception { queue.close(); - assertThat(queue.next(), is(nullValue())); + ThrowableAssertionUtils.simpleAssertThrowableClass( + IllegalStateException.class, () -> queue.next()); } @Test - public void testEnqueueShouldCloseEndpointSilentlyAfterClose() throws Exception { + public void testEnqueueShouldThrowExceptionUponClose() throws Exception { queue.close(); - final CloseableEndpointStreamPair endpoint = mock(CloseableEndpointStreamPair.class); - assertDoesNotThrow(() -> queue.enqueue(endpoint)); - verify(endpoint).close(); + ThrowableAssertionUtils.simpleAssertThrowableClass( + IllegalStateException.class, () -> queue.enqueue(mock(CloseableEndpointStreamPair.class))); } @Test From 6f72bcb2d9804e3e41c6ecb043f0d477b41841bd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Mon, 28 Sep 2026 13:57:08 +0100 Subject: [PATCH 20/25] Close remaining JDBC lifecycle gaps --- .../jdbc/ArrowFlightPreparedStatement.java | 4 +- .../driver/jdbc/ArrowFlightStatement.java | 5 +- ...lightJdbcConnectionPoolDataSourceTest.java | 98 +++++++++++++++++++ .../jdbc/ArrowFlightStatementExecuteTest.java | 5 +- .../ArrowFlightStatementProtocolTest.java | 10 ++ .../arrow/driver/jdbc/ResultSetTest.java | 38 +++++++ 6 files changed, 155 insertions(+), 5 deletions(-) diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java index 0da2fe4a48..fc83a4ec00 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java @@ -101,9 +101,9 @@ Schema getDataSetSchema() { } @Override - public synchronized void close() throws SQLException { + protected void close_() { try { - super.close(); + super.close_(); } finally { ((ArrowFlightConnection) connection).unregisterStatementOwner(this); } diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightStatement.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightStatement.java index adb887dc5a..eb170c18c6 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightStatement.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightStatement.java @@ -46,9 +46,9 @@ public class ArrowFlightStatement extends AvaticaStatement implements ArrowFligh } @Override - public synchronized void close() throws SQLException { + protected void close_() { try { - super.close(); + super.close_(); } finally { ((ArrowFlightConnection) connection).unregisterStatementOwner(this); } @@ -77,6 +77,7 @@ ExecuteResult prepareAndExecuteInternal( @Override public ResultSet executeQuery(final String sql) throws SQLException { checkOpen(); + clearOpenResultSet(); updateCount = -1; try { switchToDirectStatementMode(); diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcConnectionPoolDataSourceTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcConnectionPoolDataSourceTest.java index 18ff45bbde..14dda3d8a6 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcConnectionPoolDataSourceTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcConnectionPoolDataSourceTest.java @@ -22,10 +22,21 @@ import static org.junit.jupiter.api.Assertions.assertTrue; import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.Collections; import javax.sql.PooledConnection; import org.apache.arrow.driver.jdbc.authentication.UserPasswordAuthentication; import org.apache.arrow.driver.jdbc.utils.ConnectionWrapper; import org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.Types.MinorType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.Schema; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -33,6 +44,10 @@ public class ArrowFlightJdbcConnectionPoolDataSourceTest { + private static final String SELECT_QUERY = "SELECT * FROM POOLED_CONNECTION_RESET"; + private static final Schema QUERY_SCHEMA = + new Schema(Collections.singletonList(Field.nullable("id", MinorType.INT.getType()))); + @RegisterExtension public static final FlightServerTestExtension FLIGHT_SERVER_TEST_EXTENSION; private static final MockFlightSqlProducer PRODUCER = new MockFlightSqlProducer(); @@ -49,6 +64,25 @@ public class ArrowFlightJdbcConnectionPoolDataSourceTest { .authentication(authentication) .producer(PRODUCER) .build(); + + PRODUCER.addSelectQuery( + SELECT_QUERY, + QUERY_SCHEMA, + Collections.singletonList( + listener -> { + try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + final VectorSchemaRoot root = VectorSchemaRoot.create(QUERY_SCHEMA, allocator)) { + final IntVector vector = (IntVector) root.getVector("id"); + vector.setSafe(0, 1); + root.setRowCount(1); + listener.start(root); + listener.putNext(); + } catch (final Throwable throwable) { + listener.error(throwable); + } finally { + listener.completed(); + } + })); } private ArrowFlightJdbcConnectionPoolDataSource dataSource; @@ -114,6 +148,70 @@ public void testShouldReuseConnectionsOnPool() throws Exception { connection2.unwrap(ArrowFlightConnection.class)); } + @Test + public void testShouldResetOpenStatementExecuteResultSetWhenReusingConnection() throws Exception { + final PooledConnection pooledConnection = dataSource.getPooledConnection("user1", "pass1"); + final Connection connection; + try { + connection = pooledConnection.getConnection(); + } catch (final SQLException e) { + closePooledConnectionAfterFailedCheckout(pooledConnection, e); + throw e; + } + + try (connection) { + final ArrowFlightConnection physicalConnection = + connection.unwrap(ArrowFlightConnection.class); + final Statement statement = connection.createStatement(); + assertTrue(statement.execute(SELECT_QUERY)); + final ResultSet resultSet = statement.getResultSet(); + + assertFalse(statement.isClosed()); + assertFalse(resultSet.isClosed()); + connection.close(); + assertFalse(statement.isClosed()); + assertFalse(resultSet.isClosed()); + + final PooledConnection reusedPooledConnection; + try { + reusedPooledConnection = dataSource.getPooledConnection("user1", "pass1"); + } catch (final SQLException e) { + closePooledConnectionAfterFailedCheckout(pooledConnection, e); + throw e; + } + + final Connection reusedConnection; + try { + reusedConnection = reusedPooledConnection.getConnection(); + } catch (final SQLException e) { + closePooledConnectionAfterFailedCheckout(reusedPooledConnection, e); + throw e; + } + + try (reusedConnection) { + assertSame(pooledConnection, reusedPooledConnection); + assertSame(physicalConnection, reusedConnection.unwrap(ArrowFlightConnection.class)); + assertTrue(statement.isClosed()); + assertTrue(resultSet.isClosed()); + assertFalse(reusedConnection.isClosed()); + + try (Statement reusedStatement = reusedConnection.createStatement(); + ResultSet reusedResultSet = reusedStatement.executeQuery(SELECT_QUERY)) { + assertTrue(reusedResultSet.next()); + } + } + } + } + + private static void closePooledConnectionAfterFailedCheckout( + final PooledConnection pooledConnection, final SQLException failure) { + try { + pooledConnection.close(); + } catch (final SQLException closeFailure) { + failure.addSuppressed(closeFailure); + } + } + @Test public void testShouldNotMixConnectionsForDifferentUsers() throws Exception { PooledConnection pooledConnection = dataSource.getPooledConnection("user1", "pass1"); 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 6b4fcf0676..0145d53be2 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 @@ -242,9 +242,11 @@ public void testExecuteQueryRestoresStatementMapEntryWithStatement() throws SQLE } @Test - public void testExecuteQueryMapsPreparedStatementCloseFailureToSQLException() { + public void testExecuteQueryMapsPreparedStatementCloseFailureToSQLException() + throws SQLException { final ArrowFlightStatement arrowStatement = (ArrowFlightStatement) statement; final ArrowFlightConnection arrowConnection = (ArrowFlightConnection) connection; + final ResultSet previousResultSet = statement.executeQuery(SAMPLE_QUERY_CMD); final ArrowFlightPreparedStatement preparedStatement = mock(ArrowFlightPreparedStatement.class); final FlightRuntimeException closeFailure = CallStatus.INTERNAL.withDescription("Prepared statement close failed").toRuntimeException(); @@ -256,6 +258,7 @@ public void testExecuteQueryMapsPreparedStatementCloseFailureToSQLException() { assertThrows(SQLException.class, () -> statement.executeQuery(SAMPLE_QUERY_CMD)); assertSame(closeFailure, exception.getCause()); + assertThat(previousResultSet.isClosed(), is(true)); } finally { arrowConnection.statementMap.put(arrowStatement.handle.id, arrowStatement); } diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementProtocolTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementProtocolTest.java index 1ac31d3528..0cc533a909 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementProtocolTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementProtocolTest.java @@ -258,6 +258,11 @@ public void testStatementExecuteThenExecuteUpdateUsesStatementProtocol() throws .getCommandTypeCounter() .getOrDefault(MockFlightSqlProducer.COMMAND_STATEMENT_UPDATE, 0), is(1)); + assertThat( + PRODUCER + .getActionTypeCounter() + .getOrDefault(FlightSqlUtils.FLIGHT_SQL_CLOSE_PREPARED_STATEMENT.getType(), 0), + is(1)); } @Test @@ -286,6 +291,11 @@ public void testStatementExecuteUpdateThenExecuteQueryUsesStatementProtocol() .getCommandTypeCounter() .getOrDefault(MockFlightSqlProducer.COMMAND_STATEMENT_QUERY, 0), is(1)); + assertThat( + PRODUCER + .getActionTypeCounter() + .getOrDefault(FlightSqlUtils.FLIGHT_SQL_CLOSE_PREPARED_STATEMENT.getType(), 0), + is(1)); } @Test diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java index 3a5a39be3d..cc647dab2b 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java @@ -235,6 +235,44 @@ public void testShouldCloseStatementWhenIsCloseOnCompletion() throws Exception { } } + @Test + public void testCloseOnCompletionUnregistersDirectStatementWhenResultSetClosedEarly() + throws Exception { + try (Statement statement = connection.createStatement()) { + final ResultSet resultSet = + statement.executeQuery(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD); + statement.closeOnCompletion(); + + resultSet.close(); + + assertThat(statement.isClosed(), is(true)); + assertThat(isStatementOwner(statement), is(false)); + } + } + + @Test + public void testCloseOnCompletionUnregistersPreparedStatementWhenResultSetClosedEarly() + throws Exception { + try (PreparedStatement statement = + connection.prepareStatement(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD)) { + final ResultSet resultSet = statement.executeQuery(); + statement.closeOnCompletion(); + + resultSet.close(); + + assertThat(statement.isClosed(), is(true)); + assertThat(isStatementOwner(statement), is(false)); + } + } + + private static boolean isStatementOwner(final Statement statement) + throws ReflectiveOperationException { + final java.lang.reflect.Field statementOwnersField = + ArrowFlightConnection.class.getDeclaredField("statementOwners"); + statementOwnersField.setAccessible(true); + return ((Set) statementOwnersField.get(connection)).contains(statement); + } + /** * Tests whether the {@link ArrowFlightJdbcDriver} close the statement after complete ResultSet * with max rows limit when call {@link From b7bd7c5329424bae5eeeb71fe4d8c3101c8caa69 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Fri, 2 Oct 2026 14:30:27 +0100 Subject: [PATCH 21/25] Fix remaining JDBC lifecycle validation gaps --- .../driver/jdbc/ArrowFlightConnection.java | 18 +++-- .../driver/jdbc/ArrowFlightMetaImpl.java | 8 ++ .../jdbc/ArrowFlightPreparedStatement.java | 25 ++++-- .../ArrowFlightPreparedStatementTest.java | 80 +++++++++++++++++++ 4 files changed, 119 insertions(+), 12 deletions(-) diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java index 5f0b04ec41..87f9925e08 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java @@ -313,12 +313,16 @@ public PreparedStatement prepareStatement( final int resultSetHoldability) throws SQLException { checkOpen(); - return ArrowFlightPreparedStatement.builder(this) - .withQuery(sql) - .withGeneratedHandle() - .withResultSetType(resultSetType) - .withResultSetConcurrency(resultSetConcurrency) - .withResultSetHoldability(resultSetHoldability) - .build(); + try { + return ArrowFlightPreparedStatement.builder(this) + .withQuery(sql) + .withGeneratedHandle() + .withResultSetType(resultSetType) + .withResultSetConcurrency(resultSetConcurrency) + .withResultSetHoldability(resultSetHoldability) + .build(); + } catch (final RuntimeException e) { + throw HELPER.createException("while preparing SQL:" + sql, e); + } } } 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 0bd7789524..b085f6537b 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 @@ -23,6 +23,7 @@ import java.util.Collections; import java.util.List; import org.apache.arrow.driver.jdbc.utils.ConvertUtils; +import org.apache.arrow.util.Preconditions; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.calcite.avatica.AvaticaConnection; import org.apache.calcite.avatica.AvaticaParameter; @@ -98,6 +99,7 @@ public Frame fetch( @Override public StatementHandle prepare( final ConnectionHandle connectionHandle, final String query, final long maxRowCount) { + validateConnectionId(connectionHandle.id); try { // This is the Avatica entry point used by Connection.prepareStatement(String). ArrowFlightPreparedStatement stmt = @@ -192,6 +194,7 @@ void setDefaultConnectionProperties() { } private AvaticaStatement getStatement(final StatementHandle statementHandle) { + validateConnectionId(statementHandle.connectionId); final AvaticaStatement statement = connection.statementMap.get(statementHandle.id); if (statement instanceof ArrowFlightStatement || statement instanceof ArrowFlightPreparedStatement) { @@ -209,6 +212,11 @@ private ArrowFlightPreparedStatement getPreparedStatement(final StatementHandle "Statement operation is not supported for handle: " + statementHandle); } + private void validateConnectionId(final String connectionId) { + Preconditions.checkArgument( + connection.id.equals(connectionId), "Connection IDs are not consistent"); + } + static Signature buildDefaultSignature() { return buildSignature(null, StatementType.SELECT); } diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java index fc83a4ec00..30878ad8b9 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java @@ -102,17 +102,32 @@ Schema getDataSetSchema() { @Override protected void close_() { + RuntimeException closeFailure = null; try { super.close_(); - } finally { - ((ArrowFlightConnection) connection).unregisterStatementOwner(this); + } catch (final RuntimeException e) { + closeFailure = e; + } + try { + closePreparedResources(); + } catch (final RuntimeException e) { + if (closeFailure == null) { + closeFailure = e; + } else { + closeFailure.addSuppressed(e); + } + } + ((ArrowFlightConnection) connection).unregisterStatementOwner(this); + if (closeFailure != null) { + throw closeFailure; } } void closePreparedResources() { - if (preparedStatement != null) { - preparedStatement.close(); - preparedStatement = null; + final ArrowFlightSqlClientHandler.PreparedStatement statement = preparedStatement; + preparedStatement = null; + if (statement != null) { + statement.close(); } } 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 2ca485729d..08a6f09069 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 @@ -26,6 +26,8 @@ import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; import java.nio.charset.StandardCharsets; import java.sql.Connection; @@ -52,6 +54,8 @@ import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.arrow.vector.util.Text; +import org.apache.calcite.avatica.AvaticaResultSet; +import org.apache.calcite.avatica.AvaticaStatement; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.BeforeEach; @@ -177,6 +181,82 @@ public void testPrepareFailureIsReportedAsSQLException() { assertThat(exception.getCause(), instanceOf(FlightRuntimeException.class)); } + @Test + public void testNullQueryIsReportedAsSQLException() { + final SQLException exception = + assertThrows(SQLException.class, () -> connection.prepareStatement(null)); + + assertThat(exception.getCause(), instanceOf(NullPointerException.class)); + } + + @Test + public void testCloseReleasesPreparedResourcesWhenResultSetCloseFails() throws Exception { + try (final ArrowFlightConnection localConnection = + (ArrowFlightConnection) FLIGHT_SERVER_TEST_EXTENSION.getConnection(false)) { + final ArrowFlightPreparedStatement preparedStatement = + (ArrowFlightPreparedStatement) + localConnection.prepareStatement(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD); + final AvaticaResultSet resultSet = mock(AvaticaResultSet.class); + final RuntimeException resultSetCloseFailure = + new RuntimeException("Result set close failed"); + doThrow(resultSetCloseFailure).when(resultSet).close(); + final java.lang.reflect.Field openResultSetField = + AvaticaStatement.class.getDeclaredField("openResultSet"); + openResultSetField.setAccessible(true); + openResultSetField.set(preparedStatement, resultSet); + + final SQLException exception = assertThrows(SQLException.class, preparedStatement::close); + + assertSame(resultSetCloseFailure, exception.getCause()); + assertTrue(preparedStatement.isClosed()); + assertEquals( + 1, + PRODUCER + .getActionTypeCounter() + .getOrDefault(FlightSqlUtils.FLIGHT_SQL_CLOSE_PREPARED_STATEMENT.getType(), 0)); + } + } + + @Test + public void testMetaRejectsForeignStatementHandle() throws SQLException { + try (final ArrowFlightConnection firstConnection = + (ArrowFlightConnection) FLIGHT_SERVER_TEST_EXTENSION.getConnection(false); + final ArrowFlightConnection secondConnection = + (ArrowFlightConnection) FLIGHT_SERVER_TEST_EXTENSION.getConnection(false); + final PreparedStatement firstStatement = + firstConnection.prepareStatement(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD); + final PreparedStatement secondStatement = + secondConnection.prepareStatement( + CoreMockedSqlProducers.UUID_PREPARED_SELECT_SQL_CMD)) { + final ArrowFlightPreparedStatement firstPreparedStatement = + (ArrowFlightPreparedStatement) firstStatement; + final ArrowFlightPreparedStatement secondPreparedStatement = + (ArrowFlightPreparedStatement) secondStatement; + assertEquals(firstPreparedStatement.handle.id, secondPreparedStatement.handle.id); + + assertThrows( + IllegalArgumentException.class, + () -> firstConnection.getMeta().closeStatement(secondPreparedStatement.handle)); + assertFalse(firstPreparedStatement.isClosed()); + } + } + + @Test + public void testMetaRejectsForeignConnectionHandle() throws SQLException { + try (final ArrowFlightConnection firstConnection = + (ArrowFlightConnection) FLIGHT_SERVER_TEST_EXTENSION.getConnection(false); + final ArrowFlightConnection secondConnection = + (ArrowFlightConnection) FLIGHT_SERVER_TEST_EXTENSION.getConnection(false)) { + assertThrows( + IllegalArgumentException.class, + () -> + firstConnection + .getMeta() + .prepare( + secondConnection.handle, CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD, -1)); + } + } + @Test public void testQueryWithParameterBinding() throws SQLException { final String query = "Fake query with parameters"; From f261aeac32cabfb902019aa023719291029e08b3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Tue, 10 Nov 2026 11:26:32 +0000 Subject: [PATCH 22/25] Simplify JDBC tests and connection checks Remove redundant API visibility and interface identity tests, retain the checked prepareStatement implementation, and clarify connection ownership comments. --- .../driver/jdbc/ArrowFlightConnection.java | 5 +- .../ArrowFlightPreparedStatementTest.java | 11 --- .../jdbc/ArrowFlightStatementApiTest.java | 88 ------------------- .../jdbc/ArrowFlightStatementExecuteTest.java | 7 -- 4 files changed, 2 insertions(+), 109 deletions(-) delete mode 100644 flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementApiTest.java diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java index 87f9925e08..2c82999aec 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java @@ -52,8 +52,9 @@ public final class ArrowFlightConnection extends AvaticaConnection { private final ArrowFlightConnectionConfigImpl config; private ExecutorService executorService; private int metadataResultSetCount; - // Track result sets and statements owned by this connection so close() can release them. + // Track metadata result sets so connection close can release them. private Map metadataResultSetMap = new HashMap<>(); + // Track original statement owners when execution replaces their statement-map entries. private final Set statementOwners = Collections.newSetFromMap(new IdentityHashMap<>()); @@ -293,7 +294,6 @@ public ArrowFlightMetaImpl getMeta() { @Override public PreparedStatement prepareStatement(final String sql) throws SQLException { - checkOpen(); return prepareStatement(sql, ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY); } @@ -301,7 +301,6 @@ public PreparedStatement prepareStatement(final String sql) throws SQLException public PreparedStatement prepareStatement( final String sql, final int resultSetType, final int resultSetConcurrency) throws SQLException { - checkOpen(); return prepareStatement(sql, resultSetType, resultSetConcurrency, getHoldability()); } 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 08a6f09069..24bde53449 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 @@ -88,17 +88,6 @@ public void before() { PRODUCER.clearActionTypeCounter(); } - @Test - public void testPreparedStatementKeepsFlightInfoStatementApi() throws SQLException { - try (final PreparedStatement preparedStatement = - connection.prepareStatement(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD)) { - final ArrowFlightInfoStatement flightInfoStatement = - (ArrowFlightInfoStatement) preparedStatement; - - assertSame(preparedStatement, flightInfoStatement); - } - } - @Test public void testSimpleQueryNoParameterBinding() throws SQLException { final String query = CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD; diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementApiTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementApiTest.java deleted file mode 100644 index 49c6c3de10..0000000000 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementApiTest.java +++ /dev/null @@ -1,88 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.apache.arrow.driver.jdbc; - -import static org.junit.jupiter.api.Assertions.assertThrows; - -import java.util.List; -import org.apache.calcite.avatica.Meta.PrepareCallback; -import org.apache.calcite.avatica.Meta.StatementHandle; -import org.junit.jupiter.api.Test; - -public class ArrowFlightStatementApiTest { - - @Test - public void testMetaOrchestrationMethodsAreNotPublicStatementApis() { - assertNotPublic( - ArrowFlightStatement.class, - "prepareAndExecute", - String.class, - long.class, - int.class, - PrepareCallback.class); - assertNotPublic( - ArrowFlightPreparedStatement.class, - "prepareAndExecute", - String.class, - long.class, - int.class, - PrepareCallback.class); - assertNotPublic( - ArrowFlightPreparedStatement.class, - "execute", - StatementHandle.class, - List.class, - long.class); - assertNotPublic( - ArrowFlightPreparedStatement.class, "executeBatch", StatementHandle.class, List.class); - assertNotPublic(ArrowFlightPreparedStatement.class, "closeStatement"); - assertNotPublic( - ArrowFlightStatement.class, - "prepareAndExecuteInternal", - String.class, - long.class, - int.class, - PrepareCallback.class); - assertNotPublic( - ArrowFlightPreparedStatement.class, - "prepareAndExecuteInternal", - String.class, - long.class, - int.class, - PrepareCallback.class); - assertNotPublic( - ArrowFlightPreparedStatement.class, - "executeWithTypedValues", - StatementHandle.class, - List.class, - long.class); - assertNotPublic( - ArrowFlightPreparedStatement.class, - "executeBatchWithTypedValues", - StatementHandle.class, - List.class); - assertNotPublic(ArrowFlightPreparedStatement.class, "closePreparedResources"); - } - - private static void assertNotPublic( - final Class statementClass, final String methodName, final Class... parameterTypes) { - assertThrows( - NoSuchMethodException.class, - () -> statementClass.getMethod(methodName, parameterTypes), - methodName + " must remain an internal orchestration method"); - } -} 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 0145d53be2..704905ea47 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 @@ -152,13 +152,6 @@ public static void tearDownAfterClass() throws Exception { AutoCloseables.close(PRODUCER); } - @Test - public void testStatementKeepsFlightInfoStatementApi() { - final ArrowFlightInfoStatement flightInfoStatement = (ArrowFlightInfoStatement) statement; - - assertSame(statement, flightInfoStatement); - } - @Test public void testExecuteShouldRunSelectQuery() throws SQLException { assertThat(statement.execute(SAMPLE_QUERY_CMD), is(true)); // Means this is a SELECT query. From 0da86544e8a0b0a09be95eb5d611227ce39a2a56 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?JB=20Onofr=C3=A9?= Date: Tue, 10 Nov 2026 14:56:39 +0100 Subject: [PATCH 23/25] Make cancel-before-retrieval test deterministic testShouldInterruptFlightStreamsIfQueryIsCancelledMidQuerying called statement.cancel() right after starting the worker thread. It relied on the prepare round trip in executeQuery to make the cancel land before a result set existed, so only Avatica's cancel flag was set and next() threw "Statement canceled". Statement.executeQuery no longer prepares first, so the cancel can now land while the result set is executing and surface as "Error while executing SQL ...: exception while executing query: Statement canceled", which failed the RC Binaries job. Wait until executeQuery returns before cancelling, and replace the 300 ms sleep with a latch. --- .../apache/arrow/driver/jdbc/ResultSetTest.java | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java index cc647dab2b..183182ee1b 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java @@ -341,7 +341,8 @@ public void testShouldCancelQueryUponCancelAfterQueryingResultSet() throws SQLEx public void testShouldInterruptFlightStreamsIfQueryIsCancelledMidQuerying() throws SQLException, InterruptedException { try (final Statement statement = connection.createStatement()) { - final CountDownLatch latch = new CountDownLatch(1); + final CountDownLatch executed = new CountDownLatch(1); + final CountDownLatch canceled = new CountDownLatch(1); final Set exceptions = synchronizedSet(new HashSet<>(1)); final Thread thread = new Thread( @@ -349,19 +350,27 @@ public void testShouldInterruptFlightStreamsIfQueryIsCancelledMidQuerying() try (final ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD)) { final int cachedColumnCount = resultSet.getMetaData().getColumnCount(); - Thread.sleep(300); + executed.countDown(); + canceled.await(); while (resultSet.next()) { resultSet.getObject(RANDOM.nextInt(cachedColumnCount) + 1); } } catch (final SQLException | InterruptedException e) { exceptions.add(e); } finally { - latch.countDown(); + executed.countDown(); } }); thread.setName("Test Case: interrupt query execution before first retrieval"); thread.start(); - statement.cancel(); + // Cancel only after executeQuery returns, so the cancellation lands before the first + // retrieval instead of racing with query execution. + executed.await(); + try { + statement.cancel(); + } finally { + canceled.countDown(); + } thread.join(); assertThat( exceptions.stream() From 2561d8b1137d6a6d0a74fc75435bc76063c20d54 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Tue, 10 Nov 2026 12:15:56 +0000 Subject: [PATCH 24/25] Strengthen JDBC lifecycle tests and independent assertions --- .../driver/jdbc/ArrowFlightConnection.java | 7 ++ ...lightJdbcConnectionPoolDataSourceTest.java | 2 + ...owFlightJdbcFlightStreamResultSetTest.java | 72 ++++++++----- .../ArrowFlightPreparedStatementTest.java | 100 ++++++++++++------ .../jdbc/ArrowFlightStatementExecuteTest.java | 92 ++++++---------- ...ArrowFlightStatementExecuteUpdateTest.java | 15 +-- .../arrow/driver/jdbc/ResultSetTest.java | 14 +-- .../jdbc/utils/MockFlightSqlProducer.java | 23 ++-- 8 files changed, 174 insertions(+), 151 deletions(-) diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java index 2c82999aec..e84b54157c 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java @@ -22,6 +22,7 @@ import java.sql.PreparedStatement; import java.sql.ResultSet; import java.sql.SQLException; +import java.sql.Statement; import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; @@ -39,6 +40,7 @@ import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.util.Preconditions; +import org.apache.arrow.util.VisibleForTesting; import org.apache.calcite.avatica.AvaticaConnection; import org.apache.calcite.avatica.AvaticaFactory; import org.apache.calcite.avatica.AvaticaStatement; @@ -223,6 +225,11 @@ synchronized void unregisterStatementOwner(final AvaticaStatement statement) { statementOwners.remove(statement); } + @VisibleForTesting + synchronized boolean isStatementOwner(final Statement statement) { + return statementOwners.contains(statement); + } + private synchronized ArrayList getStatementsToClose() { final ArrayList statements = new ArrayList<>(statementOwners); final Set ownedHandles = new HashSet<>(); diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcConnectionPoolDataSourceTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcConnectionPoolDataSourceTest.java index 14dda3d8a6..a2a789b697 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcConnectionPoolDataSourceTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcConnectionPoolDataSourceTest.java @@ -113,6 +113,8 @@ public void testShouldInnerConnectionShouldIgnoreDoubleClose() throws Exception assertFalse(connection.isClosed()); connection.close(); assertTrue(connection.isClosed()); + connection.close(); + assertTrue(connection.isClosed()); } @Test diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSetTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSetTest.java index b57b2f92b6..7c8f53ea3c 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSetTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSetTest.java @@ -16,51 +16,71 @@ */ package org.apache.arrow.driver.jdbc; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; import java.sql.SQLException; +import java.util.AbstractList; import java.util.List; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; import org.apache.arrow.driver.jdbc.client.CloseableEndpointStreamPair; import org.apache.arrow.driver.jdbc.utils.FlightEndpointDataQueue; +import org.apache.arrow.flight.FlightStream; import org.junit.jupiter.api.Test; public class ArrowFlightJdbcFlightStreamResultSetTest { @Test public void testClosedQueueDuringCollectionHandoffClosesUnacceptedEndpoints() throws Exception { - final FlightEndpointDataQueue queue = mock(FlightEndpointDataQueue.class); - final CloseableEndpointStreamPair acceptedEndpoint = mock(CloseableEndpointStreamPair.class); - final CloseableEndpointStreamPair rejectedEndpoint = mock(CloseableEndpointStreamPair.class); - final CloseableEndpointStreamPair unacceptedEndpoint = mock(CloseableEndpointStreamPair.class); + final FlightStream acceptedStream = mock(FlightStream.class); + final FlightStream rejectedStream = mock(FlightStream.class); + final FlightStream unacceptedStream = mock(FlightStream.class); + final List endpoints = + List.of( + new CloseableEndpointStreamPair(acceptedStream, null), + new CloseableEndpointStreamPair(rejectedStream, null), + new CloseableEndpointStreamPair(unacceptedStream, null)); final Exception closeFailure = new Exception("Failed to close rejected endpoint"); - doThrow(new IllegalStateException("FlightEndpointDataQueue closed")) - .when(queue) - .enqueue(rejectedEndpoint); - when(queue.isClosed()).thenReturn(true); - doThrow(closeFailure).when(rejectedEndpoint).close(); + doThrow(closeFailure).when(rejectedStream).close(); + final ExecutorService executor = Executors.newSingleThreadExecutor(); + try (final FlightEndpointDataQueue queue = FlightEndpointDataQueue.createNewQueue(executor)) { + final List closingEndpoints = + new AbstractList<>() { + @Override + public CloseableEndpointStreamPair get(final int index) { + if (index == 1) { + // Close between handing off the first and second endpoint. + assertDoesNotThrow(queue::close); + } + return endpoints.get(index); + } - final SQLException exception = - assertThrows( - SQLException.class, - () -> - ArrowFlightJdbcFlightStreamResultSet.enqueueEndpointData( - queue, List.of(acceptedEndpoint, rejectedEndpoint, unacceptedEndpoint))); + @Override + public int size() { + return endpoints.size(); + } + }; + final SQLException exception = + assertThrows( + SQLException.class, + () -> + ArrowFlightJdbcFlightStreamResultSet.enqueueEndpointData( + queue, closingEndpoints)); - assertEquals("Statement canceled", exception.getMessage()); - assertEquals(1, exception.getSuppressed().length); - assertSame(closeFailure, exception.getSuppressed()[0]); - verify(acceptedEndpoint, never()).close(); - verify(rejectedEndpoint).close(); - verify(unacceptedEndpoint).close(); - verify(queue).enqueue(acceptedEndpoint); - verify(queue).enqueue(rejectedEndpoint); - verify(queue, never()).enqueue(unacceptedEndpoint); + assertEquals("Statement canceled", exception.getMessage()); + assertEquals(1, exception.getSuppressed().length); + assertSame(closeFailure, exception.getSuppressed()[0]); + verify(acceptedStream).close(); + verify(rejectedStream).close(); + verify(unacceptedStream).close(); + } finally { + executor.shutdownNow(); + } } } 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 24bde53449..f3de2b67f3 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 @@ -22,12 +22,9 @@ 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.assertNotNull; import static org.junit.jupiter.api.Assertions.assertSame; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; -import static org.mockito.Mockito.doThrow; -import static org.mockito.Mockito.mock; import java.nio.charset.StandardCharsets; import java.sql.Connection; @@ -39,6 +36,7 @@ import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.TimeZone; import org.apache.arrow.driver.jdbc.utils.CoreMockedSqlProducers; import org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer; import org.apache.arrow.flight.FlightRuntimeException; @@ -56,6 +54,8 @@ import org.apache.arrow.vector.util.Text; import org.apache.calcite.avatica.AvaticaResultSet; import org.apache.calcite.avatica.AvaticaStatement; +import org.apache.calcite.avatica.Meta; +import org.apache.calcite.avatica.QueryState; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.BeforeEach; @@ -145,18 +145,22 @@ public void testSimpleQueryNoParameterBindingWithExecuteV2() throws SQLException } @Test - public void testPrepareStatementRegistersCreatedStatementByGeneratedHandle() throws SQLException { - final String query = CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD; - final ArrowFlightConnection flightConnection = (ArrowFlightConnection) connection; + public void testConnectionCloseClosesPreparedStatementAndResultSet() throws SQLException { + try (final Connection localConnection = FLIGHT_SERVER_TEST_EXTENSION.getConnection(false); + final PreparedStatement preparedStatement = + localConnection.prepareStatement(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD); + final ResultSet resultSet = preparedStatement.executeQuery()) { + assertTrue(resultSet.next()); - try (final PreparedStatement preparedStatement = connection.prepareStatement(query)) { - final ArrowFlightPreparedStatement arrowPreparedStatement = - (ArrowFlightPreparedStatement) preparedStatement; + localConnection.close(); - assertNotNull(flightConnection.statementMap.get(arrowPreparedStatement.handle.id)); - assertSame( - arrowPreparedStatement, - flightConnection.statementMap.get(arrowPreparedStatement.handle.id)); + assertTrue(preparedStatement.isClosed()); + assertTrue(resultSet.isClosed()); + assertEquals( + 1, + PRODUCER + .getActionTypeCounter() + .getOrDefault(FlightSqlUtils.FLIGHT_SQL_CLOSE_PREPARED_STATEMENT.getType(), 0)); } } @@ -172,32 +176,58 @@ public void testPrepareFailureIsReportedAsSQLException() { @Test public void testNullQueryIsReportedAsSQLException() { - final SQLException exception = - assertThrows(SQLException.class, () -> connection.prepareStatement(null)); - - assertThat(exception.getCause(), instanceOf(NullPointerException.class)); + assertThrows(SQLException.class, () -> connection.prepareStatement(null)); } @Test public void testCloseReleasesPreparedResourcesWhenResultSetCloseFails() throws Exception { - try (final ArrowFlightConnection localConnection = - (ArrowFlightConnection) FLIGHT_SERVER_TEST_EXTENSION.getConnection(false)) { - final ArrowFlightPreparedStatement preparedStatement = - (ArrowFlightPreparedStatement) - localConnection.prepareStatement(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD); - final AvaticaResultSet resultSet = mock(AvaticaResultSet.class); - final RuntimeException resultSetCloseFailure = - new RuntimeException("Result set close failed"); - doThrow(resultSetCloseFailure).when(resultSet).close(); - final java.lang.reflect.Field openResultSetField = - AvaticaStatement.class.getDeclaredField("openResultSet"); - openResultSetField.setAccessible(true); - openResultSetField.set(preparedStatement, resultSet); + final RuntimeException resultSetCloseFailure = new RuntimeException("Result set close failed"); + final ArrowFlightJdbcFactory factory = + new ArrowFlightJdbcFactory() { + @Override + public ArrowFlightJdbcVectorSchemaRootResultSet newResultSet( + final AvaticaStatement statement, + final QueryState state, + final Meta.Signature signature, + final TimeZone timeZone, + final Meta.Frame frame) + throws SQLException { + return new ArrowFlightJdbcVectorSchemaRootResultSet( + statement, + state, + signature, + newResultSetMetaData(statement, signature), + timeZone, + frame) { + @Override + protected AvaticaResultSet execute() { + populateData(VectorSchemaRoot.of()); + return this; + } + + @Override + public void close() { + super.close(); + throw resultSetCloseFailure; + } + }; + } + }; + try (final Connection localConnection = + factory.newConnection( + new ArrowFlightJdbcDriver(), + factory, + connection.getMetaData().getURL(), + connection.getClientInfo()); + final PreparedStatement preparedStatement = + localConnection.prepareStatement(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD)) { + final ResultSet resultSet = preparedStatement.executeQuery(); final SQLException exception = assertThrows(SQLException.class, preparedStatement::close); assertSame(resultSetCloseFailure, exception.getCause()); assertTrue(preparedStatement.isClosed()); + assertTrue(resultSet.isClosed()); assertEquals( 1, PRODUCER @@ -288,8 +318,9 @@ public void testQueryWithParameterBinding() throws SQLException { 2, connection.createArrayOf("INTEGER", new Integer[] {1, 2, null})); try (final ResultSet resultSet = preparedStatement.executeQuery()) { - resultSet.next(); - assert true; + assertTrue(resultSet.next()); + assertEquals(10, resultSet.getInt(1)); + assertFalse(resultSet.next()); } } } @@ -331,8 +362,9 @@ public void testQueryWithDateMillisecondParameterBinding() throws SQLException { preparedStatement.setDate(1, Date.valueOf(date)); try (final ResultSet resultSet = preparedStatement.executeQuery()) { - resultSet.next(); - assert true; + assertTrue(resultSet.next()); + assertEquals(10, resultSet.getInt(1)); + assertFalse(resultSet.next()); } } } 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 704905ea47..602ffdc6f5 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 @@ -21,14 +21,10 @@ import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.instanceOf; import static org.hamcrest.CoreMatchers.is; -import static org.hamcrest.CoreMatchers.not; import static org.hamcrest.CoreMatchers.nullValue; 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 static org.mockito.Mockito.doThrow; -import static org.mockito.Mockito.mock; import java.sql.Connection; import java.sql.ResultSet; @@ -51,7 +47,6 @@ import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; -import org.apache.calcite.avatica.AvaticaUtils; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeAll; @@ -175,20 +170,6 @@ public void testExecuteShouldRunSelectQuery() throws SQLException { is(allOf(equalTo(statement.getLargeUpdateCount()), equalTo(-1L)))); } - @Test - public void testExecuteReplacesStatementMapEntryWithPreparedStatement() throws SQLException { - final ArrowFlightStatement arrowStatement = (ArrowFlightStatement) statement; - final ArrowFlightConnection arrowConnection = (ArrowFlightConnection) connection; - - assertThat(statement.execute(SAMPLE_QUERY_CMD), is(true)); - - final Object preparedStatement = arrowConnection.statementMap.get(arrowStatement.handle.id); - - assertNotNull(preparedStatement); - assertSame(preparedStatement, arrowConnection.statementMap.get(arrowStatement.handle.id)); - assertThat(preparedStatement, instanceOf(ArrowFlightPreparedStatement.class)); - } - @Test public void testExecutePrepareFailurePreservesServerError() { final SQLException exception = @@ -201,8 +182,6 @@ public void testExecutePrepareFailurePreservesServerError() { @Test public void testConnectionCloseClosesStatementAndResultSetAfterExecute() throws SQLException { PRODUCER.clearActionTypeCounter(); - final ArrowFlightStatement arrowStatement = (ArrowFlightStatement) statement; - final ArrowFlightConnection arrowConnection = (ArrowFlightConnection) connection; assertThat(statement.execute(SAMPLE_QUERY_CMD), is(true)); final ResultSet resultSet = statement.getResultSet(); @@ -212,7 +191,6 @@ public void testConnectionCloseClosesStatementAndResultSetAfterExecute() throws assertThat(statement.isClosed(), is(true)); assertThat(resultSet.isClosed(), is(true)); - assertThat(arrowConnection.statementMap.containsKey(arrowStatement.handle.id), is(false)); assertThat( PRODUCER .getActionTypeCounter() @@ -221,60 +199,64 @@ public void testConnectionCloseClosesStatementAndResultSetAfterExecute() throws } @Test - public void testExecuteQueryRestoresStatementMapEntryWithStatement() throws SQLException { - final ArrowFlightStatement arrowStatement = (ArrowFlightStatement) statement; - final ArrowFlightConnection arrowConnection = (ArrowFlightConnection) connection; - + public void testExecuteQueryReplacesResultSetAndRemainsConnectionOwned() throws SQLException { assertThat(statement.execute(SAMPLE_QUERY_CMD), is(true)); + final ResultSet previousResultSet = statement.getResultSet(); try (ResultSet resultSet = statement.executeQuery(SAMPLE_QUERY_CMD)) { + assertThat(previousResultSet.isClosed(), is(true)); assertThat(resultSet.next(), is(true)); - } + assertThat(resultSet.getByte(1), is((byte) 0)); + + connection.close(); - assertSame(arrowStatement, arrowConnection.statementMap.get(arrowStatement.handle.id)); + assertThat(statement.isClosed(), is(true)); + assertThat(resultSet.isClosed(), is(true)); + } } @Test public void testExecuteQueryMapsPreparedStatementCloseFailureToSQLException() throws SQLException { - final ArrowFlightStatement arrowStatement = (ArrowFlightStatement) statement; - final ArrowFlightConnection arrowConnection = (ArrowFlightConnection) connection; - final ResultSet previousResultSet = statement.executeQuery(SAMPLE_QUERY_CMD); - final ArrowFlightPreparedStatement preparedStatement = mock(ArrowFlightPreparedStatement.class); - final FlightRuntimeException closeFailure = - CallStatus.INTERNAL.withDescription("Prepared statement close failed").toRuntimeException(); - doThrow(closeFailure).when(preparedStatement).closePreparedResources(); - arrowConnection.statementMap.put(arrowStatement.handle.id, preparedStatement); + assertThat(statement.execute(SAMPLE_QUERY_CMD), is(true)); + final ResultSet previousResultSet = statement.getResultSet(); + PRODUCER.setClosePreparedStatementFailure( + CallStatus.INTERNAL.withDescription("Prepared statement close failed")); try { final SQLException exception = assertThrows(SQLException.class, () -> statement.executeQuery(SAMPLE_QUERY_CMD)); - assertSame(closeFailure, exception.getCause()); + assertThat(exception.getCause(), instanceOf(FlightRuntimeException.class)); + final FlightRuntimeException cause = (FlightRuntimeException) exception.getCause(); + assertThat(cause.status().code(), is(CallStatus.INTERNAL.code())); + assertThat(cause.status().description(), containsString("Prepared statement close failed")); assertThat(previousResultSet.isClosed(), is(true)); } finally { - arrowConnection.statementMap.put(arrowStatement.handle.id, arrowStatement); + PRODUCER.setClosePreparedStatementFailure(null); } } @Test - public void testExecuteLargeUpdateMapsPreparedStatementCloseFailureToSQLException() { - final ArrowFlightStatement arrowStatement = (ArrowFlightStatement) statement; - final ArrowFlightConnection arrowConnection = (ArrowFlightConnection) connection; - final ArrowFlightPreparedStatement preparedStatement = mock(ArrowFlightPreparedStatement.class); - final FlightRuntimeException closeFailure = - CallStatus.INTERNAL.withDescription("Prepared statement close failed").toRuntimeException(); - doThrow(closeFailure).when(preparedStatement).closePreparedResources(); - arrowConnection.statementMap.put(arrowStatement.handle.id, preparedStatement); + public void testExecuteLargeUpdateMapsPreparedStatementCloseFailureToSQLException() + throws SQLException { + assertThat(statement.execute(SAMPLE_QUERY_CMD), is(true)); + final ResultSet previousResultSet = statement.getResultSet(); + PRODUCER.setClosePreparedStatementFailure( + CallStatus.INTERNAL.withDescription("Prepared statement close failed")); try { final SQLException exception = assertThrows( SQLException.class, () -> statement.executeLargeUpdate(SAMPLE_LARGE_UPDATE_QUERY)); - assertSame(closeFailure, exception.getCause()); + assertThat(exception.getCause(), instanceOf(FlightRuntimeException.class)); + final FlightRuntimeException cause = (FlightRuntimeException) exception.getCause(); + assertThat(cause.status().code(), is(CallStatus.INTERNAL.code())); + assertThat(cause.status().description(), containsString("Prepared statement close failed")); + assertThat(previousResultSet.isClosed(), is(true)); } finally { - arrowConnection.statementMap.put(arrowStatement.handle.id, arrowStatement); + PRODUCER.setClosePreparedStatementFailure(null); } } @@ -293,12 +275,7 @@ public void testExecuteShouldRunUpdateQueryForLargeUpdate() throws SQLException 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(updateCountSmall, is((long) Integer.MAX_VALUE)); assertThat(statement.getResultSet(), is(nullValue())); } @@ -339,12 +316,7 @@ public void testExecuteShouldRunUpdateQueryForLargeUpdateV2() throws SQLExceptio 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(updateCountSmall, is((long) Integer.MAX_VALUE)); assertThat(statement.getResultSet(), is(nullValue())); } } diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementExecuteUpdateTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementExecuteUpdateTest.java index 05e85227f0..23b660aada 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementExecuteUpdateTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementExecuteUpdateTest.java @@ -17,11 +17,8 @@ package org.apache.arrow.driver.jdbc; import static java.lang.String.format; -import static org.hamcrest.CoreMatchers.allOf; -import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.instanceOf; import static org.hamcrest.CoreMatchers.is; -import static org.hamcrest.CoreMatchers.not; import static org.hamcrest.CoreMatchers.startsWith; import static org.hamcrest.MatcherAssert.assertThat; import static org.junit.jupiter.api.Assertions.assertThrows; @@ -39,7 +36,6 @@ import org.apache.arrow.vector.types.Types.MinorType; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; -import org.apache.calcite.avatica.AvaticaUtils; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeAll; @@ -116,16 +112,7 @@ public void testExecuteUpdateShouldReturnNumColsAffectedForNumRowsFittingInt() public void testExecuteUpdateShouldReturnSaturatedNumColsAffectedIfDoesNotFitInInt() throws SQLException { final long result = statement.executeUpdate(LARGE_UPDATE_SAMPLE_QUERY); - final long expectedRowCountRaw = LARGE_UPDATE_SAMPLE_QUERY_AFFECTED_COLS; - assertThat( - result, - is( - allOf( - not(equalTo(expectedRowCountRaw)), - equalTo( - (long) - AvaticaUtils.toSaturatedInt( - expectedRowCountRaw))))); // Because of long-to-integer overflow. + assertThat(result, is((long) Integer.MAX_VALUE)); } @Test diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java index 183182ee1b..7d1bf67e45 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ResultSetTest.java @@ -242,11 +242,12 @@ public void testCloseOnCompletionUnregistersDirectStatementWhenResultSetClosedEa final ResultSet resultSet = statement.executeQuery(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD); statement.closeOnCompletion(); + assertThat(((ArrowFlightConnection) connection).isStatementOwner(statement), is(true)); resultSet.close(); assertThat(statement.isClosed(), is(true)); - assertThat(isStatementOwner(statement), is(false)); + assertThat(((ArrowFlightConnection) connection).isStatementOwner(statement), is(false)); } } @@ -257,22 +258,15 @@ public void testCloseOnCompletionUnregistersPreparedStatementWhenResultSetClosed connection.prepareStatement(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD)) { final ResultSet resultSet = statement.executeQuery(); statement.closeOnCompletion(); + assertThat(((ArrowFlightConnection) connection).isStatementOwner(statement), is(true)); resultSet.close(); assertThat(statement.isClosed(), is(true)); - assertThat(isStatementOwner(statement), is(false)); + assertThat(((ArrowFlightConnection) connection).isStatementOwner(statement), is(false)); } } - private static boolean isStatementOwner(final Statement statement) - throws ReflectiveOperationException { - final java.lang.reflect.Field statementOwnersField = - ArrowFlightConnection.class.getDeclaredField("statementOwners"); - statementOwnersField.setAccessible(true); - return ((Set) statementOwnersField.get(connection)).contains(statement); - } - /** * Tests whether the {@link ArrowFlightJdbcDriver} close the statement after complete ResultSet * with max rows limit when call {@link 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 88355acdfb..36c07ce5dd 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 @@ -39,6 +39,7 @@ import java.util.Map.Entry; import java.util.Objects; import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; import java.util.function.BiConsumer; import java.util.function.Consumer; import java.util.stream.IntStream; @@ -104,8 +105,9 @@ public final class MockFlightSqlProducer implements FlightSqlProducer { private final Map>> expectedParameterValues = new HashMap<>(); private final Map isUpdateMap = new HashMap<>(); - private final Map actionTypeCounter = new HashMap<>(); - private final Map commandTypeCounter = new HashMap<>(); + private final Map actionTypeCounter = new ConcurrentHashMap<>(); + private final Map commandTypeCounter = new ConcurrentHashMap<>(); + private volatile @Nullable CallStatus closePreparedStatementFailure; public static final String COMMAND_STATEMENT_QUERY = "statement_query"; public static final String COMMAND_STATEMENT_UPDATE = "statement_update"; @@ -308,8 +310,16 @@ public void closePreparedStatement( final ActionClosePreparedStatementRequest actionClosePreparedStatementRequest, final CallContext callContext, final StreamListener streamListener) { - // TODO Implement this method. - streamListener.onCompleted(); + final CallStatus failure = closePreparedStatementFailure; + if (failure == null) { + streamListener.onCompleted(); + } else { + streamListener.onError(failure.toRuntimeException()); + } + } + + public void setClosePreparedStatementFailure(final @Nullable CallStatus failure) { + closePreparedStatementFailure = failure; } @Override @@ -692,9 +702,8 @@ public void listFlights( @Override public void doAction(CallContext context, Action action, StreamListener listener) { + actionTypeCounter.merge(action.getType(), 1, Integer::sum); FlightSqlProducer.super.doAction(context, action, listener); - actionTypeCounter.put( - action.getType(), actionTypeCounter.getOrDefault(action.getType(), 0) + 1); } /** @@ -718,7 +727,7 @@ public Map getCommandTypeCounter() { } private void incrementCommandTypeCounter(String commandType) { - commandTypeCounter.put(commandType, commandTypeCounter.getOrDefault(commandType, 0) + 1); + commandTypeCounter.merge(commandType, 1, Integer::sum); } private void getStreamCatalogFunctions( From 4d1556fd1ef414e81503278e4682bb612fe2e3dd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Tue, 10 Nov 2026 15:21:10 +0000 Subject: [PATCH 25/25] Fix JDBC cancellation visibility and owner cleanup --- .../driver/jdbc/ArrowFlightConnection.java | 1 - .../ArrowFlightJdbcFlightStreamResultSet.java | 10 +-- .../ArrowFlightStatementProtocolTest.java | 74 +++++++++++++++++++ 3 files changed, 79 insertions(+), 6 deletions(-) diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java index e84b54157c..15f722735c 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java @@ -161,7 +161,6 @@ void reset() throws SQLException { } statementMap.clear(); - statementOwners.clear(); // Reset Holdability this.setHoldability(this.metaData.getResultSetHoldability()); 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 a9537eecc1..2ab807e78f 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 @@ -54,6 +54,7 @@ public final class ArrowFlightJdbcFlightStreamResultSet private final FlightInfo flightInfo; private CloseableEndpointStreamPair currentEndpointData; private FlightEndpointDataQueue flightEndpointDataQueue; + private volatile boolean cancelled; private VectorSchemaRootTransformer transformer; private VectorSchemaRoot currentVectorSchemaRoot; @@ -213,7 +214,7 @@ public boolean next() throws SQLException { } catch (final FlightRuntimeException e) { // A concurrent statement.cancel() (or close) cancels in-flight FlightStreams, // which surface here as CANCELLED. Normalize to Avatica's "Statement canceled". - if (flightEndpointDataQueue.isClosed() + if ((cancelled || flightEndpointDataQueue.isClosed()) && e.status().code() == CallStatus.CANCELLED.code()) { throw AvaticaConnection.HELPER.createException("Statement canceled"); } @@ -225,10 +226,8 @@ public boolean next() throws SQLException { continue; } - // No more data. If the queue was closed concurrently (e.g. statement.cancel() - // racing with the reader past super.next()), surface as "Statement canceled" - // to match Avatica's cancellation semantics. - if (flightEndpointDataQueue.isClosed()) { + // Cancellation can wake the reader before queue cleanup completes. + if (cancelled || flightEndpointDataQueue.isClosed()) { throw AvaticaConnection.HELPER.createException("Statement canceled"); } @@ -242,6 +241,7 @@ public boolean next() throws SQLException { @Override protected void cancel() { + cancelled = true; super.cancel(); final CloseableEndpointStreamPair currentEndpoint = this.currentEndpointData; if (currentEndpoint != null) { diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementProtocolTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementProtocolTest.java index 0cc533a909..58a607067a 100644 --- a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementProtocolTest.java +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementProtocolTest.java @@ -18,6 +18,9 @@ import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.MatcherAssert.assertThat; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertTrue; import com.google.protobuf.Message; @@ -28,6 +31,9 @@ import java.sql.SQLException; import java.sql.Statement; import java.util.Collections; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.Consumer; import org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer; import org.apache.arrow.flight.FlightProducer.ServerStreamListener; @@ -126,6 +132,74 @@ public static void tearDownAfterClass() throws Exception { AutoCloseables.close(PRODUCER); } + @Test + public void testCancelWhileNextIsBlockedReportsSQLException() throws Exception { + final String query = "SELECT * FROM BLOCKED_CANCELLATION"; + final int endpointCount = 8; + final CountDownLatch streamsStarted = new CountDownLatch(endpointCount); + final CountDownLatch releaseStreams = new CountDownLatch(1); + PRODUCER.addSelectQuery( + query, + QUERY_SCHEMA, + Collections.nCopies( + endpointCount, + listener -> { + try (final BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + final VectorSchemaRoot root = VectorSchemaRoot.create(QUERY_SCHEMA, allocator)) { + ((IntVector) root.getVector("id")).setSafe(0, 1); + root.setRowCount(1); + listener.start(root); + listener.putNext(); + streamsStarted.countDown(); + assertTrue(releaseStreams.await(10, TimeUnit.SECONDS), "Stream was not released"); + } catch (final Throwable throwable) { + listener.error(throwable); + } finally { + listener.completed(); + } + })); + + try (Statement statement = connection.createStatement(); + ResultSet resultSet = statement.executeQuery(query)) { + assertTrue(streamsStarted.await(5, TimeUnit.SECONDS), "Streams did not start"); + assertTrue(resultSet.next()); + final AtomicReference readerFailure = new AtomicReference<>(); + final Thread reader = + new Thread( + () -> { + try { + resultSet.next(); + } catch (final Throwable throwable) { + readerFailure.set(throwable); + } + }, + "blocked-result-set-reader"); + reader.start(); + try { + final long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + while (reader.isAlive() + && reader.getState() != Thread.State.WAITING + && System.nanoTime() < deadline) { + Thread.sleep(1); + } + assertEquals(Thread.State.WAITING, reader.getState(), "Reader did not block"); + + statement.cancel(); + reader.join(5000); + + assertFalse(reader.isAlive(), "Reader did not observe cancellation"); + final SQLException exception = assertInstanceOf(SQLException.class, readerFailure.get()); + assertEquals("Statement canceled", exception.getMessage()); + } finally { + releaseStreams.countDown(); + reader.interrupt(); + reader.join(5000); + } + } finally { + releaseStreams.countDown(); + } + } + @Test public void testStatementExecuteQueryUsesStatementProtocol() throws SQLException { try (Statement statement = connection.createStatement();