feat: implement Arrow Flight "DoPut" in Frontend (#5836)

* feat: implement Arrow Flight "DoPut" in Frontend

* support auth for "do_put"

* set request_id in DoPut requests and responses

* set "db" in request header

src/common/grpc/Cargo.toml

@@ -23,6 +23,8 @@ flatbuffers = "24"
 hyper.workspace = true
 lazy_static.workspace = true
 prost.workspace = true
+serde.workspace = true
+serde_json.workspace = true
 snafu.workspace = true
 tokio.workspace = true
 tokio-util.workspace = true

src/common/grpc/src/error.rs

@@ -97,6 +97,14 @@ pub enum Error {
 
     #[snafu(display("Not supported: {}", feat))]
     NotSupported { feat: String },
+
+    #[snafu(display("Failed to serde Json"))]
+    SerdeJson {
+        #[snafu(source)]
+        error: serde_json::error::Error,
+        #[snafu(implicit)]
+        location: Location,
+    },
 }
 
 impl ErrorExt for Error {
@@ -110,7 +118,8 @@ impl ErrorExt for Error {
 
             Error::CreateChannel { .. }
             | Error::Conversion { .. }
-            | Error::DecodeFlightData { .. } => StatusCode::Internal,
+            | Error::DecodeFlightData { .. }
+            | Error::SerdeJson { .. } => StatusCode::Internal,
 
             Error::CreateRecordBatch { source, .. } => source.status_code(),
             Error::ConvertArrowSchema { source, .. } => source.status_code(),

src/common/grpc/src/flight.rs

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+pub mod do_put;
+
 use std::collections::HashMap;
 use std::sync::Arc;
 

src/common/grpc/src/flight/do_put.rs

@@ -0,0 +1,93 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use arrow_flight::PutResult;
+use common_base::AffectedRows;
+use serde::{Deserialize, Serialize};
+use snafu::ResultExt;
+
+use crate::error::{Error, SerdeJsonSnafu};
+
+/// The metadata for "DoPut" requests and responses.
+///
+/// Currently, there's only a "request_id", for coordinating requests and responses in the streams.
+/// Client can set a unique request id in this metadata, and the server will return the same id in
+/// the corresponding response. In doing so, a client can know how to do with its pending requests.
+#[derive(Serialize, Deserialize)]
+pub struct DoPutMetadata {
+    request_id: i64,
+}
+
+impl DoPutMetadata {
+    pub fn new(request_id: i64) -> Self {
+        Self { request_id }
+    }
+
+    pub fn request_id(&self) -> i64 {
+        self.request_id
+    }
+}
+
+/// The response in the "DoPut" returned stream.
+#[derive(Serialize, Deserialize)]
+pub struct DoPutResponse {
+    /// The same "request_id" in the request; see the [DoPutMetadata].
+    request_id: i64,
+    /// The successfully ingested rows number.
+    affected_rows: AffectedRows,
+}
+
+impl DoPutResponse {
+    pub fn new(request_id: i64, affected_rows: AffectedRows) -> Self {
+        Self {
+            request_id,
+            affected_rows,
+        }
+    }
+
+    pub fn request_id(&self) -> i64 {
+        self.request_id
+    }
+
+    pub fn affected_rows(&self) -> AffectedRows {
+        self.affected_rows
+    }
+}
+
+impl TryFrom<PutResult> for DoPutResponse {
+    type Error = Error;
+
+    fn try_from(value: PutResult) -> Result<Self, Self::Error> {
+        serde_json::from_slice(&value.app_metadata).context(SerdeJsonSnafu)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_serde_do_put_metadata() {
+        let serialized = r#"{"request_id":42}"#;
+        let metadata = serde_json::from_str::<DoPutMetadata>(serialized).unwrap();
+        assert_eq!(metadata.request_id(), 42);
+    }
+
+    #[test]
+    fn test_serde_do_put_response() {
+        let x = DoPutResponse::new(42, 88);
+        let serialized = serde_json::to_string(&x).unwrap();
+        assert_eq!(serialized, r#"{"request_id":42,"affected_rows":88}"#);
+    }
+}