From 64243e3a7d8007c3b281847d378c770abb069aae Mon Sep 17 00:00:00 2001
From: LFC <luofucong@greptime.com>
Date: Sun, 29 Jan 2023 11:27:13 +0800
Subject: [PATCH] refactor: accommodate java flight client (#886)

* refactor: change how AffectedRows is carried in flight stream to accommodate Java Flight client

* fix: clippy
---
 src/api/greptime/v1/database.proto |  8 ++++++--
 src/common/grpc/src/flight.rs      | 18 ++++++++++++------
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/src/api/greptime/v1/database.proto b/src/api/greptime/v1/database.proto
index 8f15925246d..56db5582937 100644
--- a/src/api/greptime/v1/database.proto
+++ b/src/api/greptime/v1/database.proto
@@ -36,6 +36,10 @@ message InsertRequest {
   uint32 region_number = 5;
 }
 
-message FlightDataExt {
-  uint32 affected_rows = 1;
+message AffectedRows {
+  uint32 value = 1;
+}
+
+message FlightMetadata {
+  AffectedRows affected_rows = 1;
 }
diff --git a/src/common/grpc/src/flight.rs b/src/common/grpc/src/flight.rs
index 77c4eee4d88..a34918639a7 100644
--- a/src/common/grpc/src/flight.rs
+++ b/src/common/grpc/src/flight.rs
@@ -15,7 +15,7 @@
 use std::collections::HashMap;
 use std::sync::Arc;
 
-use api::v1::FlightDataExt;
+use api::v1::{AffectedRows, FlightMetadata};
 use arrow_flight::utils::{flight_data_from_arrow_batch, flight_data_to_arrow_batch};
 use arrow_flight::{FlightData, IpcMessage, SchemaAsIpc};
 use common_recordbatch::{RecordBatch, RecordBatches};
@@ -66,11 +66,11 @@ impl FlightEncoder {
                 flight_batch
             }
             FlightMessage::AffectedRows(rows) => {
-                let ext_data = FlightDataExt {
-                    affected_rows: rows as _,
+                let metadata = FlightMetadata {
+                    affected_rows: Some(AffectedRows { value: rows as _ }),
                 }
                 .encode_to_vec();
-                FlightData::new(None, IpcMessage(build_none_flight_msg()), vec![], ext_data)
+                FlightData::new(None, IpcMessage(build_none_flight_msg()), metadata, vec![])
             }
         }
     }
@@ -91,9 +91,15 @@ impl FlightDecoder {
         })?;
         match message.header_type() {
             MessageHeader::NONE => {
-                let ext_data = FlightDataExt::decode(flight_data.data_body.as_slice())
+                let metadata = FlightMetadata::decode(flight_data.app_metadata.as_slice())
                     .context(DecodeFlightDataSnafu)?;
-                Ok(FlightMessage::AffectedRows(ext_data.affected_rows as _))
+                if let Some(AffectedRows { value }) = metadata.affected_rows {
+                    return Ok(FlightMessage::AffectedRows(value as _));
+                }
+                InvalidFlightDataSnafu {
+                    reason: "Expecting FlightMetadata have some meaningful content.",
+                }
+                .fail()
             }
             MessageHeader::Schema => {
                 let arrow_schema = ArrowSchema::try_from(&flight_data).map_err(|e| {