diff --git a/java/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/example/FlightSqlStatelessExample.java b/java/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/example/FlightSqlStatelessExample.java index 3c522528512ac..8cd5d622f7381 100644 --- a/java/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/example/FlightSqlStatelessExample.java +++ b/java/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/example/FlightSqlStatelessExample.java @@ -90,10 +90,8 @@ public Runnable acceptPutPreparedStatementQuery(CommandPreparedStatementQuery co FlightStream flightStream, StreamListener ackStream) { return () -> { - try { - final String query = new String(command.getPreparedStatementHandle().toStringUtf8()); - final PreparedStatement preparedStatement = createPreparedStatement(query); - + final String query = new String(command.getPreparedStatementHandle().toStringUtf8()); + try (PreparedStatement preparedStatement = createPreparedStatement(query)) { while (flightStream.next()) { final VectorSchemaRoot root = flightStream.getRoot(); final JdbcParameterBinder binder = JdbcParameterBinder.builder(preparedStatement, root).bindAll().build(); @@ -141,16 +139,16 @@ public Runnable acceptPutPreparedStatementQuery(CommandPreparedStatementQuery co @Override public void getStreamPreparedStatement(final CommandPreparedStatementQuery command, final CallContext context, final ServerStreamListener listener) { + final byte[] handle = command.getPreparedStatementHandle().toByteArray(); try { // Case where there are parameters - final byte[] handle = command.getPreparedStatementHandle().toByteArray(); try { final DoPutPreparedStatementResultPOJO doPutPreparedStatementResultPOJO = deserializePOJO(handle); final String query = doPutPreparedStatementResultPOJO.getQuery(); - final PreparedStatement statement = createPreparedStatement(query); - try (ArrowFileReader reader = new ArrowFileReader(new SeekableReadChannel( + try (PreparedStatement statement = createPreparedStatement(query); + ArrowFileReader reader = new ArrowFileReader(new SeekableReadChannel( new ByteArrayReadableSeekableByteChannel( doPutPreparedStatementResultPOJO.getParameters())), rootAllocator)) { @@ -168,8 +166,9 @@ public void getStreamPreparedStatement(final CommandPreparedStatementQuery comma } catch (StreamCorruptedException e) { // Case where there are no parameters final String query = new String(command.getPreparedStatementHandle().toStringUtf8()); - final PreparedStatement preparedStatement = createPreparedStatement(query); - executeQuery(preparedStatement, listener); + try (PreparedStatement preparedStatement = createPreparedStatement(query)) { + executeQuery(preparedStatement, listener); + } } } catch (final SQLException | IOException | ClassNotFoundException e) { LOGGER.error(format("Failed to getStreamPreparedStatement: <%s>.", e.getMessage()), e); @@ -207,24 +206,21 @@ private void executeQuery(PreparedStatement statement, public FlightInfo getFlightInfoPreparedStatement(final CommandPreparedStatementQuery command, final CallContext context, final FlightDescriptor descriptor) { - String query; + final byte[] handle = command.getPreparedStatementHandle().toByteArray(); try { - final byte[] handle = command.getPreparedStatementHandle().toByteArray(); - final DoPutPreparedStatementResultPOJO doPutPreparedStatementResultPOJO = null; + String query; try { query = deserializePOJO(handle).getQuery(); } catch (StreamCorruptedException e) { query = new String(command.getPreparedStatementHandle().toStringUtf8()); } - final PreparedStatement statement = createPreparedStatement(query); - - ResultSetMetaData metaData = statement.getMetaData(); - return getFlightInfoForSchema(command, descriptor, - jdbcToArrowSchema(metaData, DEFAULT_CALENDAR)); + try (PreparedStatement statement = createPreparedStatement(query)) { + ResultSetMetaData metaData = statement.getMetaData(); + return getFlightInfoForSchema(command, descriptor, + jdbcToArrowSchema(metaData, DEFAULT_CALENDAR)); + } } catch (final SQLException | IOException | ClassNotFoundException e) { - LOGGER.error( - format("There was a problem executing the prepared statement: <%s>.", e.getMessage()), - e); + LOGGER.error(format("There was a problem executing the prepared statement: <%s>.", e.getMessage()), e); throw CallStatus.INTERNAL.withCause(e).toRuntimeException(); } } diff --git a/java/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSqlStateless.java b/java/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSqlStateless.java index 77f140da0e32f..264caf079658f 100644 --- a/java/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSqlStateless.java +++ b/java/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSqlStateless.java @@ -81,16 +81,16 @@ public void testSimplePreparedStatementResultsWithParameterBinding() throws Exce prepare.setParameters(insertRoot); final FlightInfo flightInfo = prepare.execute(); - final FlightStream stream = sqlClient.getStream(flightInfo - .getEndpoints() - .get(0).getTicket()); - - // TODO: root is null and getSchema hangs when run as complete suite. - // This works when run as an individual test. - Assertions.assertAll( - () -> MatcherAssert.assertThat(stream.getSchema(), is(SCHEMA_INT_TABLE)), - () -> MatcherAssert.assertThat(getResults(stream), is(EXPECTED_RESULTS_FOR_PARAMETER_BINDING)) - ); + try (FlightStream stream = sqlClient.getStream(flightInfo.getEndpoints().get(0).getTicket())) { + // TODO: root is null and getSchema hangs when run as complete suite. + // This works when run as an individual test. + final VectorSchemaRoot root = stream.getRoot(); + final Schema schema = root.getSchema(); + Assertions.assertAll( + () -> MatcherAssert.assertThat(schema, is(SCHEMA_INT_TABLE)), + () -> MatcherAssert.assertThat(getResults(stream), is(EXPECTED_RESULTS_FOR_PARAMETER_BINDING)) + ); + } } } }