more error handling

rishabhmaurya · Oct 3, 2024 · 78c6fa7 · 78c6fa7
1 parent 313dfaa
commit 78c6fa7
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/modules/arrow-flight/src/main/java/org/opensearch/flight/BaseFlightProducer.java b/modules/arrow-flight/src/main/java/org/opensearch/flight/BaseFlightProducer.java
@@ -47,10 +47,10 @@ public void getStream(CallContext context, Ticket ticket, ServerStreamListener l
                 if (result.equals(BackpressureStrategy.WaitResult.READY)) {
                     listener.putNext();
                 } else if (result.equals(BackpressureStrategy.WaitResult.TIMEOUT)) {
-                    listener.error(new RuntimeException("Timeout waiting for listener"));
+                    listener.error(CallStatus.TIMED_OUT.cause());
                     throw new RuntimeException("Timeout waiting for listener" + result);
                 } else {
-                    listener.error(new RuntimeException("Error while waiting for client: " + result));
+                    listener.error(CallStatus.INTERNAL.toRuntimeException());
                     throw new RuntimeException("Error while waiting for client: " + result);
                 }
             };
@@ -60,7 +60,7 @@ public void getStream(CallContext context, Ticket ticket, ServerStreamListener l
                 root.close();
             }
         } catch (Exception e) {
-            listener.error(e);
+            listener.error(CallStatus.INTERNAL.toRuntimeException().initCause(e));
             throw e;
         } finally {
             listener.completed();

diff --git a/modules/arrow-flight/src/main/java/org/opensearch/flight/FlightStreamManager.java b/modules/arrow-flight/src/main/java/org/opensearch/flight/FlightStreamManager.java
@@ -28,6 +28,7 @@ public FlightStreamManager(FlightClient flightClient) {
 
     @Override
     public VectorSchemaRoot getVectorSchemaRoot(StreamTicket ticket) {
+        // TODO: for remote streams
         FlightStream stream = flightClient.getStream(new Ticket(ticket.getBytes()));
         return stream.getRoot();
     }