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
GreptimeTeam · Jan 22, 2024 · 3834ea7 · 3834ea7
1 parent 966875e
commit 3834ea7
Show file tree

Hide file tree

Showing 12 changed files with 459 additions and 163 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs
@@ -71,7 +71,7 @@ use session::context::QueryContextRef;
 use snafu::prelude::*;
 use sql::dialect::Dialect;
 use sql::parser::{ParseOptions, ParserContext};
-use sql::statements::copy::CopyTable;
+use sql::statements::copy::{CopyDatabase, CopyTable};
 use sql::statements::statement::Statement;
 use sqlparser::ast::ObjectName;
 pub use standalone::StandaloneDatanodeManager;
@@ -487,8 +487,11 @@ pub fn check_permission(
                 validate_param(&copy_table_from.table_name, query_ctx)?
             }
         },
-        Statement::Copy(sql::statements::copy::Copy::CopyDatabase(stmt)) => {
-            validate_param(&stmt.database_name, query_ctx)?
+        Statement::Copy(sql::statements::copy::Copy::CopyDatabase(copy_database)) => {
+            match copy_database {
+                CopyDatabase::To(stmt) => validate_param(&stmt.database_name, query_ctx)?,
+                CopyDatabase::From(stmt) => validate_param(&stmt.database_name, query_ctx)?,
+            }
         }
         Statement::TruncateTable(stmt) => {
             validate_param(stmt.table_name(), query_ctx)?;

diff --git a/src/operator/Cargo.toml b/src/operator/Cargo.toml
@@ -54,3 +54,6 @@ substrait.workspace = true
 table.workspace = true
 tokio.workspace = true
 tonic.workspace = true
+
+[dev-dependencies]
+common-test-util.workspace = true
diff --git a/src/operator/src/error.rs b/src/operator/src/error.rs
@@ -418,6 +418,9 @@ pub enum Error {
         location: Location,
     },
 
+    #[snafu(display("Invalid COPY DATABASE location, must end with '/': {}", value))]
+    InvalidCopyDatabasePath { value: String, location: Location },
+
     #[snafu(display("Table metadata manager error"))]
     TableMetadataManager {
         source: common_meta::error::Error,
@@ -596,7 +599,9 @@ impl ErrorExt for Error {
             | Error::BuildBackend { source, .. } => source.status_code(),
 
             Error::ExecuteDdl { source, .. } => source.status_code(),
-            Error::InvalidCopyParameter { .. } => StatusCode::InvalidArguments,
+            Error::InvalidCopyParameter { .. } | Error::InvalidCopyDatabasePath { .. } => {
+                StatusCode::InvalidArguments
+            }
 
             Error::ReadRecordBatch { source, .. } | Error::BuildColumnVectors { source, .. } => {
                 source.status_code()

diff --git a/src/operator/src/statement.rs b/src/operator/src/statement.rs
@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-mod backup;
+mod copy_database;
 mod copy_table_from;
 mod copy_table_to;
 mod ddl;
@@ -41,7 +41,7 @@ use query::plan::LogicalPlan;
 use query::QueryEngineRef;
 use session::context::QueryContextRef;
 use snafu::{OptionExt, ResultExt};
-use sql::statements::copy::{CopyDatabaseArgument, CopyTable, CopyTableArgument};
+use sql::statements::copy::{CopyDatabase, CopyDatabaseArgument, CopyTable, CopyTableArgument};
 use sql::statements::statement::Statement;
 use sql::statements::OptionMap;
 use sql::util::format_raw_object_name;
@@ -55,7 +55,7 @@ use crate::error::{
     PlanStatementSnafu, Result, TableNotFoundSnafu,
 };
 use crate::insert::InserterRef;
-use crate::statement::backup::{COPY_DATABASE_TIME_END_KEY, COPY_DATABASE_TIME_START_KEY};
+use crate::statement::copy_database::{COPY_DATABASE_TIME_END_KEY, COPY_DATABASE_TIME_START_KEY};
 use crate::table::table_idents_to_full_name;
 
 #[derive(Clone)]
@@ -131,9 +131,23 @@ impl StatementExecutor {
                 }
             }
 
-            Statement::Copy(sql::statements::copy::Copy::CopyDatabase(arg)) => {
-                self.copy_database(to_copy_database_request(arg, &query_ctx)?)
-                    .await
+            Statement::Copy(sql::statements::copy::Copy::CopyDatabase(copy_database)) => {
+                match copy_database {
+                    CopyDatabase::To(arg) => {
+                        self.copy_database_to(
+                            to_copy_database_request(arg, &query_ctx)?,
+                            query_ctx.clone(),
+                        )
+                        .await
+                    }
+                    CopyDatabase::From(arg) => {
+                        self.copy_database_from(
+                            to_copy_database_request(arg, &query_ctx)?,
+                            query_ctx,
+                        )
+                        .await
+                    }
+                }
             }
 
             Statement::CreateTable(stmt) => {

diff --git a/src/operator/src/statement/backup.rs b/src/operator/src/statement/backup.rs