fix: FlightData must contain valid data header (#15559)

* fix: FlightData must contain valid data header

Signed-off-by: Xuanwo <github@xuanwo.io>

* Fix toml format

Signed-off-by: Xuanwo <github@xuanwo.io>

---------

Signed-off-by: Xuanwo <github@xuanwo.io>

Author: Xuanwo

Cargo.lock

@@ -223,6 +223,8 @@ arrow-schema = { version = "51", features = ["serde"] }
 arrow-select = { version = "51" }
 parquet = { version = "51", features = ["async"] }
 parquet_rs = { package = "parquet", version = "51" }
+# Must use the same version with arrow-ipc
+flatbuffers = { version = "23" }
 
 # Crates from risingwavelabs
 arrow-udf-js = { package = "arrow-udf-js", git = "https://github.com/datafuse-extras/arrow-udf", rev = "d0a21f0" }

Cargo.toml

@@ -125,6 +125,7 @@ ctor = "0.1.26"
 dashmap = { workspace = true }
 derive-visitor = { workspace = true }
 ethnum = { workspace = true }
+flatbuffers = { workspace = true }
 futures = { workspace = true }
 futures-util = { workspace = true }
 headers = "0.4.0"

src/query/service/Cargo.toml

@@ -12,15 +12,21 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::ops::Deref;
 use std::sync::atomic::AtomicBool;
 use std::sync::atomic::Ordering;
 use std::sync::Arc;
+use std::sync::LazyLock;
 
 use arrow_flight::FlightData;
 use arrow_flight::SchemaAsIpc;
 use arrow_ipc::writer;
 use arrow_ipc::writer::IpcWriteOptions;
+use arrow_ipc::MessageBuilder;
+use arrow_ipc::MessageHeader;
+use arrow_ipc::MetadataVersion;
 use arrow_schema::Schema as ArrowSchema;
+use bytes::Bytes;
 use databend_common_base::base::tokio;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
@@ -33,6 +39,7 @@ use databend_common_sql::Planner;
 use databend_common_storages_fuse::TableContext;
 use futures::Stream;
 use futures::StreamExt;
+use prost::bytes;
 use serde::Deserialize;
 use serde::Serialize;
 use tonic::Status;
@@ -47,6 +54,29 @@ use crate::sessions::Session;
 /// A app_metakey which indicates the data is a progress type
 static H_PROGRESS: u8 = 0x01;
 
+/// The generated app metadata for our progress.
+static APP_METADATA_PROGRESS: LazyLock<Bytes> = LazyLock::new(|| Bytes::from(vec![H_PROGRESS]));
+
+/// The data header for our progress.
+///
+/// This build process is inspired from [arrow_ipc::writer::IpcDataGenerator](https://docs.rs/arrow-ipc/51.0.0/arrow_ipc/writer/struct.IpcDataGenerator.html#method.schema_to_bytes)
+static DATA_HEADER_PROGRESS: LazyLock<Bytes> = LazyLock::new(|| {
+    let mut fbb = flatbuffers::FlatBufferBuilder::new();
+
+    let mut builder = MessageBuilder::new(&mut fbb);
+    // Use the same version with arrow_ipc.
+    builder.add_version(MetadataVersion::V5);
+    // Use NONE as the header type.
+    builder.add_header_type(MessageHeader::NONE);
+    // We don't have other data to write in this message, just finish.
+    let data = builder.finish();
+
+    // finish the flat buffers.
+    fbb.finish(data, None);
+
+    Bytes::copy_from_slice(fbb.finished_data())
+});
+
 impl FlightSqlServiceImpl {
     pub(crate) fn schema_to_flight_data(data_schema: DataSchema) -> FlightData {
         let arrow_schema = ArrowSchema::from(&data_schema);
@@ -69,6 +99,16 @@ impl FlightSqlServiceImpl {
         Ok(encoded_batch.into())
     }
 
+    fn progress_to_flight_data(progress: &ProgressValue) -> Result<FlightData> {
+        let progress = serde_json::to_vec(&progress)
+            .map_err(|e| ErrorCode::Internal(format!("encode progress into json failed: {e:?}")))?;
+
+        Ok(FlightData::new()
+            .with_app_metadata(APP_METADATA_PROGRESS.deref().clone())
+            .with_data_header(DATA_HEADER_PROGRESS.deref().clone())
+            .with_data_body(Bytes::from(progress)))
+    }
+
     #[async_backtrace::framed]
     pub async fn plan_sql(
         &self,
@@ -211,12 +251,7 @@ impl FlightSqlServiceImpl {
                             progress.write_bytes = write_progress.bytes;
                         }
 
-                        let progress = serde_json::to_vec(&progress).unwrap();
-                        Some(FlightData {
-                            app_metadata: vec![H_PROGRESS].into(),
-                            data_body: progress.into(),
-                            ..Default::default()
-                        })
+                        Some(Self::progress_to_flight_data(&progress).unwrap())
                     };
 
                 while !is_finished.load(Ordering::SeqCst) {