rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-11-15 07:17:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(postgres): influxdb line protocol ingest via COPY FROM STDIN (#9982)

Browse Source 
* feat(postgres): influxdb line protocol ingest via COPY FROM STDIN

COPY FROM STDIN WITH (FORMAT influxdb [, precision 'ns']) ingests
InfluxDB line protocol over the postgres wire protocol, reusing the
HTTP ingestion pipeline wholesale for correctness and performance
parity: the whole payload is buffered and completed once at CopyDone
through the frontend's InfluxdbLineProtocolHandler::exec — conversion,
permissions, line-protocol interceptors, timestamp alignment to table
time-index units, merge mode and inserts are the same path an HTTP
write of the same bytes takes, under the connection's regular postgres
query context.

Formats are an explicit, closed enum (CopyFormat::Csv |
CopyFormat::Influxdb(options)); each variant states its own rules in
the copy lifecycle's match arms — csv requires a target table and
parses incrementally, influxdb forbids one (each line names its own
measurement) and buffers one bounded request body. The COPY layer owns
only the buffering: record boundaries, including newlines inside quoted
field values and after escapes, belong to the line-protocol parser
alone — one COPY statement is one bounded ingestion request (64 MiB;
larger imports are multiple COPY statements).

- Wiring: PostgresServer::with_influxdb_handler attaches the shared
  line-protocol handler; unwired servers reject the format when a copy
  starts.
- Grammar: COPY FROM STDIN without a table is scanned (an unquoted
  FROM cannot be a table name); table-bound formats reject it with
  guidance, influxdb rejects a named table.
- Options: precision accepts the HTTP endpoint's values via the shared
  parse_time_precision (moved to servers::influxdb), defaulting to
  nanoseconds.

Review findings (fengjiachun) are fixed at the root by removing the
pre-splitter rather than repairing it: the retired naive line splitter
broke quoted multiline records at chunk/flush boundaries (found by
removing record splitting entirely), drained the buffer quadratically
(found by append-only buffering), and capped nothing (found by an
explicit payload limit).

e2e: multi-measurement ingest with chunked feeds; a 32,768-physical-line
flush-threshold regression with a multiline quoted record split across
CopyData chunks; precision 'ms' epoch verification; invalid UTF-8 and
trailing-record-without-newline handling; payload limit; table-named
rejection; malformed-line abort with connection recovery.

Signed-off-by: Ning Sun <sunning@greptime.com>

* chore: typo

---------

Signed-off-by: Ning Sun <sunning@greptime.com>
This commit is contained in:
Ning Sun
2026-10-28 08:00:41 +00:00
committed by GitHub
parent 4135e65a76
commit c818b6ac2a
12 changed files with 851 additions and 245 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/frontend/src/server.rs
+1
View File
@@ -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<dyn Server>;
src/servers/src/http/influxdb.rs
+3 -20
View File
@@ -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> {
// 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() {
src/servers/src/influxdb.rs
+22 -1
View File
@@ -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>) -> 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> {
// 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;
src/servers/src/postgres.rs
+9 -1
View File
@@ -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<Arc<PgUserProviderSource>>,
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<InfluxdbLineProtocolHandlerRef>,
force_tls: bool,
auth_driver: AuthDriver<Arc<GreptimeDBStartupParameters>, 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<InfluxdbLineProtocolHandlerRef>,
user_provider: Option<UserProviderRef>,
#[builder(default = "Arc::new(GreptimeDBStartupParameters::new())")]
param_provider: Arc<GreptimeDBStartupParameters>,
@@ -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)),
src/servers/src/postgres/copy_in/codec.rs
+46 -72
View File
@@ -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<Arc<dyn
//! CopyInCodec>>` 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<Vec<ParsedRow>>;
}
/// 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<Box<dyn RecordParser + Send>>;
}
src/servers/src/postgres/copy_in/csv.rs
+24 -28
View File
@@ -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<Arc<dyn CopyInCodec>> {
pub(super) fn build_csv_codec(mut options: CopyOptionSet) -> PgWireResult<CopyFormat> {
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<Arc<dy
if delimiter == quote {
return Err(bad_copy_data("DELIMITER must not be equal to QUOTE"));
}
Ok(Arc::new(CsvCodec {
Ok(CopyFormat::Csv(CsvCodec {
delimiter,
quote,
escape,
@@ -83,13 +83,9 @@ pub(super) fn build_csv_codec(mut options: CopyOptionSet) -> PgWireResult<Arc<dy
}))
}
impl CopyInCodec for CsvCodec {
fn name(&self) -> &'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<dyn RecordParser + Send> {
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<dyn RecordParser + Send> {
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<i32> = 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(
src/servers/src/postgres/copy_in/influxdb.rs
+95
View File
@@ -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<Precision>,
}
pub(super) fn build_influxdb_codec(mut options: CopyOptionSet) -> PgWireResult<CopyFormat> {
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());
}
}
src/servers/src/postgres/copy_in/mod.rs
+229 -76
View File
@@ -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<Arc<dyn CopyInCodec>>;
/// 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<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, 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<Arc<dyn CopyInCodec>> {
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<Arc<dyn CopyInCodec>> {
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 codec from the parsed options.
fn resolve_codec(options: &mut options::CopyOptionSet) -> PgWireResult<Arc<dyn CopyInCodec>> {
/// 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)",
@@ -133,15 +142,15 @@ pub(crate) struct CopyFromStdin {
table: Vec<String>,
/// Requested column list; empty means all columns in table order.
columns: Vec<String>,
/// The resolved format codec, validated and ready to run.
codec: Arc<dyn CopyInCodec>,
/// 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())))
}
fn unsupported(message: impl Into<String>) -> PgWireError {
pub(super) fn unsupported(message: impl Into<String>) -> PgWireError {
copy_in_error(PgErrorCode::Ec0A000, message)
}
@@ -174,11 +183,11 @@ pub(crate) fn copy_from_stdin_from_sql(sql: &str) -> PgWireResult<Option<CopyFro
parse::CopyInScan::CopyIn(raw) => 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<api::v1::ColumnSchema>,
table_name: String,
parser: Box<dyn RecordParser + Send>,
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<GrpcRow>,
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>,
},
}
impl CopyInState {
fn new(
codec: Arc<dyn CopyInCodec>,
codec: &CsvCodec,
columns: Vec<String>,
table_info: Arc<TableInfo>,
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<RowInsertRequests> {
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<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: &dyn CopyInCodec,
codec: &CsvCodec,
columns: Vec<String>,
schema: &SchemaRef,
table_info: &Arc<TableInfo>,
@@ -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<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(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]
src/servers/src/postgres/copy_in/parse.rs
+57 -16
View File
@@ -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<String>,
/// Requested column list; empty means all columns in table order.
pub(super) columns: Vec<String>,
@@ -89,13 +91,27 @@ pub(super) fn scan_copy_from_stdin(sql: &str) -> PgWireResult<CopyInScan> {
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();
src/servers/src/postgres/server.rs
+26 -1
View File
@@ -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<ReloadableTlsServerConfig>,
keep_alive_secs: u64,
batching_enabled: bool,
influxdb_handler: Option<InfluxdbLineProtocolHandlerRef>,
bind_addr: Option<SocketAddr>,
process_manager: Option<ProcessManagerRef>,
}
@@ -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<Output = ()> + 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();
tests-integration/src/test_util.rs
+39 -30
View File
@@ -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::<SocketAddr>().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::<SocketAddr>().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::<SocketAddr>().unwrap())
tests-integration/src/tests/postgres_copy_in_test.rs
+300
View File
@@ -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<String> = 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<String> = 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<u8> = 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<String> = 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<String> = 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();
}