servers/postgres/copy_in/
codec.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//! The COPY FROM STDIN formats.
16//!
17//! [`CopyFormat`] is the closed, explicit set of supported formats — a
18//! small enum over the real implementations rather than a generic
19//! format framework. Each variant carries its validated options and
20//! its rules are stated where they apply: a variant's match arms in
21//! the copy lifecycle (see the parent module) decide whether a target
22//! table is required, how the `CopyData` bytes are consumed, and which
23//! existing ingestion pipeline completes the copy. Adding a format is
24//! a normal extension: a new variant plus its arms.
25//!
26//! # Parsing and conversion
27//!
28//! Table-bound formats parse incrementally through a [`RecordParser`]:
29//! it consumes raw `CopyData` bytes and emits complete, insertable
30//! rows — one [`ParsedRow`] per record, in the order of the targeted
31//! columns. `csv` delegates splitting and typed conversion to
32//! arrow-csv's incremental decoder (the engine behind
33//! `COPY ... FROM 'file.csv'`) and turns decoded batches into rows
34//! with the shared `api::helper::vectors_to_rows`.
35//!
36//! Formats that delegate to an existing ingestion pipeline (the
37//! influxdb line protocol) do not parse here at all: the whole payload
38//! is buffered as one bounded request body and completed through that
39//! pipeline, so record boundaries and values stay with the pipeline's
40//! own parser.
41//!
42//! # Options
43//!
44//! The `WITH (...)` clause is open: every `KEY value` entry the client
45//! writes is collected into a [`CopyOptionSet`] (see `options`)
46//! regardless of whether any format understands it — the statement is
47//! parsed textually (see `parse`) precisely so options beyond
48//! sqlparser's COPY grammar can reach the formats. A format builder
49//! takes the options it understands (`take_string`,
50//! `take_single_byte`, `take_flag`, or raw `take_value` for
51//! format-specific shapes) and rejects the rest with
52//! `reject_leftovers`.
53
54use pgwire::error::PgWireResult;
55
56use crate::postgres::copy_in::csv::CsvCodec;
57use crate::postgres::copy_in::influxdb::InfluxdbOptions;
58
59/// A resolved COPY FROM STDIN format with its validated options.
60#[derive(Debug, Clone)]
61pub(super) enum CopyFormat {
62    /// Table-bound CSV: records are parsed incrementally against the
63    /// targeted table's columns and inserted in batches.
64    Csv(CsvCodec),
65    /// Table-less influxdb line protocol: the whole payload is one
66    /// bounded request body, completed through the line-protocol
67    /// pipeline the HTTP write endpoints use (one COPY statement is one
68    /// ingestion request; streaming ingest is not supported).
69    Influxdb(InfluxdbOptions),
70}
71
72/// A parsed, insertable row: one gRPC value per targeted column, in COPY
73/// field order. Every parser emits these directly, whatever mechanism it
74/// used to produce them.
75pub(super) type ParsedRow = api::v1::Row;
76
77/// Incremental parser of one COPY FROM STDIN stream.
78///
79/// The wire protocol delivers the payload as an arbitrary sequence of
80/// `CopyData` chunks; implementations must keep state across [`feed`]
81/// calls because records (and even single fields) may span chunk
82/// boundaries. [`finish`] is called at `CopyDone`; a final record without
83/// a trailing record separator is accepted.
84pub(super) trait RecordParser: Send {
85    /// Feeds a chunk of `CopyData` bytes and returns all complete rows.
86    fn feed(&mut self, data: &[u8]) -> PgWireResult<Vec<ParsedRow>>;
87    /// Parses the buffered tail at `CopyDone`; returns no row when the
88    /// stream ended exactly at a record boundary.
89    fn finish(&mut self) -> PgWireResult<Vec<ParsedRow>>;
90}
servers/postgres/copy_in/codec.rs

servers/postgres/copy_in/
codec.rs
