client/

database.rs

// 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 std::collections::HashMap;
use std::pin::Pin;
use std::str::FromStr;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, RwLock};
use std::task::{Context, Poll};
use std::time::Duration;

use api::v1::auth_header::AuthScheme;
#[cfg(feature = "testing")]
use api::v1::ddl_request::Expr as DdlExpr;
use api::v1::greptime_database_client::GreptimeDatabaseClient;
use api::v1::greptime_request::Request;
use api::v1::query_request::Query;
#[cfg(feature = "testing")]
use api::v1::{AlterTableExpr, CreateTableExpr, DdlRequest};
use api::v1::{
    AuthHeader, Basic, GreptimeRequest, InsertRequests, QueryRequest, RequestHeader,
    RowInsertRequests,
};
use arc_swap::ArcSwapOption;
use arrow_flight::{FlightData, Ticket};
use async_stream::stream;
use base64::Engine;
use base64::prelude::BASE64_STANDARD;
use common_catalog::build_db_string;
use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
use common_error::ext::{BoxedError, ErrorExt};
use common_grpc::flight::do_put::DoPutResponse;
use common_grpc::flight::{
    FLOW_EXTENSIONS_METADATA_KEY, FlightDecoder, FlightMessage, SNAPSHOT_SEQS_METADATA_KEY,
};
use common_query::Output;
use common_recordbatch::adapter::RecordBatchMetrics;
use common_recordbatch::error::ExternalSnafu;
use common_recordbatch::{OrderOption, RecordBatch, RecordBatchStream, RecordBatchStreamWrapper};
use common_telemetry::tracing::Span;
use common_telemetry::tracing_context::W3cTrace;
use common_telemetry::{error, warn};
use futures::future;
use futures_util::{Stream, StreamExt, TryStreamExt};
use prost::Message;
use snafu::{IntoError, ResultExt};
use tokio::sync::Notify;
use tonic::metadata::{AsciiMetadataKey, AsciiMetadataValue, MetadataMap, MetadataValue};
use tonic::transport::Channel;

use crate::error::{
    ConvertFlightDataSnafu, Error, FlightGetSnafu, FlightStreamSnafu, IllegalFlightMessagesSnafu,
    InvalidTonicMetadataValueSnafu,
};
use crate::flight::{FlightMessageReader, decode_flight_data};
use crate::{Client, Result, error, from_grpc_response};

type FlightDataStream = Pin<Box<dyn Stream<Item = FlightData> + Send>>;

type DoPutResponseStream = Pin<Box<dyn Stream<Item = Result<DoPutResponse>>>>;

const HINTS_METADATA_KEY: &str = "x-greptime-hints";
/// Maximum time to wait for the optional trailing metrics message after
/// affected rows have already been delivered.
const FLIGHT_TRAILING_METRICS_TIMEOUT: Duration = Duration::from_secs(5);

/// Terminal metrics associated with a query output.
///
/// For streaming outputs, metrics are only final after the stream is fully
/// drained and [`Self::is_ready`] returns `true`. Affected-row outputs may
/// briefly await a compatibility trailing metrics message.
#[derive(Debug, Clone, Default)]
pub struct OutputMetrics {
    inner: Arc<OutputMetricsInner>,
}

#[derive(Debug)]
struct OutputMetricsInner {
    metrics: RwLock<Option<RecordBatchMetrics>>,
    completion_error: RwLock<Option<String>>,
    ready: AtomicBool,
    ready_notify: Notify,
    compatibility_task: Mutex<Option<tokio::task::AbortHandle>>,
}

impl Default for OutputMetricsInner {
    fn default() -> Self {
        Self {
            metrics: RwLock::new(None),
            completion_error: RwLock::new(None),
            ready: AtomicBool::new(false),
            ready_notify: Notify::new(),
            compatibility_task: Mutex::new(None),
        }
    }
}

impl Drop for OutputMetricsInner {
    fn drop(&mut self) {
        if let Some(handle) = self.compatibility_task.get_mut().unwrap().take() {
            handle.abort();
        }
    }
}

impl OutputMetrics {
    fn new() -> Self {
        Self::default()
    }

    /// Replaces the current terminal metrics snapshot.
    pub fn update(&self, metrics: Option<RecordBatchMetrics>) {
        *self.inner.metrics.write().unwrap() = metrics;
    }

    /// Marks the terminal metrics as final for this output.
    pub fn mark_ready(&self) {
        if self
            .inner
            .ready
            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
            .is_ok()
        {
            self.inner.ready_notify.notify_waiters();
        }
    }

    /// Waits until terminal metrics are final.
    pub async fn wait_ready(&self) {
        loop {
            let notified = self.inner.ready_notify.notified();
            if self.is_ready() {
                return;
            }
            notified.await;
        }
    }

    /// Returns an error encountered while completing the output, if any.
    pub fn completion_error(&self) -> Option<String> {
        self.inner.completion_error.read().unwrap().clone()
    }

    fn set_completion_error(&self, error: impl Into<String>) {
        *self.inner.completion_error.write().unwrap() = Some(error.into());
    }

    fn set_compatibility_task(&self, handle: tokio::task::AbortHandle) {
        let mut task = self.inner.compatibility_task.lock().unwrap();
        if !self.is_ready() {
            *task = Some(handle);
        }
    }

    fn take_compatibility_task(&self) -> Option<tokio::task::AbortHandle> {
        self.inner.compatibility_task.lock().unwrap().take()
    }

    /// Returns whether terminal metrics are final.
    ///
    /// Streaming outputs become ready only after the stream reaches EOF.
    pub fn is_ready(&self) -> bool {
        self.inner.ready.load(Ordering::Acquire)
    }

    /// Returns the latest terminal metrics snapshot, if any.
    pub fn get(&self) -> Option<RecordBatchMetrics> {
        self.inner.metrics.read().unwrap().clone()
    }

    /// Returns proved per-region watermarks.
    ///
    /// Entries whose watermark is `None` are intentionally omitted because they
    /// represent participating regions whose terminal sequence bound was not
    /// provable.
    pub fn region_watermark_map(&self) -> Option<std::collections::HashMap<u64, u64>> {
        Some(
            self.get()?
                .region_watermarks
                .into_iter()
                .filter_map(|entry| entry.watermark.map(|seq| (entry.region_id, seq)))
                .collect::<std::collections::HashMap<_, _>>(),
        )
    }

    /// Returns all regions that participated in terminal metric collection,
    /// including entries whose watermark is `None`.
    pub fn participating_regions(&self) -> Option<std::collections::BTreeSet<u64>> {
        Some(
            self.get()?
                .region_watermarks
                .into_iter()
                .map(|entry| entry.region_id)
                .collect::<std::collections::BTreeSet<_>>(),
        )
    }
}

/// Query output together with a handle for its terminal metrics.
///
/// The contained [`OutputMetrics`] lets callers read stream terminal metrics
/// after consuming `output`. For non-stream outputs, metrics are ready
/// immediately. Flight affected-row outputs without inline metrics require
/// [`OutputMetrics::wait_ready`] before compatibility trailing metrics can be read.
#[derive(Debug)]
pub struct OutputWithMetrics {
    pub output: Output,
    pub metrics: OutputMetrics,
}

impl OutputWithMetrics {
    /// Wraps an output with a terminal metrics handle.
    ///
    /// Stream outputs update the handle as the stream is consumed. Non-stream
    /// outputs are marked ready immediately.
    pub fn from_output(output: Output) -> Self {
        let terminal_metrics = OutputMetrics::new();
        let output = attach_terminal_metrics(output, &terminal_metrics);
        Self {
            output,
            metrics: terminal_metrics,
        }
    }

    /// Returns proved per-region watermarks from the terminal metrics.
    pub fn region_watermark_map(&self) -> Option<std::collections::HashMap<u64, u64>> {
        self.metrics.region_watermark_map()
    }

    /// Returns all regions participating in terminal metric collection.
    pub fn participating_regions(&self) -> Option<std::collections::BTreeSet<u64>> {
        self.metrics.participating_regions()
    }

    /// Drops the terminal metrics handle and returns the original output.
    pub fn into_output(self) -> Output {
        self.output
    }
}

fn parse_terminal_metrics(metrics_json: &str) -> Result<RecordBatchMetrics> {
    serde_json::from_str(metrics_json).map_err(|e| {
        IllegalFlightMessagesSnafu {
            reason: format!("Invalid terminal metrics message: {e}"),
        }
        .build()
    })
}

fn spawn_affected_rows_trailing_metrics_task<S>(
    terminal_metrics: &OutputMetrics,
    mut reader: FlightMessageReader<S>,
) where
    S: Stream<Item = Result<FlightMessage>> + Send + Unpin + 'static,
{
    let metrics_ref = Arc::downgrade(&terminal_metrics.inner);
    let remote_addr = reader.remote_addr().to_string();
    let task = common_runtime::spawn_global(async move {
        let result =
            tokio::time::timeout(FLIGHT_TRAILING_METRICS_TIMEOUT, reader.read_next()).await;
        let Some(inner) = metrics_ref.upgrade() else {
            return;
        };
        let metrics = OutputMetrics { inner };
        match result {
            Ok(Ok(Some(FlightMessage::Metrics(s)))) => match parse_terminal_metrics(&s) {
                Ok(metrics_json) => metrics.update(Some(metrics_json)),
                Err(error) => {
                    metrics.set_completion_error(error.to_string());
                    warn!(
                        "Failed to decode trailing Flight metrics from {}: {}",
                        remote_addr, error
                    );
                }
            },
            Ok(Ok(None)) => {}
            Ok(Ok(Some(other))) => {
                let error = format!("Unexpected trailing Flight message: {other:?}");
                metrics.set_completion_error(error.clone());
                warn!("{} from {}", error, remote_addr);
            }
            Ok(Err(error)) => {
                let error = flight_stream_error(&remote_addr, error);
                metrics.set_completion_error(error.to_string());
                warn!("{}", error);
            }
            Err(_) => {
                let error = "Timed out waiting for trailing Flight metrics";
                metrics.set_completion_error(error);
                warn!("{} from {}", error, remote_addr);
            }
        }
        metrics.mark_ready();
        metrics.take_compatibility_task();
    });
    terminal_metrics.set_compatibility_task(task.abort_handle());
}

struct StreamWithMetrics {
    stream: common_recordbatch::SendableRecordBatchStream,
    metrics: OutputMetrics,
}

impl StreamWithMetrics {
    fn new(stream: common_recordbatch::SendableRecordBatchStream, metrics: OutputMetrics) -> Self {
        Self { stream, metrics }
    }

    fn sync_terminal_metrics(&self) {
        self.metrics.update(self.stream.metrics());
    }
}

impl RecordBatchStream for StreamWithMetrics {
    fn name(&self) -> &str {
        self.stream.name()
    }

    fn schema(&self) -> datatypes::schema::SchemaRef {
        self.stream.schema()
    }

    fn output_ordering(&self) -> Option<&[OrderOption]> {
        self.stream.output_ordering()
    }

    fn metrics(&self) -> Option<RecordBatchMetrics> {
        self.sync_terminal_metrics();
        self.metrics.get()
    }
}

impl Stream for StreamWithMetrics {
    type Item = common_recordbatch::error::Result<RecordBatch>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let polled = Pin::new(&mut self.stream).poll_next(cx);
        if let Poll::Ready(Some(Err(error))) = &polled {
            self.metrics.set_completion_error(error.to_string());
        }
        if let Poll::Ready(None) = &polled {
            self.sync_terminal_metrics();
            self.metrics.mark_ready();
        }
        polled
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.stream.size_hint()
    }
}

fn attach_terminal_metrics(output: Output, terminal_metrics: &OutputMetrics) -> Output {
    let Output { data, meta } = output;
    let data = match data {
        common_query::OutputData::Stream(stream) => {
            terminal_metrics.update(stream.metrics());
            common_query::OutputData::Stream(Box::pin(StreamWithMetrics::new(
                stream,
                terminal_metrics.clone(),
            )))
        }
        other => {
            terminal_metrics.mark_ready();
            other
        }
    };
    Output::new(data, meta)
}

async fn output_from_flight_message_stream<S>(
    remote_addr: String,
    flight_message_stream: S,
) -> Result<OutputWithMetrics>
where
    S: Stream<Item = Result<FlightMessage>> + Send + Unpin + 'static,
{
    let mut reader = FlightMessageReader::new(remote_addr, flight_message_stream);
    let first_flight_message = reader
        .read_first()
        .await
        .map_err(|error| flight_stream_error(reader.remote_addr(), error))?;

    match first_flight_message {
        FlightMessage::AffectedRows { rows, metrics } => {
            let terminal_metrics = OutputMetrics::new();
            if let Some(metrics) = metrics {
                // Inline metrics are authoritative. They complete the output
                // without touching the rest of the Flight stream.
                terminal_metrics.update(Some(parse_terminal_metrics(&metrics)?));
                terminal_metrics.mark_ready();
            } else {
                spawn_affected_rows_trailing_metrics_task(&terminal_metrics, reader);
            }
            Ok(OutputWithMetrics {
                output: Output::new_with_affected_rows(rows),
                metrics: terminal_metrics,
            })
        }
        FlightMessage::RecordBatch(_) | FlightMessage::Metrics(_) => IllegalFlightMessagesSnafu {
            reason: "The first flight message cannot be a RecordBatch or Metrics message",
        }
        .fail(),
        FlightMessage::Schema(schema) => {
            let metrics = Arc::new(ArcSwapOption::from(None));
            let metrics_ref = metrics.clone();
            let schema = Arc::new(
                datatypes::schema::Schema::try_from(schema).context(error::ConvertSchemaSnafu)?,
            );
            let schema_cloned = schema.clone();
            let stream = Box::pin(stream!({
                loop {
                    let flight_message = match reader.read_next().await {
                        Ok(Some(message)) => message,
                        Ok(None) => break,
                        Err(error) => {
                            yield Err(BoxedError::new(flight_stream_error(
                                reader.remote_addr(),
                                error,
                            )))
                            .context(ExternalSnafu);
                            break;
                        }
                    };
                    match flight_message {
                        FlightMessage::RecordBatch(arrow_batch) => {
                            yield Ok(RecordBatch::from_df_record_batch(
                                schema_cloned.clone(),
                                arrow_batch,
                            ))
                        }
                        FlightMessage::Metrics(s) => {
                            match parse_terminal_metrics(&s) {
                                Ok(m) => {
                                    metrics_ref.swap(Some(Arc::new(m)));
                                }
                                Err(e) => {
                                    yield Err(BoxedError::new(e)).context(ExternalSnafu);
                                }
                            };
                        }
                        FlightMessage::AffectedRows { .. } | FlightMessage::Schema(_) => {
                            yield IllegalFlightMessagesSnafu {
                                reason: format!(
                                    "A Schema message must be succeeded exclusively by a set of RecordBatch messages, flight_message: {:?}",
                                    flight_message
                                )
                            }
                            .fail()
                            .map_err(BoxedError::new)
                            .context(ExternalSnafu);
                            break;
                        }
                    }
                }
            }));
            let record_batch_stream = RecordBatchStreamWrapper {
                schema,
                stream,
                output_ordering: None,
                metrics,
                span: Span::current(),
            };
            Ok(OutputWithMetrics::from_output(Output::new_with_stream(
                Box::pin(record_batch_stream),
            )))
        }
    }
}

fn flight_stream_error(addr: &str, error: Error) -> Error {
    let tonic_code = error.tonic_code().unwrap_or(tonic::Code::Unknown);
    let message = error.to_string();
    if error.status_code().should_log_error() {
        error!(
            error; "Failed to receive Flight data, addr: {}, code: {}",
            addr,
            tonic_code
        );
    }

    FlightStreamSnafu {
        addr: addr.to_string(),
        tonic_code,
        message,
    }
    .into_error(BoxedError::new(error))
}

#[derive(Clone, Debug, Default)]
pub struct Database {
    // The "catalog" and "schema" to be used in processing the requests at the server side.
    // They are the "hint" or "context", just like how the "database" in "USE" statement is treated in MySQL.
    // They will be carried in the request header.
    catalog: String,
    schema: String,
    // The dbname follows naming rule as out mysql, postgres and http
    // protocol. The server treat dbname in priority of catalog/schema.
    dbname: String,
    // The time zone indicates the time zone where the user is located.
    // Some queries need to be aware of the user's time zone to perform some specific actions.
    timezone: String,

    client: Client,
    ctx: FlightContext,
}

#[derive(Default)]
struct FlightRequestOptions {
    hints: Option<String>,
    flow_extensions: Option<String>,
    snapshot_seqs: Option<String>,
    timeout: Option<Duration>,
}

impl FlightRequestOptions {
    fn apply_to<T>(self, request: &mut tonic::Request<T>) -> Result<()> {
        let metadata = request.metadata_mut();
        if let Some(hints) = self.hints {
            Database::put_metadata_value(metadata, HINTS_METADATA_KEY, hints)?;
        }
        if let Some(flow_extensions) = self.flow_extensions {
            Database::put_metadata_value(metadata, FLOW_EXTENSIONS_METADATA_KEY, flow_extensions)?;
        }
        if let Some(snapshot_seqs) = self.snapshot_seqs {
            Database::put_metadata_value(metadata, SNAPSHOT_SEQS_METADATA_KEY, snapshot_seqs)?;
        }
        if let Some(timeout) = self.timeout {
            request.set_timeout(timeout);
        }
        Ok(())
    }
}

/// A single Flight DoGet request to a [`Database`].
///
/// The builder carries request-scoped metadata and options. It does not modify
/// the underlying [`Database`], so its configuration cannot affect later RPCs.
pub struct DatabaseFlightRequest<'a> {
    database: &'a Database,
    options: FlightRequestOptions,
}

pub struct DatabaseClient {
    pub addr: String,
    pub inner: GreptimeDatabaseClient<Channel>,
}

impl DatabaseClient {
    /// Returns a closure that logs the error when the request fails.
    pub fn inspect_err<'a>(&'a self, context: &'a str) -> impl Fn(&tonic::Status) + 'a {
        let addr = &self.addr;
        move |status| {
            error!("Failed to {context} request, peer: {addr}, status: {status:?}");
        }
    }
}

fn make_database_client(client: &Client) -> Result<DatabaseClient> {
    let (addr, channel) = client.find_channel()?;
    Ok(DatabaseClient {
        addr,
        inner: GreptimeDatabaseClient::new(channel)
            .max_decoding_message_size(client.max_grpc_recv_message_size())
            .max_encoding_message_size(client.max_grpc_send_message_size()),
    })
}

impl Database {
    /// Create database service client using catalog and schema
    pub fn new(catalog: impl Into<String>, schema: impl Into<String>, client: Client) -> Self {
        Self {
            catalog: catalog.into(),
            schema: schema.into(),
            dbname: String::default(),
            timezone: String::default(),
            client,
            ctx: FlightContext::default(),
        }
    }

    /// Create database service client using dbname.
    ///
    /// This API is designed for external usage. `dbname` is:
    ///
    /// - the name of database when using GreptimeDB standalone or cluster
    /// - the name provided by GreptimeCloud or other multi-tenant GreptimeDB
    ///   environment
    pub fn new_with_dbname(dbname: impl Into<String>, client: Client) -> Self {
        Self {
            catalog: String::default(),
            schema: String::default(),
            timezone: String::default(),
            dbname: dbname.into(),
            client,
            ctx: FlightContext::default(),
        }
    }

    /// Set the catalog for the database client.
    pub fn set_catalog(&mut self, catalog: impl Into<String>) {
        self.catalog = catalog.into();
    }

    fn catalog_or_default(&self) -> &str {
        if self.catalog.is_empty() {
            DEFAULT_CATALOG_NAME
        } else {
            &self.catalog
        }
    }

    /// Set the schema for the database client.
    pub fn set_schema(&mut self, schema: impl Into<String>) {
        self.schema = schema.into();
    }

    fn schema_or_default(&self) -> &str {
        if self.schema.is_empty() {
            DEFAULT_SCHEMA_NAME
        } else {
            &self.schema
        }
    }

    /// Set the timezone for the database client.
    pub fn set_timezone(&mut self, timezone: impl Into<String>) {
        self.timezone = timezone.into();
    }

    /// Set the auth scheme for the database client.
    pub fn set_auth(&mut self, auth: AuthScheme) {
        self.ctx.auth_header = Some(AuthHeader {
            auth_scheme: Some(auth),
        });
    }

    /// Creates a builder for a single Flight DoGet request.
    pub fn flight_request(&self) -> DatabaseFlightRequest<'_> {
        DatabaseFlightRequest {
            database: self,
            options: FlightRequestOptions::default(),
        }
    }

    /// Make an InsertRequests request to the database.
    pub async fn insert(&self, requests: InsertRequests) -> Result<u32> {
        self.handle(Request::Inserts(requests)).await
    }

    /// Make an InsertRequests request to the database with hints.
    pub async fn insert_with_hints(
        &self,
        requests: InsertRequests,
        hints: &[(&str, &str)],
    ) -> Result<u32> {
        let mut client = make_database_client(&self.client)?;
        let request = self.to_rpc_request(Request::Inserts(requests));

        let mut request = tonic::Request::new(request);
        let metadata = request.metadata_mut();
        Self::put_hints(metadata, hints)?;

        let response = client
            .inner
            .handle(request)
            .await
            .inspect_err(client.inspect_err("insert_with_hints"))?
            .into_inner();
        from_grpc_response(response)
    }

    /// Make a RowInsertRequests request to the database.
    pub async fn row_inserts(&self, requests: RowInsertRequests) -> Result<u32> {
        self.handle(Request::RowInserts(requests)).await
    }

    /// Make a RowInsertRequests request to the database with hints.
    pub async fn row_inserts_with_hints(
        &self,
        requests: RowInsertRequests,
        hints: &[(&str, &str)],
    ) -> Result<u32> {
        let mut client = make_database_client(&self.client)?;
        let request = self.to_rpc_request(Request::RowInserts(requests));

        let mut request = tonic::Request::new(request);
        let metadata = request.metadata_mut();
        Self::put_hints(metadata, hints)?;

        let response = client
            .inner
            .handle(request)
            .await
            .inspect_err(client.inspect_err("row_inserts_with_hints"))?
            .into_inner();
        from_grpc_response(response)
    }

    fn put_hints(metadata: &mut MetadataMap, hints: &[(&str, &str)]) -> Result<()> {
        let Some(value) = Self::encode_hints(hints) else {
            return Ok(());
        };

        Self::put_metadata_value(metadata, HINTS_METADATA_KEY, value)
    }

    fn encode_hints(hints: &[(&str, &str)]) -> Option<String> {
        hints
            .iter()
            .map(|(k, v)| format!("{}={}", k, v))
            .reduce(|a, b| format!("{},{}", a, b))
    }

    fn encode_flow_extensions(flow_extensions: &[(&str, &str)]) -> Option<String> {
        (!flow_extensions.is_empty()).then(|| {
            serde_json::to_string(&flow_extensions.to_vec())
                .expect("flow extension pairs should serialize")
        })
    }

    fn encode_snapshot_seqs(snapshot_seqs: &HashMap<u64, u64>) -> Option<String> {
        (!snapshot_seqs.is_empty()).then(|| {
            serde_json::to_string(snapshot_seqs).expect("snapshot sequence map should serialize")
        })
    }

    fn put_metadata_value(
        metadata: &mut MetadataMap,
        key: &'static str,
        value: String,
    ) -> Result<()> {
        let key = AsciiMetadataKey::from_static(key);
        let value = AsciiMetadataValue::from_str(&value).context(InvalidTonicMetadataValueSnafu)?;
        metadata.insert(key, value);
        Ok(())
    }

    /// Make a request to the database.
    pub async fn handle(&self, request: Request) -> Result<u32> {
        let mut client = make_database_client(&self.client)?;
        let request = self.to_rpc_request(request);
        let response = client
            .inner
            .handle(request)
            .await
            .inspect_err(client.inspect_err("handle"))?
            .into_inner();
        from_grpc_response(response)
    }

    /// Retry if connection fails, max_retries is the max number of retries, so the total wait time
    /// is `max_retries * GRPC_CONN_TIMEOUT`
    pub async fn handle_with_retry(
        &self,
        request: Request,
        max_retries: u32,
        hints: &[(&str, &str)],
    ) -> Result<u32> {
        let mut client = make_database_client(&self.client)?;
        let mut retries = 0;

        let request = self.to_rpc_request(request);

        loop {
            let mut tonic_request = tonic::Request::new(request.clone());
            let metadata = tonic_request.metadata_mut();
            Self::put_hints(metadata, hints)?;
            let raw_response = client
                .inner
                .handle(tonic_request)
                .await
                .inspect_err(client.inspect_err("handle"));
            match (raw_response, retries < max_retries) {
                (Ok(resp), _) => return from_grpc_response(resp.into_inner()),
                (Err(err), true) => {
                    // determine if the error is retryable
                    if is_grpc_retryable(&err) {
                        // retry
                        retries += 1;
                        warn!("Retrying {} times with error = {:?}", retries, err);
                        continue;
                    } else {
                        error!(
                            err; "Failed to send request to grpc handle, retries = {}, not retryable error, aborting",
                            retries
                        );
                        return Err(err.into());
                    }
                }
                (Err(err), false) => {
                    error!(
                        err; "Failed to send request to grpc handle after {} retries",
                        retries,
                    );
                    return Err(err.into());
                }
            }
        }
    }

    #[inline]
    fn to_rpc_request(&self, request: Request) -> GreptimeRequest {
        GreptimeRequest {
            header: Some(RequestHeader {
                catalog: self.catalog.clone(),
                schema: self.schema.clone(),
                authorization: self.ctx.auth_header.clone(),
                dbname: self.dbname.clone(),
                timezone: self.timezone.clone(),
                // TODO(Taylor-lagrange): add client grpc tracing
                tracing_context: W3cTrace::new(),
            }),
            request: Some(request),
        }
    }

    /// Executes a SQL query without any hints.
    pub async fn sql<S>(&self, sql: S) -> Result<Output>
    where
        S: AsRef<str>,
    {
        self.flight_request().sql(sql).await
    }

    /// Executes a SQL query with optional hints for query optimization.
    pub async fn sql_with_hint<S>(&self, sql: S, hints: &[(&str, &str)]) -> Result<Output>
    where
        S: AsRef<str>,
    {
        self.flight_request().with_hints(hints).sql(sql).await
    }

    /// Executes a SQL query and returns the output with terminal metrics.
    ///
    /// For stream outputs, callers must consume the stream before reading final
    /// terminal metrics from [`OutputWithMetrics::metrics`]. For affected-row
    /// outputs without inline metrics, call [`OutputMetrics::wait_ready`] when
    /// compatibility trailing metrics are required.
    pub async fn sql_with_terminal_metrics<S>(
        &self,
        sql: S,
        hints: &[(&str, &str)],
    ) -> Result<OutputWithMetrics>
    where
        S: AsRef<str>,
    {
        self.flight_request()
            .with_hints(hints)
            .sql_with_terminal_metrics(sql)
            .await
    }

    /// Executes a logical plan directly without SQL parsing.
    pub async fn logical_plan(&self, logical_plan: Vec<u8>) -> Result<Output> {
        self.flight_request().logical_plan(logical_plan).await
    }

    /// Creates a new table using the provided table expression.
    #[cfg(feature = "testing")]
    pub async fn create(&self, expr: CreateTableExpr) -> Result<Output> {
        self.flight_request().create(expr).await
    }

    /// Alters an existing table using the provided alter expression.
    #[cfg(feature = "testing")]
    pub async fn alter(&self, expr: AlterTableExpr) -> Result<Output> {
        self.flight_request().alter(expr).await
    }

    async fn do_get(
        &self,
        request: Request,
        options: FlightRequestOptions,
    ) -> Result<OutputWithMetrics> {
        let request = self.to_rpc_request(request);
        let request = Ticket {
            ticket: request.encode_to_vec().into(),
        };

        let mut request = tonic::Request::new(request);
        options.apply_to(&mut request)?;

        let mut client = self.client.make_flight_client(false, false)?;
        let remote_addr = client.addr().to_string();

        let response = client.mut_inner().do_get(request).await.or_else(|e| {
            let tonic_code = e.code();
            let e: Error = e.into();
            error!(
                "Failed to do Flight get, addr: {}, code: {}, source: {:?}",
                client.addr(),
                tonic_code,
                e
            );
            Err(BoxedError::new(e)).with_context(|_| FlightGetSnafu {
                addr: remote_addr.clone(),
                tonic_code,
            })
        })?;

        let flight_data_stream = response.into_inner();
        let mut decoder = FlightDecoder::default();

        let flight_message_stream = flight_data_stream.filter_map(move |flight_data| {
            future::ready(decode_flight_data(&mut decoder, flight_data))
        });

        output_from_flight_message_stream(remote_addr, flight_message_stream).await
    }

    /// Ingest a stream of [RecordBatch]es that belong to a table, using Arrow Flight's "`DoPut`"
    /// method. The return value is also a stream, produces [DoPutResponse]s.
    pub async fn do_put(&self, stream: FlightDataStream) -> Result<DoPutResponseStream> {
        self.do_put_with_hints(stream, &[]).await
    }

    /// Ingest a stream of [RecordBatch]es using Arrow Flight's `DoPut` with request hints.
    pub async fn do_put_with_hints(
        &self,
        stream: FlightDataStream,
        hints: &[(&str, &str)],
    ) -> Result<DoPutResponseStream> {
        let mut request = tonic::Request::new(stream);
        Self::put_hints(request.metadata_mut(), hints)?;

        if let Some(AuthHeader {
            auth_scheme: Some(AuthScheme::Basic(Basic { username, password })),
        }) = &self.ctx.auth_header
        {
            let encoded = BASE64_STANDARD.encode(format!("{username}:{password}"));
            let value = MetadataValue::from_str(&format!("Basic {encoded}"))
                .context(InvalidTonicMetadataValueSnafu)?;
            request.metadata_mut().insert("x-greptime-auth", value);
        }

        let db_to_put = if !self.dbname.is_empty() {
            &self.dbname
        } else {
            &build_db_string(self.catalog_or_default(), self.schema_or_default())
        };
        request.metadata_mut().insert(
            "x-greptime-db-name",
            MetadataValue::from_str(db_to_put).context(InvalidTonicMetadataValueSnafu)?,
        );

        let mut client = self.client.make_control_flight_client(false, false)?;
        let response = client.mut_inner().do_put(request).await?;
        let response = response
            .into_inner()
            .map_err(Into::into)
            .and_then(|x| future::ready(DoPutResponse::try_from(x).context(ConvertFlightDataSnafu)))
            .boxed();
        Ok(response)
    }
}

impl<'a> DatabaseFlightRequest<'a> {
    /// Adds query optimization hints to this Flight request.
    pub fn with_hints(mut self, hints: &[(&str, &str)]) -> Self {
        self.options.hints = Database::encode_hints(hints);
        self
    }

    /// Adds Flow extensions to this Flight request.
    pub fn with_flow_extensions(mut self, flow_extensions: &[(&str, &str)]) -> Self {
        self.options.flow_extensions = Database::encode_flow_extensions(flow_extensions);
        self
    }

    /// Adds snapshot sequence fences to this Flight request.
    pub fn with_snapshot_seqs(mut self, snapshot_seqs: &HashMap<u64, u64>) -> Self {
        self.options.snapshot_seqs = Database::encode_snapshot_seqs(snapshot_seqs);
        self
    }

    /// Sets a timeout for this Flight request only.
    pub fn with_timeout(mut self, timeout: Duration) -> Self {
        self.options.timeout = Some(timeout);
        self
    }

    /// Executes a SQL query.
    pub async fn sql<S>(self, sql: S) -> Result<Output>
    where
        S: AsRef<str>,
    {
        let request = Request::Query(QueryRequest {
            query: Some(Query::Sql(sql.as_ref().to_string())),
        });
        self.do_get(request)
            .await
            .map(OutputWithMetrics::into_output)
    }

    /// Executes a SQL query and returns terminal metrics.
    pub async fn sql_with_terminal_metrics<S>(self, sql: S) -> Result<OutputWithMetrics>
    where
        S: AsRef<str>,
    {
        self.query_with_terminal_metrics(QueryRequest {
            query: Some(Query::Sql(sql.as_ref().to_string())),
        })
        .await
    }

    /// Executes a logical plan directly without SQL parsing.
    pub async fn logical_plan(self, logical_plan: Vec<u8>) -> Result<Output> {
        self.query_with_terminal_metrics(QueryRequest {
            query: Some(Query::LogicalPlan(logical_plan)),
        })
        .await
        .map(OutputWithMetrics::into_output)
    }

    /// Executes a query and returns terminal metrics.
    pub async fn query_with_terminal_metrics(
        self,
        request: QueryRequest,
    ) -> Result<OutputWithMetrics> {
        self.do_get(Request::Query(request)).await
    }

    /// Creates a new table using the provided table expression.
    #[cfg(feature = "testing")]
    pub async fn create(self, expr: CreateTableExpr) -> Result<Output> {
        self.do_get(Request::Ddl(DdlRequest {
            expr: Some(DdlExpr::CreateTable(expr)),
        }))
        .await
        .map(OutputWithMetrics::into_output)
    }

    /// Alters an existing table using the provided alter expression.
    #[cfg(feature = "testing")]
    pub async fn alter(self, expr: AlterTableExpr) -> Result<Output> {
        self.do_get(Request::Ddl(DdlRequest {
            expr: Some(DdlExpr::AlterTable(expr)),
        }))
        .await
        .map(OutputWithMetrics::into_output)
    }

    async fn do_get(self, request: Request) -> Result<OutputWithMetrics> {
        let Self { database, options } = self;
        database.do_get(request, options).await
    }
}

/// by grpc standard, only `Unavailable` is retryable, see: https://github.com/grpc/grpc/blob/master/doc/statuscodes.md#status-codes-and-their-use-in-grpc
pub fn is_grpc_retryable(err: &tonic::Status) -> bool {
    matches!(err.code(), tonic::Code::Unavailable)
}

#[derive(Default, Debug, Clone)]
struct FlightContext {
    auth_header: Option<AuthHeader>,
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;
    use std::task::{Context, Poll};

    use api::v1::auth_header::AuthScheme;
    use api::v1::{AuthHeader, Basic};
    use common_error::ext::{ErrorExt, RetryHint};
    use common_error::status_code::StatusCode;
    use common_error::{GREPTIME_DB_HEADER_ERROR_CODE, GREPTIME_DB_HEADER_ERROR_RETRY_HINT};
    use common_query::OutputData;
    use common_recordbatch::{OrderOption, RecordBatch, RecordBatchStream};
    use datatypes::prelude::{ConcreteDataType, VectorRef};
    use datatypes::schema::{ColumnSchema, Schema};
    use datatypes::vectors::Int32Vector;
    use futures_util::StreamExt;
    use tokio::sync::oneshot;
    use tonic::codegen::http::{HeaderMap, HeaderValue};
    use tonic::metadata::MetadataMap;
    use tonic::{Code, Status};

    use super::*;
    use crate::error::TonicSnafu;

    struct MockMetricsStream {
        schema: datatypes::schema::SchemaRef,
        batch: Option<RecordBatch>,
        metrics: RecordBatchMetrics,
        terminal_metrics_only: bool,
    }

    impl Stream for MockMetricsStream {
        type Item = common_recordbatch::error::Result<RecordBatch>;

        fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
            Poll::Ready(self.batch.take().map(Ok))
        }
    }

    impl RecordBatchStream for MockMetricsStream {
        fn name(&self) -> &str {
            "MockMetricsStream"
        }

        fn schema(&self) -> datatypes::schema::SchemaRef {
            self.schema.clone()
        }

        fn output_ordering(&self) -> Option<&[OrderOption]> {
            None
        }

        fn metrics(&self) -> Option<RecordBatchMetrics> {
            if self.terminal_metrics_only && self.batch.is_some() {
                return None;
            }
            Some(self.metrics.clone())
        }
    }

    fn terminal_metrics_json() -> String {
        terminal_metrics_json_with_seq(42)
    }

    fn terminal_metrics_json_with_seq(seq: u64) -> String {
        serde_json::to_string(&RecordBatchMetrics {
            region_watermarks: vec![common_recordbatch::adapter::RegionWatermarkEntry {
                region_id: 7,
                watermark: Some(seq),
            }],
            ..Default::default()
        })
        .unwrap()
    }

    #[test]
    fn test_put_flow_extensions_preserves_comma_bearing_values() {
        let mut metadata = MetadataMap::new();
        Database::put_metadata_value(
            &mut metadata,
            FLOW_EXTENSIONS_METADATA_KEY,
            Database::encode_flow_extensions(&[
                ("flow.return_region_seq", "true"),
                ("flow.incremental_after_seqs", r#"{"1":10,"2":20}"#),
            ])
            .unwrap(),
        )
        .unwrap();

        let value = metadata
            .get(FLOW_EXTENSIONS_METADATA_KEY)
            .unwrap()
            .to_str()
            .unwrap();
        let decoded: Vec<(String, String)> = serde_json::from_str(value).unwrap();
        assert_eq!(
            decoded,
            vec![
                ("flow.return_region_seq".to_string(), "true".to_string()),
                (
                    "flow.incremental_after_seqs".to_string(),
                    r#"{"1":10,"2":20}"#.to_string()
                ),
            ]
        );
    }

    #[test]
    fn test_put_snapshot_seqs_preserves_u64_precision() {
        let mut metadata = MetadataMap::new();
        let snapshot_seqs = std::collections::HashMap::from([
            (u64::MAX, u64::MAX - 1),
            (9_007_199_254_740_993_u64, 9_007_199_254_740_995_u64),
        ]);

        Database::put_metadata_value(
            &mut metadata,
            SNAPSHOT_SEQS_METADATA_KEY,
            Database::encode_snapshot_seqs(&snapshot_seqs).unwrap(),
        )
        .unwrap();

        let value = metadata
            .get(SNAPSHOT_SEQS_METADATA_KEY)
            .unwrap()
            .to_str()
            .unwrap();
        let decoded: std::collections::HashMap<u64, u64> = serde_json::from_str(value).unwrap();
        assert_eq!(decoded, snapshot_seqs);
    }

    #[test]
    fn test_flight_request_builder_applies_request_options() {
        let database = Database::new("greptime", "public", Client::default());
        let snapshot_seqs = HashMap::from([(42, 99)]);
        let request = database
            .flight_request()
            .with_hints(&[("query_parallelism", "1")])
            .with_flow_extensions(&[("flow.return_region_seq", "true")])
            .with_snapshot_seqs(&snapshot_seqs)
            .with_timeout(Duration::from_millis(50));
        let mut tonic_request = tonic::Request::new(());

        request.options.apply_to(&mut tonic_request).unwrap();

        let metadata = tonic_request.metadata();
        assert_eq!(
            metadata.get(HINTS_METADATA_KEY).unwrap(),
            "query_parallelism=1"
        );
        assert_eq!(
            serde_json::from_str::<Vec<(String, String)>>(
                metadata
                    .get(FLOW_EXTENSIONS_METADATA_KEY)
                    .unwrap()
                    .to_str()
                    .unwrap(),
            )
            .unwrap(),
            vec![("flow.return_region_seq".to_string(), "true".to_string())]
        );
        assert_eq!(
            serde_json::from_str::<HashMap<u64, u64>>(
                metadata
                    .get(SNAPSHOT_SEQS_METADATA_KEY)
                    .unwrap()
                    .to_str()
                    .unwrap(),
            )
            .unwrap(),
            snapshot_seqs
        );
        assert!(metadata.get("grpc-timeout").is_some());
    }

    #[test]
    fn test_flight_ctx() {
        let mut ctx = FlightContext::default();
        assert!(ctx.auth_header.is_none());

        let basic = AuthScheme::Basic(Basic {
            username: "u".to_string(),
            password: "p".to_string(),
        });

        ctx.auth_header = Some(AuthHeader {
            auth_scheme: Some(basic),
        });

        assert!(matches!(
            ctx.auth_header,
            Some(AuthHeader {
                auth_scheme: Some(AuthScheme::Basic(_)),
            })
        ));
    }

    #[test]
    fn test_from_tonic_status() {
        let expected = TonicSnafu {
            code: StatusCode::Internal,
            msg: "blabla".to_string(),
            tonic_code: Code::Internal,
            retry_hint: RetryHint::NonRetryable,
        }
        .build();

        let status = Status::new(Code::Internal, "blabla");
        let actual: Error = status.into();

        assert_eq!(expected.to_string(), actual.to_string());
        assert_eq!(expected.retry_hint(), actual.retry_hint());
        assert_eq!(expected.should_retry(), actual.should_retry());
    }

    #[test]
    fn test_flight_stream_error_preserves_addr_and_message() {
        let error = flight_stream_error(
            "127.0.0.1:4001",
            Status::out_of_range("message length too large").into(),
        );

        assert!(matches!(
            &error,
            Error::FlightStream {
                addr,
                tonic_code: Code::OutOfRange,
                message,
                ..
            } if addr == "127.0.0.1:4001" && message == "message length too large"
        ));
        assert_eq!(
            "Failed to receive Flight data from 127.0.0.1:4001, code: Operation was attempted past the valid range: message length too large",
            error.to_string(),
        );
    }

    #[test]
    fn test_from_tonic_status_with_retry_hint() {
        let mut headers = HeaderMap::new();
        headers.insert(
            GREPTIME_DB_HEADER_ERROR_CODE,
            HeaderValue::from(StatusCode::Internal as u32),
        );
        headers.insert(
            GREPTIME_DB_HEADER_ERROR_RETRY_HINT,
            HeaderValue::from_static(RetryHint::Retryable.as_str()),
        );
        let status =
            Status::with_metadata(Code::Internal, "blabla", MetadataMap::from_headers(headers));

        let actual: Error = status.into();

        assert_eq!(actual.retry_hint(), RetryHint::Retryable);
        assert!(actual.should_retry());
    }

    #[test]
    fn test_from_tonic_status_fallback() {
        let mut headers = HeaderMap::new();
        headers.insert(
            GREPTIME_DB_HEADER_ERROR_CODE,
            HeaderValue::from(StatusCode::InvalidArguments as u32),
        );
        let status =
            Status::with_metadata(Code::Internal, "blabla", MetadataMap::from_headers(headers));

        let actual: Error = status.into();

        assert_eq!(actual.retry_hint(), RetryHint::NonRetryable);
        assert!(!actual.should_retry());
    }

    #[test]
    fn test_should_retry_preserves_transport_retry() {
        let status = Status::new(Code::Unavailable, "blabla");
        let actual: Error = status.into();

        assert!(actual.should_retry());
    }

    #[tokio::test]
    async fn test_query_with_terminal_metrics_tracks_terminal_only_metrics() {
        let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
            "v",
            ConcreteDataType::int32_datatype(),
            false,
        )]));
        let batch = RecordBatch::new(
            schema.clone(),
            vec![Arc::new(Int32Vector::from_slice([1, 2])) as VectorRef],
        )
        .unwrap();
        let output = Output::new_with_stream(Box::pin(MockMetricsStream {
            schema,
            batch: Some(batch),
            metrics: RecordBatchMetrics {
                region_watermarks: vec![common_recordbatch::adapter::RegionWatermarkEntry {
                    region_id: 7,
                    watermark: Some(42),
                }],
                ..Default::default()
            },
            terminal_metrics_only: true,
        }));

        let result = OutputWithMetrics::from_output(output);
        let terminal_metrics = result.metrics.clone();
        assert!(!terminal_metrics.is_ready());
        assert!(terminal_metrics.get().is_none());

        let OutputData::Stream(mut stream) = result.output.data else {
            panic!("expected stream output");
        };
        while stream.next().await.is_some() {}

        assert!(terminal_metrics.is_ready());
        assert_eq!(
            terminal_metrics.participating_regions(),
            Some(std::collections::BTreeSet::from([7_u64]))
        );
        assert_eq!(
            terminal_metrics.region_watermark_map(),
            Some(std::collections::HashMap::from([(7_u64, 42_u64)]))
        );
    }

    #[tokio::test]
    async fn test_affected_rows_inline_metrics_are_parsed() {
        let output = output_from_flight_message_stream(
            "test-peer".to_string(),
            futures_util::stream::iter(vec![Ok(FlightMessage::AffectedRows {
                rows: 3,
                metrics: Some(terminal_metrics_json()),
            })] as Vec<Result<FlightMessage>>),
        )
        .await
        .unwrap();

        assert!(matches!(output.output.data, OutputData::AffectedRows(3)));
        assert!(output.metrics.is_ready());
        assert_eq!(
            output.metrics.region_watermark_map(),
            Some(std::collections::HashMap::from([(7, 42)]))
        );
    }

    #[tokio::test]
    async fn test_affected_rows_inline_metrics_do_not_poll_trailer() {
        let metrics_json = terminal_metrics_json();
        let output = output_from_flight_message_stream(
            "test-peer".to_string(),
            futures_util::stream::iter(vec![
                Ok(FlightMessage::AffectedRows {
                    rows: 3,
                    metrics: Some(metrics_json),
                }),
                Ok(FlightMessage::Metrics(terminal_metrics_json_with_seq(99))),
            ] as Vec<Result<FlightMessage>>),
        )
        .await
        .unwrap();

        assert!(output.metrics.is_ready());
        assert_eq!(
            output.metrics.region_watermark_map(),
            Some(std::collections::HashMap::from([(7, 42)]))
        );
    }

    #[tokio::test]
    async fn test_affected_rows_without_inline_metrics_becomes_ready_after_trailer() {
        // Hold the trailer until the not-ready state has been observed.
        let (trailer_tx, trailer_rx) = oneshot::channel();
        let trailer = futures_util::stream::once(trailer_rx).map(|ready| {
            ready.unwrap();
            Ok(FlightMessage::Metrics(terminal_metrics_json()))
        });
        let output = output_from_flight_message_stream(
            "test-peer".to_string(),
            futures_util::stream::iter(vec![Ok(FlightMessage::AffectedRows {
                rows: 3,
                metrics: None,
            })])
            .chain(trailer),
        )
        .await
        .unwrap();

        assert!(!output.metrics.is_ready());
        trailer_tx.send(()).unwrap();
        tokio::time::timeout(Duration::from_secs(1), output.metrics.wait_ready())
            .await
            .expect("terminal metrics must become ready once the trailer arrives");
        assert!(output.metrics.completion_error().is_none());
        assert_eq!(
            output.metrics.region_watermark_map(),
            Some(std::collections::HashMap::from([(7, 42)]))
        );
    }

    #[tokio::test]
    async fn test_affected_rows_malformed_trailer_sets_completion_error_and_ready() {
        let output = output_from_flight_message_stream(
            "test-peer".to_string(),
            futures_util::stream::iter(vec![
                Ok(FlightMessage::AffectedRows {
                    rows: 3,
                    metrics: None,
                }),
                Ok(FlightMessage::Metrics("{not-json}".to_string())),
            ] as Vec<Result<FlightMessage>>),
        )
        .await
        .unwrap();

        output.metrics.wait_ready().await;
        let error = output.metrics.completion_error().unwrap();
        assert!(error.contains("Invalid terminal metrics message"));
        assert!(output.metrics.is_ready());
    }

    #[tokio::test]
    async fn test_affected_rows_transport_trailer_error_sets_completion_error_and_ready() {
        let output = output_from_flight_message_stream(
            "test-peer".to_string(),
            futures_util::stream::iter(vec![
                Ok(FlightMessage::AffectedRows {
                    rows: 3,
                    metrics: None,
                }),
                Err(Status::unavailable("trailer read failed").into()),
            ] as Vec<Result<FlightMessage>>),
        )
        .await
        .unwrap();

        output.metrics.wait_ready().await;
        let error = output.metrics.completion_error().unwrap();
        assert!(error.contains("trailer read failed"));
        assert!(output.metrics.is_ready());
    }

    #[tokio::test]
    async fn test_affected_rows_compatibility_reader_is_cancelled_after_second_poll_begins() {
        struct DropProbe {
            first: Option<Result<FlightMessage>>,
            polled: Option<oneshot::Sender<()>>,
            dropped: Option<oneshot::Sender<()>>,
        }

        impl Stream for DropProbe {
            type Item = Result<FlightMessage>;

            fn poll_next(
                mut self: Pin<&mut Self>,
                _cx: &mut Context<'_>,
            ) -> Poll<Option<Self::Item>> {
                if let Some(message) = self.first.take() {
                    return Poll::Ready(Some(message));
                }
                if let Some(polled) = self.polled.take() {
                    let _ = polled.send(());
                }
                Poll::Pending
            }
        }

        impl Drop for DropProbe {
            fn drop(&mut self) {
                if let Some(dropped) = self.dropped.take() {
                    let _ = dropped.send(());
                }
            }
        }

        let (polled_tx, polled_rx) = oneshot::channel();
        let (dropped_tx, dropped_rx) = oneshot::channel();
        let output = output_from_flight_message_stream(
            "test-peer".to_string(),
            DropProbe {
                first: Some(Ok(FlightMessage::AffectedRows {
                    rows: 3,
                    metrics: None,
                })),
                polled: Some(polled_tx),
                dropped: Some(dropped_tx),
            },
        )
        .await
        .unwrap();
        polled_rx.await.unwrap();
        drop(output);
        // Require cancellation before the five-second compatibility timeout.
        tokio::time::timeout(Duration::from_secs(1), dropped_rx)
            .await
            .expect("the compatibility reader must be dropped once the output is dropped")
            .unwrap();
    }

    #[tokio::test]
    async fn test_schema_record_batch_yields_before_pending_message_stream() {
        let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
            "v",
            ConcreteDataType::int32_datatype(),
            false,
        )]));
        let batch = RecordBatch::new(
            schema.clone(),
            vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef],
        )
        .unwrap();
        let messages = futures_util::stream::iter(vec![
            Ok(FlightMessage::Schema(schema.arrow_schema().clone())),
            Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())),
        ] as Vec<Result<FlightMessage>>)
        .chain(futures_util::stream::pending());
        let output = output_from_flight_message_stream("test-peer".to_string(), messages)
            .await
            .unwrap();
        let OutputData::Stream(mut stream) = output.output.data else {
            panic!("expected stream output");
        };

        let batch = tokio::time::timeout(Duration::from_secs(1), stream.next())
            .await
            .unwrap()
            .unwrap()
            .unwrap();
        assert_eq!(batch.num_rows(), 1);
        assert!(!output.metrics.is_ready());
    }

    #[tokio::test]
    async fn test_invalid_terminal_metrics_after_record_batch_yields_batch_then_error() {
        let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
            "v",
            ConcreteDataType::int32_datatype(),
            false,
        )]));
        let batch = RecordBatch::new(
            schema.clone(),
            vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef],
        )
        .unwrap();
        let output = output_from_flight_message_stream(
            "test-peer".to_string(),
            futures_util::stream::iter(vec![
                Ok(FlightMessage::Schema(schema.arrow_schema().clone())),
                Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())),
                Ok(FlightMessage::Metrics("{not-json}".to_string())),
            ] as Vec<Result<FlightMessage>>),
        )
        .await
        .unwrap();
        let terminal_metrics = output.metrics.clone();
        let OutputData::Stream(mut record_batch_stream) = output.output.data else {
            panic!("expected stream output");
        };

        let batch = record_batch_stream.next().await.unwrap().unwrap();
        assert_eq!(batch.num_rows(), 1);

        let err = record_batch_stream.next().await.unwrap().unwrap_err();
        assert_eq!("External error", err.to_string());
        assert!(
            format!("{err:?}").contains("Invalid terminal metrics message"),
            "unexpected error: {err:?}"
        );
        assert!(record_batch_stream.next().await.is_none());
        assert!(terminal_metrics.is_ready());
        assert!(terminal_metrics.get().is_none());
    }

    #[tokio::test]
    async fn test_record_batch_stream_continues_after_partial_metrics() {
        let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
            "v",
            ConcreteDataType::int32_datatype(),
            false,
        )]));
        let first_batch = RecordBatch::new(
            schema.clone(),
            vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef],
        )
        .unwrap();
        let second_batch = RecordBatch::new(
            schema.clone(),
            vec![Arc::new(Int32Vector::from_slice([2])) as VectorRef],
        )
        .unwrap();
        let output = output_from_flight_message_stream(
            "test-peer".to_string(),
            futures_util::stream::iter(vec![
                Ok(FlightMessage::Schema(schema.arrow_schema().clone())),
                Ok(FlightMessage::RecordBatch(
                    first_batch.into_df_record_batch(),
                )),
                Ok(FlightMessage::Metrics(terminal_metrics_json_with_seq(1))),
                Ok(FlightMessage::RecordBatch(
                    second_batch.into_df_record_batch(),
                )),
                Ok(FlightMessage::Metrics(terminal_metrics_json_with_seq(2))),
            ] as Vec<Result<FlightMessage>>),
        )
        .await
        .unwrap();
        let terminal_metrics = output.metrics.clone();
        let OutputData::Stream(mut record_batch_stream) = output.output.data else {
            panic!("expected stream output");
        };

        let first_batch = record_batch_stream.next().await.unwrap().unwrap();
        assert_eq!(first_batch.num_rows(), 1);
        let second_batch = record_batch_stream.next().await.unwrap().unwrap();
        assert_eq!(second_batch.num_rows(), 1);
        assert!(record_batch_stream.next().await.is_none());

        assert!(terminal_metrics.is_ready());
        assert_eq!(
            terminal_metrics.region_watermark_map(),
            Some(std::collections::HashMap::from([(7, 2)]))
        );
    }

    #[test]
    fn test_output_metrics_distinguishes_empty_region_watermarks_from_absence() {
        let metrics = OutputMetrics::default();
        metrics.update(Some(RecordBatchMetrics::default()));

        assert_eq!(
            metrics.participating_regions(),
            Some(std::collections::BTreeSet::new())
        );
        assert_eq!(
            metrics.region_watermark_map(),
            Some(std::collections::HashMap::new())
        );
    }
}