Reuse updated handle

Author: ennuite

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

@@ -1217,6 +1217,7 @@ protected void updateCommandBuilder(CommandStatementIngest.Builder builder) {
   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 static class PreparedStatement implements AutoCloseable {
       preparedStatementResult =
           FlightSqlUtils.unpackAndParseOrThrow(
               preparedStatementResults.next().getBody(), ActionCreatePreparedStatementResult.class);
+      handle = preparedStatementResult.getPreparedStatementHandle();
       isClosed = false;
     }
 
@@ -1292,8 +1294,7 @@ public SchemaResult fetchSchema(CallOption... options) {
           FlightDescriptor.command(
               Any.pack(
                       CommandPreparedStatementQuery.newBuilder()
-                          .setPreparedStatementHandle(
-                              preparedStatementResult.getPreparedStatementHandle())
+                          .setPreparedStatementHandle(handle)
                           .build())
                   .toByteArray());
       return client.getSchema(descriptor, options);
@@ -1324,8 +1325,7 @@ public FlightInfo execute(final CallOption... options) {
           FlightDescriptor.command(
               Any.pack(
                       CommandPreparedStatementQuery.newBuilder()
-                          .setPreparedStatementHandle(
-                              preparedStatementResult.getPreparedStatementHandle())
+                          .setPreparedStatementHandle(handle)
                           .build())
                   .toByteArray());
 
@@ -1339,12 +1339,16 @@ public FlightInfo execute(final CallOption... options) {
               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());
               }
@@ -1396,8 +1400,7 @@ public long executeUpdate(final CallOption... options) {
           FlightDescriptor.command(
               Any.pack(
                       CommandPreparedStatementUpdate.newBuilder()
-                          .setPreparedStatementHandle(
-                              preparedStatementResult.getPreparedStatementHandle())
+                          .setPreparedStatementHandle(handle)
                           .build())
                   .toByteArray());
       setParameters(parameterBindingRoot == null ? VectorSchemaRoot.of() : parameterBindingRoot);
@@ -1434,8 +1437,7 @@ public void close(final CallOption... options) {
               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);