From bb54fad0be8a0179179a7ff4c41e6227fd7462ea Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A9lder=20Greg=C3=B3rio?= Date: Thu, 24 Sep 2026 08:46:05 +0100 Subject: [PATCH] GH-1090: Reconcile rebased PR scope --- .../driver/jdbc/ArrowFlightInfoStatement.java | 36 +++++++++++++++++++ .../ArrowFlightJdbcFlightStreamResultSet.java | 36 +++++-------------- .../driver/jdbc/ArrowFlightMetaStatement.java | 3 +- .../jdbc/utils/FlightEndpointDataQueue.java | 15 ++------ .../utils/FlightEndpointDataQueueTest.java | 13 ++++--- 5 files changed, 54 insertions(+), 49 deletions(-) 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/ArrowFlightJdbcFlightStreamResultSet.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightJdbcFlightStreamResultSet.java index c39a6c6b6b..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 @@ -27,14 +27,11 @@ 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; 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; @@ -195,40 +192,23 @@ public boolean next() throws SQLException { return true; } - try { - if (currentEndpointData != null) { - currentEndpointData.getStream().getRoot().clear(); - if (currentEndpointData.getStream().next()) { - populateDataForCurrentFlightStream(); - continue; - } - - flightEndpointDataQueue.enqueue(currentEndpointData); + if (currentEndpointData != null) { + currentEndpointData.getStream().getRoot().clear(); + if (currentEndpointData.getStream().next()) { + populateDataForCurrentFlightStream(); + continue; } - 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; + flightEndpointDataQueue.enqueue(currentEndpointData); } + currentEndpointData = getNextEndpointStream(false); + if (currentEndpointData != null) { populateDataForCurrentFlightStream(); 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/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/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..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,8 +97,7 @@ 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. + checkOpen(); while (!futures.isEmpty()) { final Future future = endpointStreamSupplier.get(); futures.remove(future); @@ -177,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/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