From 34c4cae9bb30f98b331164c3a26f29c528161f24 Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Fri, 12 Jun 2026 02:45:45 -0700 Subject: [PATCH] feat(cli): add import-v2 progress UI (#8289) Signed-off-by: jeremyhi --- Cargo.lock | 21 +++++ Cargo.toml | 1 + src/cli/Cargo.toml | 1 + src/cli/src/data/import_v2/command.rs | 121 +++++++++++++++++++++++--- src/cli/src/data/progress.rs | 95 +++++++++++++++++++- 5 files changed, 225 insertions(+), 14 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8188a8aa494..56dbcb397da 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1977,6 +1977,7 @@ dependencies = [ "fs2", "futures", "humantime", + "indicatif", "meta-client", "meta-srv", "object-store", @@ -2996,6 +2997,7 @@ dependencies = [ "encode_unicode", "libc", "once_cell", + "unicode-width 0.2.1", "windows-sys 0.59.0", ] @@ -6809,6 +6811,19 @@ dependencies = [ "serde_core", ] +[[package]] +name = "indicatif" +version = "0.17.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "183b3088984b400f4cfac3620d5e076c84da5364016b4f49473de574b2586235" +dependencies = [ + "console", + "number_prefix", + "portable-atomic", + "unicode-width 0.2.1", + "web-time", +] + [[package]] name = "indoc" version = "2.0.6" @@ -9155,6 +9170,12 @@ dependencies = [ "libc", ] +[[package]] +name = "number_prefix" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3" + [[package]] name = "objc2" version = "0.6.4" diff --git a/Cargo.toml b/Cargo.toml index 4e95472c515..8ad88ae37c7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -166,6 +166,7 @@ humantime-serde = "1.1" hyper = "1.1" hyper-util = "0.1" icu_properties = "2.0.1" +indicatif = "0.17" itertools = "0.14" jsonb = { version = "0.4.4", default-features = false } lazy_static = "1.4" diff --git a/src/cli/Cargo.toml b/src/cli/Cargo.toml index 8c4150cd46d..27aaeed73b5 100644 --- a/src/cli/Cargo.toml +++ b/src/cli/Cargo.toml @@ -47,6 +47,7 @@ etcd-client.workspace = true fs2.workspace = true futures.workspace = true humantime.workspace = true +indicatif.workspace = true meta-client.workspace = true meta-srv.workspace = true object-store.workspace = true diff --git a/src/cli/src/data/import_v2/command.rs b/src/cli/src/data/import_v2/command.rs index 4fb84c274ce..2f8d5a7f07e 100644 --- a/src/cli/src/data/import_v2/command.rs +++ b/src/cli/src/data/import_v2/command.rs @@ -15,6 +15,7 @@ //! Import V2 CLI command. use std::collections::HashSet; +use std::io::IsTerminal; use std::time::Duration; use async_trait::async_trait; @@ -39,27 +40,72 @@ 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::progress::{IndicatifProgress, 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. +/// `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 { - /// Choose automatically (currently lightweight log progress). + /// Show an interactive bar on a TTY, otherwise log progress. #[default] Auto, - /// Always emit lightweight log progress. + /// 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 { @@ -326,12 +372,18 @@ impl Import { Ok(statements) } - /// Builds the progress reporter for this run. `never` is silent; `auto` - /// and `always` both use lightweight log progress for now. + /// 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 self.progress { - ProgressMode::Never => Box::new(NoopProgress), - ProgressMode::Auto | ProgressMode::Always => Box::new(LogProgress::new()), + 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), } } } @@ -745,6 +797,53 @@ 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 9e6945bfcd6..442cce81133 100644 --- a/src/cli/src/data/progress.rs +++ b/src/cli/src/data/progress.rs @@ -16,13 +16,15 @@ //! //! This is intentionally small and log/internal oriented. It does not touch //! 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. +//! import-v2 `--progress` flag for non-interactive runs by routing events to +//! 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::sync::Mutex; +use indicatif::{ProgressBar, ProgressDrawTarget, ProgressStyle}; + /// Receives progress events from long-running Export/Import V2 work. /// /// The trait is object-safe so callers can take `&dyn ProgressReporter` and stay @@ -118,6 +120,73 @@ impl ProgressReporter for LogProgress { } } +/// A reporter that renders an interactive progress bar via `indicatif`. +/// +/// It draws to stderr through an explicit [`ProgressDrawTarget::stderr`] so it +/// never collides with stdout (e.g. dry-run SQL). Phases with a known total get +/// a determinate bar; unknown totals fall back to an animated spinner. Each +/// phase clears itself on finish via [`ProgressBar::finish_and_clear`]. +pub(crate) struct IndicatifProgress { + bar: Mutex>, +} + +impl IndicatifProgress { + pub(crate) fn new() -> Self { + Self { + bar: Mutex::new(None), + } + } +} + +impl ProgressReporter for IndicatifProgress { + fn start_phase(&self, name: &str, total: Option) { + let Ok(mut guard) = self.bar.lock() else { + return; + }; + + // Replacing any prior phase: clear it before starting the next. + if let Some(prev) = guard.take() { + prev.finish_and_clear(); + } + + let bar = ProgressBar::with_draw_target(total, ProgressDrawTarget::stderr()); + match total { + Some(_) => { + let style = + ProgressStyle::with_template("{msg} [{bar:40}] {pos}/{len} ({percent}%)") + .unwrap_or_else(|_| ProgressStyle::default_bar()) + .progress_chars("=>-"); + bar.set_style(style); + } + None => { + let style = ProgressStyle::with_template("{spinner} {msg} {pos}") + .unwrap_or_else(|_| ProgressStyle::default_spinner()); + bar.set_style(style); + } + } + bar.set_message(name.to_string()); + *guard = Some(bar); + } + + fn inc(&self, delta: u64) { + let Ok(guard) = self.bar.lock() else { + return; + }; + if let Some(bar) = guard.as_ref() { + bar.inc(delta); + } + } + + fn finish_phase(&self) { + let Ok(mut guard) = self.bar.lock() else { + return; + }; + if let Some(bar) = guard.take() { + bar.finish_and_clear(); + } + } +} + /// RAII guard for a started progress phase. /// /// This keeps future stateful reporters safe on every early-return path after a @@ -178,6 +247,26 @@ mod tests { reporter.finish_phase(); reporter.finish_phase(); // Idempotent: finishing twice is harmless. } + + #[test] + fn test_indicatif_progress_is_safe_across_phase_lifecycle() { + // IndicatifProgress takes only `&self`, so it must survive a full + // lifecycle (including determinate and spinner phases, an out-of-phase + // `inc`, and double finish) without panicking. The draw target is + // stderr, which is non-interactive under the test harness, so nothing + // actually renders. + let progress = IndicatifProgress::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.start_phase("Streaming", None); // Spinner phase replaces the bar. + reporter.inc(5); + reporter.finish_phase(); + reporter.finish_phase(); // Idempotent: finishing twice is harmless. + } } #[cfg(test)]