lidavidm commented on code in PR #1090:
URL: https://github.com/apache/arrow-java/pull/1090#discussion_r4463124919
##########
flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java:
##########
@@ -45,7 +52,10 @@ 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.
Review Comment:
Comment in wrong place
##########
flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java:
##########
@@ -45,7 +52,10 @@ 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<Integer, ArrowFlightJdbcFlightStreamResultSet>
metadataResultSetMap = new HashMap<>();
+ private final Set<AvaticaStatement> statementOwners =
Review Comment:
I wonder, could we separate this change out into its own PR so it shows up
more clearly in the changelog/history and is easier to review?
##########
flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementApiTest.java:
##########
Review Comment:
Um, is this really necessary? Smells like overzealous AI?
##########
flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementExecuteTest.java:
##########
@@ -142,6 +152,13 @@ public static void tearDownAfterClass() throws Exception {
AutoCloseables.close(PRODUCER);
}
+ @Test
+ public void testStatementKeepsFlightInfoStatementApi() {
+ final ArrowFlightInfoStatement flightInfoStatement =
(ArrowFlightInfoStatement) statement;
+
+ assertSame(statement, flightInfoStatement);
+ }
+
Review Comment:
What is the point of this test...?
##########
flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java:
##########
@@ -79,6 +88,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);
+ }
+ }
+
Review Comment:
Same here...
##########
flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java:
##########
@@ -257,4 +290,39 @@ BufferAllocator getBufferAllocator() {
public ArrowFlightMetaImpl getMeta() {
return (ArrowFlightMetaImpl) this.meta;
}
+
+ @Override
+ public PreparedStatement prepareStatement(final String sql) throws
SQLException {
+ checkOpen();
Review Comment:
(nit, but isn't prepareStatement itself about to checkOpen?)
##########
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);
Review Comment:
If we _really_ need to test this, maybe making the property package-private
and marking it with VisibleForTesting is better?
##########
flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java:
##########
@@ -45,7 +52,10 @@ 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<Integer, ArrowFlightJdbcFlightStreamResultSet>
metadataResultSetMap = new HashMap<>();
+ private final Set<AvaticaStatement> statementOwners =
Review Comment:
Does this perhaps need to be concurrency-safe?
##########
flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementExecuteTest.java:
##########
@@ -165,6 +182,109 @@ 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 =
+ assertThrows(
+ SQLException.class, () -> statement.execute("SELECT * FROM
unregistered_table"));
+
+ assertThat(exception.getMessage(), containsString("Query not found"));
+ }
+
+ @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;
+ 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
+ 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);
+
+ try {
+ final SQLException exception =
+ 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);
Review Comment:
For the stated purpose of these tests, why do we need to manipulate the
internal state of the connection?
##########
flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightConnection.java:
##########
@@ -45,7 +52,10 @@ 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<Integer, ArrowFlightJdbcFlightStreamResultSet>
metadataResultSetMap = new HashMap<>();
+ private final Set<AvaticaStatement> statementOwners =
Review Comment:
Albeit if it can't be separated then that's fine
--
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]
