diff --git a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java index 9f36c90249..4790b61749 100644 --- a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java +++ b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlConnection.java @@ -121,7 +121,7 @@ public void commit() throws AdbcException { @Override public AdbcStatement createStatement() throws AdbcException { - return new FlightSqlStatement(allocator, client, clientCache, quirks); + return new FlightSqlStatement(allocator, client, clientCache, quirks, callOptions); } @Override @@ -152,7 +152,7 @@ public ArrowReader readPartition(ByteBuffer descriptor) throws AdbcException { public AdbcStatement bulkIngest(String targetTableName, BulkIngestMode mode) throws AdbcException { return FlightSqlStatement.ingestRoot( - allocator, client, clientCache, quirks, targetTableName, mode); + allocator, client, clientCache, quirks, targetTableName, mode, callOptions); } @Override diff --git a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlStatement.java b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlStatement.java index dd3ba1f286..557353c74a 100644 --- a/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlStatement.java +++ b/java/driver/flight-sql/src/main/java/org/apache/arrow/adbc/driver/flightsql/FlightSqlStatement.java @@ -29,6 +29,7 @@ import org.apache.arrow.adbc.core.BulkIngestMode; import org.apache.arrow.adbc.core.PartitionDescriptor; import org.apache.arrow.adbc.sql.SqlQuirks; +import org.apache.arrow.flight.CallOption; import org.apache.arrow.flight.FlightEndpoint; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightRuntimeException; @@ -36,7 +37,6 @@ import org.apache.arrow.flight.impl.Flight; import org.apache.arrow.flight.sql.FlightSqlClient; import org.apache.arrow.memory.BufferAllocator; -import org.apache.arrow.util.AutoCloseables; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; @@ -47,6 +47,7 @@ public class FlightSqlStatement implements AdbcStatement { private final FlightSqlClientWithCallOptions client; private final LoadingCache clientCache; private final SqlQuirks quirks; + private final CallOption[] connectionOptions; // State for SQL queries private @Nullable String sqlQuery; @@ -59,7 +60,8 @@ public class FlightSqlStatement implements AdbcStatement { BufferAllocator allocator, FlightSqlClientWithCallOptions client, LoadingCache clientCache, - SqlQuirks quirks) { + SqlQuirks quirks, + CallOption... connectionOptions) { this.allocator = allocator; this.client = client; this.clientCache = clientCache; @@ -68,6 +70,7 @@ public class FlightSqlStatement implements AdbcStatement { this.preparedStatement = null; this.bulkOperation = null; this.bindRoot = null; + this.connectionOptions = connectionOptions; } static FlightSqlStatement ingestRoot( @@ -76,10 +79,11 @@ static FlightSqlStatement ingestRoot( LoadingCache clientCache, SqlQuirks quirks, String targetTableName, - BulkIngestMode mode) { + BulkIngestMode mode, + CallOption... connectionOptions) { Objects.requireNonNull(targetTableName); final FlightSqlStatement statement = - new FlightSqlStatement(allocator, client, clientCache, quirks); + new FlightSqlStatement(allocator, client, clientCache, quirks, connectionOptions); statement.bulkOperation = new BulkState(mode, targetTableName); return statement; } @@ -170,7 +174,7 @@ private UpdateResult executeBulk(BulkState bulkOperation) throws AdbcException { statement.setParameters(new NonOwningRoot(bindParams)); client.executePreparedUpdate(statement); } finally { - statement.close(); + statement.close(connectionOptions); } } catch (FlightRuntimeException e) { // XXX: FlightSqlClient.executeUpdate does some extra wrapping that we need to undo @@ -313,7 +317,7 @@ public void close() throws AdbcException { // TODO(https://github.com/apache/arrow/issues/39814): this is annotated wrongly upstream if (preparedStatement != null) { try { - AutoCloseables.close(preparedStatement); + preparedStatement.close(connectionOptions); } catch (Exception e) { throw AdbcException.internal("[Flight SQL] Could not close prepared statement") .withCause(e); diff --git a/java/driver/flight-sql/src/test/java/org/apache/arrow/adbc/driver/flightsql/HeaderTest.java b/java/driver/flight-sql/src/test/java/org/apache/arrow/adbc/driver/flightsql/HeaderTest.java index 261d172e83..b02fbfc614 100644 --- a/java/driver/flight-sql/src/test/java/org/apache/arrow/adbc/driver/flightsql/HeaderTest.java +++ b/java/driver/flight-sql/src/test/java/org/apache/arrow/adbc/driver/flightsql/HeaderTest.java @@ -27,6 +27,7 @@ import org.apache.arrow.adbc.core.AdbcDriver; import org.apache.arrow.adbc.core.AdbcException; import org.apache.arrow.adbc.core.AdbcInfoCode; +import org.apache.arrow.adbc.core.AdbcStatement; import org.apache.arrow.adbc.core.AdbcStatusCode; import org.apache.arrow.adbc.drivermanager.AdbcDriverManager; import org.apache.arrow.driver.jdbc.utils.MockFlightSqlProducer; @@ -55,16 +56,18 @@ public class HeaderTest { private AdbcConnection connection; private BufferAllocator allocator; private HeaderValidator.Factory headerValidatorFactory; + private MockFlightSqlProducer producer; @BeforeEach public void setUp() { allocator = new RootAllocator(Long.MAX_VALUE); headerValidatorFactory = new HeaderValidator.Factory(); + producer = new MockFlightSqlProducer(); builder = FlightServer.builder() .middleware(HeaderValidator.KEY, headerValidatorFactory) .location(Location.forGrpcInsecure("localhost", 0)) - .producer(new MockFlightSqlProducer()); + .producer(producer); params = new HashMap<>(); } @@ -89,6 +92,40 @@ public void testArbitraryHeader() throws Exception { assertEquals(dummyValue, headers.get(dummyHeaderName)); } + /** + * The connection's call options (which carry arbitrary headers, auth, cookies, etc.) must ride on + * every RPC the driver issues on that connection's behalf, including the ClosePreparedStatement + * call made when a prepared statement is closed. This regression test guards against dropping the + * connection options on close. + */ + @Test + public void testHeaderSentWhenClosingPreparedStatement() throws Exception { + final String dummyValue = "dummy"; + final String dummyHeaderName = "test-header"; + final String query = "UPDATE the_table SET x = 1"; + params.put(FlightSqlConnectionProperties.RPC_CALL_HEADER_PREFIX + dummyHeaderName, dummyValue); + producer.addUpdateQuery(query, /*updatedRows=*/ 1L); + server = builder.build(); + server.start(); + connect(); + + // Prepare and then close a statement. Closing triggers a ClosePreparedStatement RPC, which must + // still carry the connection header. Before the fix, close() dropped the call options. + try (AdbcStatement statement = connection.createStatement()) { + statement.setSqlQuery(query); + statement.prepare(); + } + + // The connection header must appear on every RPC, including the final ClosePreparedStatement. + final int requestCount = headerValidatorFactory.getRequestCount(); + assertTrue(requestCount > 0); + for (int i = 0; i < requestCount; i++) { + CallHeaders headers = headerValidatorFactory.getHeadersReceivedAtRequest(i); + assertEquals( + dummyValue, headers.get(dummyHeaderName), "connection header missing on RPC #" + i); + } + } + @Test public void testCookies() throws Exception { builder.middleware(CookieMiddleware.KEY, new CookieMiddleware.Factory()); diff --git a/java/driver/flight-sql/src/test/java/org/apache/arrow/adbc/driver/flightsql/HeaderValidator.java b/java/driver/flight-sql/src/test/java/org/apache/arrow/adbc/driver/flightsql/HeaderValidator.java index f543dd9958..c4702eeed4 100644 --- a/java/driver/flight-sql/src/test/java/org/apache/arrow/adbc/driver/flightsql/HeaderValidator.java +++ b/java/driver/flight-sql/src/test/java/org/apache/arrow/adbc/driver/flightsql/HeaderValidator.java @@ -52,6 +52,10 @@ public CallHeaders getHeadersReceivedAtRequest(int request) { return cloneHeaders(headersReceived.get(request)); } + public int getRequestCount() { + return headersReceived.size(); + } + private static CallHeaders cloneHeaders(CallHeaders headers) { FlightCallHeaders cloneHeaders = new FlightCallHeaders(); for (String key : headers.keys()) {