From a09efcd6dee23000f02f4b71f7d3d17be5be3dc7 Mon Sep 17 00:00:00 2001 From: Xiangpeng Hao Date: Mon, 3 Mar 2025 14:21:42 -0600 Subject: [PATCH 1/3] make tonic error boxed --- arrow-flight/src/client.rs | 8 ++++---- arrow-flight/src/error.rs | 6 +++--- arrow-flight/src/streams.rs | 2 +- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/arrow-flight/src/client.rs b/arrow-flight/src/client.rs index 97d9899a9fb0..9b4c10e9a093 100644 --- a/arrow-flight/src/client.rs +++ b/arrow-flight/src/client.rs @@ -210,7 +210,7 @@ impl FlightClient { let (response_stream, trailers) = extract_lazy_trailers(response_stream); Ok(FlightRecordBatchStream::new_from_flight_data( - response_stream.map_err(FlightError::Tonic), + response_stream.map_err(|status| status.into()), ) .with_headers(md) .with_trailers(trailers)) @@ -470,7 +470,7 @@ impl FlightClient { .list_flights(request) .await? .into_inner() - .map_err(FlightError::Tonic); + .map_err(|status| status.into()); Ok(response.boxed()) } @@ -537,7 +537,7 @@ impl FlightClient { .list_actions(request) .await? .into_inner() - .map_err(FlightError::Tonic); + .map_err(|status| status.into()); Ok(action_stream.boxed()) } @@ -575,7 +575,7 @@ impl FlightClient { .do_action(request) .await? .into_inner() - .map_err(FlightError::Tonic) + .map_err(|status| status.into()) .map(|r| { r.map(|r| { // unwrap inner bytes diff --git a/arrow-flight/src/error.rs b/arrow-flight/src/error.rs index 499706e1ede7..d8698ef66ad7 100644 --- a/arrow-flight/src/error.rs +++ b/arrow-flight/src/error.rs @@ -27,7 +27,7 @@ pub enum FlightError { /// Returned when functionality is not yet available. NotYetImplemented(String), /// Error from the underlying tonic library - Tonic(tonic::Status), + Tonic(Box), /// Some unexpected message was received ProtocolError(String), /// An error occurred during decoding @@ -74,7 +74,7 @@ impl Error for FlightError { impl From for FlightError { fn from(status: tonic::Status) -> Self { - Self::Tonic(status) + Self::Tonic(Box::new(status)) } } @@ -91,7 +91,7 @@ impl From for tonic::Status { match value { FlightError::Arrow(e) => tonic::Status::internal(e.to_string()), FlightError::NotYetImplemented(e) => tonic::Status::internal(e), - FlightError::Tonic(status) => status, + FlightError::Tonic(status) => *status, FlightError::ProtocolError(e) => tonic::Status::internal(e), FlightError::DecodeError(e) => tonic::Status::internal(e), FlightError::ExternalError(e) => tonic::Status::internal(e.to_string()), diff --git a/arrow-flight/src/streams.rs b/arrow-flight/src/streams.rs index e532a80e1ebb..ab496122013d 100644 --- a/arrow-flight/src/streams.rs +++ b/arrow-flight/src/streams.rs @@ -127,7 +127,7 @@ impl Stream for FallibleTonicResponseStream { match ready!(pinned.response_stream.poll_next_unpin(cx)) { Some(Ok(res)) => Poll::Ready(Some(Ok(res))), - Some(Err(status)) => Poll::Ready(Some(Err(FlightError::Tonic(status)))), + Some(Err(status)) => Poll::Ready(Some(Err(status.into()))), None => Poll::Ready(None), } } From b3b4eee59c561b8505edfd8a89a6af3dad590ea9 Mon Sep 17 00:00:00 2001 From: Xiangpeng Hao Date: Mon, 3 Mar 2025 15:28:17 -0600 Subject: [PATCH 2/3] test with all features --- arrow-flight/src/sql/client.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/arrow-flight/src/sql/client.rs b/arrow-flight/src/sql/client.rs index 6d3ac3dbe610..6791b68b757d 100644 --- a/arrow-flight/src/sql/client.rs +++ b/arrow-flight/src/sql/client.rs @@ -309,7 +309,7 @@ impl FlightSqlServiceClient { let (response_stream, trailers) = extract_lazy_trailers(response_stream); Ok(FlightRecordBatchStream::new_from_flight_data( - response_stream.map_err(FlightError::Tonic), + response_stream.map_err(|status| status.into()), ) .with_headers(md) .with_trailers(trailers)) From 94da21556734df402f0f88dd6c59777f582e0e51 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Thu, 6 Mar 2025 13:48:59 -0500 Subject: [PATCH 3/3] Add test for FlightError size, suggested by @crepererum --- arrow-flight/src/error.rs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/arrow-flight/src/error.rs b/arrow-flight/src/error.rs index d8698ef66ad7..ac8030583299 100644 --- a/arrow-flight/src/error.rs +++ b/arrow-flight/src/error.rs @@ -147,4 +147,10 @@ mod test { let source = root_error.downcast_ref::().unwrap(); assert!(matches!(source, FlightError::DecodeError(_))); } + + #[test] + fn test_error_size() { + // use Box in variants to keep this size down + assert_eq!(std::mem::size_of::(), 32); + } }