apacheGH-37720: [Java][FlightSQL] Implement stateless prepared statem…

…ents
stevelorddremio · May 16, 2024 · f079102 · f079102
1 parent 1c546fb
commit f079102
Showing 1 changed file with 20 additions and 3 deletions.
diff --git a/java/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java b/java/flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/FlightSqlClient.java
@@ -78,6 +78,7 @@
 import org.apache.arrow.flight.SetSessionOptionsResult;
 import org.apache.arrow.flight.SyncPutListener;
 import org.apache.arrow.flight.Ticket;
+import org.apache.arrow.flight.sql.impl.FlightSql;
 import org.apache.arrow.flight.sql.impl.FlightSql.ActionCreatePreparedStatementResult;
 import org.apache.arrow.flight.sql.impl.FlightSql.CommandPreparedStatementQuery;
 import org.apache.arrow.flight.sql.util.TableRef;
@@ -1048,15 +1049,31 @@ private Schema deserializeSchema(final ByteString bytes) {
     public FlightInfo execute(final CallOption... options) {
       checkOpen();
 
-      final FlightDescriptor descriptor = FlightDescriptor
+      FlightDescriptor descriptor = FlightDescriptor
           .command(Any.pack(CommandPreparedStatementQuery.newBuilder()
                   .setPreparedStatementHandle(preparedStatementResult.getPreparedStatementHandle())
                   .build())
               .toByteArray());
 
       if (parameterBindingRoot != null && parameterBindingRoot.getRowCount() > 0) {
-        try (final SyncPutListener listener = putParameters(descriptor, options)) {
-          listener.getResult();
+        try (final SyncPutListener putListener = putParameters(descriptor, options)) {
+          final PutResult read = putListener.read();
+          if (read != null) {
+            try (final ArrowBuf metadata = read.getApplicationMetadata()) {
+              final FlightSql.DoPutPreparedStatementResult doPutPreparedStatementResult =
+                      FlightSql.DoPutPreparedStatementResult.parseFrom(metadata.nioBuffer());
+              descriptor = FlightDescriptor
+                      .command(Any.pack(CommandPreparedStatementQuery.newBuilder()
+                                      .setPreparedStatementHandle(
+                                              doPutPreparedStatementResult.getPreparedStatementHandle())
+                                      .build())
+                              .toByteArray());
+            }
+          }
+        } catch (final InterruptedException | ExecutionException e) {
+          throw CallStatus.CANCELLED.withCause(e).toRuntimeException();
+        } catch (final InvalidProtocolBufferException e) {
+          throw CallStatus.INVALID_ARGUMENT.withCause(e).toRuntimeException();
         }
       }