feat(query): add New Flight transport

.github/actions/test_query_flight_reconnect_tpch/action.yml

@@ -0,0 +1,14 @@
+name: "Test Query Flight Reconnect with TPC-H"
+description: "Verify New Flight reconnects after a random temporary network partition"
+runs:
+  using: "composite"
+  steps:
+    - uses: ./.github/actions/setup_test
+      with:
+        artifacts: meta,query,sqllogictests
+
+    - uses: ./.github/actions/setup_minio
+
+    - name: Run TPC-H Flight reconnect test
+      shell: bash
+      run: ./scripts/ci/ci-run-query-flight-reconnect-tpch.sh

.github/workflows/reuse.linux.yml

@@ -423,6 +423,24 @@ jobs:
         with:
           name: test-stateful-cluster-linux
 
+  test_query_flight_reconnect_tpch:
+    needs: [build, check]
+    runs-on:
+      - self-hosted
+      - "${{ inputs.runner_arch }}"
+      - Linux
+      - 2c
+      - "${{ inputs.runner_provider }}"
+    steps:
+      - uses: actions/checkout@v6
+      - uses: ./.github/actions/test_query_flight_reconnect_tpch
+        timeout-minutes: 20
+      - name: Upload failure
+        if: failure() || cancelled()
+        uses: ./.github/actions/artifact_failure
+        with:
+          name: test-query-flight-reconnect-tpch
+
   test_stateful_large_data:
     if: contains(github.event.pull_request.labels.*.name, 'ci-largedata')
     needs: [build, check]

scripts/ci/ci-run-query-flight-reconnect-tpch.sh

@@ -0,0 +1,34 @@
+#!/usr/bin/env bash
+# Copyright 2020-2026 The Databend Authors.
+# SPDX-License-Identifier: Apache-2.0.
+
+set -euo pipefail
+
+export STORAGE_TYPE=s3
+export STORAGE_S3_BUCKET=testbucket
+export STORAGE_S3_ROOT=admin
+export STORAGE_S3_ENDPOINT_URL=http://127.0.0.1:9900
+export STORAGE_S3_ACCESS_KEY_ID=minioadmin
+export STORAGE_S3_SECRET_ACCESS_KEY=minioadmin
+export STORAGE_ALLOW_INSECURE=true
+
+readonly BUILD_PROFILE="${BUILD_PROFILE:-debug}"
+readonly SCRIPT_PATH="$(cd "$(dirname "$0")" >/dev/null 2>&1 && pwd)"
+readonly REPO_PATH="$(cd "$SCRIPT_PATH/../.." >/dev/null 2>&1 && pwd)"
+readonly TPCH_DATA_PATH=/tmp/tpch_1
+
+python3 -m pip install --quiet mysql-connector-python requests
+sudo apt-get update -yq
+sudo apt-get install -yq iproute2 iptables lsof
+
+cd "$REPO_PATH"
+./scripts/ci/deploy/databend-query-cluster-3-nodes.sh
+
+rm -rf -- "$TPCH_DATA_PATH"
+bash tests/sqllogictests/scripts/prepare_tpch_data.sh tpch_test 1
+
+python3 tests/query-flight-reconnect/test_tpch_reconnect.py \
+	--sqllogictests "target/${BUILD_PROFILE}/databend-sqllogictests" \
+	--tpch-suite tests/sqllogictests/suites/tpch/queries.test \
+	--operation-log .databend/tpch-flight-reconnect/operations.log \
+	--repo-dir "$REPO_PATH"

src/query/service/src/schedulers/fragments/query_fragment_actions.rs

@@ -230,6 +230,8 @@ impl QueryFragmentsActions {
 
     /// unique map(target, map(source, vec(fragment_id)))
     fn fragments_connections(&self, builder: &mut DataflowDiagramBuilder) -> Result<()> {
+        let new_flight = self.ctx.get_settings().get_enable_experiment_new_flight()?;
+
         for fragment_actions in &self.fragments_actions {
             if let Some(exchange) = &fragment_actions.data_exchange {
                 let destinations = exchange.get_destinations();
@@ -249,7 +251,11 @@ impl QueryFragmentsActions {
                             )?;
                         } else {
                             for channel in exchange.get_channels(destination) {
-                                builder.add_data_edge(&source, destination, &channel)?;
+                                if new_flight && matches!(exchange, DataExchange::Merge(_)) {
+                                    builder.add_merge_edge(&source, destination, &channel)?;
+                                } else {
+                                    builder.add_data_edge(&source, destination, &channel)?;
+                                }
                             }
                         }
                     }

src/query/service/src/servers/flight/flight_client.rs

@@ -20,8 +20,6 @@ use arrow_flight::FlightData;
 use arrow_flight::Ticket;
 use arrow_flight::flight_service_client::FlightServiceClient;
 use async_channel::Receiver;
-use async_channel::Sender;
-use databend_common_base::runtime::drop_guard;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use fastrace::Span;
@@ -33,22 +31,128 @@ use serde::Deserialize;
 use serde::Serialize;
 use tokio::time::Duration;
 use tonic::Request;
-use tonic::Status;
 use tonic::Streaming;
 use tonic::metadata::AsciiMetadataKey;
 use tonic::metadata::AsciiMetadataValue;
 use tonic::transport::channel::Channel;
 
 use crate::pipelines::executor::WatchNotify;
 use crate::servers::flight::request_builder::RequestBuilder;
-use crate::servers::flight::v1::packets::DataPacket;
+use crate::servers::flight::v1::transport::legacy::LegacyInbound;
 
 /// Parameters for a do_exchange RPC call, serialized as JSON in metadata.
-#[derive(Serialize, Deserialize)]
+#[derive(Clone, Serialize, Deserialize)]
 pub struct DoExchangeParams {
     pub query_id: String,
     pub exchange_id: String,
     pub num_threads: usize,
+    /// Present only for New Flight streams. Absent keeps the existing Flight wire shape, so a
+    /// node running either transport can parse the other's requests.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub new_flight: Option<NewFlightAttachment>,
+}
+
+/// Identifies a New Flight logical stream and how long its receiver waits for a replacement
+/// connection. Keeping these together makes a partially specified attachment unrepresentable.
+#[derive(Clone, Serialize, Deserialize)]
+pub struct NewFlightAttachment {
+    pub source_id: String,
+    pub receiver_lease_secs: u64,
+    pub stream: NewFlightStream,
+}
+
+#[derive(Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum NewFlightStream {
+    Fragment,
+    Exchange,
+    Merge,
+    Statistics,
+}
+
+impl DoExchangeParams {
+    pub fn create(query_id: String, exchange_id: String, num_threads: usize) -> Self {
+        Self {
+            query_id,
+            exchange_id,
+            num_threads,
+            new_flight: None,
+        }
+    }
+
+    pub fn new_flight_fragment(
+        query_id: String,
+        exchange_id: String,
+        source_id: String,
+        num_threads: usize,
+        receiver_lease_secs: u64,
+    ) -> Self {
+        Self {
+            query_id,
+            exchange_id,
+            num_threads,
+            new_flight: Some(NewFlightAttachment {
+                source_id,
+                receiver_lease_secs,
+                stream: NewFlightStream::Fragment,
+            }),
+        }
+    }
+
+    pub fn new_flight_merge(
+        query_id: String,
+        exchange_id: String,
+        source_id: String,
+        receiver_lease_secs: u64,
+    ) -> Self {
+        Self {
+            query_id,
+            exchange_id,
+            num_threads: 1,
+            new_flight: Some(NewFlightAttachment {
+                source_id,
+                receiver_lease_secs,
+                stream: NewFlightStream::Merge,
+            }),
+        }
+    }
+
+    pub fn new_flight_exchange(
+        query_id: String,
+        exchange_id: String,
+        source_id: String,
+        num_threads: usize,
+        receiver_lease_secs: u64,
+    ) -> Self {
+        Self {
+            query_id,
+            exchange_id,
+            num_threads,
+            new_flight: Some(NewFlightAttachment {
+                source_id,
+                receiver_lease_secs,
+                stream: NewFlightStream::Exchange,
+            }),
+        }
+    }
+
+    pub fn new_flight_statistics(
+        query_id: String,
+        source_id: String,
+        receiver_lease_secs: u64,
+    ) -> Self {
+        Self {
+            // Statistics is a single stream per source, so it needs no exchange id or thread fan-out.
+            query_id,
+            exchange_id: String::new(),
+            num_threads: 1,
+            new_flight: Some(NewFlightAttachment {
+                source_id,
+                receiver_lease_secs,
+                stream: NewFlightStream::Statistics,
+            }),
+        }
+    }
 }
 
 pub struct FlightClient {
@@ -196,7 +300,7 @@ impl FlightClient {
         &mut self,
         query_id: &str,
         target: &str,
-    ) -> Result<FlightExchange> {
+    ) -> Result<LegacyInbound> {
         let streaming = self
             .get_streaming(
                 RequestBuilder::create(Ticket::default())
@@ -212,12 +316,12 @@ impl FlightClient {
             self.local_node_id.clone(),
             self.remote_node_id.clone(),
         );
-        Ok(FlightExchange::create_receiver(notify, rx))
+        Ok(LegacyInbound::create(notify, rx))
     }
 
     #[async_backtrace::framed]
     #[fastrace::trace]
-    pub async fn do_get(&mut self, query_id: &str, channel_id: &str) -> Result<FlightExchange> {
+    pub async fn do_get(&mut self, query_id: &str, channel_id: &str) -> Result<LegacyInbound> {
         let request = RequestBuilder::create(Ticket::default())
             .with_metadata("x-type", "exchange_fragment")?
             .with_metadata("x-query-id", query_id)?
@@ -232,7 +336,7 @@ impl FlightClient {
             self.local_node_id.clone(),
             self.remote_node_id.clone(),
         );
-        Ok(FlightExchange::create_receiver(notify, rx))
+        Ok(LegacyInbound::create(notify, rx))
     }
 
     fn streaming_receiver(
@@ -331,109 +435,3 @@ impl FlightClient {
         })
     }
 }
-
-pub struct FlightReceiver {
-    notify: Arc<WatchNotify>,
-    rx: Receiver<Result<FlightData>>,
-}
-
-impl Drop for FlightReceiver {
-    fn drop(&mut self) {
-        drop_guard(move || {
-            self.close();
-        })
-    }
-}
-
-impl FlightReceiver {
-    pub fn create(rx: Receiver<Result<FlightData>>) -> FlightReceiver {
-        FlightReceiver {
-            rx,
-            notify: Arc::new(WatchNotify::new()),
-        }
-    }
-
-    #[async_backtrace::framed]
-    pub async fn recv(&self) -> Result<Option<DataPacket>> {
-        match self.rx.recv().await {
-            Err(_) => Ok(None),
-            Ok(Err(error)) => Err(error),
-            Ok(Ok(message)) => Ok(Some(DataPacket::try_from(message)?)),
-        }
-    }
-
-    pub fn close(&self) {
-        self.rx.close();
-        self.notify.notify_waiters();
-    }
-}
-
-pub struct FlightSender {
-    tx: Sender<std::result::Result<FlightData, Status>>,
-}
-
-impl FlightSender {
-    pub fn create(tx: Sender<std::result::Result<FlightData, Status>>) -> FlightSender {
-        FlightSender { tx }
-    }
-
-    pub fn is_closed(&self) -> bool {
-        self.tx.is_closed()
-    }
-
-    #[async_backtrace::framed]
-    pub async fn send(&self, data: DataPacket) -> Result<()> {
-        if let Err(_cause) = self.tx.send(Ok(FlightData::try_from(data)?)).await {
-            return Err(ErrorCode::AbortedQuery(
-                "Aborted query, because the remote flight channel is closed.",
-            ));
-        }
-
-        Ok(())
-    }
-
-    pub fn close(&self) {
-        self.tx.close();
-    }
-}
-
-pub enum FlightExchange {
-    Dummy,
-    Receiver {
-        notify: Arc<WatchNotify>,
-        receiver: Receiver<Result<FlightData>>,
-    },
-    Sender(Sender<std::result::Result<FlightData, Status>>),
-}
-
-impl FlightExchange {
-    pub fn create_sender(
-        sender: Sender<std::result::Result<FlightData, Status>>,
-    ) -> FlightExchange {
-        FlightExchange::Sender(sender)
-    }
-
-    pub fn create_receiver(
-        notify: Arc<WatchNotify>,
-        receiver: Receiver<Result<FlightData>>,
-    ) -> FlightExchange {
-        FlightExchange::Receiver { notify, receiver }
-    }
-
-    pub fn convert_to_sender(self) -> FlightSender {
-        match self {
-            FlightExchange::Sender(tx) => FlightSender { tx },
-            _ => unreachable!(),
-        }
-    }
-
-    pub fn convert_to_receiver(self) -> FlightReceiver {
-        match self {
-            FlightExchange::Receiver { notify, receiver } => FlightReceiver {
-                notify,
-                rx: receiver,
-            },
-            _ => unreachable!(),
-        }
-    }
-}

src/query/service/src/servers/flight/mod.rs

@@ -20,9 +20,8 @@ pub mod v1;
 
 pub use flight_client::DoExchangeParams;
 pub use flight_client::FlightClient;
-pub use flight_client::FlightExchange;
 pub(crate) use flight_client::FlightOperation;
-pub use flight_client::FlightReceiver;
-pub use flight_client::FlightSender;
+pub use flight_client::NewFlightAttachment;
+pub use flight_client::NewFlightStream;
 pub(crate) use flight_client::add_flight_error_context;
 pub use flight_service::FlightService;