diff --git a/Cargo.lock b/Cargo.lock index 36ee6d77ffd..dc0dcf2561c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -15125,6 +15125,7 @@ dependencies = [ "catalog", "chrono", "clap", + "cli", "client", "cmd", "common-base", diff --git a/config/config.md b/config/config.md index 651cbec6e85..7385c8efb0d 100644 --- a/config/config.md +++ b/config/config.md @@ -15,6 +15,7 @@ | `default_timezone` | String | Unset | The default timezone of the server. | | `default_column_prefix` | String | Unset | The default column prefix for auto-created time index, value, and native histogram columns.
Legacy OTLP summary columns keep their historical `greptime_` prefix. | | `auto_create_table` | Bool | `true` | Server-side global switch for auto table creation on write.
When `false`, a missing table is never auto-created even if the request sets the `auto_create_table` hint to `true`. Default: `true`. | +| `experimental_metric_export` | Bool | `false` | Enables experimental Parquet COPY DATABASE using shared Metric physical scans.
Resume requires the previous export and storage writes to have ended; HTTP timeout is not confirmation. | | `user_provider` | String | Unset | The user provider for authentication.
Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd"
Password verifier formats: "plain:", "pbkdf2_sha256:::",
"mysql_native_password:",
"pg_scram_sha256::::"
"pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest, but cannot authenticate over MySQL's
native password handshake; a MySQL client must send the password in cleartext for such users.
"mysql_native_password" is MySQL-specific and cannot authenticate over PostgreSQL at all.
PostgreSQL SCRAM only covers "plain" and "pg_scram_sha256" users; if any user is "pbkdf2_sha256" or
"mysql_native_password", PostgreSQL falls back to cleartext password auth for every user.
For "pg_scram_sha256" users, keep the default iteration count (4096) and salt length (16): both are
observable in the SCRAM server-first message, and non-default values weaken resistance to username
enumeration. | | `max_in_flight_write_bytes` | String | Unset | Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight).
Set to 0 to disable the limit. Default: "0" (unlimited) | | `write_bytes_exhausted_policy` | String | Unset | Policy when write bytes quota is exhausted.
Options: "wait" (default, 10s timeout), "wait()" (e.g., "wait(30s)"), "fail" | @@ -278,6 +279,7 @@ | `default_timezone` | String | Unset | The default timezone of the server. | | `default_column_prefix` | String | Unset | The default column prefix for auto-created time index, value, and native histogram columns.
Legacy OTLP summary columns keep their historical `greptime_` prefix. | | `auto_create_table` | Bool | `true` | Server-side global switch for auto table creation on write.
When `false`, a missing table is never auto-created even if the request sets the `auto_create_table` hint to `true`. Default: `true`. | +| `experimental_metric_export` | Bool | `false` | Enables experimental Parquet COPY DATABASE using shared Metric physical scans.
Resume requires the previous export and storage writes to have ended; HTTP timeout is not confirmation. | | `user_provider` | String | Unset | The user provider for authentication.
Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd"
Password verifier formats: "plain:", "pbkdf2_sha256:::",
"mysql_native_password:",
"pg_scram_sha256::::"
"pbkdf2_sha256" and "pg_scram_sha256" protect passwords at rest, but cannot authenticate over MySQL's
native password handshake; a MySQL client must send the password in cleartext for such users.
"mysql_native_password" is MySQL-specific and cannot authenticate over PostgreSQL at all.
PostgreSQL SCRAM only covers "plain" and "pg_scram_sha256" users; if any user is "pbkdf2_sha256" or
"mysql_native_password", PostgreSQL falls back to cleartext password auth for every user.
For "pg_scram_sha256" users, keep the default iteration count (4096) and salt length (16): both are
observable in the SCRAM server-first message, and non-default values weaken resistance to username
enumeration. | | `max_in_flight_write_bytes` | String | Unset | Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight).
Set to 0 to disable the limit. Default: "0" (unlimited) | | `write_bytes_exhausted_policy` | String | Unset | Policy when write bytes quota is exhausted.
Options: "wait" (default, 10s timeout), "wait()" (e.g., "wait(30s)"), "fail" | diff --git a/config/frontend.example.toml b/config/frontend.example.toml index 95550931d46..043da57db74 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -11,6 +11,10 @@ default_column_prefix = "greptime" ## When `false`, a missing table is never auto-created even if the request sets the `auto_create_table` hint to `true`. Default: `true`. #+ auto_create_table = true +## Enables experimental Parquet COPY DATABASE using shared Metric physical scans. +## Resume requires the previous export and storage writes to have ended; HTTP timeout is not confirmation. +#+ experimental_metric_export = false + ## The user provider for authentication. ## Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd" ## Password verifier formats: "plain:", "pbkdf2_sha256:::", diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 57893881146..cc4be3f5930 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -11,6 +11,10 @@ default_column_prefix = "greptime" ## When `false`, a missing table is never auto-created even if the request sets the `auto_create_table` hint to `true`. Default: `true`. #+ auto_create_table = true +## Enables experimental Parquet COPY DATABASE using shared Metric physical scans. +## Resume requires the previous export and storage writes to have ended; HTTP timeout is not confirmation. +#+ experimental_metric_export = false + ## The user provider for authentication. ## Examples: "static_user_provider:file:/path/to/users", "static_user_provider:cmd:greptime_user=greptime_pwd" ## Password verifier formats: "plain:", "pbkdf2_sha256:::", diff --git a/src/cli/src/data/export_v2.rs b/src/cli/src/data/export_v2.rs index 9bdf3c63f5d..f9a3c2be095 100644 --- a/src/cli/src/data/export_v2.rs +++ b/src/cli/src/data/export_v2.rs @@ -37,6 +37,15 @@ //! --start-time 2025-01-01T00:00:00Z \ //! --end-time 2025-01-31T23:59:59Z //! ``` +//! +//! `--experimental-metric-export` enables shared Metric physical scans for Parquet. +//! Enable `experimental_metric_export = true` on the frontend or standalone server +//! and use an endpoint whose frontends all support and enable this option. +//! The command checks the server capability before modifying the snapshot. +//! Each snapshot path belongs to one export task. Resume requires stable source +//! schemas/data and confirmation that the previous export and storage writes ended; +//! an HTTP timeout does not establish that. Only unfinished chunks are cleaned and +//! rerun. Completed chunks and the V2 snapshot/import format remain unchanged. mod chunker; mod command; diff --git a/src/cli/src/data/export_v2/command.rs b/src/cli/src/data/export_v2/command.rs index 4df741db347..fc958e9e3e0 100644 --- a/src/cli/src/data/export_v2/command.rs +++ b/src/cli/src/data/export_v2/command.rs @@ -20,9 +20,11 @@ use std::time::Duration; use async_trait::async_trait; use clap::{Parser, Subcommand}; +use common_catalog::consts::DEFAULT_SCHEMA_NAME; use common_error::ext::BoxedError; use common_telemetry::info; use serde_json::Value; +use servers::http::{ColumnSchema, GreptimeQueryOutput, OutputSchema}; use snafu::{OptionExt, ResultExt}; use crate::Tool; @@ -289,6 +291,11 @@ pub struct ExportCreateCommand { #[clap(long, value_enum, default_value = "parquet")] format: DataFormat, + /// Use shared Metric physical scans (Parquet only). Resume requires that the previous + /// export and storage writes have ended; an HTTP timeout does not establish this. + #[clap(long)] + experimental_metric_export: bool, + /// Delete existing snapshot and recreate. #[clap(long)] force: bool, @@ -379,9 +386,6 @@ impl ExportCreateCommand { Some(self.schemas.clone()) }; - // Build storage - let storage = OpenDalStorage::from_uri(&self.to, &self.storage).map_err(BoxedError::new)?; - // Build database client let proxy = parse_proxy_opts(self.proxy.clone(), self.no_proxy)?; let database_client = DatabaseClient::new( @@ -393,12 +397,43 @@ impl ExportCreateCommand { self.no_proxy, ); + // The filesystem storage constructor can create the snapshot root. + if self.experimental_metric_export { + if self.format != DataFormat::Parquet { + return crate::data::export_v2::error::MetricExportUnavailableSnafu + .fail() + .map_err(BoxedError::new); + } + let capability = database_client + .sql_response( + "SHOW VARIABLES experimental_metric_export", + DEFAULT_SCHEMA_NAME, + ) + .await + .context(DatabaseSnafu) + .map_err(BoxedError::new)?; + let expected_schema = OutputSchema::new(vec![ColumnSchema::new( + "EXPERIMENTAL_METRIC_EXPORT".to_string(), + "String".to_string(), + )]); + if !matches!(capability.output(), [GreptimeQueryOutput::Records(records)] + if records.schema() == &expected_schema + && records.rows() == &vec![vec![Value::String("true".to_string())]]) + { + return crate::data::export_v2::error::MetricExportUnavailableSnafu + .fail() + .map_err(BoxedError::new); + } + } + let storage = OpenDalStorage::from_uri(&self.to, &self.storage).map_err(BoxedError::new)?; + Ok(Box::new(ExportCreate { config: ExportConfig { catalog: self.catalog.clone(), schemas, schema_only: self.schema_only, format: self.format, + experimental_metric_export: self.experimental_metric_export, force: self.force, time_range, chunk_time_window: self.chunk_time_window, @@ -426,6 +461,7 @@ struct ExportConfig { schemas: Option>, schema_only: bool, format: DataFormat, + experimental_metric_export: bool, force: bool, time_range: TimeRange, chunk_time_window: Option, @@ -504,6 +540,8 @@ impl ExportCreate { storage_config: &self.config.storage_config, parallelism: self.config.parallelism, chunk_parallelism: self.config.chunk_parallelism, + experimental_metric_export: self.config.experimental_metric_export, + resume: true, }, progress.as_ref(), ) @@ -533,6 +571,14 @@ impl ExportCreate { self.config.chunk_time_window, )?; + if self.config.experimental_metric_export { + for chunk in &manifest.chunks { + self.storage + .prepare_export_chunk(&schema_names, chunk.id, false) + .await?; + } + } + // 4. Write schema files self.storage.write_schema(&schema_snapshot).await?; info!("Exported {} schemas", schema_snapshot.schemas.len()); @@ -566,6 +612,8 @@ impl ExportCreate { storage_config: &self.config.storage_config, parallelism: self.config.parallelism, chunk_parallelism: self.config.chunk_parallelism, + experimental_metric_export: self.config.experimental_metric_export, + resume: false, }, progress.as_ref(), ) @@ -1744,6 +1792,7 @@ mod tests { schemas: None, schema_only: false, format: DataFormat::Parquet, + experimental_metric_export: false, force: false, time_range: TimeRange::unbounded(), chunk_time_window: None, @@ -1781,6 +1830,7 @@ mod tests { ]), schema_only: false, format: DataFormat::Parquet, + experimental_metric_export: false, force: false, time_range: TimeRange::unbounded(), chunk_time_window: None, @@ -1813,6 +1863,7 @@ mod tests { schemas: None, schema_only: false, format: DataFormat::Parquet, + experimental_metric_export: false, force: false, time_range, chunk_time_window: Some(Duration::from_secs(3600)), @@ -1846,6 +1897,7 @@ mod tests { schemas: None, schema_only: false, format: DataFormat::Csv, + experimental_metric_export: false, force: false, time_range: TimeRange::unbounded(), chunk_time_window: None, @@ -1881,6 +1933,7 @@ mod tests { schemas: None, schema_only: false, format: DataFormat::Parquet, + experimental_metric_export: false, force: false, time_range: TimeRange::new(Some(start), Some(start)), chunk_time_window: None, diff --git a/src/cli/src/data/export_v2/coordinator.rs b/src/cli/src/data/export_v2/coordinator.rs index 64aa4ce9946..b70897337a7 100644 --- a/src/cli/src/data/export_v2/coordinator.rs +++ b/src/cli/src/data/export_v2/coordinator.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use common_telemetry::info; +use common_telemetry::{error, info}; use futures::StreamExt; use futures::stream::FuturesUnordered; @@ -39,6 +39,8 @@ struct ExportContext<'a> { schemas: Vec, format: DataFormat, parallelism: usize, + experimental_metric_export: bool, + resume: bool, } pub struct ExportDataOptions<'a> { @@ -46,6 +48,8 @@ pub struct ExportDataOptions<'a> { pub storage_config: &'a ObjectStoreConfig, pub parallelism: usize, pub chunk_parallelism: usize, + pub experimental_metric_export: bool, + pub resume: bool, } pub async fn export_data( @@ -68,6 +72,8 @@ pub async fn export_data( schemas: manifest.schemas.clone(), format: manifest.format, parallelism: options.parallelism, + experimental_metric_export: options.experimental_metric_export, + resume: options.resume, }; // One progress unit per chunk. Already completed/skipped chunks from a @@ -88,11 +94,11 @@ pub async fn export_data( export_data_serial(&context, storage, manifest, progress).await } else { export_data_concurrent( - &context, storage, manifest, options.chunk_parallelism, progress, + |id, range| export_chunk(&context, id, range), ) .await }; @@ -134,11 +140,17 @@ async fn export_data_serial( }; manifest.touch(); - storage.write_manifest(manifest).await?; - // The chunk is finalized (completed, skipped, or failed) and persisted. - progress.inc(1); - - result?; + let persistence = storage.write_manifest(manifest).await; + if persistence.is_ok() { + progress.inc(1); + } + if let Err(err) = result { + if let Err(secondary) = persistence { + error!(secondary; "Failed to persist failed export chunk"); + } + return Err(err); + } + persistence?; } Ok(()) @@ -155,39 +167,44 @@ async fn export_data_serial( /// On the first chunk failure we stop scheduling new chunks but let already /// in-flight chunks finish and persist their final status, then return the /// first error. -async fn export_data_concurrent( - context: &ExportContext<'_>, +async fn export_data_concurrent( storage: &dyn SnapshotStorage, manifest: &mut Manifest, chunk_parallelism: usize, progress: &dyn ProgressReporter, -) -> Result<()> { + export: F, +) -> Result<()> +where + F: Fn(u32, TimeRange) -> Fut, + Fut: std::future::Future>>, +{ let mut pending = FuturesUnordered::new(); let mut next_idx = 0; let mut first_error: Option = None; loop { - let mut scheduled = false; - + let mut admitted = Vec::new(); // Schedule eligible chunks in order up to the parallelism limit. Once a // failure is seen, stop scheduling but keep draining in-flight chunks. - while first_error.is_none() && pending.len() < chunk_parallelism { + while first_error.is_none() && pending.len() + admitted.len() < chunk_parallelism { let Some(idx) = next_eligible_chunk(manifest, &mut next_idx) else { break; }; let (chunk_id, time_range) = mark_chunk_in_progress(manifest, idx); - scheduled = true; - - pending.push(async move { - let result = export_chunk(context, chunk_id, time_range).await; - (idx, result) - }); + admitted.push((idx, chunk_id, time_range)); } - - if scheduled { + if !admitted.is_empty() { manifest.touch(); - storage.write_manifest(manifest).await?; + match storage.write_manifest(manifest).await { + Ok(()) => { + for (idx, chunk_id, time_range) in admitted { + let export = &export; + pending.push(async move { (idx, export(chunk_id, time_range).await) }); + } + } + Err(err) => first_error = Some(err), + } } let Some((idx, export_result)) = pending.next().await else { @@ -204,9 +221,11 @@ async fn export_data_concurrent( } } manifest.touch(); - storage.write_manifest(manifest).await?; - // The chunk is finalized (completed, skipped, or failed) and persisted. - progress.inc(1); + match storage.write_manifest(manifest).await { + Ok(()) => progress.inc(1), + Err(err) if first_error.is_none() => first_error = Some(err), + Err(err) => error!(err; "Failed to persist export chunk while draining"), + } } match first_error { @@ -257,12 +276,19 @@ async fn export_chunk( chunk_id: u32, time_range: TimeRange, ) -> Result> { + if context.experimental_metric_export { + context + .storage + .prepare_export_chunk(&context.schemas, chunk_id, context.resume) + .await?; + } let scheme = StorageScheme::from_uri(context.snapshot_uri)?; let needs_dir = matches!(scheme, StorageScheme::File); let copy_options = CopyOptions { format: context.format, time_range, parallelism: context.parallelism, + experimental_metric_export: context.experimental_metric_export, }; for schema in &context.schemas { @@ -326,6 +352,113 @@ mod tests { manifest } + struct FailingManifest { + fail_at: usize, + writes: std::sync::atomic::AtomicUsize, + release: tokio::sync::Notify, + } + + #[async_trait::async_trait] + impl SnapshotStorage for FailingManifest { + async fn write_manifest(&self, manifest: &Manifest) -> Result<()> { + use std::sync::atomic::Ordering; + + let write = self.writes.fetch_add(1, Ordering::SeqCst) + 1; + if write == 1 { + assert_eq!(manifest.in_progress_count(), 2); + } + if write == self.fail_at { + self.release.notify_one(); + return crate::data::export_v2::error::IoSnafu { + operation: "injected manifest failure", + error: std::io::Error::other("manifest failure"), + } + .fail(); + } + Ok(()) + } + async fn exists(&self) -> Result { + unreachable!() + } + async fn read_manifest(&self) -> Result { + unreachable!() + } + async fn write_schema( + &self, + _: &crate::data::export_v2::schema::SchemaSnapshot, + ) -> Result<()> { + unreachable!() + } + async fn write_text(&self, _: &str, _: &str) -> Result<()> { + unreachable!() + } + async fn read_text(&self, _: &str) -> Result { + unreachable!() + } + async fn create_dir_all(&self, _: &str) -> Result<()> { + unreachable!() + } + async fn list_files_recursive(&self, _: &str) -> Result> { + unreachable!() + } + async fn delete_snapshot(&self) -> Result<()> { + unreachable!() + } + } + + #[tokio::test] + async fn test_manifest_failure_stops_admission_and_drains_started_chunks() { + use std::sync::atomic::{AtomicUsize, Ordering}; + for fail_at in [1, 2, 3] { + let storage = FailingManifest { + fail_at, + writes: AtomicUsize::new(0), + release: tokio::sync::Notify::new(), + }; + let started = AtomicUsize::new(0); + let finished = AtomicUsize::new(0); + let second_started = tokio::sync::Notify::new(); + let mut manifest = pending_manifest(4); + let result = tokio::time::timeout( + std::time::Duration::from_secs(5), + export_data_concurrent( + &storage, + &mut manifest, + 2, + &crate::data::progress::NoopProgress, + |id, _| { + let (started, finished, second_started, storage) = + (&started, &finished, &second_started, &storage); + async move { + started.fetch_add(1, Ordering::SeqCst); + match id { + 1 => second_started.notified().await, + 2 => { + second_started.notify_one(); + storage.release.notified().await; + } + _ => panic!("new chunk started after manifest persistence failed"), + } + finished.fetch_add(1, Ordering::SeqCst); + Ok(vec![format!("data/public/{id}/a.parquet")]) + } + }, + ), + ) + .await + .unwrap(); + assert!( + result + .unwrap_err() + .to_string() + .contains("injected manifest failure") + ); + let expected = if fail_at == 1 { 0 } else { 2 }; + assert_eq!(started.load(Ordering::SeqCst), expected); + assert_eq!(finished.load(Ordering::SeqCst), expected); + } + } + #[test] fn test_next_eligible_chunk_scans_in_order() { let manifest = pending_manifest(3); diff --git a/src/cli/src/data/export_v2/data.rs b/src/cli/src/data/export_v2/data.rs index dce924a717d..9faf0009829 100644 --- a/src/cli/src/data/export_v2/data.rs +++ b/src/cli/src/data/export_v2/data.rs @@ -30,6 +30,7 @@ pub(super) struct CopyOptions { pub(super) format: DataFormat, pub(super) time_range: TimeRange, pub(super) parallelism: usize, + pub(crate) experimental_metric_export: bool, } pub(super) struct CopyTarget { @@ -210,6 +211,9 @@ pub(crate) async fn execute_copy_database_from( fn build_with_options(options: &CopyOptions) -> String { let mut parts = vec![format!("FORMAT='{}'", options.format)]; + if options.experimental_metric_export { + parts.push("experimental_metric_export='true'".to_string()); + } if let Some(start) = options.time_range.start { parts.push(format!( "START_TIME='{}'", diff --git a/src/cli/src/data/export_v2/error.rs b/src/cli/src/data/export_v2/error.rs index 078f5bd2e1f..dbb4a88ae46 100644 --- a/src/cli/src/data/export_v2/error.rs +++ b/src/cli/src/data/export_v2/error.rs @@ -24,6 +24,14 @@ use snafu::{Location, Snafu}; #[snafu(visibility(pub))] #[stack_trace_debug] pub enum Error { + #[snafu(display( + "Experimental Metric export requires Parquet and a frontend that explicitly reports experimental_metric_export=true. Use an upgraded, consistently configured frontend entry point." + ))] + MetricExportUnavailable { + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Invalid URI '{}': {}", uri, reason))] InvalidUri { uri: String, @@ -211,7 +219,8 @@ pub type Result = std::result::Result; impl ErrorExt for Error { fn status_code(&self) -> StatusCode { match self { - Error::InvalidUri { .. } + Error::MetricExportUnavailable { .. } + | Error::InvalidUri { .. } | Error::UnsupportedScheme { .. } | Error::SchemaOnlyModeMismatch { .. } | Error::ResumeConfigMismatch { .. } diff --git a/src/cli/src/data/snapshot_storage.rs b/src/cli/src/data/snapshot_storage.rs index 17da8aea0b0..a5bc02ca7fb 100644 --- a/src/cli/src/data/snapshot_storage.rs +++ b/src/cli/src/data/snapshot_storage.rs @@ -40,6 +40,7 @@ use crate::data::export_v2::manifest::{MANIFEST_FILE, Manifest}; #[cfg(test)] use crate::data::export_v2::schema::SchemaDefinition; use crate::data::export_v2::schema::{SCHEMA_DIR, SCHEMAS_FILE, SchemaSnapshot}; +use crate::data::path::data_dir_for_schema_chunk; struct RemoteLocation { bucket_or_container: String, @@ -285,6 +286,21 @@ pub trait SnapshotStorage: Send + Sync { /// Lists files recursively under a relative prefix. async fn list_files_recursive(&self, prefix: &str) -> Result>; + /// Checks fresh output directories or removes a terminated, unfinished chunk's files. + async fn prepare_export_chunk( + &self, + schemas: &[String], + chunk_id: u32, + resume: bool, + ) -> Result<()> { + let _ = (schemas, chunk_id, resume); + InvalidUriSnafu { + uri: "snapshot", + reason: "storage does not support preparing export chunks", + } + .fail() + } + /// Deletes the entire snapshot (for --force). async fn delete_snapshot(&self) -> Result<()>; } @@ -696,6 +712,56 @@ impl SnapshotStorage for OpenDalStorage { Ok(files) } + async fn prepare_export_chunk( + &self, + schemas: &[String], + chunk_id: u32, + resume: bool, + ) -> Result<()> { + let mut files = Vec::new(); + for schema in schemas { + let prefix = data_dir_for_schema_chunk(schema, chunk_id); + let mut entries = match self.object_store.lister_with(&prefix).recursive(true).await { + Ok(entries) => entries, + Err(error) if error.kind() == ErrorKind::NotFound => continue, + Err(error) => { + return Err(error).context(StorageOperationSnafu { + operation: format!("list {prefix}"), + }); + } + }; + while let Some(entry) = entries.try_next().await.context(StorageOperationSnafu { + operation: format!("list {prefix}"), + })? { + let path = entry.path(); + if path == prefix && entry.metadata().is_dir() { + continue; + } + let name = path.strip_prefix(&prefix).unwrap_or(""); + if !resume || entry.metadata().is_dir() || !valid_chunk_filename(name) { + return InvalidUriSnafu { + uri: path, + reason: "expected an empty new chunk or direct Parquet files in an owned unfinished chunk", + }.fail(); + } + files.push(path.to_string()); + } + } + // Validate every schema before deleting any files; COPY starts only after this returns. + for path in files { + match self.object_store.delete(&path).await { + Ok(()) => {} + Err(error) if error.kind() == ErrorKind::NotFound => {} + Err(error) => { + return Err(error).context(StorageOperationSnafu { + operation: format!("delete unfinished chunk file {path}"), + }); + } + } + } + Ok(()) + } + async fn delete_snapshot(&self) -> Result<()> { self.object_store .delete_with("/") @@ -707,6 +773,12 @@ impl SnapshotStorage for OpenDalStorage { } } +fn valid_chunk_filename(name: &str) -> bool { + name.strip_suffix(".parquet") + .is_some_and(|stem| !stem.is_empty()) + && !name.contains(['/', '\\', '\0']) +} + #[cfg(test)] mod tests { use std::collections::HashMap; @@ -971,6 +1043,140 @@ mod tests { assert!(storage.exists().await.unwrap()); } + #[tokio::test] + async fn test_prepare_export_chunk_preserves_other_chunks_and_schema_artifacts() { + let dir = tempdir().unwrap(); + let storage = make_storage_with_rooted_fs(dir.path()); + for path in [ + "data/public/1/keep.parquet", + "data/public/2/cpu.v1.parquet", + "data/other/2/a.parquet", + "schema/ddl/public.sql", + ] { + storage.write_text(path, "original").await.unwrap(); + } + let schemas = vec!["public".to_string(), "other".to_string()]; + assert!( + storage + .prepare_export_chunk(&schemas, 2, false) + .await + .is_err() + ); + assert!( + storage + .file_exists("data/public/2/cpu.v1.parquet") + .await + .unwrap() + ); + storage + .prepare_export_chunk(&schemas, 2, true) + .await + .unwrap(); + storage + .prepare_export_chunk(&schemas, 2, true) + .await + .unwrap(); + assert!( + !storage + .file_exists("data/public/2/cpu.v1.parquet") + .await + .unwrap() + ); + assert!(!storage.file_exists("data/other/2/a.parquet").await.unwrap()); + assert_eq!( + storage + .read_text("data/public/1/keep.parquet") + .await + .unwrap(), + "original" + ); + assert_eq!( + storage.read_text("schema/ddl/public.sql").await.unwrap(), + "original" + ); + } + + #[tokio::test] + async fn test_prepare_export_chunk_rejects_foreign_or_nested_entries_before_deletion() { + for foreign in ["notes.txt", "nested/a.parquet", "empty/"] { + let dir = tempdir().unwrap(); + let storage = make_storage_with_rooted_fs(dir.path()); + storage + .write_text("data/public/2/cpu.parquet", "keep") + .await + .unwrap(); + let path = format!("data/other/2/{foreign}"); + if foreign.ends_with('/') { + storage.create_dir_all(&path).await.unwrap(); + } else { + storage.write_text(&path, "foreign").await.unwrap(); + } + assert!( + storage + .prepare_export_chunk(&["public".into(), "other".into()], 2, true) + .await + .is_err() + ); + assert_eq!( + storage + .read_text("data/public/2/cpu.parquet") + .await + .unwrap(), + "keep" + ); + } + } + + #[cfg(unix)] + #[tokio::test] + async fn test_prepare_export_chunk_reports_delete_failure() { + use std::os::unix::fs::PermissionsExt; + let dir = tempdir().unwrap(); + let storage = make_storage_with_rooted_fs(dir.path()); + storage + .write_text("data/public/2/a.parquet", "keep") + .await + .unwrap(); + let chunk_dir = dir.path().join("data/public/2"); + std::fs::set_permissions(&chunk_dir, std::fs::Permissions::from_mode(0o555)).unwrap(); + let result = storage + .prepare_export_chunk(&["public".into()], 2, true) + .await; + std::fs::set_permissions(&chunk_dir, std::fs::Permissions::from_mode(0o755)).unwrap(); + assert!( + result + .unwrap_err() + .to_string() + .contains("delete unfinished chunk file") + ); + assert_eq!( + storage.read_text("data/public/2/a.parquet").await.unwrap(), + "keep" + ); + } + + #[tokio::test] + async fn test_prepare_export_chunk_encodes_schema_boundary() { + let dir = tempdir().unwrap(); + let storage = make_storage_with_rooted_fs(dir.path()); + let schema = "../other".to_string(); + let owned = format!("{}a.parquet", data_dir_for_schema_chunk(&schema, 2)); + storage.write_text(&owned, "remove").await.unwrap(); + storage + .write_text("data/other/2/a.parquet", "keep") + .await + .unwrap(); + storage + .prepare_export_chunk(&[schema], 2, true) + .await + .unwrap(); + assert!(!storage.file_exists(&owned).await.unwrap()); + assert_eq!( + storage.read_text("data/other/2/a.parquet").await.unwrap(), + "keep" + ); + } + #[tokio::test] async fn test_delete_snapshot_only_removes_rooted_contents() { let parent = tempdir().unwrap(); diff --git a/src/cli/src/database.rs b/src/cli/src/database.rs index cba27129dd6..e295d248f79 100644 --- a/src/cli/src/database.rs +++ b/src/cli/src/database.rs @@ -99,6 +99,18 @@ impl DatabaseClient { /// Execute sql query. pub async fn sql(&self, sql: &str, schema: &str) -> Result>>> { + let body = self.sql_response(sql, schema).await?; + Ok(body.output().first().and_then(|output| match output { + GreptimeQueryOutput::Records(records) => Some(records.rows().clone()), + GreptimeQueryOutput::AffectedRows(_) => None, + })) + } + + pub(crate) async fn sql_response( + &self, + sql: &str, + schema: &str, + ) -> Result { let url = format!("http://{}/v1/sql", self.addr); let params = [ ("db", format!("{}-{}", self.catalog, schema)), @@ -138,11 +150,7 @@ impl DatabaseClient { reason: "cannot get response text".to_string(), })?; - let body = serde_json::from_str::(&text).context(SerdeJsonSnafu)?; - Ok(body.output().first().and_then(|output| match output { - GreptimeQueryOutput::Records(records) => Some(records.rows().clone()), - GreptimeQueryOutput::AffectedRows(_) => None, - })) + serde_json::from_str::(&text).context(SerdeJsonSnafu) } } diff --git a/src/frontend/src/frontend.rs b/src/frontend/src/frontend.rs index 793ab94205e..86ba2cb9bfb 100644 --- a/src/frontend/src/frontend.rs +++ b/src/frontend/src/frontend.rs @@ -50,6 +50,8 @@ pub struct FrontendOptions { /// even if a request sets the `auto_create_table` hint to `true`. When `true` /// (default), the per-request hint still applies. Default: `true`. pub auto_create_table: bool, + /// Enables experimental Parquet database exports using shared Metric scans. + pub experimental_metric_export: bool, /// Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight). /// Set to 0 to disable the limit. Default: "0" (unlimited) pub max_in_flight_write_bytes: ReadableSize, @@ -93,6 +95,7 @@ impl Default for FrontendOptions { default_timezone: None, default_column_prefix: None, auto_create_table: true, + experimental_metric_export: false, max_in_flight_write_bytes: ReadableSize(0), write_bytes_exhausted_policy: OnExhaustedPolicy::default(), http: HttpOptions::default(), @@ -273,6 +276,9 @@ max_batch_rows = 25 #[test] fn test_toml() { let opts = FrontendOptions::default(); + assert!(!opts.experimental_metric_export); + let enabled: FrontendOptions = toml::from_str("experimental_metric_export = true").unwrap(); + assert!(enabled.experimental_metric_export); let toml_string = toml::to_string(&opts).unwrap(); assert!(toml_string.contains("experimental_enable_exponential_histogram = false")); let parsed: FrontendOptions = toml::from_str(&toml_string).unwrap(); diff --git a/src/frontend/src/instance.rs b/src/frontend/src/instance.rs index e16955f2877..0d1f36ea524 100644 --- a/src/frontend/src/instance.rs +++ b/src/frontend/src/instance.rs @@ -122,6 +122,7 @@ lazy_static! { #[derive(Clone)] pub struct Instance { frontend_peer_addr: String, + experimental_metric_export: bool, catalog_manager: CatalogManagerRef, pipeline_operator: Arc, statement_executor: Arc, @@ -391,6 +392,21 @@ impl Instance { } _ => { query_interceptor.pre_execute(Some(&stmt), None, query_ctx.clone())?; + if let Statement::ShowVariables(show) = &stmt + && show + .variable + .to_string() + .eq_ignore_ascii_case("experimental_metric_export") + { + return self.show_metric_export_capability(); + } + if let Statement::Copy(sql::statements::copy::Copy::CopyDatabase(CopyDatabase::To( + arg, + ))) = &stmt + && export_database::parse_metric_export_requested(&arg.with)? + { + return self.copy_metric_database(arg.clone(), query_ctx).await; + } self.statement_executor .execute_sql(stmt, query_ctx) .await @@ -2734,6 +2750,116 @@ mod tests { ) } + struct ReleasedExportSource { + schema: GtSchemaRef, + channels: std::sync::Mutex, oneshot::Receiver<()>)>>, + } + + impl DataSource for ReleasedExportSource { + fn get_stream( + &self, + _request: ScanRequest, + ) -> std::result::Result { + let (started, release) = self.channels.lock().unwrap().take().unwrap(); + let schema = self.schema.clone(); + let stream = futures::stream::once(async move { + started.send(()).unwrap(); + release.await.unwrap(); + Ok(RecordBatch::new_empty(schema)) + }); + Ok(Box::pin(RecordBatchStreamWrapper::new( + self.schema.clone(), + Box::pin(stream), + ))) + } + } + + #[tokio::test] + async fn test_metric_export_http_timeout_drains_ordinary_writer() { + let destination = common_test_util::temp_dir::create_temp_dir("metric_export_timeout"); + let (started_tx, started_rx) = oneshot::channel(); + let (release_tx, release_rx) = oneshot::channel(); + let info = test_table_info(1024, "source").unwrap(); + let source = Arc::new(Table::new( + Arc::new(info.clone()), + FilterPushDownType::Unsupported, + Arc::new(ReleasedExportSource { + schema: info.meta.schema.clone(), + channels: std::sync::Mutex::new(Some((started_tx, release_rx))), + }), + )); + let catalog = catalog::memory::MemoryCatalogManager::new_with_table(source); + let kv = Arc::new(MemoryKvBackend::new()); + let instance = FrontendBuilder::new( + FrontendOptions { + experimental_metric_export: true, + ..Default::default() + }, + kv.clone(), + test_cache_registry(kv).unwrap(), + catalog, + Arc::new(client::client_manager::NodeClients::default()), + Arc::new(NoopProcedureExecutor), + Arc::new(ProcessManager::new("export-timeout".into(), None)), + ) + .with_local_file_access( + common_datasource::object_store::LocalFileAccess::sandboxed(destination.path()) + .unwrap(), + ) + .try_build() + .await + .unwrap(); + let server = servers::http::HttpServerBuilder::new(servers::http::HttpOptions { + timeout: Duration::from_secs(2), + ..Default::default() + }) + .with_sql_handler(Arc::new(instance)) + .build(); + let app = server.build(server.make_app()).unwrap(); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server_task = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + let sql = format!( + "COPY DATABASE greptime.public TO '{}/' WITH (experimental_metric_export='true')", + destination.path().display() + ); + let request = tokio::spawn(async move { + reqwest::Client::builder() + .no_proxy() + .build() + .unwrap() + .post(format!("http://{addr}/v1/sql")) + .form(&[("sql", sql)]) + .send() + .await + .unwrap() + }); + tokio::time::timeout(Duration::from_secs(5), started_rx) + .await + .unwrap() + .unwrap(); + let response = request.await.unwrap(); + assert_eq!(response.status(), reqwest::StatusCode::REQUEST_TIMEOUT); + // A dropped ordinary stream would close this receiver before release. + release_tx.send(()).unwrap(); + let file = destination.path().join("source.parquet"); + tokio::time::timeout(Duration::from_secs(5), async { + loop { + if std::fs::read(&file) + .is_ok_and(|bytes| bytes.len() > 8 && bytes.ends_with(b"PAR1")) + { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); + server_task.abort(); + } + fn pending_table( table_id: u32, table_name: &str, diff --git a/src/frontend/src/instance/builder.rs b/src/frontend/src/instance/builder.rs index 7a6ba474168..01e9ef034e3 100644 --- a/src/frontend/src/instance/builder.rs +++ b/src/frontend/src/instance/builder.rs @@ -387,6 +387,7 @@ impl FrontendBuilder { Ok(Instance { frontend_peer_addr, + experimental_metric_export: self.options.experimental_metric_export, catalog_manager: self.catalog_manager, pipeline_operator, statement_executor, diff --git a/src/frontend/src/instance/export_database.rs b/src/frontend/src/instance/export_database.rs index 6594490548c..f36bcf562c7 100644 --- a/src/frontend/src/instance/export_database.rs +++ b/src/frontend/src/instance/export_database.rs @@ -12,14 +12,24 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::sync::Arc; + use auth::{ PermissionAction, PermissionChecker, PermissionCheckerRef, PermissionReq, PermissionTableTarget, PermissionTableTargets, }; +use common_error::ext::BoxedError; +use common_query::Output; +use common_recordbatch::RecordBatches; +use common_telemetry::error; +use datatypes::prelude::ConcreteDataType; +use datatypes::schema::{ColumnSchema, Schema}; +use datatypes::vectors::StringVector; use operator::statement::export_database::{DatabaseExportSummary, PreparedDatabaseExport}; use session::context::QueryContextRef; use snafu::ResultExt; use sql::ast::{Ident, ObjectName}; +use sql::statements::OptionMap; use sql::statements::copy::{Copy, CopyDatabase, CopyDatabaseArgument}; use sql::statements::statement::Statement; use table::requests::CopyDatabaseRequest; @@ -28,7 +38,65 @@ use tokio_util::sync::CancellationToken; use crate::error::{PermissionSnafu, Result}; use crate::instance::Instance; +pub(crate) fn parse_metric_export_requested(options: &OptionMap) -> Result { + match options.get("experimental_metric_export") { + None | Some("false") => Ok(false), + Some("true") => Ok(true), + Some(_) => Err(operator::error::InvalidDatabaseExportSnafu { + reason: "experimental_metric_export must be true or false", + } + .build() + .into()), + } +} + impl Instance { + pub(crate) fn show_metric_export_capability(&self) -> Result { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "EXPERIMENTAL_METRIC_EXPORT", + ConcreteDataType::string_datatype(), + false, + )])); + let batches = RecordBatches::try_from_columns( + schema, + vec![Arc::new(StringVector::from(vec![ + self.experimental_metric_export.to_string(), + ])) as datatypes::vectors::VectorRef], + ) + .map_err(BoxedError::new) + .context(crate::error::ExternalSnafu)?; + Ok(Output::new_with_record_batches(batches)) + } + + pub(crate) async fn copy_metric_database( + &self, + arg: CopyDatabaseArgument, + ctx: QueryContextRef, + ) -> Result { + if !self.experimental_metric_export { + return Err(operator::error::InvalidDatabaseExportSnafu { + reason: "experimental_metric_export is disabled on this frontend", + } + .build() + .into()); + } + let req = operator::statement::to_copy_database_request(arg, &ctx)?; + let plan = self.prepare_database_export(req, None, &ctx).await?; + let cancellation = CancellationToken::new(); + let _guard = cancellation.clone().drop_guard(); + let executor = self.statement_executor.clone(); + // Dropping the request only cancels; the runtime task retains and drains started I/O. + let task = common_runtime::spawn_query(async move { + let result = executor.export_database(plan, &cancellation, ctx).await; + if let Err(err) = &result { + error!(err; "Experimental database export failed after draining started work"); + } + result + }); + let summary = task.await.context(operator::error::JoinTaskSnafu)??; + Ok(Output::new_with_affected_rows(summary.rows)) + } + #[allow(dead_code)] async fn export_database( &self, diff --git a/src/operator/src/statement.rs b/src/operator/src/statement.rs index 61804b30d4c..6822dc97b75 100644 --- a/src/operator/src/statement.rs +++ b/src/operator/src/statement.rs @@ -870,7 +870,7 @@ fn to_copy_table_request(stmt: CopyTable, query_ctx: QueryContextRef) -> Result< /// Converts [CopyDatabaseArgument] to [CopyDatabaseRequest]. /// This function extracts the necessary info including catalog/database name, time range, etc. -fn to_copy_database_request( +pub fn to_copy_database_request( arg: CopyDatabaseArgument, query_ctx: &QueryContextRef, ) -> Result { diff --git a/src/standalone/src/options.rs b/src/standalone/src/options.rs index 4884b5ee091..80ba4cf11a5 100644 --- a/src/standalone/src/options.rs +++ b/src/standalone/src/options.rs @@ -44,6 +44,8 @@ pub struct StandaloneOptions { /// Upper bound: when `false`, missing tables are never auto-created even if a /// request sets the `auto_create_table` hint to `true`. Default: `true`. pub auto_create_table: bool, + /// Enables experimental Parquet database exports using shared Metric scans. + pub experimental_metric_export: bool, /// Maximum total memory for all concurrent write request bodies and messages (HTTP, gRPC, Flight). /// Set to 0 to disable the limit. Default: "0" (unlimited) pub max_in_flight_write_bytes: ReadableSize, @@ -91,6 +93,7 @@ impl Default for StandaloneOptions { default_timezone: None, default_column_prefix: None, auto_create_table: true, + experimental_metric_export: false, max_in_flight_write_bytes: ReadableSize(0), write_bytes_exhausted_policy: OnExhaustedPolicy::default(), http: HttpOptions::default(), @@ -154,6 +157,7 @@ impl StandaloneOptions { FrontendOptions { default_timezone: cloned_opts.default_timezone, auto_create_table: cloned_opts.auto_create_table, + experimental_metric_export: cloned_opts.experimental_metric_export, max_in_flight_write_bytes: cloned_opts.max_in_flight_write_bytes, write_bytes_exhausted_policy: cloned_opts.write_bytes_exhausted_policy, http: cloned_opts.http, @@ -306,6 +310,9 @@ flow_notification_queue_capacity = 17 #[test] fn test_query_options_propagated_to_components() { let mut options = StandaloneOptions::default(); + assert!(!options.frontend_options().experimental_metric_export); + options.experimental_metric_export = true; + assert!(options.frontend_options().experimental_metric_export); options.query.parallelism = 4; assert_eq!(options.frontend_options().query.parallelism, 4); diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index d097432b828..b5603bd29df 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -29,6 +29,7 @@ cache.workspace = true catalog.workspace = true chrono.workspace = true clap.workspace = true +cli.workspace = true client.workspace = true cmd.workspace = true common-base.workspace = true diff --git a/tests-integration/src/standalone.rs b/tests-integration/src/standalone.rs index e875f2d7a01..307a7f6daa2 100644 --- a/tests-integration/src/standalone.rs +++ b/tests-integration/src/standalone.rs @@ -87,6 +87,7 @@ pub struct GreptimeDbStandaloneBuilder { slow_query_options: SlowQueryOptions, event_recorder_options: EventRecorderOptions, auto_create_table: bool, + experimental_metric_export: bool, } impl GreptimeDbStandaloneBuilder { @@ -106,9 +107,17 @@ impl GreptimeDbStandaloneBuilder { }, event_recorder_options: EventRecorderOptions::default(), auto_create_table: true, + experimental_metric_export: false, } } + /// Enables experimental Metric export for the standalone test instance. + #[must_use] + pub fn with_experimental_metric_export(mut self) -> Self { + self.experimental_metric_export = true; + self + } + #[must_use] pub fn with_auto_create_table(mut self, auto_create_table: bool) -> Self { self.auto_create_table = auto_create_table; @@ -367,6 +376,7 @@ impl GreptimeDbStandaloneBuilder { slow_query: self.slow_query_options.clone(), event_recorder: self.event_recorder_options.clone(), auto_create_table: self.auto_create_table, + experimental_metric_export: self.experimental_metric_export, // Tests cover the descriptor, so they run with it enabled. otlp: frontend::service_config::OtlpOptions { experimental_enable_resource_info: true, diff --git a/tests-integration/tests/export_logical_tables.rs b/tests-integration/tests/export_logical_tables.rs index 8d5ea3e42bd..e87bf387979 100644 --- a/tests-integration/tests/export_logical_tables.rs +++ b/tests-integration/tests/export_logical_tables.rs @@ -746,3 +746,487 @@ async fn database_export_preserves_valid_table_names() { } } } + +#[derive(clap::Parser)] +struct ExportDataCli { + #[command(subcommand)] + command: cli::DataCommand, +} + +async fn run_data_cli(args: &[&str]) -> Result<(), common_error::ext::BoxedError> { + use clap::Parser; + ExportDataCli::try_parse_from(std::iter::once("greptime-data").chain(args.iter().copied())) + .unwrap() + .command + .build() + .await? + .do_work() + .await +} + +async fn export_http(instance: Arc) -> (String, tokio::task::JoinHandle<()>) { + let server = servers::http::HttpServerBuilder::new(servers::http::HttpOptions::default()) + .with_sql_handler(instance) + .build(); + let app = server.build(server.make_app()).unwrap(); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap().to_string(); + let task = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + (addr, task) +} + +struct FailSecondChunk { + fail: std::sync::atomic::AtomicBool, + copies: std::sync::atomic::AtomicUsize, +} + +impl servers::interceptor::SqlQueryInterceptor for FailSecondChunk { + type Error = frontend::error::Error; + + fn pre_execute( + &self, + statement: Option<&sql::statements::statement::Statement>, + _plan: Option<&datafusion_expr::LogicalPlan>, + _ctx: session::context::QueryContextRef, + ) -> Result<(), Self::Error> { + use sql::statements::copy::{Copy, CopyDatabase}; + use sql::statements::statement::Statement; + if let Some(Statement::Copy(Copy::CopyDatabase(CopyDatabase::To(arg)))) = statement { + self.copies + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + if arg.location.ends_with("/z_later/2/") + && self.fail.swap(false, std::sync::atomic::Ordering::SeqCst) + { + return Err(operator::error::InvalidDatabaseExportSnafu { + reason: "injected terminated COPY failure", + } + .build() + .into()); + } + } + Ok(()) + } +} + +#[tokio::test] +async fn metric_export_v2_cli_resume_roundtrip() { + metric_export_v2_cli_roundtrip(false).await; +} + +#[tokio::test] +#[ignore = "requires GT_S3_* credentials and an S3-compatible test bucket"] +async fn metric_export_v2_cli_s3_resume_roundtrip() { + metric_export_v2_cli_roundtrip(true).await; +} + +async fn metric_export_v2_cli_roundtrip(s3: bool) { + use servers::interceptor::SqlQueryInterceptorRef; + let plugins = common_base::Plugins::new(); + let faults = Arc::new(FailSecondChunk { + fail: std::sync::atomic::AtomicBool::new(true), + copies: std::sync::atomic::AtomicUsize::new(0), + }); + plugins.insert::>(faults.clone()); + let source = GreptimeDbStandaloneBuilder::new("metric_v2_source") + .with_experimental_metric_export() + .with_plugin(plugins) + .build() + .await; + let instance = source.fe_instance(); + let mut names = Vec::new(); + for physical in ["v2_a", "v2_b"] { + let (logical, _, renamed) = + create_metric_export_source_tables(instance, physical, "dense").await; + sql( + instance, + &format!("ALTER TABLE {renamed} RENAME {physical}"), + ) + .await; + names.extend(logical); + names.push(format!("{physical}_excluded")); + } + sql( + instance, + "CREATE TABLE audit (host STRING PRIMARY KEY, val DOUBLE, ts TIMESTAMP TIME INDEX)", + ) + .await; + sql(instance, "INSERT INTO audit VALUES ('a',1,1),('z',NULL,3)").await; + sql(instance, "CREATE VIEW dashboard AS SELECT * FROM audit").await; + sql(instance, "CREATE DATABASE z_later").await; + sql( + instance, + "CREATE TABLE z_later.events (ts TIMESTAMP TIME INDEX, val BIGINT)", + ) + .await; + sql(instance, "INSERT INTO z_later.events VALUES (1,10),(3,20)").await; + names.push("audit".into()); + if s3 { + sql( + instance, + "CREATE TABLE bulk (host STRING PRIMARY KEY, payload STRING, ts TIMESTAMP TIME INDEX)", + ) + .await; + for batch in 0..64 { + let rows = (0..64) + .map(|row| { + let payload = (0..128) + .map(|_| uuid::Uuid::new_v4().simple().to_string()) + .collect::(); + format!("('{}','{payload}',1)", batch * 64 + row) + }) + .collect::>() + .join(","); + sql(instance, &format!("INSERT INTO bulk VALUES {rows}")).await; + } + names.push("bulk".into()); + } + let (addr, server) = export_http(instance.clone()).await; + let destination = tempfile::tempdir_in(common_test_util::find_workspace_path(".")).unwrap(); + let (uri, store, storage_args) = if s3 { + let endpoint = std::env::var("GT_S3_ENDPOINT_URL").unwrap(); + let bucket = std::env::var("GT_S3_BUCKET").unwrap(); + let region = std::env::var("GT_S3_REGION").unwrap(); + let key = std::env::var("GT_S3_ACCESS_KEY_ID").unwrap(); + let secret = std::env::var("GT_S3_ACCESS_KEY").unwrap(); + let root = format!("pr04b-{}", uuid::Uuid::new_v4()); + let store = object_store::ObjectStore::new( + object_store::services::S3::default() + .endpoint(&endpoint) + .bucket(&bucket) + .root(&root) + .region(®ion) + .access_key_id(&key) + .secret_access_key(&secret), + ) + .unwrap(); + ( + format!("s3://{bucket}/{root}"), + store, + vec![ + "--s3".into(), + "--s3-endpoint".into(), + endpoint, + "--s3-region".into(), + region, + "--s3-access-key-id".into(), + key, + "--s3-secret-access-key".into(), + secret, + ], + ) + } else { + ( + url::Url::from_file_path(destination.path()) + .unwrap() + .to_string(), + object_store::ObjectStore::new( + object_store::services::Fs::default().root(destination.path().to_str().unwrap()), + ) + .unwrap(), + Vec::::new(), + ) + }; + let mut args = vec![ + "export-v2", + "create", + "--addr", + &addr, + "--to", + &uri, + "--schemas", + "public,z_later", + "--experimental-metric-export", + "--no-proxy", + "--start-time", + "1970-01-01T00:00:00Z", + "--end-time", + "1970-01-01T00:00:00.006Z", + "--chunk-time-window", + "3ms", + "--progress", + "never", + ]; + args.extend(storage_args.iter().map(String::as_str)); + assert!(run_data_cli(&args).await.is_err()); + let before: cli::export_v2::manifest::Manifest = + serde_json::from_slice(&store.read("manifest.json").await.unwrap().to_vec()).unwrap(); + assert_eq!( + before.chunks[0].status, + cli::export_v2::manifest::ChunkStatus::Completed + ); + assert_eq!( + before.chunks[1].status, + cli::export_v2::manifest::ChunkStatus::Failed + ); + if s3 { + assert!( + store + .stat("data/public/1/bulk.parquet") + .await + .unwrap() + .content_length() + > 5 * 1024 * 1024 + ); + } + let mut preserved = Vec::new(); + for path in &before.chunks[0].files { + preserved.push((path.clone(), store.read(path).await.unwrap().to_vec())); + } + assert!(store.exists("data/public/2/audit.parquet").await.unwrap()); + store + .write("data/public/2/unknown.txt", "keep") + .await + .unwrap(); + let copies = faults.copies.load(std::sync::atomic::Ordering::SeqCst); + assert!(run_data_cli(&args).await.is_err()); + assert_eq!( + faults.copies.load(std::sync::atomic::Ordering::SeqCst), + copies + ); + assert_eq!( + store + .read("data/public/2/unknown.txt") + .await + .unwrap() + .to_vec(), + b"keep" + ); + store.delete("data/public/2/unknown.txt").await.unwrap(); + run_data_cli(&args).await.unwrap(); + let after: cli::export_v2::manifest::Manifest = + serde_json::from_slice(&store.read("manifest.json").await.unwrap().to_vec()).unwrap(); + assert!(after.is_complete()); + assert_eq!( + serde_json::to_value(&before.chunks[0]).unwrap(), + serde_json::to_value(&after.chunks[0]).unwrap() + ); + for (path, bytes) in preserved { + assert_eq!(bytes, store.read(&path).await.unwrap().to_vec()); + } + let mut verify = vec!["export-v2", "verify", "--snapshot", &uri]; + verify.extend(storage_args.iter().map(String::as_str)); + run_data_cli(&verify).await.unwrap(); + let target = GreptimeDbStandaloneBuilder::new("metric_v2_target") + .build() + .await; + for id in 0..24 { + sql( + target.fe_instance(), + &format!("CREATE TABLE occupied_{id} (ts TIMESTAMP TIME INDEX)"), + ) + .await; + } + let (target_addr, target_server) = export_http(target.fe_instance().clone()).await; + let state = destination.path().join("restore-state.json"); + let mut import = vec![ + "import-v2", + "--addr", + &target_addr, + "--from", + &uri, + "--no-proxy", + "--state-path", + state.to_str().unwrap(), + "--progress", + "never", + ]; + import.extend(storage_args.iter().map(String::as_str)); + run_data_cli(&import).await.unwrap(); + for name in names { + let source_table = table(instance, &name).await; + let target_table = table(target.fe_instance(), &name).await; + assert_ne!( + source_table.table_info().table_id(), + target_table.table_info().table_id() + ); + assert_eq!( + source_table.schema().column_schemas(), + target_table.schema().column_schemas() + ); + let query = format!("SELECT * FROM \"{name}\" ORDER BY ts,host"); + assert_eq!( + values(instance, &query).await, + values(target.fe_instance(), &query).await + ); + } + assert_eq!( + values(instance, "SELECT * FROM z_later.events ORDER BY ts").await, + values( + target.fe_instance(), + "SELECT * FROM z_later.events ORDER BY ts" + ) + .await + ); + server.abort(); + target_server.abort(); + if s3 { + store.delete_with("/").recursive(true).await.unwrap(); + } +} + +#[tokio::test] +async fn metric_export_v2_disabled_and_legacy_cli() { + let source = GreptimeDbStandaloneBuilder::new("metric_v2_disabled") + .build() + .await; + let (addr, server) = export_http(source.fe_instance().clone()).await; + let destination = tempfile::tempdir_in(common_test_util::find_workspace_path(".")).unwrap(); + let uri = url::Url::from_file_path(destination.path()) + .unwrap() + .to_string(); + let manifest = destination.path().join("manifest.json"); + std::fs::write(&manifest, b"preserve-before-capability-check").unwrap(); + let args = [ + "export-v2", + "create", + "--addr", + &addr, + "--to", + &uri, + "--force", + "--experimental-metric-export", + "--no-proxy", + "--progress", + "never", + ]; + assert!(run_data_cli(&args).await.is_err()); + assert_eq!( + std::fs::read(&manifest).unwrap(), + b"preserve-before-capability-check" + ); + std::fs::remove_file(&manifest).unwrap(); + sql( + source.fe_instance(), + "CREATE TABLE ordinary (ts TIMESTAMP TIME INDEX, val BIGINT)", + ) + .await; + sql(source.fe_instance(), "INSERT INTO ordinary VALUES (1,2)").await; + run_data_cli(&[ + "export-v2", + "create", + "--addr", + &addr, + "--to", + &uri, + "--schemas", + "public", + "--no-proxy", + "--progress", + "never", + ]) + .await + .unwrap(); + run_data_cli(&["export-v2", "verify", "--snapshot", &uri]) + .await + .unwrap(); + server.abort(); +} + +#[tokio::test] +async fn metric_export_v2_refuses_missing_or_malformed_capability_before_force() { + use axum::http::StatusCode; + use serde_json::json; + let column = json!({"name": "EXPERIMENTAL_METRIC_EXPORT", "data_type": "String"}); + let records = |rows, columns| { + json!({"records": { + "schema": {"column_schemas": columns}, + "rows": rows + }}) + }; + let valid = records(json!([["true"]]), json!([column.clone()])); + let mut responses = [ + (StatusCode::BAD_REQUEST, json!([])), + (StatusCode::OK, json!([])), + (StatusCode::OK, json!([[true]])), + (StatusCode::OK, json!([["true", "extra"]])), + (StatusCode::OK, json!([["true"], ["true"]])), + ] + .map(|(status, rows)| (status, json!([records(rows, json!([column.clone()]))]))) + .to_vec(); + responses.extend([ + ( + StatusCode::OK, + json!([ + valid.clone(), + records(json!([["false"]]), json!([column.clone()])) + ]), + ), + ( + StatusCode::OK, + json!([records(json!([["true"]]), json!([]))]), + ), + ( + StatusCode::OK, + json!([records(json!([["true"]]), json!([column.clone(), column]))]), + ), + ( + StatusCode::OK, + json!([records( + json!([["true"]]), + json!([{"name": "EXPERIMENTAL_METRIC_EXPORT", "data_type": "Boolean"}]) + )]), + ), + ( + StatusCode::OK, + json!([records( + json!([["true"]]), + json!([{"name": "OTHER", "data_type": "String"}]) + )]), + ), + ]); + for (status, output) in responses { + let app = + axum::Router::new().route( + "/v1/sql", + axum::routing::post( + move |axum::Form(form): axum::Form< + std::collections::HashMap, + >| async move { + assert_eq!(form["sql"], "SHOW VARIABLES experimental_metric_export"); + ( + status, + axum::Json(json!({ + "execution_time_ms": 0, + "output": output + })), + ) + }, + ), + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap().to_string(); + let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + let destination = tempfile::tempdir().unwrap(); + let snapshot = destination.path().join("snapshot"); + let uri = url::Url::from_file_path(&snapshot).unwrap().to_string(); + let args = [ + "export-v2", + "create", + "--addr", + &addr, + "--to", + &uri, + "--experimental-metric-export", + "--force", + "--no-proxy", + "--progress", + "never", + ]; + let error = run_data_cli(&args).await.unwrap_err(); + if status == StatusCode::OK { + assert!( + error.to_string().contains("Metric export requires"), + "{error}" + ); + } + assert!(!snapshot.exists()); + std::fs::create_dir(&snapshot).unwrap(); + let manifest = snapshot.join("manifest.json"); + std::fs::write(&manifest, b"preserve").unwrap(); + assert!(run_data_cli(&args).await.is_err()); + assert_eq!(std::fs::read(manifest).unwrap(), b"preserve"); + server.abort(); + } +} diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index 1ad39ae1726..e4e30baeba8 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -2360,6 +2360,7 @@ mem_threshold_on_create = "auto" r#" enable_telemetry = true auto_create_table = true +experimental_metric_export = false max_in_flight_write_bytes = "0KiB" write_bytes_exhausted_policy = "wait" init_regions_in_background = false