This is an automated email from the ASF dual-hosted git repository.
lidavidm pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-java.git
The following commit(s) were added to refs/heads/main by this push:
new 94b0d71e3 GH-1117: Use updated handle in prepared stmt (#1120)
94b0d71e3 is described below
commit 94b0d71e3cff7dfdc4293a8517b0debdf7c856e7
Author: Pedro Matias <[email protected]>
AuthorDate: Wed Aug 26 00:31:03 2026 +0100
GH-1117: Use updated handle in prepared stmt (#1120)
## What's Changed
When a server sent an updated handle for a prepared statement after
parameters were bound,
the client would use the new handle for the subsequent call to
`GetFlightInfo` and then discard it. Any other
future requests for that prepared statement would use the original
handle.
With these changes, the updated handle is stored and reused for future
requests.
This change was created with AI assistance (Claude Code). All lines were
manually reviewed by a human. The output is not copyrightable subject
matter.
Closes #1117.
---
.../apache/arrow/flight/sql/FlightSqlClient.java | 22 +--
.../arrow/flight/sql/test/TestFlightSql.java | 189 +++++++++++++++++++++
2 files changed, 201 insertions(+), 10 deletions(-)
diff --git
a/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java
b/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java
index 0af09faee..69422b8cc 100644
---
a/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java
+++
b/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java
@@ -1217,6 +1217,7 @@ public class FlightSqlClient implements AutoCloseable {
public static class PreparedStatement implements AutoCloseable {
private final FlightClient client;
private final ActionCreatePreparedStatementResult preparedStatementResult;
+ private ByteString handle;
private VectorSchemaRoot parameterBindingRoot;
private boolean isClosed;
private Schema resultSetSchema;
@@ -1229,6 +1230,7 @@ public class FlightSqlClient implements AutoCloseable {
preparedStatementResult =
FlightSqlUtils.unpackAndParseOrThrow(
preparedStatementResults.next().getBody(),
ActionCreatePreparedStatementResult.class);
+ handle = preparedStatementResult.getPreparedStatementHandle();
isClosed = false;
}
@@ -1305,8 +1307,7 @@ public class FlightSqlClient implements AutoCloseable {
FlightDescriptor.command(
Any.pack(
CommandPreparedStatementQuery.newBuilder()
- .setPreparedStatementHandle(
-
preparedStatementResult.getPreparedStatementHandle())
+ .setPreparedStatementHandle(handle)
.build())
.toByteArray());
return client.getSchema(descriptor, options);
@@ -1337,8 +1338,7 @@ public class FlightSqlClient implements AutoCloseable {
FlightDescriptor.command(
Any.pack(
CommandPreparedStatementQuery.newBuilder()
- .setPreparedStatementHandle(
-
preparedStatementResult.getPreparedStatementHandle())
+ .setPreparedStatementHandle(handle)
.build())
.toByteArray());
@@ -1352,12 +1352,16 @@ public class FlightSqlClient implements AutoCloseable {
try (final ArrowBuf metadata = read.getApplicationMetadata()) {
final FlightSql.DoPutPreparedStatementResult
doPutPreparedStatementResult =
FlightSql.DoPutPreparedStatementResult.parseFrom(metadata.nioBuffer());
+ final ByteString updatedHandle =
+ doPutPreparedStatementResult.getPreparedStatementHandle();
+ if (!updatedHandle.isEmpty()) {
+ handle = updatedHandle;
+ }
descriptor =
FlightDescriptor.command(
Any.pack(
CommandPreparedStatementQuery.newBuilder()
- .setPreparedStatementHandle(
-
doPutPreparedStatementResult.getPreparedStatementHandle())
+ .setPreparedStatementHandle(handle)
.build())
.toByteArray());
}
@@ -1409,8 +1413,7 @@ public class FlightSqlClient implements AutoCloseable {
FlightDescriptor.command(
Any.pack(
CommandPreparedStatementUpdate.newBuilder()
- .setPreparedStatementHandle(
-
preparedStatementResult.getPreparedStatementHandle())
+ .setPreparedStatementHandle(handle)
.build())
.toByteArray());
setParameters(parameterBindingRoot == null ? VectorSchemaRoot.of() :
parameterBindingRoot);
@@ -1447,8 +1450,7 @@ public class FlightSqlClient implements AutoCloseable {
FlightSqlUtils.FLIGHT_SQL_CLOSE_PREPARED_STATEMENT.getType(),
Any.pack(
ActionClosePreparedStatementRequest.newBuilder()
- .setPreparedStatementHandle(
-
preparedStatementResult.getPreparedStatementHandle())
+ .setPreparedStatementHandle(handle)
.build())
.toByteArray());
final Iterator<Result> closePreparedStatementResults =
client.doAction(action, options);
diff --git
a/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSql.java
b/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSql.java
index e2934ab1e..167bd689c 100644
---
a/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSql.java
+++
b/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSql.java
@@ -27,9 +27,13 @@ import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertThrows;
import com.google.common.collect.ImmutableList;
+import com.google.protobuf.Any;
+import com.google.protobuf.ByteString;
+import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.PipedInputStream;
import java.io.PipedOutputStream;
+import java.nio.channels.Channels;
import java.nio.charset.StandardCharsets;
import java.sql.SQLException;
import java.util.ArrayList;
@@ -42,17 +46,21 @@ import java.util.Optional;
import java.util.stream.IntStream;
import org.apache.arrow.flight.CancelFlightInfoRequest;
import org.apache.arrow.flight.FlightClient;
+import org.apache.arrow.flight.FlightDescriptor;
import org.apache.arrow.flight.FlightInfo;
import org.apache.arrow.flight.FlightRuntimeException;
import org.apache.arrow.flight.FlightServer;
import org.apache.arrow.flight.FlightStatusCode;
import org.apache.arrow.flight.FlightStream;
import org.apache.arrow.flight.Location;
+import org.apache.arrow.flight.PutResult;
import org.apache.arrow.flight.RenewFlightEndpointRequest;
+import org.apache.arrow.flight.Result;
import org.apache.arrow.flight.sql.FlightSqlClient;
import org.apache.arrow.flight.sql.FlightSqlClient.PreparedStatement;
import org.apache.arrow.flight.sql.FlightSqlColumnMetadata;
import org.apache.arrow.flight.sql.FlightSqlProducer;
+import org.apache.arrow.flight.sql.NoOpFlightSqlProducer;
import org.apache.arrow.flight.sql.example.FlightSqlExample;
import org.apache.arrow.flight.sql.impl.FlightSql;
import
org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementIngest.TableDefinitionOptions;
@@ -60,6 +68,7 @@ import
org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementIngest.TableDe
import
org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementIngest.TableDefinitionOptions.TableNotExistOption;
import org.apache.arrow.flight.sql.impl.FlightSql.SqlSupportedCaseSensitivity;
import org.apache.arrow.flight.sql.util.TableRef;
+import org.apache.arrow.memory.ArrowBuf;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.IntVector;
@@ -67,6 +76,8 @@ import org.apache.arrow.vector.VarCharVector;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.ipc.ArrowStreamReader;
import org.apache.arrow.vector.ipc.ArrowStreamWriter;
+import org.apache.arrow.vector.ipc.WriteChannel;
+import org.apache.arrow.vector.ipc.message.MessageSerializer;
import org.apache.arrow.vector.types.Types.MinorType;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.Field;
@@ -1594,4 +1605,182 @@ public class TestFlightSql {
new
RenewFlightEndpointRequest(info.getEndpoints().get(0))));
assertEquals(FlightStatusCode.UNIMPLEMENTED, fre.status().code());
}
+
+ @Test
+ public void testPreparedStatementUsesUpdatedHandleAfterDoPut() throws
Exception {
+ final ByteString originalHandle =
ByteString.copyFromUtf8("original-handle");
+ final ByteString updatedHandle = ByteString.copyFromUtf8("updated-handle");
+
+ try (BufferAllocator testAllocator = new RootAllocator(Integer.MAX_VALUE))
{
+ final Schema paramSchema =
+ new Schema(singletonList(Field.nullable("id",
MinorType.INT.getType())));
+ final UpdatedHandleFlightSqlProducer mockProducer =
+ new UpdatedHandleFlightSqlProducer(
+ testAllocator, originalHandle, updatedHandle, paramSchema);
+
+ try (FlightServer testServer =
+ FlightServer.builder(
+ testAllocator, Location.forGrpcInsecure(LOCALHOST, 0),
mockProducer)
+ .build()
+ .start();
+ FlightSqlClient testClient =
+ new FlightSqlClient(
+ FlightClient.builder(
+ testAllocator, Location.forGrpcInsecure(LOCALHOST,
testServer.getPort()))
+ .build())) {
+
+ try (PreparedStatement ps = testClient.prepare("test query with
param=?");
+ VectorSchemaRoot params = VectorSchemaRoot.create(paramSchema,
testAllocator)) {
+ final IntVector v = (IntVector) params.getVector(0);
+ v.setSafe(0, 42);
+ params.setRowCount(1);
+ ps.setParameters(params);
+ ps.execute(); // DoPut → server returns updatedHandle in
DoPutPreparedStatementResult
+ }
+
+ assertAll(
+ () ->
+ assertThat(mockProducer.executeHandle)
+ .as("getFlightInfoPreparedStatement must use the updated
handle")
+ .isEqualTo(updatedHandle),
+ () ->
+ assertThat(mockProducer.closeHandle)
+ .as("ClosePreparedStatement must use the updated handle")
+ .isEqualTo(updatedHandle));
+ }
+ }
+ }
+
+ @Test
+ public void testPreparedStatementHandleUnchangedWithoutDoPut() throws
Exception {
+ final ByteString originalHandle =
ByteString.copyFromUtf8("original-handle");
+ final ByteString updatedHandle = ByteString.copyFromUtf8("updated-handle");
+
+ try (BufferAllocator testAllocator = new RootAllocator(Integer.MAX_VALUE))
{
+ final UpdatedHandleFlightSqlProducer mockProducer =
+ new UpdatedHandleFlightSqlProducer(
+ testAllocator, originalHandle, updatedHandle, new
Schema(emptyList()));
+
+ try (FlightServer testServer =
+ FlightServer.builder(
+ testAllocator, Location.forGrpcInsecure(LOCALHOST, 0),
mockProducer)
+ .build()
+ .start();
+ FlightSqlClient testClient =
+ new FlightSqlClient(
+ FlightClient.builder(
+ testAllocator, Location.forGrpcInsecure(LOCALHOST,
testServer.getPort()))
+ .build())) {
+
+ try (PreparedStatement ps = testClient.prepare("SELECT 1")) {
+ ps.execute();
+ }
+
+ assertAll(
+ () ->
+ assertThat(mockProducer.executeHandle)
+ .as("getFlightInfoPreparedStatement must use the original
handle")
+ .isEqualTo(originalHandle),
+ () ->
+ assertThat(mockProducer.closeHandle)
+ .as("ClosePreparedStatement must use the original handle")
+ .isEqualTo(originalHandle));
+ }
+ }
+ }
+
+ /**
+ * Minimal producer that returns an updated prepared-statement handle in the
{@code
+ * CommandPreparedStatementQuery} used with {@code DoPut} and records which
handle is used in
+ * subsequent operations, allowing the test to verify that the client
propagates the updated
+ * handle correctly.
+ */
+ private static final class UpdatedHandleFlightSqlProducer extends
NoOpFlightSqlProducer {
+
+ private final BufferAllocator allocator;
+ private final ByteString originalHandle;
+ private final ByteString updatedHandle;
+ private final ByteString serializedParamSchema;
+ ByteString executeHandle;
+ ByteString closeHandle;
+
+ UpdatedHandleFlightSqlProducer(
+ BufferAllocator allocator,
+ ByteString originalHandle,
+ ByteString updatedHandle,
+ Schema paramSchema) {
+ this.allocator = allocator;
+ this.originalHandle = originalHandle;
+ this.updatedHandle = updatedHandle;
+ this.serializedParamSchema = serializeSchema(paramSchema);
+ }
+
+ private static ByteString serializeSchema(Schema schema) {
+ try {
+ final ByteArrayOutputStream out = new ByteArrayOutputStream();
+ MessageSerializer.serialize(new
WriteChannel(Channels.newChannel(out)), schema);
+ return ByteString.copyFrom(out.toByteArray());
+ } catch (IOException e) {
+ throw new RuntimeException(e);
+ }
+ }
+
+ @Override
+ public void createPreparedStatement(
+ FlightSql.ActionCreatePreparedStatementRequest request,
+ CallContext context,
+ StreamListener<Result> listener) {
+ listener.onNext(
+ new Result(
+ Any.pack(
+
FlightSql.ActionCreatePreparedStatementResult.newBuilder()
+ .setPreparedStatementHandle(originalHandle)
+ .setParameterSchema(serializedParamSchema)
+ .build())
+ .toByteArray()));
+ listener.onCompleted();
+ }
+
+ @Override
+ public Runnable acceptPutPreparedStatementQuery(
+ FlightSql.CommandPreparedStatementQuery command,
+ CallContext context,
+ FlightStream flightStream,
+ StreamListener<PutResult> ackStream) {
+ return () -> {
+ while (flightStream.next()) {
+ // consume parameter batches
+ }
+ final byte[] responseBytes =
+ FlightSql.DoPutPreparedStatementResult.newBuilder()
+ .setPreparedStatementHandle(updatedHandle)
+ .build()
+ .toByteArray();
+ final ArrowBuf buf = allocator.buffer(responseBytes.length);
+ buf.writeBytes(responseBytes);
+ try (PutResult putResult = PutResult.metadata(buf)) {
+ ackStream.onNext(putResult);
+ ackStream.onCompleted();
+ }
+ };
+ }
+
+ @Override
+ public FlightInfo getFlightInfoPreparedStatement(
+ FlightSql.CommandPreparedStatementQuery command,
+ CallContext context,
+ FlightDescriptor descriptor) {
+ executeHandle = command.getPreparedStatementHandle();
+ return new FlightInfo(new Schema(emptyList()), descriptor, emptyList(),
-1, -1);
+ }
+
+ @Override
+ public void closePreparedStatement(
+ FlightSql.ActionClosePreparedStatementRequest request,
+ CallContext context,
+ StreamListener<Result> listener) {
+ closeHandle = request.getPreparedStatementHandle();
+ listener.onCompleted();
+ }
+ }
}
