Refine Display and Source implementation for error types (#5439)

* refine error format and source Signed-off-by: Bugen Zhao <[email protected]> * fix unit tests Signed-off-by: Bugen Zhao <[email protected]> * fix deref style Signed-off-by: Bugen Zhao <[email protected]> * update source impl Signed-off-by: Bugen Zhao <[email protected]> --------- Signed-off-by: Bugen Zhao <[email protected]>
apache · Feb 29, 2024 · 877e870 · 877e870
1 parent ed996da
commit 877e870
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 15 deletions.
diff --git a/arrow-flight/src/error.rs b/arrow-flight/src/error.rs
@@ -49,17 +49,24 @@ impl FlightError {
 
 impl std::fmt::Display for FlightError {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        // TODO better format / error
-        write!(f, "{self:?}")
+        match self {
+            FlightError::Arrow(source) => write!(f, "Arrow error: {}", source),
+            FlightError::NotYetImplemented(desc) => write!(f, "Not yet implemented: {}", desc),
+            FlightError::Tonic(source) => write!(f, "Tonic error: {}", source),
+            FlightError::ProtocolError(desc) => write!(f, "Protocol error: {}", desc),
+            FlightError::DecodeError(desc) => write!(f, "Decode error: {}", desc),
+            FlightError::ExternalError(source) => write!(f, "External error: {}", source),
+        }
     }
 }
 
 impl Error for FlightError {
     fn source(&self) -> Option<&(dyn Error + 'static)> {
-        if let Self::ExternalError(e) = self {
-            Some(e.as_ref())
-        } else {
-            None
+        match self {
+            FlightError::Arrow(source) => Some(source),
+            FlightError::Tonic(source) => Some(source),
+            FlightError::ExternalError(source) => Some(source.as_ref()),
+            _ => None,
         }
     }
 }

diff --git a/arrow-flight/tests/client.rs b/arrow-flight/tests/client.rs
@@ -935,7 +935,7 @@ async fn test_cancel_flight_info_error_no_response() {
 
         assert_eq!(
             err.to_string(),
-            "ProtocolError(\"Received no response for cancel_flight_info call\")"
+            "Protocol error: Received no response for cancel_flight_info call"
         );
         // server still got the request
         let expected_request = Action::new("CancelFlightInfo", request.encode_to_vec());
@@ -985,7 +985,7 @@ async fn test_renew_flight_endpoint_error_no_response() {
 
         assert_eq!(
             err.to_string(),
-            "ProtocolError(\"Received no response for renew_flight_endpoint call\")"
+            "Protocol error: Received no response for renew_flight_endpoint call"
         );
         // server still got the request
         let expected_request = Action::new("RenewFlightEndpoint", request.encode_to_vec());

diff --git a/arrow-flight/tests/encode_decode.rs b/arrow-flight/tests/encode_decode.rs
@@ -57,7 +57,7 @@ async fn test_error() {
     let result: Result<Vec<_>, _> = decode_stream.try_collect().await;
 
     let result = result.unwrap_err();
-    assert_eq!(result.to_string(), r#"NotYetImplemented("foo")"#);
+    assert_eq!(result.to_string(), "Not yet implemented: foo");
 }
 
 #[tokio::test]
@@ -287,7 +287,7 @@ async fn test_mismatched_record_batch_schema() {
     let err = result.unwrap_err();
     assert_eq!(
         err.to_string(),
-        "Arrow(InvalidArgumentError(\"number of columns(1) must match number of fields(2) in schema\"))"
+        "Arrow error: Invalid argument error: number of columns(1) must match number of fields(2) in schema"
     );
 }
 
@@ -312,7 +312,7 @@ async fn test_chained_streams_batch_decoder() {
     let err = result.unwrap_err();
     assert_eq!(
         err.to_string(),
-        "ProtocolError(\"Unexpectedly saw multiple Schema messages in FlightData stream\")"
+        "Protocol error: Unexpectedly saw multiple Schema messages in FlightData stream"
     );
 }
 

diff --git a/arrow-schema/src/error.rs b/arrow-schema/src/error.rs
@@ -114,10 +114,10 @@ impl Display for ArrowError {
 
 impl Error for ArrowError {
     fn source(&self) -> Option<&(dyn Error + 'static)> {
-        if let Self::ExternalError(e) = self {
-            Some(e.as_ref())
-        } else {
-            None
+        match self {
+            ArrowError::ExternalError(source) => Some(source.as_ref()),
+            ArrowError::IoError(_, source) => Some(source),
+            _ => None,
         }
     }
 }