feat: impl COPY a query resultset to external file (#5250)

* feat: impl COPY a query resultset to external file

* chore: add more tests for parse `copy_table_to`

* chore: add more tests for parse `copy_table_to`

src/sql/src/parsers/copy_parser.rs

@@ -16,11 +16,15 @@ use std::collections::HashMap;
 
 use snafu::ResultExt;
 use sqlparser::keywords::Keyword;
+use sqlparser::tokenizer::Token;
 use sqlparser::tokenizer::Token::Word;
 
 use crate::error::{self, Result};
 use crate::parser::ParserContext;
-use crate::statements::copy::{CopyDatabase, CopyDatabaseArgument, CopyTable, CopyTableArgument};
+use crate::statements::copy::{
+    CopyDatabase, CopyDatabaseArgument, CopyQueryTo, CopyQueryToArgument, CopyTable,
+    CopyTableArgument,
+};
 use crate::statements::statement::Statement;
 use crate::util::parse_option_string;
 
@@ -32,7 +36,11 @@ impl ParserContext<'_> {
     pub(crate) fn parse_copy(&mut self) -> Result<Statement> {
         let _ = self.parser.next_token();
         let next = self.parser.peek_token();
-        let copy = if let Word(word) = next.token
+        let copy = if next.token == Token::LParen {
+            let copy_query = self.parse_copy_query_to()?;
+
+            crate::statements::copy::Copy::CopyQueryTo(copy_query)
+        } else if let Word(word) = next.token
             && word.keyword == Keyword::DATABASE
         {
             let _ = self.parser.next_token();
@@ -126,6 +134,40 @@ impl ParserContext<'_> {
         }
     }
 
+    fn parse_copy_query_to(&mut self) -> Result<CopyQueryTo> {
+        self.parser
+            .expect_token(&Token::LParen)
+            .with_context(|_| error::UnexpectedSnafu {
+                expected: "'('",
+                actual: self.peek_token_as_string(),
+            })?;
+        let query = self.parse_query()?;
+        self.parser
+            .expect_token(&Token::RParen)
+            .with_context(|_| error::UnexpectedSnafu {
+                expected: "')'",
+                actual: self.peek_token_as_string(),
+            })?;
+        self.parser
+            .expect_keyword(Keyword::TO)
+            .context(error::SyntaxSnafu)?;
+        let (with, connection, location, limit) = self.parse_copy_parameters()?;
+        if limit.is_some() {
+            return error::InvalidSqlSnafu {
+                msg: "limit is not supported",
+            }
+            .fail();
+        }
+        Ok(CopyQueryTo {
+            query: Box::new(query),
+            arg: CopyQueryToArgument {
+                with: with.into(),
+                connection: connection.into(),
+                location,
+            },
+        })
+    }
+
     fn parse_copy_parameters(&mut self) -> Result<(With, Connection, String, Option<u64>)> {
         let location =
             self.parser
@@ -453,4 +495,75 @@ mod tests {
             stmt.connection.to_str_map()
         );
     }
+
+    #[test]
+    fn test_copy_query_to() {
+        let sql = "COPY (SELECT * FROM tbl WHERE ts > 10) TO 'tbl_file.parquet' WITH (FORMAT = 'parquet') CONNECTION (FOO='Bar', ONE='two')";
+        let stmt =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap()
+                .pop()
+                .unwrap();
+
+        let Copy(crate::statements::copy::Copy::CopyQueryTo(stmt)) = stmt else {
+            unreachable!()
+        };
+
+        let query = ParserContext::create_with_dialect(
+            "SELECT * FROM tbl WHERE ts > 10",
+            &GreptimeDbDialect {},
+            ParseOptions::default(),
+        )
+        .unwrap()
+        .remove(0);
+
+        assert_eq!(&query, stmt.query.as_ref());
+        assert_eq!(
+            [("format", "parquet")]
+                .into_iter()
+                .collect::<HashMap<_, _>>(),
+            stmt.arg.with.to_str_map()
+        );
+
+        assert_eq!(
+            [("foo", "Bar"), ("one", "two")]
+                .into_iter()
+                .collect::<HashMap<_, _>>(),
+            stmt.arg.connection.to_str_map()
+        );
+    }
+
+    #[test]
+    fn test_invalid_copy_query_to() {
+        {
+            let sql = "COPY SELECT * FROM tbl WHERE ts > 10 TO 'tbl_file.parquet' WITH (FORMAT = 'parquet') CONNECTION (FOO='Bar', ONE='two')";
+
+            assert!(ParserContext::create_with_dialect(
+                sql,
+                &GreptimeDbDialect {},
+                ParseOptions::default()
+            )
+            .is_err())
+        }
+        {
+            let sql = "COPY SELECT * FROM tbl WHERE ts > 10) TO 'tbl_file.parquet' WITH (FORMAT = 'parquet') CONNECTION (FOO='Bar', ONE='two')";
+
+            assert!(ParserContext::create_with_dialect(
+                sql,
+                &GreptimeDbDialect {},
+                ParseOptions::default()
+            )
+            .is_err())
+        }
+        {
+            let sql = "COPY (SELECT * FROM tbl WHERE ts > 10 TO 'tbl_file.parquet' WITH (FORMAT = 'parquet') CONNECTION (FOO='Bar', ONE='two')";
+
+            assert!(ParserContext::create_with_dialect(
+                sql,
+                &GreptimeDbDialect {},
+                ParseOptions::default()
+            )
+            .is_err())
+        }
+    }
 }

src/sql/src/statements/copy.rs

@@ -18,12 +18,14 @@ use serde::Serialize;
 use sqlparser::ast::ObjectName;
 use sqlparser_derive::{Visit, VisitMut};
 
+use crate::statements::statement::Statement;
 use crate::statements::OptionMap;
 
 #[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut, Serialize)]
 pub enum Copy {
     CopyTable(CopyTable),
     CopyDatabase(CopyDatabase),
+    CopyQueryTo(CopyQueryTo),
 }
 
 impl Display for Copy {
@@ -31,10 +33,32 @@ impl Display for Copy {
         match self {
             Copy::CopyTable(s) => s.fmt(f),
             Copy::CopyDatabase(s) => s.fmt(f),
+            Copy::CopyQueryTo(s) => s.fmt(f),
         }
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut, Serialize)]
+pub struct CopyQueryTo {
+    pub query: Box<Statement>,
+    pub arg: CopyQueryToArgument,
+}
+
+impl Display for CopyQueryTo {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "COPY ({}) TO {}", &self.query, &self.arg.location)?;
+        if !self.arg.with.is_empty() {
+            let options = self.arg.with.kv_pairs();
+            write!(f, " WITH ({})", options.join(", "))?;
+        }
+        if !self.arg.connection.is_empty() {
+            let options = self.arg.connection.kv_pairs();
+            write!(f, " CONNECTION ({})", options.join(", "))?;
+        }
+        Ok(())
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut, Serialize)]
 pub enum CopyTable {
     To(CopyTableArgument),
@@ -115,6 +139,13 @@ pub struct CopyTableArgument {
     pub limit: Option<u64>,
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut, Serialize)]
+pub struct CopyQueryToArgument {
+    pub with: OptionMap,
+    pub connection: OptionMap,
+    pub location: String,
+}
+
 #[cfg(test)]
 impl CopyTableArgument {
     pub fn format(&self) -> Option<String> {
@@ -238,4 +269,29 @@ mod tests {
             }
         }
     }
+
+    #[test]
+    fn test_display_copy_query_to() {
+        let sql = r"copy (SELECT * FROM tbl WHERE ts > 10) to 's3://my-bucket/data.parquet'
+            with (format = 'parquet', pattern = '.*parquet.*')
+            connection(region = 'us-west-2', secret_access_key = '12345678');";
+        let stmts =
+            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
+                .unwrap();
+        assert_eq!(1, stmts.len());
+        assert_matches!(&stmts[0], Statement::Copy { .. });
+
+        match &stmts[0] {
+            Statement::Copy(copy) => {
+                let new_sql = format!("{}", copy);
+                assert_eq!(
+                    r#"COPY (SELECT * FROM tbl WHERE ts > 10) TO s3://my-bucket/data.parquet WITH (format = 'parquet', pattern = '.*parquet.*') CONNECTION (region = 'us-west-2', secret_access_key = '******')"#,
+                    &new_sql
+                );
+            }
+            _ => {
+                unreachable!();
+            }
+        }
+    }
 }