jbonofre commented on code in PR #1090:
URL: https://github.com/apache/arrow-java/pull/1090#discussion_r4523606700
##########
flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java:
##########
@@ -16,66 +16,263 @@
*/
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.flight.FlightRuntimeException;
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;
+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 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 newPreparedStatement(
- final ArrowFlightConnection connection,
- final ArrowFlightSqlClientHandler.PreparedStatement preparedStmt,
- final StatementHandle statementHandle,
- final Signature signature,
- final int resultSetType,
- final int resultSetConcurrency,
- final int resultSetHoldability)
- throws SQLException {
- return new ArrowFlightPreparedStatement(
- connection,
- preparedStmt,
- statementHandle,
- signature,
- resultSetType,
- resultSetConcurrency,
- resultSetHoldability);
+ static Builder builder(final ArrowFlightConnection connection) {
+ return new Builder(connection);
}
@Override
public ArrowFlightConnection getConnection() throws SQLException {
return (ArrowFlightConnection) super.getConnection();
}
+ ExecuteResult prepareAndExecuteInternal(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));
+ }
+
+ ExecuteResult prepareAndExecuteInternal(
+ final String query,
+ final long maxRowCount,
+ final int maxRowsInFirstFrame,
+ final PrepareCallback callback)
+ throws SQLException {
+
+ return ArrowFlightPreparedStatement.builder(getConnection())
Review Comment:
Calling `Statement.execute(String)` again on the same statement leaks the
previous server-side prepared statement.
On the second call, `statementMap` holds the internal
`ArrowFlightPreparedStatement` created by the first call, so
`ArrowFlightMetaImpl.prepareAndExecute` routes here. We build a new prepared
statement under the same handle, and its constructor replaces the
`statementMap` entry, but `this.preparedStatement` is never closed. No
`ClosePreparedStatement` is sent for it. A loop of N `execute()` calls leaves
N-1 prepared statements open on the server, and `stmt.close()` only closes the
last one.
The old `statementHandlePreparedStatementMap.put(...)` dropped the previous
entry the same way, so this is not a regression. Since this method is new,
could we close the previous one once the replacement is built?
```java
final ArrowFlightPreparedStatement replacement =
ArrowFlightPreparedStatement.builder(getConnection())
.withQuery(query)
.withExistingStatement(this)
.build();
closePreparedResources();
return replacement.prepareAndExecuteInternal(callback);
```
A test that calls `execute()` a few times on one statement and checks how
many prepared statements the mock producer closed would cover this.
##########
flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSet.java:
##########
@@ -192,23 +198,40 @@ 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;
+ }
+
+ enqueueEndpointData(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()
Review Comment:
This check can miss a cancel. `cancel()` cancels the current stream first
and only then calls `flightEndpointDataQueue.close()`.
`FlightEndpointDataQueue.close()` sets `closed` in its `finally` block, after
it has cancelled the other streams, waited on the futures and closed the
endpoints.
A reader blocked in `currentEndpointData.getStream().next()` is woken by the
stream cancel and can reach this line while `isClosed()` is still `false`. It
then rethrows the raw `FlightRuntimeException` instead of `Statement canceled`.
The `isClosed()` check at line 231 has the same race.
`AvaticaResultSet.cancel()` sets the statement cancel flag first, but that
field is protected on `AvaticaStatement`, so it is not reachable from here. A
simple option is a `volatile boolean cancelled` field that `cancel()` sets
before it cancels the stream, checked here and at line 231.
It would also be good to add a test where one thread is blocked in `next()`
while another calls `statement.cancel()`. The current test only covers the
enqueue path.
##########
flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java:
##########
@@ -16,66 +16,263 @@
*/
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.flight.FlightRuntimeException;
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;
+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 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 newPreparedStatement(
- final ArrowFlightConnection connection,
- final ArrowFlightSqlClientHandler.PreparedStatement preparedStmt,
- final StatementHandle statementHandle,
- final Signature signature,
- final int resultSetType,
- final int resultSetConcurrency,
- final int resultSetHoldability)
- throws SQLException {
- return new ArrowFlightPreparedStatement(
- connection,
- preparedStmt,
- statementHandle,
- signature,
- resultSetType,
- resultSetConcurrency,
- resultSetHoldability);
+ static Builder builder(final ArrowFlightConnection connection) {
+ return new Builder(connection);
}
@Override
public ArrowFlightConnection getConnection() throws SQLException {
return (ArrowFlightConnection) super.getConnection();
}
+ ExecuteResult prepareAndExecuteInternal(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));
+ }
+
+ ExecuteResult prepareAndExecuteInternal(
+ final String query,
+ final long maxRowCount,
+ final int maxRowsInFirstFrame,
+ final PrepareCallback callback)
+ throws SQLException {
+
+ return ArrowFlightPreparedStatement.builder(getConnection())
+ .withQuery(query)
+ .withExistingStatement(this)
+ .build()
+ .prepareAndExecuteInternal(callback);
+ }
+
+ Schema getDataSetSchema() {
+ ensurePrepared();
+ return preparedStatement.getDataSetSchema();
+ }
+
@Override
- public synchronized void close() throws SQLException {
- this.preparedStatement.close();
- super.close();
+ protected void close_() {
+ RuntimeException closeFailure = null;
+ try {
+ super.close_();
+ } 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() {
+ final ArrowFlightSqlClientHandler.PreparedStatement statement =
preparedStatement;
+ preparedStatement = null;
+ if (statement != null) {
+ statement.close();
+ }
+ }
+
+ ExecuteResult executeWithTypedValues(
+ final StatementHandle statementHandle,
+ final List<TypedValue> 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<List<TypedValue>>
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.");
+ }
+ }
+
+ 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;
+ try {
+ preparedStatement = connection.getClientHandler().prepare(query);
+ } catch (final FlightRuntimeException e) {
+ throw new SQLException(e.getMessage(), e);
+ }
+ final Signature signature =
Review Comment:
If anything throws between `prepare()` and the end of `build()`, the server
side prepared statement is never cloud. One concrete case: a parameter schema
with a `ListView`, `LargeListView` or `RunEndEncoded` field makes
`ConvertUtils.convertArrowFieldsToAvaticaParameters` throw
`UnsupportedOperationException` from `buildSignature`.
The old `prepareForHandle` had the same gap. Now that `build()` owns the
whole lifecycle, could we close it on failure?
```java
final ArrowFlightPreparedStatement statement;
try {
final Signature signature =
ArrowFlightMetaImpl.buildSignature(
query,
preparedStatement.getDataSetSchema(),
preparedStatement.getParameterSchema(),
preparedStatement.isUpdate());
statement =
new ArrowFlightPreparedStatement(
connection,
generateHandle ? null : handle,
signature,
preparedStatement,
resultSetType,
resultSetConcurrency,
resultSetHoldability);
} catch (final Exception e) {
try {
preparedStatement.close();
} catch (final Exception closeException) {
e.addSuppressed(closeException);
}
throw e;
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
