sql/parsers/
copy_parser.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
15use snafu::ResultExt;
16use sqlparser::keywords::Keyword;
17use sqlparser::tokenizer::Token;
18use sqlparser::tokenizer::Token::Word;
19
20use crate::error::{self, Result};
21use crate::parser::ParserContext;
22use crate::statements::OptionMap;
23use crate::statements::copy::{
24    CopyDatabase, CopyDatabaseArgument, CopyQueryTo, CopyQueryToArgument, CopyTable,
25    CopyTableArgument,
26};
27use crate::statements::statement::Statement;
28use crate::util::parse_option_string;
29
30// COPY tbl TO 'output.parquet';
31impl ParserContext<'_> {
32    pub(crate) fn parse_copy(&mut self) -> Result<Statement> {
33        let _ = self.parser.next_token();
34        let next = self.parser.peek_token();
35        if next.token == Token::LParen {
36            let copy_query = self.parse_copy_query_to()?;
37            // the COPY ... TO STDOUT is a special case for postgres wire protocol
38            // the logic is completely identical to query, but with an alternative data encoding on transport
39            //
40            // so at the query engine level, we simple parse the command as it's inner query
41            // we will deal with the encoding and its format options in server/src/postgres/handler.rs
42            if copy_query.arg.location == "STDOUT" {
43                Ok(*copy_query.query)
44            } else {
45                Ok(Statement::Copy(crate::statements::copy::Copy::CopyQueryTo(
46                    copy_query,
47                )))
48            }
49        } else if let Word(word) = next.token
50            && word.keyword == Keyword::DATABASE
51        {
52            let _ = self.parser.next_token();
53            let copy_database = self.parser_copy_database()?;
54            Ok(Statement::Copy(
55                crate::statements::copy::Copy::CopyDatabase(copy_database),
56            ))
57        } else {
58            let copy_table = self.parse_copy_table()?;
59            Ok(Statement::Copy(crate::statements::copy::Copy::CopyTable(
60                copy_table,
61            )))
62        }
63    }
64
65    fn parser_copy_database(&mut self) -> Result<CopyDatabase> {
66        let database_name = self
67            .parse_object_name()
68            .with_context(|_| error::UnexpectedSnafu {
69                expected: "a database name",
70                actual: self.peek_token_as_string(),
71            })?;
72
73        let req = if self.parser.parse_keyword(Keyword::TO) {
74            let (with, connection, location, limit) = self.parse_copy_parameters()?;
75            if limit.is_some() {
76                return error::InvalidSqlSnafu {
77                    msg: "limit is not supported",
78                }
79                .fail();
80            }
81
82            let argument = CopyDatabaseArgument {
83                database_name,
84                with,
85                connection,
86                location,
87            };
88            CopyDatabase::To(argument)
89        } else {
90            self.parser
91                .expect_keyword(Keyword::FROM)
92                .context(error::SyntaxSnafu)?;
93            let (with, connection, location, limit) = self.parse_copy_parameters()?;
94            if limit.is_some() {
95                return error::InvalidSqlSnafu {
96                    msg: "limit is not supported",
97                }
98                .fail();
99            }
100
101            let argument = CopyDatabaseArgument {
102                database_name,
103                with,
104                connection,
105                location,
106            };
107            CopyDatabase::From(argument)
108        };
109        Ok(req)
110    }
111
112    fn parse_copy_table(&mut self) -> Result<CopyTable> {
113        let raw_table_name = self
114            .parse_object_name()
115            .with_context(|_| error::UnexpectedSnafu {
116                expected: "a table name",
117                actual: self.peek_token_as_string(),
118            })?;
119        let table_name = Self::canonicalize_object_name(raw_table_name)?;
120
121        if self.parser.parse_keyword(Keyword::TO) {
122            let (with, connection, location, limit) = self.parse_copy_parameters()?;
123            Ok(CopyTable::To(CopyTableArgument {
124                table_name,
125                with,
126                connection,
127                location,
128                limit,
129            }))
130        } else {
131            self.parser
132                .expect_keyword(Keyword::FROM)
133                .context(error::SyntaxSnafu)?;
134            let (with, connection, location, limit) = self.parse_copy_parameters()?;
135            Ok(CopyTable::From(CopyTableArgument {
136                table_name,
137                with,
138                connection,
139                location,
140                limit,
141            }))
142        }
143    }
144
145    fn parse_copy_query_to(&mut self) -> Result<CopyQueryTo> {
146        self.parser
147            .expect_token(&Token::LParen)
148            .with_context(|_| error::UnexpectedSnafu {
149                expected: "'('",
150                actual: self.peek_token_as_string(),
151            })?;
152        let query = self.parse_query()?;
153        self.parser
154            .expect_token(&Token::RParen)
155            .with_context(|_| error::UnexpectedSnafu {
156                expected: "')'",
157                actual: self.peek_token_as_string(),
158            })?;
159        self.parser
160            .expect_keyword(Keyword::TO)
161            .context(error::SyntaxSnafu)?;
162
163        if self.parser.parse_keyword(Keyword::STDOUT) {
164            // early return without parsing options
165            // we will deal with copy to stdout on postgres protocol layer
166            // consume [WITH] (...) options if present (they will be ignored)
167            // we support both "WITH (FORMAT binary)" and "(FORMAT binary)"
168            // for PostgreSQL compatibility
169
170            // Check for optional WITH keyword or direct LParen (PostgreSQL syntax)
171            // Both "WITH (...)" and "(...)" are valid after STDOUT
172            let _ = self.parser.parse_keyword(Keyword::WITH);
173            if self.parser.peek_token().token == Token::LParen {
174                let _ = self.parser.next_token();
175                // consume all tokens until we find matching RParen
176                let mut depth = 1;
177                while depth > 0 {
178                    match self.parser.next_token().token {
179                        Token::LParen => depth += 1,
180                        Token::RParen => depth -= 1,
181                        Token::EOF => {
182                            return error::UnexpectedTokenSnafu {
183                                expected: ")",
184                                actual: "EOF",
185                            }
186                            .fail();
187                        }
188                        _ => {}
189                    }
190                }
191            }
192
193            Ok(CopyQueryTo {
194                query: Box::new(query),
195                arg: CopyQueryToArgument {
196                    with: OptionMap::default(),
197                    connection: OptionMap::default(),
198                    location: "STDOUT".to_string(),
199                },
200            })
201        } else {
202            let (with, connection, location, limit) = self.parse_copy_parameters()?;
203            if limit.is_some() {
204                return error::InvalidSqlSnafu {
205                    msg: "limit is not supported",
206                }
207                .fail();
208            }
209            Ok(CopyQueryTo {
210                query: Box::new(query),
211                arg: CopyQueryToArgument {
212                    with,
213                    connection,
214                    location,
215                },
216            })
217        }
218    }
219
220    fn parse_copy_parameters(&mut self) -> Result<(OptionMap, OptionMap, String, Option<u64>)> {
221        let location =
222            self.parser
223                .parse_literal_string()
224                .with_context(|_| error::UnexpectedSnafu {
225                    expected: "a file name",
226                    actual: self.peek_token_as_string(),
227                })?;
228
229        let options = self
230            .parser
231            .parse_options(Keyword::WITH)
232            .context(error::SyntaxSnafu)?;
233
234        let with = options
235            .into_iter()
236            .map(parse_option_string)
237            .collect::<Result<Vec<_>>>()?;
238        let with = OptionMap::new(with);
239
240        let connection_options = self
241            .parser
242            .parse_options(Keyword::CONNECTION)
243            .context(error::SyntaxSnafu)?;
244
245        let connection = connection_options
246            .into_iter()
247            .map(parse_option_string)
248            .collect::<Result<Vec<_>>>()?;
249        let connection = OptionMap::new(connection);
250
251        let limit = if self.parser.parse_keyword(Keyword::LIMIT) {
252            Some(
253                self.parser
254                    .parse_literal_uint()
255                    .with_context(|_| error::UnexpectedSnafu {
256                        expected: "the number of maximum rows",
257                        actual: self.peek_token_as_string(),
258                    })?,
259            )
260        } else {
261            None
262        };
263
264        Ok((with, connection, location, limit))
265    }
266}
267
268#[cfg(test)]
269mod tests {
270    use std::assert_matches::assert_matches;
271    use std::collections::HashMap;
272
273    use sqlparser::ast::{Ident, ObjectName};
274
275    use super::*;
276    use crate::dialect::GreptimeDbDialect;
277    use crate::parser::ParseOptions;
278    use crate::statements::statement::Statement::Copy;
279
280    #[test]
281    fn test_parse_copy_table() {
282        let sql0 = "COPY catalog0.schema0.tbl TO 'tbl_file.parquet'";
283        let sql1 = "COPY catalog0.schema0.tbl TO 'tbl_file.parquet' WITH (FORMAT = 'parquet')";
284        let result0 = ParserContext::create_with_dialect(
285            sql0,
286            &GreptimeDbDialect {},
287            ParseOptions::default(),
288        )
289        .unwrap();
290        let result1 = ParserContext::create_with_dialect(
291            sql1,
292            &GreptimeDbDialect {},
293            ParseOptions::default(),
294        )
295        .unwrap();
296
297        for mut result in [result0, result1] {
298            assert_eq!(1, result.len());
299
300            let statement = result.remove(0);
301            assert_matches!(statement, Statement::Copy { .. });
302            match statement {
303                Copy(copy) => {
304                    let crate::statements::copy::Copy::CopyTable(CopyTable::To(copy_table)) = copy
305                    else {
306                        unreachable!()
307                    };
308                    let table = copy_table.table_name.to_string();
309                    assert_eq!("catalog0.schema0.tbl", table);
310
311                    let file_name = &copy_table.location;
312                    assert_eq!("tbl_file.parquet", file_name);
313
314                    let format = copy_table.format().unwrap();
315                    assert_eq!("parquet", format.to_lowercase());
316                }
317                _ => unreachable!(),
318            }
319        }
320    }
321
322    #[test]
323    fn test_parse_copy_table_from_basic() {
324        let results = [
325            "COPY catalog0.schema0.tbl FROM 'tbl_file.parquet'",
326            "COPY catalog0.schema0.tbl FROM 'tbl_file.parquet' WITH (FORMAT = 'parquet')",
327        ]
328        .iter()
329        .map(|sql| {
330            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
331                .unwrap()
332        })
333        .collect::<Vec<_>>();
334
335        for mut result in results {
336            assert_eq!(1, result.len());
337
338            let statement = result.remove(0);
339            assert_matches!(statement, Statement::Copy { .. });
340            match statement {
341                Statement::Copy(crate::statements::copy::Copy::CopyTable(CopyTable::From(
342                    copy_table,
343                ))) => {
344                    let table = copy_table.table_name.to_string();
345                    assert_eq!("catalog0.schema0.tbl", table);
346
347                    let file_name = &copy_table.location;
348                    assert_eq!("tbl_file.parquet", file_name);
349
350                    let format = copy_table.format().unwrap();
351                    assert_eq!("parquet", format.to_lowercase());
352                }
353                _ => unreachable!(),
354            }
355        }
356    }
357
358    #[test]
359    fn test_parse_copy_table_from() {
360        struct Test<'a> {
361            sql: &'a str,
362            expected_pattern: Option<String>,
363            expected_connection: HashMap<&'a str, &'a str>,
364        }
365
366        let tests = [
367            Test {
368                sql: "COPY catalog0.schema0.tbl FROM 'tbl_file.parquet' WITH (PATTERN = 'demo.*')",
369                expected_pattern: Some("demo.*".into()),
370                expected_connection: HashMap::new(),
371            },
372            Test {
373                sql: "COPY catalog0.schema0.tbl FROM 'tbl_file.parquet' WITH (PATTERN = 'demo.*') CONNECTION (FOO='Bar', ONE='two')",
374                expected_pattern: Some("demo.*".into()),
375                expected_connection: HashMap::from([("foo", "Bar"), ("one", "two")]),
376            },
377        ];
378
379        for test in tests {
380            let mut result = ParserContext::create_with_dialect(
381                test.sql,
382                &GreptimeDbDialect {},
383                ParseOptions::default(),
384            )
385            .unwrap();
386            assert_eq!(1, result.len());
387
388            let statement = result.remove(0);
389            assert_matches!(statement, Statement::Copy { .. });
390            match statement {
391                Statement::Copy(crate::statements::copy::Copy::CopyTable(CopyTable::From(
392                    copy_table,
393                ))) => {
394                    if let Some(expected_pattern) = test.expected_pattern {
395                        assert_eq!(copy_table.pattern().unwrap(), expected_pattern);
396                    }
397                    assert_eq!(copy_table.connection.to_str_map(), test.expected_connection);
398                }
399                _ => unreachable!(),
400            }
401        }
402    }
403
404    #[test]
405    fn test_parse_copy_table_to() {
406        struct Test<'a> {
407            sql: &'a str,
408            expected_connection: HashMap<&'a str, &'a str>,
409        }
410
411        let tests = [
412            Test {
413                sql: "COPY catalog0.schema0.tbl TO 'tbl_file.parquet' ",
414                expected_connection: HashMap::new(),
415            },
416            Test {
417                sql: "COPY catalog0.schema0.tbl TO 'tbl_file.parquet' CONNECTION (FOO='Bar', ONE='two')",
418                expected_connection: HashMap::from([("foo", "Bar"), ("one", "two")]),
419            },
420            Test {
421                sql: "COPY catalog0.schema0.tbl TO 'tbl_file.parquet' WITH (FORMAT = 'parquet') CONNECTION (FOO='Bar', ONE='two')",
422                expected_connection: HashMap::from([("foo", "Bar"), ("one", "two")]),
423            },
424        ];
425
426        for test in tests {
427            let mut result = ParserContext::create_with_dialect(
428                test.sql,
429                &GreptimeDbDialect {},
430                ParseOptions::default(),
431            )
432            .unwrap();
433            assert_eq!(1, result.len());
434
435            let statement = result.remove(0);
436            assert_matches!(statement, Statement::Copy { .. });
437            match statement {
438                Statement::Copy(crate::statements::copy::Copy::CopyTable(CopyTable::To(
439                    copy_table,
440                ))) => {
441                    assert_eq!(copy_table.connection.to_str_map(), test.expected_connection);
442                }
443                _ => unreachable!(),
444            }
445        }
446    }
447
448    #[test]
449    fn test_copy_database_to() {
450        let sql = "COPY DATABASE catalog0.schema0 TO 'tbl_file.parquet' WITH (FORMAT = 'parquet') CONNECTION (FOO='Bar', ONE='two')";
451        let stmt =
452            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
453                .unwrap()
454                .pop()
455                .unwrap();
456
457        let Copy(crate::statements::copy::Copy::CopyDatabase(stmt)) = stmt else {
458            unreachable!()
459        };
460
461        let CopyDatabase::To(stmt) = stmt else {
462            unreachable!()
463        };
464
465        assert_eq!(
466            ObjectName::from(vec![Ident::new("catalog0"), Ident::new("schema0")]),
467            stmt.database_name
468        );
469        assert_eq!(
470            [("format", "parquet")]
471                .into_iter()
472                .collect::<HashMap<_, _>>(),
473            stmt.with.to_str_map()
474        );
475
476        assert_eq!(
477            [("foo", "Bar"), ("one", "two")]
478                .into_iter()
479                .collect::<HashMap<_, _>>(),
480            stmt.connection.to_str_map()
481        );
482    }
483
484    #[test]
485    fn test_copy_database_from() {
486        let sql = "COPY DATABASE catalog0.schema0 FROM '/a/b/c/' WITH (FORMAT = 'parquet') CONNECTION (FOO='Bar', ONE='two')";
487        let stmt =
488            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
489                .unwrap()
490                .pop()
491                .unwrap();
492
493        let Copy(crate::statements::copy::Copy::CopyDatabase(stmt)) = stmt else {
494            unreachable!()
495        };
496
497        let CopyDatabase::From(stmt) = stmt else {
498            unreachable!()
499        };
500
501        assert_eq!(
502            ObjectName::from(vec![Ident::new("catalog0"), Ident::new("schema0")]),
503            stmt.database_name
504        );
505        assert_eq!(
506            [("format", "parquet")]
507                .into_iter()
508                .collect::<HashMap<_, _>>(),
509            stmt.with.to_str_map()
510        );
511
512        assert_eq!(
513            [("foo", "Bar"), ("one", "two")]
514                .into_iter()
515                .collect::<HashMap<_, _>>(),
516            stmt.connection.to_str_map()
517        );
518    }
519
520    #[test]
521    fn test_copy_query_to() {
522        let sql = "COPY (SELECT * FROM tbl WHERE ts > 10) TO 'tbl_file.parquet' WITH (FORMAT = 'parquet') CONNECTION (FOO='Bar', ONE='two')";
523        let stmt =
524            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
525                .unwrap()
526                .pop()
527                .unwrap();
528
529        let Copy(crate::statements::copy::Copy::CopyQueryTo(stmt)) = stmt else {
530            unreachable!()
531        };
532
533        let query = ParserContext::create_with_dialect(
534            "SELECT * FROM tbl WHERE ts > 10",
535            &GreptimeDbDialect {},
536            ParseOptions::default(),
537        )
538        .unwrap()
539        .remove(0);
540
541        assert_eq!(&query, stmt.query.as_ref());
542        assert_eq!(
543            [("format", "parquet")]
544                .into_iter()
545                .collect::<HashMap<_, _>>(),
546            stmt.arg.with.to_str_map()
547        );
548
549        assert_eq!(
550            [("foo", "Bar"), ("one", "two")]
551                .into_iter()
552                .collect::<HashMap<_, _>>(),
553            stmt.arg.connection.to_str_map()
554        );
555    }
556
557    #[test]
558    fn test_invalid_copy_query_to() {
559        {
560            let sql = "COPY SELECT * FROM tbl WHERE ts > 10 TO 'tbl_file.parquet' WITH (FORMAT = 'parquet') CONNECTION (FOO='Bar', ONE='two')";
561
562            assert!(
563                ParserContext::create_with_dialect(
564                    sql,
565                    &GreptimeDbDialect {},
566                    ParseOptions::default()
567                )
568                .is_err()
569            )
570        }
571        {
572            let sql = "COPY SELECT * FROM tbl WHERE ts > 10) TO 'tbl_file.parquet' WITH (FORMAT = 'parquet') CONNECTION (FOO='Bar', ONE='two')";
573
574            assert!(
575                ParserContext::create_with_dialect(
576                    sql,
577                    &GreptimeDbDialect {},
578                    ParseOptions::default()
579                )
580                .is_err()
581            )
582        }
583        {
584            let sql = "COPY (SELECT * FROM tbl WHERE ts > 10 TO 'tbl_file.parquet' WITH (FORMAT = 'parquet') CONNECTION (FOO='Bar', ONE='two')";
585
586            assert!(
587                ParserContext::create_with_dialect(
588                    sql,
589                    &GreptimeDbDialect {},
590                    ParseOptions::default()
591                )
592                .is_err()
593            )
594        }
595    }
596
597    #[test]
598    fn test_copy_query_to_stdout() {
599        let sql = "COPY (SELECT * FROM tbl WHERE ts > 10) TO STDOUT WITH (FORMAT = 'csv')";
600        let stmt =
601            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
602                .unwrap()
603                .pop()
604                .unwrap();
605
606        let expected_query = ParserContext::create_with_dialect(
607            "SELECT * FROM tbl WHERE ts > 10",
608            &GreptimeDbDialect {},
609            ParseOptions::default(),
610        )
611        .unwrap()
612        .remove(0);
613
614        assert_eq!(&expected_query, &stmt);
615    }
616
617    #[test]
618    fn test_copy_query_to_stdout_without_format() {
619        let sql = "COPY (SELECT generate_series(1, 2), generate_series(2, 3)) TO STDOUT";
620        let stmt =
621            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
622                .unwrap()
623                .pop()
624                .unwrap();
625
626        let query_str = "SELECT generate_series(1, 2), generate_series(2, 3)";
627        let expected_query = ParserContext::create_with_dialect(
628            query_str,
629            &GreptimeDbDialect {},
630            ParseOptions::default(),
631        )
632        .unwrap()
633        .remove(0);
634
635        assert_eq!(&expected_query, &stmt);
636    }
637
638    #[test]
639    fn test_copy_query_to_stdout_with_binary_format() {
640        let sql = "COPY (SELECT * FROM test_table) TO STDOUT WITH (FORMAT binary)";
641        let result =
642            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
643
644        if let Err(e) = &result {
645            panic!(
646                "COPY TO STDOUT WITH (FORMAT binary) should parse without error, got: {:?}",
647                e
648            );
649        }
650
651        let stmt = result.unwrap().pop().unwrap();
652
653        let expected_query = ParserContext::create_with_dialect(
654            "SELECT * FROM test_table",
655            &GreptimeDbDialect {},
656            ParseOptions::default(),
657        )
658        .unwrap()
659        .remove(0);
660
661        assert_eq!(&expected_query, &stmt);
662    }
663
664    #[test]
665    fn test_copy_query_to_stdout_with_csv_format() {
666        let sql = "COPY (SELECT * FROM test_table) TO STDOUT WITH (FORMAT csv)";
667        let result =
668            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
669
670        if let Err(e) = &result {
671            panic!(
672                "COPY TO STDOUT WITH (FORMAT csv) should parse without error, got: {:?}",
673                e
674            );
675        }
676
677        let stmt = result.unwrap().pop().unwrap();
678
679        let expected_query = ParserContext::create_with_dialect(
680            "SELECT * FROM test_table",
681            &GreptimeDbDialect {},
682            ParseOptions::default(),
683        )
684        .unwrap()
685        .remove(0);
686
687        assert_eq!(&expected_query, &stmt);
688    }
689
690    #[test]
691    fn test_copy_query_to_stdout_with_equals_format() {
692        let sql = "COPY (SELECT * FROM test_table) TO STDOUT WITH (FORMAT = 'binary')";
693        let result =
694            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
695
696        if let Err(e) = &result {
697            panic!(
698                "COPY TO STDOUT WITH (FORMAT = 'binary') should parse without error, got: {:?}",
699                e
700            );
701        }
702
703        let stmt = result.unwrap().pop().unwrap();
704
705        let expected_query = ParserContext::create_with_dialect(
706            "SELECT * FROM test_table",
707            &GreptimeDbDialect {},
708            ParseOptions::default(),
709        )
710        .unwrap()
711        .remove(0);
712
713        assert_eq!(&expected_query, &stmt);
714    }
715
716    #[test]
717    fn test_copy_query_to_stdout_with_multiple_options() {
718        let sql =
719            "COPY (SELECT * FROM test_table) TO STDOUT WITH (FORMAT csv, DELIMITER ',', HEADER)";
720        let result =
721            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
722
723        if let Err(e) = &result {
724            panic!(
725                "COPY TO STDOUT WITH multiple options should parse without error, got: {:?}",
726                e
727            );
728        }
729
730        let stmt = result.unwrap().pop().unwrap();
731
732        let expected_query = ParserContext::create_with_dialect(
733            "SELECT * FROM test_table",
734            &GreptimeDbDialect {},
735            ParseOptions::default(),
736        )
737        .unwrap()
738        .remove(0);
739
740        assert_eq!(&expected_query, &stmt);
741    }
742
743    #[test]
744    fn test_copy_query_to_stdout_without_with_keyword() {
745        let sql = "COPY (SELECT * FROM test_table) TO STDOUT (FORMAT binary)";
746        let result =
747            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
748
749        if let Err(e) = &result {
750            panic!(
751                "COPY TO STDOUT (FORMAT binary) without WITH keyword should parse without error, got: {:?}",
752                e
753            );
754        }
755
756        let stmt = result.unwrap().pop().unwrap();
757
758        let expected_query = ParserContext::create_with_dialect(
759            "SELECT * FROM test_table",
760            &GreptimeDbDialect {},
761            ParseOptions::default(),
762        )
763        .unwrap()
764        .remove(0);
765
766        assert_eq!(&expected_query, &stmt);
767    }
768
769    #[test]
770    fn test_copy_query_to_stdout_without_with_csv_format() {
771        let sql = "COPY (SELECT * FROM test_table) TO STDOUT (FORMAT csv)";
772        let result =
773            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
774
775        if let Err(e) = &result {
776            panic!(
777                "COPY TO STDOUT (FORMAT csv) without WITH keyword should parse without error, got: {:?}",
778                e
779            );
780        }
781
782        let stmt = result.unwrap().pop().unwrap();
783
784        let expected_query = ParserContext::create_with_dialect(
785            "SELECT * FROM test_table",
786            &GreptimeDbDialect {},
787            ParseOptions::default(),
788        )
789        .unwrap()
790        .remove(0);
791
792        assert_eq!(&expected_query, &stmt);
793    }
794
795    #[test]
796    fn test_copy_query_to_stdout_without_with_multiple_options() {
797        let sql = "COPY (SELECT * FROM test_table) TO STDOUT (FORMAT csv, DELIMITER ',', HEADER)";
798        let result =
799            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
800
801        if let Err(e) = &result {
802            panic!(
803                "COPY TO STDOUT (FORMAT csv, ...) without WITH keyword should parse without error, got: {:?}",
804                e
805            );
806        }
807
808        let stmt = result.unwrap().pop().unwrap();
809
810        let expected_query = ParserContext::create_with_dialect(
811            "SELECT * FROM test_table",
812            &GreptimeDbDialect {},
813            ParseOptions::default(),
814        )
815        .unwrap()
816        .remove(0);
817
818        assert_eq!(&expected_query, &stmt);
819    }
820
821    #[test]
822    fn test_invalid_copy_query() {
823        let sql = "COPY (SELECT * FROM test_table) TO STDOUT (FORMAT csv";
824        let result =
825            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
826
827        assert!(result.is_err());
828
829        let sql = "COPY (SELECT * FROM test_table) TO STDOUT (FORMAT csv))";
830        let result =
831            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
832
833        assert!(result.is_err());
834    }
835}
sql/parsers/copy_parser.rs

sql/parsers/
copy_parser.rs
