servers/grpc/

flight.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.

mod stream;

use std::collections::HashMap;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};

use api::v1::GreptimeRequest;
use arrow_flight::flight_service_server::FlightService;
use arrow_flight::{
    Action, ActionType, Criteria, Empty, FlightData, FlightDescriptor, FlightInfo,
    HandshakeRequest, HandshakeResponse, PollInfo, PutResult, SchemaResult, Ticket,
};
use async_trait::async_trait;
use bytes::{self, Bytes};
use common_error::ext::ErrorExt;
use common_grpc::flight::do_put::{DoPutMetadata, DoPutResponse};
use common_grpc::flight::{
    FLOW_EXTENSIONS_METADATA_KEY, FlightDecoder, FlightMessage, SNAPSHOT_SEQS_METADATA_KEY,
};
use common_memory_manager::MemoryGuard;
use common_query::{Output, OutputData};
use common_recordbatch::DfRecordBatch;
use common_telemetry::debug;
use common_telemetry::tracing::info_span;
use common_telemetry::tracing_context::{FutureExt, TracingContext};
use datatypes::arrow::datatypes::SchemaRef;
use futures::{Stream, future, ready};
use futures_util::{StreamExt, TryStreamExt};
use prost::Message;
use query::metrics::terminal_recordbatch_metrics_from_plan_if_requested;
use query::options::FlowQueryExtensions;
use session::context::{Channel, QueryContextRef};
use snafu::{IntoError, OptionExt, ResultExt, ensure};
use table::table_name::TableName;
use tokio::sync::mpsc;
use tokio_stream::wrappers::ReceiverStream;
use tonic::{Request, Response, Status, Streaming};

use crate::error::{InvalidParameterSnafu, InvalidQuerySnafu, Result, ToJsonSnafu};
pub use crate::grpc::flight::stream::{
    FlightRecordBatchSource, FlightRecordBatchStream, FlightRecordBatchStreamInput,
};
use crate::grpc::greptime_handler::{
    GreptimeRequestHandler, create_query_context, get_request_type,
};
use crate::grpc::{FlightCompression, TonicResult, context_auth};
use crate::request_memory_limiter::ServerMemoryLimiter;
use crate::request_memory_metrics::RequestMemoryMetrics;
use crate::{error, hint_headers};

pub type TonicStream<T> = Pin<Box<dyn Stream<Item = TonicResult<T>> + Send + 'static>>;

/// A subset of [FlightService]
#[async_trait]
pub trait FlightCraft: Send + Sync + 'static {
    async fn do_get(
        &self,
        request: Request<Ticket>,
    ) -> TonicResult<Response<TonicStream<FlightData>>>;

    async fn do_put(
        &self,
        request: Request<Streaming<FlightData>>,
    ) -> TonicResult<Response<TonicStream<PutResult>>> {
        let _ = request;
        Err(Status::unimplemented("Not yet implemented"))
    }
}

pub type FlightCraftRef = Arc<dyn FlightCraft>;

pub struct FlightCraftWrapper<T: FlightCraft>(pub T);

impl<T: FlightCraft> From<T> for FlightCraftWrapper<T> {
    fn from(t: T) -> Self {
        Self(t)
    }
}

#[async_trait]
impl FlightCraft for FlightCraftRef {
    async fn do_get(
        &self,
        request: Request<Ticket>,
    ) -> TonicResult<Response<TonicStream<FlightData>>> {
        (**self).do_get(request).await
    }

    async fn do_put(
        &self,
        request: Request<Streaming<FlightData>>,
    ) -> TonicResult<Response<TonicStream<PutResult>>> {
        self.as_ref().do_put(request).await
    }
}

#[async_trait]
impl<T: FlightCraft> FlightService for FlightCraftWrapper<T> {
    type HandshakeStream = TonicStream<HandshakeResponse>;

    async fn handshake(
        &self,
        _: Request<Streaming<HandshakeRequest>>,
    ) -> TonicResult<Response<Self::HandshakeStream>> {
        Err(Status::unimplemented("Not yet implemented"))
    }

    type ListFlightsStream = TonicStream<FlightInfo>;

    async fn list_flights(
        &self,
        _: Request<Criteria>,
    ) -> TonicResult<Response<Self::ListFlightsStream>> {
        Err(Status::unimplemented("Not yet implemented"))
    }

    async fn get_flight_info(
        &self,
        _: Request<FlightDescriptor>,
    ) -> TonicResult<Response<FlightInfo>> {
        Err(Status::unimplemented("Not yet implemented"))
    }

    async fn poll_flight_info(
        &self,
        _: Request<FlightDescriptor>,
    ) -> TonicResult<Response<PollInfo>> {
        Err(Status::unimplemented("Not yet implemented"))
    }

    async fn get_schema(
        &self,
        _: Request<FlightDescriptor>,
    ) -> TonicResult<Response<SchemaResult>> {
        Err(Status::unimplemented("Not yet implemented"))
    }

    type DoGetStream = TonicStream<FlightData>;

    async fn do_get(&self, request: Request<Ticket>) -> TonicResult<Response<Self::DoGetStream>> {
        self.0.do_get(request).await
    }

    type DoPutStream = TonicStream<PutResult>;

    async fn do_put(
        &self,
        request: Request<Streaming<FlightData>>,
    ) -> TonicResult<Response<Self::DoPutStream>> {
        self.0.do_put(request).await
    }

    type DoExchangeStream = TonicStream<FlightData>;

    async fn do_exchange(
        &self,
        _: Request<Streaming<FlightData>>,
    ) -> TonicResult<Response<Self::DoExchangeStream>> {
        Err(Status::unimplemented("Not yet implemented"))
    }

    type DoActionStream = TonicStream<arrow_flight::Result>;

    async fn do_action(&self, _: Request<Action>) -> TonicResult<Response<Self::DoActionStream>> {
        Err(Status::unimplemented("Not yet implemented"))
    }

    type ListActionsStream = TonicStream<ActionType>;

    async fn list_actions(
        &self,
        _: Request<Empty>,
    ) -> TonicResult<Response<Self::ListActionsStream>> {
        Err(Status::unimplemented("Not yet implemented"))
    }
}

#[async_trait]
impl FlightCraft for GreptimeRequestHandler {
    async fn do_get(
        &self,
        request: Request<Ticket>,
    ) -> TonicResult<Response<TonicStream<FlightData>>> {
        let mut hints = hint_headers::extract_hints(request.metadata());
        hints.extend(extract_flow_extensions(request.metadata())?);
        let snapshot_seqs = extract_snapshot_seqs(request.metadata())?;

        let ticket = request.into_inner().ticket;
        let request =
            GreptimeRequest::decode(ticket.as_ref()).context(error::InvalidFlightTicketSnafu)?;
        let query_ctx =
            create_query_context(Channel::Grpc, request.header.as_ref(), hints, snapshot_seqs)?;
        // Validate flow hint syntax at the transport boundary before dispatching the request.
        // This does not authorize or execute anything; `handle_request()` below still performs
        // the normal frontend handling and auth checks before query execution.
        let flow_extensions = FlowQueryExtensions::parse_flow_extensions(&query_ctx.extensions())
            .map_err(|e| Status::invalid_argument(e.output_msg()))?;
        let should_emit_terminal_metrics = flow_extensions
            .as_ref()
            .is_some_and(|extensions| extensions.should_collect_region_watermark());

        // The Grpc protocol pass query by Flight. It needs to be wrapped under a span, in order to record stream
        let span = info_span!(
            "GreptimeRequestHandler::do_get",
            protocol = "grpc",
            request_type = get_request_type(&request)
        );
        let flight_compression = self.flight_compression;
        async {
            let query = request.request.context(InvalidQuerySnafu {
                reason: "Expecting non-empty GreptimeRequest.",
            })?;
            self.authenticate_request_with_query_ctx(request.header.as_ref(), &query_ctx)
                .await?;
            let output = self.handle_request_with_query_ctx(query, query_ctx.clone());
            let stream = to_flight_data_stream(
                output,
                TracingContext::from_current_span(),
                flight_compression,
                query_ctx,
                should_emit_terminal_metrics,
            );
            Ok(Response::new(stream))
        }
        .trace(span)
        .await
    }

    async fn do_put(
        &self,
        request: Request<Streaming<FlightData>>,
    ) -> TonicResult<Response<TonicStream<PutResult>>> {
        let (headers, extensions, stream) = request.into_parts();

        let limiter = extensions.get::<ServerMemoryLimiter>().cloned();

        let query_ctx = context_auth::create_query_context_from_grpc_metadata(&headers)?;
        context_auth::check_auth(self.user_provider.clone(), &headers, query_ctx.clone()).await?;

        const MAX_PENDING_RESPONSES: usize = 32;
        let (tx, rx) = mpsc::channel::<TonicResult<DoPutResponse>>(MAX_PENDING_RESPONSES);

        let stream = PutRecordBatchRequestStream::new(
            stream,
            query_ctx.current_catalog().to_string(),
            query_ctx.current_schema(),
            limiter,
        )
        .await?;
        // Ack immediately when stream is created successfully (in Init state)
        let _ = tx.send(Ok(DoPutResponse::new(0, 0, 0.0))).await;
        self.put_record_batches(stream, tx, query_ctx).await;

        let response = ReceiverStream::new(rx)
            .and_then(|response| {
                future::ready({
                    serde_json::to_vec(&response)
                        .context(ToJsonSnafu)
                        .map(|x| PutResult {
                            app_metadata: Bytes::from(x),
                        })
                        .map_err(Into::into)
                })
            })
            .boxed();
        Ok(Response::new(response))
    }
}

pub struct PutRecordBatchRequest {
    pub table_name: TableName,
    pub request_id: i64,
    pub timestamp_range: Option<(i64, i64)>,
    pub record_batch: DfRecordBatch,
    pub schema_bytes: Bytes,
    pub flight_data: FlightData,
    pub(crate) _guard: Option<MemoryGuard<RequestMemoryMetrics>>,
}

impl PutRecordBatchRequest {
    fn try_new(
        table_name: TableName,
        record_batch: DfRecordBatch,
        request_id: i64,
        timestamp_range: Option<(i64, i64)>,
        schema_bytes: Bytes,
        flight_data: FlightData,
        limiter: Option<&ServerMemoryLimiter>,
    ) -> Result<Self> {
        let memory_usage = flight_data.data_body.len()
            + flight_data.app_metadata.len()
            + flight_data.data_header.len();

        let _guard = if let Some(limiter) = limiter {
            let guard = limiter.try_acquire(memory_usage as u64).ok_or_else(|| {
                let inner_err = common_memory_manager::Error::MemoryLimitExceeded {
                    requested_bytes: memory_usage as u64,
                    limit_bytes: limiter.limit_bytes(),
                };
                error::MemoryLimitExceededSnafu.into_error(inner_err)
            })?;
            Some(guard)
        } else {
            None
        };

        Ok(Self {
            table_name,
            request_id,
            timestamp_range,
            record_batch,
            schema_bytes,
            flight_data,
            _guard,
        })
    }
}

pub struct PutRecordBatchRequestStream {
    flight_data_stream: Streaming<FlightData>,
    catalog: String,
    schema_name: String,
    limiter: Option<ServerMemoryLimiter>,
    // Client now lazily sends schema data so we cannot eagerly wait for it.
    // Instead, we need to decode while receiving record batches.
    state: StreamState,
}

enum StreamState {
    Init,
    Ready {
        table_name: TableName,
        schema: SchemaRef,
        schema_bytes: Bytes,
        decoder: FlightDecoder,
    },
}

impl PutRecordBatchRequestStream {
    /// Creates a new `PutRecordBatchRequestStream` in Init state.
    /// The stream will transition to Ready state when it receives the schema message.
    pub async fn new(
        flight_data_stream: Streaming<FlightData>,
        catalog: String,
        schema: String,
        limiter: Option<ServerMemoryLimiter>,
    ) -> TonicResult<Self> {
        Ok(Self {
            flight_data_stream,
            catalog,
            schema_name: schema,
            limiter,
            state: StreamState::Init,
        })
    }

    /// Returns the table name extracted from the flight descriptor.
    /// Returns None if the stream is still in Init state.
    pub fn table_name(&self) -> Option<&TableName> {
        match &self.state {
            StreamState::Init => None,
            StreamState::Ready { table_name, .. } => Some(table_name),
        }
    }

    /// Returns the Arrow schema decoded from the first flight message.
    /// Returns None if the stream is still in Init state.
    pub fn schema(&self) -> Option<&SchemaRef> {
        match &self.state {
            StreamState::Init => None,
            StreamState::Ready { schema, .. } => Some(schema),
        }
    }

    /// Returns the raw schema bytes in IPC format.
    /// Returns None if the stream is still in Init state.
    pub fn schema_bytes(&self) -> Option<&Bytes> {
        match &self.state {
            StreamState::Init => None,
            StreamState::Ready { schema_bytes, .. } => Some(schema_bytes),
        }
    }

    fn extract_table_name(mut descriptor: FlightDescriptor) -> Result<String> {
        ensure!(
            descriptor.r#type == arrow_flight::flight_descriptor::DescriptorType::Path as i32,
            InvalidParameterSnafu {
                reason: "expect FlightDescriptor::type == 'Path' only",
            }
        );
        ensure!(
            descriptor.path.len() == 1,
            InvalidParameterSnafu {
                reason: "expect FlightDescriptor::path has only one table name",
            }
        );
        Ok(descriptor.path.remove(0))
    }
}

impl Stream for PutRecordBatchRequestStream {
    type Item = TonicResult<PutRecordBatchRequest>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        loop {
            let poll = ready!(self.flight_data_stream.poll_next_unpin(cx));

            match poll {
                Some(Ok(flight_data)) => {
                    let limiter = self.limiter.clone();

                    match &mut self.state {
                        StreamState::Init => {
                            // First message - expecting schema
                            let flight_descriptor = match flight_data.flight_descriptor.as_ref() {
                                Some(descriptor) => descriptor.clone(),
                                None => {
                                    return Poll::Ready(Some(Err(Status::failed_precondition(
                                        "table to put is not found in flight descriptor",
                                    ))));
                                }
                            };

                            let table_name_str = match Self::extract_table_name(flight_descriptor) {
                                Ok(name) => name,
                                Err(e) => {
                                    return Poll::Ready(Some(Err(Status::invalid_argument(
                                        e.to_string(),
                                    ))));
                                }
                            };
                            let table_name = TableName::new(
                                self.catalog.clone(),
                                self.schema_name.clone(),
                                table_name_str,
                            );

                            // Decode the schema
                            let mut decoder = FlightDecoder::default();
                            let schema_message = decoder.try_decode(&flight_data).map_err(|e| {
                                Status::invalid_argument(format!("Failed to decode schema: {}", e))
                            })?;

                            match schema_message {
                                Some(FlightMessage::Schema(schema)) => {
                                    let schema_bytes = decoder.schema_bytes().ok_or_else(|| {
                                        Status::internal(
                                            "decoder should have schema bytes after decoding schema",
                                        )
                                    })?;

                                    // Transition to Ready state with all necessary data
                                    self.state = StreamState::Ready {
                                        table_name,
                                        schema,
                                        schema_bytes,
                                        decoder,
                                    };
                                    // Continue to next iteration to process RecordBatch messages
                                    continue;
                                }
                                _ => {
                                    return Poll::Ready(Some(Err(Status::failed_precondition(
                                        "first message must be a Schema message",
                                    ))));
                                }
                            }
                        }
                        StreamState::Ready {
                            table_name,
                            schema: _,
                            schema_bytes,
                            decoder,
                        } => {
                            // Extract request_id and time range from FlightData before decoding
                            let metadata = if !flight_data.app_metadata.is_empty() {
                                serde_json::from_slice::<DoPutMetadata>(&flight_data.app_metadata)
                                    .ok()
                            } else {
                                None
                            };
                            let request_id = metadata
                                .as_ref()
                                .map(|meta| meta.request_id())
                                .unwrap_or_default();
                            let timestamp_range = metadata.and_then(|meta| meta.timestamp_range());

                            // Decode FlightData to RecordBatch
                            match decoder.try_decode(&flight_data) {
                                Ok(Some(FlightMessage::RecordBatch(record_batch))) => {
                                    let table_name = table_name.clone();
                                    let schema_bytes = schema_bytes.clone();
                                    return Poll::Ready(Some(
                                        PutRecordBatchRequest::try_new(
                                            table_name,
                                            record_batch,
                                            request_id,
                                            timestamp_range,
                                            schema_bytes,
                                            flight_data,
                                            limiter.as_ref(),
                                        )
                                        .map_err(|e| Status::invalid_argument(e.to_string())),
                                    ));
                                }
                                Ok(Some(other)) => {
                                    debug!("Unexpected flight message: {:?}", other);
                                    return Poll::Ready(Some(Err(Status::invalid_argument(
                                        "Expected RecordBatch message, got other message type",
                                    ))));
                                }
                                Ok(None) => {
                                    // Dictionary batch - processed internally by decoder, continue polling
                                    continue;
                                }
                                Err(e) => {
                                    return Poll::Ready(Some(Err(Status::invalid_argument(
                                        format!("Failed to decode RecordBatch: {}", e),
                                    ))));
                                }
                            }
                        }
                    }
                }
                Some(Err(e)) => {
                    return Poll::Ready(Some(Err(e)));
                }
                None => {
                    return Poll::Ready(None);
                }
            }
        }
    }
}

fn extract_flow_extensions(
    metadata: &tonic::metadata::MetadataMap,
) -> TonicResult<Vec<(String, String)>> {
    Ok(extract_json_metadata(metadata, FLOW_EXTENSIONS_METADATA_KEY)?.unwrap_or_default())
}

fn extract_snapshot_seqs(
    metadata: &tonic::metadata::MetadataMap,
) -> TonicResult<HashMap<u64, u64>> {
    Ok(extract_json_metadata(metadata, SNAPSHOT_SEQS_METADATA_KEY)?.unwrap_or_default())
}

fn extract_json_metadata<T: serde::de::DeserializeOwned>(
    metadata: &tonic::metadata::MetadataMap,
    key: &'static str,
) -> TonicResult<Option<T>> {
    let Some(value) = metadata.get(key) else {
        return Ok(None);
    };

    let value = value
        .to_str()
        .map_err(|e| Status::invalid_argument(format!("Invalid {key} metadata value: {e}")))?;

    let parsed = serde_json::from_str::<T>(value)
        .map_err(|e| Status::invalid_argument(format!("Invalid {key} metadata JSON: {e}")))?;
    Ok(Some(parsed))
}

fn to_flight_data_stream<F>(
    output: F,
    tracing_context: TracingContext,
    flight_compression: FlightCompression,
    query_ctx: QueryContextRef,
    should_emit_terminal_metrics: bool,
) -> TonicStream<FlightData>
where
    F: std::future::Future<Output = Result<Output>> + Send + 'static,
{
    let initializer = async move {
        let output = output.await.map_err(Status::from)?;
        output_to_flight_record_batch_source(output, should_emit_terminal_metrics)
    };
    let stream = FlightRecordBatchStream::new(
        FlightRecordBatchStreamInput::initializer(initializer),
        tracing_context,
        flight_compression,
        query_ctx,
    );
    Box::pin(stream) as _
}

fn output_to_flight_record_batch_source(
    output: Output,
    should_emit_terminal_metrics: bool,
) -> TonicResult<FlightRecordBatchSource> {
    match output.data {
        OutputData::Stream(stream) => Ok(FlightRecordBatchSource::RecordBatches(stream)),
        OutputData::RecordBatches(x) => Ok(FlightRecordBatchSource::RecordBatches(x.as_stream())),
        OutputData::AffectedRows(rows) => {
            let terminal_metrics = match terminal_recordbatch_metrics_from_plan_if_requested(
                output.meta.plan,
                should_emit_terminal_metrics,
            ) {
                Some(metrics) => match serde_json::to_string(&metrics) {
                    Ok(metrics) => Some(metrics),
                    Err(e) => {
                        return Err(Status::internal(format!(
                            "Failed to serialize terminal metrics: {e}"
                        )));
                    }
                },
                None => None,
            };
            Ok(FlightRecordBatchSource::AffectedRows {
                rows,
                metrics: terminal_metrics,
            })
        }
    }
}

#[cfg(test)]
mod tests {
    use query::options::FLOW_SCHEDULED_TIME_MILLIS;
    use tonic::metadata::{AsciiMetadataValue, MetadataMap};

    use super::*;

    #[test]
    fn test_extract_flow_extensions_preserves_comma_bearing_values() {
        let mut metadata = MetadataMap::new();
        metadata.insert(
            FLOW_EXTENSIONS_METADATA_KEY,
            AsciiMetadataValue::try_from(
                r#"[["flow.return_region_seq","true"],["flow.incremental_after_seqs","{\"1\":10,\"2\":20}"]]"#,
            )
            .unwrap(),
        );

        let extensions = extract_flow_extensions(&metadata).unwrap();
        assert_eq!(
            extensions,
            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_flow_extensions_can_carry_scheduled_time() {
        let mut metadata = MetadataMap::new();
        metadata.insert(
            FLOW_EXTENSIONS_METADATA_KEY,
            AsciiMetadataValue::try_from(r#"[["flow.scheduled_time_millis","1700000000000"]]"#)
                .unwrap(),
        );

        let flow_extensions = extract_flow_extensions(&metadata).unwrap();
        let query_ctx =
            create_query_context(Channel::Grpc, None, flow_extensions, HashMap::new()).unwrap();

        assert_eq!(
            query_ctx.extension(FLOW_SCHEDULED_TIME_MILLIS),
            Some("1700000000000")
        );
    }

    #[test]
    fn test_extract_flow_extensions_rejects_invalid_json() {
        let mut metadata = MetadataMap::new();
        metadata.insert(
            FLOW_EXTENSIONS_METADATA_KEY,
            AsciiMetadataValue::try_from("not-json").unwrap(),
        );

        let err = extract_flow_extensions(&metadata).unwrap_err();
        assert_eq!(err.code(), tonic::Code::InvalidArgument);
    }
}