rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-07-06 14:00:40 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(cli): add export-v2 progress reporting (#8294)

Browse Source 
Signed-off-by: jeremyhi <fengjiachun@gmail.com>
This commit is contained in:
jeremyhi
2026-06-15 02:56:31 -07:00
committed by GitHub
parent 15597047c9
commit 90606af070
4 changed files with 253 additions and 148 deletions
Download Patch File Download Diff File Expand all files Collapse all files

88
src/cli/src/data/export_v2/command.rs
View File

@@ -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<Duration>,
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(),
};

58
src/cli/src/data/export_v2/coordinator.rs
View File

@@ -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 {

129
src/cli/src/data/import_v2/command.rs
View File

@@ -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<dyn ProgressReporter> {
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!(

126
src/cli/src/data/progress.rs
View File

@@ -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<dyn ProgressReporter> {
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