[ArrowFlight] Fix ArrowStream close() (#534)

Author: nikolaobradovic-db

rust/NEXT_CHANGELOG.md

@@ -8,6 +8,8 @@
 
 ### Bug Fixes
 
+- **Arrow Flight — `close()` now propagates flush errors** (Beta): `ZerobusArrowStream::close()` previously swallowed a failed final `flush()` and always returned `Ok(())`, contradicting its documentation and diverging from the proto stream's `close()`. It now returns the flush error after still tearing down the stream and moving pending batches to the failed set (retrievable via `get_unacked_batches()`).
+
 ### Documentation
 
 ### Internal Changes

rust/sdk/src/arrow_stream.rs

@@ -1566,8 +1566,10 @@ impl ZerobusArrowStream {
             "Closing Arrow Flight stream"
         );
 
-        // Flush pending batches.
-        if let Err(e) = self.flush().await {
+        // Flush pending batches. Capture the result so the stream is still torn
+        // down on failure, then propagate it to the caller.
+        let flush_result = self.flush().await;
+        if let Err(e) = &flush_result {
             warn!(
                 "Flush failed during close: {}. Moving pending batches to failed.",
                 e
@@ -1593,7 +1595,7 @@ impl ZerobusArrowStream {
             }
         }
 
-        Ok(())
+        flush_result
     }
 
     /// Returns all batches that were ingested but not acknowledged by the server.

rust/tests/src/arrow_tests.rs

@@ -523,6 +523,60 @@ mod arrow_flight_tests {
             Ok(())
         }
 
+        #[tokio::test]
+        async fn test_close_propagates_flush_error() -> Result<(), Box<dyn std::error::Error>> {
+            setup_tracing();
+            info!("Starting test_close_propagates_flush_error");
+
+            let (mock_server, server_url) = start_mock_flight_server().await?;
+            let schema = create_test_arrow_schema();
+
+            // Ack is delayed well past the flush timeout, so the flush performed by
+            // close() times out while the stream is still open (no server error, so the
+            // idempotent close guard does not short-circuit).
+            mock_server
+                .inject_responses(
+                    TABLE_NAME,
+                    vec![MockFlightResponse::BatchAck {
+                        ack_up_to_offset: 0,
+                        delay_ms: 5000,
+                        ack_up_to_records: 1,
+                    }],
+                )
+                .await;
+
+            let sdk = ZerobusSdk::builder()
+                .endpoint(server_url.clone())
+                .unity_catalog_url("https://mock-uc.com")
+                .tls_config(Arc::new(NoTlsConfig))
+                .build()?;
+
+            let mut stream = sdk
+                .stream_builder()
+                .table(TABLE_NAME)
+                .headers_provider(Arc::new(TestHeadersProvider::default()))
+                .arrow(schema.clone())
+                .flush_timeout_ms(100)
+                .build_arrow()
+                .await?;
+
+            let batch = create_test_record_batch(schema, vec![1], vec![Some("unacked")]);
+            let _offset = stream.ingest_batch(batch).await?;
+
+            let close_result = stream.close().await;
+            assert!(
+                close_result.is_err(),
+                "close() must propagate the error when its flush fails"
+            );
+
+            // The stream is still torn down, and the unacked batch is recoverable.
+            assert!(stream.is_closed());
+            let unacked = stream.get_unacked_batches().await?;
+            assert_eq!(unacked.len(), 1, "Should have 1 unacked batch");
+
+            Ok(())
+        }
+
         #[tokio::test]
         async fn test_empty_flush() -> Result<(), Box<dyn std::error::Error>> {
             setup_tracing();
@@ -756,6 +810,9 @@ mod arrow_flight_tests {
             let offset2 = stream.ingest_batch(batch2).await?;
             assert!(stream.wait_for_offset(offset2).await.is_err());
 
+            // The permanent error already closed the stream, so close() short-circuits
+            // to Ok via its idempotent guard; the unacked batch was moved to failed by
+            // the error path, not by this close().
             let _ = stream.close().await;
 
             let unacked = stream.get_unacked_batches().await?;