provide flight client through stream manager

Signed-off-by: bowenlan-amzn <[email protected]>

libs/arrow/src/main/java/org/opensearch/arrow/StreamManager.java

@@ -20,6 +20,9 @@
  */
 @ExperimentalApi
 public abstract class StreamManager implements AutoCloseable {
+
+    public abstract void setFlightClient(Object flightClient);
+
     private final ConcurrentHashMap<StreamTicket, ArrowStreamProvider> streams;
 
     /**

modules/arrow-flight/src/main/java/org/opensearch/flight/FlightService.java

@@ -117,6 +117,7 @@ protected void doStart() {
             final Location location = Location.forGrpcInsecure(host, port);
             server = FlightServer.builder(allocator, location, producer).build();
             client = FlightClient.builder(allocator, location).build();
+            streamManager.setFlightClient(client);
             server.start();
             logger.info("Arrow Flight server started successfully");
         } catch (IOException e) {

modules/arrow-flight/src/main/java/org/opensearch/flight/FlightStreamManager.java

@@ -25,7 +25,7 @@
  */
 public class FlightStreamManager extends StreamManager {
 
-    private final FlightClient flightClient;
+    private FlightClient flightClient;
 
     /**
      * Constructs a new FlightStreamManager.
@@ -42,6 +42,12 @@ public FlightStreamManager(FlightClient flightClient) {
      * @param ticket The StreamTicket identifying the desired stream.
      * @return The VectorSchemaRoot associated with the given ticket.
      */
+    @Override
+    public void setFlightClient(Object flightClient) {
+        assert flightClient instanceof FlightClient;
+        this.flightClient = (FlightClient) flightClient;
+    }
+
     @Override
     public VectorSchemaRoot getVectorSchemaRoot(StreamTicket ticket) {
         // TODO: for remote streams, register streams in cluster state with node details