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

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

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<E>(
     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<E>(
+    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));
+    }
 }

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<u64>);
+
+    /// 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<u64>) {}
+    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<u64>,
+    ) -> 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<u64> },
+        Inc { delta: u64 },
+        FinishPhase,
+    }
+
+    /// A reporter that records every event in order for later assertions.
+    #[derive(Default)]
+    pub(crate) struct RecordingProgress {
+        events: Mutex<Vec<ProgressEvent>>,
+    }
+
+    impl RecordingProgress {
+        pub(crate) fn events(&self) -> Vec<ProgressEvent> {
+            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<u64>) {
+            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);
+        }
+    }
+}