From dca01654bc77dab55f86badd6a9e8e79fbdfb1e1 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Thu, 17 Sep 2026 12:02:39 +0000 Subject: [PATCH] feat: prepare and execute database Metric exports (#9180) * feat: prepare and authorize captured database exports Signed-off-by: jeremyhi * feat: bound database export jobs and drain failures Signed-off-by: jeremyhi * test: validate database export identity and restore equivalence Signed-off-by: jeremyhi * fix: validate database export directory URLs Signed-off-by: jeremyhi * fix: preserve Windows export directory paths Signed-off-by: jeremyhi * fix: reject local database export filename aliases Signed-off-by: jeremyhi * refactor: consolidate database export planning policies Signed-off-by: jeremyhi * fix: restore escaped database export filenames Signed-off-by: jeremyhi * test: align database restore assertions with shared policies Signed-off-by: jeremyhi * refactor: clarify database export boundaries and names Signed-off-by: jeremyhi --------- Signed-off-by: jeremyhi --- Cargo.lock | 1 + Cargo.toml | 1 + src/frontend/src/instance.rs | 91 +++ src/frontend/src/instance/export_database.rs | 113 ++++ src/operator/Cargo.toml | 3 +- src/operator/src/error.rs | 8 + src/operator/src/statement.rs | 2 + src/operator/src/statement/copy_database.rs | 144 +---- src/operator/src/statement/copy_table_to.rs | 14 +- src/operator/src/statement/database_copy.rs | 313 +++++++++ src/operator/src/statement/export_database.rs | 400 ++++++++++++ .../src/statement/export_logical_tables.rs | 25 +- .../statement/export_logical_tables/tests.rs | 2 - src/query/src/dist_plan/planner.rs | 53 +- .../tests/export_logical_tables.rs | 602 +++++++++++++++++- 15 files changed, 1594 insertions(+), 178 deletions(-) create mode 100644 src/frontend/src/instance/export_database.rs create mode 100644 src/operator/src/statement/database_copy.rs create mode 100644 src/operator/src/statement/export_database.rs diff --git a/Cargo.lock b/Cargo.lock index d6c94a404e..e8c03e616c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10074,6 +10074,7 @@ dependencies = [ "object-store", "partition", "path-slash", + "percent-encoding", "prometheus 0.14.0", "prost 0.14.1", "prost-types 0.14.1", diff --git a/Cargo.toml b/Cargo.toml index 96ce122962..52b98de555 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -202,6 +202,7 @@ parquet-variant = "=59.2.0" parquet-variant-compute = "=59.2.0" parquet-variant-json = "=59.2.0" paste = "1.0" +percent-encoding = "2.3" pin-project = "1.0" pretty_assertions = "1.4.0" prometheus = { version = "0.14", features = ["process"] } diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index 7225987d75..44949f932f 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -15,6 +15,7 @@ pub mod builder; mod dashboard; mod entity_graph; +mod export_database; mod grpc; mod influxdb; mod jaeger; @@ -2877,6 +2878,96 @@ mod tests { assert_eq!(CheckedAction { action, targets }, checker.take_check()); } + #[tokio::test] + async fn database_export_authorizes_all_tables_before_preparation() -> TestResult<()> { + struct ExportAcl(std::sync::Mutex>); + impl PermissionChecker for ExportAcl { + fn check_permission( + &self, + _: UserInfoRef, + req: PermissionReq, + ) -> auth::error::Result { + assert!(matches!( + req, + PermissionReq::SqlStatement(Statement::Copy( + sql::statements::copy::Copy::CopyDatabase( + sql::statements::copy::CopyDatabase::To(_) + ) + )) + )); + Ok(PermissionResp::Allow) + } + fn check_permission_with_table_targets( + &self, + _: UserInfoRef, + req: PermissionReq, + targets: PermissionTableTargets, + ) -> auth::error::Result { + self.0.lock().unwrap().push(targets.clone()); + let PermissionTableTargets::Resolved(tables) = targets else { + panic!("unresolved export") + }; + if req.is_readonly() { + // The first table is readable; the later table has only write access. + Ok(if tables.iter().any(|t| t.table == "target") { + PermissionResp::Reject + } else { + PermissionResp::Allow + }) + } else { + self.check_permission(QueryContext::arc().current_user(), req) + } + } + } + let checker = Arc::new(ExportAcl(Default::default())); + let plugins = Plugins::new(); + plugins.insert::(checker.clone()); + let instance = test_instance_with_plugins( + test_logical_table(1024, "source")?, + test_table(1025, "target")?, + plugins, + ) + .await?; + let req = table::requests::CopyDatabaseRequest { + catalog_name: "greptime".into(), + schema_name: "public".into(), + location: "invalid-destination".into(), + with: Default::default(), + connection: Default::default(), + time_range: None, + }; + let result = instance + .export_database_for_test( + req.clone(), + None, + &tokio_util::sync::CancellationToken::new(), + QueryContext::arc(), + ) + .await; + assert!(matches!(result, Err(Error::Permission { .. }))); + let expected = PermissionTableTargets::resolved(vec![ + PermissionTableTarget::new("greptime", "public", "source"), + PermissionTableTarget::new("greptime", "public", "target"), + ]); + assert_eq!(*checker.0.lock().unwrap(), vec![expected.clone(), expected]); + // An empty selection must still check the operation privilege. + instance + .plugins + .map_mut::(|checker| { + *checker.unwrap() = Arc::new(WriteOnlyPermissionChecker) + }); + let result = instance + .export_database_for_test( + req, + Some(&[]), + &tokio_util::sync::CancellationToken::new(), + QueryContext::arc(), + ) + .await; + assert!(matches!(result, Err(Error::Permission { .. }))); + Ok(()) + } + #[tokio::test] async fn test_prom_remote_read_with_custom_timestamp_and_value_columns() -> TestResult<()> { let schema = Arc::new(GtSchema::new(vec![ diff --git a/src/frontend/src/instance/export_database.rs b/src/frontend/src/instance/export_database.rs new file mode 100644 index 0000000000..6594490548 --- /dev/null +++ b/src/frontend/src/instance/export_database.rs @@ -0,0 +1,113 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use auth::{ + PermissionAction, PermissionChecker, PermissionCheckerRef, PermissionReq, + PermissionTableTarget, PermissionTableTargets, +}; +use operator::statement::export_database::{DatabaseExportSummary, PreparedDatabaseExport}; +use session::context::QueryContextRef; +use snafu::ResultExt; +use sql::ast::{Ident, ObjectName}; +use sql::statements::copy::{Copy, CopyDatabase, CopyDatabaseArgument}; +use sql::statements::statement::Statement; +use table::requests::CopyDatabaseRequest; +use tokio_util::sync::CancellationToken; + +use crate::error::{PermissionSnafu, Result}; +use crate::instance::Instance; + +impl Instance { + #[allow(dead_code)] + async fn export_database( + &self, + req: CopyDatabaseRequest, + names: Option<&[String]>, + cancellation: &CancellationToken, + ctx: QueryContextRef, + ) -> Result { + if cancellation.is_cancelled() { + return Err(operator::error::DatabaseExportCancelledSnafu.build().into()); + } + let plan = self.prepare_database_export(req, names, &ctx).await?; + Ok(self + .statement_executor + .export_database(plan, cancellation, ctx) + .await?) + } + + /// Exercise the internal authenticated entry without exposing SQL/CLI activation. + #[cfg(any(test, feature = "testing"))] + pub async fn export_database_for_test( + &self, + req: CopyDatabaseRequest, + names: Option<&[String]>, + cancellation: &CancellationToken, + ctx: QueryContextRef, + ) -> Result { + self.export_database(req, names, cancellation, ctx).await + } + + async fn prepare_database_export( + &self, + req: CopyDatabaseRequest, + names: Option<&[String]>, + ctx: &QueryContextRef, + ) -> Result { + let stmt = Statement::Copy(Copy::CopyDatabase(CopyDatabase::To(CopyDatabaseArgument { + database_name: ObjectName::from(vec![ + Ident::new(&req.catalog_name), + Ident::new(&req.schema_name), + ]), + with: req.with.clone().into(), + connection: req.connection.clone().into(), + location: req.location.clone(), + }))); + self.plugins + .get::() + .as_ref() + .check_permission_with_context( + ctx.current_user(), + PermissionReq::SqlStatement(&stmt), + Some(&ctx.current_schema()), + ) + .context(PermissionSnafu)?; + let tables = self + .statement_executor + .capture_database_export_tables(&req, names, ctx) + .await?; + let targets = PermissionTableTargets::resolved( + tables + .iter() + .map(|table| { + let info = table.table_info(); + PermissionTableTarget::new(&info.catalog_name, &info.schema_name, &info.name) + }) + .collect(), + ); + self.check_table_permission(ctx, PermissionReq::SqlStatement(&stmt), targets.clone()) + .context(PermissionSnafu)?; + // COPY is classified as a write by existing permission checkers. + self.check_table_permission( + ctx, + PermissionReq::Action(PermissionAction::read("database.export")), + targets, + ) + .context(PermissionSnafu)?; + Ok(self + .statement_executor + .prepare_database_export(req, tables) + .await?) + } +} diff --git a/src/operator/Cargo.toml b/src/operator/Cargo.toml index 4ed4927b41..ea1bc990f7 100644 --- a/src/operator/Cargo.toml +++ b/src/operator/Cargo.toml @@ -64,6 +64,7 @@ meter-macros.workspace = true moka = { workspace = true, features = ["future"] } object-store.workspace = true partition.workspace = true +percent-encoding.workspace = true prometheus.workspace = true prost.workspace = true prost-types = { workspace = true, optional = true } @@ -83,6 +84,7 @@ tokio.workspace = true tokio-util.workspace = true tonic.workspace = true tracing.workspace = true +url.workspace = true [dev-dependencies] axum.workspace = true @@ -91,4 +93,3 @@ common-meta = { workspace = true, features = ["testing"] } common-test-util.workspace = true object-store = { workspace = true, features = ["testing"] } path-slash = "0.2" -url.workspace = true diff --git a/src/operator/src/error.rs b/src/operator/src/error.rs index c6b84543c4..5a005250b9 100644 --- a/src/operator/src/error.rs +++ b/src/operator/src/error.rs @@ -86,6 +86,12 @@ pub enum Error { source: common_meta::error::Error, }, + #[snafu(display("Invalid database export: {reason}"))] + InvalidDatabaseExport { reason: String }, + + #[snafu(display("Database export cancelled"))] + DatabaseExportCancelled {}, + #[snafu(display("Invalid logical table export: {reason}"))] InvalidLogicalTableExport { reason: String }, @@ -1089,6 +1095,8 @@ impl ErrorExt for Error { Error::InvalidTimeIndexType { .. } | Error::InvalidTimezone { .. } => { StatusCode::InvalidArguments } + Error::InvalidDatabaseExport { .. } => StatusCode::InvalidArguments, + Error::DatabaseExportCancelled { .. } => StatusCode::Cancelled, Error::InvalidLogicalTableExport { .. } => StatusCode::InvalidArguments, Error::LogicalTableExportResource { .. } => StatusCode::Suspended, Error::LogicalTableExportCancelled { .. } => StatusCode::Cancelled, diff --git a/src/operator/src/statement.rs b/src/operator/src/statement.rs index 69d387de20..61804b30d4 100644 --- a/src/operator/src/statement.rs +++ b/src/operator/src/statement.rs @@ -19,9 +19,11 @@ mod copy_query_to; mod copy_table_from; mod copy_table_to; mod cursor; +mod database_copy; pub mod ddl; mod describe; mod dml; +pub mod export_database; pub mod export_logical_tables; mod kill; pub mod semantic_graph; diff --git a/src/operator/src/statement/copy_database.rs b/src/operator/src/statement/copy_database.rs index 43b2692cdd..7a9dd84717 100644 --- a/src/operator/src/statement/copy_database.rs +++ b/src/operator/src/statement/copy_database.rs @@ -12,8 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::HashMap; -use std::path::Path; use std::str::FromStr; use std::sync::Arc; @@ -21,56 +19,26 @@ use client::{Output, OutputData, OutputMeta}; use common_catalog::format_full_table_name; use common_datasource::file_format::Format; use common_datasource::lister::{Lister, Source}; -#[cfg(windows)] -use common_datasource::object_store::{FS_SCHEMA, parse_url}; use common_datasource::object_store::{LocalFileAccess, build_backend, build_backend_for_write}; -use common_stat::get_total_cpu_cores; use common_telemetry::{debug, error, info, tracing}; use futures::future::try_join_all; use object_store::Entry; use regex::Regex; use session::context::QueryContextRef; -use snafu::{OptionExt, ResultExt, ensure}; -use store_api::metric_engine_consts::{LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME}; +use snafu::ResultExt; use table::requests::{CopyDatabaseRequest, CopyDirection, CopyTableRequest}; -use table::table_reference::TableReference; use tokio::sync::Semaphore; use crate::error; -use crate::error::{CatalogSnafu, InvalidCopyDatabasePathSnafu}; use crate::statement::StatementExecutor; +use crate::statement::database_copy::{ + DatabaseExportFile, database_import_source, parse_parallelism_from_option_map, + validate_database_directory, +}; pub(crate) const COPY_DATABASE_TIME_START_KEY: &str = "start_time"; pub(crate) const COPY_DATABASE_TIME_END_KEY: &str = "end_time"; pub(crate) const CONTINUE_ON_ERROR_KEY: &str = "continue_on_error"; -pub(crate) const PARALLELISM_KEY: &str = "parallelism"; - -fn is_directory_location(location: &str) -> bool { - if location.ends_with('/') { - return true; - } - - #[cfg(windows)] - { - location.ends_with(std::path::MAIN_SEPARATOR) - && matches!( - parse_url(location), - Ok((schema, _, _)) if schema.eq_ignore_ascii_case(FS_SCHEMA) - ) - } - - #[cfg(not(windows))] - false -} - -/// Get parallelism from options, default to total CPU cores. -fn parse_parallelism_from_option_map(options: &HashMap) -> usize { - options - .get(PARALLELISM_KEY) - .and_then(|v| v.parse::().ok()) - .unwrap_or_else(get_total_cpu_cores) - .max(1) -} impl StatementExecutor { #[tracing::instrument(skip_all)] @@ -79,13 +47,7 @@ impl StatementExecutor { req: CopyDatabaseRequest, ctx: QueryContextRef, ) -> error::Result { - // Location must end with a separator so that every table is exported to a file. - ensure!( - is_directory_location(&req.location), - InvalidCopyDatabasePathSnafu { - value: req.location, - } - ); + validate_database_directory(&req.location)?; build_backend_for_write(&req.location, &req.connection, &self.local_file_access) .await .context(error::BuildBackendSnafu)?; @@ -95,12 +57,10 @@ impl StatementExecutor { "Copy database {}.{} to dir: {}, time: {:?}, parallelism: {}", req.catalog_name, req.schema_name, req.location, req.time_range, parallelism ); - let table_names = self - .catalog_manager - .table_names(&req.catalog_name, &req.schema_name, Some(&ctx)) - .await - .context(CatalogSnafu)?; - let num_tables = table_names.len(); + let tables = self + .capture_database_export_tables(&req, None, &ctx) + .await?; + let num_tables = tables.len(); let suffix = Format::try_from(&req.with) .context(error::ParseFileFormatSnafu)? @@ -109,34 +69,10 @@ impl StatementExecutor { let mut tasks = Vec::with_capacity(num_tables); let semaphore = Arc::new(Semaphore::new(parallelism)); - for (i, table_name) in table_names.into_iter().enumerate() { - let table = self - .get_table(&TableReference { - catalog: &req.catalog_name, - schema: &req.schema_name, - table: &table_name, - }) - .await?; - // Only base tables, ignores views and temporary tables. - if table.table_type() != table::metadata::TableType::Base { - continue; - } - // Ignores physical tables of metric engine. - if table.table_info().meta.engine == METRIC_ENGINE_NAME - && !table - .table_info() - .meta - .options - .extra_options - .contains_key(LOGICAL_TABLE_METADATA_KEY) - { - continue; - } - + for (i, table) in tables.into_iter().enumerate() { + let table_name = table.table_info().name.clone(); let semaphore_moved = semaphore.clone(); - let mut table_file = req.location.clone(); - table_file.push_str(&table_name); - table_file.push_str(suffix); + let table_file = DatabaseExportFile::new(&req.location, &table_name, suffix)?.location; let table_no = i + 1; let moved_ctx = ctx.clone(); let full_table_name = @@ -160,7 +96,8 @@ impl StatementExecutor { "Copy table({}/{}): {} to {}", table_no, num_tables, full_table_name, table_file ); - self.copy_table_to(copy_table_req, moved_ctx).await + self.copy_captured_table_to(table, copy_table_req, moved_ctx) + .await }); } @@ -176,13 +113,7 @@ impl StatementExecutor { req: CopyDatabaseRequest, ctx: QueryContextRef, ) -> error::Result { - // Location must end with a directory separator. - ensure!( - is_directory_location(&req.location), - InvalidCopyDatabasePathSnafu { - value: req.location, - } - ); + validate_database_directory(&req.location)?; let parallelism = parse_parallelism_from_option_map(&req.with); info!( @@ -205,8 +136,8 @@ impl StatementExecutor { let semaphore = Arc::new(Semaphore::new(parallelism)); for e in entries { - let table_name = match parse_file_name_to_copy(&e) { - Ok(table_name) => table_name, + let (table_name, location) = match database_import_source(&req.location, e.path()) { + Ok(source) => source, Err(err) => { if continue_on_error { error!(err; "Failed to import table from file: {:?}", e); @@ -221,7 +152,7 @@ impl StatementExecutor { catalog_name: req.catalog_name.clone(), schema_name: req.schema_name.clone(), table_name: table_name.clone(), - location: format!("{}{}", req.location, e.path()), + location, with: req.with.clone(), connection: req.connection.clone(), pattern: None, @@ -266,17 +197,6 @@ impl StatementExecutor { } } -/// Parses table names from files' names. -fn parse_file_name_to_copy(e: &Entry) -> error::Result { - Path::new(e.name()) - .file_stem() - .and_then(|os_str| os_str.to_str()) - .map(|s| s.to_string()) - .context(error::InvalidTableNameSnafu { - table_name: e.name().to_string(), - }) -} - /// Lists all files with expected suffix that can be imported to database. async fn list_files_to_copy( req: &CopyDatabaseRequest, @@ -299,10 +219,9 @@ async fn list_files_to_copy( #[cfg(test)] mod tests { - use std::collections::{HashMap, HashSet}; + use std::collections::HashSet; use common_datasource::object_store::LocalFileAccess; - use common_stat::get_total_cpu_cores; use object_store::ObjectStore; use object_store::services::Fs; use object_store::util::normalize_dir; @@ -310,9 +229,8 @@ mod tests { use path_slash::PathExt; use table::requests::CopyDatabaseRequest; - use crate::statement::copy_database::{ - list_files_to_copy, parse_file_name_to_copy, parse_parallelism_from_option_map, - }; + use crate::statement::copy_database::list_files_to_copy; + use crate::statement::database_copy::database_import_source; #[tokio::test] async fn test_list_files_and_parse_table_name() { @@ -345,7 +263,11 @@ mod tests { .await .unwrap() .into_iter() - .map(|e| parse_file_name_to_copy(&e).unwrap()) + .map(|e| { + database_import_source(&request.location, e.path()) + .unwrap() + .0 + }) .collect::>(); assert_eq!( @@ -355,16 +277,4 @@ mod tests { listed ); } - - #[test] - fn test_parse_parallelism_from_option_map() { - let options = HashMap::new(); - assert_eq!( - parse_parallelism_from_option_map(&options), - get_total_cpu_cores() - ); - - let options = HashMap::from([("parallelism".to_string(), "0".to_string())]); - assert_eq!(parse_parallelism_from_option_map(&options), 1); - } } diff --git a/src/operator/src/statement/copy_table_to.rs b/src/operator/src/statement/copy_table_to.rs index faa06b8821..21b82f51a3 100644 --- a/src/operator/src/statement/copy_table_to.rs +++ b/src/operator/src/statement/copy_table_to.rs @@ -35,6 +35,7 @@ use datafusion_expr::LogicalPlanBuilder; use object_store::ObjectStore; use session::context::QueryContextRef; use snafu::{OptionExt, ResultExt}; +use table::TableRef; use table::requests::CopyTableRequest; use table::table::adapter::DfTableProviderAdapter; use table::table_reference::TableReference; @@ -108,7 +109,18 @@ impl StatementExecutor { ) -> Result { let table_ref = TableReference::full(&req.catalog_name, &req.schema_name, &req.table_name); let table = self.get_table(&table_ref).await?; - let table_id = table.table_info().table_id(); + self.copy_captured_table_to(table, req, query_ctx).await + } + + pub(crate) async fn copy_captured_table_to( + &self, + table: TableRef, + req: CopyTableRequest, + query_ctx: QueryContextRef, + ) -> Result { + let info = table.table_info(); + let table_ref = TableReference::full(&info.catalog_name, &info.schema_name, &info.name); + let table_id = info.table_id(); let format = Format::try_from(&req.with).context(error::ParseFileFormatSnafu)?; let df_table_ref = DfTableReference::from(table_ref); diff --git a/src/operator/src/statement/database_copy.rs b/src/operator/src/statement/database_copy.rs new file mode 100644 index 0000000000..94ff1df11b --- /dev/null +++ b/src/operator/src/statement/database_copy.rs @@ -0,0 +1,313 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Shared selection and destination rules for database COPY and prepared exports. + +use std::collections::HashMap; + +use common_datasource::object_store::FILE_SCHEMA; +#[cfg(windows)] +use common_datasource::object_store::{FS_SCHEMA, parse_url}; +use common_stat::get_total_cpu_cores; +use session::context::QueryContextRef; +use snafu::{OptionExt, ResultExt, ensure}; +use store_api::metric_engine_consts::{LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME}; +use table::TableRef; +use table::metadata::TableType; +use table::requests::CopyDatabaseRequest; +use table::table_reference::TableReference; +use url::Url; + +use crate::error::{self, Result}; +use crate::statement::StatementExecutor; + +fn is_directory_location(location: &str) -> bool { + if location.ends_with('/') { + return true; + } + + #[cfg(windows)] + { + location.ends_with(std::path::MAIN_SEPARATOR) + && matches!( + parse_url(location), + Ok((schema, _, _)) if schema.eq_ignore_ascii_case(FS_SCHEMA) + ) + } + + #[cfg(not(windows))] + false +} + +/// Get parallelism from options, default to total CPU cores. +pub(crate) fn parse_parallelism_from_option_map(options: &HashMap) -> usize { + options + .get("parallelism") + .and_then(|v| v.parse::().ok()) + .unwrap_or_else(get_total_cpu_cores) + .max(1) +} + +pub(crate) fn validate_database_directory(location: &str) -> Result<()> { + ensure!( + is_directory_location(location), + error::InvalidCopyDatabasePathSnafu { value: location } + ); + #[cfg(windows)] + if common_datasource::object_store::handle_windows_path(location).is_some() { + return Ok(()); + } + let parsed_directory = match Url::parse(location) { + Ok(url) => { + url.query().is_none() && url.fragment().is_none() && is_directory_location(url.path()) + } + Err(_) => true, + }; + ensure!( + parsed_directory, + error::InvalidCopyDatabasePathSnafu { value: location } + ); + Ok(()) +} + +/// The writer key and its externally reported location, resolved together. +pub(crate) struct DatabaseExportFile { + pub(crate) path: String, + pub(crate) location: String, +} + +impl DatabaseExportFile { + pub(crate) fn new(directory: &str, name: &str, suffix: &str) -> Result { + let filename = format!("{name}{suffix}"); + #[cfg(windows)] + if common_datasource::object_store::handle_windows_path(directory).is_some() { + return Ok(Self { + location: format!("{directory}{filename}"), + path: filename, + }); + } + match Url::parse(directory) { + Ok(mut url) => { + url.path_segments_mut() + .map_err(|_| error::InvalidCopyDatabasePathSnafu { value: directory }.build())? + .pop_if_empty() + .push(&filename); + // File URLs are decoded by the filesystem backend; object-store + // backends use the encoded URL path as their key. + let path = if url.scheme().eq_ignore_ascii_case(FILE_SCHEMA) { + filename + } else { + url.path() + .rsplit('/') + .next() + .unwrap_or_default() + .to_string() + }; + Ok(Self { + path, + location: url.into(), + }) + } + Err(url::ParseError::RelativeUrlWithoutBase) => Ok(Self { + location: format!("{directory}{filename}"), + path: filename, + }), + Err(source) => Err(source) + .context(common_datasource::error::InvalidUrlSnafu { url: directory }) + .context(error::BuildBackendSnafu), + } + } +} + +/// Resolve a listed writer key back to its table name and COPY input location. +pub(crate) fn database_import_source(directory: &str, path: &str) -> Result<(String, String)> { + let mut filename = path.rsplit('/').next().unwrap_or(path).to_string(); + let mut location = format!("{directory}{path}"); + #[cfg(windows)] + let literal_path = common_datasource::object_store::handle_windows_path(directory).is_some(); + #[cfg(not(windows))] + let literal_path = false; + if !literal_path && let Ok(mut url) = Url::parse(directory) { + if url.scheme().eq_ignore_ascii_case(FILE_SCHEMA) { + url.path_segments_mut() + .map_err(|_| error::InvalidCopyDatabasePathSnafu { value: directory }.build())? + .pop_if_empty() + .extend(path.split('/')); + } else { + // Listed object keys already contain the export URL's escaping. + url.set_path(&format!("{}{path}", url.path())); + filename = percent_encoding::percent_decode_str(&filename) + .decode_utf8() + .ok() + .context(error::InvalidTableNameSnafu { table_name: path })? + .into_owned(); + } + location = url.into(); + } + let table_name = filename + .rsplit_once('.') + .map(|(stem, _)| stem) + .filter(|stem| !stem.is_empty()) + .context(error::InvalidTableNameSnafu { table_name: path })? + .to_string(); + Ok((table_name, location)) +} + +impl StatementExecutor { + /// Capture each selected data table once. Views, temporary and Metric physical + /// tables do not have data outputs. `None` selects the whole schema. + pub async fn capture_database_export_tables( + &self, + req: &CopyDatabaseRequest, + names: Option<&[String]>, + ctx: &QueryContextRef, + ) -> Result> { + let mut names = match names { + Some(names) => names.to_vec(), + None => self + .catalog_manager + .table_names(&req.catalog_name, &req.schema_name, Some(ctx)) + .await + .context(error::CatalogSnafu)?, + }; + names.sort(); + let mut tables = Vec::new(); + for name in names { + let table = self + .get_table(&TableReference::full( + &req.catalog_name, + &req.schema_name, + &name, + )) + .await?; + let info = table.table_info(); + if table.table_type() == TableType::Base + && (info.meta.engine != METRIC_ENGINE_NAME + || info + .meta + .options + .extra_options + .contains_key(LOGICAL_TABLE_METADATA_KEY)) + { + tables.push(table); + } + } + Ok(tables) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn directory_url_components_cannot_capture_output_names() { + for location in [ + "file:///copy/fresh?attempt=/", + "file:///copy/fresh#attempt/", + "s3://bucket/fresh?attempt=/", + "s3://bucket/fresh#attempt/", + "file:///copy/fresh", + ] { + assert!(validate_database_directory(location).is_err(), "{location}"); + } + for location in ["/copy/fresh/", "file:///copy/fresh/", "s3://bucket/fresh/"] { + validate_database_directory(location).unwrap(); + } + } + + #[cfg(windows)] + #[test] + fn windows_directory_names_are_literal_paths() { + for location in ["C:/copy/fresh#1/", r"C:\copy\fresh#1\"] { + validate_database_directory(location).unwrap(); + } + assert!(validate_database_directory("C:/copy/fresh#1").is_err()); + } + + #[tokio::test] + async fn output_locations_resolve_to_writer_keys() { + use common_datasource::object_store::{LocalFileAccess, build_backend_for_write_with_path}; + + let dir = common_test_util::temp_dir::create_temp_dir("database_export_paths"); + let access = LocalFileAccess::sandboxed(dir.path()).unwrap(); + let file_url = Url::from_directory_path(dir.path()).unwrap().to_string(); + let connection = HashMap::from([ + ("region".into(), "us-east-1".into()), + ("access_key_id".into(), "test-key".into()), + ("secret_access_key".into(), "test-secret".into()), + ]); + for directory in [ + format!("{}/", dir.path().display()), + file_url, + "s3://export-bucket/data/".into(), + ] { + for name in ["a#b", "a:b"] { + if cfg!(windows) && name.contains(':') && !directory.starts_with("s3:") { + continue; + } + let file = DatabaseExportFile::new(&directory, name, ".parquet").unwrap(); + let backend = + build_backend_for_write_with_path(&file.location, &connection, &access) + .await + .unwrap(); + assert_eq!(backend.object_path.as_deref(), Some(file.path.as_str())); + let (table_name, input_location) = + database_import_source(&directory, &file.path).unwrap(); + assert_eq!(table_name, name); + assert_eq!(input_location, file.location); + let (table_name, nested_location) = + database_import_source(&directory, &format!("nested/{}", file.path)).unwrap(); + assert_eq!(table_name, name); + assert_eq!( + nested_location, + file.location + .replace(&directory, &format!("{directory}nested/")) + ); + if !directory.starts_with("s3:") { + backend + .object_store + .write(&file.path, "test") + .await + .unwrap(); + assert_eq!( + std::fs::read(dir.path().join(format!("{name}.parquet"))).unwrap(), + b"test" + ); + } else { + assert_eq!( + file.path, + if name == "a#b" { + "a%23b.parquet" + } else { + "a:b.parquet" + } + ); + } + } + } + } + + #[test] + fn test_parse_parallelism_from_option_map() { + let options = HashMap::new(); + assert_eq!( + parse_parallelism_from_option_map(&options), + get_total_cpu_cores() + ); + + let options = HashMap::from([("parallelism".to_string(), "0".to_string())]); + assert_eq!(parse_parallelism_from_option_map(&options), 1); + } +} diff --git a/src/operator/src/statement/export_database.rs b/src/operator/src/statement/export_database.rs new file mode 100644 index 0000000000..c9c698a356 --- /dev/null +++ b/src/operator/src/statement/export_database.rs @@ -0,0 +1,400 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Request-scoped database export preparation. The frontend authorizes the entire +//! captured selection before preparation; these trusted methods do not check ACLs. + +use std::collections::{BTreeMap, HashSet}; +use std::future::Future; + +use common_datasource::file_format::Format; +use common_datasource::object_store::{FILE_SCHEMA, FS_SCHEMA, build_backend_for_write, parse_url}; +use common_meta::key::table_route::TableRouteValue; +use futures::StreamExt; +use futures::stream::FuturesUnordered; +use session::context::QueryContextRef; +use snafu::{OptionExt, ResultExt, ensure}; +use store_api::metric_engine_consts::{LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME}; +use table::TableRef; +use table::metadata::TableType; +use table::requests::{CopyDatabaseRequest, CopyDirection, CopyTableRequest}; +use tokio_util::sync::CancellationToken; + +use crate::error::{self, InvalidDatabaseExportSnafu, Result}; +use crate::statement::StatementExecutor; +use crate::statement::database_copy::{ + DatabaseExportFile, parse_parallelism_from_option_map, validate_database_directory, +}; +use crate::statement::export_logical_tables::{LogicalTableExport, LogicalTableExportLimits}; + +/// A validated request-scoped selection, not a metadata snapshot or an ACL token. +pub struct PreparedDatabaseExport { + request: CopyDatabaseRequest, + jobs: Vec, +} + +impl PreparedDatabaseExport { + /// Inspect physical grouping through the integration testing adapter. + #[cfg(feature = "testing")] + pub fn job_count_for_test(&self) -> usize { + self.jobs.len() + } +} + +enum DatabaseExportJob { + Ordinary { + table: TableRef, + output: DatabaseExportFile, + }, + Metric(LogicalTableExport), +} + +impl StatementExecutor { + /// Validate all jobs and destinations before executing any query or writer. + /// Callers must authorize every captured table before calling this method. + pub async fn prepare_database_export( + &self, + req: CopyDatabaseRequest, + tables: Vec, + ) -> Result { + validate_database_directory(&req.location)?; + let format = Format::try_from(&req.with).context(error::ParseFileFormatSnafu)?; + ensure!( + matches!(format, Format::Parquet(_)), + error::UnsupportedFormatSnafu { format } + ); + let (scheme, _, _) = parse_url(&req.location).context(error::BuildBackendSnafu)?; + let local = + scheme.eq_ignore_ascii_case(FS_SCHEMA) || scheme.eq_ignore_ascii_case(FILE_SCHEMA); + let mut filenames = HashSet::new(); + let mut logical = Vec::new(); + let mut jobs = Vec::new(); + for table in tables { + let info = table.table_info(); + let name = &info.name; + let output = DatabaseExportFile::new(&req.location, name, ".parquet")?; + ensure!( + filenames.insert(if local { + output.path.to_ascii_lowercase() + } else { + output.path.clone() + }), + InvalidDatabaseExportSnafu { + reason: format!("duplicate output name: {name}") + } + ); + ensure!( + info.catalog_name == req.catalog_name + && info.schema_name == req.schema_name + && table.table_type() == TableType::Base, + InvalidDatabaseExportSnafu { + reason: "expected base tables in the selected schema" + } + ); + if info.meta.engine == METRIC_ENGINE_NAME { + ensure!( + info.meta + .options + .extra_options + .contains_key(LOGICAL_TABLE_METADATA_KEY), + InvalidDatabaseExportSnafu { + reason: "expected a Metric logical table" + } + ); + logical.push(table); + } else { + jobs.push(DatabaseExportJob::Ordinary { table, output }); + } + } + let ids = logical + .iter() + .map(|t| t.table_info().table_id()) + .collect::>(); + let routes = self + .table_metadata_manager + .table_route_manager() + .table_route_storage() + .batch_get(&ids) + .await + .context(error::TableMetadataManagerSnafu)?; + let mut groups = BTreeMap::<_, Vec>::new(); + for (table, route) in logical.into_iter().zip(routes) { + let Some(TableRouteValue::Logical(route)) = route else { + return InvalidDatabaseExportSnafu { + reason: format!( + "missing or non-logical route for {}", + table.table_info().table_id() + ), + } + .fail(); + }; + groups + .entry(route.physical_table_id()) + .or_default() + .push(table); + } + let physical_ids = groups.keys().copied().collect::>(); + let mut physical = self + .catalog_manager + .tables_by_ids(&req.catalog_name, &req.schema_name, &physical_ids) + .await + .context(error::CatalogSnafu)? + .into_iter() + .map(|t| (t.table_info().table_id(), t)) + .collect::>(); + for (id, tables) in groups { + let table = physical + .remove(&id) + .with_context(|| InvalidDatabaseExportSnafu { + reason: format!("missing physical table {id} in the selected schema"), + })?; + jobs.push(DatabaseExportJob::Metric( + LogicalTableExport::try_new_in_directory(table, &tables, &req.location)?, + )); + } + build_backend_for_write(&req.location, &req.connection, &self.local_file_access) + .await + .context(error::BuildBackendSnafu)?; + Ok(PreparedDatabaseExport { request: req, jobs }) + } +} + +/// Returned only after every started job has closed or drained its owned I/O. +#[derive(Debug)] +pub struct DatabaseExportSummary { + pub rows: usize, + pub output_files: Vec, +} + +impl StatementExecutor { + /// Execute with one job budget. The caller must keep this future alive until + /// it returns, including after cooperative cancellation. Closed files remain + /// owned by the caller's attempt; this method does not publish completion. + pub async fn export_database( + &self, + plan: PreparedDatabaseExport, + cancellation: &CancellationToken, + ctx: QueryContextRef, + ) -> Result { + let mut output_files = Vec::new(); + for job in &plan.jobs { + match job { + DatabaseExportJob::Ordinary { output, .. } => { + output_files.push(output.location.clone()) + } + DatabaseExportJob::Metric(unit) => { + output_files.extend(unit.output_files().map(|file| file.location.clone())) + } + } + } + output_files.sort(); + let req = &plan.request; + let rows = run_database_export_jobs( + plan.jobs, + parse_parallelism_from_option_map(&req.with), + cancellation, + |job, token| { + let ctx = ctx.clone(); + async move { + match job { + DatabaseExportJob::Metric(unit) => self + .export_logical_tables( + &unit, + &req.location, + &req.connection, + req.time_range.as_ref(), + LogicalTableExportLimits::default(), + &token, + ctx, + ) + .await + .map(|summary| summary.rows), + DatabaseExportJob::Ordinary { table, output } => { + let info = table.table_info(); + let copy = CopyTableRequest { + catalog_name: info.catalog_name.clone(), + schema_name: info.schema_name.clone(), + table_name: info.name.clone(), + location: output.location, + with: req.with.clone(), + connection: req.connection.clone(), + pattern: None, + direction: CopyDirection::Export, + timestamp_range: req.time_range, + limit: None, + }; + self.copy_captured_table_to(table, copy, ctx).await + } + } + } + }, + ) + .await?; + Ok(DatabaseExportSummary { rows, output_files }) + } +} + +async fn run_database_export_jobs>>( + jobs: impl IntoIterator, + parallelism: usize, + cancellation: &CancellationToken, + mut run: impl FnMut(J, CancellationToken) -> F, +) -> Result { + let token = CancellationToken::new(); + let mut jobs = jobs.into_iter(); + let mut active = FuturesUnordered::new(); + let mut first_error = None; + let mut rows = 0; + loop { + while first_error.is_none() + && !cancellation.is_cancelled() + && active.len() < parallelism.max(1) + { + let Some(job) = jobs.next() else { break }; + active.push(run(job, token.clone())); + } + if first_error.is_none() && cancellation.is_cancelled() { + first_error = Some(error::DatabaseExportCancelledSnafu.build()); + token.cancel(); + } + if active.is_empty() { + break; + } + let result = tokio::select! { + biased; + _ = cancellation.cancelled(), if first_error.is_none() => { + first_error = Some(error::DatabaseExportCancelledSnafu.build()); + token.cancel(); + continue; + } + result = active.next() => result, + }; + match result { + Some(Ok(count)) => rows += count, + Some(Err(err)) if first_error.is_none() => { + first_error = Some(err); + token.cancel(); + } + Some(Err(err)) => common_telemetry::warn!(err; "Failed to drain database export job"), + None => break, + } + } + match first_error { + Some(err) => Err(err), + None => Ok(rows), + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + + use tokio::sync::{Semaphore, mpsc}; + + use super::*; + + #[tokio::test] + async fn bounded_admission_and_drain() { + for cancel in [false, true] { + let cancellation = CancellationToken::new(); + let start_error = Arc::new(Semaphore::new(0)); + let finish_io = Arc::new(Semaphore::new(0)); + let (started, mut receiver) = mpsc::unbounded_channel(); + let finished = Arc::new(AtomicUsize::new(0)); + let task = tokio::spawn({ + let cancellation = cancellation.clone(); + let start_error = start_error.clone(); + let finish_io = finish_io.clone(); + let finished = finished.clone(); + async move { + run_database_export_jobs(0..4, 2, &cancellation, |job, token| { + let started = started.clone(); + let start_error = start_error.clone(); + let finish_io = finish_io.clone(); + let finished = finished.clone(); + async move { + started.send(job).unwrap(); + if job == 0 { + if cancel { + token.cancelled().await; + } else { + start_error.acquire().await.unwrap().forget(); + } + return InvalidDatabaseExportSnafu { + reason: "first error", + } + .fail(); + } + // Ordinary COPY continues its I/O even when Metric jobs cancel. + token.cancelled().await; + started.send(10).unwrap(); + finish_io.acquire().await.unwrap().forget(); + finished.fetch_add(1, Ordering::SeqCst); + InvalidDatabaseExportSnafu { + reason: "drain error", + } + .fail() + } + }) + .await + } + }); + assert_eq!(receiver.recv().await, Some(0)); + assert_eq!(receiver.recv().await, Some(1)); + assert!(receiver.try_recv().is_err()); + if cancel { + cancellation.cancel(); + } else { + start_error.add_permits(1); + } + assert_eq!(receiver.recv().await, Some(10)); + assert!(!task.is_finished()); + finish_io.add_permits(1); + let err = task.await.unwrap().unwrap_err(); + if cancel { + assert!(matches!(err, error::Error::DatabaseExportCancelled { .. })); + } else { + assert!( + matches!(err, error::Error::InvalidDatabaseExport { reason } if reason == "first error") + ); + } + assert_eq!(finished.load(Ordering::SeqCst), 1); + assert_eq!(receiver.recv().await, None); + } + } + + #[tokio::test] + async fn cancellation_before_admission_and_successful_refill() { + let token = CancellationToken::new(); + token.cancel(); + let result = run_database_export_jobs(0..4, 2, &token, |_, _| async { + panic!("cancelled job admitted") + }) + .await; + assert!(matches!( + result, + Err(error::Error::DatabaseExportCancelled { .. }) + )); + let started = AtomicUsize::new(0); + let result = run_database_export_jobs(0..7, 2, &CancellationToken::new(), |job, _| { + started.fetch_add(1, Ordering::SeqCst); + async move { Ok(job) } + }) + .await + .unwrap(); + assert_eq!(result, 21); + assert_eq!(started.load(Ordering::SeqCst), 7); + } +} diff --git a/src/operator/src/statement/export_logical_tables.rs b/src/operator/src/statement/export_logical_tables.rs index 7c38af14ba..7547a3f54d 100644 --- a/src/operator/src/statement/export_logical_tables.rs +++ b/src/operator/src/statement/export_logical_tables.rs @@ -51,6 +51,7 @@ use tokio_util::sync::CancellationToken; use crate::error::{self, InvalidLogicalTableExportSnafu, LogicalTableExportResourceSnafu, Result}; use crate::statement::StatementExecutor; +use crate::statement::database_copy::DatabaseExportFile; /// Export preprocessing and per-file limits. Query memory and spill remain /// governed by the query engine. @@ -101,7 +102,7 @@ pub struct LogicalTableExport { } struct LogicalTableProjection { - name: String, + output: DatabaseExportFile, schema: SchemaRef, projection: Vec, } @@ -110,6 +111,14 @@ impl LogicalTableExport { /// Capture schemas from selected Metric table references. /// Export validates their physical-table association against table routes. pub fn try_new(physical: TableRef, tables: &[TableRef]) -> Result { + Self::try_new_in_directory(physical, tables, "") + } + + pub(crate) fn try_new_in_directory( + physical: TableRef, + tables: &[TableRef], + directory: &str, + ) -> Result { let physical_info = physical.table_info(); ensure!( physical_info.meta.engine == METRIC_ENGINE_NAME @@ -134,12 +143,6 @@ impl LogicalTableExport { for table in tables { let info = table.table_info(); let name = &info.name; - ensure!( - !name.contains('/') && !name.contains('\\'), - InvalidLogicalTableExportSnafu { - reason: "logical table names must not contain path separators" - } - ); ensure!( info.catalog_name == physical_info.catalog_name && info.schema_name == physical_info.schema_name @@ -187,7 +190,7 @@ impl LogicalTableExport { .insert( info.table_id(), LogicalTableProjection { - name: name.clone(), + output: DatabaseExportFile::new(directory, name, ".parquet")?, schema, projection: indices, } @@ -216,6 +219,10 @@ impl LogicalTableExport { }) } + pub(crate) fn output_files(&self) -> impl Iterator { + self.logical_tables.values().map(|table| &table.output) + } + async fn validate_table_routes(&self, manager: &TableRouteManager) -> Result<()> { let table_ids = self.logical_tables.keys().copied().collect::>(); let routes = manager @@ -485,7 +492,7 @@ impl ActiveWriter { store: &ObjectStore, limits: LogicalTableExportLimits, ) -> Result { - let path = format!("{}.parquet", table.name); + let path = table.output.path.clone(); ensure!( !store .exists(&path) diff --git a/src/operator/src/statement/export_logical_tables/tests.rs b/src/operator/src/statement/export_logical_tables/tests.rs index a85418fa2e..4576c1e395 100644 --- a/src/operator/src/statement/export_logical_tables/tests.rs +++ b/src/operator/src/statement/export_logical_tables/tests.rs @@ -296,8 +296,6 @@ fn validates_selected_schemas_and_projects_only_selected_columns() { LogicalTableExport::try_new(unit.physical_table.clone(), &[selected.clone(), selected]) .is_err() ); - let unsafe_name = table(1030, "a/b", vec![], false); - assert!(LogicalTableExport::try_new(unit.physical_table.clone(), &[unsafe_name]).is_err()); let wrong_type = table( 1030, "wrong", diff --git a/src/query/src/dist_plan/planner.rs b/src/query/src/dist_plan/planner.rs index 130a123fc6..b2a2ade9b4 100644 --- a/src/query/src/dist_plan/planner.rs +++ b/src/query/src/dist_plan/planner.rs @@ -39,6 +39,7 @@ use partition::manager::{PartitionRuleManagerRef, create_partitions_from_region_ use session::context::QueryContext; use snafu::{OptionExt, ResultExt}; use store_api::storage::RegionId; +use table::TableRef; use table::metadata::TableInfo; pub use table::metadata::TableType; use table::table::adapter::DfTableProviderAdapter; @@ -238,7 +239,7 @@ impl ExtensionPlanner for DistExtensionPlanner { impl DistExtensionPlanner { /// Extract fully resolved table name from logical plan fn extract_full_table_name(plan: &LogicalPlan) -> Result> { - let mut extractor = TableNameExtractor::default(); + let mut extractor = TableScanExtractor::default(); let _ = plan.visit(&mut extractor)?; Ok(extractor.table_name) } @@ -248,19 +249,25 @@ impl DistExtensionPlanner { table_name: &TableName, logical_plan: &LogicalPlan, ) -> Result> { - let table = self - .catalog_manager - .table( - &table_name.catalog_name, - &table_name.schema_name, - &table_name.table_name, - None, - ) - .await - .context(CatalogSnafu)? - .with_context(|| TableNotFoundSnafu { - table: table_name.to_string(), - })?; + let mut extractor = TableScanExtractor::default(); + let _ = logical_plan.visit(&mut extractor)?; + // Resolving by name again could bind an authorized scan to a replacement table. + let table = match extractor.captured_table { + Some(table) => table, + None => self + .catalog_manager + .table( + &table_name.catalog_name, + &table_name.schema_name, + &table_name.table_name, + None, + ) + .await + .context(CatalogSnafu)? + .with_context(|| TableNotFoundSnafu { + table: table_name.to_string(), + })?, + }; let table_info = table.table_info(); let (physical_table_id, physical_table_route) = self @@ -464,13 +471,14 @@ fn partition_column_types(table_info: &TableInfo) -> Vec<(String, ConcreteDataTy .collect() } -/// Visitor to extract table name from logical plan (TableScan node) +/// Extract the scan name and captured table identity from a logical plan. #[derive(Default)] -struct TableNameExtractor { +struct TableScanExtractor { pub table_name: Option, + captured_table: Option, } -impl TreeNodeVisitor<'_> for TableNameExtractor { +impl TreeNodeVisitor<'_> for TableScanExtractor { type Node = LogicalPlan; fn f_down(&mut self, node: &Self::Node) -> Result { @@ -482,6 +490,7 @@ impl TreeNodeVisitor<'_> for TableNameExtractor { .downcast_ref::() { if provider.table().table_type() == TableType::Base { + self.captured_table = Some(provider.table()); let info = provider.table().table_info(); self.table_name = Some(TableName::new( info.catalog_name.clone(), @@ -772,6 +781,16 @@ mod tests { ] } + #[tokio::test] + async fn region_routing_uses_captured_table() { + let (mut planner, plan, table_name) = planner_and_plan(vec![0], vec![None]).await; + planner.catalog_manager = MemoryCatalogManager::with_default_setup(); + assert_eq!( + vec![RegionId::new(LOGICAL_TABLE_ID, 1)], + planner.get_regions(&table_name, &plan).await.unwrap() + ); + } + #[tokio::test] async fn logical_table_pruning_uses_physical_partition_datatypes() { let (planner, plan, table_name) = diff --git a/tests-integration/tests/export_logical_tables.rs b/tests-integration/tests/export_logical_tables.rs index 4aeb745b88..9b4ea97a4a 100644 --- a/tests-integration/tests/export_logical_tables.rs +++ b/tests-integration/tests/export_logical_tables.rs @@ -54,40 +54,52 @@ async fn values(instance: &Arc, query: &str) -> Vec, + physical: &str, + encoding: &str, +) -> ([String; 3], Vec, String) { + sql(instance, &format!("CREATE TABLE {physical} (ts TIMESTAMP TIME INDEX, val DOUBLE, host STRING PRIMARY KEY) PARTITION ON COLUMNS (host) (host < 'm', host >= 'm') ENGINE=metric WITH (physical_metric_table='', primary_key_encoding='{encoding}')")).await; + for (suffix, extra, key) in [ + ("cpu.v1", "zone_tag STRING,", ", zone_tag"), + ("requests", "service_tag STRING,", ", service_tag"), + ("empty", "", ""), + ] { + let name = format!("{physical}_{suffix}"); + sql(instance, &format!("CREATE TABLE \"{name}\" (host STRING, {extra} val DOUBLE, ts TIMESTAMP TIME INDEX, PRIMARY KEY(host{key})) ENGINE=metric WITH (on_physical_table='{physical}')")).await; + if suffix != "empty" { + sql(instance, &format!("INSERT INTO \"{name}\" (host,val,ts) VALUES ('a',1,1),('a',NULL,2),('z',3,3),('z',4,4)")).await; + } + } + // This live but unselected table models rows outside the routing whitelist. + sql(instance, &format!("CREATE TABLE {physical}_excluded (host STRING, huge_tag STRING, val DOUBLE, ts TIMESTAMP TIME INDEX, PRIMARY KEY(host, huge_tag)) ENGINE=metric WITH (on_physical_table='{physical}')")).await; + sql( + instance, + &format!( + "INSERT INTO {physical}_excluded (host, huge_tag, val, ts) VALUES ('z','ignore',9,2)" + ), + ) + .await; + let names = ["cpu.v1", "requests", "empty"].map(|suffix| format!("{physical}_{suffix}")); + let tables = vec![ + table(instance, &names[0]).await, + table(instance, &names[1]).await, + table(instance, &names[2]).await, + ]; + let renamed = format!("renamed_{physical}"); + sql( + instance, + &format!("ALTER TABLE {physical} RENAME {renamed}"), + ) + .await; + (names, tables, renamed) +} + async fn roundtrip(instance: &Arc) { let destination = tempfile::tempdir_in(common_test_util::find_workspace_path(".")).unwrap(); for (physical, encoding) in [("phy", "dense"), ("other_phy", "sparse")] { - sql(instance, &format!("CREATE TABLE {physical} (ts TIMESTAMP TIME INDEX, val DOUBLE, host STRING PRIMARY KEY) PARTITION ON COLUMNS (host) (host < 'm', host >= 'm') ENGINE=metric WITH (physical_metric_table='', primary_key_encoding='{encoding}')")).await; - for (suffix, extra, key) in [ - ("cpu.v1", "zone_tag STRING,", ", zone_tag"), - ("requests", "service_tag STRING,", ", service_tag"), - ("empty", "", ""), - ] { - let name = format!("{physical}_{suffix}"); - sql(instance, &format!("CREATE TABLE \"{name}\" (host STRING, {extra} val DOUBLE, ts TIMESTAMP TIME INDEX, PRIMARY KEY(host{key})) ENGINE=metric WITH (on_physical_table='{physical}')")).await; - if suffix != "empty" { - sql(instance, &format!("INSERT INTO \"{name}\" (host,val,ts) VALUES ('a',1,1),('a',NULL,2),('z',3,3),('z',4,4)")).await; - } - } - // This live but unselected table models rows outside the routing whitelist. - sql(instance, &format!("CREATE TABLE {physical}_excluded (host STRING, huge_tag STRING, val DOUBLE, ts TIMESTAMP TIME INDEX, PRIMARY KEY(host, huge_tag)) ENGINE=metric WITH (on_physical_table='{physical}')")).await; - sql( - instance, - &format!("INSERT INTO {physical}_excluded (host, huge_tag, val, ts) VALUES ('z','ignore',9,2)"), - ) - .await; - let names = ["cpu.v1", "requests", "empty"].map(|suffix| format!("{physical}_{suffix}")); - let tables = vec![ - table(instance, &names[0]).await, - table(instance, &names[1]).await, - table(instance, &names[2]).await, - ]; - let renamed = format!("renamed_{physical}"); - sql( - instance, - &format!("ALTER TABLE {physical} RENAME {renamed}"), - ) - .await; + let (names, tables, renamed) = + create_metric_export_source_tables(instance, physical, encoding).await; let unit = LogicalTableExport::try_new(table(instance, &renamed).await, &tables).unwrap(); let range = TimestampRange::new(Timestamp::new_millisecond(2), Timestamp::new_millisecond(4)) @@ -196,3 +208,531 @@ async fn physical_export_distributed_roundtrip() { .await; roundtrip(cluster.fe_instance()).await; } + +fn database_export_request(directory: &std::path::Path) -> table::requests::CopyDatabaseRequest { + table::requests::CopyDatabaseRequest { + catalog_name: "greptime".into(), + schema_name: "public".into(), + location: format!("{}/", directory.display()), + with: [ + ("format".into(), "parquet".into()), + ("parallelism".into(), "2".into()), + ] + .into(), + connection: Default::default(), + time_range: Some( + TimestampRange::new(Timestamp::new_millisecond(2), Timestamp::new_millisecond(4)) + .unwrap(), + ), + } +} + +async fn database_export_roundtrip(instance: &Arc) { + let destination = tempfile::tempdir_in(common_test_util::find_workspace_path(".")).unwrap(); + let (first_logical_table_names, _, renamed_physical_table) = + create_metric_export_source_tables(instance, "db_a", "dense").await; + let (second_logical_table_names, _, _) = + create_metric_export_source_tables(instance, "db_b", "sparse").await; + sql( + instance, + "CREATE TABLE audit (host STRING, val DOUBLE, ts TIMESTAMP TIME INDEX, PRIMARY KEY(host))", + ) + .await; + sql( + instance, + "INSERT INTO audit VALUES ('a',1,1),('z',NULL,2),('z',4,3)", + ) + .await; + sql(instance, "CREATE VIEW dashboard AS SELECT * FROM audit").await; + let selected = vec![ + first_logical_table_names[0].clone(), + first_logical_table_names[2].clone(), + second_logical_table_names[1].clone(), + "audit".into(), + ]; + let mut names = selected.clone(); + names.extend([renamed_physical_table, "dashboard".into()]); + let req = database_export_request(&destination.path().join("data")); + let executor = instance.statement_executor(); + let captured = executor + .capture_database_export_tables(&req, None, &QueryContext::arc()) + .await + .unwrap(); + assert_eq!(captured.len(), 9); + let captured = executor + .capture_database_export_tables(&req, Some(&names), &QueryContext::arc()) + .await + .unwrap(); + assert_eq!(captured.len(), 4); + let plan = executor + .prepare_database_export(req.clone(), captured) + .await + .unwrap(); + assert_eq!(plan.job_count_for_test(), 3); + assert_eq!( + std::fs::read_dir(destination.path().join("data")) + .unwrap() + .count(), + 0 + ); + let summary = instance + .export_database_for_test( + req, + Some(&names), + &CancellationToken::new(), + QueryContext::arc(), + ) + .await + .unwrap(); + assert_eq!(summary.rows, 6); + let expected = selected + .iter() + .map(|name| { + destination + .path() + .join("data") + .join(format!("{name}.parquet")) + .to_str() + .unwrap() + .to_string() + }) + .collect::>(); + assert_eq!( + summary + .output_files + .into_iter() + .collect::>(), + expected + ); + assert_eq!( + std::fs::read_dir(destination.path().join("data")) + .unwrap() + .count(), + 4 + ); + sql(instance, "CREATE TABLE restored_phy (ts TIMESTAMP TIME INDEX, val DOUBLE, host STRING PRIMARY KEY) ENGINE=metric WITH (physical_metric_table='')").await; + for (index, name) in selected.iter().enumerate() { + let restored = format!("restored_{index}"); + if index < 3 { + let (extra, key) = [ + ("zone_tag STRING,", ", zone_tag"), + ("", ""), + ("service_tag STRING,", ", service_tag"), + ][index]; + sql(instance, &format!("CREATE TABLE {restored} (host STRING, {extra} val DOUBLE, ts TIMESTAMP TIME INDEX, PRIMARY KEY(host{key})) ENGINE=metric WITH (on_physical_table='restored_phy')")).await; + } else { + sql( + instance, + &format!("CREATE TABLE {restored} LIKE \"{name}\""), + ) + .await; + } + sql( + instance, + &format!( + "COPY {restored} FROM '{}/data/{name}.parquet' WITH (FORMAT='parquet')", + destination.path().display() + ), + ) + .await; + assert_eq!( + table(instance, name).await.schema().column_schemas(), + table(instance, &restored).await.schema().column_schemas() + ); + assert_eq!( + values( + instance, + &format!("SELECT * FROM {restored} ORDER BY host, ts") + ) + .await, + values( + instance, + &format!("SELECT * FROM \"{name}\" WHERE ts >= 2 AND ts < 4 ORDER BY host, ts") + ) + .await + ); + } + let ordinary_req = database_export_request(&destination.path().join("captured")); + let captured = executor + .capture_database_export_tables( + &ordinary_req, + Some(&["audit".into()]), + &QueryContext::arc(), + ) + .await + .unwrap(); + let plan = executor + .prepare_database_export(ordinary_req, captured) + .await + .unwrap(); + sql(instance, "ALTER TABLE audit RENAME original_audit").await; + sql(instance, "CREATE TABLE audit LIKE original_audit").await; + sql(instance, "INSERT INTO audit VALUES ('replacement',99,2)").await; + let result = executor + .export_database(plan, &CancellationToken::new(), QueryContext::arc()) + .await + .unwrap(); + assert_eq!(result.rows, 2); + sql( + instance, + "CREATE TABLE captured_restore LIKE original_audit", + ) + .await; + sql( + instance, + &format!( + "COPY captured_restore FROM '{}/captured/audit.parquet' WITH (FORMAT='parquet')", + destination.path().display() + ), + ) + .await; + assert_eq!( + values(instance, "SELECT * FROM captured_restore ORDER BY host,ts").await, + values( + instance, + "SELECT * FROM original_audit WHERE ts >= 2 AND ts < 4 ORDER BY host,ts" + ) + .await + ); + for suffix in ["?attempt=/", "#attempt/"] { + let path = destination.path().join("invalid_destination"); + let mut req = database_export_request(&path); + req.location = format!("{}{suffix}", url::Url::from_file_path(&path).unwrap()); + req.with.insert("parallelism".into(), "1".into()); + let result = instance + .export_database_for_test( + req, + Some(&["audit".into(), "original_audit".into()]), + &CancellationToken::new(), + QueryContext::arc(), + ) + .await; + assert!(matches!( + result, + Err(frontend::error::Error::TableOperation { + source: operator::error::Error::InvalidCopyDatabasePath { .. }, + .. + }) + )); + assert!(!path.exists()); + assert!( + tests_integration::test_util::try_execute_sql( + instance, + &format!( + "COPY DATABASE public TO '{}{suffix}' WITH (FORMAT='parquet')", + url::Url::from_file_path(&path).unwrap() + ) + ) + .await + .is_err() + ); + assert!(!path.exists()); + } + let token = CancellationToken::new(); + token.cancel(); + let req = database_export_request(&destination.path().join("cancelled")); + let result = instance + .export_database_for_test(req, Some(&names), &token, QueryContext::arc()) + .await; + assert!(matches!( + result, + Err(frontend::error::Error::TableOperation { + source: operator::error::Error::DatabaseExportCancelled { .. }, + .. + }) + )); + assert!(!destination.path().join("cancelled").exists()); +} + +#[tokio::test(flavor = "multi_thread")] +async fn database_export_standalone_roundtrip() { + let standalone = GreptimeDbStandaloneBuilder::new("database_export") + .build() + .await; + database_export_roundtrip(standalone.fe_instance()).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn database_export_distributed_roundtrip() { + let cluster = GreptimeDbClusterBuilder::new("database_export") + .await + .with_datanodes(2) + .with_local_file_access( + common_datasource::object_store::LocalFileAccess::sandboxed( + common_test_util::find_workspace_path("."), + ) + .unwrap(), + ) + .build(false) + .await; + database_export_roundtrip(cluster.fe_instance()).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn database_export_rejects_invalid_members_before_output() { + use common_meta::key::table_info::TableInfoKey; + use common_meta::key::table_route::{TableRouteKey, TableRouteValue}; + use common_meta::key::{MetadataKey, MetadataValue}; + use common_meta::rpc::store::PutRequest; + + let standalone = GreptimeDbStandaloneBuilder::new("invalid_database_export") + .build() + .await; + let instance = standalone.fe_instance(); + let (_, tables, physical) = + create_metric_export_source_tables(instance, "invalid_phy", "dense").await; + let executor = instance.statement_executor(); + let destination = tempfile::tempdir_in(common_test_util::find_workspace_path(".")).unwrap(); + for name in ["Foo", "foo"] { + sql( + instance, + &format!("CREATE TABLE \"{name}\" (ts TIMESTAMP TIME INDEX)"), + ) + .await; + } + let case_tables = vec![table(instance, "Foo").await, table(instance, "foo").await]; + assert_ne!( + case_tables[0].table_info().table_id(), + case_tables[1].table_info().table_id() + ); + let case_path = destination.path().join("case_aliases"); + let mut case_req = database_export_request(&case_path); + case_req.with.insert("parallelism".into(), "1".into()); + let result = instance + .export_database_for_test( + case_req.clone(), + Some(&["Foo".into(), "foo".into()]), + &CancellationToken::new(), + QueryContext::arc(), + ) + .await; + assert!(matches!(result, + Err(frontend::error::Error::TableOperation { + source: operator::error::Error::InvalidDatabaseExport { reason }, .. + }) if reason == "duplicate output name: foo")); + assert!(!case_path.exists()); + case_req.location = "s3://export-bucket/data/".into(); + case_req.connection.extend([ + ("region".into(), "us-east-1".into()), + ("access_key_id".into(), "test-key".into()), + ("secret_access_key".into(), "test-secret".into()), + ]); + let plan = executor + .prepare_database_export(case_req, case_tables) + .await + .unwrap(); + assert_eq!(plan.job_count_for_test(), 2); + let req = database_export_request(&destination.path().join("data")); + let key = TableRouteKey::new(tables[2].table_info().table_id()).to_bytes(); + let original = standalone + .kv_backend + .get(&key) + .await + .unwrap() + .unwrap() + .value; + for route in [ + None, + Some(TableRouteValue::physical(vec![])), + Some(TableRouteValue::logical(u32::MAX)), + ] { + if let Some(route) = route { + standalone + .kv_backend + .put( + PutRequest::new() + .with_key(key.clone()) + .with_value(route.try_as_raw_value().unwrap()), + ) + .await + .unwrap(); + } else { + standalone.kv_backend.delete(&key, false).await.unwrap(); + } + let result = executor + .prepare_database_export(req.clone(), tables.clone()) + .await; + assert!(matches!( + result, + Err(operator::error::Error::InvalidDatabaseExport { .. }) + )); + assert!(!destination.path().join("data").exists()); + } + standalone + .kv_backend + .put(PutRequest::new().with_key(key.clone()).with_value(original)) + .await + .unwrap(); + let plan = executor + .prepare_database_export(req.clone(), tables.clone()) + .await + .unwrap(); + // Preparation does not waive PR3's membership revalidation before a scan. + standalone.kv_backend.delete(&key, false).await.unwrap(); + let result = executor + .export_database(plan, &CancellationToken::new(), QueryContext::arc()) + .await; + assert!(matches!( + result, + Err(operator::error::Error::InvalidLogicalTableExport { .. }) + )); + assert_eq!( + std::fs::read_dir(destination.path().join("data")) + .unwrap() + .count(), + 0 + ); + let physical_id = table(instance, &physical).await.table_info().table_id(); + let route = TableRouteValue::logical(physical_id) + .try_as_raw_value() + .unwrap(); + standalone + .kv_backend + .put(PutRequest::new().with_key(key).with_value(route)) + .await + .unwrap(); + standalone + .kv_backend + .delete(&TableInfoKey::new(physical_id).to_bytes(), false) + .await + .unwrap(); + let result = executor + .prepare_database_export(req.clone(), tables.clone()) + .await; + assert!(matches!( + result, + Err(operator::error::Error::InvalidDatabaseExport { .. }) + )); + let duplicate_name = tables[0].table_info().name.clone(); + let result = executor + .prepare_database_export(req, vec![tables[0].clone(), tables[0].clone()]) + .await; + assert!(matches!(result, + Err(operator::error::Error::InvalidDatabaseExport { reason }) + if reason == format!("duplicate output name: {duplicate_name}"))); + assert_eq!( + std::fs::read_dir(destination.path().join("data")) + .unwrap() + .count(), + 0 + ); +} + +#[tokio::test(flavor = "multi_thread")] +async fn database_export_preserves_valid_table_names() { + let standalone = GreptimeDbStandaloneBuilder::new("database_export_names") + .build() + .await; + let instance = standalone.fe_instance(); + let destination = tempfile::tempdir_in(common_test_util::find_workspace_path(".")).unwrap(); + sql(instance, "CREATE TABLE names_physical (ts TIMESTAMP TIME INDEX, val DOUBLE, host STRING PRIMARY KEY) ENGINE=metric WITH (physical_metric_table='')").await; + let mut names = vec!["ordinary#b".to_string(), "metric#b".to_string()]; + if !cfg!(windows) { + names.extend(["ordinary:b".to_string(), "metric:b".to_string()]); + } + for name in &names { + let engine = if name.starts_with("metric") { + "ENGINE=metric WITH (on_physical_table='names_physical')" + } else { + "" + }; + sql(instance, &format!("CREATE TABLE \"{name}\" (ts TIMESTAMP TIME INDEX, val DOUBLE, host STRING PRIMARY KEY) {engine}")).await; + sql( + instance, + &format!("INSERT INTO \"{name}\" (ts,val,host) VALUES (1,42,'h')"), + ) + .await; + } + sql( + instance, + "CREATE VIEW names_view AS SELECT * FROM \"ordinary#b\"", + ) + .await; + sql(instance, "CREATE DATABASE names_restored").await; + for name in &names { + sql(instance, &format!("CREATE TABLE names_restored.\"{name}\" (ts TIMESTAMP TIME INDEX, val DOUBLE, host STRING PRIMARY KEY)")).await; + } + for (attempt, file_url, legacy) in [ + ("plain", false, false), + ("url", true, false), + ("legacy", true, true), + ] { + let directory = destination.path().join(attempt); + let mut req = database_export_request(&directory); + req.time_range = None; + if file_url { + req.location = url::Url::from_directory_path(&directory) + .unwrap() + .to_string(); + } + if legacy { + let output = sql( + instance, + &format!( + "COPY DATABASE public TO '{}' WITH (FORMAT='parquet')", + req.location + ), + ) + .await; + assert!(matches!(output.data, OutputData::AffectedRows(rows) if rows == names.len())); + } else { + let summary = instance + .export_database_for_test( + req.clone(), + None, + &CancellationToken::new(), + QueryContext::arc(), + ) + .await + .unwrap(); + assert_eq!(summary.rows, names.len()); + let expected = names + .iter() + .map(|name| { + let path = directory.join(format!("{name}.parquet")); + if file_url { + url::Url::from_file_path(path).unwrap().to_string() + } else { + path.to_str().unwrap().to_string() + } + }) + .collect::>(); + assert_eq!( + summary + .output_files + .into_iter() + .collect::>(), + expected + ); + } + for name in &names { + let path = directory.join(format!("{name}.parquet")); + assert!(path.is_file()); + } + assert_eq!(std::fs::read_dir(&directory).unwrap().count(), names.len()); + let output = sql( + instance, + &format!( + "COPY DATABASE names_restored FROM '{}' WITH (FORMAT='parquet')", + req.location + ), + ) + .await; + assert!(matches!(output.data, OutputData::AffectedRows(rows) if rows == names.len())); + for name in &names { + assert_eq!( + values( + instance, + &format!("SELECT ts, val, host FROM names_restored.\"{name}\"") + ) + .await, + values(instance, &format!("SELECT ts, val, host FROM \"{name}\"")).await, + ); + sql( + instance, + &format!("TRUNCATE TABLE names_restored.\"{name}\""), + ) + .await; + } + } +}