qusql_parse/
copy.rs

1// Licensed under the Apache License, Version 2.0 (the "License");
2// you may not use this file except in compliance with the License.
3// You may obtain a copy of the License at
4//
5// http://www.apache.org/licenses/LICENSE-2.0
6//
7// Unless required by applicable law or agreed to in writing, software
8// distributed under the License is distributed on an "AS IS" BASIS,
9// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10// See the License for the specific language governing permissions and
11// limitations under the License.
12
13//! PostgreSQL `COPY` statement AST and parser.
14//!
15//! Implements the full COPY syntax including:
16//! - Table or subquery as source
17//! - FROM/TO direction with file, PROGRAM, STDIN, or STDOUT
18//! - Modern `WITH ( options )` syntax
19//! - Legacy bare-keyword options (pre-PostgreSQL 9.0 syntax)
20//!
21//! Reference: <https://www.postgresql.org/docs/current/sql-copy.html>
22//!
23//! ```text
24//! COPY table_name [ ( column_name [, ...] ) ]
25//!     FROM { 'filename' | PROGRAM 'command' | STDIN }
26//!     [ [ WITH ] ( option [, ...] ) ]
27//!     [ WHERE condition ]
28//!
29//! COPY { table_name [ ( column_name [, ...] ) ] | ( query ) }
30//!     TO { 'filename' | PROGRAM 'command' | STDOUT }
31//!     [ [ WITH ] ( option [, ...] ) ]
32//! ```
33
34use alloc::{boxed::Box, vec::Vec};
35
36use crate::{
37    Identifier, QualifiedName, SString, Span, Spanned, Statement,
38    expression::{Expression, PRIORITY_MAX, parse_expression_unreserved},
39    keywords::Keyword,
40    lexer::Token,
41    parser::{ParseError, Parser},
42    qualified_name::parse_qualified_name_unreserved,
43    statement::parse_statement,
44};
45
46/// A column list in a COPY option: either `*` (all columns) or `( col, ... )`.
47///
48/// Used by `FORCE_QUOTE`, `FORCE_NOT_NULL`, and `FORCE_NULL` options.
49#[derive(Clone, Debug)]
50pub enum CopyColumnList<'a> {
51    /// `*` — applies to all columns
52    All(Span),
53    /// `( col1, col2, ... )` — applies to named columns
54    Columns {
55        lparen_span: Span,
56        columns: Vec<Identifier<'a>>,
57        rparen_span: Span,
58    },
59}
60
61impl<'a> Spanned for CopyColumnList<'a> {
62    fn span(&self) -> Span {
63        match self {
64            CopyColumnList::All(s) => s.clone(),
65            CopyColumnList::Columns {
66                lparen_span,
67                rparen_span,
68                ..
69            } => lparen_span.join_span(rparen_span),
70        }
71    }
72}
73
74fn parse_copy_column_list<'a>(
75    parser: &mut Parser<'a, '_>,
76) -> Result<CopyColumnList<'a>, ParseError> {
77    if let Some(star) = parser.skip_token(Token::Mul) {
78        return Ok(CopyColumnList::All(star));
79    }
80    let lparen_span = parser.consume_token(Token::LParen)?;
81    let mut columns = Vec::new();
82    if !matches!(parser.token, Token::RParen) {
83        loop {
84            parser.recovered(
85                "')' or ','",
86                &|t| matches!(t, Token::RParen | Token::Comma),
87                |parser| {
88                    columns.push(parser.consume_plain_identifier_unreserved()?);
89                    Ok(())
90                },
91            )?;
92            if matches!(parser.token, Token::RParen) {
93                break;
94            }
95            parser.consume_token(Token::Comma)?;
96        }
97    }
98    let rparen_span = parser.consume_token(Token::RParen)?;
99    Ok(CopyColumnList::Columns {
100        lparen_span,
101        columns,
102        rparen_span,
103    })
104}
105
106/// Value for the `HEADER` option: a boolean or `MATCH`.
107#[derive(Clone, Debug)]
108pub enum CopyHeaderValue {
109    /// `TRUE`, `ON`, or `1` — include the header line
110    True(Span),
111    /// `FALSE`, `OFF`, or `0` — no header line
112    False(Span),
113    /// `MATCH` — validate that the header matches the table columns (FROM only)
114    Match(Span),
115}
116
117impl Spanned for CopyHeaderValue {
118    fn span(&self) -> Span {
119        match self {
120            CopyHeaderValue::True(s) | CopyHeaderValue::False(s) | CopyHeaderValue::Match(s) => {
121                s.clone()
122            }
123        }
124    }
125}
126
127/// A single `WITH ( ... )` option in a COPY statement.
128#[derive(Clone, Debug)]
129pub enum CopyOption<'a> {
130    /// `FORMAT format_name` — data format: `text`, `csv`, or `binary`
131    Format { span: Span, name: Identifier<'a> },
132    /// `FREEZE [ boolean ]` — copy with rows already frozen
133    Freeze { span: Span, value: Option<bool> },
134    /// `DELIMITER 'delimiter_character'` — column separator
135    Delimiter { span: Span, value: SString<'a> },
136    /// `NULL 'null_string'` — string representing NULL
137    Null { span: Span, value: SString<'a> },
138    /// `DEFAULT 'default_string'` — string representing DEFAULT (FROM only)
139    Default { span: Span, value: SString<'a> },
140    /// `HEADER [ boolean | MATCH ]` — header line handling
141    Header {
142        span: Span,
143        value: Option<CopyHeaderValue>,
144    },
145    /// `QUOTE 'quote_character'` — CSV quoting character
146    Quote { span: Span, value: SString<'a> },
147    /// `ESCAPE 'escape_character'` — CSV escape character
148    Escape { span: Span, value: SString<'a> },
149    /// `FORCE_QUOTE { ( col, ... ) | * }` — force quoting (TO CSV only)
150    ForceQuote {
151        span: Span,
152        columns: CopyColumnList<'a>,
153    },
154    /// `FORCE_NOT_NULL { ( col, ... ) | * }` — never match null string (FROM CSV only)
155    ForceNotNull {
156        span: Span,
157        columns: CopyColumnList<'a>,
158    },
159    /// `FORCE_NULL { ( col, ... ) | * }` — match null string even when quoted (FROM CSV only)
160    ForceNull {
161        span: Span,
162        columns: CopyColumnList<'a>,
163    },
164    /// `ON_ERROR error_action` — behaviour when input value conversion fails (FROM only)
165    OnError { span: Span, action: Identifier<'a> },
166    /// `REJECT_LIMIT maxerror` — maximum conversion errors before failing (FROM only)
167    RejectLimit { span: Span, limit: Expression<'a> },
168    /// `ENCODING 'encoding_name'` — file encoding
169    Encoding { span: Span, value: SString<'a> },
170    /// `LOG_VERBOSITY verbosity` — message verbosity: `default`, `verbose`, or `silent`
171    LogVerbosity {
172        span: Span,
173        verbosity: Identifier<'a>,
174    },
175}
176
177impl<'a> Spanned for CopyOption<'a> {
178    fn span(&self) -> Span {
179        match self {
180            CopyOption::Format { span, name } => span.join_span(name),
181            CopyOption::Freeze { span, .. } => span.clone(),
182            CopyOption::Delimiter { span, value } => span.join_span(value),
183            CopyOption::Null { span, value } => span.join_span(value),
184            CopyOption::Default { span, value } => span.join_span(value),
185            CopyOption::Header { span, value } => span.join_span(value),
186            CopyOption::Quote { span, value } => span.join_span(value),
187            CopyOption::Escape { span, value } => span.join_span(value),
188            CopyOption::ForceQuote { span, columns } => span.join_span(columns),
189            CopyOption::ForceNotNull { span, columns } => span.join_span(columns),
190            CopyOption::ForceNull { span, columns } => span.join_span(columns),
191            CopyOption::OnError { span, action } => span.join_span(action),
192            CopyOption::RejectLimit { span, limit } => span.join_span(limit),
193            CopyOption::Encoding { span, value } => span.join_span(value),
194            CopyOption::LogVerbosity { span, verbosity } => span.join_span(verbosity),
195        }
196    }
197}
198
199/// The source / destination of a COPY: either a table (with optional column list)
200/// or a parenthesised query (only valid with `TO`).
201#[derive(Clone, Debug)]
202pub enum CopySource<'a> {
203    /// Plain table reference, e.g. `public.actor (col1, col2)`
204    Table {
205        name: QualifiedName<'a>,
206        columns: Vec<Identifier<'a>>,
207    },
208    /// Subquery source, e.g. `(SELECT * FROM t)` — only valid with `TO`
209    Query {
210        lparen_span: Span,
211        query: Box<Statement<'a>>,
212        rparen_span: Span,
213    },
214}
215
216impl<'a> Spanned for CopySource<'a> {
217    fn span(&self) -> Span {
218        match self {
219            CopySource::Table { name, columns } => name.span().join_span(columns),
220            CopySource::Query {
221                lparen_span,
222                query,
223                rparen_span,
224            } => lparen_span.join_span(query.as_ref()).join_span(rparen_span),
225        }
226    }
227}
228
229/// Where to read from / write to in a COPY statement.
230#[derive(Clone, Debug)]
231pub enum CopyLocation<'a> {
232    /// A file path: `'path/to/file'`
233    Filename(SString<'a>),
234    /// A shell command: `PROGRAM 'command'`
235    Program {
236        program_span: Span,
237        command: SString<'a>,
238    },
239    /// Standard input (only valid with `FROM`)
240    Stdin(Span),
241    /// Standard output (only valid with `TO`)
242    Stdout(Span),
243}
244
245impl<'a> Spanned for CopyLocation<'a> {
246    fn span(&self) -> Span {
247        match self {
248            CopyLocation::Filename(s) => s.span(),
249            CopyLocation::Program {
250                program_span,
251                command,
252            } => program_span.join_span(command),
253            CopyLocation::Stdin(s) | CopyLocation::Stdout(s) => s.clone(),
254        }
255    }
256}
257
258/// A PostgreSQL `COPY ... FROM` statement.
259///
260/// ```sql
261/// COPY table_name FROM STDIN;
262/// COPY country FROM '/tmp/data.csv' WITH (FORMAT csv, HEADER);
263/// ```
264#[derive(Clone, Debug)]
265pub struct CopyFrom<'a> {
266    /// Span of `COPY`
267    pub copy_span: Span,
268    /// Table source (a subquery is invalid for FROM)
269    pub source: CopySource<'a>,
270    /// Span of `FROM`
271    pub from_span: Span,
272    /// Location: file, PROGRAM, or STDIN
273    pub location: CopyLocation<'a>,
274    /// Span of `WITH` keyword if present
275    pub with_span: Option<Span>,
276    /// Options specified after `WITH ( ... )` or as legacy bare keywords
277    pub options: Vec<CopyOption<'a>>,
278    /// Optional `WHERE condition`
279    pub where_: Option<(Span, Expression<'a>)>,
280}
281
282impl<'a> Spanned for CopyFrom<'a> {
283    fn span(&self) -> Span {
284        self.copy_span
285            .join_span(&self.source)
286            .join_span(&self.from_span)
287            .join_span(&self.location)
288            .join_span(&self.with_span)
289            .join_span(&self.options)
290            .join_span(&self.where_)
291    }
292}
293
294impl<'a> CopyFrom<'a> {
295    /// Returns `true` when this COPY reads its data from stdin, signalling the
296    /// parser to consume the following inline data block up to `\.`.
297    pub(crate) fn reads_from_stdin(&self) -> bool {
298        matches!(&self.location, CopyLocation::Stdin(_))
299    }
300}
301
302/// A PostgreSQL `COPY ... TO` statement.
303///
304/// ```sql
305/// COPY (SELECT * FROM t) TO '/tmp/out.csv' WITH (FORMAT csv, HEADER);
306/// COPY country TO PROGRAM 'gzip > /tmp/country.gz';
307/// ```
308#[derive(Clone, Debug)]
309pub struct CopyTo<'a> {
310    /// Span of `COPY`
311    pub copy_span: Span,
312    /// Table or subquery source
313    pub source: CopySource<'a>,
314    /// Span of `TO`
315    pub to_span: Span,
316    /// Location: file, PROGRAM, or STDOUT
317    pub location: CopyLocation<'a>,
318    /// Span of `WITH` keyword if present
319    pub with_span: Option<Span>,
320    /// Options specified after `WITH ( ... )` or as legacy bare keywords
321    pub options: Vec<CopyOption<'a>>,
322}
323
324impl<'a> Spanned for CopyTo<'a> {
325    fn span(&self) -> Span {
326        self.copy_span
327            .join_span(&self.source)
328            .join_span(&self.to_span)
329            .join_span(&self.location)
330            .join_span(&self.with_span)
331            .join_span(&self.options)
332    }
333}
334
335/// Parse the location (`'file'`, `PROGRAM 'cmd'`, `STDIN`, or `STDOUT`)
336/// that follows `FROM` or `TO`.
337fn parse_copy_location<'a>(parser: &mut Parser<'a, '_>) -> Result<CopyLocation<'a>, ParseError> {
338    match &parser.token {
339        Token::String(_, _) => Ok(CopyLocation::Filename(parser.consume_string()?)),
340        Token::Ident(_, Keyword::PROGRAM) => {
341            let program_span = parser.consume_keyword(Keyword::PROGRAM)?;
342            let command = parser.consume_string()?;
343            Ok(CopyLocation::Program {
344                program_span,
345                command,
346            })
347        }
348        Token::Ident(_, Keyword::STDIN) => {
349            Ok(CopyLocation::Stdin(parser.consume_keyword(Keyword::STDIN)?))
350        }
351        Token::Ident(_, Keyword::STDOUT) => Ok(CopyLocation::Stdout(
352            parser.consume_keyword(Keyword::STDOUT)?,
353        )),
354        _ => parser.expected_failure("'filename', PROGRAM, STDIN, or STDOUT"),
355    }
356}
357
358/// Parse a single option from a modern `WITH ( option, ... )` option list.
359fn parse_copy_option_modern<'a>(parser: &mut Parser<'a, '_>) -> Result<CopyOption<'a>, ParseError> {
360    match &parser.token {
361        Token::Ident(_, Keyword::FORMAT) => {
362            let span = parser.consume_keyword(Keyword::FORMAT)?;
363            let name = parser.consume_plain_identifier_unreserved()?;
364            Ok(CopyOption::Format { span, name })
365        }
366        Token::Ident(_, Keyword::FREEZE) => {
367            let span = parser.consume_keyword(Keyword::FREEZE)?;
368            let value = parser.try_parse_bool().map(|(b, _)| b);
369            Ok(CopyOption::Freeze { span, value })
370        }
371        Token::Ident(_, Keyword::DELIMITER) => {
372            let span = parser.consume_keyword(Keyword::DELIMITER)?;
373            let value = parser.consume_string()?;
374            Ok(CopyOption::Delimiter { span, value })
375        }
376        Token::Ident(_, Keyword::NULL) => {
377            let span = parser.consume_keyword(Keyword::NULL)?;
378            let value = parser.consume_string()?;
379            Ok(CopyOption::Null { span, value })
380        }
381        Token::Ident(_, Keyword::DEFAULT) => {
382            let span = parser.consume_keyword(Keyword::DEFAULT)?;
383            let value = parser.consume_string()?;
384            Ok(CopyOption::Default { span, value })
385        }
386        Token::Ident(_, Keyword::HEADER) => {
387            let span = parser.consume_keyword(Keyword::HEADER)?;
388            let value = match &parser.token {
389                Token::Ident(_, Keyword::MATCH) => Some(CopyHeaderValue::Match(
390                    parser.consume_keyword(Keyword::MATCH)?,
391                )),
392                _ => parser.try_parse_bool().map(|(b, s)| {
393                    if b {
394                        CopyHeaderValue::True(s)
395                    } else {
396                        CopyHeaderValue::False(s)
397                    }
398                }),
399            };
400            Ok(CopyOption::Header { span, value })
401        }
402        Token::Ident(_, Keyword::QUOTE) => {
403            let span = parser.consume_keyword(Keyword::QUOTE)?;
404            let value = parser.consume_string()?;
405            Ok(CopyOption::Quote { span, value })
406        }
407        Token::Ident(_, Keyword::ESCAPE) => {
408            let span = parser.consume_keyword(Keyword::ESCAPE)?;
409            let value = parser.consume_string()?;
410            Ok(CopyOption::Escape { span, value })
411        }
412        Token::Ident(_, Keyword::FORCE_QUOTE) => {
413            let span = parser.consume_keyword(Keyword::FORCE_QUOTE)?;
414            let columns = parse_copy_column_list(parser)?;
415            Ok(CopyOption::ForceQuote { span, columns })
416        }
417        Token::Ident(_, Keyword::FORCE_NOT_NULL) => {
418            let span = parser.consume_keyword(Keyword::FORCE_NOT_NULL)?;
419            let columns = parse_copy_column_list(parser)?;
420            Ok(CopyOption::ForceNotNull { span, columns })
421        }
422        Token::Ident(_, Keyword::FORCE_NULL) => {
423            let span = parser.consume_keyword(Keyword::FORCE_NULL)?;
424            let columns = parse_copy_column_list(parser)?;
425            Ok(CopyOption::ForceNull { span, columns })
426        }
427        Token::Ident(_, Keyword::ON_ERROR) => {
428            let span = parser.consume_keyword(Keyword::ON_ERROR)?;
429            let action = parser.consume_plain_identifier_unreserved()?;
430            Ok(CopyOption::OnError { span, action })
431        }
432        Token::Ident(_, Keyword::REJECT_LIMIT) => {
433            let span = parser.consume_keyword(Keyword::REJECT_LIMIT)?;
434            let limit = parse_expression_unreserved(parser, PRIORITY_MAX)?;
435            Ok(CopyOption::RejectLimit { span, limit })
436        }
437        Token::Ident(_, Keyword::ENCODING) => {
438            let span = parser.consume_keyword(Keyword::ENCODING)?;
439            let value = parser.consume_string()?;
440            Ok(CopyOption::Encoding { span, value })
441        }
442        Token::Ident(_, Keyword::LOG_VERBOSITY) => {
443            let span = parser.consume_keyword(Keyword::LOG_VERBOSITY)?;
444            let verbosity = parser.consume_plain_identifier_unreserved()?;
445            Ok(CopyOption::LogVerbosity { span, verbosity })
446        }
447        _ => parser.expected_failure(
448            "COPY option (FORMAT, FREEZE, DELIMITER, NULL, DEFAULT, HEADER, QUOTE, ESCAPE, \
449             FORCE_QUOTE, FORCE_NOT_NULL, FORCE_NULL, ON_ERROR, REJECT_LIMIT, ENCODING, \
450             LOG_VERBOSITY)",
451        ),
452    }
453}
454
455/// Parse a comma-separated list of options inside `( ... )`.
456fn parse_copy_options_modern<'a>(
457    parser: &mut Parser<'a, '_>,
458) -> Result<(Span, Vec<CopyOption<'a>>), ParseError> {
459    let lparen = parser.consume_token(Token::LParen)?;
460    let mut options = Vec::new();
461    if !matches!(parser.token, Token::RParen) {
462        loop {
463            parser.recovered(
464                "')' or ','",
465                &|t| matches!(t, Token::RParen | Token::Comma),
466                |parser| {
467                    options.push(parse_copy_option_modern(parser)?);
468                    Ok(())
469                },
470            )?;
471            if matches!(parser.token, Token::RParen) {
472                break;
473            }
474            parser.consume_token(Token::Comma)?;
475        }
476    }
477    let rparen = parser.consume_token(Token::RParen)?;
478    Ok((lparen.join_span(&rparen), options))
479}
480
481/// Parse legacy bare-keyword options (pre-PostgreSQL 9.0 syntax):
482///
483/// ```text
484/// [ [ WITH ]
485///     [ BINARY ]
486///     [ DELIMITER [ AS ] 'char' ]
487///     [ NULL [ AS ] 'string' ]
488///     [ CSV
489///         [ HEADER ]
490///         [ QUOTE [ AS ] 'char' ]
491///         [ ESCAPE [ AS ] 'char' ]
492///         [ FORCE NOT NULL col, ...   (FROM only) ]
493///         [ FORCE QUOTE { col, ... | * }  (TO only) ]
494///     ]
495/// ]
496/// ```
497fn parse_copy_options_legacy<'a>(
498    parser: &mut Parser<'a, '_>,
499) -> Result<Vec<CopyOption<'a>>, ParseError> {
500    let mut options = Vec::new();
501
502    // BINARY — sets FORMAT to binary
503    if let Some(span) = parser.skip_keyword(Keyword::BINARY) {
504        options.push(CopyOption::Format {
505            span: span.clone(),
506            name: Identifier::new("binary", span),
507        });
508    }
509
510    // DELIMITER [ AS ] 'char'
511    if let Some(span) = parser.skip_keyword(Keyword::DELIMITER) {
512        parser.skip_keyword(Keyword::AS);
513        let value = parser.consume_string()?;
514        options.push(CopyOption::Delimiter { span, value });
515    }
516
517    // NULL [ AS ] 'string'
518    if let Some(span) = parser.skip_keyword(Keyword::NULL) {
519        parser.skip_keyword(Keyword::AS);
520        let value = parser.consume_string()?;
521        options.push(CopyOption::Null { span, value });
522    }
523
524    // CSV [ HEADER ] [ QUOTE [ AS ] 'char' ] [ ESCAPE [ AS ] 'char' ]
525    //     [ FORCE NOT NULL col, ... | FORCE QUOTE { col, ... | * } ]
526    if let Some(csv_span) = parser.skip_keyword(Keyword::CSV) {
527        // Record FORMAT csv
528        options.push(CopyOption::Format {
529            span: csv_span.clone(),
530            name: Identifier::new("csv", csv_span),
531        });
532
533        // HEADER
534        if let Some(span) = parser.skip_keyword(Keyword::HEADER) {
535            options.push(CopyOption::Header { span, value: None });
536        }
537
538        // QUOTE [ AS ] 'char'
539        if let Some(span) = parser.skip_keyword(Keyword::QUOTE) {
540            parser.skip_keyword(Keyword::AS);
541            let value = parser.consume_string()?;
542            options.push(CopyOption::Quote { span, value });
543        }
544
545        // ESCAPE [ AS ] 'char'
546        if let Some(span) = parser.skip_keyword(Keyword::ESCAPE) {
547            parser.skip_keyword(Keyword::AS);
548            let value = parser.consume_string()?;
549            options.push(CopyOption::Escape { span, value });
550        }
551
552        // FORCE NOT NULL col1, col2, ...   (FROM CSV)
553        // FORCE QUOTE { col1, col2, ... | * }   (TO CSV)
554        if let Some(force_span) = parser.skip_keyword(Keyword::FORCE) {
555            match &parser.token {
556                Token::Ident(_, Keyword::NOT) => {
557                    // FORCE NOT NULL col, ...
558                    let span = force_span
559                        .join_span(&parser.consume_keyword(Keyword::NOT)?)
560                        .join_span(&parser.consume_keyword(Keyword::NULL)?);
561                    let mut cols = Vec::new();
562                    loop {
563                        cols.push(parser.consume_plain_identifier_unreserved()?);
564                        if parser.skip_token(Token::Comma).is_none() {
565                            break;
566                        }
567                    }
568                    let col_span = if let Some(last) = cols.last() {
569                        span.clone().join_span(last)
570                    } else {
571                        span.clone()
572                    };
573                    options.push(CopyOption::ForceNotNull {
574                        span,
575                        columns: CopyColumnList::Columns {
576                            lparen_span: col_span.clone(),
577                            columns: cols,
578                            rparen_span: col_span,
579                        },
580                    });
581                }
582                Token::Ident(_, Keyword::QUOTE) => {
583                    // FORCE QUOTE { * | col, ... }
584                    let span = force_span.join_span(&parser.consume_keyword(Keyword::QUOTE)?);
585                    let columns = if let Some(star) = parser.skip_token(Token::Mul) {
586                        CopyColumnList::All(star)
587                    } else {
588                        let mut cols = Vec::new();
589                        loop {
590                            cols.push(parser.consume_plain_identifier_unreserved()?);
591                            if parser.skip_token(Token::Comma).is_none() {
592                                break;
593                            }
594                        }
595                        let col_span = if let Some(last) = cols.last() {
596                            span.clone().join_span(last)
597                        } else {
598                            span.clone()
599                        };
600                        CopyColumnList::Columns {
601                            lparen_span: col_span.clone(),
602                            columns: cols,
603                            rparen_span: col_span,
604                        }
605                    };
606                    options.push(CopyOption::ForceQuote { span, columns });
607                }
608                _ => {
609                    parser
610                        .expected_failure("NOT NULL or QUOTE after FORCE in COPY legacy options")?;
611                }
612            }
613        }
614    }
615
616    Ok(options)
617}
618
619/// Shared helper: parse the optional `[ WITH ] ( options )` or legacy bare
620/// options that follow the location in both COPY FROM and COPY TO.
621fn parse_copy_options_clause<'a>(
622    parser: &mut Parser<'a, '_>,
623) -> Result<(Option<Span>, Vec<CopyOption<'a>>), ParseError> {
624    if let Some(with_span) = parser.skip_keyword(Keyword::WITH) {
625        if matches!(parser.token, Token::LParen) {
626            let (_, opts) = parse_copy_options_modern(parser)?;
627            Ok((Some(with_span), opts))
628        } else {
629            let opts = parse_copy_options_legacy(parser)?;
630            Ok((Some(with_span), opts))
631        }
632    } else if matches!(parser.token, Token::LParen) {
633        let (_, opts) = parse_copy_options_modern(parser)?;
634        Ok((None, opts))
635    } else if {
636        let parser: &Parser<'_, '_> = parser;
637        matches!(
638            &parser.token,
639            Token::Ident(
640                _,
641                Keyword::BINARY | Keyword::DELIMITER | Keyword::NULL | Keyword::CSV
642            )
643        )
644    } {
645        let opts = parse_copy_options_legacy(parser)?;
646        Ok((None, opts))
647    } else {
648        Ok((None, Vec::new()))
649    }
650}
651
652/// Parse the source (table or parenthesised query) that immediately follows
653/// the `COPY` keyword.
654fn parse_copy_source<'a>(parser: &mut Parser<'a, '_>) -> Result<CopySource<'a>, ParseError> {
655    if matches!(parser.token, Token::LParen) {
656        let lparen_span = parser.consume_token(Token::LParen)?;
657        let query =
658            parser.recovered(
659                "')'",
660                &|t| t == &Token::RParen,
661                |parser| match parse_statement(parser)? {
662                    Some(s) => Ok(Some(s)),
663                    None => {
664                        parser.expected_error("query");
665                        Ok(None)
666                    }
667                },
668            )?;
669        let rparen_span = parser.consume_token(Token::RParen)?;
670        let query = match query {
671            Some(s) => Box::new(s),
672            None => return Err(crate::parser::ParseError::Unrecovered),
673        };
674        Ok(CopySource::Query {
675            lparen_span,
676            query,
677            rparen_span,
678        })
679    } else {
680        let name = parse_qualified_name_unreserved(parser)?;
681        let columns = if matches!(parser.token, Token::LParen) {
682            parser.consume_token(Token::LParen)?;
683            let mut cols = Vec::new();
684            if !matches!(parser.token, Token::RParen) {
685                loop {
686                    parser.recovered(
687                        "')' or ','",
688                        &|t| matches!(t, Token::RParen | Token::Comma),
689                        |parser| {
690                            cols.push(parser.consume_plain_identifier_unreserved()?);
691                            Ok(())
692                        },
693                    )?;
694                    if matches!(parser.token, Token::RParen) {
695                        break;
696                    }
697                    parser.consume_token(Token::Comma)?;
698                }
699            }
700            parser.consume_token(Token::RParen)?;
701            cols
702        } else {
703            Vec::new()
704        };
705        Ok(CopySource::Table { name, columns })
706    }
707}
708
709/// Parse `COPY ... FROM location [options] [WHERE]`.
710fn parse_copy_from<'a>(
711    parser: &mut Parser<'a, '_>,
712    copy_span: Span,
713    source: CopySource<'a>,
714) -> Result<CopyFrom<'a>, ParseError> {
715    let from_span = parser.consume_keyword(Keyword::FROM)?;
716    let location = parse_copy_location(parser)?;
717
718    if matches!(source, CopySource::Query { .. }) {
719        parser.err(
720            "Subquery source is only valid with COPY ... TO, not FROM",
721            &from_span,
722        );
723    }
724
725    let (with_span, options) = parse_copy_options_clause(parser)?;
726
727    let where_ = if let Some(where_span) = parser.skip_keyword(Keyword::WHERE) {
728        let cond = parse_expression_unreserved(parser, PRIORITY_MAX)?;
729        Some((where_span, cond))
730    } else {
731        None
732    };
733
734    Ok(CopyFrom {
735        copy_span,
736        source,
737        from_span,
738        location,
739        with_span,
740        options,
741        where_,
742    })
743}
744
745/// Parse `COPY ... TO location [options]`.
746fn parse_copy_to<'a>(
747    parser: &mut Parser<'a, '_>,
748    copy_span: Span,
749    source: CopySource<'a>,
750) -> Result<CopyTo<'a>, ParseError> {
751    let to_span = parser.consume_keyword(Keyword::TO)?;
752    let location = parse_copy_location(parser)?;
753    let (with_span, options) = parse_copy_options_clause(parser)?;
754
755    Ok(CopyTo {
756        copy_span,
757        source,
758        to_span,
759        location,
760        with_span,
761        options,
762    })
763}
764
765/// Parse a PostgreSQL `COPY` statement, returning either a
766/// [`Statement::CopyFrom`] or [`Statement::CopyTo`].
767pub(crate) fn parse_copy_statement<'a>(
768    parser: &mut Parser<'a, '_>,
769) -> Result<Statement<'a>, ParseError> {
770    let copy_span = parser.consume_keyword(Keyword::COPY)?;
771    let source = parse_copy_source(parser)?;
772
773    match &parser.token {
774        Token::Ident(_, Keyword::FROM) => Ok(Statement::CopyFrom(Box::new(parse_copy_from(
775            parser, copy_span, source,
776        )?))),
777        Token::Ident(_, Keyword::TO) => Ok(Statement::CopyTo(Box::new(parse_copy_to(
778            parser, copy_span, source,
779        )?))),
780        _ => parser.expected_failure("FROM or TO"),
781    }
782}
qusql_parse/copy.rs

qusql_parse/
copy.rs
