servers/postgres/copy_in/

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

//! PostgreSQL `COPY ... FROM STDIN` (copy-in) support.
//!
//! The wire handshake is driven by pgwire: returning [`Response::CopyIn`]
//! from a query handler makes the server send a `CopyInResponse` and enter
//! the copy-in sub-protocol; subsequent `CopyData` messages from the client
//! are dispatched to [`CopyHandler`], whose implementation below parses
//! the data stream incrementally and flushes it into the query handler as
//! [`RowInsertRequests`] batches.
//!
//! # Formats
//!
//! The supported formats are the explicit [`codec::CopyFormat`] enum —
//! `csv` (table-bound, incremental row parsing) and `influxdb`
//! (table-less, one bounded whole-payload request). Each format states
//! its own target rule and ingestion shape in the match arms of the
//! copy lifecycle below; adding a format is a new variant plus its
//! arms. Builders resolve a `FORMAT` name (and its aliases) through
//! [`REGISTERED_FORMATS`]. See the `codec` module docs.

mod codec;
mod csv;
mod influxdb;
mod options;
mod otlp;
mod parse;

use std::fmt::Debug;
use std::sync::Arc;

use crate::influxdb::InfluxdbRequest;
use crate::query_handler::PipelineHandlerRef;
use api::helper::ColumnDataTypeWrapper;
use api::v1::column_def::options_from_column_schema;
use api::v1::{
    Row as GrpcRow, RowInsertRequest, RowInsertRequests, Rows as GrpcRows, SemanticType,
};
use async_trait::async_trait;
use codec::{CopyFormat, RecordParser};
use common_catalog::consts::{TRACE_TABLE_NAME, TRACE_TABLE_NAME_SESSION_KEY};
use common_grpc::precision::Precision;
use common_query::OutputData;
use common_telemetry::{info, warn};
use csv::CsvCodec;
use datafusion::sql::sqlparser::ast::{CopyTarget, Statement as SqlParserStatement};
use datatypes::data_type::DataType;
use datatypes::schema::{ColumnSchema, SchemaRef};
use futures::{Sink, SinkExt, stream};
use influxdb::InfluxdbOptions;
use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetricsServiceRequest;
use otlp::{OtlpOptions, OtlpSignal};
use pgwire::api::copy::CopyHandler;
use pgwire::api::results::{CopyResponse, Response, Tag};
use pgwire::api::{ClientInfo, PgWireConnectionState};
use pgwire::error::{ErrorInfo, PgWireError, PgWireResult};
use pgwire::messages::PgWireBackendMessage;
use pgwire::messages::copy::{CopyData, CopyDone, CopyFail};
use prost::Message;
use session::context::QueryContextRef;
use table::metadata::TableInfo;

use crate::metrics::METRIC_POSTGRES_COPY_IN_ROWS;
use crate::postgres::PostgresServerHandlerInner;
use crate::postgres::types::PgErrorCode;
use crate::postgres::utils::convert_err;

/// Number of parsed rows buffered before flushing a batch to the handler.
///
/// Each flush runs the full insert pipeline (catalog lookup, on-demand
/// schema check, partition split, region write) and, without the
/// pending-rows batcher, waits for it — so larger batches mean fewer
/// round trips at the cost of linearly more memory per connection
/// (decoder batch + rows buffer). The constant doubles as the arrow
/// decoder's batch size. High-rate COPY users should enable the
/// pending-rows batcher (`global-table-batch` / `BatcherOptions.protocols`
/// containing `Postgres`), which turns flushes into enqueues and makes
/// the per-batch fixed costs negligible.
pub(crate) const COPY_IN_FLUSH_ROWS: usize = 32768;

/// Upper bound of a table-less COPY FROM STDIN payload (influxdb line
/// protocol, OTLP). One COPY statement is one bounded ingestion request
/// (the whole payload is buffered like one HTTP request body); a larger
/// import is multiple COPY statements.
pub(crate) const COPY_IN_MAX_PAYLOAD_BYTES: usize = 64 * 1024 * 1024;

/// Builds a format from the format-neutral option set.
type FormatBuilder = fn(options::CopyOptionSet) -> PgWireResult<CopyFormat>;

/// The registered COPY FROM STDIN formats: accepted `FORMAT` name to
/// builder. `csv` is the only implemented format; the others are
/// registered so they report a clear "not supported yet" instead of an
/// unknown-format error.
const REGISTERED_FORMATS: &[(&str, FormatBuilder)] = &[
    ("csv", csv::build_csv_codec),
    ("influxdb", influxdb::build_influxdb_codec),
    ("influxdb_line", influxdb::build_influxdb_codec),
    ("otlp", otlp::build_otlp_codec),
    ("otlp_metrics", otlp::build_otlp_codec),
    ("text", build_unimplemented_text_codec),
    ("txt", build_unimplemented_text_codec),
    ("binary", build_unimplemented_binary_codec),
];

fn build_unimplemented_text_codec(_options: options::CopyOptionSet) -> PgWireResult<CopyFormat> {
    Err(unsupported(
        "the text format is not supported yet; use WITH (FORMAT csv)",
    ))
}

fn build_unimplemented_binary_codec(_options: options::CopyOptionSet) -> PgWireResult<CopyFormat> {
    Err(unsupported(
        "COPY FROM STDIN WITH (FORMAT binary) is not supported yet",
    ))
}

/// Resolves the format from the parsed options.
fn resolve_format(options: &mut options::CopyOptionSet) -> PgWireResult<CopyFormat> {
    let name = options.take_format()?.ok_or_else(|| {
        unsupported(
            "COPY FROM STDIN requires an explicit format in this version; use WITH (FORMAT csv)",
        )
    })?;
    let builder = REGISTERED_FORMATS
        .iter()
        .find(|(registered, _)| *registered == name)
        .map(|(_, builder)| *builder)
        .ok_or_else(|| bad_copy_data(format!("unknown COPY format \"{}\"", name)))?;
    builder(std::mem::take(options))
}

/// The parsed `COPY tbl [(columns)] FROM STDIN [(options)]` statement.
#[derive(Debug, Clone)]
pub(crate) struct CopyFromStdin {
    /// Table name parts as written (`t`, `schema.t`, `catalog.schema.t`);
    /// quoted names keep their case.
    table: Vec<String>,
    /// Requested column list; empty means all columns in table order.
    columns: Vec<String>,
    /// The resolved format with its validated options.
    format: CopyFormat,
}

fn copy_in_error(code: PgErrorCode, message: impl Into<String>) -> PgWireError {
    PgWireError::UserError(Box::new(code.to_err_info(message.into())))
}

pub(super) fn unsupported(message: impl Into<String>) -> PgWireError {
    copy_in_error(PgErrorCode::Ec0A000, message)
}

fn bad_copy_data(message: impl Into<String>) -> PgWireError {
    copy_in_error(PgErrorCode::Ec22P04, message)
}

fn invalid_text_repr(message: impl Into<String>) -> PgWireError {
    copy_in_error(PgErrorCode::Ec22P02, message)
}

pub(crate) fn copy_in_must_run_alone_error() -> PgWireError {
    PgWireError::UserError(Box::new(ErrorInfo::new(
        "ERROR".to_string(),
        PgErrorCode::Ec0A000.code(),
        "COPY FROM STDIN must be executed alone".to_string(),
    )))
}

/// Parses a leading `COPY ... FROM STDIN` statement from SQL text.
///
/// The statement is parsed textually — not through sqlparser — because the
/// `WITH (...)` option list is open: formats may accept options beyond
/// sqlparser's COPY grammar, and unknown options must reach the format
/// codec instead of failing statement parsing. Returns `Ok(None)` when
/// the string does not lead with a copy-in statement.
pub(crate) fn copy_from_stdin_from_sql(sql: &str) -> PgWireResult<Option<CopyFromStdin>> {
    let raw = match parse::scan_copy_from_stdin(sql)? {
        parse::CopyInScan::NotCopyIn => return Ok(None),
        parse::CopyInScan::CopyIn(raw) => raw,
    };
    let mut option_set = raw.options;
    let format = resolve_format(&mut option_set)?;
    Ok(Some(CopyFromStdin {
        table: raw.table,
        columns: raw.columns,
        format,
    }))
}

/// Whether an already-parsed statement is a `COPY ... FROM STDIN`; used to
/// reject copy-in statements trailing other statements, where the textual
/// scanner (which only sees statement-leading COPYs) does not see them.
pub(crate) fn is_copy_from_stdin(statement: &SqlParserStatement) -> bool {
    matches!(
        statement,
        SqlParserStatement::Copy {
            to: false,
            target: CopyTarget::Stdin,
            ..
        }
    )
}

/// Resolves table name parts against the query context defaults.
fn resolve_table_name(
    query_ctx: &QueryContextRef,
    parts: &[String],
) -> PgWireResult<(String, String, String)> {
    match parts {
        [table] => Ok((
            query_ctx.current_catalog().to_string(),
            query_ctx.current_schema(),
            table.clone(),
        )),
        [schema, table] => Ok((
            query_ctx.current_catalog().to_string(),
            schema.clone(),
            table.clone(),
        )),
        [catalog, schema, table] => Ok((catalog.clone(), schema.clone(), table.clone())),
        _ => Err(unsupported(format!(
            "invalid table name for COPY: {}",
            parts.join(".")
        ))),
    }
}

/// Per-connection state of an in-progress COPY FROM STDIN.
///
/// The payload shape follows the codec: table-bound formats (csv) parse
/// records into rows against the targeted table's columns; table-less
/// formats (influxdb line protocol) split the stream into complete
/// lines that each route to their own measurement, and flush through
/// the handler's line-protocol pipeline — the same one the HTTP write
/// endpoint uses.
pub(crate) struct CopyInState {
    query_ctx: QueryContextRef,
    payload: Payload,
    /// The resolved codec's name, for logging.
    codec_name: &'static str,
    /// The wire format code of the resolved codec.
    codec_format_code: i8,
    rows_written: u64,
}

enum Payload {
    /// Table-bound rows: parsed against `grpc_schema` and flushed to
    /// `table_name` via `copy_in_insert`. The schema is cloned into each
    /// flushed batch: `Rows.schema` is an owned `Vec` in the proto
    /// request, so one copy per request is inherent.
    Table {
        grpc_schema: Vec<api::v1::ColumnSchema>,
        table_name: String,
        parser: Box<dyn RecordParser + Send>,
        /// Parsed rows not yet flushed to the handler.
        rows: Vec<GrpcRow>,
    },
    /// The whole buffered influxdb line-protocol payload with its
    /// request precision. One COPY statement is one bounded ingestion
    /// request, buffered until `CopyDone` — exactly like one HTTP write
    /// body — and record boundaries belong to the line-protocol parser,
    /// never to the COPY layer.
    Influxdb {
        data: Vec<u8>,
        precision: Option<Precision>,
    },
    /// The whole buffered OTLP payload with its validated options,
    /// decoded and completed as one request at `CopyDone` through the
    /// matching OTLP pipeline.
    Otlp { data: Vec<u8>, options: OtlpOptions },
}

impl CopyInState {
    fn new(
        codec: &CsvCodec,
        columns: Vec<String>,
        table_info: Arc<TableInfo>,
        query_ctx: QueryContextRef,
        table_name: String,
    ) -> PgWireResult<Self> {
        let schema = table_info.meta.schema.clone();
        let timestamp_index = schema.timestamp_index();

        let columns = resolve_columns(codec, columns, &schema, &table_info)?;
        if let Some(ts_index) = timestamp_index {
            let ts_column = &schema.column_schemas()[ts_index];
            if !columns.iter().any(|column| column.name == ts_column.name) {
                return Err(unsupported(format!(
                    "COPY FROM STDIN requires the time index column \"{}\" in the column list",
                    ts_column.name
                )));
            }
        }
        let parser = codec.create_parser(&columns, query_ctx.timezone().to_string().as_str())?;

        let grpc_schema = columns
            .iter()
            .map(|column| build_grpc_column_schema(column, &table_info))
            .collect::<PgWireResult<Vec<_>>>()?;

        Ok(Self {
            query_ctx,
            payload: Payload::Table {
                grpc_schema,
                table_name,
                parser,
                rows: Vec::with_capacity(COPY_IN_FLUSH_ROWS),
            },
            codec_name: "csv",
            codec_format_code: 0,
            rows_written: 0,
        })
    }

    /// State for the influxdb format: the whole payload is buffered and
    /// completed once at `CopyDone` through the line-protocol pipeline.
    fn new_influxdb(options: &InfluxdbOptions, query_ctx: QueryContextRef) -> Self {
        Self {
            query_ctx,
            payload: Payload::Influxdb {
                data: Vec::new(),
                precision: options.precision,
            },
            codec_name: "influxdb",
            codec_format_code: 0,
            rows_written: 0,
        }
    }

    /// State for the OTLP format: the whole binary payload is buffered
    /// and completed once at `CopyDone` through the matching OTLP
    /// pipeline. Binary: the advertised wire format code is 1.
    fn new_otlp(options: &OtlpOptions, query_ctx: QueryContextRef) -> Self {
        Self {
            query_ctx,
            payload: Payload::Otlp {
                data: Vec::new(),
                options: options.clone(),
            },
            codec_name: "otlp",
            codec_format_code: 1,
            rows_written: 0,
        }
    }

    /// The wire format code advertised to the client.
    pub(crate) fn format_code(&self) -> i8 {
        self.codec_format_code
    }

    /// Feeds a `CopyData` chunk into the parser and buffers parsed rows
    /// (csv) or payload bytes (influxdb).
    fn feed(&mut self, data: &[u8]) -> PgWireResult<()> {
        match &mut self.payload {
            Payload::Table { parser, rows, .. } => rows.extend(parser.feed(data)?),
            Payload::Influxdb { data: buffered, .. } | Payload::Otlp { data: buffered, .. } => {
                if buffered.len() + data.len() > COPY_IN_MAX_PAYLOAD_BYTES {
                    return Err(bad_copy_data(format!(
                        "COPY FROM STDIN payload exceeds the {} limit; split the import into multiple COPY statements",
                        COPY_IN_MAX_PAYLOAD_BYTES
                    )));
                }
                buffered.extend_from_slice(data);
            }
        }
        Ok(())
    }

    /// Parses the buffered tail at `CopyDone`; a final record without a
    /// trailing newline is accepted.
    fn finish(&mut self) -> PgWireResult<()> {
        match &mut self.payload {
            Payload::Table { parser, rows, .. } => rows.extend(parser.finish()?),
            // The whole payload is one request body; nothing to finalize.
            Payload::Influxdb { .. } | Payload::Otlp { .. } => {}
        }
        Ok(())
    }

    fn should_flush(&self) -> bool {
        match &self.payload {
            Payload::Table { rows, .. } => rows.len() >= COPY_IN_FLUSH_ROWS,
            // Completed once at CopyDone, not mid-stream.
            Payload::Influxdb { .. } | Payload::Otlp { .. } => false,
        }
    }

    /// Takes the buffered rows as an insert request, if any.
    fn take_requests(&mut self) -> Option<RowInsertRequests> {
        let Payload::Table {
            grpc_schema,
            table_name,
            rows,
            ..
        } = &mut self.payload
        else {
            return None;
        };
        if rows.is_empty() {
            return None;
        }
        let rows = GrpcRows {
            schema: grpc_schema.clone(),
            rows: std::mem::take(rows),
        };
        Some(RowInsertRequests {
            inserts: vec![RowInsertRequest {
                table_name: table_name.clone(),
                rows: Some(rows),
            }],
        })
    }

    /// Takes the buffered OTLP payload and its options, if any.
    fn take_otlp(&mut self) -> Option<(OtlpOptions, Vec<u8>)> {
        let Payload::Otlp { options, data } = &mut self.payload else {
            return None;
        };
        if data.is_empty() {
            return None;
        }
        Some((options.clone(), std::mem::take(data)))
    }

    /// Takes the buffered payload as one line-protocol request body, if
    /// any — the same bytes an HTTP write of the same records would
    /// carry.
    fn take_lines(&mut self) -> PgWireResult<Option<InfluxdbRequest>> {
        let Payload::Influxdb { data, precision } = &mut self.payload else {
            return Ok(None);
        };
        if data.is_empty() {
            return Ok(None);
        }
        let lines = String::from_utf8(std::mem::take(data))
            .map_err(|e| bad_copy_data(format!("COPY FROM STDIN data is not valid UTF-8: {e}")))?;
        Ok(Some(InfluxdbRequest {
            precision: *precision,
            lines,
        }))
    }
}

fn resolve_columns(
    codec: &CsvCodec,
    columns: Vec<String>,
    schema: &SchemaRef,
    table_info: &Arc<TableInfo>,
) -> PgWireResult<Vec<ColumnSchema>> {
    let table_name = &table_info.name;
    let selected: Vec<ColumnSchema> = if columns.is_empty() {
        schema.column_schemas().to_vec()
    } else {
        let mut selected = Vec::with_capacity(columns.len());
        for name in &columns {
            let column = schema
                .column_schema_by_name(name)
                .ok_or_else(|| {
                    copy_in_error(
                        PgErrorCode::Ec42703,
                        format!(
                            "column \"{}\" of relation \"{}\" does not exist",
                            name, table_name
                        ),
                    )
                })?
                .clone();
            if selected.iter().any(|c: &ColumnSchema| &c.name == name) {
                return Err(bad_copy_data(format!(
                    "column \"{}\" specified more than once",
                    name
                )));
            }
            selected.push(column);
        }
        selected
    };

    for column in &selected {
        codec.check_column_type(&column.data_type)?;
    }

    Ok(selected)
}

/// Derives a column's semantic type the same way the flow adapter does
/// (`flow::adapter::util::column_schemas_to_proto`): the time-index flag
/// for timestamps, primary-key membership for tags.
fn semantic_type_of(table_info: &TableInfo, column: &ColumnSchema) -> PgWireResult<SemanticType> {
    let meta = &table_info.meta;
    if column.is_time_index() {
        Ok(SemanticType::Timestamp)
    } else if meta
        .primary_key_indices
        .iter()
        .any(|&index| meta.schema.column_schemas()[index].name == column.name)
    {
        Ok(SemanticType::Tag)
    } else {
        Ok(SemanticType::Field)
    }
}

fn build_grpc_column_schema(
    column: &ColumnSchema,
    table_info: &Arc<TableInfo>,
) -> PgWireResult<api::v1::ColumnSchema> {
    let (datatype, datatype_extension) = ColumnDataTypeWrapper::try_from(column.data_type.clone())
        .map_err(|_| {
            unsupported(format!(
                "column type \"{}\" is not supported for COPY FROM STDIN yet",
                column.data_type.name()
            ))
        })?
        .to_parts();
    Ok(api::v1::ColumnSchema {
        column_name: column.name.clone(),
        datatype: datatype as i32,
        semantic_type: semantic_type_of(table_info, column)? as i32,
        datatype_extension,
        options: options_from_column_schema(column),
    })
}

// ---------------------------------------------------------------------------
// pgwire handlers
// ---------------------------------------------------------------------------

impl PostgresServerHandlerInner {
    /// Prepares a COPY FROM STDIN: resolves the target table, validates
    /// columns and stashes the per-connection copy state.
    pub(crate) async fn begin_copy_in(&self, stmt: CopyFromStdin) -> PgWireResult<Response> {
        // The batching-aware context routes copy-in flushes through the
        // pending-rows batcher like SQL INSERT when it is enabled for
        // this connection.
        let query_ctx = self.new_query_context();

        // Each format states its own target rule and ingestion shape
        // explicitly.
        let state = match &stmt.format {
            // csv requires a target table: records are parsed against
            // its columns and inserted in batches.
            CopyFormat::Csv(codec) => {
                if stmt.table.is_empty() {
                    return Err(unsupported(
                        "format csv requires a table: COPY tbl [(cols)] FROM STDIN WITH (FORMAT csv)",
                    ));
                }
                let (catalog, schema, table) = resolve_table_name(&query_ctx, &stmt.table)?;

                let Some(table_info) = self
                    .copy_in_handler
                    .copy_in_table(&catalog, &schema, &table, query_ctx.clone())
                    .await
                    .map_err(convert_err)?
                else {
                    return Err(copy_in_error(
                        PgErrorCode::Ec42P01,
                        format!("relation \"{}\" does not exist", stmt.table.join(".")),
                    ));
                };

                // The insert path resolves the bare table name against the
                // query context, so point the context at the resolved
                // catalog and schema when the COPY target is qualified.
                let query_ctx = adjust_query_context(query_ctx, &catalog, &schema);

                CopyInState::new(codec, stmt.columns, table_info, query_ctx, table)?
            }
            // influxdb forbids a target table — each line names its own
            // measurement — and completes through the line-protocol
            // pipeline the HTTP write endpoints use; the server must be
            // wired for it.
            CopyFormat::Influxdb(options) => {
                if !stmt.table.is_empty() {
                    return Err(unsupported(format!(
                        "format influxdb routes each record to its own table; \"COPY {} FROM STDIN\" must not name one",
                        stmt.table.join(".")
                    )));
                }
                if self.influxdb_handler.is_none() {
                    return Err(unsupported(
                        "this server is not wired for the influxdb line protocol format",
                    ));
                }
                CopyInState::new_influxdb(options, query_ctx)
            }
            // OTLP forbids a target table too — the payload's resource
            // metrics route to their own tables — and completes through
            // the OTLP metrics pipeline; the server must be wired.
            CopyFormat::Otlp(options) => {
                if !stmt.table.is_empty() {
                    return Err(unsupported(format!(
                        "format otlp routes each record to its own table; \"COPY {} FROM STDIN\" must not name one",
                        stmt.table.join(".")
                    )));
                }
                if self.otlp_handler.is_none() {
                    return Err(unsupported("this server is not wired for the otlp format"));
                }
                CopyInState::new_otlp(options, query_ctx)
            }
        };
        let num_columns = match &state.payload {
            Payload::Table { grpc_schema, .. } => grpc_schema.len(),
            // No fixed column set: the format routes per record.
            Payload::Influxdb { .. } | Payload::Otlp { .. } => 0,
        };
        let format_code = state.format_code();

        *self.copy_in_state.lock().await = Some(state);

        Ok(Response::CopyIn(CopyResponse::new(
            format_code,
            num_columns,
            stream::empty(),
        )))
    }

    /// Flushes the buffered payload of an in-progress copy. Returns the
    /// number of rows written by the flushed batch.
    async fn flush_copy_in(&self, state: &mut CopyInState) -> PgWireResult<u64> {
        let query_ctx = state.query_ctx.clone();
        // The buffered influxdb payload is completed through the
        // line-protocol pipeline — the same handler the HTTP write
        // endpoints use, so conversion, permissions, interceptors,
        // timestamp alignment and inserts are identical to an HTTP
        // write of the same bytes. The context stays this connection's
        // postgres context.
        if let Some((options, bytes)) = state.take_otlp() {
            let Some(otlp_handler) = &self.otlp_handler else {
                return Err(unsupported("this server is not wired for the otlp format"));
            };
            return self
                .flush_otlp(otlp_handler, options, bytes, query_ctx)
                .await;
        }
        if let Some(request) = state.take_lines()? {
            let Some(influxdb_handler) = &self.influxdb_handler else {
                return Err(unsupported(
                    "this server is not wired for the influxdb line protocol format",
                ));
            };
            let output = influxdb_handler
                .exec(request, query_ctx)
                .await
                .map_err(convert_err)?;
            return match output.data {
                OutputData::AffectedRows(rows) => Ok(rows as u64),
                OutputData::RecordBatches(_) | OutputData::Stream(_) => Ok(0),
            };
        }
        let Some(requests) = state.take_requests() else {
            return Ok(0);
        };
        let output = self
            .copy_in_handler
            .copy_in_insert(requests, query_ctx)
            .await
            .map_err(convert_err)?;
        match output.data {
            OutputData::AffectedRows(rows) => Ok(rows as u64),
            OutputData::RecordBatches(_) | OutputData::Stream(_) => {
                warn!("unexpected output from COPY FROM STDIN insert");
                Ok(0)
            }
        }
    }

    /// Completes an OTLP copy through the pipeline the matching
    /// OTLP/HTTP endpoint uses — the payload is decoded with the same
    /// protobuf types and handed to the same handler methods with the
    /// same context preparation, so both ingestion routes are
    /// identical. The context stays this connection's postgres context
    /// apart from the signal-specific preparation below.
    async fn flush_otlp(
        &self,
        otlp_handler: &crate::query_handler::OpenTelemetryProtocolHandlerRef,
        options: OtlpOptions,
        bytes: Vec<u8>,
        query_ctx: QueryContextRef,
    ) -> PgWireResult<u64> {
        use opentelemetry_proto::tonic::collector::logs::v1::ExportLogsServiceRequest;
        use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest;
        use pipeline::{GreptimePipelineParams, PipelineWay, SelectInfo};
        use session::protocol_ctx::{MetricType, OtlpMetricCtx, ProtocolCtx};

        // The HTTP routes derive the pipeline handler from the same
        // protocol handler via trait upcasting; so does the copy.
        let pipeline_handler: PipelineHandlerRef = otlp_handler.clone();
        match options.signal {
            OtlpSignal::Metrics {
                promote_all_resource_attrs,
                resource_attrs,
                promote_scope_attrs,
                translation_strategy,
            } => {
                // Same context preparation as the HTTP metrics route.
                let mut ctx = (*query_ctx).clone();
                ctx.set_protocol_ctx(ProtocolCtx::OtlpMetric(OtlpMetricCtx {
                    promote_all_resource_attrs,
                    resource_attrs,
                    promote_scope_attrs,
                    with_metric_engine: self.otlp_with_metric_engine,
                    // Set by the frontend from its config.
                    resource_info: false,
                    metric_type: MetricType::Init,
                    metric_translation_strategy: translation_strategy,
                }));
                let request = ExportMetricsServiceRequest::decode(bytes.as_slice())
                    .map_err(|e| bad_copy_data(format!("invalid OTLP payload: {e}")))?;
                let outcome = otlp_handler
                    .metrics(
                        request,
                        Arc::new(ctx),
                        Some((self.memory_limiter.clone(), 0)),
                    )
                    .await
                    .map_err(convert_err)?;
                // Same failure semantics as the HTTP route: a payload
                // whose points were all rejected is an error, not a
                // successful `COPY 0`.
                if outcome.accepted_data_points == 0 && outcome.rejected_data_points > 0 {
                    return Err(all_rejected_error(
                        "data points",
                        outcome.rejected_data_points,
                        outcome.error_message.as_deref(),
                    ));
                }
                // Rows are the accepted data points.
                Ok(outcome.accepted_data_points.max(0) as u64)
            }
            OtlpSignal::Traces {
                table_name,
                pipeline_name,
                pipeline_version,
                pipeline_params,
            } => {
                let table_name = table_name.unwrap_or_else(|| TRACE_TABLE_NAME.to_string());
                // Same context preparation as the HTTP traces route.
                let mut ctx = (*query_ctx).clone();
                ctx.set_extension(TRACE_TABLE_NAME_SESSION_KEY, &table_name);
                let pipeline = PipelineWay::from_name_and_default(
                    pipeline_name.as_deref(),
                    pipeline_version.as_deref(),
                    None,
                )
                .map_err(convert_err)?;
                let request = ExportTraceServiceRequest::decode(bytes.as_slice())
                    .map_err(|e| bad_copy_data(format!("invalid OTLP payload: {e}")))?;
                let outcome = otlp_handler
                    .traces(
                        pipeline_handler,
                        request,
                        pipeline,
                        GreptimePipelineParams::from_params(pipeline_params.as_deref()),
                        table_name,
                        Arc::new(ctx),
                    )
                    .await
                    .map_err(convert_err)?;
                // Same failure semantics as the HTTP route.
                if outcome.accepted_spans == 0 && outcome.rejected_spans > 0 {
                    return Err(all_rejected_error(
                        "spans",
                        outcome.rejected_spans as i64,
                        outcome.error_message.as_deref(),
                    ));
                }
                // Rows are the accepted spans.
                Ok(outcome.accepted_spans as u64)
            }
            OtlpSignal::Logs {
                table_name,
                pipeline_name,
                pipeline_version,
                pipeline_params,
                extract_keys,
            } => {
                let table_name = table_name.unwrap_or_else(|| "opentelemetry_logs".to_string());
                // Same select semantics as the HTTP extract-keys header.
                let select_info = match extract_keys.as_deref() {
                    Some(keys) if !keys.is_empty() => SelectInfo::from(keys.to_string()),
                    _ => Default::default(),
                };
                let pipeline = PipelineWay::from_name_and_default(
                    pipeline_name.as_deref(),
                    pipeline_version.as_deref(),
                    Some(PipelineWay::OtlpLogDirect(Box::new(select_info))),
                )
                .map_err(convert_err)?;
                let request = ExportLogsServiceRequest::decode(bytes.as_slice())
                    .map_err(|e| bad_copy_data(format!("invalid OTLP payload: {e}")))?;
                let outputs = otlp_handler
                    .logs(
                        pipeline_handler,
                        request,
                        pipeline,
                        GreptimePipelineParams::from_params(pipeline_params.as_deref()),
                        table_name,
                        query_ctx,
                    )
                    .await
                    .map_err(convert_err)?;
                // Rows are the affected rows across the outputs.
                Ok(outputs
                    .iter()
                    .map(|o| match &o.data {
                        OutputData::AffectedRows(rows) => *rows as u64,
                        _ => 0,
                    })
                    .sum())
            }
        }
    }
}

/// The error for an OTLP payload whose points/spans were all rejected,
/// carrying the pipeline's reason when it provided one.
fn all_rejected_error(kind: &str, rejected: i64, reason: Option<&str>) -> PgWireError {
    bad_copy_data(match reason {
        Some(reason) => format!("OTLP ingestion rejected all {rejected} {kind}: {reason}"),
        None => format!("OTLP ingestion rejected all {rejected} {kind}"),
    })
}

fn adjust_query_context(
    query_ctx: QueryContextRef,
    catalog: &str,
    schema: &str,
) -> QueryContextRef {
    if query_ctx.current_catalog() == catalog && query_ctx.current_schema() == schema {
        return query_ctx;
    }
    let mut forked = query_ctx.fork();
    forked.set_current_catalog(catalog);
    forked.set_current_schema(schema);
    Arc::new(forked)
}

#[async_trait]
impl CopyHandler for PostgresServerHandlerInner {
    async fn on_copy_data<C>(&self, _client: &mut C, copy_data: CopyData) -> PgWireResult<()>
    where
        C: ClientInfo + Sink<PgWireBackendMessage> + Unpin + Send + Sync,
        C::Error: Debug,
        PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
    {
        let mut guard = self.copy_in_state.lock().await;
        let Some(mut state) = guard.take() else {
            return Err(copy_in_error(
                PgErrorCode::Ec08P01,
                "CopyData message received outside COPY FROM STDIN",
            ));
        };
        // The state is owned here: any parse, conversion or insert error
        // below aborts the whole copy and drops the buffered data —
        // otherwise a pooled connection would retain failed-import
        // buffers until its next COPY. On success it is put back.
        state.feed(&copy_data.data)?;
        if state.should_flush() {
            let rows = self.flush_copy_in(&mut state).await?;
            state.rows_written += rows;
            METRIC_POSTGRES_COPY_IN_ROWS.inc_by(rows);
        }
        *guard = Some(state);
        Ok(())
    }

    async fn on_copy_done<C>(&self, client: &mut C, _done: CopyDone) -> PgWireResult<()>
    where
        C: ClientInfo + Sink<PgWireBackendMessage> + Unpin + Send + Sync,
        C::Error: Debug,
        PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
    {
        let mut guard = self.copy_in_state.lock().await;
        // Take ownership up front: even finalization failures (parse
        // errors in the trailing record, insert errors) drop the state
        // instead of retaining it on a long-lived connection.
        let Some(mut state) = guard.take() else {
            return Err(copy_in_error(
                PgErrorCode::Ec08P01,
                "CopyDone message received outside COPY FROM STDIN",
            ));
        };
        state.finish()?;
        let rows = self.flush_copy_in(&mut state).await?;
        state.rows_written += rows;
        METRIC_POSTGRES_COPY_IN_ROWS.inc_by(rows);

        let rows_written = state.rows_written;
        let codec_name = state.codec_name;
        info!(
            "PostgreSQL COPY FROM STDIN (format {}) finished: {} rows written",
            codec_name, rows_written
        );
        client
            .send(PgWireBackendMessage::CommandComplete(
                Tag::new("COPY").with_rows(rows_written as usize).into(),
            ))
            .await?;

        // After a successful extended-protocol copy-in, pgwire leaves the
        // connection in `CopyInProgress(true)` and swallows the client's
        // terminating `Sync` message, so `ReadyForQuery` would never be
        // sent and the connection would deadlock. Route the state to
        // `AwaitingSync` so the `Sync` is dispatched to `on_sync`, which
        // completes the round trip.
        if let PgWireConnectionState::CopyInProgress(true) = client.state() {
            client.set_state(PgWireConnectionState::AwaitingSync);
        }
        Ok(())
    }

    async fn on_copy_fail<C>(&self, _client: &mut C, fail: CopyFail) -> PgWireError
    where
        C: ClientInfo + Sink<PgWireBackendMessage> + Unpin + Send + Sync,
        C::Error: Debug,
        PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
    {
        // Drop the buffered state; partial data of a failed copy is not
        // written.
        self.copy_in_state.lock().await.take();
        copy_in_error(
            PgErrorCode::Ec57014,
            format!("COPY FROM STDIN mode terminated by user: {}", fail.message),
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_parse_copy_from_stdin_requires_format() {
        let err = copy_from_stdin_from_sql("COPY t FROM STDIN").unwrap_err();
        assert!(
            err.to_string()
                .contains("requires an explicit format in this version"),
            "{err}"
        );

        let copy = copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT csv)")
            .unwrap()
            .unwrap();
        assert_eq!(copy.table, vec!["t"]);
        assert!(copy.columns.is_empty());
        assert!(matches!(copy.format, CopyFormat::Csv(_)));
        // Both formats are textual: format code 0.
    }

    #[test]
    fn test_parse_copy_from_stdin_csv() {
        let copy = copy_from_stdin_from_sql(
            "COPY t (a, b) FROM STDIN WITH (FORMAT csv, HEADER, DELIMITER ';')",
        )
        .unwrap()
        .unwrap();
        assert_eq!(copy.columns, vec!["a", "b"]);
        assert!(matches!(copy.format, CopyFormat::Csv(_)));
    }

    #[test]
    fn test_parse_copy_from_stdin_text_unsupported() {
        for sql in [
            "COPY t FROM STDIN WITH (FORMAT text)",
            "COPY t FROM STDIN WITH (FORMAT txt)",
        ] {
            let err = copy_from_stdin_from_sql(sql).unwrap_err();
            assert!(
                err.to_string().contains("text format is not supported yet"),
                "{err}"
            );
        }
    }

    #[test]
    fn test_parse_copy_from_stdin_binary_unsupported() {
        let err = copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT binary)").unwrap_err();
        assert!(err.to_string().contains("not supported"));
    }

    #[test]
    fn test_parse_copy_from_stdin_unknown_format() {
        let err = copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT parquet)").unwrap_err();
        assert!(err.to_string().contains("unknown COPY format"));
    }

    #[test]
    fn test_parse_copy_from_stdin_not_copy_in() {
        assert!(
            copy_from_stdin_from_sql("COPY (SELECT 1) TO STDOUT")
                .unwrap()
                .is_none()
        );
        assert!(
            copy_from_stdin_from_sql("COPY t TO '/tmp/x.csv'")
                .unwrap()
                .is_none()
        );
        // GreptimeDB's own file-based COPY FROM.
        assert!(
            copy_from_stdin_from_sql("COPY t FROM '/tmp/x.csv' WITH (FORMAT csv)")
                .unwrap()
                .is_none()
        );
    }

    #[test]
    fn test_redundant_options() {
        let err = copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT csv, FORMAT text)")
            .unwrap_err();
        assert!(err.to_string().contains("redundant FORMAT"));

        let err = copy_from_stdin_from_sql("COPY t FROM STDIN WITH (DELIMITER ',', DELIMITER ';')")
            .unwrap_err();
        assert!(err.to_string().contains("redundant DELIMITER"));
    }

    #[test]
    fn test_format_specific_options_rejected() {
        // The open WITH section parses anything, but the format rejects
        // the options it does not understand.
        let err =
            copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT csv, PRECISION 'ns', FREEZE)")
                .unwrap_err();
        assert!(
            err.to_string().contains("not supported by format csv"),
            "{err}"
        );
    }

    #[test]
    fn test_multi_statement_rejected() {
        let err = copy_from_stdin_from_sql("COPY t FROM STDIN; SELECT 1").unwrap_err();
        assert!(err.to_string().contains("alone"));
        // Trailing separators and quoted semicolons are fine.
        assert!(
            copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT csv);")
                .unwrap()
                .is_some()
        );
        assert!(
            copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT csv, DELIMITER ';')")
                .unwrap()
                .is_some()
        );
    }
}