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;
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_from_csv_options() {
406        let sql =
407            "COPY my_table FROM '/tmp/test.csv' WITH (FORMAT = 'CSV', SKIP_BAD_RECORDS = 'false')";
408        let mut result =
409            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
410                .unwrap();
411        assert_eq!(1, result.len());
412
413        let statement = result.remove(0);
414        assert_matches!(statement, Statement::Copy { .. });
415        match statement {
416            Statement::Copy(crate::statements::copy::Copy::CopyTable(CopyTable::From(
417                copy_table,
418            ))) => {
419                assert_eq!(copy_table.with.get("format"), Some("CSV"));
420                assert_eq!(copy_table.with.get("skip_bad_records"), Some("false"));
421            }
422            _ => unreachable!(),
423        }
424    }
425
426    #[test]
427    fn test_parse_copy_table_to() {
428        struct Test<'a> {
429            sql: &'a str,
430            expected_connection: HashMap<&'a str, &'a str>,
431        }
432
433        let tests = [
434            Test {
435                sql: "COPY catalog0.schema0.tbl TO 'tbl_file.parquet' ",
436                expected_connection: HashMap::new(),
437            },
438            Test {
439                sql: "COPY catalog0.schema0.tbl TO 'tbl_file.parquet' CONNECTION (FOO='Bar', ONE='two')",
440                expected_connection: HashMap::from([("foo", "Bar"), ("one", "two")]),
441            },
442            Test {
443                sql: "COPY catalog0.schema0.tbl TO 'tbl_file.parquet' WITH (FORMAT = 'parquet') CONNECTION (FOO='Bar', ONE='two')",
444                expected_connection: HashMap::from([("foo", "Bar"), ("one", "two")]),
445            },
446        ];
447
448        for test in tests {
449            let mut result = ParserContext::create_with_dialect(
450                test.sql,
451                &GreptimeDbDialect {},
452                ParseOptions::default(),
453            )
454            .unwrap();
455            assert_eq!(1, result.len());
456
457            let statement = result.remove(0);
458            assert_matches!(statement, Statement::Copy { .. });
459            match statement {
460                Statement::Copy(crate::statements::copy::Copy::CopyTable(CopyTable::To(
461                    copy_table,
462                ))) => {
463                    assert_eq!(copy_table.connection.to_str_map(), test.expected_connection);
464                }
465                _ => unreachable!(),
466            }
467        }
468    }
469
470    #[test]
471    fn test_copy_database_to() {
472        let sql = "COPY DATABASE catalog0.schema0 TO 'tbl_file.parquet' WITH (FORMAT = 'parquet') CONNECTION (FOO='Bar', ONE='two')";
473        let stmt =
474            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
475                .unwrap()
476                .pop()
477                .unwrap();
478
479        let Copy(crate::statements::copy::Copy::CopyDatabase(stmt)) = stmt else {
480            unreachable!()
481        };
482
483        let CopyDatabase::To(stmt) = stmt else {
484            unreachable!()
485        };
486
487        assert_eq!(
488            ObjectName::from(vec![Ident::new("catalog0"), Ident::new("schema0")]),
489            stmt.database_name
490        );
491        assert_eq!(
492            [("format", "parquet")]
493                .into_iter()
494                .collect::<HashMap<_, _>>(),
495            stmt.with.to_str_map()
496        );
497
498        assert_eq!(
499            [("foo", "Bar"), ("one", "two")]
500                .into_iter()
501                .collect::<HashMap<_, _>>(),
502            stmt.connection.to_str_map()
503        );
504    }
505
506    #[test]
507    fn test_copy_database_from() {
508        let sql = "COPY DATABASE catalog0.schema0 FROM '/a/b/c/' WITH (FORMAT = 'parquet') CONNECTION (FOO='Bar', ONE='two')";
509        let stmt =
510            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
511                .unwrap()
512                .pop()
513                .unwrap();
514
515        let Copy(crate::statements::copy::Copy::CopyDatabase(stmt)) = stmt else {
516            unreachable!()
517        };
518
519        let CopyDatabase::From(stmt) = stmt else {
520            unreachable!()
521        };
522
523        assert_eq!(
524            ObjectName::from(vec![Ident::new("catalog0"), Ident::new("schema0")]),
525            stmt.database_name
526        );
527        assert_eq!(
528            [("format", "parquet")]
529                .into_iter()
530                .collect::<HashMap<_, _>>(),
531            stmt.with.to_str_map()
532        );
533
534        assert_eq!(
535            [("foo", "Bar"), ("one", "two")]
536                .into_iter()
537                .collect::<HashMap<_, _>>(),
538            stmt.connection.to_str_map()
539        );
540    }
541
542    #[test]
543    fn test_copy_query_to() {
544        let sql = "COPY (SELECT * FROM tbl WHERE ts > 10) TO 'tbl_file.parquet' WITH (FORMAT = 'parquet') CONNECTION (FOO='Bar', ONE='two')";
545        let stmt =
546            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
547                .unwrap()
548                .pop()
549                .unwrap();
550
551        let Copy(crate::statements::copy::Copy::CopyQueryTo(stmt)) = stmt else {
552            unreachable!()
553        };
554
555        let query = ParserContext::create_with_dialect(
556            "SELECT * FROM tbl WHERE ts > 10",
557            &GreptimeDbDialect {},
558            ParseOptions::default(),
559        )
560        .unwrap()
561        .remove(0);
562
563        assert_eq!(&query, stmt.query.as_ref());
564        assert_eq!(
565            [("format", "parquet")]
566                .into_iter()
567                .collect::<HashMap<_, _>>(),
568            stmt.arg.with.to_str_map()
569        );
570
571        assert_eq!(
572            [("foo", "Bar"), ("one", "two")]
573                .into_iter()
574                .collect::<HashMap<_, _>>(),
575            stmt.arg.connection.to_str_map()
576        );
577    }
578
579    #[test]
580    fn test_invalid_copy_query_to() {
581        {
582            let sql = "COPY SELECT * FROM tbl WHERE ts > 10 TO 'tbl_file.parquet' WITH (FORMAT = 'parquet') CONNECTION (FOO='Bar', ONE='two')";
583
584            assert!(
585                ParserContext::create_with_dialect(
586                    sql,
587                    &GreptimeDbDialect {},
588                    ParseOptions::default()
589                )
590                .is_err()
591            )
592        }
593        {
594            let sql = "COPY SELECT * FROM tbl WHERE ts > 10) TO 'tbl_file.parquet' WITH (FORMAT = 'parquet') CONNECTION (FOO='Bar', ONE='two')";
595
596            assert!(
597                ParserContext::create_with_dialect(
598                    sql,
599                    &GreptimeDbDialect {},
600                    ParseOptions::default()
601                )
602                .is_err()
603            )
604        }
605        {
606            let sql = "COPY (SELECT * FROM tbl WHERE ts > 10 TO 'tbl_file.parquet' WITH (FORMAT = 'parquet') CONNECTION (FOO='Bar', ONE='two')";
607
608            assert!(
609                ParserContext::create_with_dialect(
610                    sql,
611                    &GreptimeDbDialect {},
612                    ParseOptions::default()
613                )
614                .is_err()
615            )
616        }
617    }
618
619    #[test]
620    fn test_copy_query_to_stdout() {
621        let sql = "COPY (SELECT * FROM tbl WHERE ts > 10) TO STDOUT WITH (FORMAT = 'csv')";
622        let stmt =
623            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
624                .unwrap()
625                .pop()
626                .unwrap();
627
628        let expected_query = ParserContext::create_with_dialect(
629            "SELECT * FROM tbl WHERE ts > 10",
630            &GreptimeDbDialect {},
631            ParseOptions::default(),
632        )
633        .unwrap()
634        .remove(0);
635
636        assert_eq!(&expected_query, &stmt);
637    }
638
639    #[test]
640    fn test_copy_query_to_stdout_without_format() {
641        let sql = "COPY (SELECT generate_series(1, 2), generate_series(2, 3)) TO STDOUT";
642        let stmt =
643            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
644                .unwrap()
645                .pop()
646                .unwrap();
647
648        let query_str = "SELECT generate_series(1, 2), generate_series(2, 3)";
649        let expected_query = ParserContext::create_with_dialect(
650            query_str,
651            &GreptimeDbDialect {},
652            ParseOptions::default(),
653        )
654        .unwrap()
655        .remove(0);
656
657        assert_eq!(&expected_query, &stmt);
658    }
659
660    #[test]
661    fn test_copy_query_to_stdout_with_binary_format() {
662        let sql = "COPY (SELECT * FROM test_table) TO STDOUT WITH (FORMAT binary)";
663        let result =
664            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
665
666        if let Err(e) = &result {
667            panic!(
668                "COPY TO STDOUT WITH (FORMAT binary) should parse without error, got: {:?}",
669                e
670            );
671        }
672
673        let stmt = result.unwrap().pop().unwrap();
674
675        let expected_query = ParserContext::create_with_dialect(
676            "SELECT * FROM test_table",
677            &GreptimeDbDialect {},
678            ParseOptions::default(),
679        )
680        .unwrap()
681        .remove(0);
682
683        assert_eq!(&expected_query, &stmt);
684    }
685
686    #[test]
687    fn test_copy_query_to_stdout_with_csv_format() {
688        let sql = "COPY (SELECT * FROM test_table) TO STDOUT WITH (FORMAT csv)";
689        let result =
690            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
691
692        if let Err(e) = &result {
693            panic!(
694                "COPY TO STDOUT WITH (FORMAT csv) should parse without error, got: {:?}",
695                e
696            );
697        }
698
699        let stmt = result.unwrap().pop().unwrap();
700
701        let expected_query = ParserContext::create_with_dialect(
702            "SELECT * FROM test_table",
703            &GreptimeDbDialect {},
704            ParseOptions::default(),
705        )
706        .unwrap()
707        .remove(0);
708
709        assert_eq!(&expected_query, &stmt);
710    }
711
712    #[test]
713    fn test_copy_query_to_stdout_with_equals_format() {
714        let sql = "COPY (SELECT * FROM test_table) TO STDOUT WITH (FORMAT = 'binary')";
715        let result =
716            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
717
718        if let Err(e) = &result {
719            panic!(
720                "COPY TO STDOUT WITH (FORMAT = 'binary') should parse without error, got: {:?}",
721                e
722            );
723        }
724
725        let stmt = result.unwrap().pop().unwrap();
726
727        let expected_query = ParserContext::create_with_dialect(
728            "SELECT * FROM test_table",
729            &GreptimeDbDialect {},
730            ParseOptions::default(),
731        )
732        .unwrap()
733        .remove(0);
734
735        assert_eq!(&expected_query, &stmt);
736    }
737
738    #[test]
739    fn test_copy_query_to_stdout_with_multiple_options() {
740        let sql =
741            "COPY (SELECT * FROM test_table) TO STDOUT WITH (FORMAT csv, DELIMITER ',', HEADER)";
742        let result =
743            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
744
745        if let Err(e) = &result {
746            panic!(
747                "COPY TO STDOUT WITH multiple options should parse without error, got: {:?}",
748                e
749            );
750        }
751
752        let stmt = result.unwrap().pop().unwrap();
753
754        let expected_query = ParserContext::create_with_dialect(
755            "SELECT * FROM test_table",
756            &GreptimeDbDialect {},
757            ParseOptions::default(),
758        )
759        .unwrap()
760        .remove(0);
761
762        assert_eq!(&expected_query, &stmt);
763    }
764
765    #[test]
766    fn test_copy_query_to_stdout_without_with_keyword() {
767        let sql = "COPY (SELECT * FROM test_table) TO STDOUT (FORMAT binary)";
768        let result =
769            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
770
771        if let Err(e) = &result {
772            panic!(
773                "COPY TO STDOUT (FORMAT binary) without WITH keyword should parse without error, got: {:?}",
774                e
775            );
776        }
777
778        let stmt = result.unwrap().pop().unwrap();
779
780        let expected_query = ParserContext::create_with_dialect(
781            "SELECT * FROM test_table",
782            &GreptimeDbDialect {},
783            ParseOptions::default(),
784        )
785        .unwrap()
786        .remove(0);
787
788        assert_eq!(&expected_query, &stmt);
789    }
790
791    #[test]
792    fn test_copy_query_to_stdout_without_with_csv_format() {
793        let sql = "COPY (SELECT * FROM test_table) TO STDOUT (FORMAT csv)";
794        let result =
795            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
796
797        if let Err(e) = &result {
798            panic!(
799                "COPY TO STDOUT (FORMAT csv) without WITH keyword should parse without error, got: {:?}",
800                e
801            );
802        }
803
804        let stmt = result.unwrap().pop().unwrap();
805
806        let expected_query = ParserContext::create_with_dialect(
807            "SELECT * FROM test_table",
808            &GreptimeDbDialect {},
809            ParseOptions::default(),
810        )
811        .unwrap()
812        .remove(0);
813
814        assert_eq!(&expected_query, &stmt);
815    }
816
817    #[test]
818    fn test_copy_query_to_stdout_without_with_multiple_options() {
819        let sql = "COPY (SELECT * FROM test_table) TO STDOUT (FORMAT csv, DELIMITER ',', HEADER)";
820        let result =
821            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
822
823        if let Err(e) = &result {
824            panic!(
825                "COPY TO STDOUT (FORMAT csv, ...) without WITH keyword should parse without error, got: {:?}",
826                e
827            );
828        }
829
830        let stmt = result.unwrap().pop().unwrap();
831
832        let expected_query = ParserContext::create_with_dialect(
833            "SELECT * FROM test_table",
834            &GreptimeDbDialect {},
835            ParseOptions::default(),
836        )
837        .unwrap()
838        .remove(0);
839
840        assert_eq!(&expected_query, &stmt);
841    }
842
843    #[test]
844    fn test_invalid_copy_query() {
845        let sql = "COPY (SELECT * FROM test_table) TO STDOUT (FORMAT csv";
846        let result =
847            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
848
849        assert!(result.is_err());
850
851        let sql = "COPY (SELECT * FROM test_table) TO STDOUT (FORMAT csv))";
852        let result =
853            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
854
855        assert!(result.is_err());
856    }
857}
sql/parsers/copy_parser.rs

sql/parsers/
copy_parser.rs
