Skip to main content

akar_common/
progress_bar.rs

1use indicatif::{ProgressBar as IndiProgressBar, ProgressStyle};
2use std::sync::Arc;
3use std::sync::atomic::{AtomicBool, Ordering};
4use std::time::Duration;
5
6/// A thin wrapper around `indicatif::ProgressBar` for Akar operations.
7///
8/// Provides factory methods for common progress styles (spinner, count-based,
9/// bytes-transferred) used during bulk ingest, export, and long-running queries.
10///
11/// # Example
12///
13/// ```no_run
14/// use akar_common::progress_bar::AkarProgress;
15///
16/// let pb = AkarProgress::new("Loading data…", Some(1000));
17/// for i in 0..1000 {
18///     pb.inc();
19///     std::thread::sleep(std::time::Duration::from_millis(1));
20/// }
21/// pb.finish("Done.");
22/// ```
23pub struct AkarProgress {
24    inner: Option<IndiProgressBar>,
25    cancelled: Arc<AtomicBool>,
26}
27
28impl AkarProgress {
29    /// Create a new progress bar with an optional total count.
30    ///
31    /// When `total` is `None`, a spinner is used (indeterminate progress).
32    /// When `total` is `Some(n)`, a count-based bar is shown.
33    pub fn new(msg: &str, total: Option<u64>) -> Self {
34        let cancelled = Arc::new(AtomicBool::new(false));
35        let inner = match total {
36            Some(n) => {
37                let pb = IndiProgressBar::new(n);
38                pb.set_style(
39                    ProgressStyle::default_bar()
40                        .template("{msg} [{bar:40}] {pos}/{len} ({eta})")
41                        .unwrap()
42                        .progress_chars("=> "),
43                );
44                pb.set_message(msg.to_owned());
45                Some(pb)
46            }
47            None => {
48                let pb = IndiProgressBar::new_spinner();
49                pb.set_style(
50                    ProgressStyle::default_spinner()
51                        .template("{spinner:.green} {msg}")
52                        .unwrap(),
53                );
54                pb.set_message(msg.to_owned());
55                pb.enable_steady_tick(Duration::from_millis(100));
56                Some(pb)
57            }
58        };
59        Self { inner, cancelled }
60    }
61
62    /// Advance the progress bar by one step.
63    pub fn inc(&self) {
64        if let Some(ref pb) = self.inner {
65            pb.inc(1);
66        }
67    }
68
69    /// Advance by `delta` steps.
70    pub fn inc_by(&self, delta: u64) {
71        if let Some(ref pb) = self.inner {
72            pb.inc(delta);
73        }
74    }
75
76    /// Set the current position.
77    pub fn set_pos(&self, pos: u64) {
78        if let Some(ref pb) = self.inner {
79            pb.set_position(pos);
80        }
81    }
82
83    /// Update the message displayed alongside the bar.
84    pub fn set_message(&self, msg: &str) {
85        if let Some(ref pb) = self.inner {
86            pb.set_message(msg.to_owned());
87        }
88    }
89
90    /// Mark as finished with a final message.
91    pub fn finish(&self, msg: &str) {
92        if let Some(ref pb) = self.inner {
93            pb.finish_with_message(msg.to_owned());
94        }
95    }
96
97    /// Abort and clear the progress display.
98    pub fn abort(&self) {
99        if let Some(ref pb) = self.inner {
100            pb.abandon();
101        }
102    }
103
104    /// Mark as cancelled. Callers should check [`Self::is_cancelled`] in their
105    /// work loop to abort early.
106    pub fn cancel(&self) {
107        self.cancelled.store(true, Ordering::Relaxed);
108        self.finish("Cancelled.");
109    }
110
111    /// Whether a cancellation has been requested.
112    pub fn is_cancelled(&self) -> bool {
113        self.cancelled.load(Ordering::Relaxed)
114    }
115
116    /// Get a reference to the cancellation flag for sharing across threads.
117    pub fn cancelled_flag(&self) -> Arc<AtomicBool> {
118        self.cancelled.clone()
119    }
120}
121
122impl Drop for AkarProgress {
123    fn drop(&mut self) {
124        // If the progress bar was not explicitly finished, clear it.
125        if let Some(ref pb) = self.inner {
126            pb.finish_and_clear();
127        }
128    }
129}