add udf fixes

risingwavelabs · Mar 7, 2024 · acaa3ed · acaa3ed
1 parent e5a64f8
commit acaa3ed
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 6 deletions.
diff --git a/src/expr/udf/src/error.rs b/src/expr/udf/src/error.rs
@@ -23,9 +23,6 @@ pub type Result<T, E = Error> = std::result::Result<T, E>;
 #[derive(Error, Debug, Box, Construct)]
 #[thiserror_ext(newtype(name = Error))]
 pub enum ErrorInner {
-    #[error("failed to connect to UDF service: {0}")]
-    Connect(#[from] tonic::transport::Error),
-
     #[error("failed to send requests to UDF service: {0}")]
     Tonic(#[from] tonic::Status),
 
@@ -58,6 +55,10 @@ impl Error {
             _ => false,
         }
     }
+
+    pub fn is_tonic_error(&self) -> bool {
+        matches!(self.inner(), ErrorInner::Tonic(_))
+    }
 }
 
 static_assertions::const_assert_eq!(std::mem::size_of::<Error>(), 8);
diff --git a/src/expr/udf/src/external.rs b/src/expr/udf/src/external.rs
@@ -213,10 +213,15 @@ impl ArrowFlightUdfClient {
         let mut backoff = Duration::from_millis(100);
         loop {
             match self.call(id, input.clone()).await {
-                Err(err) if err.is_connection_error() => {
-                    tracing::error!(error = %err.as_report(), "UDF connection error. retry...");
+                Err(err) if err.is_tonic_error() => {
+                    tracing::error!(error = %err.as_report(), "UDF tonic error. retry...");
+                }
+                ret => {
+                    if ret.is_err() {
+                        tracing::error!(error = %ret.as_ref().unwrap_err().as_report(), "UDF error. exiting...");
+                    }
+                    return ret;
                 }
-                ret => return ret,
             }
             tokio::time::sleep(backoff).await;
             backoff *= 2;