Skip to content

Commit 94b0d71

Browse files
ennuiteennuite
authored
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.
1 parent 180e4ec commit 94b0d71

2 files changed

Lines changed: 201 additions & 10 deletions

File tree

‎flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java‎

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1217,6 +1217,7 @@ protected void updateCommandBuilder(CommandStatementIngest.Builder builder) {
12171217
public static class PreparedStatement implements AutoCloseable {
12181218
private final FlightClient client;
12191219
private final ActionCreatePreparedStatementResult preparedStatementResult;
1220+
private ByteString handle;
12201221
private VectorSchemaRoot parameterBindingRoot;
12211222
private boolean isClosed;
12221223
private Schema resultSetSchema;
@@ -1229,6 +1230,7 @@ public static class PreparedStatement implements AutoCloseable {
12291230
preparedStatementResult =
12301231
FlightSqlUtils.unpackAndParseOrThrow(
12311232
preparedStatementResults.next().getBody(), ActionCreatePreparedStatementResult.class);
1233+
handle = preparedStatementResult.getPreparedStatementHandle();
12321234
isClosed = false;
12331235
}
12341236

@@ -1305,8 +1307,7 @@ public SchemaResult fetchSchema(CallOption... options) {
13051307
FlightDescriptor.command(
13061308
Any.pack(
13071309
CommandPreparedStatementQuery.newBuilder()
1308-
.setPreparedStatementHandle(
1309-
preparedStatementResult.getPreparedStatementHandle())
1310+
.setPreparedStatementHandle(handle)
13101311
.build())
13111312
.toByteArray());
13121313
return client.getSchema(descriptor, options);
@@ -1337,8 +1338,7 @@ public FlightInfo execute(final CallOption... options) {
13371338
FlightDescriptor.command(
13381339
Any.pack(
13391340
CommandPreparedStatementQuery.newBuilder()
1340-
.setPreparedStatementHandle(
1341-
preparedStatementResult.getPreparedStatementHandle())
1341+
.setPreparedStatementHandle(handle)
13421342
.build())
13431343
.toByteArray());
13441344

@@ -1352,12 +1352,16 @@ public FlightInfo execute(final CallOption... options) {
13521352
try (final ArrowBuf metadata = read.getApplicationMetadata()) {
13531353
final FlightSql.DoPutPreparedStatementResult doPutPreparedStatementResult =
13541354
FlightSql.DoPutPreparedStatementResult.parseFrom(metadata.nioBuffer());
1355+
final ByteString updatedHandle =
1356+
doPutPreparedStatementResult.getPreparedStatementHandle();
1357+
if (!updatedHandle.isEmpty()) {
1358+
handle = updatedHandle;
1359+
}
13551360
descriptor =
13561361
FlightDescriptor.command(
13571362
Any.pack(
13581363
CommandPreparedStatementQuery.newBuilder()
1359-
.setPreparedStatementHandle(
1360-
doPutPreparedStatementResult.getPreparedStatementHandle())
1364+
.setPreparedStatementHandle(handle)
13611365
.build())
13621366
.toByteArray());
13631367
}
@@ -1409,8 +1413,7 @@ public long executeUpdate(final CallOption... options) {
14091413
FlightDescriptor.command(
14101414
Any.pack(
14111415
CommandPreparedStatementUpdate.newBuilder()
1412-
.setPreparedStatementHandle(
1413-
preparedStatementResult.getPreparedStatementHandle())
1416+
.setPreparedStatementHandle(handle)
14141417
.build())
14151418
.toByteArray());
14161419
setParameters(parameterBindingRoot == null ? VectorSchemaRoot.of() : parameterBindingRoot);
@@ -1447,8 +1450,7 @@ public void close(final CallOption... options) {
14471450
FlightSqlUtils.FLIGHT_SQL_CLOSE_PREPARED_STATEMENT.getType(),
14481451
Any.pack(
14491452
ActionClosePreparedStatementRequest.newBuilder()
1450-
.setPreparedStatementHandle(
1451-
preparedStatementResult.getPreparedStatementHandle())
1453+
.setPreparedStatementHandle(handle)
14521454
.build())
14531455
.toByteArray());
14541456
final Iterator<Result> closePreparedStatementResults = client.doAction(action, options);

‎flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSql.java‎

Lines changed: 189 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,9 +27,13 @@
2727
import static org.junit.jupiter.api.Assertions.assertThrows;
2828

2929
import com.google.common.collect.ImmutableList;
30+
import com.google.protobuf.Any;
31+
import com.google.protobuf.ByteString;
32+
import java.io.ByteArrayOutputStream;
3033
import java.io.IOException;
3134
import java.io.PipedInputStream;
3235
import java.io.PipedOutputStream;
36+
import java.nio.channels.Channels;
3337
import java.nio.charset.StandardCharsets;
3438
import java.sql.SQLException;
3539
import java.util.ArrayList;
@@ -42,31 +46,38 @@
4246
import java.util.stream.IntStream;
4347
import org.apache.arrow.flight.CancelFlightInfoRequest;
4448
import org.apache.arrow.flight.FlightClient;
49+
import org.apache.arrow.flight.FlightDescriptor;
4550
import org.apache.arrow.flight.FlightInfo;
4651
import org.apache.arrow.flight.FlightRuntimeException;
4752
import org.apache.arrow.flight.FlightServer;
4853
import org.apache.arrow.flight.FlightStatusCode;
4954
import org.apache.arrow.flight.FlightStream;
5055
import org.apache.arrow.flight.Location;
56+
import org.apache.arrow.flight.PutResult;
5157
import org.apache.arrow.flight.RenewFlightEndpointRequest;
58+
import org.apache.arrow.flight.Result;
5259
import org.apache.arrow.flight.sql.FlightSqlClient;
5360
import org.apache.arrow.flight.sql.FlightSqlClient.PreparedStatement;
5461
import org.apache.arrow.flight.sql.FlightSqlColumnMetadata;
5562
import org.apache.arrow.flight.sql.FlightSqlProducer;
63+
import org.apache.arrow.flight.sql.NoOpFlightSqlProducer;
5664
import org.apache.arrow.flight.sql.example.FlightSqlExample;
5765
import org.apache.arrow.flight.sql.impl.FlightSql;
5866
import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementIngest.TableDefinitionOptions;
5967
import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementIngest.TableDefinitionOptions.TableExistsOption;
6068
import org.apache.arrow.flight.sql.impl.FlightSql.CommandStatementIngest.TableDefinitionOptions.TableNotExistOption;
6169
import org.apache.arrow.flight.sql.impl.FlightSql.SqlSupportedCaseSensitivity;
6270
import org.apache.arrow.flight.sql.util.TableRef;
71+
import org.apache.arrow.memory.ArrowBuf;
6372
import org.apache.arrow.memory.BufferAllocator;
6473
import org.apache.arrow.memory.RootAllocator;
6574
import org.apache.arrow.vector.IntVector;
6675
import org.apache.arrow.vector.VarCharVector;
6776
import org.apache.arrow.vector.VectorSchemaRoot;
6877
import org.apache.arrow.vector.ipc.ArrowStreamReader;
6978
import org.apache.arrow.vector.ipc.ArrowStreamWriter;
79+
import org.apache.arrow.vector.ipc.WriteChannel;
80+
import org.apache.arrow.vector.ipc.message.MessageSerializer;
7081
import org.apache.arrow.vector.types.Types.MinorType;
7182
import org.apache.arrow.vector.types.pojo.ArrowType;
7283
import org.apache.arrow.vector.types.pojo.Field;
@@ -1594,4 +1605,182 @@ public void testRenewEndpoint() {
15941605
new RenewFlightEndpointRequest(info.getEndpoints().get(0))));
15951606
assertEquals(FlightStatusCode.UNIMPLEMENTED, fre.status().code());
15961607
}
1608+
1609+
@Test
1610+
public void testPreparedStatementUsesUpdatedHandleAfterDoPut() throws Exception {
1611+
final ByteString originalHandle = ByteString.copyFromUtf8("original-handle");
1612+
final ByteString updatedHandle = ByteString.copyFromUtf8("updated-handle");
1613+
1614+
try (BufferAllocator testAllocator = new RootAllocator(Integer.MAX_VALUE)) {
1615+
final Schema paramSchema =
1616+
new Schema(singletonList(Field.nullable("id", MinorType.INT.getType())));
1617+
final UpdatedHandleFlightSqlProducer mockProducer =
1618+
new UpdatedHandleFlightSqlProducer(
1619+
testAllocator, originalHandle, updatedHandle, paramSchema);
1620+
1621+
try (FlightServer testServer =
1622+
FlightServer.builder(
1623+
testAllocator, Location.forGrpcInsecure(LOCALHOST, 0), mockProducer)
1624+
.build()
1625+
.start();
1626+
FlightSqlClient testClient =
1627+
new FlightSqlClient(
1628+
FlightClient.builder(
1629+
testAllocator, Location.forGrpcInsecure(LOCALHOST, testServer.getPort()))
1630+
.build())) {
1631+
1632+
try (PreparedStatement ps = testClient.prepare("test query with param=?");
1633+
VectorSchemaRoot params = VectorSchemaRoot.create(paramSchema, testAllocator)) {
1634+
final IntVector v = (IntVector) params.getVector(0);
1635+
v.setSafe(0, 42);
1636+
params.setRowCount(1);
1637+
ps.setParameters(params);
1638+
ps.execute(); // DoPut → server returns updatedHandle in DoPutPreparedStatementResult
1639+
}
1640+
1641+
assertAll(
1642+
() ->
1643+
assertThat(mockProducer.executeHandle)
1644+
.as("getFlightInfoPreparedStatement must use the updated handle")
1645+
.isEqualTo(updatedHandle),
1646+
() ->
1647+
assertThat(mockProducer.closeHandle)
1648+
.as("ClosePreparedStatement must use the updated handle")
1649+
.isEqualTo(updatedHandle));
1650+
}
1651+
}
1652+
}
1653+
1654+
@Test
1655+
public void testPreparedStatementHandleUnchangedWithoutDoPut() throws Exception {
1656+
final ByteString originalHandle = ByteString.copyFromUtf8("original-handle");
1657+
final ByteString updatedHandle = ByteString.copyFromUtf8("updated-handle");
1658+
1659+
try (BufferAllocator testAllocator = new RootAllocator(Integer.MAX_VALUE)) {
1660+
final UpdatedHandleFlightSqlProducer mockProducer =
1661+
new UpdatedHandleFlightSqlProducer(
1662+
testAllocator, originalHandle, updatedHandle, new Schema(emptyList()));
1663+
1664+
try (FlightServer testServer =
1665+
FlightServer.builder(
1666+
testAllocator, Location.forGrpcInsecure(LOCALHOST, 0), mockProducer)
1667+
.build()
1668+
.start();
1669+
FlightSqlClient testClient =
1670+
new FlightSqlClient(
1671+
FlightClient.builder(
1672+
testAllocator, Location.forGrpcInsecure(LOCALHOST, testServer.getPort()))
1673+
.build())) {
1674+
1675+
try (PreparedStatement ps = testClient.prepare("SELECT 1")) {
1676+
ps.execute();
1677+
}
1678+
1679+
assertAll(
1680+
() ->
1681+
assertThat(mockProducer.executeHandle)
1682+
.as("getFlightInfoPreparedStatement must use the original handle")
1683+
.isEqualTo(originalHandle),
1684+
() ->
1685+
assertThat(mockProducer.closeHandle)
1686+
.as("ClosePreparedStatement must use the original handle")
1687+
.isEqualTo(originalHandle));
1688+
}
1689+
}
1690+
}
1691+
1692+
/**
1693+
* Minimal producer that returns an updated prepared-statement handle in the {@code
1694+
* CommandPreparedStatementQuery} used with {@code DoPut} and records which handle is used in
1695+
* subsequent operations, allowing the test to verify that the client propagates the updated
1696+
* handle correctly.
1697+
*/
1698+
private static final class UpdatedHandleFlightSqlProducer extends NoOpFlightSqlProducer {
1699+
1700+
private final BufferAllocator allocator;
1701+
private final ByteString originalHandle;
1702+
private final ByteString updatedHandle;
1703+
private final ByteString serializedParamSchema;
1704+
ByteString executeHandle;
1705+
ByteString closeHandle;
1706+
1707+
UpdatedHandleFlightSqlProducer(
1708+
BufferAllocator allocator,
1709+
ByteString originalHandle,
1710+
ByteString updatedHandle,
1711+
Schema paramSchema) {
1712+
this.allocator = allocator;
1713+
this.originalHandle = originalHandle;
1714+
this.updatedHandle = updatedHandle;
1715+
this.serializedParamSchema = serializeSchema(paramSchema);
1716+
}
1717+
1718+
private static ByteString serializeSchema(Schema schema) {
1719+
try {
1720+
final ByteArrayOutputStream out = new ByteArrayOutputStream();
1721+
MessageSerializer.serialize(new WriteChannel(Channels.newChannel(out)), schema);
1722+
return ByteString.copyFrom(out.toByteArray());
1723+
} catch (IOException e) {
1724+
throw new RuntimeException(e);
1725+
}
1726+
}
1727+
1728+
@Override
1729+
public void createPreparedStatement(
1730+
FlightSql.ActionCreatePreparedStatementRequest request,
1731+
CallContext context,
1732+
StreamListener<Result> listener) {
1733+
listener.onNext(
1734+
new Result(
1735+
Any.pack(
1736+
FlightSql.ActionCreatePreparedStatementResult.newBuilder()
1737+
.setPreparedStatementHandle(originalHandle)
1738+
.setParameterSchema(serializedParamSchema)
1739+
.build())
1740+
.toByteArray()));
1741+
listener.onCompleted();
1742+
}
1743+
1744+
@Override
1745+
public Runnable acceptPutPreparedStatementQuery(
1746+
FlightSql.CommandPreparedStatementQuery command,
1747+
CallContext context,
1748+
FlightStream flightStream,
1749+
StreamListener<PutResult> ackStream) {
1750+
return () -> {
1751+
while (flightStream.next()) {
1752+
// consume parameter batches
1753+
}
1754+
final byte[] responseBytes =
1755+
FlightSql.DoPutPreparedStatementResult.newBuilder()
1756+
.setPreparedStatementHandle(updatedHandle)
1757+
.build()
1758+
.toByteArray();
1759+
final ArrowBuf buf = allocator.buffer(responseBytes.length);
1760+
buf.writeBytes(responseBytes);
1761+
try (PutResult putResult = PutResult.metadata(buf)) {
1762+
ackStream.onNext(putResult);
1763+
ackStream.onCompleted();
1764+
}
1765+
};
1766+
}
1767+
1768+
@Override
1769+
public FlightInfo getFlightInfoPreparedStatement(
1770+
FlightSql.CommandPreparedStatementQuery command,
1771+
CallContext context,
1772+
FlightDescriptor descriptor) {
1773+
executeHandle = command.getPreparedStatementHandle();
1774+
return new FlightInfo(new Schema(emptyList()), descriptor, emptyList(), -1, -1);
1775+
}
1776+
1777+
@Override
1778+
public void closePreparedStatement(
1779+
FlightSql.ActionClosePreparedStatementRequest request,
1780+
CallContext context,
1781+
StreamListener<Result> listener) {
1782+
closeHandle = request.getPreparedStatementHandle();
1783+
listener.onCompleted();
1784+
}
1785+
}
15971786
}

0 commit comments

Comments
 (0)