feat: copy database from (#3164)

* wip: impl COPY DATABASE FROM parser

* wip: impl copy database from

* wip: add some ut

* wip: add continue_on_error option

* test: add sqlness cases for copy database

* fix: trailing newline

* fix: typo

* fix: some cr comments

* chore: resolve confilicts

* fix: some cr comments

src/sql/src/parsers/copy_parser.rs

@@ -15,13 +15,12 @@
 use std::collections::HashMap;
 
 use snafu::ResultExt;
-use sqlparser::ast::ObjectName;
 use sqlparser::keywords::Keyword;
 use sqlparser::tokenizer::Token::Word;
 
 use crate::error::{self, Result};
 use crate::parser::ParserContext;
-use crate::statements::copy::{CopyDatabaseArgument, CopyTable, CopyTableArgument};
+use crate::statements::copy::{CopyDatabase, CopyDatabaseArgument, CopyTable, CopyTableArgument};
 use crate::statements::statement::Statement;
 use crate::util::parse_option_string;
 
@@ -47,7 +46,7 @@ impl<'a> ParserContext<'a> {
         Ok(Statement::Copy(copy))
     }
 
-    fn parser_copy_database(&mut self) -> Result<CopyDatabaseArgument> {
+    fn parser_copy_database(&mut self) -> Result<CopyDatabase> {
         let database_name =
             self.parser
                 .parse_object_name()
@@ -57,17 +56,29 @@ impl<'a> ParserContext<'a> {
                     actual: self.peek_token_as_string(),
                 })?;
 
-        self.parser
-            .expect_keyword(Keyword::TO)
-            .context(error::SyntaxSnafu)?;
-
-        let (with, connection, location) = self.parse_copy_to()?;
-        Ok(CopyDatabaseArgument {
-            database_name,
-            with: with.into(),
-            connection: connection.into(),
-            location,
-        })
+        let req = if self.parser.parse_keyword(Keyword::TO) {
+            let (with, connection, location) = self.parse_copy_parameters()?;
+            let argument = CopyDatabaseArgument {
+                database_name,
+                with: with.into(),
+                connection: connection.into(),
+                location,
+            };
+            CopyDatabase::To(argument)
+        } else {
+            self.parser
+                .expect_keyword(Keyword::FROM)
+                .context(error::SyntaxSnafu)?;
+            let (with, connection, location) = self.parse_copy_parameters()?;
+            let argument = CopyDatabaseArgument {
+                database_name,
+                with: with.into(),
+                connection: connection.into(),
+                location,
+            };
+            CopyDatabase::From(argument)
+        };
+        Ok(req)
     }
 
     fn parse_copy_table(&mut self) -> Result<CopyTable> {
@@ -82,7 +93,7 @@ impl<'a> ParserContext<'a> {
         let table_name = Self::canonicalize_object_name(raw_table_name);
 
         if self.parser.parse_keyword(Keyword::TO) {
-            let (with, connection, location) = self.parse_copy_to()?;
+            let (with, connection, location) = self.parse_copy_parameters()?;
             Ok(CopyTable::To(CopyTableArgument {
                 table_name,
                 with: with.into(),
@@ -93,52 +104,17 @@ impl<'a> ParserContext<'a> {
             self.parser
                 .expect_keyword(Keyword::FROM)
                 .context(error::SyntaxSnafu)?;
-            Ok(CopyTable::From(self.parse_copy_table_from(table_name)?))
+            let (with, connection, location) = self.parse_copy_parameters()?;
+            Ok(CopyTable::From(CopyTableArgument {
+                table_name,
+                with: with.into(),
+                connection: connection.into(),
+                location,
+            }))
         }
     }
 
-    fn parse_copy_table_from(&mut self, table_name: ObjectName) -> Result<CopyTableArgument> {
-        let location =
-            self.parser
-                .parse_literal_string()
-                .with_context(|_| error::UnexpectedSnafu {
-                    sql: self.sql,
-                    expected: "a uri",
-                    actual: self.peek_token_as_string(),
-                })?;
-
-        let options = self
-            .parser
-            .parse_options(Keyword::WITH)
-            .context(error::SyntaxSnafu)?;
-
-        let with = options
-            .into_iter()
-            .filter_map(|option| {
-                parse_option_string(option.value).map(|v| (option.name.value.to_lowercase(), v))
-            })
-            .collect();
-
-        let connection_options = self
-            .parser
-            .parse_options(Keyword::CONNECTION)
-            .context(error::SyntaxSnafu)?;
-
-        let connection = connection_options
-            .into_iter()
-            .filter_map(|option| {
-                parse_option_string(option.value).map(|v| (option.name.value.to_lowercase(), v))
-            })
-            .collect();
-        Ok(CopyTableArgument {
-            table_name,
-            with,
-            connection,
-            location,
-        })
-    }
-
-    fn parse_copy_to(&mut self) -> Result<(With, Connection, String)> {
+    fn parse_copy_parameters(&mut self) -> Result<(With, Connection, String)> {
         let location =
             self.parser
                 .parse_literal_string()
@@ -181,7 +157,7 @@ mod tests {
     use std::assert_matches::assert_matches;
     use std::collections::HashMap;
 
-    use sqlparser::ast::Ident;
+    use sqlparser::ast::{Ident, ObjectName};
 
     use super::*;
     use crate::dialect::GreptimeDbDialect;
@@ -398,6 +374,50 @@ mod tests {
         let Copy(crate::statements::copy::Copy::CopyDatabase(stmt)) = stmt else {
             unreachable!()
         };
+
+        let CopyDatabase::To(stmt) = stmt else {
+            unreachable!()
+        };
+
+        assert_eq!(
+            ObjectName(vec![Ident::new("catalog0"), Ident::new("schema0")]),
+            stmt.database_name
+        );
+        assert_eq!(
+            [("format".to_string(), "parquet".to_string())]
+                .into_iter()
+                .collect::<HashMap<_, _>>(),
+            stmt.with.map
+        );
+
+        assert_eq!(
+            [
+                ("foo".to_string(), "Bar".to_string()),
+                ("one".to_string(), "two".to_string())
+            ]
+            .into_iter()
+            .collect::<HashMap<_, _>>(),
+            stmt.connection.map
+        );
+    }
+
+    #[test]
+    fn test_copy_database_from() {
+        let sql = "COPY DATABASE catalog0.schema0 FROM '/a/b/c/' 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::CopyDatabase(stmt)) = stmt else {
+            unreachable!()
+        };
+
+        let CopyDatabase::From(stmt) = stmt else {
+            unreachable!()
+        };
+
         assert_eq!(
             ObjectName(vec![Ident::new("catalog0"), Ident::new("schema0")]),
             stmt.database_name

src/sql/src/statements/copy.rs

@@ -20,7 +20,7 @@ use crate::statements::OptionMap;
 #[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut)]
 pub enum Copy {
     CopyTable(CopyTable),
-    CopyDatabase(CopyDatabaseArgument),
+    CopyDatabase(CopyDatabase),
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut)]
@@ -29,6 +29,12 @@ pub enum CopyTable {
     From(CopyTableArgument),
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut)]
+pub enum CopyDatabase {
+    To(CopyDatabaseArgument),
+    From(CopyDatabaseArgument),
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut)]
 pub struct CopyDatabaseArgument {
     pub database_name: ObjectName,