servers/postgres/copy_in/
parse.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//! Textual parser for `COPY ... FROM STDIN` statements.
16//!
17//! COPY FROM STDIN is parsed here instead of through sqlparser because the
18//! `WITH (...)` option list is open: formats may accept options beyond
19//! sqlparser's COPY grammar (for example a line-protocol format taking
20//! `precision`), and unknown options must reach the format codec rather
21//! than fail statement parsing.
22//!
23//! The parser commits to error reporting only after `FROM STDIN` has been
24//! seen; anything else (like `COPY ... TO STDOUT`) is reported as
25//! [`CopyInScan::NotCopyIn`] so the caller falls back to the regular SQL
26//! path.
27
28use datafusion::sql::sqlparser::dialect::GenericDialect;
29use datafusion::sql::sqlparser::tokenizer::{Token, Tokenizer};
30use pgwire::error::PgWireResult;
31
32use crate::postgres::copy_in::options::{CopyOptionSet, CopyOptionValue};
33use crate::postgres::copy_in::{bad_copy_data, copy_in_must_run_alone_error};
34
35/// A scanned `COPY [tbl [(cols)]] FROM STDIN [WITH (...)]` statement.
36///
37/// Unquoted identifiers are folded to lowercase and quoted ones keep their
38/// case (`COPY Metrics` addresses the same table as `CREATE TABLE Metrics`,
39/// namely `metrics`), matching `ParserContext::canonicalize_identifier` in
40/// the SQL parser and PostgreSQL's own folding rule.
41#[derive(Debug)]
42pub(super) struct RawCopyFromStdin {
43    /// Table name parts, as written (`t`, `schema.t`,
44    /// `catalog.schema.t`); quoted names keep their case. Empty for
45    /// table-less formats whose records carry their own routing (e.g.
46    /// the influxdb line protocol names its measurement per line).
47    pub(super) table: Vec<String>,
48    /// Requested column list; empty means all columns in table order.
49    pub(super) columns: Vec<String>,
50    /// The open option list, handed to the resolved format codec.
51    pub(super) options: CopyOptionSet,
52}
53
54pub(super) enum CopyInScan {
55    /// Not a `COPY ... FROM STDIN` statement.
56    NotCopyIn,
57    /// A complete copy-in statement with only trailing separators left.
58    CopyIn(RawCopyFromStdin),
59}
60
61/// Scans a SQL string for a leading `COPY ... FROM STDIN` statement.
62pub(super) fn scan_copy_from_stdin(sql: &str) -> PgWireResult<CopyInScan> {
63    // Fast path: the statement must start with the COPY keyword, possibly
64    // behind comments; only then is a full tokenization worth its cost.
65    let trimmed = sql.trim_start();
66    let starts_like_copy = trimmed
67        .as_bytes()
68        .get(..4)
69        .is_some_and(|prefix| prefix.eq_ignore_ascii_case(b"copy"));
70    let starts_with_comment = trimmed.starts_with("--") || trimmed.starts_with("/*");
71    if !starts_like_copy && !starts_with_comment {
72        return Ok(CopyInScan::NotCopyIn);
73    }
74
75    let tokens = match Tokenizer::new(&GenericDialect {}, sql).tokenize() {
76        Ok(tokens) => tokens,
77        // Let the regular SQL path produce the parse error.
78        Err(_) => return Ok(CopyInScan::NotCopyIn),
79    };
80    // The tokenizer keeps whitespace; the scanner works on tokens only.
81    let tokens: Vec<Token> = tokens
82        .into_iter()
83        .filter(|token| !matches!(token, Token::Whitespace(_)))
84        .collect();
85    let mut cursor = Tokens {
86        tokens: &tokens,
87        pos: 0,
88    };
89
90    if !cursor.take_keyword("copy") {
91        return Ok(CopyInScan::NotCopyIn);
92    }
93
94    // Table-less statements (`COPY FROM STDIN WITH (FORMAT influxdb) ...`)
95    // are table-format specific: the codec decides whether a missing
96    // table is legal. An unquoted `FROM` cannot be a table name (SQL
97    // would require quoting it), so peek instead of letting the
98    // identifier parser swallow the keyword. A column list without a
99    // table is not.
100    let table = if cursor.peek_keyword("from") {
101        Vec::new()
102    } else {
103        match parse_object_name(&mut cursor) {
104            Some(table) => table,
105            None => return Ok(CopyInScan::NotCopyIn),
106        }
107    };
108    let columns = if table.is_empty() {
109        Vec::new()
110    } else {
111        match parse_column_list(&mut cursor) {
112            Some(columns) => columns,
113            None => return Ok(CopyInScan::NotCopyIn),
114        }
115    };
116
117    if !cursor.take_keyword("from") || !cursor.take_keyword("stdin") {
118        // A COPY with another direction or source; not ours to handle.
119        return Ok(CopyInScan::NotCopyIn);
120    }
121
122    let mut options = CopyOptionSet::default();
123    // PostgreSQL allows the WITH keyword to be omitted.
124    let has_options = if cursor.take_keyword("with") {
125        if !matches!(cursor.peek(), Some(Token::LParen)) {
126            return Err(bad_copy_data("expected '(' after WITH in COPY FROM STDIN"));
127        }
128        true
129    } else {
130        matches!(cursor.peek(), Some(Token::LParen))
131    };
132    if has_options {
133        // Consume the opening '('.
134        cursor.next();
135        parse_option_list(&mut cursor, &mut options)?;
136    }
137
138    // Only statement separators may follow.
139    while let Some(token) = cursor.next() {
140        if !matches!(token, Token::SemiColon) {
141            return Err(copy_in_must_run_alone_error());
142        }
143    }
144
145    Ok(CopyInScan::CopyIn(RawCopyFromStdin {
146        table,
147        columns,
148        options,
149    }))
150}
151
152struct Tokens<'a> {
153    tokens: &'a [Token],
154    pos: usize,
155}
156
157impl<'a> Tokens<'a> {
158    fn next(&mut self) -> Option<&'a Token> {
159        let token = self.tokens.get(self.pos)?;
160        self.pos += 1;
161        Some(token)
162    }
163
164    fn peek(&mut self) -> Option<&'a Token> {
165        self.tokens.get(self.pos)
166    }
167
168    /// Consumes the next token if it is the unquoted keyword `keyword`.
169    fn take_keyword(&mut self, keyword: &str) -> bool {
170        if self.peek_keyword(keyword) {
171            self.pos += 1;
172            true
173        } else {
174            false
175        }
176    }
177
178    /// Whether the next token is the unquoted keyword `keyword`, without
179    /// consuming it.
180    fn peek_keyword(&self, keyword: &str) -> bool {
181        matches!(self.tokens.get(self.pos), Some(Token::Word(w))
182            if w.quote_style.is_none() && w.value.eq_ignore_ascii_case(keyword))
183    }
184}
185
186/// Returns the identifier with the SQL parser's canonicalization: quoted
187/// identifiers keep their case, unquoted ones fold to lowercase.
188fn expect_ident(cursor: &mut Tokens<'_>) -> Option<String> {
189    match cursor.next() {
190        Some(Token::Word(w)) => Some(if w.quote_style.is_none() {
191            w.value.to_lowercase()
192        } else {
193            w.value.clone()
194        }),
195        _ => None,
196    }
197}
198
199fn parse_object_name(cursor: &mut Tokens<'_>) -> Option<Vec<String>> {
200    let mut parts = vec![expect_ident(cursor)?];
201    while matches!(cursor.peek(), Some(Token::Period)) {
202        cursor.next();
203        parts.push(expect_ident(cursor)?);
204    }
205    Some(parts)
206}
207
208/// Parses an optional `(column, ...)` list. `None` means the statement is
209/// not a copy-in shape (either no list, or a malformed list better left to
210/// the regular parser to report).
211fn parse_column_list(cursor: &mut Tokens<'_>) -> Option<Vec<String>> {
212    if !matches!(cursor.peek(), Some(Token::LParen)) {
213        return Some(Vec::new());
214    }
215    cursor.next();
216    let mut columns = vec![expect_ident(cursor)?];
217    loop {
218        match cursor.next() {
219            Some(Token::Comma) => columns.push(expect_ident(cursor)?),
220            Some(Token::RParen) => break,
221            _ => return None,
222        }
223    }
224    Some(columns)
225}
226
227/// Parses `NAME [value] (, NAME [value])*` until the closing `)`.
228///
229/// Option values are open-ended: a string literal, a number, an identifier
230/// or a boolean word. Whether an option makes sense is decided by the
231/// format codec, not here.
232fn parse_option_list(cursor: &mut Tokens<'_>, options: &mut CopyOptionSet) -> PgWireResult<()> {
233    loop {
234        let key = match cursor.next() {
235            Some(Token::Word(w)) if w.quote_style.is_none() => w.value.to_lowercase(),
236            _ => return Err(bad_copy_data("expected an option name after '(' or ','")),
237        };
238        let value = match cursor.peek() {
239            Some(Token::Comma) | Some(Token::RParen) => CopyOptionValue::Flag,
240            Some(Token::SingleQuotedString(s)) => {
241                let value = CopyOptionValue::String(s.clone());
242                cursor.next();
243                value
244            }
245            Some(Token::EscapedStringLiteral(s)) => {
246                // E'...' escape-string constants, e.g. DELIMITER E'\t':
247                // sqlparser's tokenizer already decoded the escapes, so
248                // the payload is used as-is, like SingleQuotedString.
249                let value = CopyOptionValue::String(s.clone());
250                cursor.next();
251                value
252            }
253            Some(Token::Number(n, _)) => {
254                let value = CopyOptionValue::Number(n.clone());
255                cursor.next();
256                value
257            }
258            Some(Token::Word(w)) => {
259                let value = if w.quote_style.is_none() {
260                    match w.value.to_ascii_lowercase().as_str() {
261                        "true" => CopyOptionValue::Boolean(true),
262                        "false" => CopyOptionValue::Boolean(false),
263                        _ => CopyOptionValue::Ident(w.value.clone()),
264                    }
265                } else {
266                    CopyOptionValue::Ident(w.value.clone())
267                };
268                cursor.next();
269                value
270            }
271            _ => return Err(bad_copy_data("expected an option value or ',' or ')'")),
272        };
273        options.insert(key, value)?;
274        match cursor.next() {
275            Some(Token::Comma) => continue,
276            Some(Token::RParen) => break,
277            _ => {
278                return Err(bad_copy_data(
279                    "expected ',' or ')' in COPY FROM STDIN options",
280                ));
281            }
282        }
283    }
284    Ok(())
285}
286
287#[cfg(test)]
288mod tests {
289    use super::*;
290
291    fn scan(sql: &str) -> PgWireResult<Option<RawCopyFromStdin>> {
292        Ok(match scan_copy_from_stdin(sql)? {
293            CopyInScan::CopyIn(raw) => Some(raw),
294            CopyInScan::NotCopyIn => None,
295        })
296    }
297
298    #[test]
299    fn test_scan_tableless() {
300        // Table-less statements: the codec decides whether a missing
301        // table is legal (the influxdb line protocol names its
302        // measurement per line).
303        let raw = scan("COPY FROM STDIN WITH (FORMAT influxdb, precision 'ns')")
304            .unwrap()
305            .unwrap();
306        assert!(raw.table.is_empty());
307        assert!(raw.columns.is_empty());
308
309        let raw = scan("copy from stdin (format influxdb)").unwrap().unwrap();
310        assert!(raw.table.is_empty());
311
312        // A column list without a table is not a copy-in statement.
313        assert!(scan("COPY (a, b) FROM STDIN").unwrap().is_none());
314        // COPY OUT and file COPY are unaffected.
315        assert!(scan("COPY (SELECT 1) TO STDOUT").unwrap().is_none());
316        assert!(scan("COPY t FROM '/tmp/x.csv'").unwrap().is_none());
317    }
318
319    #[test]
320    fn test_scan_basic() {
321        let raw = scan("COPY t FROM STDIN").unwrap().unwrap();
322        assert_eq!(raw.table, vec!["t"]);
323        assert!(raw.columns.is_empty());
324    }
325
326    #[test]
327    fn test_scan_qualified_table_and_columns() {
328        let raw = scan("COPY db.public.Metrics (ts, \"Host\") FROM STDIN")
329            .unwrap()
330            .unwrap();
331        assert_eq!(raw.table, vec!["db", "public", "metrics"]);
332        assert_eq!(raw.columns, vec!["ts", "Host"]);
333    }
334
335    #[test]
336    fn test_scan_not_copy_in() {
337        assert!(scan("COPY (SELECT 1) TO STDOUT").unwrap().is_none());
338        assert!(
339            scan("COPY t TO '/tmp/x.csv' WITH (FORMAT csv)")
340                .unwrap()
341                .is_none()
342        );
343        assert!(scan("SELECT 1").unwrap().is_none());
344        // GreptimeDB's own file-based COPY FROM.
345        assert!(scan("COPY t FROM '/tmp/x.csv'").unwrap().is_none());
346    }
347
348    #[test]
349    fn test_scan_leading_comment() {
350        let raw = scan("-- a comment\nCOPY t FROM STDIN").unwrap().unwrap();
351        assert_eq!(raw.table, vec!["t"]);
352    }
353
354    #[test]
355    fn test_scan_folds_unquoted_identifiers() {
356        // Unquoted identifiers fold to lowercase, like the SQL parser.
357        let raw = scan("copy Metrics (Host, ts) from stdin").unwrap().unwrap();
358        assert_eq!(raw.table, vec!["metrics"]);
359        assert_eq!(raw.columns, vec!["host", "ts"]);
360    }
361
362    #[test]
363    fn test_scan_preserves_quoted_identifiers() {
364        let raw = scan("COPY \"Metrics\" (\"Host\") FROM STDIN")
365            .unwrap()
366            .unwrap();
367        assert_eq!(raw.table, vec!["Metrics"]);
368        assert_eq!(raw.columns, vec!["Host"]);
369    }
370
371    #[test]
372    fn test_scan_options_kinds() {
373        let raw = scan(
374            "COPY t FROM STDIN WITH (FORMAT csv, DELIMITER ';', HEADER, \
375             skip_bad_records true, precision 3, mode strict)",
376        )
377        .unwrap()
378        .unwrap();
379        let mut options = raw.options;
380        assert_eq!(options.take_format().unwrap().as_deref(), Some("csv"));
381        assert_eq!(options.take_single_byte("DELIMITER").unwrap(), Some(b';'));
382        assert_eq!(options.take_flag("HEADER").unwrap(), Some(true));
383        assert_eq!(
384            options.take_value("skip_bad_records"),
385            Some(CopyOptionValue::Boolean(true))
386        );
387        assert_eq!(
388            options.take_value("precision"),
389            Some(CopyOptionValue::Number("3".into()))
390        );
391        assert_eq!(
392            options.take_value("mode"),
393            Some(CopyOptionValue::Ident("strict".into()))
394        );
395    }
396
397    #[test]
398    fn test_scan_escaped_string_constants() {
399        // sqlparser's tokenizer already decodes E'...' escapes, so the
400        // payload is used as-is: E'\t' is a tab.
401        let mut raw = scan("COPY t FROM STDIN WITH (FORMAT csv, DELIMITER E'\\t')")
402            .unwrap()
403            .unwrap();
404        assert_eq!(
405            raw.options.take_single_byte("DELIMITER").unwrap(),
406            Some(b'\t')
407        );
408
409        // E'\\' is a single literal backslash (a legal one-byte
410        // delimiter), NOT an escaped backslash-t.
411        let mut raw = scan("COPY t FROM STDIN WITH (FORMAT csv, DELIMITER E'\\\\')")
412            .unwrap()
413            .unwrap();
414        assert_eq!(
415            raw.options.take_single_byte("DELIMITER").unwrap(),
416            Some(b'\\')
417        );
418
419        // E'\\t' decodes to TWO characters — backslash followed by t —
420        // and must be rejected as a multi-character delimiter instead of
421        // being unescaped a second time into a tab.
422        // The scanner leaves option validation to the codec; use the
423        // full resolution path to assert the delimiter rejection.
424        let err = match crate::postgres::copy_in::copy_from_stdin_from_sql(
425            "COPY t FROM STDIN WITH (FORMAT csv, DELIMITER E'\\\\t')",
426        ) {
427            Ok(Some(_)) => panic!("expected the multi-character delimiter to be rejected"),
428            Ok(None) => panic!("expected the statement to scan as copy-in"),
429            Err(e) => e,
430        };
431        assert!(
432            err.to_string().contains("single one-byte character"),
433            "{err}"
434        );
435    }
436
437    #[test]
438    fn test_scan_semicolon_inside_option_value() {
439        // A quoted semicolon is not a statement separator.
440        let mut raw = scan("COPY t FROM STDIN WITH (NULL ';')").unwrap().unwrap();
441        assert_eq!(
442            raw.options.take_string("NULL").unwrap().as_deref(),
443            Some(";")
444        );
445    }
446
447    #[test]
448    fn test_scan_trailing_statements_rejected() {
449        let err = scan("COPY t FROM STDIN; SELECT 1").unwrap_err();
450        assert!(err.to_string().contains("alone"));
451
452        // Trailing separators alone are fine.
453        assert!(scan("COPY t FROM STDIN;").unwrap().is_some());
454        assert!(scan("COPY t FROM STDIN ; ;").unwrap().is_some());
455    }
456
457    #[test]
458    fn test_scan_options_without_with_keyword() {
459        let raw = scan("COPY t FROM STDIN (FORMAT csv, HEADER)")
460            .unwrap()
461            .unwrap();
462        let mut options = raw.options;
463        assert_eq!(options.take_format().unwrap().as_deref(), Some("csv"));
464        assert_eq!(options.take_flag("HEADER").unwrap(), Some(true));
465        assert!(options.reject_leftovers("csv").is_ok());
466    }
467
468    #[test]
469    fn test_scan_malformed_after_stdin() {
470        let err = scan("COPY t FROM STDIN WITH").unwrap_err();
471        assert!(err.to_string().contains("'('"), "{err}");
472        let err = scan("COPY t FROM STDIN WITH (").unwrap_err();
473        assert!(err.to_string().contains("option"), "{err}");
474        // A malformed column list never commits to the copy-in shape and
475        // falls through to the regular parser.
476        assert!(scan("COPY t (a,").unwrap().is_none());
477    }
478
479    #[test]
480    fn test_scan_redundant_option() {
481        let err = scan("COPY t FROM STDIN WITH (FORMAT csv, FORMAT text)").unwrap_err();
482        assert!(err.to_string().contains("redundant FORMAT"));
483    }
484}
servers/postgres/copy_in/parse.rs

servers/postgres/copy_in/
parse.rs
