From 90606af0706f27ead6bf2155b4238df8b3af0fd1 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Mon, 15 Jun 2026 02:56:31 -0700 Subject: [PATCH] feat(cli): add export-v2 progress reporting (#8294) Signed-off-by: jeremyhi --- src/cli/src/data/export_v2/command.rs | 88 +++++++++++++-- src/cli/src/data/export_v2/coordinator.rs | 58 ++++++++-- src/cli/src/data/import_v2/command.rs | 129 +--------------------- src/cli/src/data/progress.rs | 126 ++++++++++++++++++++- 4 files changed, 253 insertions(+), 148 deletions(-) diff --git a/src/cli/src/data/export_v2/command.rs b/src/cli/src/data/export_v2/command.rs index a264a0f293..4df741db34 100644 --- a/src/cli/src/data/export_v2/command.rs +++ b/src/cli/src/data/export_v2/command.rs @@ -27,7 +27,7 @@ use snafu::{OptionExt, ResultExt}; use crate::Tool; use crate::common::ObjectStoreConfig; -use crate::data::export_v2::coordinator::export_data; +use crate::data::export_v2::coordinator::{ExportDataOptions, export_data}; use crate::data::export_v2::error::{ ChunkTimeWindowRequiresBoundsSnafu, DatabaseSnafu, EmptyResultSnafu, IoSnafu, ManifestVersionMismatchSnafu, Result, ResumeConfigMismatchSnafu, SchemaOnlyArgsNotAllowedSnafu, @@ -39,6 +39,7 @@ use crate::data::export_v2::manifest::{ }; use crate::data::export_v2::schema::{DDL_DIR, SCHEMA_DIR, SCHEMAS_FILE}; use crate::data::path::{data_dir_for_schema_chunk, ddl_path_for_schema}; +use crate::data::progress::{ProgressMode, build_progress_reporter}; use crate::data::snapshot_storage::{ OpenDalStorage, SnapshotStorage, validate_snapshot_uri, validate_uri, }; @@ -321,6 +322,10 @@ pub struct ExportCreateCommand { #[clap(long)] no_proxy: bool, + /// Progress reporting mode. + #[clap(long, value_enum, default_value_t = ProgressMode::Auto)] + progress: ProgressMode, + /// Object store configuration for remote storage backends. #[clap(flatten)] storage: ObjectStoreConfig, @@ -399,6 +404,7 @@ impl ExportCreateCommand { chunk_time_window: self.chunk_time_window, parallelism: self.parallelism, chunk_parallelism: self.chunk_parallelism, + progress: self.progress, snapshot_uri: self.to.clone(), storage_config: self.storage.clone(), }, @@ -425,6 +431,7 @@ struct ExportConfig { chunk_time_window: Option, parallelism: usize, chunk_parallelism: usize, + progress: ProgressMode, snapshot_uri: String, storage_config: ObjectStoreConfig, } @@ -487,14 +494,18 @@ impl ExportCreate { return Ok(()); } + let progress = build_progress_reporter(self.config.progress); export_data( self.storage.as_ref(), &self.database_client, - &self.config.snapshot_uri, - &self.config.storage_config, &mut manifest, - self.config.parallelism, - self.config.chunk_parallelism, + ExportDataOptions { + snapshot_uri: &self.config.snapshot_uri, + storage_config: &self.config.storage_config, + parallelism: self.config.parallelism, + chunk_parallelism: self.config.chunk_parallelism, + }, + progress.as_ref(), ) .await?; return Ok(()); @@ -545,14 +556,18 @@ impl ExportCreate { info!("Snapshot created: {}", manifest.snapshot_id); if !self.config.schema_only { + let progress = build_progress_reporter(self.config.progress); export_data( self.storage.as_ref(), &self.database_client, - &self.config.snapshot_uri, - &self.config.storage_config, &mut manifest, - self.config.parallelism, - self.config.chunk_parallelism, + ExportDataOptions { + snapshot_uri: &self.config.snapshot_uri, + storage_config: &self.config.storage_config, + parallelism: self.config.parallelism, + chunk_parallelism: self.config.chunk_parallelism, + }, + progress.as_ref(), ) .await?; } @@ -1607,6 +1622,56 @@ mod tests { assert_eq!(1, cmd.chunk_parallelism); } + #[test] + fn test_progress_mode_defaults_to_auto() { + let cmd = ExportCreateCommand::parse_from([ + "export-v2-create", + "--addr", + "127.0.0.1:4000", + "--to", + "file:///tmp/export-v2-test", + ]); + + assert_eq!(ProgressMode::Auto, cmd.progress); + } + + #[test] + fn test_progress_mode_parses_explicit_values() { + for (value, expected) in [ + ("auto", ProgressMode::Auto), + ("always", ProgressMode::Always), + ("never", ProgressMode::Never), + ] { + let cmd = ExportCreateCommand::parse_from([ + "export-v2-create", + "--addr", + "127.0.0.1:4000", + "--to", + "file:///tmp/export-v2-test", + "--progress", + value, + ]); + + assert_eq!(expected, cmd.progress); + } + } + + #[test] + fn test_progress_mode_rejects_unknown_value() { + assert!( + ExportCreateCommand::try_parse_from([ + "export-v2-create", + "--addr", + "127.0.0.1:4000", + "--to", + "file:///tmp/export-v2-test", + "--progress", + "bogus", + ]) + .is_err() + ); + } + #[test] fn test_chunk_parallelism_parses_valid_value() { let cmd = ExportCreateCommand::parse_from([ @@ -1684,6 +1749,7 @@ mod tests { chunk_time_window: None, parallelism: 1, chunk_parallelism: 1, + progress: ProgressMode::Auto, snapshot_uri: "file:///tmp/snapshot".to_string(), storage_config: ObjectStoreConfig::default(), }; @@ -1720,6 +1786,7 @@ mod tests { chunk_time_window: None, parallelism: 1, chunk_parallelism: 1, + progress: ProgressMode::Auto, snapshot_uri: "file:///tmp/snapshot".to_string(), storage_config: ObjectStoreConfig::default(), }; @@ -1751,6 +1818,7 @@ mod tests { chunk_time_window: Some(Duration::from_secs(3600)), parallelism: 1, chunk_parallelism: 1, + progress: ProgressMode::Auto, snapshot_uri: "file:///tmp/snapshot".to_string(), storage_config: ObjectStoreConfig::default(), }; @@ -1783,6 +1851,7 @@ mod tests { chunk_time_window: None, parallelism: 1, chunk_parallelism: 1, + progress: ProgressMode::Auto, snapshot_uri: "file:///tmp/snapshot".to_string(), storage_config: ObjectStoreConfig::default(), }; @@ -1817,6 +1886,7 @@ mod tests { chunk_time_window: None, parallelism: 1, chunk_parallelism: 1, + progress: ProgressMode::Auto, snapshot_uri: "file:///tmp/snapshot".to_string(), storage_config: ObjectStoreConfig::default(), }; diff --git a/src/cli/src/data/export_v2/coordinator.rs b/src/cli/src/data/export_v2/coordinator.rs index 1f63b6705e..64aa4ce994 100644 --- a/src/cli/src/data/export_v2/coordinator.rs +++ b/src/cli/src/data/export_v2/coordinator.rs @@ -21,6 +21,7 @@ use crate::data::export_v2::data::{CopyOptions, build_copy_target, execute_copy_ use crate::data::export_v2::error::{Error, Result}; use crate::data::export_v2::manifest::{ChunkStatus, DataFormat, Manifest, TimeRange}; use crate::data::path::data_dir_for_schema_chunk; +use crate::data::progress::{ProgressPhase, ProgressReporter}; use crate::data::snapshot_storage::{SnapshotStorage, StorageScheme}; use crate::database::DatabaseClient; @@ -40,14 +41,19 @@ struct ExportContext<'a> { parallelism: usize, } +pub struct ExportDataOptions<'a> { + pub snapshot_uri: &'a str, + pub storage_config: &'a ObjectStoreConfig, + pub parallelism: usize, + pub chunk_parallelism: usize, +} + pub async fn export_data( storage: &dyn SnapshotStorage, database_client: &DatabaseClient, - snapshot_uri: &str, - storage_config: &ObjectStoreConfig, manifest: &mut Manifest, - parallelism: usize, - chunk_parallelism: usize, + options: ExportDataOptions<'_>, + progress: &dyn ProgressReporter, ) -> Result<()> { if manifest.chunks.is_empty() { return Ok(()); @@ -56,19 +62,43 @@ pub async fn export_data( let context = ExportContext { storage, database_client, - snapshot_uri, - storage_config, + snapshot_uri: options.snapshot_uri, + storage_config: options.storage_config, catalog: manifest.catalog.clone(), schemas: manifest.schemas.clone(), format: manifest.format, - parallelism, + parallelism: options.parallelism, }; - if chunk_parallelism <= 1 { - export_data_serial(&context, storage, manifest).await - } else { - export_data_concurrent(&context, storage, manifest, chunk_parallelism).await + // One progress unit per chunk. Already completed/skipped chunks from a + // previous run count up front so the reported total matches the chunk plan + // regardless of resume position; each chunk finalized in this run (completed, + // skipped, or failed) then increments exactly once. + let progress_phase = ProgressPhase::start( + progress, + "Export data chunks", + Some(manifest.chunks.len() as u64), + ); + let already_done = (manifest.completed_count() + manifest.skipped_count()) as u64; + if already_done > 0 { + progress.inc(already_done); } + + let result = if options.chunk_parallelism <= 1 { + export_data_serial(&context, storage, manifest, progress).await + } else { + export_data_concurrent( + &context, + storage, + manifest, + options.chunk_parallelism, + progress, + ) + .await + }; + + progress_phase.finish(); + result } /// Exports chunks one at a time, preserving the original serial behavior. @@ -76,6 +106,7 @@ async fn export_data_serial( context: &ExportContext<'_>, storage: &dyn SnapshotStorage, manifest: &mut Manifest, + progress: &dyn ProgressReporter, ) -> Result<()> { for idx in 0..manifest.chunks.len() { if matches!( @@ -104,6 +135,8 @@ 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?; } @@ -127,6 +160,7 @@ async fn export_data_concurrent( storage: &dyn SnapshotStorage, manifest: &mut Manifest, chunk_parallelism: usize, + progress: &dyn ProgressReporter, ) -> Result<()> { let mut pending = FuturesUnordered::new(); let mut next_idx = 0; @@ -171,6 +205,8 @@ 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 first_error { diff --git a/src/cli/src/data/import_v2/command.rs b/src/cli/src/data/import_v2/command.rs index 2f8d5a7f07..ed9727cd8a 100644 --- a/src/cli/src/data/import_v2/command.rs +++ b/src/cli/src/data/import_v2/command.rs @@ -15,7 +15,6 @@ //! Import V2 CLI command. use std::collections::HashSet; -use std::io::IsTerminal; use std::time::Duration; use async_trait::async_trait; @@ -40,72 +39,10 @@ use crate::data::import_v2::error::{ use crate::data::import_v2::executor::{DdlExecutor, DdlStatement}; use crate::data::import_v2::state::{ImportTaskKey, default_state_path}; use crate::data::path::{data_dir_for_schema_chunk, ddl_path_for_schema}; -use crate::data::progress::{IndicatifProgress, LogProgress, NoopProgress, ProgressReporter}; +use crate::data::progress::{ProgressMode, build_progress_reporter}; use crate::data::snapshot_storage::{OpenDalStorage, SnapshotStorage, validate_uri}; use crate::database::{DatabaseClient, parse_proxy_opts}; -/// Controls progress reporting for import-v2. -/// -/// `auto` shows an interactive bar only on a TTY and falls back to lightweight -/// log progress otherwise; `always` always emits progress, using the bar on a -/// TTY and lightweight logs otherwise; `never` is silent. -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, clap::ValueEnum)] -#[value(rename_all = "lowercase")] -pub(crate) enum ProgressMode { - /// Show an interactive bar on a TTY, otherwise log progress. - #[default] - Auto, - /// Always emit progress: a bar on TTY, lightweight logs otherwise. - Always, - /// Never emit progress. - Never, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum ProgressOutputKind { - Bar, - Log, - Silent, -} - -/// Selects the progress output style for `mode` given whether stderr is a TTY -/// and whether the terminal environment is suitable for progress-bar control -/// sequences. -/// -/// The interactive `indicatif` stderr target hides itself when redirected or -/// when the terminal is not user-attended, so forced progress falls back to -/// [`LogProgress`] instead of a hidden bar in those cases. -fn progress_output_kind( - mode: ProgressMode, - stderr_is_tty: bool, - term_supports_progress_bar: bool, -) -> ProgressOutputKind { - let can_show_bar = stderr_is_tty && term_supports_progress_bar; - match mode { - ProgressMode::Never => ProgressOutputKind::Silent, - ProgressMode::Always if can_show_bar => ProgressOutputKind::Bar, - ProgressMode::Always => ProgressOutputKind::Log, - ProgressMode::Auto if can_show_bar => ProgressOutputKind::Bar, - ProgressMode::Auto => ProgressOutputKind::Log, - } -} - -fn progress_bar_supported(term: Option<&std::ffi::OsStr>, no_color: bool) -> bool { - if no_color { - return false; - } - - term.map(|term| !term.is_empty() && term != "dumb") - .unwrap_or(false) -} - -fn term_supports_progress_bar() -> bool { - progress_bar_supported( - std::env::var_os("TERM").as_deref(), - std::env::var_os("NO_COLOR").is_some(), - ) -} - /// Import from a snapshot. #[derive(Debug, Parser)] pub struct ImportV2Command { @@ -336,7 +273,7 @@ impl Import { import: self, format: manifest.format, }; - let progress = self.progress_reporter(); + let progress = build_progress_reporter(self.progress); import_with_resume_session_with_progress(resume_session, &executor, progress.as_ref()) .await?; } @@ -371,21 +308,6 @@ impl Import { Ok(statements) } - - /// Builds the progress reporter for this run. `never` is silent; otherwise - /// an interactive bar is used on TTY stderr, falling back to lightweight log - /// progress when stderr is redirected. - fn progress_reporter(&self) -> Box { - match progress_output_kind( - self.progress, - std::io::stderr().is_terminal(), - term_supports_progress_bar(), - ) { - ProgressOutputKind::Bar => Box::new(IndicatifProgress::new()), - ProgressOutputKind::Log => Box::new(LogProgress::new()), - ProgressOutputKind::Silent => Box::new(NoopProgress), - } - } } struct CopyDatabaseImportTaskExecutor<'a> { @@ -797,53 +719,6 @@ mod tests { ); } - #[test] - fn test_progress_output_kind_visibility_matrix() { - // auto follows the TTY for the bar and falls back to log progress; - // always emits progress even when stderr is redirected; never is silent. - assert_eq!( - ProgressOutputKind::Bar, - progress_output_kind(ProgressMode::Auto, true, true) - ); - assert_eq!( - ProgressOutputKind::Log, - progress_output_kind(ProgressMode::Auto, true, false) - ); - assert_eq!( - ProgressOutputKind::Log, - progress_output_kind(ProgressMode::Auto, false, true) - ); - assert_eq!( - ProgressOutputKind::Log, - progress_output_kind(ProgressMode::Always, false, true) - ); - assert_eq!( - ProgressOutputKind::Log, - progress_output_kind(ProgressMode::Always, true, false) - ); - assert_eq!( - ProgressOutputKind::Bar, - progress_output_kind(ProgressMode::Always, true, true) - ); - assert_eq!( - ProgressOutputKind::Silent, - progress_output_kind(ProgressMode::Never, true, true) - ); - assert_eq!( - ProgressOutputKind::Silent, - progress_output_kind(ProgressMode::Never, false, true) - ); - } - - #[test] - fn test_progress_bar_supported_respects_terminal_environment() { - assert!(progress_bar_supported(Some("xterm".as_ref()), false)); - assert!(!progress_bar_supported(Some("dumb".as_ref()), false)); - assert!(!progress_bar_supported(Some("".as_ref()), false)); - assert!(!progress_bar_supported(None, false)); - assert!(!progress_bar_supported(Some("xterm".as_ref()), true)); - } - #[test] fn test_progress_mode_rejects_unknown_value() { assert!( diff --git a/src/cli/src/data/progress.rs b/src/cli/src/data/progress.rs index 442cce8113..7bff9dfeab 100644 --- a/src/cli/src/data/progress.rs +++ b/src/cli/src/data/progress.rs @@ -20,11 +20,88 @@ //! stderr, while [`IndicatifProgress`] renders an interactive bar on a TTY. //! Both implement [`ProgressReporter`], so call sites stay agnostic. -use std::io::{self, Write}; +use std::io::{self, IsTerminal, Write}; use std::sync::Mutex; use indicatif::{ProgressBar, ProgressDrawTarget, ProgressStyle}; +/// Controls progress reporting for Export/Import V2. +/// +/// `auto` shows an interactive bar only on a TTY and falls back to lightweight +/// log progress otherwise; `always` always emits progress, using the bar on a +/// TTY and lightweight logs otherwise; `never` is silent. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, clap::ValueEnum)] +#[value(rename_all = "lowercase")] +pub(crate) enum ProgressMode { + /// Show an interactive bar on a TTY, otherwise log progress. + #[default] + Auto, + /// Always emit progress: a bar on TTY, lightweight logs otherwise. + Always, + /// Never emit progress. + Never, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum ProgressOutputKind { + Bar, + Log, + Silent, +} + +/// Selects the progress output style for `mode` given whether stderr is a TTY +/// and whether the terminal environment is suitable for progress-bar control +/// sequences. +/// +/// The interactive `indicatif` stderr target hides itself when redirected or +/// when the terminal is not user-attended, so forced progress falls back to +/// [`LogProgress`] instead of a hidden bar in those cases. +pub(crate) fn progress_output_kind( + mode: ProgressMode, + stderr_is_tty: bool, + term_supports_progress_bar: bool, +) -> ProgressOutputKind { + let can_show_bar = stderr_is_tty && term_supports_progress_bar; + match mode { + ProgressMode::Never => ProgressOutputKind::Silent, + ProgressMode::Always if can_show_bar => ProgressOutputKind::Bar, + ProgressMode::Always => ProgressOutputKind::Log, + ProgressMode::Auto if can_show_bar => ProgressOutputKind::Bar, + ProgressMode::Auto => ProgressOutputKind::Log, + } +} + +fn progress_bar_supported(term: Option<&std::ffi::OsStr>, no_color: bool) -> bool { + if no_color { + return false; + } + + term.map(|term| !term.is_empty() && term != "dumb") + .unwrap_or(false) +} + +fn term_supports_progress_bar() -> bool { + progress_bar_supported( + std::env::var_os("TERM").as_deref(), + std::env::var_os("NO_COLOR").is_some(), + ) +} + +/// Builds the progress reporter for `mode`. `never` is silent; otherwise an +/// interactive bar is used on TTY stderr, falling back to lightweight log +/// progress when stderr is redirected. +pub(crate) fn build_progress_reporter(mode: ProgressMode) -> Box { + match progress_output_kind( + mode, + std::io::stderr().is_terminal(), + term_supports_progress_bar(), + ) { + ProgressOutputKind::Bar => Box::new(IndicatifProgress::new()), + ProgressOutputKind::Log => Box::new(LogProgress::new()), + ProgressOutputKind::Silent => Box::new(NoopProgress), + } +} + /// Receives progress events from long-running Export/Import V2 work. /// /// The trait is object-safe so callers can take `&dyn ProgressReporter` and stay @@ -248,6 +325,53 @@ mod tests { reporter.finish_phase(); // Idempotent: finishing twice is harmless. } + #[test] + fn test_progress_output_kind_visibility_matrix() { + // auto follows the TTY for the bar and falls back to log progress; + // always emits progress even when stderr is redirected; never is silent. + assert_eq!( + ProgressOutputKind::Bar, + progress_output_kind(ProgressMode::Auto, true, true) + ); + assert_eq!( + ProgressOutputKind::Log, + progress_output_kind(ProgressMode::Auto, true, false) + ); + assert_eq!( + ProgressOutputKind::Log, + progress_output_kind(ProgressMode::Auto, false, true) + ); + assert_eq!( + ProgressOutputKind::Log, + progress_output_kind(ProgressMode::Always, false, true) + ); + assert_eq!( + ProgressOutputKind::Log, + progress_output_kind(ProgressMode::Always, true, false) + ); + assert_eq!( + ProgressOutputKind::Bar, + progress_output_kind(ProgressMode::Always, true, true) + ); + assert_eq!( + ProgressOutputKind::Silent, + progress_output_kind(ProgressMode::Never, true, true) + ); + assert_eq!( + ProgressOutputKind::Silent, + progress_output_kind(ProgressMode::Never, false, true) + ); + } + + #[test] + fn test_progress_bar_supported_respects_terminal_environment() { + assert!(progress_bar_supported(Some("xterm".as_ref()), false)); + assert!(!progress_bar_supported(Some("dumb".as_ref()), false)); + assert!(!progress_bar_supported(Some("".as_ref()), false)); + assert!(!progress_bar_supported(None, false)); + assert!(!progress_bar_supported(Some("xterm".as_ref()), true)); + } + #[test] fn test_indicatif_progress_is_safe_across_phase_lifecycle() { // IndicatifProgress takes only `&self`, so it must survive a full