diff --git a/src/frontend/src/server.rs b/src/frontend/src/server.rs index 0b1e2b9713e..d84f01fce95 100644 --- a/src/frontend/src/server.rs +++ b/src/frontend/src/server.rs @@ -429,6 +429,7 @@ where user_provider.clone(), Some(self.instance.process_manager().clone()), ) + .with_influxdb_handler(instance.clone()) .with_batching_enabled(postgres_batching), ) as Box; diff --git a/src/servers/src/http/influxdb.rs b/src/servers/src/http/influxdb.rs index 671beb7bae2..44815cae3a1 100644 --- a/src/servers/src/http/influxdb.rs +++ b/src/servers/src/http/influxdb.rs @@ -24,7 +24,7 @@ use common_grpc::precision::Precision; use common_telemetry::tracing; use session::context::{Channel, QueryContext, QueryContextRef}; -use crate::error::{Result, TimePrecisionSnafu}; +use crate::error::Result; use crate::http::header::write_cost_header_map; use crate::influxdb::InfluxdbRequest; use crate::query_handler::InfluxdbLineProtocolHandlerRef; @@ -107,30 +107,13 @@ pub async fn influxdb_write( )) } -fn parse_time_precision(value: &str) -> Result { - // Precision conversion needs to be compatible with influxdb v1 v2 api. - // For details, see the Influxdb documents. - // https://docs.influxdata.com/influxdb/v1.8/tools/api/#apiv2write-http-endpoint - // https://docs.influxdata.com/influxdb/v1.8/tools/api/#write-http-endpoint - match value { - "n" | "ns" => Ok(Precision::Nanosecond), - "u" | "us" => Ok(Precision::Microsecond), - "ms" => Ok(Precision::Millisecond), - "s" => Ok(Precision::Second), - "m" => Ok(Precision::Minute), - "h" => Ok(Precision::Hour), - unknown => TimePrecisionSnafu { - name: unknown.to_string(), - } - .fail(), - } -} +pub(crate) use crate::influxdb::parse_time_precision; #[cfg(test)] mod tests { use common_grpc::precision::Precision; - use crate::http::influxdb::parse_time_precision; + use super::parse_time_precision; #[test] fn test_parse_time_precision() { diff --git a/src/servers/src/influxdb.rs b/src/servers/src/influxdb.rs index 9bff0bbc6e4..9dd26d9cda7 100644 --- a/src/servers/src/influxdb.rs +++ b/src/servers/src/influxdb.rs @@ -20,7 +20,7 @@ use hyper::Request; use influxdb_line_protocol::{FieldValue, parse_lines}; use snafu::ResultExt; -use crate::error::{Error, InfluxdbLineProtocolSnafu}; +use crate::error::{Error, InfluxdbLineProtocolSnafu, TimePrecisionSnafu}; use crate::row_writer::{self, MultiTableData}; const INFLUXDB_API_PATH_NAME: &str = "influxdb"; @@ -113,6 +113,27 @@ fn unwrap_or_default_precision(precision: Option) -> Precision { } } +/// Parses an InfluxDB time-precision name shared by the HTTP write +/// endpoints (v1/v2) and the `precision` option of COPY FROM STDIN. +pub(crate) fn parse_time_precision(value: &str) -> crate::error::Result { + // Precision conversion needs to be compatible with influxdb v1 v2 api. + // For details, see the Influxdb documents. + // https://docs.influxdata.com/influxdb/v1.8/tools/api/#apiv2write-http-endpoint + // https://docs.influxdata.com/influxdb/v1.8/tools/api/#write-http-endpoint + match value { + "n" | "ns" => Ok(Precision::Nanosecond), + "u" | "us" => Ok(Precision::Microsecond), + "ms" => Ok(Precision::Millisecond), + "s" => Ok(Precision::Second), + "m" => Ok(Precision::Minute), + "h" => Ok(Precision::Hour), + unknown => TimePrecisionSnafu { + name: unknown.to_string(), + } + .fail(), + } +} + #[cfg(test)] mod tests { use api::v1::value::ValueData; diff --git a/src/servers/src/postgres.rs b/src/servers/src/postgres.rs index e13ec7392d5..fbef2e9ac4c 100644 --- a/src/servers/src/postgres.rs +++ b/src/servers/src/postgres.rs @@ -43,8 +43,8 @@ use session::context::Channel; use self::auth_handler::{PgUserProviderSource, ZeroPidSecretKeyGenerator}; use self::handler::DefaultQueryParser; -use crate::query_handler::CopyInHandlerRef; use crate::query_handler::sql::ServerSqlQueryHandlerRef; +use crate::query_handler::{CopyInHandlerRef, InfluxdbLineProtocolHandlerRef}; pub(crate) struct GreptimeDBStartupParameters { version: String, @@ -81,6 +81,11 @@ pub struct PostgresServerHandlerInner { query_handler: ServerSqlQueryHandlerRef, auth_source: Option>, copy_in_handler: CopyInHandlerRef, + /// Executes influxdb line-protocol batches flushed by a table-less + /// COPY FROM STDIN — the same handler the HTTP write endpoints use. + /// Absent on servers not wired for it; such connections reject the + /// influxdb format when a copy starts. + influxdb_handler: Option, force_tls: bool, auth_driver: AuthDriver, UserInfoRef>, /// SCRAM selected; failures get the uniform metric label. @@ -99,6 +104,8 @@ pub struct PostgresServerHandlerInner { pub(crate) struct MakePostgresServerHandler { query_handler: ServerSqlQueryHandlerRef, copy_in_handler: CopyInHandlerRef, + #[builder(default)] + influxdb_handler: Option, user_provider: Option, #[builder(default = "Arc::new(GreptimeDBStartupParameters::new())")] param_provider: Arc, @@ -149,6 +156,7 @@ impl MakePostgresServerHandler { .clone() .map(|p| Arc::new(PgUserProviderSource::new(p))), copy_in_handler: self.copy_in_handler.clone(), + influxdb_handler: self.influxdb_handler.clone(), force_tls: self.force_tls, auth_driver: AuthDriver::new(self.param_provider.clone()) .with_pid_secret_key_generator(Arc::new(ZeroPidSecretKeyGenerator)), diff --git a/src/servers/src/postgres/copy_in/codec.rs b/src/servers/src/postgres/copy_in/codec.rs index 37a94811309..a7672c619ef 100644 --- a/src/servers/src/postgres/copy_in/codec.rs +++ b/src/servers/src/postgres/copy_in/codec.rs @@ -12,58 +12,63 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! The COPY FROM STDIN codec framework. +//! The COPY FROM STDIN formats. //! -//! A *codec* binds one COPY data format to its implementation. Codecs are -//! stateless strategies: they validate the options they accept at -//! statement-parse time, describe which column types they can ingest, and -//! create the stateful [`RecordParser`] per stream. -//! -//! Formats are not limited to PostgreSQL's own — GreptimeDB-specific -//! formats (a line protocol, a JSON-lines dialect, ...) register exactly -//! like `csv` does, and may define their own options through the open -//! `WITH (...)` clause. +//! [`CopyFormat`] is the closed, explicit set of supported formats — a +//! small enum over the real implementations rather than a generic +//! format framework. Each variant carries its validated options and +//! its rules are stated where they apply: a variant's match arms in +//! the copy lifecycle (see the parent module) decide whether a target +//! table is required, how the `CopyData` bytes are consumed, and which +//! existing ingestion pipeline completes the copy. Adding a format is +//! a normal extension: a new variant plus its arms. //! //! # Parsing and conversion //! -//! A [`RecordParser`] consumes raw `CopyData` bytes and emits complete, -//! insertable rows: one [`ParsedRow`] per record, in the order of the -//! targeted columns. Parsers validate record shape (field counts) and -//! convert field values themselves — `csv` delegates both splitting and -//! typed conversion to arrow-csv's incremental decoder (the engine -//! behind `COPY ... FROM 'file.csv'`) and turns decoded batches into -//! rows with the shared `api::helper::vectors_to_rows`. A future format -//! that carries values as text (e.g. a line protocol) would reintroduce -//! shared field-conversion helpers here; see the removed `text` codec in -//! git history for the shape they had. +//! Table-bound formats parse incrementally through a [`RecordParser`]: +//! it consumes raw `CopyData` bytes and emits complete, insertable +//! rows — one [`ParsedRow`] per record, in the order of the targeted +//! columns. `csv` delegates splitting and typed conversion to +//! arrow-csv's incremental decoder (the engine behind +//! `COPY ... FROM 'file.csv'`) and turns decoded batches into rows +//! with the shared `api::helper::vectors_to_rows`. +//! +//! Formats that delegate to an existing ingestion pipeline (the +//! influxdb line protocol) do not parse here at all: the whole payload +//! is buffered as one bounded request body and completed through that +//! pipeline, so record boundaries and values stay with the pipeline's +//! own parser. //! //! # Options //! //! The `WITH (...)` clause is open: every `KEY value` entry the client -//! writes is collected into a [`CopyOptionSet`] (see `options`) regardless -//! of whether any format understands it — the statement is parsed -//! textually (see `parse`) precisely so options beyond sqlparser's COPY -//! grammar can reach the codecs. A codec builder takes the options it -//! understands (`take_string`, `take_single_byte`, `take_flag`, or raw -//! `take_value` for codec-specific shapes) and rejects the rest with -//! `reject_leftovers`, so a future format like -//! `COPY t FROM STDIN WITH (FORMAT influxdb_line, precision 'ns')` only -//! needs its own codec and registry entry. -//! -//! To add a format: -//! -//! 1. implement [`CopyInCodec`] for a config struct holding the format's -//! validated options (see `csv::CsvCodec` for the example), -//! 2. write a builder `fn(CopyOptionSet) -> PgWireResult>` that applies defaults and rejects options the format -//! does not accept, -//! 3. register the builder under every accepted `FORMAT` name (including -//! aliases) in `REGISTERED_FORMATS` in the parent module. +//! writes is collected into a [`CopyOptionSet`] (see `options`) +//! regardless of whether any format understands it — the statement is +//! parsed textually (see `parse`) precisely so options beyond +//! sqlparser's COPY grammar can reach the formats. A format builder +//! takes the options it understands (`take_string`, +//! `take_single_byte`, `take_flag`, or raw `take_value` for +//! format-specific shapes) and rejects the rest with +//! `reject_leftovers`. -use datatypes::prelude::ConcreteDataType; -use datatypes::schema::ColumnSchema; use pgwire::error::PgWireResult; +use crate::postgres::copy_in::csv::CsvCodec; +use crate::postgres::copy_in::influxdb::InfluxdbOptions; + +/// A resolved COPY FROM STDIN format with its validated options. +#[derive(Debug, Clone)] +pub(super) enum CopyFormat { + /// Table-bound CSV: records are parsed incrementally against the + /// targeted table's columns and inserted in batches. + Csv(CsvCodec), + /// Table-less influxdb line protocol: the whole payload is one + /// bounded request body, completed through the line-protocol + /// pipeline the HTTP write endpoints use (one COPY statement is one + /// ingestion request; streaming ingest is not supported). + Influxdb(InfluxdbOptions), +} + /// A parsed, insertable row: one gRPC value per targeted column, in COPY /// field order. Every parser emits these directly, whatever mechanism it /// used to produce them. @@ -83,34 +88,3 @@ pub(super) trait RecordParser: Send { /// stream ended exactly at a record boundary. fn finish(&mut self) -> PgWireResult>; } - -/// A COPY FROM STDIN data format implementation. -/// -/// Implementations are shared per statement (hence `Send + Sync`) and must -/// be stateless; per-stream state belongs to the [`RecordParser`] they -/// create. Every method is tailored to the format: `check_column_type` -/// declares the types the format can decode (csv is bounded by what -/// arrow-csv supports). -pub(super) trait CopyInCodec: Send + Sync + std::fmt::Debug { - /// The canonical format name, used in error messages. - fn name(&self) -> &'static str; - /// The format code advertised in the `CopyInResponse` header: 0 when - /// the client sends textual data (the default), 1 for binary. - fn format_code(&self) -> i8 { - 0 - } - /// Validates that a column of `data_type` can be ingested by this - /// format. - fn check_column_type(&self, data_type: &ConcreteDataType) -> PgWireResult<()>; - /// Creates the stateful parser for one copy stream over the targeted - /// columns (in COPY field order). - /// - /// `timezone` is the session timezone (an IANA name or a UTC offset): - /// formats that parse offset-free timestamp strings from text should - /// interpret them in this zone, matching SQL INSERT. - fn create_parser( - &self, - columns: &[ColumnSchema], - timezone: &str, - ) -> PgWireResult>; -} diff --git a/src/servers/src/postgres/copy_in/csv.rs b/src/servers/src/postgres/copy_in/csv.rs index b2919210e3f..ab8f521098d 100644 --- a/src/servers/src/postgres/copy_in/csv.rs +++ b/src/servers/src/postgres/copy_in/csv.rs @@ -40,12 +40,12 @@ use datatypes::schema::{ColumnSchema, Schema as GreptimeSchema}; use datatypes::vectors::{BooleanVector, Helper, StringVector, Vector}; use pgwire::error::{PgWireError, PgWireResult}; -use crate::postgres::copy_in::codec::{CopyInCodec, ParsedRow, RecordParser}; +use crate::postgres::copy_in::codec::{CopyFormat, ParsedRow, RecordParser}; use crate::postgres::copy_in::options::CopyOptionSet; use crate::postgres::copy_in::{COPY_IN_FLUSH_ROWS, bad_copy_data, invalid_text_repr, unsupported}; /// The validated `csv` format configuration. -#[derive(Debug)] +#[derive(Debug, Clone)] pub(super) struct CsvCodec { delimiter: u8, quote: u8, @@ -54,7 +54,7 @@ pub(super) struct CsvCodec { } /// Builds the CSV codec from the parsed options. -pub(super) fn build_csv_codec(mut options: CopyOptionSet) -> PgWireResult> { +pub(super) fn build_csv_codec(mut options: CopyOptionSet) -> PgWireResult { let delimiter = options.take_single_byte("DELIMITER")?.unwrap_or(b','); let quote = options.take_single_byte("QUOTE")?.unwrap_or(b'"'); // In PostgreSQL the default escape character for CSV is the quote @@ -75,7 +75,7 @@ pub(super) fn build_csv_codec(mut options: CopyOptionSet) -> PgWireResult PgWireResult &'static str { - "csv" - } - +impl CsvCodec { /// The column types arrow-csv can decode against. - fn check_column_type(&self, data_type: &ConcreteDataType) -> PgWireResult<()> { + pub(super) fn check_column_type(&self, data_type: &ConcreteDataType) -> PgWireResult<()> { let supported = matches!( data_type, ConcreteDataType::Boolean(_) @@ -117,7 +113,7 @@ impl CopyInCodec for CsvCodec { Ok(()) } - fn create_parser( + pub(super) fn create_parser( &self, columns: &[ColumnSchema], timezone: &str, @@ -292,9 +288,15 @@ mod tests { ] } + fn build_codec(set: CopyOptionSet) -> CsvCodec { + match build_csv_codec(set).unwrap() { + CopyFormat::Csv(codec) => codec, + other => panic!("expected csv format, got {other:?}"), + } + } + fn parser() -> Box { - build_csv_codec(CopyOptionSet::default()) - .unwrap() + build_codec(CopyOptionSet::default()) .create_parser(&test_columns(), "UTC") .unwrap() } @@ -308,8 +310,7 @@ mod tests { fn rows_with(mut set: CopyOptionSet) -> Box { set.take_format().unwrap(); - build_csv_codec(set) - .unwrap() + build_codec(set) .create_parser(&test_columns(), "UTC") .unwrap() } @@ -329,8 +330,8 @@ mod tests { #[test] fn test_build_defaults() { - let codec = build_csv_codec(CopyOptionSet::default()).unwrap(); - assert_eq!(codec.name(), "csv"); + let codec = build_codec(CopyOptionSet::default()); + assert!(matches!(codec, CsvCodec { .. })); } #[test] @@ -368,7 +369,7 @@ mod tests { #[test] fn test_check_column_type() { - let codec = build_csv_codec(CopyOptionSet::default()).unwrap(); + let codec = build_codec(CopyOptionSet::default()); assert!( codec .check_column_type(&ConcreteDataType::json_datatype()) @@ -524,8 +525,7 @@ mod tests { .into_bytes(); let check = |chunk: usize| { - let mut parser = build_csv_codec(CopyOptionSet::default()) - .unwrap() + let mut parser = build_codec(CopyOptionSet::default()) .create_parser(&sweep_columns, "UTC") .unwrap(); let mut ids: Vec = Vec::new(); @@ -560,8 +560,7 @@ mod tests { check(chunk); let mut shifted = payload.clone(); shifted.drain(..37); - let mut parser = build_csv_codec(CopyOptionSet::default()) - .unwrap() + let mut parser = build_codec(CopyOptionSet::default()) .create_parser(&sweep_columns, "UTC") .unwrap(); let mut count = 0usize; @@ -585,8 +584,7 @@ mod tests { false, ), ]; - let mut parser = build_csv_codec(CopyOptionSet::default()) - .unwrap() + let mut parser = build_codec(CopyOptionSet::default()) .create_parser(&columns, "UTC") .unwrap(); let rows = feed_all( @@ -628,8 +626,7 @@ mod tests { false, ), ]; - let mut parser = build_csv_codec(CopyOptionSet::default()) - .unwrap() + let mut parser = build_codec(CopyOptionSet::default()) .create_parser(&columns, "UTC") .unwrap(); parser.feed(b"maybe,2024-01-01 00:00:00\n").unwrap(); @@ -656,8 +653,7 @@ mod tests { false, ), ]; - let mut parser = build_csv_codec(CopyOptionSet::default()) - .unwrap() + let mut parser = build_codec(CopyOptionSet::default()) .create_parser(&columns, "Asia/Shanghai") .unwrap(); let rows = feed_all( diff --git a/src/servers/src/postgres/copy_in/influxdb.rs b/src/servers/src/postgres/copy_in/influxdb.rs new file mode 100644 index 00000000000..b4e2a9add30 --- /dev/null +++ b/src/servers/src/postgres/copy_in/influxdb.rs @@ -0,0 +1,95 @@ +// 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. + +//! The `influxdb` line protocol format for COPY FROM STDIN. +//! +//! `COPY FROM STDIN WITH (FORMAT influxdb, precision 'ns')` ingests +//! [InfluxDB line protocol](https://docs.influxdata.com/influxdb/v1.8/write_protocols/line_protocol_reference/) +//! text. The format is table-less — each line names its own measurement +//! — and does not parse here at all: the COPY layer buffers the whole +//! payload (bounded by [`crate::postgres::copy_in::COPY_IN_MAX_PAYLOAD_BYTES`]) +//! and completes the copy by handing it to the frontend's line-protocol +//! write pipeline as one body, exactly like the HTTP +//! `/v1/influxdb/v1/write` endpoint receives one request body: +//! conversion, permissions, interceptors, timestamp alignment, merge +//! mode and inserts are the same path, and record boundaries (including +//! newlines inside quoted field values and after escapes) are decided +//! by the line-protocol parser alone — the COPY layer never splits +//! records. +//! +//! One COPY statement is therefore one bounded ingestion request; a +//! larger import is multiple COPY statements. +//! +//! The `precision` option carries the same values as the HTTP +//! `precision` query parameter (`ns`, `us`, `ms`, `s`, ...), defaults to +//! nanoseconds, and is resolved by the shared `parse_time_precision`. + +use common_grpc::precision::Precision; +use pgwire::error::PgWireResult; + +use crate::postgres::copy_in::bad_copy_data; +use crate::postgres::copy_in::codec::CopyFormat; +use crate::postgres::copy_in::options::CopyOptionSet; + +/// The format's validated options. A `None` precision carries the +/// pipeline's default (nanoseconds), exactly like an HTTP write without +/// a `precision` parameter. +#[derive(Debug, Clone)] +pub(super) struct InfluxdbOptions { + pub(super) precision: Option, +} + +pub(super) fn build_influxdb_codec(mut options: CopyOptionSet) -> PgWireResult { + let precision = match options.take_string("precision")? { + Some(value) => Some( + crate::influxdb::parse_time_precision(&value) + .map_err(|e| bad_copy_data(format!("invalid precision option: {e}")))?, + ), + None => None, + }; + options.reject_leftovers("influxdb")?; + Ok(CopyFormat::Influxdb(InfluxdbOptions { precision })) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::postgres::copy_in::options::CopyOptionValue; + + #[test] + fn test_builder_options() { + let format = build_influxdb_codec(CopyOptionSet::default()).unwrap(); + assert!(matches!(format, CopyFormat::Influxdb(_))); + assert!(matches!(format, CopyFormat::Influxdb(_))); + + let mut set = CopyOptionSet::default(); + set.insert("precision".into(), CopyOptionValue::String("ms".into())) + .unwrap(); + assert!(matches!( + build_influxdb_codec(set).unwrap(), + CopyFormat::Influxdb(InfluxdbOptions { + precision: Some(Precision::Millisecond) + }) + )); + + let mut set = CopyOptionSet::default(); + set.insert("precision".into(), CopyOptionValue::String("bogus".into())) + .unwrap(); + assert!(build_influxdb_codec(set).is_err()); + + let mut set = CopyOptionSet::default(); + set.insert("unknown".into(), CopyOptionValue::Flag).unwrap(); + assert!(build_influxdb_codec(set).is_err()); + } +} diff --git a/src/servers/src/postgres/copy_in/mod.rs b/src/servers/src/postgres/copy_in/mod.rs index a3959a23830..92528d4ab2f 100644 --- a/src/servers/src/postgres/copy_in/mod.rs +++ b/src/servers/src/postgres/copy_in/mod.rs @@ -21,38 +21,42 @@ //! the data stream incrementally and flushes it into the query handler as //! [`RowInsertRequests`] batches. //! -//! # Format framework +//! # Formats //! -//! Each COPY data format is implemented as a [`codec::CopyInCodec`]: a -//! stateless strategy that validates its options at statement-parse time, -//! describes the column types it can ingest, and creates the stateful -//! record parser for one stream. Formats register a builder under every -//! accepted `FORMAT` name in [`REGISTERED_FORMATS`] — `csv` is the only -//! implemented format today; PostgreSQL's `text`/`binary` and -//! GreptimeDB-specific formats alike can be added the same way. See the -//! `codec` module docs. +//! 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 parse; use std::fmt::Debug; use std::sync::Arc; +use crate::influxdb::InfluxdbRequest; 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::{CopyInCodec, RecordParser}; +use codec::{CopyFormat, RecordParser}; +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 pgwire::api::copy::CopyHandler; use pgwire::api::results::{CopyResponse, Response, Tag}; use pgwire::api::{ClientInfo, PgWireConnectionState}; @@ -80,38 +84,43 @@ use crate::postgres::utils::convert_err; /// the per-batch fixed costs negligible. pub(crate) const COPY_IN_FLUSH_ROWS: usize = 32768; -/// Builds a codec from the format-neutral option set. -type CodecBuilder = fn(options::CopyOptionSet) -> PgWireResult>; +/// Upper bound of an influxdb COPY FROM STDIN payload. One COPY +/// statement is one bounded ingestion request (the whole payload is +/// buffered like one HTTP write body); a larger import is multiple COPY +/// statements. Streaming ingest may arrive later, with framing owned by +/// the shared line-protocol parser rather than the COPY layer. +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; /// 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, CodecBuilder)] = &[ +const REGISTERED_FORMATS: &[(&str, FormatBuilder)] = &[ ("csv", csv::build_csv_codec), + ("influxdb", influxdb::build_influxdb_codec), + ("influxdb_line", influxdb::build_influxdb_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> { +fn build_unimplemented_text_codec(_options: options::CopyOptionSet) -> PgWireResult { Err(unsupported( "the text format is not supported yet; use WITH (FORMAT csv)", )) } -fn build_unimplemented_binary_codec( - _options: options::CopyOptionSet, -) -> PgWireResult> { +fn build_unimplemented_binary_codec(_options: options::CopyOptionSet) -> PgWireResult { Err(unsupported( "COPY FROM STDIN WITH (FORMAT binary) is not supported yet", )) } -/// Resolves the format codec from the parsed options. -fn resolve_codec(options: &mut options::CopyOptionSet) -> PgWireResult> { +/// Resolves the format from the parsed options. +fn resolve_format(options: &mut options::CopyOptionSet) -> PgWireResult { let name = options.take_format()?.ok_or_else(|| { unsupported( "COPY FROM STDIN requires an explicit format in this version; use WITH (FORMAT csv)", @@ -133,15 +142,15 @@ pub(crate) struct CopyFromStdin { table: Vec, /// Requested column list; empty means all columns in table order. columns: Vec, - /// The resolved format codec, validated and ready to run. - codec: Arc, + /// The resolved format with its validated options. + format: CopyFormat, } fn copy_in_error(code: PgErrorCode, message: impl Into) -> PgWireError { PgWireError::UserError(Box::new(code.to_err_info(message.into()))) } -fn unsupported(message: impl Into) -> PgWireError { +pub(super) fn unsupported(message: impl Into) -> PgWireError { copy_in_error(PgErrorCode::Ec0A000, message) } @@ -174,11 +183,11 @@ pub(crate) fn copy_from_stdin_from_sql(sql: &str) -> PgWireResult raw, }; let mut option_set = raw.options; - let codec = resolve_codec(&mut option_set)?; + let format = resolve_format(&mut option_set)?; Ok(Some(CopyFromStdin { table: raw.table, columns: raw.columns, - codec, + format, })) } @@ -221,26 +230,49 @@ fn resolve_table_name( } /// 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, - /// The gRPC column schemas of the targeted columns. Cloned into each - /// flushed batch: `Rows.schema` is an owned `Vec` in the proto - /// request, so one copy per request is inherent. - grpc_schema: Vec, - table_name: String, - parser: Box, + 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, - /// Parsed rows not yet flushed to the handler. - rows: Vec, 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, + table_name: String, + parser: Box, + /// Parsed rows not yet flushed to the handler. + rows: Vec, + }, + /// 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, + precision: Option, + }, +} + impl CopyInState { fn new( - codec: Arc, + codec: &CsvCodec, columns: Vec, table_info: Arc, query_ctx: QueryContextRef, @@ -249,7 +281,7 @@ impl CopyInState { let schema = table_info.meta.schema.clone(); let timestamp_index = schema.timestamp_index(); - let columns = resolve_columns(codec.as_ref(), columns, &schema, &table_info)?; + 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) { @@ -268,58 +300,122 @@ impl CopyInState { Ok(Self { query_ctx, - grpc_schema, - table_name, - parser, - codec_name: codec.name(), - codec_format_code: codec.format_code(), - rows: Vec::with_capacity(COPY_IN_FLUSH_ROWS), + 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, + } + } + /// 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. + /// Feeds a `CopyData` chunk into the parser and buffers parsed rows + /// (csv) or payload bytes (influxdb). fn feed(&mut self, data: &[u8]) -> PgWireResult<()> { - self.rows.extend(self.parser.feed(data)?); + match &mut self.payload { + Payload::Table { parser, rows, .. } => rows.extend(parser.feed(data)?), + Payload::Influxdb { 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<()> { - self.rows.extend(self.parser.finish()?); + match &mut self.payload { + Payload::Table { parser, rows, .. } => rows.extend(parser.finish()?), + // The whole payload is one request body; nothing to finalize. + Payload::Influxdb { .. } => {} + } Ok(()) } fn should_flush(&self) -> bool { - self.rows.len() >= COPY_IN_FLUSH_ROWS + match &self.payload { + Payload::Table { rows, .. } => rows.len() >= COPY_IN_FLUSH_ROWS, + // Completed once at CopyDone, not mid-stream. + Payload::Influxdb { .. } => false, + } } /// Takes the buffered rows as an insert request, if any. fn take_requests(&mut self) -> Option { - if self.rows.is_empty() { + let Payload::Table { + grpc_schema, + table_name, + rows, + .. + } = &mut self.payload + else { + return None; + }; + if rows.is_empty() { return None; } let rows = GrpcRows { - schema: self.grpc_schema.clone(), - rows: std::mem::take(&mut self.rows), + schema: grpc_schema.clone(), + rows: std::mem::take(rows), }; Some(RowInsertRequests { inserts: vec![RowInsertRequest { - table_name: self.table_name.clone(), + table_name: table_name.clone(), rows: Some(rows), }], }) } + + /// 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> { + 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: &dyn CopyInCodec, + codec: &CsvCodec, columns: Vec, schema: &SchemaRef, table_info: &Arc, @@ -411,27 +507,63 @@ impl PostgresServerHandlerInner { // pending-rows batcher like SQL INSERT when it is enabled for // this connection. let query_ctx = self.new_query_context(); - 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(".")), - )); + // 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) + } + }; + let num_columns = match &state.payload { + Payload::Table { grpc_schema, .. } => grpc_schema.len(), + // No fixed column set: the format routes per record. + Payload::Influxdb { .. } => 0, }; - - // 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); - - let state = CopyInState::new(stmt.codec, stmt.columns, table_info, query_ctx, table)?; - let num_columns = state.grpc_schema.len(); let format_code = state.format_code(); *self.copy_in_state.lock().await = Some(state); @@ -443,13 +575,34 @@ impl PostgresServerHandlerInner { ))) } - /// Flushes buffered rows of an in-progress copy. Returns the number of - /// rows written by the flushed batch. + /// 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 { + 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(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 query_ctx = state.query_ctx.clone(); let output = self .copy_in_handler .copy_in_insert(requests, query_ctx) @@ -587,8 +740,8 @@ mod tests { .unwrap(); assert_eq!(copy.table, vec!["t"]); assert!(copy.columns.is_empty()); - assert_eq!(copy.codec.name(), "csv"); - assert_eq!(copy.codec.format_code(), 0); + assert!(matches!(copy.format, CopyFormat::Csv(_))); + // Both formats are textual: format code 0. } #[test] @@ -599,7 +752,7 @@ mod tests { .unwrap() .unwrap(); assert_eq!(copy.columns, vec!["a", "b"]); - assert_eq!(copy.codec.name(), "csv"); + assert!(matches!(copy.format, CopyFormat::Csv(_))); } #[test] diff --git a/src/servers/src/postgres/copy_in/parse.rs b/src/servers/src/postgres/copy_in/parse.rs index 8139960bfa4..e9be443eea5 100644 --- a/src/servers/src/postgres/copy_in/parse.rs +++ b/src/servers/src/postgres/copy_in/parse.rs @@ -32,7 +32,7 @@ use pgwire::error::PgWireResult; use crate::postgres::copy_in::options::{CopyOptionSet, CopyOptionValue}; use crate::postgres::copy_in::{bad_copy_data, copy_in_must_run_alone_error}; -/// A scanned `COPY tbl [(cols)] FROM STDIN [WITH (...)]` statement. +/// A scanned `COPY [tbl [(cols)]] FROM STDIN [WITH (...)]` statement. /// /// Unquoted identifiers are folded to lowercase and quoted ones keep their /// case (`COPY Metrics` addresses the same table as `CREATE TABLE Metrics`, @@ -41,7 +41,9 @@ use crate::postgres::copy_in::{bad_copy_data, copy_in_must_run_alone_error}; #[derive(Debug)] pub(super) struct RawCopyFromStdin { /// Table name parts, as written (`t`, `schema.t`, - /// `catalog.schema.t`); quoted names keep their case. + /// `catalog.schema.t`); quoted names keep their case. Empty for + /// table-less formats whose records carry their own routing (e.g. + /// the influxdb line protocol names its measurement per line). pub(super) table: Vec, /// Requested column list; empty means all columns in table order. pub(super) columns: Vec, @@ -89,13 +91,27 @@ pub(super) fn scan_copy_from_stdin(sql: &str) -> PgWireResult { return Ok(CopyInScan::NotCopyIn); } - let table = match parse_object_name(&mut cursor) { - Some(table) => table, - None => return Ok(CopyInScan::NotCopyIn), + // Table-less statements (`COPY FROM STDIN WITH (FORMAT influxdb) ...`) + // are table-format specific: the codec decides whether a missing + // table is legal. An unquoted `FROM` cannot be a table name (SQL + // would require quoting it), so peek instead of letting the + // identifier parser swallow the keyword. A column list without a + // table is not. + let table = if cursor.peek_keyword("from") { + Vec::new() + } else { + match parse_object_name(&mut cursor) { + Some(table) => table, + None => return Ok(CopyInScan::NotCopyIn), + } }; - let columns = match parse_column_list(&mut cursor) { - Some(columns) => columns, - None => return Ok(CopyInScan::NotCopyIn), + let columns = if table.is_empty() { + Vec::new() + } else { + match parse_column_list(&mut cursor) { + Some(columns) => columns, + None => return Ok(CopyInScan::NotCopyIn), + } }; if !cursor.take_keyword("from") || !cursor.take_keyword("stdin") { @@ -151,16 +167,20 @@ impl<'a> Tokens<'a> { /// Consumes the next token if it is the unquoted keyword `keyword`. fn take_keyword(&mut self, keyword: &str) -> bool { - match self.peek() { - Some(Token::Word(w)) - if w.quote_style.is_none() && w.value.eq_ignore_ascii_case(keyword) => - { - self.pos += 1; - true - } - _ => false, + if self.peek_keyword(keyword) { + self.pos += 1; + true + } else { + false } } + + /// Whether the next token is the unquoted keyword `keyword`, without + /// consuming it. + fn peek_keyword(&self, keyword: &str) -> bool { + matches!(self.tokens.get(self.pos), Some(Token::Word(w)) + if w.quote_style.is_none() && w.value.eq_ignore_ascii_case(keyword)) + } } /// Returns the identifier with the SQL parser's canonicalization: quoted @@ -275,6 +295,27 @@ mod tests { }) } + #[test] + fn test_scan_tableless() { + // Table-less statements: the codec decides whether a missing + // table is legal (the influxdb line protocol names its + // measurement per line). + let raw = scan("COPY FROM STDIN WITH (FORMAT influxdb, precision 'ns')") + .unwrap() + .unwrap(); + assert!(raw.table.is_empty()); + assert!(raw.columns.is_empty()); + + let raw = scan("copy from stdin (format influxdb)").unwrap().unwrap(); + assert!(raw.table.is_empty()); + + // A column list without a table is not a copy-in statement. + assert!(scan("COPY (a, b) FROM STDIN").unwrap().is_none()); + // COPY OUT and file COPY are unaffected. + assert!(scan("COPY (SELECT 1) TO STDOUT").unwrap().is_none()); + assert!(scan("COPY t FROM '/tmp/x.csv'").unwrap().is_none()); + } + #[test] fn test_scan_basic() { let raw = scan("COPY t FROM STDIN").unwrap().unwrap(); diff --git a/src/servers/src/postgres/server.rs b/src/servers/src/postgres/server.rs index 7e3b3e2b7bd..3919081ba62 100644 --- a/src/servers/src/postgres/server.rs +++ b/src/servers/src/postgres/server.rs @@ -30,6 +30,7 @@ use tokio_rustls::TlsAcceptor; use crate::error::Result; use crate::postgres::{MakePostgresServerHandler, MakePostgresServerHandlerBuilder}; use crate::query_handler::CopyInHandlerRef; +use crate::query_handler::InfluxdbLineProtocolHandlerRef; use crate::query_handler::sql::ServerSqlQueryHandlerRef; use crate::server::{AbortableStream, BaseTcpServer, Server}; use crate::tls::ReloadableTlsServerConfig; @@ -40,6 +41,7 @@ pub struct PostgresServer { tls_server_config: Arc, keep_alive_secs: u64, batching_enabled: bool, + influxdb_handler: Option, bind_addr: Option, process_manager: Option, } @@ -77,6 +79,7 @@ impl PostgresServer { tls_server_config, keep_alive_secs, batching_enabled: false, + influxdb_handler: None, bind_addr: None, process_manager, } @@ -88,12 +91,34 @@ impl PostgresServer { self } + /// Wires the influxdb line-protocol handler used by table-less + /// `COPY FROM STDIN WITH (FORMAT influxdb)` flushes — the same + /// handler the HTTP influxdb write endpoints use, so a flushed batch + /// runs through the identical pipeline. + pub fn with_influxdb_handler(mut self, handler: InfluxdbLineProtocolHandlerRef) -> Self { + self.influxdb_handler = Some(handler); + self + } + fn accept( &self, io_runtime: Runtime, accepting_stream: AbortableStream, ) -> impl Future + use<> { - let handler_maker = self.make_handler.clone(); + // Assemble the per-connection handler directly, including any + // `with_*` wiring applied after construction (the influxdb + // line-protocol handler). Constructing the struct — instead of + // going through the builder — keeps this infallible: no + // `.build().unwrap()` on the hot accept path. + let make_handler = &self.make_handler; + let handler_maker = Arc::new(MakePostgresServerHandler { + query_handler: make_handler.query_handler.clone(), + copy_in_handler: make_handler.copy_in_handler.clone(), + influxdb_handler: self.influxdb_handler.clone(), + user_provider: make_handler.user_provider.clone(), + param_provider: make_handler.param_provider.clone(), + force_tls: make_handler.force_tls, + }); let batching_enabled = self.batching_enabled; let tls_server_config = self.tls_server_config.clone(); let process_manager = self.process_manager.clone(); diff --git a/tests-integration/src/test_util.rs b/tests-integration/src/test_util.rs index 3cc820619c6..3dd451b4716 100644 --- a/tests-integration/src/test_util.rs +++ b/tests-integration/src/test_util.rs @@ -1204,16 +1204,19 @@ pub async fn setup_pg_server_with_slow_query_threshold( .expect("Failed to load certificates and keys"), ); - let mut pg_server = Box::new(PostgresServer::new( - fe_instance_ref.clone(), - fe_instance_ref.clone(), - opts.tls.should_force_tls(), - tls_server_config, - 0, - runtime, - None, - None, - )); + let mut pg_server = Box::new( + PostgresServer::new( + fe_instance_ref.clone(), + fe_instance_ref.clone(), + opts.tls.should_force_tls(), + tls_server_config, + 0, + runtime, + None, + None, + ) + .with_influxdb_handler(fe_instance_ref.clone()), + ); pg_server .start(fe_pg_addr.parse::().unwrap()) @@ -1248,16 +1251,19 @@ pub async fn setup_pg_server_with_user_provider( .expect("Failed to load certificates and keys"), ); - let mut pg_server = Box::new(PostgresServer::new( - fe_instance_ref.clone(), - fe_instance_ref.clone(), - opts.tls.should_force_tls(), - tls_server_config, - 0, - runtime, - user_provider, - None, - )); + let mut pg_server = Box::new( + PostgresServer::new( + fe_instance_ref.clone(), + fe_instance_ref.clone(), + opts.tls.should_force_tls(), + tls_server_config, + 0, + runtime, + user_provider, + None, + ) + .with_influxdb_handler(fe_instance_ref.clone()), + ); pg_server .start(fe_pg_addr.parse::().unwrap()) @@ -1304,16 +1310,19 @@ pub async fn setup_pg_server_with_prom_native_histogram( .expect("Failed to load certificates and keys"), ); - let mut pg_server = Box::new(PostgresServer::new( - instance.fe_instance().clone(), - instance.fe_instance().clone(), - opts.tls.should_force_tls(), - tls_server_config, - 0, - runtime, - None, - None, - )); + let mut pg_server = Box::new( + PostgresServer::new( + instance.fe_instance().clone(), + instance.fe_instance().clone(), + opts.tls.should_force_tls(), + tls_server_config, + 0, + runtime, + None, + None, + ) + .with_influxdb_handler(instance.fe_instance().clone()), + ); pg_server .start(fe_pg_addr.parse::().unwrap()) diff --git a/tests-integration/src/tests/postgres_copy_in_test.rs b/tests-integration/src/tests/postgres_copy_in_test.rs index 82e076e7a65..94b19c1315d 100644 --- a/tests-integration/src/tests/postgres_copy_in_test.rs +++ b/tests-integration/src/tests/postgres_copy_in_test.rs @@ -316,6 +316,7 @@ async fn test_postgres_copy_in_failed_copy_then_new_copy() { Err(e) => e, }; assert!(err.as_db_error().is_some(), "{err}"); + tokio::time::sleep(std::time::Duration::from_millis(200)).await; // The connection is usable and a NEW copy succeeds on it. let rows = feed_copy_in( @@ -519,3 +520,302 @@ async fn test_postgres_copy_in_empty_field_is_null() { server.shutdown().await.unwrap(); } + +/// `COPY FROM STDIN WITH (FORMAT influxdb)` ingests line protocol through +/// the same pipeline as the HTTP influxdb write endpoint: lines route to +/// their own measurements, tables are auto-created, and a batch behaves +/// like an HTTP write of the same bytes. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_in_influxdb_line_protocol() { + init_default_ut_logging(); + let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_influxdb").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + let payload = b"monitor,host=hosta cpu=66.6,memory=1024 1663840496100023100\n\ + monitor,host=hostb memory=1027 1663840496400340001\n\ + weather,location=us-midwest temperature=82 1663840496100023102\n"; + + // Feed byte-by-byte-ish: lines and UTF-8 may span CopyData chunks. + let rows = feed_copy_in( + &client, + "COPY FROM STDIN WITH (FORMAT influxdb)", + payload.chunks(7).map(Bytes::copy_from_slice).collect(), + ) + .await + .unwrap(); + assert_eq!(rows, 3); + + let messages = client + .simple_query("SELECT host, cpu, memory FROM monitor ORDER BY host") + .await + .unwrap(); + let rows: Vec<(String, String, String)> = messages + .iter() + .filter_map(|m| { + let tokio_postgres::SimpleQueryMessage::Row(row) = m else { + return None; + }; + Some(( + row.get(0).unwrap_or("").to_string(), + row.get(1).unwrap_or("").to_string(), + row.get(2).unwrap_or("").to_string(), + )) + }) + .collect(); + assert_eq!( + rows, + vec![ + ("hosta".into(), "66.6".into(), "1024.0".into()), + ("hostb".into(), "".into(), "1027.0".into()), + ] + ); + + let messages = client + .simple_query("SELECT temperature FROM weather") + .await + .unwrap(); + let temps: Vec = messages + .iter() + .filter_map(|m| match m { + tokio_postgres::SimpleQueryMessage::Row(row) => row.get(0).map(str::to_string), + _ => None, + }) + .collect(); + assert_eq!(temps, vec!["82.0".to_string()]); + + server.shutdown().await.unwrap(); +} + +/// The `precision` option carries the HTTP endpoint's semantics: +/// nanosecond timestamps by default, other units on request. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_in_influxdb_precision() { + init_default_ut_logging(); + let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_influxdb_prec").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + let rows = feed_copy_in( + &client, + "COPY FROM STDIN WITH (FORMAT influxdb, precision 'ms')", + vec![Bytes::from_static(b"prec,host=h v=1 1465839830100\n")], + ) + .await + .unwrap(); + assert_eq!(rows, 1); + + let messages = client + .simple_query("SELECT CAST(greptime_timestamp AS BIGINT) FROM prec") + .await + .unwrap(); + let ts: Vec = messages + .iter() + .filter_map(|m| match m { + tokio_postgres::SimpleQueryMessage::Row(row) => row.get(0).map(str::to_string), + _ => None, + }) + .collect(); + // 1465839830100 ms == 1465839830100000000 ns + assert_eq!(ts.len(), 1); + let ts_ns: i64 = ts[0].parse().unwrap(); + assert_eq!(ts_ns, 1465839830100000000); + + server.shutdown().await.unwrap(); +} + +/// Records with newlines INSIDE quoted field values are one line-protocol +/// record: the parser owns record boundaries, so a CopyData chunk boundary +/// inside the quoted string — the case a naive line splitter breaks on — +/// must not split it. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_in_influxdb_multiline_quoted_field() { + init_default_ut_logging(); + let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_influxdb_ml").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + // A record whose string field value contains a physical newline, and + // one with a backslash right before the newline (the parser skips the + // escaped newline; also not a record boundary). + let payload: Vec = b"mlog,host=a line=\"first\nsecond\" 100\n\ + mlog,host=b line=\"tail\\\nmore\" 200\n" + .to_vec(); + + // Split exactly on the internal newline of the first record: chunk 1 + // ends after `first`, chunk 2 starts with the newline. + let split_at = payload.windows(6).position(|w| w == b"first\n").unwrap() + b"first".len(); + let (head, tail) = payload.split_at(split_at); + let rows = feed_copy_in( + &client, + "COPY FROM STDIN WITH (FORMAT influxdb)", + vec![Bytes::copy_from_slice(head), Bytes::copy_from_slice(tail)], + ) + .await + .unwrap(); + assert_eq!(rows, 2); + + let messages = client + .simple_query("SELECT line FROM mlog ORDER BY host") + .await + .unwrap(); + let lines: Vec = messages + .iter() + .filter_map(|m| match m { + tokio_postgres::SimpleQueryMessage::Row(row) => row.get(0).map(str::to_string), + _ => None, + }) + .collect(); + // The quoted values keep their physical newlines: `first`, newline, + // `second`; and `tail`, backslash, newline, `more` (a backslash only + // escapes `\` and `"` in line-protocol strings, so it stays literal + // before the newline). + assert_eq!( + lines, + vec!["first\nsecond".to_string(), "tail\\\nmore".to_string()] + ); + + server.shutdown().await.unwrap(); +} + +/// The payload is bounded by the HTTP write endpoints' own body limit. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_in_influxdb_payload_limit() { + init_default_ut_logging(); + let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_influxdb_cap").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + // The HTTP write endpoints' default body limit (64 MiB) — kept in + // sync with servers::http::DEFAULT_BODY_LIMIT. + let limit = 64 * 1024 * 1024; + // Several chunks over the limit, none containing a record boundary. + let chunk = vec![b'x'; 1024 * 1024]; + let mut sent = 0usize; + let mut sink = Box::pin( + client + .copy_in("COPY FROM STDIN WITH (FORMAT influxdb)") + .await + .unwrap(), + ); + while sent <= limit { + sink.as_mut() + .send(Bytes::from(chunk.clone())) + .await + .unwrap(); + sent += chunk.len(); + } + // The cap fires mid-stream; the error surfaces when the copy + // completes (the sink buffers sends locally). + let error = sink + .as_mut() + .finish() + .await + .expect_err("expected the oversized payload to fail"); + assert!( + error + .as_db_error() + .map(|d| d.message().contains("exceeds")) + .unwrap_or(false), + "unexpected error: {error}" + ); + + // The connection stays usable. + let messages = client.simple_query("SELECT 1").await.unwrap(); + assert!(!messages.is_empty()); + + server.shutdown().await.unwrap(); +} + +/// The regression for the retired line-splitter's failure mode: with +/// 32,767 ordinary records plus a multiline quoted record whose first +/// physical line lands exactly on the old 32,768-line flush threshold — +/// and a CopyData chunk boundary at the quoted value's internal newline — +/// the whole payload must still ingest as one request. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_in_influxdb_flush_threshold_multiline() { + init_default_ut_logging(); + let (_guard, server) = + setup_pg_server(StorageType::File, "pg_copy_in_influxdb_threshold").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + // 32,767 ordinary records, then a record whose quoted value spans + // three physical lines. + let mut head = String::new(); + for i in 0..32_767 { + head.push_str(&format!( + "m,host=h{i} v={i} {} +", + 1_000 + i + )); + } + head.push_str("mlog,host=z line=\"first\nsecond\nthird\" 99999"); + let tail = " +" + .to_string(); + + let rows = feed_copy_in( + &client, + "COPY FROM STDIN WITH (FORMAT influxdb)", + vec![ + Bytes::from(head.into_bytes()), + Bytes::from(tail.into_bytes()), + ], + ) + .await + .unwrap(); + assert_eq!(rows, 32_768); + + let messages = client.simple_query("SELECT line FROM mlog").await.unwrap(); + let lines: Vec = messages + .iter() + .filter_map(|m| match m { + tokio_postgres::SimpleQueryMessage::Row(row) => row.get(0).map(str::to_string), + _ => None, + }) + .collect(); + assert_eq!(lines, vec!["first\nsecond\nthird".to_string()]); + + server.shutdown().await.unwrap(); +} + +/// Invalid UTF-8 fails the copy with a clear error and the connection +/// stays usable; a trailing record without a newline is still ingested. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_in_influxdb_invalid_utf8_and_trailing_record() { + init_default_ut_logging(); + let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_influxdb_utf8").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + // Trailing record without a newline: accepted. + let rows = feed_copy_in( + &client, + "COPY FROM STDIN WITH (FORMAT influxdb)", + vec![Bytes::from_static(b"m2,host=a v=7 300")], + ) + .await + .unwrap(); + assert_eq!(rows, 1); + + // Invalid UTF-8: the copy fails, the connection stays usable. + let error = feed_copy_in( + &client, + "COPY FROM STDIN WITH (FORMAT influxdb)", + vec![Bytes::from_static(b"m3 v=1 100\n\xff\xfe")], + ) + .await + .unwrap_err(); + assert!( + error + .as_db_error() + .map(|d| d.message().contains("not valid UTF-8")) + .unwrap_or(false), + "unexpected error: {error}" + ); + let messages = client.simple_query("SELECT 1").await.unwrap(); + assert!(!messages.is_empty()); + + server.shutdown().await.unwrap(); +}