From 0a2dc6d3f255e3f47abedcbeb56dd9396a10cfff Mon Sep 17 00:00:00 2001 From: jeremyhi Date: Wed, 10 Jun 2026 04:35:32 -0700 Subject: [PATCH] feat(cli): add export import progress abstraction (#8270) Signed-off-by: fengjiachun <3860496+fengjiachun@users.noreply.github.com> Co-authored-by: fengjiachun <3860496+fengjiachun@users.noreply.github.com> --- src/cli/src/data.rs | 1 + src/cli/src/data/import_v2/coordinator.rs | 114 ++++++++++++++++ src/cli/src/data/progress.rs | 150 ++++++++++++++++++++++ 3 files changed, 265 insertions(+) create mode 100644 src/cli/src/data/progress.rs diff --git a/src/cli/src/data.rs b/src/cli/src/data.rs index 114886542e8..c880bd81c73 100644 --- a/src/cli/src/data.rs +++ b/src/cli/src/data.rs @@ -17,6 +17,7 @@ pub mod export_v2; mod import; pub mod import_v2; pub(crate) mod path; +pub(crate) mod progress; pub mod snapshot_storage; pub(crate) mod sql; mod storage_export; diff --git a/src/cli/src/data/import_v2/coordinator.rs b/src/cli/src/data/import_v2/coordinator.rs index fe2d0164642..82a8585416a 100644 --- a/src/cli/src/data/import_v2/coordinator.rs +++ b/src/cli/src/data/import_v2/coordinator.rs @@ -28,6 +28,7 @@ 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}; #[async_trait] pub(crate) trait ImportTaskExecutor { @@ -130,6 +131,20 @@ pub(crate) async fn import_with_resume_session( session: ImportResumeSession, executor: &E, ) -> Result<()> +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 +} + +pub(crate) async fn import_with_resume_session_with_progress( + session: ImportResumeSession, + executor: &E, + progress: &dyn ProgressReporter, +) -> Result<()> where E: ImportTaskExecutor + Sync, { @@ -162,6 +177,14 @@ where config.state_path.display() ); + // One progress unit per import task. Already-completed tasks still count so + // the reported total matches the task plan regardless of resume position. + let progress_phase = ProgressPhase::start( + progress, + "Import data tasks", + Some(config.tasks.len() as u64), + ); + let import_start = Instant::now(); for (idx, task) in config.tasks.iter().enumerate() { if state.task_status(task.chunk_id, &task.schema) == Some(ImportTaskStatus::Completed) { @@ -172,6 +195,7 @@ where task.chunk_id, task.schema ); + progress.inc(1); continue; } @@ -216,6 +240,7 @@ where task.schema, task_start.elapsed() ); + progress.inc(1); } Err(task_error) => { // Persist Failed best-effort, but always surface the original @@ -240,6 +265,7 @@ where } } + progress_phase.finish(); delete_import_state(&config.state_path).await?; info!("Data import finished in {:?}", import_start.elapsed()); drop(lock); @@ -365,6 +391,7 @@ mod tests { use super::*; use crate::data::export_v2::manifest::{ChunkMeta, TimeRange}; use crate::data::import_v2::error::TestTaskFailedSnafu; + use crate::data::progress::test_util::{ProgressEvent, RecordingProgress}; #[derive(Debug, Clone, Copy)] enum FailureMode { @@ -692,4 +719,91 @@ mod tests { Some(ImportTaskStatus::Failed) ); } + + #[tokio::test] + async fn test_import_progress_counts_skipped_and_completed_tasks() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("import_state.json"); + let tasks = vec![ + ImportTaskKey::new(1, "public"), + ImportTaskKey::new(2, "analytics"), + ]; + let mut state = ImportState::new( + "snapshot-1", + "127.0.0.1:4000", + "greptime", + &["public".to_string(), "analytics".to_string()], + tasks.clone(), + ); + state.mark_ddl_completed(); + // First task is already completed: it should still count as one progress + // unit so the reported total matches the task plan on resume. + state + .set_task_status(1, "public", ImportTaskStatus::Completed, None) + .unwrap(); + save_import_state(&path, &state).await.unwrap(); + + let imported = Arc::new(Mutex::new(Vec::new())); + let executor = recording_executor(imported.clone()); + let progress = RecordingProgress::default(); + + let session = prepare_import_resume(config(path, tasks)).await.unwrap(); + import_with_resume_session_with_progress(session, &executor, &progress) + .await + .unwrap(); + + // The skipped-completed task is not re-imported. + assert_eq!( + imported.lock().unwrap().clone(), + vec![ImportTaskKey::new(2, "analytics")] + ); + + let events = progress.events(); + assert_eq!( + events.first(), + Some(&ProgressEvent::StartPhase { + name: "Import data tasks".to_string(), + total: Some(2), + }) + ); + assert_eq!(events.last(), Some(&ProgressEvent::FinishPhase)); + // Both the skipped-completed task and the newly imported task increment. + assert_eq!(progress.total_inc(), 2); + } + + #[tokio::test] + async fn test_import_progress_stops_incrementing_on_task_failure() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("import_state.json"); + let failed_task = ImportTaskKey::new(1, "public"); + let tasks = vec![failed_task.clone(), ImportTaskKey::new(2, "analytics")]; + let executor = RecordingExecutor { + imported: Arc::new(Mutex::new(Vec::new())), + fail_task: Some(failed_task.clone()), + failure_mode: Some(FailureMode::Fatal), + attempts: Arc::new(AtomicUsize::new(0)), + }; + + let mut state = ImportState::new( + "snapshot-1", + "127.0.0.1:4000", + "greptime", + &["public".to_string(), "analytics".to_string()], + tasks.clone(), + ); + state.mark_ddl_completed(); + save_import_state(&path, &state).await.unwrap(); + + let progress = RecordingProgress::default(); + let session = prepare_import_resume(config(path, tasks)).await.unwrap(); + import_with_resume_session_with_progress(session, &executor, &progress) + .await + .unwrap_err(); + + // The first task fails, so neither task increments. The phase is still + // finished before returning the task error so future UI reporters can + // clean up their state. + assert_eq!(progress.total_inc(), 0); + assert!(progress.events().contains(&ProgressEvent::FinishPhase)); + } } diff --git a/src/cli/src/data/progress.rs b/src/cli/src/data/progress.rs new file mode 100644 index 00000000000..0851d0a0959 --- /dev/null +++ b/src/cli/src/data/progress.rs @@ -0,0 +1,150 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! 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. + +/// Receives progress events from long-running Export/Import V2 work. +/// +/// The trait is object-safe so callers can take `&dyn ProgressReporter` and stay +/// agnostic about the concrete implementation (no-op in production today, a +/// recording fake in tests). +pub(crate) trait ProgressReporter: Send + Sync { + /// Begins a phase with an optional known total number of units. + fn start_phase(&self, name: &str, total: Option); + + /// Advances the current phase by `delta` units. + fn inc(&self, delta: u64); + + /// Marks the current phase as finished. + fn finish_phase(&self); +} + +/// A reporter that discards every event. Used as the production default and in +/// tests that do not care about progress. +pub(crate) struct NoopProgress; + +impl ProgressReporter for NoopProgress { + fn start_phase(&self, _name: &str, _total: Option) {} + fn inc(&self, _delta: u64) {} + fn finish_phase(&self) {} +} + +/// RAII guard for a started progress phase. +/// +/// This keeps future stateful reporters safe on every early-return path after a +/// phase starts. Call [`Self::finish`] to end the phase at a deliberate point; +/// otherwise the guard finishes it when dropped. +#[must_use = "dropping the guard immediately finishes the phase"] +pub(crate) struct ProgressPhase<'a> { + reporter: &'a dyn ProgressReporter, + finished: bool, +} + +impl<'a> ProgressPhase<'a> { + pub(crate) fn start( + reporter: &'a dyn ProgressReporter, + name: &str, + total: Option, + ) -> Self { + reporter.start_phase(name, total); + Self { + reporter, + finished: false, + } + } + + pub(crate) fn finish(mut self) { + self.finish_once(); + } + + fn finish_once(&mut self) { + if !self.finished { + self.reporter.finish_phase(); + self.finished = true; + } + } +} + +impl Drop for ProgressPhase<'_> { + fn drop(&mut self) { + self.finish_once(); + } +} + +#[cfg(test)] +pub(crate) mod test_util { + use std::sync::Mutex; + + use super::ProgressReporter; + + /// A single recorded progress event, used to assert progress behavior in + /// unit tests. + #[derive(Debug, Clone, PartialEq, Eq)] + pub(crate) enum ProgressEvent { + StartPhase { name: String, total: Option }, + Inc { delta: u64 }, + FinishPhase, + } + + /// A reporter that records every event in order for later assertions. + #[derive(Default)] + pub(crate) struct RecordingProgress { + events: Mutex>, + } + + impl RecordingProgress { + pub(crate) fn events(&self) -> Vec { + self.events.lock().unwrap().clone() + } + + /// Sum of all `inc` deltas recorded so far. + pub(crate) fn total_inc(&self) -> u64 { + self.events + .lock() + .unwrap() + .iter() + .filter_map(|event| match event { + ProgressEvent::Inc { delta } => Some(*delta), + _ => None, + }) + .sum() + } + + fn push(&self, event: ProgressEvent) { + self.events.lock().unwrap().push(event); + } + } + + impl ProgressReporter for RecordingProgress { + fn start_phase(&self, name: &str, total: Option) { + self.push(ProgressEvent::StartPhase { + name: name.to_string(), + total, + }); + } + + fn inc(&self, delta: u64) { + self.push(ProgressEvent::Inc { delta }); + } + + fn finish_phase(&self) { + self.push(ProgressEvent::FinishPhase); + } + } +}