servers/postgres/copy_in/
mod.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! PostgreSQL `COPY ... FROM STDIN` (copy-in) support.
16//!
17//! The wire handshake is driven by pgwire: returning [`Response::CopyIn`]
18//! from a query handler makes the server send a `CopyInResponse` and enter
19//! the copy-in sub-protocol; subsequent `CopyData` messages from the client
20//! are dispatched to [`CopyHandler`], whose implementation below parses
21//! the data stream incrementally and flushes it into the query handler as
22//! [`RowInsertRequests`] batches.
23//!
24//! # Formats
25//!
26//! The supported formats are the explicit [`codec::CopyFormat`] enum —
27//! `csv` (table-bound, incremental row parsing) and `influxdb`
28//! (table-less, one bounded whole-payload request). Each format states
29//! its own target rule and ingestion shape in the match arms of the
30//! copy lifecycle below; adding a format is a new variant plus its
31//! arms. Builders resolve a `FORMAT` name (and its aliases) through
32//! [`REGISTERED_FORMATS`]. See the `codec` module docs.
33
34mod codec;
35mod csv;
36mod influxdb;
37mod options;
38mod parse;
39
40use std::fmt::Debug;
41use std::sync::Arc;
42
43use crate::influxdb::InfluxdbRequest;
44use api::helper::ColumnDataTypeWrapper;
45use api::v1::column_def::options_from_column_schema;
46use api::v1::{
47    Row as GrpcRow, RowInsertRequest, RowInsertRequests, Rows as GrpcRows, SemanticType,
48};
49use async_trait::async_trait;
50use codec::{CopyFormat, RecordParser};
51use common_grpc::precision::Precision;
52use common_query::OutputData;
53use common_telemetry::{info, warn};
54use csv::CsvCodec;
55use datafusion::sql::sqlparser::ast::{CopyTarget, Statement as SqlParserStatement};
56use datatypes::data_type::DataType;
57use datatypes::schema::{ColumnSchema, SchemaRef};
58use futures::{Sink, SinkExt, stream};
59use influxdb::InfluxdbOptions;
60use pgwire::api::copy::CopyHandler;
61use pgwire::api::results::{CopyResponse, Response, Tag};
62use pgwire::api::{ClientInfo, PgWireConnectionState};
63use pgwire::error::{ErrorInfo, PgWireError, PgWireResult};
64use pgwire::messages::PgWireBackendMessage;
65use pgwire::messages::copy::{CopyData, CopyDone, CopyFail};
66use session::context::QueryContextRef;
67use table::metadata::TableInfo;
68
69use crate::metrics::METRIC_POSTGRES_COPY_IN_ROWS;
70use crate::postgres::PostgresServerHandlerInner;
71use crate::postgres::types::PgErrorCode;
72use crate::postgres::utils::convert_err;
73
74/// Number of parsed rows buffered before flushing a batch to the handler.
75///
76/// Each flush runs the full insert pipeline (catalog lookup, on-demand
77/// schema check, partition split, region write) and, without the
78/// pending-rows batcher, waits for it — so larger batches mean fewer
79/// round trips at the cost of linearly more memory per connection
80/// (decoder batch + rows buffer). The constant doubles as the arrow
81/// decoder's batch size. High-rate COPY users should enable the
82/// pending-rows batcher (`global-table-batch` / `BatcherOptions.protocols`
83/// containing `Postgres`), which turns flushes into enqueues and makes
84/// the per-batch fixed costs negligible.
85pub(crate) const COPY_IN_FLUSH_ROWS: usize = 32768;
86
87/// Upper bound of an influxdb COPY FROM STDIN payload. One COPY
88/// statement is one bounded ingestion request (the whole payload is
89/// buffered like one HTTP write body); a larger import is multiple COPY
90/// statements. Streaming ingest may arrive later, with framing owned by
91/// the shared line-protocol parser rather than the COPY layer.
92pub(crate) const COPY_IN_MAX_PAYLOAD_BYTES: usize = 64 * 1024 * 1024;
93
94/// Builds a format from the format-neutral option set.
95type FormatBuilder = fn(options::CopyOptionSet) -> PgWireResult<CopyFormat>;
96
97/// The registered COPY FROM STDIN formats: accepted `FORMAT` name to
98/// builder. `csv` is the only implemented format; the others are
99/// registered so they report a clear "not supported yet" instead of an
100/// unknown-format error.
101const REGISTERED_FORMATS: &[(&str, FormatBuilder)] = &[
102    ("csv", csv::build_csv_codec),
103    ("influxdb", influxdb::build_influxdb_codec),
104    ("influxdb_line", influxdb::build_influxdb_codec),
105    ("text", build_unimplemented_text_codec),
106    ("txt", build_unimplemented_text_codec),
107    ("binary", build_unimplemented_binary_codec),
108];
109
110fn build_unimplemented_text_codec(_options: options::CopyOptionSet) -> PgWireResult<CopyFormat> {
111    Err(unsupported(
112        "the text format is not supported yet; use WITH (FORMAT csv)",
113    ))
114}
115
116fn build_unimplemented_binary_codec(_options: options::CopyOptionSet) -> PgWireResult<CopyFormat> {
117    Err(unsupported(
118        "COPY FROM STDIN WITH (FORMAT binary) is not supported yet",
119    ))
120}
121
122/// Resolves the format from the parsed options.
123fn resolve_format(options: &mut options::CopyOptionSet) -> PgWireResult<CopyFormat> {
124    let name = options.take_format()?.ok_or_else(|| {
125        unsupported(
126            "COPY FROM STDIN requires an explicit format in this version; use WITH (FORMAT csv)",
127        )
128    })?;
129    let builder = REGISTERED_FORMATS
130        .iter()
131        .find(|(registered, _)| *registered == name)
132        .map(|(_, builder)| *builder)
133        .ok_or_else(|| bad_copy_data(format!("unknown COPY format \"{}\"", name)))?;
134    builder(std::mem::take(options))
135}
136
137/// The parsed `COPY tbl [(columns)] FROM STDIN [(options)]` statement.
138#[derive(Debug, Clone)]
139pub(crate) struct CopyFromStdin {
140    /// Table name parts as written (`t`, `schema.t`, `catalog.schema.t`);
141    /// quoted names keep their case.
142    table: Vec<String>,
143    /// Requested column list; empty means all columns in table order.
144    columns: Vec<String>,
145    /// The resolved format with its validated options.
146    format: CopyFormat,
147}
148
149fn copy_in_error(code: PgErrorCode, message: impl Into<String>) -> PgWireError {
150    PgWireError::UserError(Box::new(code.to_err_info(message.into())))
151}
152
153pub(super) fn unsupported(message: impl Into<String>) -> PgWireError {
154    copy_in_error(PgErrorCode::Ec0A000, message)
155}
156
157fn bad_copy_data(message: impl Into<String>) -> PgWireError {
158    copy_in_error(PgErrorCode::Ec22P04, message)
159}
160
161fn invalid_text_repr(message: impl Into<String>) -> PgWireError {
162    copy_in_error(PgErrorCode::Ec22P02, message)
163}
164
165pub(crate) fn copy_in_must_run_alone_error() -> PgWireError {
166    PgWireError::UserError(Box::new(ErrorInfo::new(
167        "ERROR".to_string(),
168        PgErrorCode::Ec0A000.code(),
169        "COPY FROM STDIN must be executed alone".to_string(),
170    )))
171}
172
173/// Parses a leading `COPY ... FROM STDIN` statement from SQL text.
174///
175/// The statement is parsed textually — not through sqlparser — because the
176/// `WITH (...)` option list is open: formats may accept options beyond
177/// sqlparser's COPY grammar, and unknown options must reach the format
178/// codec instead of failing statement parsing. Returns `Ok(None)` when
179/// the string does not lead with a copy-in statement.
180pub(crate) fn copy_from_stdin_from_sql(sql: &str) -> PgWireResult<Option<CopyFromStdin>> {
181    let raw = match parse::scan_copy_from_stdin(sql)? {
182        parse::CopyInScan::NotCopyIn => return Ok(None),
183        parse::CopyInScan::CopyIn(raw) => raw,
184    };
185    let mut option_set = raw.options;
186    let format = resolve_format(&mut option_set)?;
187    Ok(Some(CopyFromStdin {
188        table: raw.table,
189        columns: raw.columns,
190        format,
191    }))
192}
193
194/// Whether an already-parsed statement is a `COPY ... FROM STDIN`; used to
195/// reject copy-in statements trailing other statements, where the textual
196/// scanner (which only sees statement-leading COPYs) does not see them.
197pub(crate) fn is_copy_from_stdin(statement: &SqlParserStatement) -> bool {
198    matches!(
199        statement,
200        SqlParserStatement::Copy {
201            to: false,
202            target: CopyTarget::Stdin,
203            ..
204        }
205    )
206}
207
208/// Resolves table name parts against the query context defaults.
209fn resolve_table_name(
210    query_ctx: &QueryContextRef,
211    parts: &[String],
212) -> PgWireResult<(String, String, String)> {
213    match parts {
214        [table] => Ok((
215            query_ctx.current_catalog().to_string(),
216            query_ctx.current_schema(),
217            table.clone(),
218        )),
219        [schema, table] => Ok((
220            query_ctx.current_catalog().to_string(),
221            schema.clone(),
222            table.clone(),
223        )),
224        [catalog, schema, table] => Ok((catalog.clone(), schema.clone(), table.clone())),
225        _ => Err(unsupported(format!(
226            "invalid table name for COPY: {}",
227            parts.join(".")
228        ))),
229    }
230}
231
232/// Per-connection state of an in-progress COPY FROM STDIN.
233///
234/// The payload shape follows the codec: table-bound formats (csv) parse
235/// records into rows against the targeted table's columns; table-less
236/// formats (influxdb line protocol) split the stream into complete
237/// lines that each route to their own measurement, and flush through
238/// the handler's line-protocol pipeline — the same one the HTTP write
239/// endpoint uses.
240pub(crate) struct CopyInState {
241    query_ctx: QueryContextRef,
242    payload: Payload,
243    /// The resolved codec's name, for logging.
244    codec_name: &'static str,
245    /// The wire format code of the resolved codec.
246    codec_format_code: i8,
247    rows_written: u64,
248}
249
250enum Payload {
251    /// Table-bound rows: parsed against `grpc_schema` and flushed to
252    /// `table_name` via `copy_in_insert`. The schema is cloned into each
253    /// flushed batch: `Rows.schema` is an owned `Vec` in the proto
254    /// request, so one copy per request is inherent.
255    Table {
256        grpc_schema: Vec<api::v1::ColumnSchema>,
257        table_name: String,
258        parser: Box<dyn RecordParser + Send>,
259        /// Parsed rows not yet flushed to the handler.
260        rows: Vec<GrpcRow>,
261    },
262    /// The whole buffered influxdb line-protocol payload with its
263    /// request precision. One COPY statement is one bounded ingestion
264    /// request, buffered until `CopyDone` — exactly like one HTTP write
265    /// body — and record boundaries belong to the line-protocol parser,
266    /// never to the COPY layer.
267    Influxdb {
268        data: Vec<u8>,
269        precision: Option<Precision>,
270    },
271}
272
273impl CopyInState {
274    fn new(
275        codec: &CsvCodec,
276        columns: Vec<String>,
277        table_info: Arc<TableInfo>,
278        query_ctx: QueryContextRef,
279        table_name: String,
280    ) -> PgWireResult<Self> {
281        let schema = table_info.meta.schema.clone();
282        let timestamp_index = schema.timestamp_index();
283
284        let columns = resolve_columns(codec, columns, &schema, &table_info)?;
285        if let Some(ts_index) = timestamp_index {
286            let ts_column = &schema.column_schemas()[ts_index];
287            if !columns.iter().any(|column| column.name == ts_column.name) {
288                return Err(unsupported(format!(
289                    "COPY FROM STDIN requires the time index column \"{}\" in the column list",
290                    ts_column.name
291                )));
292            }
293        }
294        let parser = codec.create_parser(&columns, query_ctx.timezone().to_string().as_str())?;
295
296        let grpc_schema = columns
297            .iter()
298            .map(|column| build_grpc_column_schema(column, &table_info))
299            .collect::<PgWireResult<Vec<_>>>()?;
300
301        Ok(Self {
302            query_ctx,
303            payload: Payload::Table {
304                grpc_schema,
305                table_name,
306                parser,
307                rows: Vec::with_capacity(COPY_IN_FLUSH_ROWS),
308            },
309            codec_name: "csv",
310            codec_format_code: 0,
311            rows_written: 0,
312        })
313    }
314
315    /// State for the influxdb format: the whole payload is buffered and
316    /// completed once at `CopyDone` through the line-protocol pipeline.
317    fn new_influxdb(options: &InfluxdbOptions, query_ctx: QueryContextRef) -> Self {
318        Self {
319            query_ctx,
320            payload: Payload::Influxdb {
321                data: Vec::new(),
322                precision: options.precision,
323            },
324            codec_name: "influxdb",
325            codec_format_code: 0,
326            rows_written: 0,
327        }
328    }
329
330    /// The wire format code advertised to the client.
331    pub(crate) fn format_code(&self) -> i8 {
332        self.codec_format_code
333    }
334
335    /// Feeds a `CopyData` chunk into the parser and buffers parsed rows
336    /// (csv) or payload bytes (influxdb).
337    fn feed(&mut self, data: &[u8]) -> PgWireResult<()> {
338        match &mut self.payload {
339            Payload::Table { parser, rows, .. } => rows.extend(parser.feed(data)?),
340            Payload::Influxdb { data: buffered, .. } => {
341                if buffered.len() + data.len() > COPY_IN_MAX_PAYLOAD_BYTES {
342                    return Err(bad_copy_data(format!(
343                        "COPY FROM STDIN payload exceeds the {} limit; split the import into multiple COPY statements",
344                        COPY_IN_MAX_PAYLOAD_BYTES
345                    )));
346                }
347                buffered.extend_from_slice(data);
348            }
349        }
350        Ok(())
351    }
352
353    /// Parses the buffered tail at `CopyDone`; a final record without a
354    /// trailing newline is accepted.
355    fn finish(&mut self) -> PgWireResult<()> {
356        match &mut self.payload {
357            Payload::Table { parser, rows, .. } => rows.extend(parser.finish()?),
358            // The whole payload is one request body; nothing to finalize.
359            Payload::Influxdb { .. } => {}
360        }
361        Ok(())
362    }
363
364    fn should_flush(&self) -> bool {
365        match &self.payload {
366            Payload::Table { rows, .. } => rows.len() >= COPY_IN_FLUSH_ROWS,
367            // Completed once at CopyDone, not mid-stream.
368            Payload::Influxdb { .. } => false,
369        }
370    }
371
372    /// Takes the buffered rows as an insert request, if any.
373    fn take_requests(&mut self) -> Option<RowInsertRequests> {
374        let Payload::Table {
375            grpc_schema,
376            table_name,
377            rows,
378            ..
379        } = &mut self.payload
380        else {
381            return None;
382        };
383        if rows.is_empty() {
384            return None;
385        }
386        let rows = GrpcRows {
387            schema: grpc_schema.clone(),
388            rows: std::mem::take(rows),
389        };
390        Some(RowInsertRequests {
391            inserts: vec![RowInsertRequest {
392                table_name: table_name.clone(),
393                rows: Some(rows),
394            }],
395        })
396    }
397
398    /// Takes the buffered payload as one line-protocol request body, if
399    /// any — the same bytes an HTTP write of the same records would
400    /// carry.
401    fn take_lines(&mut self) -> PgWireResult<Option<InfluxdbRequest>> {
402        let Payload::Influxdb { data, precision } = &mut self.payload else {
403            return Ok(None);
404        };
405        if data.is_empty() {
406            return Ok(None);
407        }
408        let lines = String::from_utf8(std::mem::take(data))
409            .map_err(|e| bad_copy_data(format!("COPY FROM STDIN data is not valid UTF-8: {e}")))?;
410        Ok(Some(InfluxdbRequest {
411            precision: *precision,
412            lines,
413        }))
414    }
415}
416
417fn resolve_columns(
418    codec: &CsvCodec,
419    columns: Vec<String>,
420    schema: &SchemaRef,
421    table_info: &Arc<TableInfo>,
422) -> PgWireResult<Vec<ColumnSchema>> {
423    let table_name = &table_info.name;
424    let selected: Vec<ColumnSchema> = if columns.is_empty() {
425        schema.column_schemas().to_vec()
426    } else {
427        let mut selected = Vec::with_capacity(columns.len());
428        for name in &columns {
429            let column = schema
430                .column_schema_by_name(name)
431                .ok_or_else(|| {
432                    copy_in_error(
433                        PgErrorCode::Ec42703,
434                        format!(
435                            "column \"{}\" of relation \"{}\" does not exist",
436                            name, table_name
437                        ),
438                    )
439                })?
440                .clone();
441            if selected.iter().any(|c: &ColumnSchema| &c.name == name) {
442                return Err(bad_copy_data(format!(
443                    "column \"{}\" specified more than once",
444                    name
445                )));
446            }
447            selected.push(column);
448        }
449        selected
450    };
451
452    for column in &selected {
453        codec.check_column_type(&column.data_type)?;
454    }
455
456    Ok(selected)
457}
458
459/// Derives a column's semantic type the same way the flow adapter does
460/// (`flow::adapter::util::column_schemas_to_proto`): the time-index flag
461/// for timestamps, primary-key membership for tags.
462fn semantic_type_of(table_info: &TableInfo, column: &ColumnSchema) -> PgWireResult<SemanticType> {
463    let meta = &table_info.meta;
464    if column.is_time_index() {
465        Ok(SemanticType::Timestamp)
466    } else if meta
467        .primary_key_indices
468        .iter()
469        .any(|&index| meta.schema.column_schemas()[index].name == column.name)
470    {
471        Ok(SemanticType::Tag)
472    } else {
473        Ok(SemanticType::Field)
474    }
475}
476
477fn build_grpc_column_schema(
478    column: &ColumnSchema,
479    table_info: &Arc<TableInfo>,
480) -> PgWireResult<api::v1::ColumnSchema> {
481    let (datatype, datatype_extension) = ColumnDataTypeWrapper::try_from(column.data_type.clone())
482        .map_err(|_| {
483            unsupported(format!(
484                "column type \"{}\" is not supported for COPY FROM STDIN yet",
485                column.data_type.name()
486            ))
487        })?
488        .to_parts();
489    Ok(api::v1::ColumnSchema {
490        column_name: column.name.clone(),
491        datatype: datatype as i32,
492        semantic_type: semantic_type_of(table_info, column)? as i32,
493        datatype_extension,
494        options: options_from_column_schema(column),
495    })
496}
497
498// ---------------------------------------------------------------------------
499// pgwire handlers
500// ---------------------------------------------------------------------------
501
502impl PostgresServerHandlerInner {
503    /// Prepares a COPY FROM STDIN: resolves the target table, validates
504    /// columns and stashes the per-connection copy state.
505    pub(crate) async fn begin_copy_in(&self, stmt: CopyFromStdin) -> PgWireResult<Response> {
506        // The batching-aware context routes copy-in flushes through the
507        // pending-rows batcher like SQL INSERT when it is enabled for
508        // this connection.
509        let query_ctx = self.new_query_context();
510
511        // Each format states its own target rule and ingestion shape
512        // explicitly.
513        let state = match &stmt.format {
514            // csv requires a target table: records are parsed against
515            // its columns and inserted in batches.
516            CopyFormat::Csv(codec) => {
517                if stmt.table.is_empty() {
518                    return Err(unsupported(
519                        "format csv requires a table: COPY tbl [(cols)] FROM STDIN WITH (FORMAT csv)",
520                    ));
521                }
522                let (catalog, schema, table) = resolve_table_name(&query_ctx, &stmt.table)?;
523
524                let Some(table_info) = self
525                    .copy_in_handler
526                    .copy_in_table(&catalog, &schema, &table, query_ctx.clone())
527                    .await
528                    .map_err(convert_err)?
529                else {
530                    return Err(copy_in_error(
531                        PgErrorCode::Ec42P01,
532                        format!("relation \"{}\" does not exist", stmt.table.join(".")),
533                    ));
534                };
535
536                // The insert path resolves the bare table name against the
537                // query context, so point the context at the resolved
538                // catalog and schema when the COPY target is qualified.
539                let query_ctx = adjust_query_context(query_ctx, &catalog, &schema);
540
541                CopyInState::new(codec, stmt.columns, table_info, query_ctx, table)?
542            }
543            // influxdb forbids a target table — each line names its own
544            // measurement — and completes through the line-protocol
545            // pipeline the HTTP write endpoints use; the server must be
546            // wired for it.
547            CopyFormat::Influxdb(options) => {
548                if !stmt.table.is_empty() {
549                    return Err(unsupported(format!(
550                        "format influxdb routes each record to its own table; \"COPY {} FROM STDIN\" must not name one",
551                        stmt.table.join(".")
552                    )));
553                }
554                if self.influxdb_handler.is_none() {
555                    return Err(unsupported(
556                        "this server is not wired for the influxdb line protocol format",
557                    ));
558                }
559                CopyInState::new_influxdb(options, query_ctx)
560            }
561        };
562        let num_columns = match &state.payload {
563            Payload::Table { grpc_schema, .. } => grpc_schema.len(),
564            // No fixed column set: the format routes per record.
565            Payload::Influxdb { .. } => 0,
566        };
567        let format_code = state.format_code();
568
569        *self.copy_in_state.lock().await = Some(state);
570
571        Ok(Response::CopyIn(CopyResponse::new(
572            format_code,
573            num_columns,
574            stream::empty(),
575        )))
576    }
577
578    /// Flushes the buffered payload of an in-progress copy. Returns the
579    /// number of rows written by the flushed batch.
580    async fn flush_copy_in(&self, state: &mut CopyInState) -> PgWireResult<u64> {
581        let query_ctx = state.query_ctx.clone();
582        // The buffered influxdb payload is completed through the
583        // line-protocol pipeline — the same handler the HTTP write
584        // endpoints use, so conversion, permissions, interceptors,
585        // timestamp alignment and inserts are identical to an HTTP
586        // write of the same bytes. The context stays this connection's
587        // postgres context.
588        if let Some(request) = state.take_lines()? {
589            let Some(influxdb_handler) = &self.influxdb_handler else {
590                return Err(unsupported(
591                    "this server is not wired for the influxdb line protocol format",
592                ));
593            };
594            let output = influxdb_handler
595                .exec(request, query_ctx)
596                .await
597                .map_err(convert_err)?;
598            return match output.data {
599                OutputData::AffectedRows(rows) => Ok(rows as u64),
600                OutputData::RecordBatches(_) | OutputData::Stream(_) => Ok(0),
601            };
602        }
603        let Some(requests) = state.take_requests() else {
604            return Ok(0);
605        };
606        let output = self
607            .copy_in_handler
608            .copy_in_insert(requests, query_ctx)
609            .await
610            .map_err(convert_err)?;
611        match output.data {
612            OutputData::AffectedRows(rows) => Ok(rows as u64),
613            OutputData::RecordBatches(_) | OutputData::Stream(_) => {
614                warn!("unexpected output from COPY FROM STDIN insert");
615                Ok(0)
616            }
617        }
618    }
619}
620
621fn adjust_query_context(
622    query_ctx: QueryContextRef,
623    catalog: &str,
624    schema: &str,
625) -> QueryContextRef {
626    if query_ctx.current_catalog() == catalog && query_ctx.current_schema() == schema {
627        return query_ctx;
628    }
629    let mut forked = query_ctx.fork();
630    forked.set_current_catalog(catalog);
631    forked.set_current_schema(schema);
632    Arc::new(forked)
633}
634
635#[async_trait]
636impl CopyHandler for PostgresServerHandlerInner {
637    async fn on_copy_data<C>(&self, _client: &mut C, copy_data: CopyData) -> PgWireResult<()>
638    where
639        C: ClientInfo + Sink<PgWireBackendMessage> + Unpin + Send + Sync,
640        C::Error: Debug,
641        PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
642    {
643        let mut guard = self.copy_in_state.lock().await;
644        let Some(mut state) = guard.take() else {
645            return Err(copy_in_error(
646                PgErrorCode::Ec08P01,
647                "CopyData message received outside COPY FROM STDIN",
648            ));
649        };
650        // The state is owned here: any parse, conversion or insert error
651        // below aborts the whole copy and drops the buffered data —
652        // otherwise a pooled connection would retain failed-import
653        // buffers until its next COPY. On success it is put back.
654        state.feed(&copy_data.data)?;
655        if state.should_flush() {
656            let rows = self.flush_copy_in(&mut state).await?;
657            state.rows_written += rows;
658            METRIC_POSTGRES_COPY_IN_ROWS.inc_by(rows);
659        }
660        *guard = Some(state);
661        Ok(())
662    }
663
664    async fn on_copy_done<C>(&self, client: &mut C, _done: CopyDone) -> PgWireResult<()>
665    where
666        C: ClientInfo + Sink<PgWireBackendMessage> + Unpin + Send + Sync,
667        C::Error: Debug,
668        PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
669    {
670        let mut guard = self.copy_in_state.lock().await;
671        // Take ownership up front: even finalization failures (parse
672        // errors in the trailing record, insert errors) drop the state
673        // instead of retaining it on a long-lived connection.
674        let Some(mut state) = guard.take() else {
675            return Err(copy_in_error(
676                PgErrorCode::Ec08P01,
677                "CopyDone message received outside COPY FROM STDIN",
678            ));
679        };
680        state.finish()?;
681        let rows = self.flush_copy_in(&mut state).await?;
682        state.rows_written += rows;
683        METRIC_POSTGRES_COPY_IN_ROWS.inc_by(rows);
684
685        let rows_written = state.rows_written;
686        let codec_name = state.codec_name;
687        info!(
688            "PostgreSQL COPY FROM STDIN (format {}) finished: {} rows written",
689            codec_name, rows_written
690        );
691        client
692            .send(PgWireBackendMessage::CommandComplete(
693                Tag::new("COPY").with_rows(rows_written as usize).into(),
694            ))
695            .await?;
696
697        // After a successful extended-protocol copy-in, pgwire leaves the
698        // connection in `CopyInProgress(true)` and swallows the client's
699        // terminating `Sync` message, so `ReadyForQuery` would never be
700        // sent and the connection would deadlock. Route the state to
701        // `AwaitingSync` so the `Sync` is dispatched to `on_sync`, which
702        // completes the round trip.
703        if let PgWireConnectionState::CopyInProgress(true) = client.state() {
704            client.set_state(PgWireConnectionState::AwaitingSync);
705        }
706        Ok(())
707    }
708
709    async fn on_copy_fail<C>(&self, _client: &mut C, fail: CopyFail) -> PgWireError
710    where
711        C: ClientInfo + Sink<PgWireBackendMessage> + Unpin + Send + Sync,
712        C::Error: Debug,
713        PgWireError: From<<C as Sink<PgWireBackendMessage>>::Error>,
714    {
715        // Drop the buffered state; partial data of a failed copy is not
716        // written.
717        self.copy_in_state.lock().await.take();
718        copy_in_error(
719            PgErrorCode::Ec57014,
720            format!("COPY FROM STDIN mode terminated by user: {}", fail.message),
721        )
722    }
723}
724
725#[cfg(test)]
726mod tests {
727    use super::*;
728
729    #[test]
730    fn test_parse_copy_from_stdin_requires_format() {
731        let err = copy_from_stdin_from_sql("COPY t FROM STDIN").unwrap_err();
732        assert!(
733            err.to_string()
734                .contains("requires an explicit format in this version"),
735            "{err}"
736        );
737
738        let copy = copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT csv)")
739            .unwrap()
740            .unwrap();
741        assert_eq!(copy.table, vec!["t"]);
742        assert!(copy.columns.is_empty());
743        assert!(matches!(copy.format, CopyFormat::Csv(_)));
744        // Both formats are textual: format code 0.
745    }
746
747    #[test]
748    fn test_parse_copy_from_stdin_csv() {
749        let copy = copy_from_stdin_from_sql(
750            "COPY t (a, b) FROM STDIN WITH (FORMAT csv, HEADER, DELIMITER ';')",
751        )
752        .unwrap()
753        .unwrap();
754        assert_eq!(copy.columns, vec!["a", "b"]);
755        assert!(matches!(copy.format, CopyFormat::Csv(_)));
756    }
757
758    #[test]
759    fn test_parse_copy_from_stdin_text_unsupported() {
760        for sql in [
761            "COPY t FROM STDIN WITH (FORMAT text)",
762            "COPY t FROM STDIN WITH (FORMAT txt)",
763        ] {
764            let err = copy_from_stdin_from_sql(sql).unwrap_err();
765            assert!(
766                err.to_string().contains("text format is not supported yet"),
767                "{err}"
768            );
769        }
770    }
771
772    #[test]
773    fn test_parse_copy_from_stdin_binary_unsupported() {
774        let err = copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT binary)").unwrap_err();
775        assert!(err.to_string().contains("not supported"));
776    }
777
778    #[test]
779    fn test_parse_copy_from_stdin_unknown_format() {
780        let err = copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT parquet)").unwrap_err();
781        assert!(err.to_string().contains("unknown COPY format"));
782    }
783
784    #[test]
785    fn test_parse_copy_from_stdin_not_copy_in() {
786        assert!(
787            copy_from_stdin_from_sql("COPY (SELECT 1) TO STDOUT")
788                .unwrap()
789                .is_none()
790        );
791        assert!(
792            copy_from_stdin_from_sql("COPY t TO '/tmp/x.csv'")
793                .unwrap()
794                .is_none()
795        );
796        // GreptimeDB's own file-based COPY FROM.
797        assert!(
798            copy_from_stdin_from_sql("COPY t FROM '/tmp/x.csv' WITH (FORMAT csv)")
799                .unwrap()
800                .is_none()
801        );
802    }
803
804    #[test]
805    fn test_redundant_options() {
806        let err = copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT csv, FORMAT text)")
807            .unwrap_err();
808        assert!(err.to_string().contains("redundant FORMAT"));
809
810        let err = copy_from_stdin_from_sql("COPY t FROM STDIN WITH (DELIMITER ',', DELIMITER ';')")
811            .unwrap_err();
812        assert!(err.to_string().contains("redundant DELIMITER"));
813    }
814
815    #[test]
816    fn test_format_specific_options_rejected() {
817        // The open WITH section parses anything, but the format rejects
818        // the options it does not understand.
819        let err =
820            copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT csv, PRECISION 'ns', FREEZE)")
821                .unwrap_err();
822        assert!(
823            err.to_string().contains("not supported by format csv"),
824            "{err}"
825        );
826    }
827
828    #[test]
829    fn test_multi_statement_rejected() {
830        let err = copy_from_stdin_from_sql("COPY t FROM STDIN; SELECT 1").unwrap_err();
831        assert!(err.to_string().contains("alone"));
832        // Trailing separators and quoted semicolons are fine.
833        assert!(
834            copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT csv);")
835                .unwrap()
836                .is_some()
837        );
838        assert!(
839            copy_from_stdin_from_sql("COPY t FROM STDIN WITH (FORMAT csv, DELIMITER ';')")
840                .unwrap()
841                .is_some()
842        );
843    }
844}
servers/postgres/copy_in/mod.rs

servers/postgres/copy_in/
mod.rs
