fix: fix incorrect matches

GreptimeTeam · Sep 18, 2023 · 906ca48 · 906ca48
1 parent e7e254c
commit 906ca48
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 6 deletions.
diff --git a/src/client/src/error.rs b/src/client/src/error.rs
@@ -72,6 +72,9 @@ pub enum Error {
         source: common_grpc::error::Error,
     },
 
+    #[snafu(display("Failed to request RegionServer, code: {}, source: {}", code, source))]
+    RegionServer { code: Code, source: BoxedError },
+
     // Server error carried in Tonic Status's metadata.
     #[snafu(display("{}", msg))]
     Server { code: StatusCode, msg: String },
@@ -95,9 +98,9 @@ impl ErrorExt for Error {
             | Error::ClientStreaming { .. } => StatusCode::Internal,
 
             Error::Server { code, .. } => *code,
-            Error::FlightGet { source, .. } | Error::HandleRequest { source, .. } => {
-                source.status_code()
-            }
+            Error::FlightGet { source, .. }
+            | Error::HandleRequest { source, .. }
+            | Error::RegionServer { source, .. } => source.status_code(),
             Error::CreateChannel { source, .. } | Error::ConvertFlightData { source, .. } => {
                 source.status_code()
             }

diff --git a/src/client/src/region.rs b/src/client/src/region.rs
@@ -29,7 +29,7 @@ use prost::Message;
 use snafu::{location, Location, OptionExt, ResultExt};
 use tokio_stream::StreamExt;
 
-use crate::error::Error::FlightGet;
+use crate::error::Error::RegionServer;
 use crate::error::{
     self, ConvertFlightDataSnafu, IllegalDatabaseResponseSnafu, IllegalFlightMessagesSnafu,
     MissingFieldSnafu, Result, ServerSnafu,
@@ -45,7 +45,7 @@ pub struct RegionRequester {
 impl Datanode for RegionRequester {
     async fn handle(&self, request: RegionRequest) -> MetaResult<AffectedRows> {
         self.handle_inner(request).await.map_err(|err| {
-            if matches!(err, FlightGet { .. }) {
+            if matches!(err, RegionServer { .. }) {
                 meta_error::Error::RetryLater {
                     source: BoxedError::new(err),
                 }
@@ -163,7 +163,19 @@ impl RegionRequester {
         let RegionResponse {
             header,
             affected_rows,
-        } = client.handle(request).await?.into_inner();
+        } = client
+            .handle(request)
+            .await
+            .map_err(|e| {
+                let code = e.code();
+                let err: error::Error = e.into();
+                // Uses `Error::RegionServer` instead of `Error::Server`
+                error::Error::RegionServer {
+                    code,
+                    source: BoxedError::new(err),
+                }
+            })?
+            .into_inner();
 
         check_response_header(header)?;