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

servers/postgres/copy_in/
parse.rs
