linkrs_query/parser/parsing/dml_parser/
copy.rs

1use super::*;
2
3impl DmlParser {
4    /// Parse COPY statement
5    /// Syntax:
6    ///   COPY VERTEX <tag> FROM 'path' [WITH (HEADER [true|false], DELIMITER ',' , BATCH_SIZE n)]
7    ///   COPY EDGE <edge_type> FROM 'path' [WITH ...]
8    ///   COPY <tag> FROM 'path' // defaults to VERTEX
9    pub fn parse_copy_statement(&mut self, ctx: &mut ParseContext) -> Result<Stmt, ParseError> {
10        use crate::parser::ast::stmt::{CopyDirection, CopyStmt, CopyTarget};
11
12        let start_span = ctx.current_span();
13        ctx.expect_token(TokenKind::Copy)?;
14
15        // Parse optional VERTEX/EDGE keyword and target name
16        let target = if ctx.match_token(TokenKind::Vertex) {
17            let tag = ctx.expect_identifier()?;
18            CopyTarget::Vertex(tag)
19        } else if ctx.match_token(TokenKind::Edge) {
20            let edge = ctx.expect_identifier()?;
21            CopyTarget::Edge(edge)
22        } else if ctx.match_token(TokenKind::Tag) {
23            let tag = ctx.expect_identifier()?;
24            CopyTarget::Vertex(tag)
25        } else {
26            // Bare identifier defaults to vertex tag
27            let name = ctx.expect_identifier()?;
28            CopyTarget::Vertex(name)
29        };
30
31        let direction = if ctx.match_token(TokenKind::To) {
32            CopyDirection::To
33        } else {
34            ctx.expect_token(TokenKind::From)?;
35            CopyDirection::From
36        };
37        let mut file_paths = Vec::new();
38        if ctx.check_token(TokenKind::LParen) {
39            ctx.expect_token(TokenKind::LParen)?;
40            loop {
41                if ctx.check_token(TokenKind::RParen) {
42                    ctx.next_token();
43                    break;
44                }
45                file_paths.push(ctx.expect_string_literal()?);
46                if ctx.match_token(TokenKind::Comma) {
47                    continue;
48                }
49                ctx.expect_token(TokenKind::RParen)?;
50                break;
51            }
52            if file_paths.is_empty() {
53                return Err(ParseError::new(
54                    crate::parser::core::error::ParseErrorKind::SyntaxError,
55                    "COPY FROM file list must contain at least one file".to_string(),
56                    ctx.current_position(),
57                ));
58            }
59        } else {
60            file_paths.push(ctx.expect_string_literal()?);
61        }
62        // Optional `BY COLUMN` marker right after the file list: column-wise
63        // merge for multi-file imports. Also accepted as a trailing marker
64        // after the option list below.
65        let mut by_column = false;
66        if ctx.check_token(TokenKind::By) || ctx.check_keyword("BY") {
67            if !ctx.match_token(TokenKind::By) {
68                let _ = ctx.consume_keyword("BY");
69            }
70            if ctx.check_keyword("COLUMN") {
71                let _ = ctx.consume_keyword("COLUMN");
72                by_column = true;
73            } else {
74                return Err(ParseError::new(
75                    crate::parser::core::error::ParseErrorKind::SyntaxError,
76                    "Expected COLUMN after BY in COPY statement".to_string(),
77                    ctx.current_position(),
78                ));
79            }
80        }
81        if direction == CopyDirection::To && file_paths.len() > 1 {
82            return Err(ParseError::new(
83                crate::parser::core::error::ParseErrorKind::SyntaxError,
84                "COPY TO supports a single file path".to_string(),
85                ctx.current_position(),
86            ));
87        }
88        if direction == CopyDirection::To && by_column {
89            return Err(ParseError::new(
90                crate::parser::core::error::ParseErrorKind::SyntaxError,
91                "BY COLUMN is only valid for COPY FROM".to_string(),
92                ctx.current_position(),
93            ));
94        }
95
96        // Defaults
97        let mut header = true;
98        let mut delimiter = ',';
99        let mut batch_size: Option<usize> = None;
100
101        // Consume the option name token (dedicated keyword or bare identifier).
102        fn consume_option_name(ctx: &mut ParseContext, kind: TokenKind, name: &str) {
103            if !ctx.match_token(kind) {
104                let _ = ctx.consume_keyword(name);
105            }
106        }
107
108        // Consume an optional `=` / `EQ` between an option name and its value.
109        fn consume_optional_assign(ctx: &mut ParseContext) {
110            if !ctx.match_token(TokenKind::Assign) {
111                let _ = ctx.match_token(TokenKind::Eq);
112            }
113        }
114
115        let parse_delimiter_value = |ctx: &mut ParseContext| -> Result<char, ParseError> {
116            consume_optional_assign(ctx);
117            match ctx.current_token().kind.clone() {
118                TokenKind::StringLiteral(s) | TokenKind::Identifier(s) => {
119                    ctx.next_token();
120                    s.chars().next().ok_or_else(|| {
121                        ParseError::new(
122                            crate::parser::core::error::ParseErrorKind::SyntaxError,
123                            "DELIMITER must be a single character".to_string(),
124                            ctx.current_position(),
125                        )
126                    })
127                }
128                other => Err(ParseError::new(
129                    crate::parser::core::error::ParseErrorKind::UnexpectedToken,
130                    format!("Expected delimiter string, found {other:?}"),
131                    ctx.current_position(),
132                )),
133            }
134        };
135
136        let parse_bool_value = |ctx: &mut ParseContext| -> Option<bool> {
137            match ctx.current_token().kind.clone() {
138                TokenKind::BooleanLiteral(b) => {
139                    ctx.next_token();
140                    Some(b)
141                }
142                TokenKind::Identifier(s) if s.eq_ignore_ascii_case("true") => {
143                    ctx.next_token();
144                    Some(true)
145                }
146                TokenKind::Identifier(s) if s.eq_ignore_ascii_case("false") => {
147                    ctx.next_token();
148                    Some(false)
149                }
150                _ => None,
151            }
152        };
153
154        // Option list: `WITH (OPT v, OPT v)` / `WITH OPT OPT ...` / `(OPT, ...)`.
155        let mut in_parens = false;
156        if (ctx.match_token(TokenKind::With) || !ctx.check_token(TokenKind::With))
157            && ctx.match_token(TokenKind::LParen)
158        {
159            in_parens = true;
160        }
161
162        loop {
163            if ctx.match_token(TokenKind::Comma) {
164                continue;
165            }
166            if in_parens && ctx.match_token(TokenKind::RParen) {
167                break;
168            }
169
170            if ctx.check_token(TokenKind::Header) || ctx.check_keyword("HEADER") {
171                consume_option_name(ctx, TokenKind::Header, "HEADER");
172                header = parse_bool_value(ctx).unwrap_or(true);
173                continue;
174            }
175            if ctx.check_token(TokenKind::Delimiter) || ctx.check_keyword("DELIMITER") {
176                consume_option_name(ctx, TokenKind::Delimiter, "DELIMITER");
177                delimiter = parse_delimiter_value(ctx)?;
178                continue;
179            }
180            if ctx.check_keyword("BATCH_SIZE") || ctx.check_keyword("BATCH") {
181                if ctx.check_keyword("BATCH_SIZE") {
182                    let _ = ctx.consume_keyword("BATCH_SIZE");
183                } else {
184                    let _ = ctx.consume_keyword("BATCH");
185                    let _ = ctx.consume_keyword("SIZE");
186                }
187                consume_optional_assign(ctx);
188                let batch = ctx.expect_integer_literal()?;
189                if batch < 0 {
190                    return Err(ParseError::new(
191                        crate::parser::core::error::ParseErrorKind::SyntaxError,
192                        "BATCH_SIZE must be positive".to_string(),
193                        ctx.current_position(),
194                    ));
195                }
196                batch_size = Some(batch as usize);
197                continue;
198            }
199            if ctx.check_token(TokenKind::Csv) || ctx.check_keyword("CSV") {
200                consume_option_name(ctx, TokenKind::Csv, "CSV");
201                continue;
202            }
203            // `NO HEADER` is the explicit spelling of `HEADER false`.
204            if ctx.check_token(TokenKind::No) || ctx.check_keyword("NO") {
205                ctx.next_token();
206                if ctx.check_token(TokenKind::Header) || ctx.check_keyword("HEADER") {
207                    consume_option_name(ctx, TokenKind::Header, "HEADER");
208                    header = false;
209                    continue;
210                }
211                break;
212            }
213
214            if in_parens {
215                if matches!(
216                    ctx.current_token().kind,
217                    TokenKind::Eof | TokenKind::Semicolon
218                ) {
219                    break;
220                }
221                // Unknown option inside parens: stop consuming rather than
222                // guessing; the statement-level parser reports the leftover.
223                break;
224            }
225            // A second `WITH` continues the option list (`WITH a WITH b`).
226            if ctx.match_token(TokenKind::With) {
227                if ctx.match_token(TokenKind::LParen) {
228                    in_parens = true;
229                }
230                continue;
231            }
232            break;
233        }
234
235        // Trailing `BY COLUMN` marker after the option list.
236        if !by_column && (ctx.check_token(TokenKind::By) || ctx.check_keyword("BY")) {
237            if !ctx.match_token(TokenKind::By) {
238                let _ = ctx.consume_keyword("BY");
239            }
240            if ctx.check_keyword("COLUMN") {
241                let _ = ctx.consume_keyword("COLUMN");
242                by_column = true;
243            } else {
244                return Err(ParseError::new(
245                    crate::parser::core::error::ParseErrorKind::SyntaxError,
246                    "Expected COLUMN after BY in COPY statement".to_string(),
247                    ctx.current_position(),
248                ));
249            }
250        }
251        if direction == CopyDirection::To && by_column {
252            return Err(ParseError::new(
253                crate::parser::core::error::ParseErrorKind::SyntaxError,
254                "BY COLUMN is only valid for COPY FROM".to_string(),
255                ctx.current_position(),
256            ));
257        }
258
259        let span = ctx.span_since(start_span.start);
260
261        Ok(Stmt::Copy(CopyStmt {
262            span,
263            target,
264            direction,
265            file_paths,
266            by_column,
267            header,
268            delimiter,
269            batch_size,
270        }))
271    }
272}
linkrs_query/parser/parsing/dml_parser/copy.rs

linkrs_query/parser/parsing/dml_parser/
copy.rs
