diff --git a/java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightClient.java b/java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightClient.java index 155e373bda2f3..a43f43078df46 100644 --- a/java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightClient.java +++ b/java/flight/flight-core/src/main/java/org/apache/arrow/flight/FlightClient.java @@ -412,6 +412,16 @@ public ClientStreamListener getWriter() { return writer; } + /** + * Make sure stream is drained. You must call this to be notified of any errors that may have + * happened after the exchange is complete. This should be called after `getWriter().completed()` + * instead of `getWriter().getResult`. + */ + public void getResult() { + // After exchange is complete, make sure stream is drained to propagate errors through reader + reader.next(); + } + /** Shut down the streams in this call. */ @Override public void close() throws Exception { diff --git a/java/flight/flight-core/src/test/java/org/apache/arrow/flight/TestDoExchange.java b/java/flight/flight-core/src/test/java/org/apache/arrow/flight/TestDoExchange.java index 8c99cd496611a..f9db9bfd23a88 100644 --- a/java/flight/flight-core/src/test/java/org/apache/arrow/flight/TestDoExchange.java +++ b/java/flight/flight-core/src/test/java/org/apache/arrow/flight/TestDoExchange.java @@ -392,7 +392,7 @@ public void testDoExchangeError() throws Exception { stream.getWriter().completed(); // Must call reader.next() to get any errors after exchange, will return false if no error - final FlightRuntimeException fre = assertThrows(FlightRuntimeException.class, reader::next); + final FlightRuntimeException fre = assertThrows(FlightRuntimeException.class, stream::getResult); assertEquals("error completing exchange", fre.status().description()); } } @@ -438,7 +438,7 @@ public void testCloseWithMetadata() throws Exception { // Close our write channel and ensure the server also closes theirs stream.getWriter().completed(); - assertFalse(reader.next()); + stream.getResult(); // Not necessary to close reader here, but check closing twice doesn't lead to negative refcnt from metadata stream.getReader().close();