GH 41262:[Java][FlightSQL] Implement stateless prepared statement

PR updates
stevelorddremio · Apr 19, 2024 · dca6707 · dca6707
1 parent c3ca784
commit dca6707
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 23 deletions.
diff --git a/...ight-sql/src/test/java/org/apache/arrow/flight/sql/example/FlightSqlStatelessExample.java b/...ight-sql/src/test/java/org/apache/arrow/flight/sql/example/FlightSqlStatelessExample.java
@@ -42,14 +42,11 @@
 import org.apache.arrow.flight.CallStatus;
 import org.apache.arrow.flight.FlightDescriptor;
 import org.apache.arrow.flight.FlightInfo;
-import org.apache.arrow.flight.FlightServer;
 import org.apache.arrow.flight.FlightStream;
 import org.apache.arrow.flight.Location;
 import org.apache.arrow.flight.PutResult;
 import org.apache.arrow.flight.sql.FlightSqlProducer;
 import org.apache.arrow.memory.ArrowBuf;
-import org.apache.arrow.memory.BufferAllocator;
-import org.apache.arrow.memory.RootAllocator;
 import org.apache.arrow.vector.VectorLoader;
 import org.apache.arrow.vector.VectorSchemaRoot;
 import org.apache.arrow.vector.VectorUnloader;
@@ -70,17 +67,6 @@
 public class FlightSqlStatelessExample extends FlightSqlExample {
   private static final Logger LOGGER = getLogger(FlightSqlStatelessExample.class);
 
-  public static void main(String[] args) throws Exception {
-    Location location = Location.forGrpcInsecure("localhost", 55555);
-    final FlightSqlStatelessExample example = new FlightSqlStatelessExample(location);
-    Location listenLocation = Location.forGrpcInsecure("0.0.0.0", 55555);
-    try (final BufferAllocator allocator = new RootAllocator();
-         final FlightServer server = FlightServer.builder(allocator, listenLocation, example).build()) {
-      server.start();
-      server.awaitTermination();
-    }
-  }
-
   public FlightSqlStatelessExample(final Location location) {
     super(location);
   }

diff --git a/...ght/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSqlStateless.java b/...ght/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSqlStateless.java
@@ -22,6 +22,7 @@
 import static org.hamcrest.CoreMatchers.*;
 
 import org.apache.arrow.flight.FlightClient;
+import org.apache.arrow.flight.FlightEndpoint;
 import org.apache.arrow.flight.FlightInfo;
 import org.apache.arrow.flight.FlightServer;
 import org.apache.arrow.flight.FlightStream;
@@ -81,15 +82,17 @@ public void testSimplePreparedStatementResultsWithParameterBinding() throws Exce
         prepare.setParameters(insertRoot);
         final FlightInfo flightInfo = prepare.execute();
 
-        try (FlightStream stream = sqlClient.getStream(flightInfo.getEndpoints().get(0).getTicket())) {
-          // TODO: root is null and getSchema hangs when run as complete suite.
-          // This works when run as an individual test.
-          final VectorSchemaRoot root = stream.getRoot();
-          final Schema schema = root.getSchema();
-          Assertions.assertAll(
-              () -> MatcherAssert.assertThat(schema, is(SCHEMA_INT_TABLE)),
-              () -> MatcherAssert.assertThat(getResults(stream), is(EXPECTED_RESULTS_FOR_PARAMETER_BINDING))
-          );
+        for (FlightEndpoint endpoint: flightInfo.getEndpoints()) {
+          try (FlightStream stream = sqlClient.getStream(endpoint.getTicket())) {
+            // TODO: root is null and getSchema hangs when run as complete suite.
+            // This works when run as an individual test.
+            final VectorSchemaRoot root = stream.getRoot();
+            final Schema schema = root.getSchema();
+            Assertions.assertAll(
+                () -> MatcherAssert.assertThat(schema, is(SCHEMA_INT_TABLE)),
+                () -> MatcherAssert.assertThat(getResults(stream), is(EXPECTED_RESULTS_FOR_PARAMETER_BINDING))
+            );
+          }
         }
       }
     }