feat: status_code in response header (#1982)

* feat: status_code in response header

* chore: parese grpc response

* fix: sqlness failed

* chore: fix sqlness

src/client/src/database.rs

@@ -17,9 +17,9 @@ use api::v1::ddl_request::Expr as DdlExpr;
 use api::v1::greptime_request::Request;
 use api::v1::query_request::Query;
 use api::v1::{
-    greptime_response, AffectedRows, AlterExpr, AuthHeader, CompactTableExpr, CreateTableExpr,
-    DdlRequest, DeleteRequest, DropTableExpr, FlushTableExpr, GreptimeRequest, InsertRequests,
-    PromRangeQuery, QueryRequest, RequestHeader,
+    AlterExpr, AuthHeader, CompactTableExpr, CreateTableExpr, DdlRequest, DeleteRequest,
+    DropTableExpr, FlushTableExpr, GreptimeRequest, InsertRequests, PromRangeQuery, QueryRequest,
+    RequestHeader,
 };
 use arrow_flight::{FlightData, Ticket};
 use common_error::ext::{BoxedError, ErrorExt};
@@ -28,12 +28,10 @@ use common_query::Output;
 use common_telemetry::{logging, timer};
 use futures_util::{TryFutureExt, TryStreamExt};
 use prost::Message;
-use snafu::{ensure, OptionExt, ResultExt};
+use snafu::{ensure, ResultExt};
 
-use crate::error::{
-    ConvertFlightDataSnafu, IllegalDatabaseResponseSnafu, IllegalFlightMessagesSnafu,
-};
-use crate::{error, metrics, Client, Result, StreamInserter};
+use crate::error::{ConvertFlightDataSnafu, IllegalFlightMessagesSnafu, ServerSnafu};
+use crate::{error, from_grpc_response, metrics, Client, Result, StreamInserter};
 
 #[derive(Clone, Debug, Default)]
 pub struct Database {
@@ -142,16 +140,8 @@ impl Database {
     async fn handle(&self, request: Request) -> Result<u32> {
         let mut client = self.client.make_database_client()?.inner;
         let request = self.to_rpc_request(request);
-        let response = client
-            .handle(request)
-            .await?
-            .into_inner()
-            .response
-            .context(IllegalDatabaseResponseSnafu {
-                err_msg: "GreptimeResponse is empty",
-            })?;
-        let greptime_response::Response::AffectedRows(AffectedRows { value }) = response;
-        Ok(value)
+        let response = client.handle(request).await?.into_inner();
+        from_grpc_response(response)
     }
 
     #[inline]
@@ -264,7 +254,7 @@ impl Database {
                 let e: error::Error = e.into();
                 let code = e.status_code();
                 let msg = e.to_string();
-                error::ServerSnafu { code, msg }
+                ServerSnafu { code, msg }
                     .fail::<()>()
                     .map_err(BoxedError::new)
                     .context(error::FlightGetSnafu {

src/client/src/lib.rs

@@ -21,9 +21,42 @@ mod metrics;
 mod stream_insert;
 
 pub use api;
+use api::v1::greptime_response::Response;
+use api::v1::{AffectedRows, GreptimeResponse};
 pub use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
+use common_error::status_code::StatusCode;
+use snafu::OptionExt;
 
 pub use self::client::Client;
 pub use self::database::Database;
 pub use self::error::{Error, Result};
 pub use self::stream_insert::StreamInserter;
+use crate::error::{IllegalDatabaseResponseSnafu, ServerSnafu};
+
+pub fn from_grpc_response(response: GreptimeResponse) -> Result<u32> {
+    let header = response.header.context(IllegalDatabaseResponseSnafu {
+        err_msg: "missing header",
+    })?;
+    let status = header.status.context(IllegalDatabaseResponseSnafu {
+        err_msg: "missing status",
+    })?;
+
+    if StatusCode::is_success(status.status_code) {
+        let res = response.response.context(IllegalDatabaseResponseSnafu {
+            err_msg: "missing response",
+        })?;
+        match res {
+            Response::AffectedRows(AffectedRows { value }) => Ok(value),
+        }
+    } else {
+        let status_code =
+            StatusCode::from_u32(status.status_code).context(IllegalDatabaseResponseSnafu {
+                err_msg: format!("invalid status: {:?}", status),
+            })?;
+        ServerSnafu {
+            code: status_code,
+            msg: status.err_msg,
+        }
+        .fail()
+    }
+}

src/client/src/stream_insert.rs

@@ -15,17 +15,16 @@
 use api::v1::greptime_database_client::GreptimeDatabaseClient;
 use api::v1::greptime_request::Request;
 use api::v1::{
-    greptime_response, AffectedRows, AuthHeader, GreptimeRequest, GreptimeResponse, InsertRequest,
-    InsertRequests, RequestHeader,
+    AuthHeader, GreptimeRequest, GreptimeResponse, InsertRequest, InsertRequests, RequestHeader,
 };
-use snafu::OptionExt;
 use tokio::sync::mpsc;
 use tokio::task::JoinHandle;
 use tokio_stream::wrappers::ReceiverStream;
 use tonic::transport::Channel;
 use tonic::{Response, Status};
 
-use crate::error::{self, IllegalDatabaseResponseSnafu, Result};
+use crate::error::{self, Result};
+use crate::from_grpc_response;
 
 /// A structure that provides some methods for streaming data insert.
 ///
@@ -89,17 +88,8 @@ impl StreamInserter {
         drop(self.sender);
 
         let response = self.join.await.unwrap()?;
-
-        let response = response
-            .into_inner()
-            .response
-            .context(IllegalDatabaseResponseSnafu {
-                err_msg: "GreptimeResponse is empty",
-            })?;
-
-        let greptime_response::Response::AffectedRows(AffectedRows { value }) = response;
-
-        Ok(value)
+        let response = response.into_inner();
+        from_grpc_response(response)
     }
 
     fn to_rpc_request(&self, request: Request) -> GreptimeRequest {