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 codec framework.
16//!
17//! A *codec* binds one COPY data format to its implementation. Codecs are
18//! stateless strategies: they validate the options they accept at
19//! statement-parse time, describe which column types they can ingest, and
20//! create the stateful [`RecordParser`] per stream.
21//!
22//! Formats are not limited to PostgreSQL's own — GreptimeDB-specific
23//! formats (a line protocol, a JSON-lines dialect, ...) register exactly
24//! like `csv` does, and may define their own options through the open
25//! `WITH (...)` clause.
26//!
27//! # Parsing and conversion
28//!
29//! A [`RecordParser`] consumes raw `CopyData` bytes and emits complete,
30//! insertable rows: one [`ParsedRow`] per record, in the order of the
31//! targeted columns. Parsers validate record shape (field counts) and
32//! convert field values themselves — `csv` delegates both splitting and
33//! typed conversion to arrow-csv's incremental decoder (the engine
34//! behind `COPY ... FROM 'file.csv'`) and turns decoded batches into
35//! rows with the shared `api::helper::vectors_to_rows`. A future format
36//! that carries values as text (e.g. a line protocol) would reintroduce
37//! shared field-conversion helpers here; see the removed `text` codec in
38//! git history for the shape they had.
39//!
40//! # Options
41//!
42//! The `WITH (...)` clause is open: every `KEY value` entry the client
43//! writes is collected into a [`CopyOptionSet`] (see `options`) regardless
44//! of whether any format understands it — the statement is parsed
45//! textually (see `parse`) precisely so options beyond sqlparser's COPY
46//! grammar can reach the codecs. A codec builder takes the options it
47//! understands (`take_string`, `take_single_byte`, `take_flag`, or raw
48//! `take_value` for codec-specific shapes) and rejects the rest with
49//! `reject_leftovers`, so a future format like
50//! `COPY t FROM STDIN WITH (FORMAT influxdb_line, precision 'ns')` only
51//! needs its own codec and registry entry.
52//!
53//! To add a format:
54//!
55//! 1. implement [`CopyInCodec`] for a config struct holding the format's
56//!    validated options (see `csv::CsvCodec` for the example),
57//! 2. write a builder `fn(CopyOptionSet) -> PgWireResult<Arc<dyn
58//!    CopyInCodec>>` that applies defaults and rejects options the format
59//!    does not accept,
60//! 3. register the builder under every accepted `FORMAT` name (including
61//!    aliases) in `REGISTERED_FORMATS` in the parent module.
62
63use datatypes::prelude::ConcreteDataType;
64use datatypes::schema::ColumnSchema;
65use pgwire::error::PgWireResult;
66
67/// A parsed, insertable row: one gRPC value per targeted column, in COPY
68/// field order. Every parser emits these directly, whatever mechanism it
69/// used to produce them.
70pub(super) type ParsedRow = api::v1::Row;
71
72/// Incremental parser of one COPY FROM STDIN stream.
73///
74/// The wire protocol delivers the payload as an arbitrary sequence of
75/// `CopyData` chunks; implementations must keep state across [`feed`]
76/// calls because records (and even single fields) may span chunk
77/// boundaries. [`finish`] is called at `CopyDone`; a final record without
78/// a trailing record separator is accepted.
79pub(super) trait RecordParser: Send {
80    /// Feeds a chunk of `CopyData` bytes and returns all complete rows.
81    fn feed(&mut self, data: &[u8]) -> PgWireResult<Vec<ParsedRow>>;
82    /// Parses the buffered tail at `CopyDone`; returns no row when the
83    /// stream ended exactly at a record boundary.
84    fn finish(&mut self) -> PgWireResult<Vec<ParsedRow>>;
85}
86
87/// A COPY FROM STDIN data format implementation.
88///
89/// Implementations are shared per statement (hence `Send + Sync`) and must
90/// be stateless; per-stream state belongs to the [`RecordParser`] they
91/// create. Every method is tailored to the format: `check_column_type`
92/// declares the types the format can decode (csv is bounded by what
93/// arrow-csv supports).
94pub(super) trait CopyInCodec: Send + Sync + std::fmt::Debug {
95    /// The canonical format name, used in error messages.
96    fn name(&self) -> &'static str;
97    /// The format code advertised in the `CopyInResponse` header: 0 when
98    /// the client sends textual data (the default), 1 for binary.
99    fn format_code(&self) -> i8 {
100        0
101    }
102    /// Validates that a column of `data_type` can be ingested by this
103    /// format.
104    fn check_column_type(&self, data_type: &ConcreteDataType) -> PgWireResult<()>;
105    /// Creates the stateful parser for one copy stream over the targeted
106    /// columns (in COPY field order).
107    ///
108    /// `timezone` is the session timezone (an IANA name or a UTC offset):
109    /// formats that parse offset-free timestamp strings from text should
110    /// interpret them in this zone, matching SQL INSERT.
111    fn create_parser(
112        &self,
113        columns: &[ColumnSchema],
114        timezone: &str,
115    ) -> PgWireResult<Box<dyn RecordParser + Send>>;
116}
servers/postgres/copy_in/codec.rs

servers/postgres/copy_in/
codec.rs
