diff --git a/Cargo.lock b/Cargo.lock index c8dfd19620a..16739107f16 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -15394,6 +15394,7 @@ dependencies = [ "auth", "axum 0.8.4", "base64 0.22.1", + "bytes", "cache", "catalog", "chrono", diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 2e49bf0b7cc..67f1e6003a0 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -13,6 +13,7 @@ // limitations under the License. pub mod builder; +mod copy_in; mod dashboard; mod entity_graph; mod export_database; @@ -4699,4 +4700,69 @@ mod tests { assert!(matches!(&submitted[0], DdlTask::CreateTable(_))); Ok(()) } + + /// COPY FROM STDIN must enforce the same cross-catalog restriction as + /// SQL INSERT when `disallow_cross_catalog_query` is configured. + #[tokio::test] + async fn test_copy_in_cross_catalog_restriction() -> TestResult<()> { + use catalog::memory::MemoryCatalogManager; + use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; + use servers::query_handler::CopyInHandler; + + let catalog = MemoryCatalogManager::new_with_table(test_table(1024, "source")?); + // A table in another catalog. + catalog.register_catalog_sync("other").unwrap(); + catalog + .register_schema_sync(catalog::RegisterSchemaRequest { + catalog: "other".to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + }) + .unwrap(); + catalog + .register_table_sync(catalog::RegisterTableRequest { + catalog: "other".to_string(), + schema: DEFAULT_SCHEMA_NAME.to_string(), + table_name: "copy_guard".to_string(), + table_id: 2048, + table: test_table(2048, "copy_guard")?, + }) + .unwrap(); + let executor = Arc::new(MockProcedureExecutor::new(catalog.clone())); + let plugins = Plugins::new(); + plugins.insert(QueryOptions { + disallow_cross_catalog_query: true, + }); + let instance = test_instance_with_catalog_manager( + catalog, + test_table(1025, "target")?, + plugins, + None, + executor, + ) + .await?; + + let ctx = Arc::new(QueryContext::with( + DEFAULT_CATALOG_NAME, + DEFAULT_SCHEMA_NAME, + )); + + // Cross-catalog COPY is rejected before any data is streamed. + let err = instance + .copy_in_table("other", DEFAULT_SCHEMA_NAME, "copy_guard", ctx.clone()) + .await + .unwrap_err(); + assert!( + format!("{err:?}").contains("Illegal access to catalog"), + "{err}" + ); + + // Same-catalog COPY resolves normally. + let found = instance + .copy_in_table(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, "target", ctx) + .await + .unwrap(); + assert!(found.is_some()); + + Ok(()) + } } diff --git a/src/frontend/src/instance/copy_in.rs b/src/frontend/src/instance/copy_in.rs new file mode 100644 index 00000000000..6580cc44b9d --- /dev/null +++ b/src/frontend/src/instance/copy_in.rs @@ -0,0 +1,87 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use api::v1::RowInsertRequests; +use async_trait::async_trait; +use auth::{PermissionChecker, PermissionCheckerRef, PermissionReq}; +use common_error::ext::BoxedError; +use common_query::Output; +use query::query_engine::options::{QueryOptions, validate_catalog_and_schema}; +use servers::error::{AuthSnafu, ExecuteGrpcRequestSnafu, ExecuteQuerySnafu, Result}; +use servers::query_handler::CopyInHandler; +use session::context::QueryContextRef; +use snafu::ResultExt; +use table::metadata::TableInfoRef; + +use crate::instance::Instance; + +#[async_trait] +impl CopyInHandler for Instance { + async fn copy_in_table( + &self, + catalog: &str, + schema: &str, + table: &str, + query_ctx: QueryContextRef, + ) -> Result> { + // Mirror SQL INSERT's catalog restriction (see `check_permission`): + // when cross-catalog queries are disallowed, a COPY may only + // target the current catalog. This must run against the ORIGINAL + // query context — `begin_copy_in` switches the context to the + // target catalog afterwards, so nothing downstream would catch + // it. + let disallow_cross_catalog_query = self + .plugins + .get::() + .map(|opts| opts.disallow_cross_catalog_query) + .unwrap_or_default(); + if disallow_cross_catalog_query { + validate_catalog_and_schema(catalog, schema, &query_ctx) + .map_err(BoxedError::new) + .context(ExecuteQuerySnafu)?; + } + + // A COPY writes rows like a SQL INSERT; check the write permission + // before the client starts streaming data. + self.plugins + .get::() + .as_ref() + .check_permission( + query_ctx.current_user(), + PermissionReq::BulkInsert { + catalog, + schema, + table, + }, + ) + .context(AuthSnafu)?; + + self.catalog_manager() + .table(catalog, schema, table, None) + .await + .map(|table| table.map(|table| table.table_info())) + .map_err(servers::error::Error::from) + } + + async fn copy_in_insert( + &self, + requests: RowInsertRequests, + query_ctx: QueryContextRef, + ) -> Result { + self.handle_row_inserts(requests, query_ctx, false, false) + .await + .map_err(BoxedError::new) + .context(ExecuteGrpcRequestSnafu) + } +} diff --git a/src/frontend/src/server.rs b/src/frontend/src/server.rs index 8de7bfe340f..0b1e2b9713e 100644 --- a/src/frontend/src/server.rs +++ b/src/frontend/src/server.rs @@ -420,6 +420,7 @@ where let pg_server = Box::new( PostgresServer::new( + instance.clone(), instance.clone(), opts.tls.should_force_tls(), tls_server_config, diff --git a/src/servers/src/metrics.rs b/src/servers/src/metrics.rs index 888bdc42cdf..d928c60c741 100644 --- a/src/servers/src/metrics.rs +++ b/src/servers/src/metrics.rs @@ -320,6 +320,11 @@ lazy_static! { "servers postgres prepared count" ) .unwrap(); + pub static ref METRIC_POSTGRES_COPY_IN_ROWS: IntCounter = register_int_counter!( + "greptime_servers_postgres_copy_in_rows", + "servers postgres COPY FROM STDIN rows written" + ) + .unwrap(); pub static ref METRIC_SERVER_GRPC_DB_REQUEST_TIMER: HistogramVec = register_histogram_vec!( "greptime_servers_grpc_db_request_elapsed", "servers grpc db request elapsed", diff --git a/src/servers/src/postgres.rs b/src/servers/src/postgres.rs index 15aae48a494..e13ec7392d5 100644 --- a/src/servers/src/postgres.rs +++ b/src/servers/src/postgres.rs @@ -13,6 +13,7 @@ // limitations under the License. mod auth_handler; +mod copy_in; mod fixtures; mod handler; mod server; @@ -42,6 +43,7 @@ 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; pub(crate) struct GreptimeDBStartupParameters { @@ -78,6 +80,7 @@ impl ServerParameterProvider for GreptimeDBStartupParameters { pub struct PostgresServerHandlerInner { query_handler: ServerSqlQueryHandlerRef, auth_source: Option>, + copy_in_handler: CopyInHandlerRef, force_tls: bool, auth_driver: AuthDriver, UserInfoRef>, /// SCRAM selected; failures get the uniform metric label. @@ -86,11 +89,16 @@ pub struct PostgresServerHandlerInner { batching_enabled: bool, session: Arc, query_parser: Arc, + + /// State of an in-progress `COPY ... FROM STDIN`; set when the copy + /// starts and taken when it finishes, fails or the connection drops. + copy_in_state: tokio::sync::Mutex>, } #[derive(Builder)] pub(crate) struct MakePostgresServerHandler { query_handler: ServerSqlQueryHandlerRef, + copy_in_handler: CopyInHandlerRef, user_provider: Option, #[builder(default = "Arc::new(GreptimeDBStartupParameters::new())")] param_provider: Arc, @@ -112,6 +120,10 @@ impl PgWireServerHandlers for PostgresServerHandler { self.0.clone() } + fn copy_handler(&self) -> Arc { + self.0.clone() + } + fn error_handler(&self) -> Arc { self.0.clone() } @@ -136,6 +148,7 @@ impl MakePostgresServerHandler { .user_provider .clone() .map(|p| Arc::new(PgUserProviderSource::new(p))), + copy_in_handler: self.copy_in_handler.clone(), force_tls: self.force_tls, auth_driver: AuthDriver::new(self.param_provider.clone()) .with_pid_secret_key_generator(Arc::new(ZeroPidSecretKeyGenerator)), @@ -144,6 +157,7 @@ impl MakePostgresServerHandler { batching_enabled, session: session.clone(), query_parser: Arc::new(DefaultQueryParser::new(self.query_handler.clone(), session)), + copy_in_state: tokio::sync::Mutex::new(None), }; PostgresServerHandler(Arc::new(handler)) } diff --git a/src/servers/src/postgres/copy_in/codec.rs b/src/servers/src/postgres/copy_in/codec.rs new file mode 100644 index 00000000000..37a94811309 --- /dev/null +++ b/src/servers/src/postgres/copy_in/codec.rs @@ -0,0 +1,116 @@ +// 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 COPY FROM STDIN codec framework. +//! +//! 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. +//! +//! # 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. +//! +//! # 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. + +use datatypes::prelude::ConcreteDataType; +use datatypes::schema::ColumnSchema; +use pgwire::error::PgWireResult; + +/// 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. +pub(super) type ParsedRow = api::v1::Row; + +/// Incremental parser of one COPY FROM STDIN stream. +/// +/// The wire protocol delivers the payload as an arbitrary sequence of +/// `CopyData` chunks; implementations must keep state across [`feed`] +/// calls because records (and even single fields) may span chunk +/// boundaries. [`finish`] is called at `CopyDone`; a final record without +/// a trailing record separator is accepted. +pub(super) trait RecordParser: Send { + /// Feeds a chunk of `CopyData` bytes and returns all complete rows. + fn feed(&mut self, data: &[u8]) -> PgWireResult>; + /// Parses the buffered tail at `CopyDone`; returns no row when the + /// 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 new file mode 100644 index 00000000000..b2919210e3f --- /dev/null +++ b/src/servers/src/postgres/copy_in/csv.rs @@ -0,0 +1,703 @@ +// 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 `csv` format — currently the only implemented COPY FROM STDIN +//! format. +//! +//! Parsing and typed conversion both reuse arrow-csv's incremental +//! [`Decoder`](arrow::csv::reader::Decoder) — the same engine backing +//! GreptimeDB's file-based `COPY ... FROM 'file.csv'` (see +//! `common-datasource`'s `csv` file format) — so both COPY routes accept +//! identical data. The decoder splits records (quoting, escaped quotes, +//! custom delimiters, `\n`/`\r\n` separators, optional headers) and +//! converts fields against the targeted columns' schema in one pass. +//! +//! Consequences of sharing the engine: JSON and Binary columns are +//! rejected, timestamps must be date/time strings (offset-free ones +//! parse in the session timezone), and an empty field — quoted or not — +//! is NULL, exactly like file-based COPY. + +use std::sync::Arc; + +use api::helper::vectors_to_rows; +use arrow::csv::ReaderBuilder; +use arrow::csv::reader::Decoder as CsvDecoder; +use arrow_schema::{ArrowError, DataType as ArrowDataType, Field, Schema as ArrowSchema}; +use datatypes::data_type::DataType; +use datatypes::prelude::ConcreteDataType; +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::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)] +pub(super) struct CsvCodec { + delimiter: u8, + quote: u8, + escape: u8, + header: bool, +} + +/// Builds the CSV codec from the parsed options. +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 + // character itself (doubled quotes inside quoted fields). + let escape = options.take_single_byte("ESCAPE")?.unwrap_or(quote); + // An empty string is the default NULL marker (an empty field is NULL); + // a custom NULL string is not supported because arrow-csv's null + // matching does not distinguish quoted from unquoted fields. + if let Some(null) = options.take_string("NULL")? + && !null.is_empty() + { + return Err(unsupported( + "custom NULL string is not supported by format csv; an empty field is NULL", + )); + } + let header = options.take_flag("HEADER")?.unwrap_or(false); + options.reject_leftovers("csv")?; + if delimiter == quote { + return Err(bad_copy_data("DELIMITER must not be equal to QUOTE")); + } + Ok(Arc::new(CsvCodec { + delimiter, + quote, + escape, + header, + })) +} + +impl CopyInCodec for CsvCodec { + fn name(&self) -> &'static str { + "csv" + } + + /// The column types arrow-csv can decode against. + fn check_column_type(&self, data_type: &ConcreteDataType) -> PgWireResult<()> { + let supported = matches!( + data_type, + ConcreteDataType::Boolean(_) + | ConcreteDataType::Int8(_) + | ConcreteDataType::Int16(_) + | ConcreteDataType::Int32(_) + | ConcreteDataType::Int64(_) + | ConcreteDataType::UInt8(_) + | ConcreteDataType::UInt16(_) + | ConcreteDataType::UInt32(_) + | ConcreteDataType::UInt64(_) + | ConcreteDataType::Float32(_) + | ConcreteDataType::Float64(_) + | ConcreteDataType::String(_) + | ConcreteDataType::Date(_) + | ConcreteDataType::Timestamp(_) + | ConcreteDataType::Decimal128(_) + ); + if !supported { + return Err(unsupported(format!( + "column type \"{}\" is not supported by format csv", + data_type.name() + ))); + } + Ok(()) + } + + fn create_parser( + &self, + columns: &[ColumnSchema], + timezone: &str, + ) -> PgWireResult> { + // Decode schema: the greptime schema's arrow projection, except + // Boolean columns decode as Utf8 — arrow-csv only accepts + // `true`/`false`, while PostgreSQL COPY IN accepts its full set + // of boolean spellings (`t`, `f`, `yes`, `on`, ...), which COPY + // OUT emits — and Timestamp fields carry the SESSION timezone so + // offset-free strings parse like SQL INSERT does instead of as + // UTC. The substitutions are applied in `flush_batch` and by the + // decoder's tz-aware parsing respectively; explicit offsets in + // the input always win. + let boolean_columns: Vec = columns + .iter() + .enumerate() + .filter(|(_, c)| matches!(c.data_type, ConcreteDataType::Boolean(_))) + .map(|(i, _)| i) + .collect(); + let arrow_schema = GreptimeSchema::try_new(columns.to_vec()) + .map_err(|e| unsupported(format!("invalid columns for format csv: {e}")))? + .arrow_schema() + .clone(); + let timezone = timezone.to_string(); + let fields = arrow_schema + .fields() + .iter() + .zip(columns) + .map(|(field, column)| match &column.data_type { + ConcreteDataType::Boolean(_) => { + Field::new(field.name(), ArrowDataType::Utf8, field.is_nullable()) + } + ConcreteDataType::Timestamp(_) => match field.data_type() { + ArrowDataType::Timestamp(unit, None) => Field::new( + field.name(), + ArrowDataType::Timestamp(*unit, Some(Arc::from(timezone.as_str()))), + field.is_nullable(), + ), + _ => field.as_ref().clone(), + }, + _ => field.as_ref().clone(), + }) + .collect::>(); + let mut builder = ReaderBuilder::new(Arc::new(ArrowSchema::new(fields))) + .with_header(self.header) + .with_delimiter(self.delimiter) + .with_quote(self.quote) + .with_batch_size(COPY_IN_FLUSH_ROWS); + // An escape character equal to the quote is the doubled-quote + // default; csv-core handles that through quoting alone. + if self.escape != self.quote { + builder = builder.with_escape(self.escape); + } + Ok(Box::new(CsvRecordParser { + decoder: builder.build_decoder(), + boolean_columns, + })) + } +} + +/// Incremental parser backed by arrow-csv's decoder. +struct CsvRecordParser { + decoder: CsvDecoder, + /// Indices of Boolean columns, decoded as Utf8 and converted by + /// [`flush_batch`] to accept PostgreSQL's boolean spellings. + boolean_columns: Vec, +} + +impl RecordParser for CsvRecordParser { + fn feed(&mut self, data: &[u8]) -> PgWireResult> { + let mut buf = data; + let mut rows = Vec::new(); + while !buf.is_empty() { + let consumed = self.decoder.decode(buf).map_err(decode_error)?; + if consumed == 0 { + while let Some(batch) = self.decoder.flush().map_err(decode_error)? { + rows.extend(self.flush_batch(&batch)?); + } + } else { + buf = &buf[consumed..]; + } + } + Ok(rows) + } + + fn finish(&mut self) -> PgWireResult> { + // An empty decode signals EOF, finalizing a trailing record that + // lacks its record separator. + self.decoder.decode(&[]).map_err(decode_error)?; + let mut rows = Vec::new(); + while let Some(batch) = self.decoder.flush().map_err(decode_error)? { + rows.extend(self.flush_batch(&batch)?); + } + Ok(rows) + } +} + +impl CsvRecordParser { + fn flush_batch( + &mut self, + batch: &arrow::record_batch::RecordBatch, + ) -> PgWireResult> { + let mut vectors = Helper::try_into_vectors(batch.columns()) + .map_err(|e| bad_copy_data(format!("failed to convert decoded CSV batch: {e}")))?; + for &index in &self.boolean_columns { + let Some(vector) = vectors[index].as_any().downcast_ref::() else { + return Err(bad_copy_data( + "expected boolean column to decode as a string vector", + )); + }; + let mut values = Vec::with_capacity(batch.num_rows()); + for row_index in 0..batch.num_rows() { + values.push(match vector.get(row_index) { + datatypes::value::Value::Null => None, + datatypes::value::Value::String(s) => Some(parse_pg_bool(s.as_utf8())?), + other => { + return Err(bad_copy_data(format!( + "expected string values in a boolean column, got {other:?}" + ))); + } + }); + } + vectors[index] = Arc::new(BooleanVector::from(values)); + } + Ok(vectors_to_rows(vectors.iter(), batch.num_rows())) + } +} + +/// Parses PostgreSQL's boolean spellings: `true`/`false`, `t`/`f`, +/// `yes`/`no`, `on`/`off` and `1`/`0`, case-insensitive with surrounding +/// whitespace allowed (see the BOOLEAN type's input grammar). COPY OUT +/// emits `t`/`f`, so accepting them keeps both COPY directions +/// round-trippable. +fn parse_pg_bool(input: &str) -> PgWireResult { + match input.trim().to_ascii_lowercase().as_str() { + "true" | "t" | "yes" | "y" | "on" | "1" => Ok(true), + "false" | "f" | "no" | "n" | "off" | "0" => Ok(false), + _ => Err(invalid_text_repr(format!( + "invalid input syntax for type boolean: \"{input}\"" + ))), + } +} + +fn decode_error(error: ArrowError) -> PgWireError { + match &error { + ArrowError::CsvError(message) => { + bad_copy_data(format!("failed to parse CSV data: {message}")) + } + ArrowError::ParseError(message) | ArrowError::CastError(message) => { + invalid_text_repr(format!("failed to parse CSV data: {message}")) + } + _ => bad_copy_data(format!("failed to parse CSV data: {error}")), + } +} + +#[cfg(test)] +mod tests { + use datatypes::schema::ColumnSchema; + + use super::*; + use crate::postgres::copy_in::options::CopyOptionValue; + + fn test_columns() -> Vec { + vec![ + ColumnSchema::new("host", ConcreteDataType::string_datatype(), true), + ColumnSchema::new("val", ConcreteDataType::float64_datatype(), true), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ] + } + + fn parser() -> Box { + build_csv_codec(CopyOptionSet::default()) + .unwrap() + .create_parser(&test_columns(), "UTC") + .unwrap() + } + + fn string_at(row: &ParsedRow, index: usize) -> Option { + row.values[index].value_data.as_ref().map(|v| match v { + api::v1::value::ValueData::StringValue(s) => s.clone(), + other => panic!("expected string, got {other:?}"), + }) + } + + fn rows_with(mut set: CopyOptionSet) -> Box { + set.take_format().unwrap(); + build_csv_codec(set) + .unwrap() + .create_parser(&test_columns(), "UTC") + .unwrap() + } + + /// Feeds chunks and then the EOF signal, returning all rows: with the + /// strict decoder protocol, feeds only emit full + /// `COPY_IN_FLUSH_ROWS` batches; a trailing partial batch surfaces at + /// `finish`. + fn feed_all(parser: &mut Box, chunks: &[&[u8]]) -> Vec { + let mut rows = Vec::new(); + for chunk in chunks { + rows.extend(parser.feed(chunk).unwrap()); + } + rows.extend(parser.finish().unwrap()); + rows + } + + #[test] + fn test_build_defaults() { + let codec = build_csv_codec(CopyOptionSet::default()).unwrap(); + assert_eq!(codec.name(), "csv"); + } + + #[test] + fn test_build_rejects_bad_options() { + let mut set = CopyOptionSet::default(); + set.insert("delimiter".into(), CopyOptionValue::String("\"".into())) + .unwrap(); + let err = build_csv_codec(set).unwrap_err(); + assert!( + err.to_string() + .contains("DELIMITER must not be equal to QUOTE") + ); + + let mut set = CopyOptionSet::default(); + set.insert("null".into(), CopyOptionValue::String("NULL".into())) + .unwrap(); + assert!( + build_csv_codec(set) + .unwrap_err() + .to_string() + .contains("custom NULL string") + ); + + let mut set = CopyOptionSet::default(); + set.take_format().unwrap(); + set.insert("precision".into(), CopyOptionValue::Ident("ns".into())) + .unwrap(); + assert!( + build_csv_codec(set) + .unwrap_err() + .to_string() + .contains("not supported by format csv") + ); + } + + #[test] + fn test_check_column_type() { + let codec = build_csv_codec(CopyOptionSet::default()).unwrap(); + assert!( + codec + .check_column_type(&ConcreteDataType::json_datatype()) + .is_err() + ); + assert!( + codec + .check_column_type(&ConcreteDataType::binary_datatype()) + .is_err() + ); + assert!( + codec + .check_column_type(&ConcreteDataType::string_datatype()) + .is_ok() + ); + } + + #[test] + fn test_csv_parse_records() { + let mut parser = parser(); + let rows = feed_all( + &mut parser, + &[b"host1,1.5,2023-11-14 22:13:20\n\ + ,2.5,2023-11-14 22:13:21\n\ + \"quoted, comma\",3.5,2023-11-14 22:13:22\n"], + ); + assert_eq!(rows.len(), 3); + assert_eq!(string_at(&rows[0], 0), Some("host1".into())); + // An empty field is NULL. + assert_eq!(string_at(&rows[1], 0), None); + assert_eq!(string_at(&rows[2], 0), Some("quoted, comma".into())); + } + + #[test] + fn test_csv_escaped_quotes() { + let mut parser = parser(); + let rows = feed_all( + &mut parser, + &[b"\"say \"\"hi\"\"\",1.0,2023-11-14 22:13:20\n"], + ); + assert_eq!(string_at(&rows[0], 0), Some("say \"hi\"".into())); + } + + #[test] + fn test_csv_custom_escape() { + let mut set = CopyOptionSet::default(); + set.insert("escape".into(), CopyOptionValue::String("\\".into())) + .unwrap(); + let mut parser = rows_with(set); + let rows = feed_all(&mut parser, &[b"\"a\\\"b\",1.0,2023-11-14 22:13:20\n"]); + assert_eq!(string_at(&rows[0], 0), Some("a\"b".into())); + } + + #[test] + fn test_csv_quoted_newline_and_chunk_boundaries() { + let mut parser = parser(); + // Split inside a quoted field containing a newline, and a CRLF + // record separator split across chunks. No batch is full, so all + // feeds return nothing and both rows surface at the EOF signal. + assert!(parser.feed(b"\"multi").unwrap().is_empty()); + assert!( + parser + .feed(b"line\nfield\",1.0,2023-11-14 22:13:20") + .unwrap() + .is_empty() + ); + assert!( + parser + .feed(b"\r\nhost2,2.0,2023-11-14 22:13:21\n") + .unwrap() + .is_empty() + ); + let rows = parser.finish().unwrap(); + assert_eq!(rows.len(), 2); + assert_eq!(string_at(&rows[0], 0), Some("multiline\nfield".into())); + assert_eq!(string_at(&rows[1], 0), Some("host2".into())); + } + + #[test] + fn test_csv_header() { + let mut set = CopyOptionSet::default(); + set.insert("header".into(), CopyOptionValue::Flag).unwrap(); + let mut parser = rows_with(set); + let rows = feed_all( + &mut parser, + &[b"host,val,ts\nhost1,1.0,2023-11-14 22:13:20\n"], + ); + assert_eq!(rows.len(), 1); + assert_eq!(string_at(&rows[0], 0), Some("host1".into())); + } + + #[test] + fn test_csv_trailing_record_without_newline() { + let mut parser = parser(); + let fed = parser + .feed(b"host1,1.0,2023-11-14 22:13:20\nhost2,2.0,2023-11-14 22:13:21") + .unwrap(); + // While the trailing record is pending, flush cannot emit the + // completed one either; both surface at the EOF signal. + assert!(fed.is_empty()); + let rows = parser.finish().unwrap(); + assert_eq!(rows.len(), 2); + assert_eq!(string_at(&rows[0], 0), Some("host1".into())); + assert_eq!(string_at(&rows[1], 0), Some("host2".into())); + } + + #[test] + fn test_csv_finish_at_record_boundary() { + let mut parser = parser(); + // The complete record stays buffered (no full batch); the feed + // emits nothing and `finish` surfaces it. + assert!( + parser + .feed(b"host1,1.0,2023-11-14 22:13:20\n") + .unwrap() + .is_empty() + ); + let rows = parser.finish().unwrap(); + assert_eq!(rows.len(), 1); + assert_eq!(string_at(&rows[0], 0), Some("host1".into())); + } + + /// Regression test for chunk-boundary corruption: CopyData chunking is + /// arbitrary (records may split mid-field), and an earlier feed loop + /// flushed partial batches at record-boundary chunk ends, which + /// arrow-csv does not permit mid-stream — fields then parse + /// ("00001h,ost-001u,..."). Feed the same payload at many chunk sizes + /// and verify EVERY cell. + #[test] + fn test_csv_chunked_feeds_are_alignment_safe() { + // (int32 id, timestamp) rows of near-equal length. + let sweep_columns = vec![ + ColumnSchema::new("i", ConcreteDataType::int32_datatype(), true), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ]; + // A payload wider than one decoder batch so batch-full flushes + // interleave with chunk boundaries too. + let total = 2 * COPY_IN_FLUSH_ROWS + 137; + let payload: Vec = (0..total) + .map(|i| { + format!( + "{:06},2024-01-01 00:00:{:02}.{:03}\n", + i, + (i / 1000) % 60, + i % 1000 + ) + }) + .collect::() + .into_bytes(); + + let check = |chunk: usize| { + let mut parser = build_csv_codec(CopyOptionSet::default()) + .unwrap() + .create_parser(&sweep_columns, "UTC") + .unwrap(); + let mut ids: Vec = Vec::new(); + for c in payload.chunks(chunk) { + for row in parser.feed(c).unwrap() { + ids.push(match row.values[0].value_data.as_ref().unwrap() { + api::v1::value::ValueData::I32Value(v) => *v, + other => panic!("expected i32, got {other:?}"), + }); + } + } + for row in parser.finish().unwrap() { + ids.push(match row.values[0].value_data.as_ref().unwrap() { + api::v1::value::ValueData::I32Value(v) => *v, + other => panic!("expected i32, got {other:?}"), + }); + } + assert_eq!(ids.len(), total, "chunk={chunk}"); + for (i, v) in ids.iter().enumerate() { + assert_eq!(*v, i as i32, "chunk={chunk} row {i} corrupted"); + } + }; + + // Small chunk sizes hit every possible boundary alignment, + // including multiple decode calls per record. + for chunk in [1usize, 2, 3, 7, 8, 13, 31, 64, 97, 128, 1024] { + check(chunk); + } + // Client-natural sizes with a shifted start to vary where + // boundaries land relative to records. + for chunk in [4096usize, 8192, 65531] { + check(chunk); + let mut shifted = payload.clone(); + shifted.drain(..37); + let mut parser = build_csv_codec(CopyOptionSet::default()) + .unwrap() + .create_parser(&sweep_columns, "UTC") + .unwrap(); + let mut count = 0usize; + for c in shifted.chunks(chunk) { + count += parser.feed(c).unwrap().len(); + } + count += parser.finish().unwrap().len(); + assert_eq!(count, total - 1, "chunk={chunk} shifted"); + } + } + + #[test] + fn test_csv_boolean_spellings() { + // PostgreSQL COPY OUT emits `t`/`f`; COPY IN accepts the full + // boolean input grammar. + let columns = vec![ + ColumnSchema::new("flag", ConcreteDataType::boolean_datatype(), true), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ]; + let mut parser = build_csv_codec(CopyOptionSet::default()) + .unwrap() + .create_parser(&columns, "UTC") + .unwrap(); + let rows = feed_all( + &mut parser, + &[b"t,2024-01-01 00:00:00\n\ + f,2024-01-01 00:00:01\n\ + TRUE,2024-01-01 00:00:02\n\ + no,2024-01-01 00:00:03\n\ + on,2024-01-01 00:00:04\n\ + 0,2024-01-01 00:00:05\n\ + ,2024-01-01 00:00:06\n"], + ); + let expected = [ + Some(true), + Some(false), + Some(true), + Some(false), + Some(true), + Some(false), + None, + ]; + for (row, want) in rows.iter().zip(expected) { + assert_eq!( + row.values[0].value_data, + want.map(api::v1::value::ValueData::BoolValue), + "{:?}", + row.values[0].value_data + ); + } + } + + #[test] + fn test_csv_boolean_rejects_garbage() { + let columns = vec![ + ColumnSchema::new("flag", ConcreteDataType::boolean_datatype(), true), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ]; + let mut parser = build_csv_codec(CopyOptionSet::default()) + .unwrap() + .create_parser(&columns, "UTC") + .unwrap(); + parser.feed(b"maybe,2024-01-01 00:00:00\n").unwrap(); + let err = match parser.finish() { + Ok(_) => panic!("expected parse error"), + Err(e) => e, + }; + assert!( + err.to_string() + .contains("invalid input syntax for type boolean"), + "{err}" + ); + } + + #[test] + fn test_csv_timestamp_session_timezone() { + // Offset-free timestamp strings must parse in the session + // timezone (parity with SQL INSERT); explicit offsets win. + let columns = vec![ + ColumnSchema::new("host", ConcreteDataType::string_datatype(), true), + ColumnSchema::new( + "ts", + ConcreteDataType::timestamp_millisecond_datatype(), + false, + ), + ]; + let mut parser = build_csv_codec(CopyOptionSet::default()) + .unwrap() + .create_parser(&columns, "Asia/Shanghai") + .unwrap(); + let rows = feed_all( + &mut parser, + &[b"host1,2024-01-01 00:00:00\n\ + host2,2024-01-01T00:00:00+08:00\n\ + host3,2024-01-01T00:00:00Z\n"], + ); + let millis = |row: &ParsedRow| match row.values[1].value_data.as_ref().unwrap() { + api::v1::value::ValueData::TimestampMillisecondValue(v) => *v, + other => panic!("expected ts, got {other:?}"), + }; + // 2024-01-01T00:00:00+08:00 == 2023-12-31T16:00:00Z + assert_eq!(millis(&rows[0]), 1704038400000); + // Explicit +08:00 offset: same instant as the naive parse. + assert_eq!(millis(&rows[1]), 1704038400000); + // Explicit Z: parsed as UTC regardless of the session zone. + assert_eq!(millis(&rows[2]), 1704067200000); + } + + #[test] + fn test_csv_errors() { + let mut p = parser(); + let err = match p.feed(b"host1,1.0\n") { + Ok(_) => panic!("expected parse error"), + Err(e) => e, + }; + assert!( + err.to_string().contains("incorrect number of fields"), + "{err}" + ); + + // Value conversion runs at flush time, so the bad double + // surfaces at the EOF signal, not at feed time. + let mut p = parser(); + p.feed(b"host1,not_a_double,2023-11-14 22:13:20\n").unwrap(); + let err = match p.finish() { + Ok(_) => panic!("expected parse error"), + Err(e) => e, + }; + assert!(err.to_string().contains("failed to parse"), "{err}"); + } +} diff --git a/src/servers/src/postgres/copy_in/mod.rs b/src/servers/src/postgres/copy_in/mod.rs new file mode 100644 index 00000000000..a3959a23830 --- /dev/null +++ b/src/servers/src/postgres/copy_in/mod.rs @@ -0,0 +1,691 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! PostgreSQL `COPY ... FROM STDIN` (copy-in) support. +//! +//! The wire handshake is driven by pgwire: returning [`Response::CopyIn`] +//! from a query handler makes the server send a `CopyInResponse` and enter +//! the copy-in sub-protocol; subsequent `CopyData` messages from the client +//! are dispatched to [`CopyHandler`], whose implementation below parses +//! the data stream incrementally and flushes it into the query handler as +//! [`RowInsertRequests`] batches. +//! +//! # Format framework +//! +//! 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. + +mod codec; +mod csv; +mod options; +mod parse; + +use std::fmt::Debug; +use std::sync::Arc; + +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 common_query::OutputData; +use common_telemetry::{info, warn}; +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 pgwire::api::copy::CopyHandler; +use pgwire::api::results::{CopyResponse, Response, Tag}; +use pgwire::api::{ClientInfo, PgWireConnectionState}; +use pgwire::error::{ErrorInfo, PgWireError, PgWireResult}; +use pgwire::messages::PgWireBackendMessage; +use pgwire::messages::copy::{CopyData, CopyDone, CopyFail}; +use session::context::QueryContextRef; +use table::metadata::TableInfo; + +use crate::metrics::METRIC_POSTGRES_COPY_IN_ROWS; +use crate::postgres::PostgresServerHandlerInner; +use crate::postgres::types::PgErrorCode; +use crate::postgres::utils::convert_err; + +/// Number of parsed rows buffered before flushing a batch to the handler. +/// +/// Each flush runs the full insert pipeline (catalog lookup, on-demand +/// schema check, partition split, region write) and, without the +/// pending-rows batcher, waits for it — so larger batches mean fewer +/// round trips at the cost of linearly more memory per connection +/// (decoder batch + rows buffer). The constant doubles as the arrow +/// decoder's batch size. High-rate COPY users should enable the +/// pending-rows batcher (`global-table-batch` / `BatcherOptions.protocols` +/// containing `Postgres`), which turns flushes into enqueues and makes +/// the per-batch fixed costs negligible. +pub(crate) const COPY_IN_FLUSH_ROWS: usize = 32768; + +/// Builds a codec from the format-neutral option set. +type CodecBuilder = 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)] = &[ + ("csv", csv::build_csv_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> { + Err(unsupported( + "the text format is not supported yet; use WITH (FORMAT csv)", + )) +} + +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> { + let name = options.take_format()?.ok_or_else(|| { + unsupported( + "COPY FROM STDIN requires an explicit format in this version; use WITH (FORMAT csv)", + ) + })?; + let builder = REGISTERED_FORMATS + .iter() + .find(|(registered, _)| *registered == name) + .map(|(_, builder)| *builder) + .ok_or_else(|| bad_copy_data(format!("unknown COPY format \"{}\"", name)))?; + builder(std::mem::take(options)) +} + +/// The parsed `COPY tbl [(columns)] FROM STDIN [(options)]` statement. +#[derive(Debug, Clone)] +pub(crate) struct CopyFromStdin { + /// Table name parts as written (`t`, `schema.t`, `catalog.schema.t`); + /// quoted names keep their case. + table: Vec, + /// Requested column list; empty means all columns in table order. + columns: Vec, + /// The resolved format codec, validated and ready to run. + codec: Arc, +} + +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 { + copy_in_error(PgErrorCode::Ec0A000, message) +} + +fn bad_copy_data(message: impl Into) -> PgWireError { + copy_in_error(PgErrorCode::Ec22P04, message) +} + +fn invalid_text_repr(message: impl Into) -> PgWireError { + copy_in_error(PgErrorCode::Ec22P02, message) +} + +pub(crate) fn copy_in_must_run_alone_error() -> PgWireError { + PgWireError::UserError(Box::new(ErrorInfo::new( + "ERROR".to_string(), + PgErrorCode::Ec0A000.code(), + "COPY FROM STDIN must be executed alone".to_string(), + ))) +} + +/// Parses a leading `COPY ... FROM STDIN` statement from SQL text. +/// +/// The statement is parsed textually — not through sqlparser — because the +/// `WITH (...)` option list is open: formats may accept options beyond +/// sqlparser's COPY grammar, and unknown options must reach the format +/// codec instead of failing statement parsing. Returns `Ok(None)` when +/// the string does not lead with a copy-in statement. +pub(crate) fn copy_from_stdin_from_sql(sql: &str) -> PgWireResult> { + let raw = match parse::scan_copy_from_stdin(sql)? { + parse::CopyInScan::NotCopyIn => return Ok(None), + parse::CopyInScan::CopyIn(raw) => raw, + }; + let mut option_set = raw.options; + let codec = resolve_codec(&mut option_set)?; + Ok(Some(CopyFromStdin { + table: raw.table, + columns: raw.columns, + codec, + })) +} + +/// Whether an already-parsed statement is a `COPY ... FROM STDIN`; used to +/// reject copy-in statements trailing other statements, where the textual +/// scanner (which only sees statement-leading COPYs) does not see them. +pub(crate) fn is_copy_from_stdin(statement: &SqlParserStatement) -> bool { + matches!( + statement, + SqlParserStatement::Copy { + to: false, + target: CopyTarget::Stdin, + .. + } + ) +} + +/// Resolves table name parts against the query context defaults. +fn resolve_table_name( + query_ctx: &QueryContextRef, + parts: &[String], +) -> PgWireResult<(String, String, String)> { + match parts { + [table] => Ok(( + query_ctx.current_catalog().to_string(), + query_ctx.current_schema(), + table.clone(), + )), + [schema, table] => Ok(( + query_ctx.current_catalog().to_string(), + schema.clone(), + table.clone(), + )), + [catalog, schema, table] => Ok((catalog.clone(), schema.clone(), table.clone())), + _ => Err(unsupported(format!( + "invalid table name for COPY: {}", + parts.join(".") + ))), + } +} + +/// Per-connection state of an in-progress COPY FROM STDIN. +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, + /// 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, +} + +impl CopyInState { + fn new( + codec: Arc, + columns: Vec, + table_info: Arc, + query_ctx: QueryContextRef, + table_name: String, + ) -> PgWireResult { + let schema = table_info.meta.schema.clone(); + let timestamp_index = schema.timestamp_index(); + + let columns = resolve_columns(codec.as_ref(), columns, &schema, &table_info)?; + if let Some(ts_index) = timestamp_index { + let ts_column = &schema.column_schemas()[ts_index]; + if !columns.iter().any(|column| column.name == ts_column.name) { + return Err(unsupported(format!( + "COPY FROM STDIN requires the time index column \"{}\" in the column list", + ts_column.name + ))); + } + } + let parser = codec.create_parser(&columns, query_ctx.timezone().to_string().as_str())?; + + let grpc_schema = columns + .iter() + .map(|column| build_grpc_column_schema(column, &table_info)) + .collect::>>()?; + + 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), + 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. + fn feed(&mut self, data: &[u8]) -> PgWireResult<()> { + self.rows.extend(self.parser.feed(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()?); + Ok(()) + } + + fn should_flush(&self) -> bool { + self.rows.len() >= COPY_IN_FLUSH_ROWS + } + + /// Takes the buffered rows as an insert request, if any. + fn take_requests(&mut self) -> Option { + if self.rows.is_empty() { + return None; + } + let rows = GrpcRows { + schema: self.grpc_schema.clone(), + rows: std::mem::take(&mut self.rows), + }; + Some(RowInsertRequests { + inserts: vec![RowInsertRequest { + table_name: self.table_name.clone(), + rows: Some(rows), + }], + }) + } +} + +fn resolve_columns( + codec: &dyn CopyInCodec, + columns: Vec, + schema: &SchemaRef, + table_info: &Arc, +) -> PgWireResult> { + let table_name = &table_info.name; + let selected: Vec = if columns.is_empty() { + schema.column_schemas().to_vec() + } else { + let mut selected = Vec::with_capacity(columns.len()); + for name in &columns { + let column = schema + .column_schema_by_name(name) + .ok_or_else(|| { + copy_in_error( + PgErrorCode::Ec42703, + format!( + "column \"{}\" of relation \"{}\" does not exist", + name, table_name + ), + ) + })? + .clone(); + if selected.iter().any(|c: &ColumnSchema| &c.name == name) { + return Err(bad_copy_data(format!( + "column \"{}\" specified more than once", + name + ))); + } + selected.push(column); + } + selected + }; + + for column in &selected { + codec.check_column_type(&column.data_type)?; + } + + Ok(selected) +} + +/// Derives a column's semantic type the same way the flow adapter does +/// (`flow::adapter::util::column_schemas_to_proto`): the time-index flag +/// for timestamps, primary-key membership for tags. +fn semantic_type_of(table_info: &TableInfo, column: &ColumnSchema) -> PgWireResult { + let meta = &table_info.meta; + if column.is_time_index() { + Ok(SemanticType::Timestamp) + } else if meta + .primary_key_indices + .iter() + .any(|&index| meta.schema.column_schemas()[index].name == column.name) + { + Ok(SemanticType::Tag) + } else { + Ok(SemanticType::Field) + } +} + +fn build_grpc_column_schema( + column: &ColumnSchema, + table_info: &Arc, +) -> PgWireResult { + let (datatype, datatype_extension) = ColumnDataTypeWrapper::try_from(column.data_type.clone()) + .map_err(|_| { + unsupported(format!( + "column type \"{}\" is not supported for COPY FROM STDIN yet", + column.data_type.name() + )) + })? + .to_parts(); + Ok(api::v1::ColumnSchema { + column_name: column.name.clone(), + datatype: datatype as i32, + semantic_type: semantic_type_of(table_info, column)? as i32, + datatype_extension, + options: options_from_column_schema(column), + }) +} + +// --------------------------------------------------------------------------- +// pgwire handlers +// --------------------------------------------------------------------------- + +impl PostgresServerHandlerInner { + /// Prepares a COPY FROM STDIN: resolves the target table, validates + /// columns and stashes the per-connection copy state. + pub(crate) async fn begin_copy_in(&self, stmt: CopyFromStdin) -> PgWireResult { + // The batching-aware context routes copy-in flushes through the + // pending-rows batcher like SQL INSERT when it is enabled for + // this connection. + let query_ctx = self.new_query_context(); + 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); + + 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); + + Ok(Response::CopyIn(CopyResponse::new( + format_code, + num_columns, + stream::empty(), + ))) + } + + /// Flushes buffered rows 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 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) + .await + .map_err(convert_err)?; + match output.data { + OutputData::AffectedRows(rows) => Ok(rows as u64), + OutputData::RecordBatches(_) | OutputData::Stream(_) => { + warn!("unexpected output from COPY FROM STDIN insert"); + Ok(0) + } + } + } +} + +fn adjust_query_context( + query_ctx: QueryContextRef, + catalog: &str, + schema: &str, +) -> QueryContextRef { + if query_ctx.current_catalog() == catalog && query_ctx.current_schema() == schema { + return query_ctx; + } + let mut forked = query_ctx.fork(); + forked.set_current_catalog(catalog); + forked.set_current_schema(schema); + Arc::new(forked) +} + +#[async_trait] +impl CopyHandler for PostgresServerHandlerInner { + async fn on_copy_data(&self, _client: &mut C, copy_data: CopyData) -> PgWireResult<()> + where + C: ClientInfo + Sink + Unpin + Send + Sync, + C::Error: Debug, + PgWireError: From<>::Error>, + { + let mut guard = self.copy_in_state.lock().await; + let Some(mut state) = guard.take() else { + return Err(copy_in_error( + PgErrorCode::Ec08P01, + "CopyData message received outside COPY FROM STDIN", + )); + }; + // The state is owned here: any parse, conversion or insert error + // below aborts the whole copy and drops the buffered data — + // otherwise a pooled connection would retain failed-import + // buffers until its next COPY. On success it is put back. + state.feed(©_data.data)?; + if state.should_flush() { + let rows = self.flush_copy_in(&mut state).await?; + state.rows_written += rows; + METRIC_POSTGRES_COPY_IN_ROWS.inc_by(rows); + } + *guard = Some(state); + Ok(()) + } + + async fn on_copy_done(&self, client: &mut C, _done: CopyDone) -> PgWireResult<()> + where + C: ClientInfo + Sink + Unpin + Send + Sync, + C::Error: Debug, + PgWireError: From<>::Error>, + { + let mut guard = self.copy_in_state.lock().await; + // Take ownership up front: even finalization failures (parse + // errors in the trailing record, insert errors) drop the state + // instead of retaining it on a long-lived connection. + let Some(mut state) = guard.take() else { + return Err(copy_in_error( + PgErrorCode::Ec08P01, + "CopyDone message received outside COPY FROM STDIN", + )); + }; + state.finish()?; + let rows = self.flush_copy_in(&mut state).await?; + state.rows_written += rows; + METRIC_POSTGRES_COPY_IN_ROWS.inc_by(rows); + + let rows_written = state.rows_written; + let codec_name = state.codec_name; + info!( + "PostgreSQL COPY FROM STDIN (format {}) finished: {} rows written", + codec_name, rows_written + ); + client + .send(PgWireBackendMessage::CommandComplete( + Tag::new("COPY").with_rows(rows_written as usize).into(), + )) + .await?; + + // After a successful extended-protocol copy-in, pgwire leaves the + // connection in `CopyInProgress(true)` and swallows the client's + // terminating `Sync` message, so `ReadyForQuery` would never be + // sent and the connection would deadlock. Route the state to + // `AwaitingSync` so the `Sync` is dispatched to `on_sync`, which + // completes the round trip. + if let PgWireConnectionState::CopyInProgress(true) = client.state() { + client.set_state(PgWireConnectionState::AwaitingSync); + } + Ok(()) + } + + async fn on_copy_fail(&self, _client: &mut C, fail: CopyFail) -> PgWireError + where + C: ClientInfo + Sink + Unpin + Send + Sync, + C::Error: Debug, + PgWireError: From<>::Error>, + { + // Drop the buffered state; partial data of a failed copy is not + // written. + self.copy_in_state.lock().await.take(); + copy_in_error( + PgErrorCode::Ec57014, + format!("COPY FROM STDIN mode terminated by user: {}", fail.message), + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_parse_copy_from_stdin_requires_format() { + let err = copy_from_stdin_from_sql("COPY t FROM STDIN").unwrap_err(); + assert!( + err.to_string() + .contains("requires an explicit format in this version"), + "{err}" + ); + + let copy = copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT csv)") + .unwrap() + .unwrap(); + assert_eq!(copy.table, vec!["t"]); + assert!(copy.columns.is_empty()); + assert_eq!(copy.codec.name(), "csv"); + assert_eq!(copy.codec.format_code(), 0); + } + + #[test] + fn test_parse_copy_from_stdin_csv() { + let copy = copy_from_stdin_from_sql( + "COPY t (a, b) FROM STDIN WITH (FORMAT csv, HEADER, DELIMITER ';')", + ) + .unwrap() + .unwrap(); + assert_eq!(copy.columns, vec!["a", "b"]); + assert_eq!(copy.codec.name(), "csv"); + } + + #[test] + fn test_parse_copy_from_stdin_text_unsupported() { + for sql in [ + "COPY t FROM STDIN WITH (FORMAT text)", + "COPY t FROM STDIN WITH (FORMAT txt)", + ] { + let err = copy_from_stdin_from_sql(sql).unwrap_err(); + assert!( + err.to_string().contains("text format is not supported yet"), + "{err}" + ); + } + } + + #[test] + fn test_parse_copy_from_stdin_binary_unsupported() { + let err = copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT binary)").unwrap_err(); + assert!(err.to_string().contains("not supported")); + } + + #[test] + fn test_parse_copy_from_stdin_unknown_format() { + let err = copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT parquet)").unwrap_err(); + assert!(err.to_string().contains("unknown COPY format")); + } + + #[test] + fn test_parse_copy_from_stdin_not_copy_in() { + assert!( + copy_from_stdin_from_sql("COPY (SELECT 1) TO STDOUT") + .unwrap() + .is_none() + ); + assert!( + copy_from_stdin_from_sql("COPY t TO '/tmp/x.csv'") + .unwrap() + .is_none() + ); + // GreptimeDB's own file-based COPY FROM. + assert!( + copy_from_stdin_from_sql("COPY t FROM '/tmp/x.csv' WITH (FORMAT csv)") + .unwrap() + .is_none() + ); + } + + #[test] + fn test_redundant_options() { + let err = copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT csv, FORMAT text)") + .unwrap_err(); + assert!(err.to_string().contains("redundant FORMAT")); + + let err = copy_from_stdin_from_sql("COPY t FROM STDIN WITH (DELIMITER ',', DELIMITER ';')") + .unwrap_err(); + assert!(err.to_string().contains("redundant DELIMITER")); + } + + #[test] + fn test_format_specific_options_rejected() { + // The open WITH section parses anything, but the format rejects + // the options it does not understand. + let err = + copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT csv, PRECISION 'ns', FREEZE)") + .unwrap_err(); + assert!( + err.to_string().contains("not supported by format csv"), + "{err}" + ); + } + + #[test] + fn test_multi_statement_rejected() { + let err = copy_from_stdin_from_sql("COPY t FROM STDIN; SELECT 1").unwrap_err(); + assert!(err.to_string().contains("alone")); + // Trailing separators and quoted semicolons are fine. + assert!( + copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT csv);") + .unwrap() + .is_some() + ); + assert!( + copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT csv, DELIMITER ';')") + .unwrap() + .is_some() + ); + } +} diff --git a/src/servers/src/postgres/copy_in/options.rs b/src/servers/src/postgres/copy_in/options.rs new file mode 100644 index 00000000000..d710a0ca374 --- /dev/null +++ b/src/servers/src/postgres/copy_in/options.rs @@ -0,0 +1,281 @@ +// 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. + +//! Format-neutral collection of COPY `WITH (...)` options. +//! +//! The option list is **open**: every `KEY value` entry the client writes +//! is collected here, regardless of whether any format understands it. +//! Each format codec then takes the options it knows ([`CopyOptionSet::take_string`] +//! and friends) and rejects everything else via +//! [`CopyOptionSet::reject_leftovers`]. This way a format can introduce +//! its own options — say a line-protocol format accepting +//! `WITH (FORMAT influxdb_line, precision 'ns')` — without touching the +//! statement parser or any other codec. + +use pgwire::error::{PgWireError, PgWireResult}; + +use crate::postgres::copy_in::bad_copy_data; +use crate::postgres::types::PgErrorCode; + +/// The value of one COPY option. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum CopyOptionValue { + /// A single-quoted string literal. + String(String), + /// A numeric literal, kept as written. + Number(String), + /// A bare identifier. + Ident(String), + /// `TRUE` / `FALSE`. + Boolean(bool), + /// The option appears without a value. + Flag, +} + +impl CopyOptionValue { + /// The value as a descriptive word for error messages. + pub(crate) fn kind(&self) -> &'static str { + match self { + CopyOptionValue::String(_) => "a string value", + CopyOptionValue::Number(_) => "a numeric value", + CopyOptionValue::Ident(_) => "an identifier value", + CopyOptionValue::Boolean(_) => "a boolean value", + CopyOptionValue::Flag => "no value", + } + } +} + +/// The options of a `COPY ... FROM STDIN WITH (...)` statement. +/// +/// Entries keep statement order; keys are lowercased for lookup and +/// displayed uppercased in error messages. +#[derive(Debug, Default)] +pub(crate) struct CopyOptionSet { + entries: Vec<(String, CopyOptionValue)>, +} + +impl CopyOptionSet { + /// Adds an option; duplicate keys are rejected like PostgreSQL does. + pub(crate) fn insert(&mut self, key: String, value: CopyOptionValue) -> PgWireResult<()> { + if self.entries.iter().any(|(k, _)| *k == key) { + return Err(bad_copy_data(format!( + "conflicting or redundant {} option", + key.to_uppercase() + ))); + } + self.entries.push((key, value)); + Ok(()) + } + + /// Takes the raw value of `key`; for codec-specific options whose + /// shape this module does not prescribe. + pub(crate) fn take_value(&mut self, key: &str) -> Option { + let key = key.to_lowercase(); + let index = self.entries.iter().position(|(k, _)| *k == key)?; + Some(self.entries.remove(index).1) + } + + /// Takes the `FORMAT` option, lowercased. Accepts identifier or string + /// values. + pub(crate) fn take_format(&mut self) -> PgWireResult> { + let Some(value) = self.take_value("format") else { + return Ok(None); + }; + match value { + CopyOptionValue::Ident(name) | CopyOptionValue::String(name) => { + Ok(Some(name.to_lowercase())) + } + other => Err(bad_copy_data(format!( + "FORMAT requires an identifier value, got {}", + other.kind() + ))), + } + } + + /// Takes a string-valued option. + pub(crate) fn take_string(&mut self, key: &str) -> PgWireResult> { + let Some(value) = self.take_value(key) else { + return Ok(None); + }; + match value { + CopyOptionValue::String(s) => Ok(Some(s)), + other => Err(bad_copy_data(format!( + "COPY option {} requires a string value, got {}", + key.to_uppercase(), + other.kind() + ))), + } + } + + /// Takes a single one-byte character option. A tab is a legal (and + /// common) CSV delimiter; newlines and other control characters are + /// not, matching PostgreSQL's rules. + pub(crate) fn take_single_byte(&mut self, key: &str) -> PgWireResult> { + let Some(value) = self.take_string(key)? else { + return Ok(None); + }; + let mut chars = value.chars(); + let (c, rest) = (chars.next(), chars.as_str()); + match c { + Some('\t') if rest.is_empty() => Ok(Some(b'\t')), + Some(c) if rest.is_empty() && c.is_ascii() && !c.is_ascii_control() => { + Ok(Some(c as u8)) + } + _ => Err(bad_copy_data(format!( + "COPY option {} must be a single one-byte character other than a newline or carriage return", + key.to_uppercase() + ))), + } + } + + /// Takes a boolean option; a valueless option means `TRUE`. + pub(crate) fn take_flag(&mut self, key: &str) -> PgWireResult> { + let Some(value) = self.take_value(key) else { + return Ok(None); + }; + match value { + CopyOptionValue::Flag => Ok(Some(true)), + CopyOptionValue::Boolean(b) => Ok(Some(b)), + other => Err(bad_copy_data(format!( + "COPY option {} requires a boolean value, got {}", + key.to_uppercase(), + other.kind() + ))), + } + } + + /// Rejects every option not consumed by the format, with one message + /// naming them all. Codecs call this after their `take_*` calls so a + /// typo cannot be silently ignored. + pub(crate) fn reject_leftovers(self, format: &str) -> PgWireResult<()> { + if self.entries.is_empty() { + return Ok(()); + } + let unknown = self + .entries + .iter() + .map(|(key, _)| key.to_uppercase()) + .collect::>() + .join(", "); + Err(PgWireError::UserError(Box::new( + PgErrorCode::Ec0A000.to_err_info(format!( + "COPY option(s) {unknown} are not supported by format {format}" + )), + ))) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn new_set(entries: &[(&str, CopyOptionValue)]) -> CopyOptionSet { + let mut set = CopyOptionSet::default(); + for (key, value) in entries { + set.insert((*key).to_string(), value.clone()).unwrap(); + } + set + } + + #[test] + fn test_insert_rejects_duplicates() { + let mut set = CopyOptionSet::default(); + set.insert("format".to_string(), CopyOptionValue::Ident("csv".into())) + .unwrap(); + let err = set + .insert("format".to_string(), CopyOptionValue::Ident("text".into())) + .unwrap_err(); + assert!(err.to_string().contains("redundant FORMAT")); + } + + #[test] + fn test_take_single_byte_tab_allowed() { + // A literal tab is a legal delimiter (PostgreSQL permits it). + let mut set = new_set(&[("delimiter", CopyOptionValue::String("\t".into()))]); + assert_eq!(set.take_single_byte("DELIMITER").unwrap(), Some(b'\t')); + + // Newlines and other control characters stay rejected. + for bad in ["\n", "\r", "\0"] { + let mut set = new_set(&[("delimiter", CopyOptionValue::String(bad.into()))]); + assert!( + set.take_single_byte("DELIMITER").is_err(), + "{bad:?} must be rejected" + ); + } + } + + #[test] + fn test_take_string_and_single_byte() { + let mut set = new_set(&[ + ("null", CopyOptionValue::String("nil".into())), + ("delimiter", CopyOptionValue::String("|".into())), + ]); + assert_eq!(set.take_string("null").unwrap().as_deref(), Some("nil")); + assert_eq!(set.take_single_byte("delimiter").unwrap(), Some(b'|')); + + let mut set = new_set(&[("delimiter", CopyOptionValue::Number("42".into()))]); + let err = set.take_single_byte("delimiter").unwrap_err(); + assert!(err.to_string().contains("requires a string value")); + + let mut set = new_set(&[("delimiter", CopyOptionValue::String("ab".into()))]); + let err = set.take_single_byte("delimiter").unwrap_err(); + assert!(err.to_string().contains("single one-byte")); + } + + #[test] + fn test_take_flag() { + let mut set = new_set(&[ + ("header", CopyOptionValue::Flag), + ("frozen", CopyOptionValue::Boolean(false)), + ]); + assert_eq!(set.take_flag("header").unwrap(), Some(true)); + assert_eq!(set.take_flag("frozen").unwrap(), Some(false)); + assert_eq!(set.take_flag("absent").unwrap(), None); + } + + #[test] + fn test_take_format_normalizes() { + let mut set = new_set(&[("format", CopyOptionValue::Ident("CSV".into()))]); + assert_eq!(set.take_format().unwrap().as_deref(), Some("csv")); + } + + #[test] + fn test_reject_leftovers_lists_all() { + let mut set = new_set(&[ + ("precision", CopyOptionValue::String("ns".into())), + ("weird", CopyOptionValue::Flag), + ]); + set.take_string("precision").unwrap(); + let err = set.reject_leftovers("text").unwrap_err(); + assert!( + err.to_string().contains("WEIRD"), + "leftover option must be listed: {err}" + ); + assert!( + err.to_string().contains("not supported by format text"), + "{err}" + ); + // The consumed option is not listed. + assert!(!err.to_string().contains("PRECISION")); + } + + #[test] + fn test_raw_value_escape_hatch() { + let mut set = new_set(&[("precision", CopyOptionValue::Ident("ns".into()))]); + assert_eq!( + set.take_value("precision"), + Some(CopyOptionValue::Ident("ns".into())) + ); + } +} diff --git a/src/servers/src/postgres/copy_in/parse.rs b/src/servers/src/postgres/copy_in/parse.rs new file mode 100644 index 00000000000..8139960bfa4 --- /dev/null +++ b/src/servers/src/postgres/copy_in/parse.rs @@ -0,0 +1,443 @@ +// 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. + +//! Textual parser for `COPY ... FROM STDIN` statements. +//! +//! COPY FROM STDIN is parsed here instead of through sqlparser because the +//! `WITH (...)` option list is open: formats may accept options beyond +//! sqlparser's COPY grammar (for example a line-protocol format taking +//! `precision`), and unknown options must reach the format codec rather +//! than fail statement parsing. +//! +//! The parser commits to error reporting only after `FROM STDIN` has been +//! seen; anything else (like `COPY ... TO STDOUT`) is reported as +//! [`CopyInScan::NotCopyIn`] so the caller falls back to the regular SQL +//! path. + +use datafusion::sql::sqlparser::dialect::GenericDialect; +use datafusion::sql::sqlparser::tokenizer::{Token, Tokenizer}; +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. +/// +/// Unquoted identifiers are folded to lowercase and quoted ones keep their +/// case (`COPY Metrics` addresses the same table as `CREATE TABLE Metrics`, +/// namely `metrics`), matching `ParserContext::canonicalize_identifier` in +/// the SQL parser and PostgreSQL's own folding rule. +#[derive(Debug)] +pub(super) struct RawCopyFromStdin { + /// Table name parts, as written (`t`, `schema.t`, + /// `catalog.schema.t`); quoted names keep their case. + pub(super) table: Vec, + /// Requested column list; empty means all columns in table order. + pub(super) columns: Vec, + /// The open option list, handed to the resolved format codec. + pub(super) options: CopyOptionSet, +} + +pub(super) enum CopyInScan { + /// Not a `COPY ... FROM STDIN` statement. + NotCopyIn, + /// A complete copy-in statement with only trailing separators left. + CopyIn(RawCopyFromStdin), +} + +/// Scans a SQL string for a leading `COPY ... FROM STDIN` statement. +pub(super) fn scan_copy_from_stdin(sql: &str) -> PgWireResult { + // Fast path: the statement must start with the COPY keyword, possibly + // behind comments; only then is a full tokenization worth its cost. + let trimmed = sql.trim_start(); + let starts_like_copy = trimmed + .as_bytes() + .get(..4) + .is_some_and(|prefix| prefix.eq_ignore_ascii_case(b"copy")); + let starts_with_comment = trimmed.starts_with("--") || trimmed.starts_with("/*"); + if !starts_like_copy && !starts_with_comment { + return Ok(CopyInScan::NotCopyIn); + } + + let tokens = match Tokenizer::new(&GenericDialect {}, sql).tokenize() { + Ok(tokens) => tokens, + // Let the regular SQL path produce the parse error. + Err(_) => return Ok(CopyInScan::NotCopyIn), + }; + // The tokenizer keeps whitespace; the scanner works on tokens only. + let tokens: Vec = tokens + .into_iter() + .filter(|token| !matches!(token, Token::Whitespace(_))) + .collect(); + let mut cursor = Tokens { + tokens: &tokens, + pos: 0, + }; + + if !cursor.take_keyword("copy") { + return Ok(CopyInScan::NotCopyIn); + } + + let table = 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), + }; + + if !cursor.take_keyword("from") || !cursor.take_keyword("stdin") { + // A COPY with another direction or source; not ours to handle. + return Ok(CopyInScan::NotCopyIn); + } + + let mut options = CopyOptionSet::default(); + // PostgreSQL allows the WITH keyword to be omitted. + let has_options = if cursor.take_keyword("with") { + if !matches!(cursor.peek(), Some(Token::LParen)) { + return Err(bad_copy_data("expected '(' after WITH in COPY FROM STDIN")); + } + true + } else { + matches!(cursor.peek(), Some(Token::LParen)) + }; + if has_options { + // Consume the opening '('. + cursor.next(); + parse_option_list(&mut cursor, &mut options)?; + } + + // Only statement separators may follow. + while let Some(token) = cursor.next() { + if !matches!(token, Token::SemiColon) { + return Err(copy_in_must_run_alone_error()); + } + } + + Ok(CopyInScan::CopyIn(RawCopyFromStdin { + table, + columns, + options, + })) +} + +struct Tokens<'a> { + tokens: &'a [Token], + pos: usize, +} + +impl<'a> Tokens<'a> { + fn next(&mut self) -> Option<&'a Token> { + let token = self.tokens.get(self.pos)?; + self.pos += 1; + Some(token) + } + + fn peek(&mut self) -> Option<&'a Token> { + self.tokens.get(self.pos) + } + + /// 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, + } + } +} + +/// Returns the identifier with the SQL parser's canonicalization: quoted +/// identifiers keep their case, unquoted ones fold to lowercase. +fn expect_ident(cursor: &mut Tokens<'_>) -> Option { + match cursor.next() { + Some(Token::Word(w)) => Some(if w.quote_style.is_none() { + w.value.to_lowercase() + } else { + w.value.clone() + }), + _ => None, + } +} + +fn parse_object_name(cursor: &mut Tokens<'_>) -> Option> { + let mut parts = vec![expect_ident(cursor)?]; + while matches!(cursor.peek(), Some(Token::Period)) { + cursor.next(); + parts.push(expect_ident(cursor)?); + } + Some(parts) +} + +/// Parses an optional `(column, ...)` list. `None` means the statement is +/// not a copy-in shape (either no list, or a malformed list better left to +/// the regular parser to report). +fn parse_column_list(cursor: &mut Tokens<'_>) -> Option> { + if !matches!(cursor.peek(), Some(Token::LParen)) { + return Some(Vec::new()); + } + cursor.next(); + let mut columns = vec![expect_ident(cursor)?]; + loop { + match cursor.next() { + Some(Token::Comma) => columns.push(expect_ident(cursor)?), + Some(Token::RParen) => break, + _ => return None, + } + } + Some(columns) +} + +/// Parses `NAME [value] (, NAME [value])*` until the closing `)`. +/// +/// Option values are open-ended: a string literal, a number, an identifier +/// or a boolean word. Whether an option makes sense is decided by the +/// format codec, not here. +fn parse_option_list(cursor: &mut Tokens<'_>, options: &mut CopyOptionSet) -> PgWireResult<()> { + loop { + let key = match cursor.next() { + Some(Token::Word(w)) if w.quote_style.is_none() => w.value.to_lowercase(), + _ => return Err(bad_copy_data("expected an option name after '(' or ','")), + }; + let value = match cursor.peek() { + Some(Token::Comma) | Some(Token::RParen) => CopyOptionValue::Flag, + Some(Token::SingleQuotedString(s)) => { + let value = CopyOptionValue::String(s.clone()); + cursor.next(); + value + } + Some(Token::EscapedStringLiteral(s)) => { + // E'...' escape-string constants, e.g. DELIMITER E'\t': + // sqlparser's tokenizer already decoded the escapes, so + // the payload is used as-is, like SingleQuotedString. + let value = CopyOptionValue::String(s.clone()); + cursor.next(); + value + } + Some(Token::Number(n, _)) => { + let value = CopyOptionValue::Number(n.clone()); + cursor.next(); + value + } + Some(Token::Word(w)) => { + let value = if w.quote_style.is_none() { + match w.value.to_ascii_lowercase().as_str() { + "true" => CopyOptionValue::Boolean(true), + "false" => CopyOptionValue::Boolean(false), + _ => CopyOptionValue::Ident(w.value.clone()), + } + } else { + CopyOptionValue::Ident(w.value.clone()) + }; + cursor.next(); + value + } + _ => return Err(bad_copy_data("expected an option value or ',' or ')'")), + }; + options.insert(key, value)?; + match cursor.next() { + Some(Token::Comma) => continue, + Some(Token::RParen) => break, + _ => { + return Err(bad_copy_data( + "expected ',' or ')' in COPY FROM STDIN options", + )); + } + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn scan(sql: &str) -> PgWireResult> { + Ok(match scan_copy_from_stdin(sql)? { + CopyInScan::CopyIn(raw) => Some(raw), + CopyInScan::NotCopyIn => None, + }) + } + + #[test] + fn test_scan_basic() { + let raw = scan("COPY t FROM STDIN").unwrap().unwrap(); + assert_eq!(raw.table, vec!["t"]); + assert!(raw.columns.is_empty()); + } + + #[test] + fn test_scan_qualified_table_and_columns() { + let raw = scan("COPY db.public.Metrics (ts, \"Host\") FROM STDIN") + .unwrap() + .unwrap(); + assert_eq!(raw.table, vec!["db", "public", "metrics"]); + assert_eq!(raw.columns, vec!["ts", "Host"]); + } + + #[test] + fn test_scan_not_copy_in() { + assert!(scan("COPY (SELECT 1) TO STDOUT").unwrap().is_none()); + assert!( + scan("COPY t TO '/tmp/x.csv' WITH (FORMAT csv)") + .unwrap() + .is_none() + ); + assert!(scan("SELECT 1").unwrap().is_none()); + // GreptimeDB's own file-based COPY FROM. + assert!(scan("COPY t FROM '/tmp/x.csv'").unwrap().is_none()); + } + + #[test] + fn test_scan_leading_comment() { + let raw = scan("-- a comment\nCOPY t FROM STDIN").unwrap().unwrap(); + assert_eq!(raw.table, vec!["t"]); + } + + #[test] + fn test_scan_folds_unquoted_identifiers() { + // Unquoted identifiers fold to lowercase, like the SQL parser. + let raw = scan("copy Metrics (Host, ts) from stdin").unwrap().unwrap(); + assert_eq!(raw.table, vec!["metrics"]); + assert_eq!(raw.columns, vec!["host", "ts"]); + } + + #[test] + fn test_scan_preserves_quoted_identifiers() { + let raw = scan("COPY \"Metrics\" (\"Host\") FROM STDIN") + .unwrap() + .unwrap(); + assert_eq!(raw.table, vec!["Metrics"]); + assert_eq!(raw.columns, vec!["Host"]); + } + + #[test] + fn test_scan_options_kinds() { + let raw = scan( + "COPY t FROM STDIN WITH (FORMAT csv, DELIMITER ';', HEADER, \ + skip_bad_records true, precision 3, mode strict)", + ) + .unwrap() + .unwrap(); + let mut options = raw.options; + assert_eq!(options.take_format().unwrap().as_deref(), Some("csv")); + assert_eq!(options.take_single_byte("DELIMITER").unwrap(), Some(b';')); + assert_eq!(options.take_flag("HEADER").unwrap(), Some(true)); + assert_eq!( + options.take_value("skip_bad_records"), + Some(CopyOptionValue::Boolean(true)) + ); + assert_eq!( + options.take_value("precision"), + Some(CopyOptionValue::Number("3".into())) + ); + assert_eq!( + options.take_value("mode"), + Some(CopyOptionValue::Ident("strict".into())) + ); + } + + #[test] + fn test_scan_escaped_string_constants() { + // sqlparser's tokenizer already decodes E'...' escapes, so the + // payload is used as-is: E'\t' is a tab. + let mut raw = scan("COPY t FROM STDIN WITH (FORMAT csv, DELIMITER E'\\t')") + .unwrap() + .unwrap(); + assert_eq!( + raw.options.take_single_byte("DELIMITER").unwrap(), + Some(b'\t') + ); + + // E'\\' is a single literal backslash (a legal one-byte + // delimiter), NOT an escaped backslash-t. + let mut raw = scan("COPY t FROM STDIN WITH (FORMAT csv, DELIMITER E'\\\\')") + .unwrap() + .unwrap(); + assert_eq!( + raw.options.take_single_byte("DELIMITER").unwrap(), + Some(b'\\') + ); + + // E'\\t' decodes to TWO characters — backslash followed by t — + // and must be rejected as a multi-character delimiter instead of + // being unescaped a second time into a tab. + // The scanner leaves option validation to the codec; use the + // full resolution path to assert the delimiter rejection. + let err = match crate::postgres::copy_in::copy_from_stdin_from_sql( + "COPY t FROM STDIN WITH (FORMAT csv, DELIMITER E'\\\\t')", + ) { + Ok(Some(_)) => panic!("expected the multi-character delimiter to be rejected"), + Ok(None) => panic!("expected the statement to scan as copy-in"), + Err(e) => e, + }; + assert!( + err.to_string().contains("single one-byte character"), + "{err}" + ); + } + + #[test] + fn test_scan_semicolon_inside_option_value() { + // A quoted semicolon is not a statement separator. + let mut raw = scan("COPY t FROM STDIN WITH (NULL ';')").unwrap().unwrap(); + assert_eq!( + raw.options.take_string("NULL").unwrap().as_deref(), + Some(";") + ); + } + + #[test] + fn test_scan_trailing_statements_rejected() { + let err = scan("COPY t FROM STDIN; SELECT 1").unwrap_err(); + assert!(err.to_string().contains("alone")); + + // Trailing separators alone are fine. + assert!(scan("COPY t FROM STDIN;").unwrap().is_some()); + assert!(scan("COPY t FROM STDIN ; ;").unwrap().is_some()); + } + + #[test] + fn test_scan_options_without_with_keyword() { + let raw = scan("COPY t FROM STDIN (FORMAT csv, HEADER)") + .unwrap() + .unwrap(); + let mut options = raw.options; + assert_eq!(options.take_format().unwrap().as_deref(), Some("csv")); + assert_eq!(options.take_flag("HEADER").unwrap(), Some(true)); + assert!(options.reject_leftovers("csv").is_ok()); + } + + #[test] + fn test_scan_malformed_after_stdin() { + let err = scan("COPY t FROM STDIN WITH").unwrap_err(); + assert!(err.to_string().contains("'('"), "{err}"); + let err = scan("COPY t FROM STDIN WITH (").unwrap_err(); + assert!(err.to_string().contains("option"), "{err}"); + // A malformed column list never commits to the copy-in shape and + // falls through to the regular parser. + assert!(scan("COPY t (a,").unwrap().is_none()); + } + + #[test] + fn test_scan_redundant_option() { + let err = scan("COPY t FROM STDIN WITH (FORMAT csv, FORMAT text)").unwrap_err(); + assert!(err.to_string().contains("redundant FORMAT")); + } +} diff --git a/src/servers/src/postgres/handler.rs b/src/servers/src/postgres/handler.rs index 9d169811984..d03747f1b50 100644 --- a/src/servers/src/postgres/handler.rs +++ b/src/servers/src/postgres/handler.rs @@ -56,11 +56,14 @@ use crate::SqlPlan; use crate::error::{DataFusionSnafu, InferParameterTypesSnafu, Result}; use crate::postgres::types::*; use crate::postgres::utils::convert_err; -use crate::postgres::{PostgresServerHandlerInner, fixtures}; +use crate::postgres::{PostgresServerHandlerInner, copy_in, fixtures}; use crate::query_handler::sql::ServerSqlQueryHandlerRef; impl PostgresServerHandlerInner { - fn new_query_context(&self) -> QueryContextRef { + /// Creates a query context carrying this connection's batching + /// selection, so write paths can route through the pending-rows + /// batcher when it is enabled. + pub(super) fn new_query_context(&self) -> QueryContextRef { let mut ctx = self.session.new_query_context(); Arc::make_mut(&mut ctx).set_batching_enabled(self.batching_enabled); ctx @@ -87,8 +90,25 @@ impl SimpleQueryHandler for PostgresServerHandlerInner { return Ok(vec![Response::EmptyQuery]); } + // A leading COPY FROM STDIN takes over the connection for the copy + // sub-protocol; it is parsed textually so the WITH clause may carry + // options beyond sqlparser's COPY grammar. It must run alone. + if let Some(copy_from_stdin) = copy_in::copy_from_stdin_from_sql(query)? { + let response = self.begin_copy_in(copy_from_stdin).await?; + send_warning_opt(client, query_ctx).await?; + return Ok(vec![response]); + } + let parsed_query = self.query_parser.compatibility_parser.parse(query); + // A COPY FROM STDIN trailing other statements must be rejected; + // the textual scanner above only sees statement-leading COPYs. + if let Ok(statements) = &parsed_query + && statements.iter().any(copy_in::is_copy_from_stdin) + { + return Err(copy_in::copy_in_must_run_alone_error()); + } + let query = if let Ok(statements) = &parsed_query { // Comments, whitespace and empty statements also require EmptyQueryResponse. if statements.is_empty() { @@ -282,6 +302,36 @@ where _ => 0, }; + // Render timestamps with an explicit session-timezone offset so the + // exported text carries the instant unambiguously: the timezone-aware + // COPY IN decoder honors embedded offsets, keeping COPY OUT -> COPY IN + // round trips exact in non-UTC sessions (matching SQL INSERT parity). + // This touches both halves of the encoder: the arrow schema drives + // value conversion into a tz-aware DateTime, and the pg field type + // selects the offset-bearing text format. + let arrow_schema_override = copy_out_schema_with_timezone(&schema, &query_ctx); + let pg_fields = if arrow_schema_override.is_some() { + pg_fields + .into_iter() + .map(|mut field| { + if field.datatype() == &Type::TIMESTAMP { + field = FieldInfo::new( + field.name().to_string(), + None, + None, + Type::TIMESTAMPTZ, + field.format(), + ) + .with_type_size(field.type_size()) + .with_type_modifier(field.type_modifier()); + } + field + }) + .collect::>() + } else { + pg_fields + }; + let pg_schema = Arc::new(pg_fields); let num_columns = pg_schema.len(); @@ -297,7 +347,8 @@ where schema.clone(), recordbatches_stream, copy_encoder, - ); + ) + .with_arrow_schema_override(arrow_schema_override); let copy_stream: RowStream = Box::pin( row_stream @@ -336,6 +387,7 @@ impl DefaultQueryParser { pub struct PgSqlPlan { pub(crate) plan: SqlPlan, pub(crate) copy_to_stdout_format: Option, + pub(crate) copy_from_stdin: Option, } #[async_trait] @@ -360,6 +412,19 @@ impl QueryParser for DefaultQueryParser { return Ok(Some(PgSqlPlan { plan: SqlPlan::Shortcut(sql.to_string()), copy_to_stdout_format: None, + copy_from_stdin: None, + })); + } + + // A COPY FROM STDIN is parsed textually — not through sqlparser — + // so the WITH clause may carry options beyond sqlparser's COPY + // grammar; the extended query handler starts the copy sub-protocol + // when the statement is executed. + if let Some(copy_from_stdin) = copy_in::copy_from_stdin_from_sql(sql)? { + return Ok(Some(PgSqlPlan { + plan: SqlPlan::Shortcut(sql.to_string()), + copy_to_stdout_format: None, + copy_from_stdin: Some(copy_from_stdin), })); } @@ -370,7 +435,7 @@ impl QueryParser for DefaultQueryParser { } let first_stmt = statements.remove(0); let format = check_copy_to_stdout(&first_stmt); - (statement_to_query(&first_stmt), format) + (first_stmt.to_string(), format) } else { // bypass the error: it can run into error because of different // versions of sqlparser @@ -400,11 +465,13 @@ impl QueryParser for DefaultQueryParser { Ok(Some(PgSqlPlan { plan: SqlPlan::Plan(logical_plan, stmt), copy_to_stdout_format, + copy_from_stdin: None, })) } else { Ok(Some(PgSqlPlan { plan: SqlPlan::Statement(stmt, sql), copy_to_stdout_format, + copy_from_stdin: None, })) } } @@ -460,6 +527,14 @@ impl ExtendedQueryHandler for PostgresServerHandlerInner { let pg_sql_plan = &portal.statement.statement; let sql_plan = &pg_sql_plan.plan; + // A COPY FROM STDIN takes over the connection for the copy + // sub-protocol. + if let Some(copy_from_stdin) = &pg_sql_plan.copy_from_stdin { + let response = self.begin_copy_in(copy_from_stdin.clone()).await?; + send_warning_opt(client, query_ctx).await?; + return Ok(response); + } + let output = match sql_plan { SqlPlan::Empty => { // early return if query is empty @@ -782,6 +857,38 @@ impl ErrorHandler for PostgresServerHandlerInner { } } +/// Rewrites timestamp fields of the table's arrow schema to carry the +/// session timezone, so COPY OUT renders them with an explicit offset +/// (`2024-01-01 08:00:00+08:00`). Returns `None` when there is nothing +/// to rewrite (no naive timestamps or a UTC session leaving output as-is +/// is fine either way — offsets are always unambiguous on re-import). +fn copy_out_schema_with_timezone( + schema: &SchemaRef, + query_ctx: &QueryContextRef, +) -> Option> { + let timezone = query_ctx.timezone().to_string(); + let arrow_schema = schema.arrow_schema(); + let mut changed = false; + let fields = arrow_schema + .fields() + .iter() + .map(|field| match field.data_type() { + arrow::datatypes::DataType::Timestamp(unit, None) => { + changed = true; + field + .as_ref() + .clone() + .with_data_type(arrow::datatypes::DataType::Timestamp( + *unit, + Some(Arc::from(timezone.as_str())), + )) + } + _ => field.as_ref().clone(), + }) + .collect::>(); + changed.then(|| Arc::new(arrow::datatypes::Schema::new(fields))) +} + fn check_copy_to_stdout(statement: &SqlParserStatement) -> Option { if let SqlParserStatement::Copy { target, options, .. diff --git a/src/servers/src/postgres/server.rs b/src/servers/src/postgres/server.rs index 6e4aa942f9d..7e3b3e2b7bd 100644 --- a/src/servers/src/postgres/server.rs +++ b/src/servers/src/postgres/server.rs @@ -29,6 +29,7 @@ use tokio_rustls::TlsAcceptor; use crate::error::Result; use crate::postgres::{MakePostgresServerHandler, MakePostgresServerHandlerBuilder}; +use crate::query_handler::CopyInHandlerRef; use crate::query_handler::sql::ServerSqlQueryHandlerRef; use crate::server::{AbortableStream, BaseTcpServer, Server}; use crate::tls::ReloadableTlsServerConfig; @@ -44,9 +45,16 @@ pub struct PostgresServer { } impl PostgresServer { - /// Creates a new Postgres server with provided query_handler and async runtime + /// Creates a new Postgres server with provided query_handler and async runtime. + /// + /// `query_handler` and `copy_in_handler` are expected to be handles to + /// the same frontend instance (the SQL and COPY FROM STDIN facets of + /// one object); they are separate parameters only because they are + /// separate traits. + #[allow(clippy::too_many_arguments)] pub fn new( query_handler: ServerSqlQueryHandlerRef, + copy_in_handler: CopyInHandlerRef, force_tls: bool, tls_server_config: Arc, keep_alive_secs: u64, @@ -57,6 +65,7 @@ impl PostgresServer { let make_handler = Arc::new( MakePostgresServerHandlerBuilder::default() .query_handler(query_handler.clone()) + .copy_in_handler(copy_in_handler) .user_provider(user_provider.clone()) .force_tls(force_tls) .build() diff --git a/src/servers/src/postgres/types.rs b/src/servers/src/postgres/types.rs index 25cb94662c0..b0621d0a7b9 100644 --- a/src/servers/src/postgres/types.rs +++ b/src/servers/src/postgres/types.rs @@ -285,6 +285,11 @@ where schema: SchemaRef, record_batches: Pin>, encoder: S, + /// When set, replaces each batch's arrow schema for value encoding. + /// The COPY OUT path uses this to render timestamps with an explicit + /// session-timezone offset, so exported text re-imports identically + /// by the timezone-aware COPY IN decoder. + arrow_schema_override: Option>, } impl Stream for RecordBatchRowStream @@ -304,7 +309,10 @@ where return Poll::Ready(Some(Ok(vec![]))); } - let arrow_schema = record_batch.schema(); + let arrow_schema = self + .arrow_schema_override + .clone() + .unwrap_or_else(|| record_batch.schema()); let query_ctx = self.query_ctx.clone(); let pg_schema = self.pg_schema.clone(); let schema = self.schema.clone(); @@ -352,9 +360,19 @@ where schema, record_batches: Box::pin(record_batches), encoder, + arrow_schema_override: None, } } + /// See [`RecordBatchRowStream::arrow_schema_override`]. + pub(crate) fn with_arrow_schema_override( + mut self, + schema: Option>, + ) -> Self { + self.arrow_schema_override = schema; + self + } + fn encode_row( query_ctx: &QueryContextRef, pg_schema: &Arc>, @@ -2464,6 +2482,7 @@ mod test { PgSqlPlan { plan: SqlPlan::Empty, copy_to_stdout_format: None, + copy_from_stdin: None, }, client_param_types, )); diff --git a/src/servers/src/postgres/types/error.rs b/src/servers/src/postgres/types/error.rs index 00d1bbbb1c1..86317b91570 100644 --- a/src/servers/src/postgres/types/error.rs +++ b/src/servers/src/postgres/types/error.rs @@ -67,6 +67,9 @@ pub enum PgErrorCode { /// connection_exception #[snafu(display("connection_exception"))] Ec08000 = 400, + /// protocol_violation + #[snafu(display("protocol_violation"))] + Ec08P01 = 401, // === End of Class 08 — Connection Exception ===== // === Begin of Class 09 — Triggered Action Exception === @@ -130,6 +133,12 @@ pub enum PgErrorCode { /// invalid_parameter_value #[snafu(display("invalid_parameter_value"))] Ec22023 = 1401, + /// invalid_text_representation + #[snafu(display("invalid_text_representation"))] + Ec22P02 = 1402, + /// bad_copy_file_format + #[snafu(display("bad_copy_file_format"))] + Ec22P04 = 1403, // === End of Class 22 — Data Exception ===== // === Begin of Class 23 — Integrity Constraint Violation === @@ -295,6 +304,9 @@ pub enum PgErrorCode { /// operator_intervention #[snafu(display("operator_intervention"))] Ec57000 = 3600, + /// query_canceled + #[snafu(display("query_canceled"))] + Ec57014 = 3602, /// cannot_connect_now #[snafu(display("cannot_connect_now"))] diff --git a/src/servers/src/query_handler.rs b/src/servers/src/query_handler.rs index 12a56076d6f..57f55855ac1 100644 --- a/src/servers/src/query_handler.rs +++ b/src/servers/src/query_handler.rs @@ -43,6 +43,7 @@ use otel_arrow_rust::proto::opentelemetry::collector::metrics::v1::ExportMetrics use pipeline::{GreptimePipelineParams, Pipeline, PipelineInfo, PipelineVersion, PipelineWay}; use serde_json::Value; use session::context::{QueryContext, QueryContextRef}; +use table::metadata::TableInfoRef; #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] pub struct DashboardDefinition { @@ -290,3 +291,33 @@ pub trait JaegerQueryHandler { query_params: QueryTraceParams, ) -> Result; } + +/// Handler for the PostgreSQL `COPY ... FROM STDIN` (copy-in) sub-protocol. +/// +/// The servers layer owns the wire protocol and data parsing; this trait is +/// the boundary to catalog lookup and ingestion, implemented by the +/// frontend instance. +#[async_trait] +pub trait CopyInHandler: Send + Sync { + /// Resolves the target table of a COPY and returns its metadata, or + /// `None` when the table does not exist. Implementations should also + /// enforce the write permission for the table here so unauthorized + /// copies fail before any data is transferred. + async fn copy_in_table( + &self, + catalog: &str, + schema: &str, + table: &str, + query_ctx: QueryContextRef, + ) -> Result>; + + /// Inserts a parsed batch produced by a COPY stream. Returns the + /// number of affected rows. + async fn copy_in_insert( + &self, + requests: RowInsertRequests, + query_ctx: QueryContextRef, + ) -> Result; +} + +pub type CopyInHandlerRef = Arc; diff --git a/src/servers/tests/mod.rs b/src/servers/tests/mod.rs index 699f3f76f1e..d78f57b7b83 100644 --- a/src/servers/tests/mod.rs +++ b/src/servers/tests/mod.rs @@ -14,13 +14,14 @@ use std::sync::Arc; +use api::v1::RowInsertRequests; use api::v1::greptime_request::Request; use api::v1::query_request::Query; use async_trait::async_trait; use catalog::memory::MemoryCatalogManager; use catalog::system_schema::SystemSchemaProvider; use catalog::system_schema::pg_catalog::PGCatalogProvider; -use catalog::{CatalogManager, RegisterTableRequest}; +use catalog::{CatalogManager, CatalogManagerRef, RegisterTableRequest}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, PG_CATALOG_NAME}; use common_error::ext::BoxedError; use common_grpc::flight::do_put::DoPutResponse; @@ -32,6 +33,7 @@ use query::parser::{PromQuery, QueryStatement}; use query::query_engine::DescribeResult; use query::{QueryEngineFactory, QueryEngineRef}; use servers::error::{ExecuteQuerySnafu, NotSupportedSnafu, Result}; +use servers::query_handler::CopyInHandler; use servers::query_handler::grpc::GrpcQueryHandler; use servers::query_handler::sql::{ServerSqlQueryHandlerRef, SqlQueryHandler}; use session::context::QueryContextRef; @@ -39,6 +41,7 @@ use snafu::{ResultExt, ensure}; use sql::parser::{ParseOptions, ParserContext}; use sql::statements::statement::Statement; use table::TableRef; +use table::metadata::TableInfoRef; mod http; mod interceptor; @@ -54,6 +57,64 @@ impl DummyInstance { fn new(query_engine: QueryEngineRef) -> Self { Self { query_engine } } + + pub fn catalog_manager(&self) -> CatalogManagerRef { + self.query_engine.engine_state().catalog_manager().clone() + } +} + +/// A [`CopyInHandler`] test double: resolves tables through the catalog +/// manager and records the insert requests instead of writing them. +pub struct RecordingCopyInHandler { + catalog_manager: CatalogManagerRef, + pub requests: std::sync::Mutex>, +} + +impl RecordingCopyInHandler { + pub fn new(catalog_manager: CatalogManagerRef) -> Self { + Self { + catalog_manager, + requests: Default::default(), + } + } +} + +#[async_trait] +impl CopyInHandler for RecordingCopyInHandler { + async fn copy_in_table( + &self, + catalog: &str, + schema: &str, + table: &str, + _query_ctx: QueryContextRef, + ) -> Result> { + Ok(self + .catalog_manager + .table(catalog, schema, table, None) + .await + .unwrap() + .map(|table| table.table_info())) + } + + async fn copy_in_insert( + &self, + requests: RowInsertRequests, + _query_ctx: QueryContextRef, + ) -> Result { + let rows: usize = requests + .inserts + .iter() + .map(|insert| { + insert + .rows + .as_ref() + .map(|rows| rows.rows.len()) + .unwrap_or(0) + }) + .sum(); + self.requests.lock().unwrap().push(requests); + Ok(Output::new_with_affected_rows(rows)) + } } #[async_trait] diff --git a/src/servers/tests/postgres/mod.rs b/src/servers/tests/postgres/mod.rs index 23971145a86..1bc5e79c236 100644 --- a/src/servers/tests/postgres/mod.rs +++ b/src/servers/tests/postgres/mod.rs @@ -229,7 +229,18 @@ fn create_postgres_server_with_user_provider( tls: TlsOption, user_provider: Option, ) -> Result> { + create_postgres_server_inner(table, tls, user_provider).map(|(server, _)| server) +} + +fn create_postgres_server_inner( + table: TableRef, + tls: TlsOption, + user_provider: Option, +) -> Result<(Box, Arc)> { let instance = Arc::new(create_testing_instance(table)); + let copy_in_handler = Arc::new(crate::RecordingCopyInHandler::new( + instance.catalog_manager(), + )); let io_runtime = RuntimeBuilder::default() .worker_threads(4) .thread_name("postgres-io-handlers") @@ -241,15 +252,19 @@ fn create_postgres_server_with_user_provider( .expect("Failed to load certificates and keys"), ); - Ok(Box::new(PostgresServer::new( - instance, - tls.should_force_tls(), - tls_server_config, - 0, - io_runtime, - user_provider, - None, - ))) + Ok(( + Box::new(PostgresServer::new( + instance, + copy_in_handler.clone(), + tls.should_force_tls(), + tls_server_config, + 0, + io_runtime, + user_provider, + None, + )), + copy_in_handler, + )) } async fn start_test_server_with_user_provider( @@ -1149,3 +1164,237 @@ async fn do_simple_query_with_secure_server( do_simple_query(server_tls, client_tls).await } + +// --------------------------------------------------------------------------- +// COPY FROM STDIN tests +// --------------------------------------------------------------------------- + +use futures::SinkExt; + +fn copy_in_test_table() -> TableRef { + use common_recordbatch::RecordBatch; + use datatypes::prelude::ConcreteDataType; + use datatypes::schema::ColumnSchema; + use datatypes::types::TimestampType; + use datatypes::vectors::{Float64Vector, StringVector, TimestampMillisecondVector, VectorRef}; + + let column_schemas = vec![ + ColumnSchema::new( + "ts", + ConcreteDataType::Timestamp(TimestampType::Millisecond(Default::default())), + false, + ) + .with_time_index(true), + ColumnSchema::new("host", ConcreteDataType::string_datatype(), false), + ColumnSchema::new("val", ConcreteDataType::float64_datatype(), true), + ColumnSchema::new("note", ConcreteDataType::string_datatype(), true), + ]; + let schema = Arc::new(datatypes::schema::Schema::new(column_schemas)); + let columns: Vec = vec![ + Arc::new(TimestampMillisecondVector::from_slice([0i64])), + Arc::new(StringVector::from_slice(&["localhost"])), + Arc::new(Float64Vector::from_slice([1.0f64])), + Arc::new(StringVector::from_slice(&["seed"])), + ]; + let recordbatch = RecordBatch::new(schema, columns).unwrap(); + MemTable::table("metrics", recordbatch) +} + +async fn start_copy_in_test_server( + table: TableRef, +) -> Result<(Box, u16, Arc)> { + common_telemetry::init_default_ut_logging(); + let _ = install_default_crypto_provider(); + let (mut server, handler) = create_postgres_server_inner(table, Default::default(), None)?; + let listening = "127.0.0.1:0".parse::().unwrap(); + server.start(listening).await.unwrap(); + let port = server.bind_addr().unwrap().port(); + Ok((server, port, handler)) +} + +fn recorded_rows(handler: &crate::RecordingCopyInHandler) -> Vec { + handler + .requests + .lock() + .unwrap() + .iter() + .flat_map(|request| { + request.inserts.iter().flat_map(|insert| { + insert + .rows + .clone() + .map(|rows| rows.rows) + .unwrap_or_default() + }) + }) + .collect() +} + +async fn feed_copy_in( + client: &Client, + sql: &str, + chunks: Vec, +) -> std::result::Result { + let mut sink = Box::pin(client.copy_in(sql).await?); + for chunk in chunks { + sink.as_mut().send(chunk).await?; + } + sink.as_mut().finish().await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_copy_in_csv() -> Result<()> { + let (server, port, handler) = start_copy_in_test_server(copy_in_test_table()).await?; + let client = create_plain_connection(port, false).await.unwrap(); + + let rows = feed_copy_in( + &client, + "COPY metrics FROM STDIN WITH (FORMAT csv)", + vec![bytes::Bytes::from_static( + b"2023-11-14 22:13:20,host1,1.5,hello\n2023-11-14 22:13:21,host2,,\"quoted, comma\"\n", + )], + ) + .await + .unwrap(); + assert_eq!(rows, 2); + + let rows = recorded_rows(&handler); + assert_eq!(rows.len(), 2); + assert_eq!( + rows[0].values[1].value_data, + Some(api::v1::value::ValueData::StringValue("host1".to_string())) + ); + assert_eq!( + rows[0].values[2].value_data, + Some(api::v1::value::ValueData::F64Value(1.5)) + ); + // Unquoted empty field is NULL, quoted field keeps its content. + assert_eq!(rows[1].values[2].value_data, None); + assert_eq!( + rows[1].values[3].value_data, + Some(api::v1::value::ValueData::StringValue( + "quoted, comma".to_string() + )) + ); + + server.shutdown().await.unwrap(); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_copy_in_column_list_and_batches() -> Result<()> { + let (server, port, handler) = start_copy_in_test_server(copy_in_test_table()).await?; + let client = create_plain_connection(port, false).await.unwrap(); + + let mut chunks = Vec::new(); + let mut payload = String::new(); + for i in 0..70000u64 { + payload.push_str(&format!("{i},2023-11-14 22:13:20,host{i},0.5\n")); + if (i + 1) % 100 == 0 { + chunks.push(bytes::Bytes::from(std::mem::take(&mut payload))); + } + } + let rows = feed_copy_in( + &client, + "COPY metrics (note, ts, host, val) FROM STDIN (FORMAT csv)", + chunks, + ) + .await + .unwrap(); + assert_eq!(rows, 70000); + + let recorded = handler.requests.lock().unwrap().clone(); + // The 70000 rows must have been flushed in multiple batches of at + // most COPY_IN_FLUSH_ROWS. + assert!(recorded.len() > 1); + let total: usize = recorded + .iter() + .map(|request| { + request.inserts[0] + .rows + .as_ref() + .map(|rows| rows.rows.len()) + .unwrap_or(0) + }) + .sum(); + assert_eq!(total, 70000); + + let rows = recorded_rows(&handler); + // The column list reorders the fields: first field lands in `note`. + assert_eq!( + rows[0].values[0].value_data, + Some(api::v1::value::ValueData::StringValue("0".to_string())) + ); + + server.shutdown().await.unwrap(); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_copy_in_unknown_table() -> Result<()> { + let (server, port, _handler) = start_copy_in_test_server(copy_in_test_table()).await?; + let client = create_plain_connection(port, false).await.unwrap(); + + let err = match client + .copy_in::<_, bytes::Bytes>("COPY not_exist FROM STDIN WITH (FORMAT csv)") + .await + { + Ok(_) => panic!("expected COPY to unknown table to fail"), + Err(e) => e, + }; + let message = err.as_db_error().expect("expected db error").message(); + assert!(message.contains("does not exist"), "{message}"); + + server.shutdown().await.unwrap(); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_copy_in_bad_value() -> Result<()> { + let (server, port, handler) = start_copy_in_test_server(copy_in_test_table()).await?; + let client = create_plain_connection(port, false).await.unwrap(); + + let err = match feed_copy_in( + &client, + "COPY metrics FROM STDIN WITH (FORMAT csv)", + vec![bytes::Bytes::from_static(b"not_a_timestamp,host1,1.5,x\n")], + ) + .await + { + Ok(_) => panic!("expected COPY to fail"), + Err(e) => e, + }; + let message = err.as_db_error().expect("expected db error").message(); + assert!(message.contains("failed to parse"), "{message}"); + // The failed batch must not be written. + assert!(handler.requests.lock().unwrap().is_empty()); + + server.shutdown().await.unwrap(); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_copy_in_multi_statement_rejected() -> Result<()> { + let (server, port, _handler) = start_copy_in_test_server(copy_in_test_table()).await?; + let client = create_plain_connection(port, false).await.unwrap(); + + let err = match feed_copy_in( + &client, + "COPY metrics FROM STDIN WITH (FORMAT csv); SELECT 1", + vec![bytes::Bytes::from_static(b"1,2,3,4\n")], + ) + .await + { + Ok(_) => panic!("expected multi-statement COPY to fail"), + Err(e) => e, + }; + let message = err + .as_db_error() + .expect("expected db error") + .message() + .to_string(); + assert!(message.contains("alone"), "{message}"); + + server.shutdown().await.unwrap(); + Ok(()) +} diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index 4ddd9ee5a46..ebf009a5211 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -77,6 +77,7 @@ uuid.workspace = true [dev-dependencies] base64.workspace = true +bytes.workspace = true catalog = { workspace = true, features = ["testing"] } chrono.workspace = true clap.workspace = true diff --git a/tests-integration/src/test_util.rs b/tests-integration/src/test_util.rs index 0729e6ea96b..3cc820619c6 100644 --- a/tests-integration/src/test_util.rs +++ b/tests-integration/src/test_util.rs @@ -1205,7 +1205,8 @@ pub async fn setup_pg_server_with_slow_query_threshold( ); let mut pg_server = Box::new(PostgresServer::new( - fe_instance_ref, + fe_instance_ref.clone(), + fe_instance_ref.clone(), opts.tls.should_force_tls(), tls_server_config, 0, @@ -1248,7 +1249,8 @@ pub async fn setup_pg_server_with_user_provider( ); let mut pg_server = Box::new(PostgresServer::new( - fe_instance_ref, + fe_instance_ref.clone(), + fe_instance_ref.clone(), opts.tls.should_force_tls(), tls_server_config, 0, @@ -1303,6 +1305,7 @@ pub async fn setup_pg_server_with_prom_native_histogram( ); let mut pg_server = Box::new(PostgresServer::new( + instance.fe_instance().clone(), instance.fe_instance().clone(), opts.tls.should_force_tls(), tls_server_config, diff --git a/tests-integration/src/tests.rs b/tests-integration/src/tests.rs index 4a5b2f0c6c5..1984b849a7e 100644 --- a/tests-integration/src/tests.rs +++ b/tests-integration/src/tests.rs @@ -17,6 +17,7 @@ mod instance_kafka_wal_test; mod instance_noop_wal_test; mod instance_test; mod nested_merge_scan_capability_test; +mod postgres_copy_in_test; mod promql_test; mod reconcile_table; mod remote_dyn_filter_test; diff --git a/tests-integration/src/tests/postgres_copy_in_test.rs b/tests-integration/src/tests/postgres_copy_in_test.rs new file mode 100644 index 00000000000..82e076e7a65 --- /dev/null +++ b/tests-integration/src/tests/postgres_copy_in_test.rs @@ -0,0 +1,521 @@ +// 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. + +//! End-to-end tests for `COPY ... FROM STDIN` over the PostgreSQL protocol, +//! covering the full ingestion path (parser -> insert -> storage). + +use bytes::Bytes; +use common_telemetry::init_default_ut_logging; +use futures::SinkExt; +use tokio_postgres::NoTls; + +use crate::test_util::{StorageType, setup_pg_server}; + +async fn connect(port: u16) -> tokio_postgres::Client { + let url = format!("host=127.0.0.1 port={port} dbname=public connect_timeout=5"); + let (client, conn) = tokio_postgres::connect(&url, NoTls).await.unwrap(); + let _handle = tokio::spawn(conn); + client +} + +async fn feed_copy_in( + client: &tokio_postgres::Client, + sql: &str, + chunks: Vec, +) -> std::result::Result { + let mut sink = Box::pin(client.copy_in(sql).await?); + for chunk in chunks { + sink.as_mut().send(chunk).await?; + } + sink.as_mut().finish().await +} + +fn host_of(message: &tokio_postgres::SimpleQueryMessage, column: usize) -> String { + let tokio_postgres::SimpleQueryMessage::Row(row) = message else { + panic!("expected row, got {message:?}"); + }; + // simple_query renders NULL as `None`, not empty string. + row.get(column).map(str::to_string).unwrap_or_default() +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_in_end_to_end() { + init_default_ut_logging(); + let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_e2e").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + client + .simple_query( + "CREATE TABLE copy_demo (host STRING, val DOUBLE, note STRING, ts TIMESTAMP, TIME INDEX(ts))", + ) + .await + .unwrap(); + + // CSV format with quoted field and NULL (unquoted empty) value. + let rows = feed_copy_in( + &client, + "COPY copy_demo FROM STDIN WITH (FORMAT csv)", + vec![Bytes::from_static( + b"host1,1.5,\"hello, world\",2023-11-14 22:13:20\nhost2,,x,2023-11-14 22:13:21\n", + )], + ) + .await + .unwrap(); + assert_eq!(rows, 2); + + // A column list targeting a subset. + let rows = feed_copy_in( + &client, + "COPY copy_demo (note, ts) FROM STDIN WITH (FORMAT csv)", + vec![Bytes::from_static( + b"from_column_list,2023-11-14 22:13:24\n", + )], + ) + .await + .unwrap(); + assert_eq!(rows, 1); + + // Verify data written through the real insert path. + let messages = client + .simple_query("SELECT host, val, note FROM copy_demo ORDER BY ts") + .await + .unwrap(); + let rows: Vec<_> = messages + .iter() + .filter(|m| matches!(m, tokio_postgres::SimpleQueryMessage::Row(_))) + .collect(); + assert_eq!(rows.len(), 3); + assert_eq!(host_of(rows[0], 0), "host1"); + assert_eq!(host_of(rows[0], 2), "hello, world"); + assert_eq!(host_of(rows[1], 0), "host2"); + // NULL val and note round-trip as NULL. + assert_eq!(host_of(rows[2], 0), ""); + assert_eq!(host_of(rows[2], 2), "from_column_list"); + + server.shutdown().await.unwrap(); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_in_error_keeps_table_clean() { + init_default_ut_logging(); + let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_error").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + client + .simple_query( + "CREATE TABLE copy_err (host STRING, val DOUBLE, ts TIMESTAMP, TIME INDEX(ts))", + ) + .await + .unwrap(); + + // Bad double value aborts the copy. + let err = match feed_copy_in( + &client, + "COPY copy_err FROM STDIN WITH (FORMAT csv)", + vec![Bytes::from_static( + b"host1,not_a_double,2023-11-14 22:13:20\n", + )], + ) + .await + { + Ok(_) => panic!("expected COPY to fail"), + Err(e) => e, + }; + assert!( + err.as_db_error() + .expect("expected db error") + .message() + .contains("failed to parse"), + "{err}" + ); + + // A bare COPY (PostgreSQL's text default) fails with guidance. + let err = match client.copy_in::<_, Bytes>("COPY copy_err FROM STDIN").await { + Ok(_) => panic!("expected bare COPY to fail"), + Err(e) => e, + }; + let message = err + .as_db_error() + .expect("expected db error") + .message() + .to_string(); + assert!(message.contains("requires an explicit format"), "{message}"); + + // The connection stays usable and no rows were written. + let messages = client + .simple_query("SELECT host FROM copy_err") + .await + .unwrap(); + assert!( + messages + .iter() + .all(|m| !matches!(m, tokio_postgres::SimpleQueryMessage::Row(_))), + "expected no rows, got {messages:?}" + ); + + server.shutdown().await.unwrap(); +} + +/// Exercises the CSV format's options end-to-end through a Postgres +/// client driver: HEADER, custom delimiter, quoted fields with embedded +/// newlines and escaped quotes, CRLF record separators and NULL fields. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_in_csv_options() { + init_default_ut_logging(); + let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_csv_options").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + client + .simple_query( + "CREATE TABLE copy_csv_opts (host STRING, val DOUBLE, note STRING, ts TIMESTAMP, TIME INDEX(ts))", + ) + .await + .unwrap(); + + let rows = feed_copy_in( + &client, + // WITH omitted, matching PostgreSQL; HEADER skips the first record. + "COPY copy_csv_opts FROM STDIN (FORMAT csv, HEADER, DELIMITER ';')", + vec![Bytes::from_static( + b"host;val;note;ts\r\n\ + h1;1.5;\"line one\nline two\";2023-11-14 22:13:20\r\n\ + h2;;\"say \"\"hi\"\"\";2023-11-14 22:13:21\r\n", + )], + ) + .await + .unwrap(); + assert_eq!(rows, 2); + + // Verify the parsed values came through the real insert path. + let messages = client + .simple_query("SELECT host, val, note FROM copy_csv_opts ORDER BY ts") + .await + .unwrap(); + let rows: Vec<_> = messages + .iter() + .filter(|m| matches!(m, tokio_postgres::SimpleQueryMessage::Row(_))) + .collect(); + assert_eq!(rows.len(), 2); + assert_eq!(host_of(rows[0], 0), "h1"); + assert_eq!(host_of(rows[0], 1), "1.5"); + // Quoted field with an embedded newline survives the round trip. + assert_eq!(host_of(rows[0], 2), "line one\nline two"); + assert_eq!(host_of(rows[1], 0), "h2"); + // Unquoted empty field is NULL; escaped quotes decode to one quote. + assert_eq!(host_of(rows[1], 1), ""); + assert_eq!(host_of(rows[1], 2), "say \"hi\""); + + server.shutdown().await.unwrap(); +} + +/// COPY OUT → COPY IN round trip: PostgreSQL's COPY OUT emits `t`/`f` +/// for booleans, which COPY IN must accept back. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_in_bool_round_trip() { + init_default_ut_logging(); + let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_bool_rt").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + client + .simple_query( + "CREATE TABLE bool_src (flag BOOLEAN, ts TIMESTAMP, TIME INDEX(ts)); + CREATE TABLE bool_dst (flag BOOLEAN, ts TIMESTAMP, TIME INDEX(ts)); + INSERT INTO bool_src VALUES (true, '2024-01-01 00:00:00'), (false, '2024-01-01 00:00:01'), (true, '2024-01-01 00:00:02');", + ) + .await + .unwrap(); + + // COPY OUT: collect the exported bytes. + let mut exported = Vec::new(); + let mut out = Box::pin( + client + .copy_out( + "COPY (SELECT flag, ts FROM bool_src ORDER BY ts) TO STDOUT WITH (FORMAT csv)", + ) + .await + .unwrap(), + ); + use futures::StreamExt; + while let Some(chunk) = out.as_mut().next().await { + exported.extend_from_slice(&chunk.unwrap()); + } + // The wire encoder emits PostgreSQL boolean spellings. + let text = String::from_utf8(exported.clone()).unwrap(); + assert!(text.contains("t,"), "COPY OUT should emit t/f, got: {text}"); + assert!(text.contains("f,"), "COPY OUT should emit t/f, got: {text}"); + + // COPY IN: the exact bytes must import back. + let rows = feed_copy_in( + &client, + "COPY bool_dst FROM STDIN WITH (FORMAT csv)", + vec![Bytes::from(exported)], + ) + .await + .unwrap(); + assert_eq!(rows, 3); + + let messages = client + .simple_query("SELECT flag FROM bool_dst ORDER BY ts") + .await + .unwrap(); + let flags: Vec<&str> = messages + .iter() + .filter_map(|m| match m { + tokio_postgres::SimpleQueryMessage::Row(row) => row.get(0), + _ => None, + }) + .collect(); + // The simple-query path also renders booleans in the t/f spelling. + assert_eq!(flags, vec!["t", "f", "t"]); + + server.shutdown().await.unwrap(); +} + +/// A failed COPY must release its state so a later COPY on the same +/// connection works; this guards the error-path state drop. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_in_failed_copy_then_new_copy() { + init_default_ut_logging(); + let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_fail_reuse").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + client + .simple_query( + "CREATE TABLE copy_reuse (host STRING, val DOUBLE, ts TIMESTAMP, TIME INDEX(ts))", + ) + .await + .unwrap(); + + // A copy that fails at CopyDone (bad double parsed at flush time). + let err = match feed_copy_in( + &client, + "COPY copy_reuse FROM STDIN WITH (FORMAT csv)", + vec![Bytes::from_static( + b"host1,not_a_double,2024-01-01 00:00:00\n", + )], + ) + .await + { + Ok(_) => panic!("expected COPY to fail"), + Err(e) => e, + }; + assert!(err.as_db_error().is_some(), "{err}"); + + // The connection is usable and a NEW copy succeeds on it. + let rows = feed_copy_in( + &client, + "COPY copy_reuse FROM STDIN WITH (FORMAT csv)", + vec![Bytes::from_static(b"host1,1.5,2024-01-01 00:00:01\n")], + ) + .await + .unwrap(); + assert_eq!(rows, 1); + + let messages = client + .simple_query("SELECT host, val FROM copy_reuse") + .await + .unwrap(); + let found: Vec<_> = messages + .iter() + .filter(|m| matches!(m, tokio_postgres::SimpleQueryMessage::Row(_))) + .collect(); + assert_eq!(found.len(), 1); + + server.shutdown().await.unwrap(); +} + +/// COPY/INSERT parity in a non-UTC session: offset-free timestamp strings +/// must be interpreted in the session timezone by both paths. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_in_session_timezone_parity() { + init_default_ut_logging(); + let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_tz").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + client + .simple_query( + "CREATE TABLE copy_tz (case_id STRING PRIMARY KEY, ts TIMESTAMP, TIME INDEX(ts))", + ) + .await + .unwrap(); + + client + .simple_query("SET TIMEZONE TO 'Asia/Shanghai'") + .await + .unwrap(); + + client + .simple_query("INSERT INTO copy_tz VALUES ('insert', '2024-01-01 00:00:00')") + .await + .unwrap(); + let rows = feed_copy_in( + &client, + "COPY copy_tz (case_id, ts) FROM STDIN WITH (FORMAT csv)", + vec![Bytes::from_static(b"copy,2024-01-01 00:00:00\n")], + ) + .await + .unwrap(); + assert_eq!(rows, 1); + + let messages = client + .simple_query("SELECT case_id, CAST(ts AS BIGINT) FROM copy_tz ORDER BY case_id") + .await + .unwrap(); + let results: Vec<(String, String)> = messages + .iter() + .filter_map(|m| match m { + tokio_postgres::SimpleQueryMessage::Row(row) => { + Some((row.get(0)?.to_string(), row.get(1)?.to_string())) + } + _ => None, + }) + .collect(); + assert_eq!(results.len(), 2, "{results:?}"); + assert_eq!(results[0].0, "copy"); + assert_eq!(results[1].0, "insert"); + // The whole point: both paths store the same instant. + assert_eq!( + results[0].1, results[1].1, + "COPY and INSERT must agree on offset-free timestamps: {results:?}" + ); + // And it is the Asia/Shanghai interpretation, not UTC. + assert_eq!(results[0].1, "1704038400000"); + + server.shutdown().await.unwrap(); +} + +/// COPY OUT -> COPY IN must preserve timestamp instants in non-UTC +/// sessions: COPY OUT renders timestamps with an explicit session-tz +/// offset, which the timezone-aware COPY IN decoder honors. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_out_in_timestamp_round_trip() { + init_default_ut_logging(); + let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_out_in_tz").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + client + .simple_query( + "CREATE TABLE tz_src (host STRING, ts TIMESTAMP, TIME INDEX(ts)); + CREATE TABLE tz_dst (host STRING, ts TIMESTAMP, TIME INDEX(ts)); + INSERT INTO tz_src VALUES ('a', '2024-01-01 00:00:00'), ('b', '2024-07-01 12:30:00');", + ) + .await + .unwrap(); + + client + .simple_query("SET TIMEZONE TO 'Asia/Shanghai'") + .await + .unwrap(); + + let mut exported = Vec::new(); + let mut out = Box::pin( + client + .copy_out("COPY (SELECT host, ts FROM tz_src ORDER BY ts) TO STDOUT WITH (FORMAT csv)") + .await + .unwrap(), + ); + use futures::StreamExt; + while let Some(chunk) = out.as_mut().next().await { + exported.extend_from_slice(&chunk.unwrap()); + } + // The export now carries explicit offsets. + let text = String::from_utf8(exported.clone()).unwrap(); + assert!( + text.contains("+08"), + "COPY OUT should render timestamps with the session-tz offset, got: {text}" + ); + + let rows = feed_copy_in( + &client, + "COPY tz_dst (host, ts) FROM STDIN WITH (FORMAT csv)", + vec![Bytes::from(exported)], + ) + .await + .unwrap(); + assert_eq!(rows, 2); + + // Epochs must match the source exactly. + let compare = client + .simple_query( + "SELECT (SELECT CAST(ts AS BIGINT) FROM tz_src WHERE host='a') = (SELECT CAST(ts AS BIGINT) FROM tz_dst WHERE host='a') \ + AND (SELECT CAST(ts AS BIGINT) FROM tz_src WHERE host='b') = (SELECT CAST(ts AS BIGINT) FROM tz_dst WHERE host='b')", + ) + .await + .unwrap(); + let ok: Vec<_> = compare + .iter() + .filter_map(|m| match m { + tokio_postgres::SimpleQueryMessage::Row(row) => row.get(0).map(str::to_string), + _ => None, + }) + .collect(); + assert_eq!(ok, vec!["t"], "epochs must round-trip exactly"); + + server.shutdown().await.unwrap(); +} + +/// Empty fields — quoted or not — import as NULL, matching the +/// file-based COPY ... FROM 'file.csv' behavior exactly. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_postgres_copy_in_empty_field_is_null() { + init_default_ut_logging(); + let (_guard, server) = setup_pg_server(StorageType::File, "pg_copy_in_empty").await; + let port = server.bind_addr().unwrap().port(); + let client = connect(port).await; + + client + .simple_query("CREATE TABLE empty_str (note STRING, ts TIMESTAMP, TIME INDEX(ts))") + .await + .unwrap(); + + let rows = feed_copy_in( + &client, + "COPY empty_str (note, ts) FROM STDIN WITH (FORMAT csv)", + vec![Bytes::from_static( + b"plain,2024-01-01 00:00:00\n\"\",2024-01-01 00:00:01\n,2024-01-01 00:00:02\n", + )], + ) + .await + .unwrap(); + assert_eq!(rows, 3); + + let messages = client + .simple_query("SELECT note IS NULL AS is_null, note FROM empty_str ORDER BY ts") + .await + .unwrap(); + let rows: Vec<(bool, Option)> = messages + .iter() + .filter_map(|m| match m { + tokio_postgres::SimpleQueryMessage::Row(row) => { + Some((row.get(0) == Some("t"), row.get(1).map(str::to_string))) + } + _ => None, + }) + .collect(); + assert_eq!(rows.len(), 3); + assert_eq!(rows[0], (false, Some("plain".to_string()))); + // Quoted empty string: NULL, exactly like the file-based COPY route. + assert_eq!(rows[1], (true, None)); + // Unquoted empty field: NULL. + assert_eq!(rows[2], (true, None)); + + server.shutdown().await.unwrap(); +} diff --git a/tests-integration/tests/sql.rs b/tests-integration/tests/sql.rs index 4a01b0209f6..31a1c50e5d3 100644 --- a/tests-integration/tests/sql.rs +++ b/tests-integration/tests/sql.rs @@ -2908,3 +2908,112 @@ async fn test_sql_batcher_alignment() { instance.guard.remove_all().await; } } + +/// COPY FROM STDIN must route through the pending-rows batcher (with +/// synchronous acknowledgement) when PostgreSQL batching is enabled, and +/// keep direct writes when it is not. +#[tokio::test(flavor = "multi_thread")] +async fn test_postgres_copy_in_with_batching() { + use bytes::Bytes; + use common_telemetry::{dump_metrics, init_default_ut_logging}; + use frontend::server::Services; + use frontend::service_config::BatcherOptions; + use futures::SinkExt; + use servers::batcher::BatchingProtocol; + use tests_integration::standalone::GreptimeDbStandaloneBuilder; + + fn flushes() -> u64 { + dump_metrics() + .unwrap() + .lines() + .find_map(|line| { + line.strip_prefix("greptime_table_batcher_flush_total ") + .map(|value| value.parse().unwrap()) + }) + .unwrap_or(0) + } + + async fn feed_copy_in( + client: &Client, + sql: &str, + chunks: Vec, + ) -> std::result::Result { + let mut sink = Box::pin(client.copy_in(sql).await?); + for chunk in chunks { + sink.as_mut().send(chunk).await?; + } + sink.as_mut().finish().await + } + + init_default_ut_logging(); + for pg_enabled in [false, true] { + let protocols = if pg_enabled { + vec![BatchingProtocol::Postgres] + } else { + vec![] + }; + let mut instance = GreptimeDbStandaloneBuilder::new("pg_copy_in_batching") + .with_table_batcher(BatcherOptions { + protocols, + pending_rows_flush_interval: Duration::from_millis(10), + ..Default::default() + }) + .build() + .await; + let mut opts = instance.opts.clone(); + opts.http.addr = "127.0.0.1:0".into(); + opts.grpc.bind_addr = "127.0.0.1:0".into(); + opts.postgres.addr = "127.0.0.1:0".into(); + let mut servers = Services::new(opts, instance.fe_instance().clone(), Default::default()) + .build() + .unwrap(); + servers.start_all().await.unwrap(); + let (pg, connection) = tokio_postgres::connect( + &format!( + "postgres://{}/public", + servers.addr("POSTGRES_SERVER").unwrap() + ), + NoTls, + ) + .await + .unwrap(); + let pg_task = tokio::spawn(async move { connection.await.unwrap() }); + + pg.simple_query( + "CREATE TABLE batch_copy_in (ts TIMESTAMP TIME INDEX, host STRING, val INT)", + ) + .await + .unwrap(); + + let before = flushes(); + let rows = feed_copy_in( + &pg, + "COPY batch_copy_in FROM STDIN WITH (FORMAT csv)", + vec![Bytes::from_static( + b"2023-11-14 22:13:20,host1,1\n2023-11-14 22:13:21,host2,2\n", + )], + ) + .await + .unwrap(); + // Synchronous acknowledgement: the reported count is exact whether + // or not the batcher is active. + assert_eq!(rows, 2); + // The batch must have been flushed through the pending-rows batcher. + assert_eq!(flushes() - before, if pg_enabled { 1 } else { 0 }); + + // The flushed rows are visible immediately after the copy. + let rows: Vec = pg + .query("SELECT host FROM batch_copy_in ORDER BY host", &[]) + .await + .unwrap() + .iter() + .map(|row| row.get(0)) + .collect(); + assert_eq!(rows, vec!["host1", "host2"]); + + drop(pg); + pg_task.await.unwrap(); + servers.shutdown_all().await.unwrap(); + instance.guard.remove_all().await; + } +}