rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-09-23 13:45:42 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: prepare and execute database Metric exports (#9180)

Browse Source 
* feat: prepare and authorize captured database exports

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* feat: bound database export jobs and drain failures

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* test: validate database export identity and restore equivalence

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* fix: validate database export directory URLs

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* fix: preserve Windows export directory paths

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* fix: reject local database export filename aliases

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* refactor: consolidate database export planning policies

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* fix: restore escaped database export filenames

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* test: align database restore assertions with shared policies

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

* refactor: clarify database export boundaries and names

Signed-off-by: jeremyhi <fengjiachun@gmail.com>

---------

Signed-off-by: jeremyhi <fengjiachun@gmail.com>
This commit is contained in:
jeremyhi
2026-09-17 12:02:39 +00:00
committed by GitHub
parent ab21f684bf
commit dca01654bc
15 changed files with 1594 additions and 178 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+1
View File
@@ -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",
Cargo.toml
+1
View File
@@ -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"] }
src/frontend/src/instance.rs
+91
View File
@@ -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<Vec<PermissionTableTargets>>);
impl PermissionChecker for ExportAcl {
fn check_permission(
&self,
_: UserInfoRef,
req: PermissionReq,
) -> auth::error::Result<PermissionResp> {
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<PermissionResp> {
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::<PermissionCheckerRef>(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::<PermissionCheckerRef, _, _>(|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![
src/frontend/src/instance/export_database.rs
+113
View File
@@ -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<DatabaseExportSummary> {
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<DatabaseExportSummary> {
self.export_database(req, names, cancellation, ctx).await
}
async fn prepare_database_export(
&self,
req: CopyDatabaseRequest,
names: Option<&[String]>,
ctx: &QueryContextRef,
) -> Result<PreparedDatabaseExport> {
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::<PermissionCheckerRef>()
.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?)
}
}
src/operator/Cargo.toml
+2 -1
View File
@@ -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
src/operator/src/error.rs
+8
View File
@@ -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,
src/operator/src/statement.rs
+2
View File
@@ -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;
src/operator/src/statement/copy_database.rs
+27 -117
View File
@@ -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<String, String>) -> usize {
options
.get(PARALLELISM_KEY)
.and_then(|v| v.parse::<usize>().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<Output> {
// 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<Output> {
// 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<String> {
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::<HashSet<_>>();
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);
}
}
src/operator/src/statement/copy_table_to.rs
+13 -1
View File
@@ -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<usize> {
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<usize> {
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);
src/operator/src/statement/database_copy.rs
+313
View File
@@ -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<String, String>) -> usize {
options
.get("parallelism")
.and_then(|v| v.parse::<usize>().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<Self> {
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<Vec<TableRef>> {
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);
}
}
src/operator/src/statement/export_database.rs
+400
View File
@@ -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<DatabaseExportJob>,
}
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<TableRef>,
) -> Result<PreparedDatabaseExport> {
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::<Vec<_>>();
let routes = self
.table_metadata_manager
.table_route_manager()
.table_route_storage()
.batch_get(&ids)
.await
.context(error::TableMetadataManagerSnafu)?;
let mut groups = BTreeMap::<_, Vec<TableRef>>::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::<Vec<_>>();
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::<BTreeMap<_, _>>();
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<String>,
}
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<DatabaseExportSummary> {
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<J, F: Future<Output = Result<usize>>>(
jobs: impl IntoIterator<Item = J>,
parallelism: usize,
cancellation: &CancellationToken,
mut run: impl FnMut(J, CancellationToken) -> F,
) -> Result<usize> {
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);
}
}
src/operator/src/statement/export_logical_tables.rs
+16 -9
View File
@@ -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<usize>,
}
@@ -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> {
Self::try_new_in_directory(physical, tables, "")
}
pub(crate) fn try_new_in_directory(
physical: TableRef,
tables: &[TableRef],
directory: &str,
) -> Result<Self> {
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<Item = &DatabaseExportFile> {
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::<Vec<_>>();
let routes = manager
@@ -485,7 +492,7 @@ impl ActiveWriter {
store: &ObjectStore,
limits: LogicalTableExportLimits,
) -> Result<Self> {
let path = format!("{}.parquet", table.name);
let path = table.output.path.clone();
ensure!(
!store
.exists(&path)
src/operator/src/statement/export_logical_tables/tests.rs
-2
View File
@@ -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",
src/query/src/dist_plan/planner.rs
+36 -17
View File
@@ -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<Option<TableName>> {
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<Vec<RegionId>> {
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<TableName>,
captured_table: Option<TableRef>,
}
impl TreeNodeVisitor<'_> for TableNameExtractor {
impl TreeNodeVisitor<'_> for TableScanExtractor {
type Node = LogicalPlan;
fn f_down(&mut self, node: &Self::Node) -> Result<TreeNodeRecursion> {
@@ -482,6 +490,7 @@ impl TreeNodeVisitor<'_> for TableNameExtractor {
.downcast_ref::<DfTableProviderAdapter>()
{
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) =
tests-integration/tests/export_logical_tables.rs
+571 -31
View File
@@ -54,40 +54,52 @@ async fn values(instance: &Arc<Instance>, query: &str) -> Vec<Vec<datatypes::val
.collect()
}
async fn create_metric_export_source_tables(
instance: &Arc<Instance>,
physical: &str,
encoding: &str,
) -> ([String; 3], Vec<table::TableRef>, 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<Instance>) {
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<Instance>) {
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::<std::collections::BTreeSet<_>>();
assert_eq!(
summary
.output_files
.into_iter()
.collect::<std::collections::BTreeSet<_>>(),
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::<std::collections::BTreeSet<_>>();
assert_eq!(
summary
.output_files
.into_iter()
.collect::<std::collections::BTreeSet<_>>(),
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;
}
}
}