From c567a6562a44daec066606c141b21dc4409faea0 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Thu, 11 Jun 2026 23:13:27 -0700 Subject: [PATCH] feat(cli): add import-v2 progress mode (#8283) Signed-off-by: jeremyhi --- src/cli/src/data/import_v2/command.rs | 88 +++++++++++++++++++- src/cli/src/data/import_v2/coordinator.rs | 7 +- src/cli/src/data/progress.rs | 99 ++++++++++++++++++++++- 3 files changed, 186 insertions(+), 8 deletions(-) diff --git a/src/cli/src/data/import_v2/command.rs b/src/cli/src/data/import_v2/command.rs index 9ee60252a0..4fb84c274c 100644 --- a/src/cli/src/data/import_v2/command.rs +++ b/src/cli/src/data/import_v2/command.rs @@ -29,7 +29,7 @@ use crate::data::export_v2::data::{build_copy_source, execute_copy_database_from use crate::data::export_v2::manifest::{ChunkMeta, ChunkStatus, DataFormat, MANIFEST_VERSION}; use crate::data::import_v2::coordinator::{ ImportResumeConfig, ImportTaskExecutor, build_import_tasks, chunk_has_schema_files, - import_with_resume_session, prepare_import_resume, + import_with_resume_session_with_progress, prepare_import_resume, }; use crate::data::import_v2::error::{ ChunkImportFailedSnafu, EmptyChunkManifestSnafu, ImportStatePathUnavailableSnafu, @@ -39,9 +39,27 @@ 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::{LogProgress, NoopProgress, ProgressReporter}; use crate::data::snapshot_storage::{OpenDalStorage, SnapshotStorage, validate_uri}; use crate::database::{DatabaseClient, parse_proxy_opts}; +/// Controls progress reporting for import-v2. +/// +/// For now `auto` and `always` both emit lightweight log progress; only +/// `never` is silent. A richer interactive bar can later distinguish `auto` +/// from `always` without changing call sites. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, clap::ValueEnum)] +#[value(rename_all = "lowercase")] +pub(crate) enum ProgressMode { + /// Choose automatically (currently lightweight log progress). + #[default] + Auto, + /// Always emit lightweight log progress. + Always, + /// Never emit progress. + Never, +} + /// Import from a snapshot. #[derive(Debug, Parser)] pub struct ImportV2Command { @@ -66,6 +84,10 @@ pub struct ImportV2Command { #[clap(long)] dry_run: bool, + /// Progress reporting mode. + #[clap(long, value_enum, default_value_t = ProgressMode::Auto)] + progress: ProgressMode, + /// Basic authentication (user:password). #[clap(long)] auth_basic: Option, @@ -126,6 +148,7 @@ impl ImportV2Command { catalog: self.catalog.clone(), schemas, dry_run: self.dry_run, + progress: self.progress, snapshot_uri: self.from.clone(), storage_config: self.storage.clone(), storage: Box::new(storage), @@ -139,6 +162,7 @@ pub struct Import { catalog: String, schemas: Option>, dry_run: bool, + progress: ProgressMode, snapshot_uri: String, storage_config: ObjectStoreConfig, storage: Box, @@ -266,7 +290,9 @@ impl Import { import: self, format: manifest.format, }; - import_with_resume_session(resume_session, &executor).await?; + let progress = self.progress_reporter(); + import_with_resume_session_with_progress(resume_session, &executor, progress.as_ref()) + .await?; } if ddl_executed { @@ -299,6 +325,15 @@ impl Import { Ok(statements) } + + /// Builds the progress reporter for this run. `never` is silent; `auto` + /// and `always` both use lightweight log progress for now. + fn progress_reporter(&self) -> Box { + match self.progress { + ProgressMode::Never => Box::new(NoopProgress), + ProgressMode::Auto | ProgressMode::Always => Box::new(LogProgress::new()), + } + } } struct CopyDatabaseImportTaskExecutor<'a> { @@ -677,6 +712,55 @@ mod tests { } } + fn parse_command(extra: &[&str]) -> ImportV2Command { + let mut args = vec![ + "import-v2", + "--addr", + "127.0.0.1:4000", + "--from", + "file:///tmp/snapshot", + ]; + args.extend_from_slice(extra); + ImportV2Command::try_parse_from(args).expect("command should parse") + } + + #[test] + fn test_progress_mode_defaults_to_auto() { + assert_eq!(parse_command(&[]).progress, ProgressMode::Auto); + } + + #[test] + fn test_progress_mode_parses_explicit_values() { + assert_eq!( + parse_command(&["--progress", "always"]).progress, + ProgressMode::Always + ); + assert_eq!( + parse_command(&["--progress", "never"]).progress, + ProgressMode::Never + ); + assert_eq!( + parse_command(&["--progress", "auto"]).progress, + ProgressMode::Auto + ); + } + + #[test] + fn test_progress_mode_rejects_unknown_value() { + assert!( + ImportV2Command::try_parse_from([ + "import-v2", + "--addr", + "127.0.0.1:4000", + "--from", + "file:///tmp/snapshot", + "--progress", + "bogus", + ]) + .is_err() + ); + } + #[test] fn test_parse_ddl_statements() { let content = r#" diff --git a/src/cli/src/data/import_v2/coordinator.rs b/src/cli/src/data/import_v2/coordinator.rs index 82a8585416..58b1bdf4d7 100644 --- a/src/cli/src/data/import_v2/coordinator.rs +++ b/src/cli/src/data/import_v2/coordinator.rs @@ -28,7 +28,9 @@ use crate::data::import_v2::state::{ delete_import_state, load_import_state, save_import_state, try_acquire_import_state_lock, }; use crate::data::path::data_dir_for_schema_chunk; -use crate::data::progress::{NoopProgress, ProgressPhase, ProgressReporter}; +#[cfg(test)] +use crate::data::progress::NoopProgress; +use crate::data::progress::{ProgressPhase, ProgressReporter}; #[async_trait] pub(crate) trait ImportTaskExecutor { @@ -127,6 +129,7 @@ pub(crate) async fn prepare_import_resume( }) } +#[cfg(test)] pub(crate) async fn import_with_resume_session( session: ImportResumeSession, executor: &E, @@ -134,8 +137,6 @@ pub(crate) async fn import_with_resume_session( where E: ImportTaskExecutor + Sync, { - // Production callers do not care about progress yet; keep their call site - // unchanged and route them through the no-op reporter. let progress = NoopProgress; import_with_resume_session_with_progress(session, executor, &progress).await } diff --git a/src/cli/src/data/progress.rs b/src/cli/src/data/progress.rs index 0851d0a095..9e6945bfcd 100644 --- a/src/cli/src/data/progress.rs +++ b/src/cli/src/data/progress.rs @@ -15,9 +15,13 @@ //! Minimal internal progress abstraction for Export/Import V2. //! //! This is intentionally small and log/internal oriented. It does not touch -//! stdout and is safe for non-interactive runs. A TTY progress bar (e.g. -//! `indicatif`) and a `--progress` CLI flag are deliberately out of scope; they -//! can layer on top of this trait later without changing call sites. +//! stdout and is safe for non-interactive runs. [`LogProgress`] backs the +//! import-v2 `--progress` flag by routing events to stderr. A TTY progress bar +//! (e.g. `indicatif`) is deliberately out of scope; it can layer on top of this +//! trait later without changing call sites. + +use std::io::{self, Write}; +use std::sync::Mutex; /// Receives progress events from long-running Export/Import V2 work. /// @@ -45,6 +49,75 @@ impl ProgressReporter for NoopProgress { fn finish_phase(&self) {} } +/// A lightweight reporter that logs phase lifecycle and progress through the +/// stderr. It never touches stdout, so it is safe for non-interactive runs and +/// keeps dry-run output clean. +pub(crate) struct LogProgress { + phase: Mutex>, +} + +struct PhaseState { + name: String, + total: Option, + done: u64, +} + +impl LogProgress { + pub(crate) fn new() -> Self { + Self { + phase: Mutex::new(None), + } + } +} + +fn write_progress_line(line: String) { + let _ = writeln!(io::stderr().lock(), "{line}"); +} + +impl ProgressReporter for LogProgress { + fn start_phase(&self, name: &str, total: Option) { + let Ok(mut phase) = self.phase.lock() else { + return; + }; + *phase = Some(PhaseState { + name: name.to_string(), + total, + done: 0, + }); + match total { + Some(total) => write_progress_line(format!("Starting phase '{name}' ({total} units)")), + None => write_progress_line(format!("Starting phase '{name}'")), + } + } + + fn inc(&self, delta: u64) { + let Ok(mut guard) = self.phase.lock() else { + return; + }; + if let Some(phase) = guard.as_mut() { + phase.done += delta; + match phase.total { + Some(total) => { + write_progress_line(format!("Phase '{}': {}/{}", phase.name, phase.done, total)) + } + None => write_progress_line(format!("Phase '{}': {}", phase.name, phase.done)), + } + } + } + + fn finish_phase(&self) { + let Ok(mut guard) = self.phase.lock() else { + return; + }; + if let Some(phase) = guard.take() { + write_progress_line(format!( + "Finished phase '{}' ({} units)", + phase.name, phase.done + )); + } + } +} + /// RAII guard for a started progress phase. /// /// This keeps future stateful reporters safe on every early-return path after a @@ -87,6 +160,26 @@ impl Drop for ProgressPhase<'_> { } } +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_log_progress_is_safe_across_phase_lifecycle() { + // LogProgress takes only `&self`, so it must drive a full phase + // lifecycle (including an out-of-phase `inc`) without panicking. + let progress = LogProgress::new(); + let reporter: &dyn ProgressReporter = &progress; + + reporter.inc(1); // No active phase yet: must be a no-op, not a panic. + reporter.start_phase("Import data tasks", Some(2)); + reporter.inc(1); + reporter.inc(1); + reporter.finish_phase(); + reporter.finish_phase(); // Idempotent: finishing twice is harmless. + } +} + #[cfg(test)] pub(crate) mod test_util { use std::sync::Mutex;