asimov_dataset_cli/ui/
mod.rs

1// This is free and unencumbered software released into the public domain.
2
3use crossbeam::channel::Receiver;
4use eyre::Result;
5
6mod prepare;
7mod publish;
8
9use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget, ProgressStyle};
10pub use prepare::{PrepareProgress, PrepareState, ReaderProgress};
11pub use publish::{PublishProgress, PublishState};
12
13pub enum UIEvent {
14    Resize,
15}
16
17#[derive(Debug)]
18pub enum Event {
19    Reader(ReaderProgress),
20    Prepare(PrepareProgress),
21    Publish(PublishProgress),
22}
23
24pub fn run_prepare(
25    verbosity: u8,
26    mut state: PrepareState,
27    progress_rx: Receiver<Event>,
28) -> Result<()> {
29    let parsing_style =
30        ProgressStyle::with_template("{msg:10} [{bar:40}] {binary_bytes} / {binary_total_bytes}")
31            .unwrap()
32            .progress_chars("##-");
33
34    let prepare_style =
35        ProgressStyle::with_template("{msg:10} [{bar:40}] {human_pos} / {human_len} statements")
36            .unwrap()
37            .progress_chars("##-");
38
39    let multi = MultiProgress::new();
40    if verbosity < 1 {
41        // only show bars for `-v`
42        multi.set_draw_target(ProgressDrawTarget::hidden());
43    }
44    let reader_bar = ProgressBar::new(state.total_bytes as u64)
45        .with_message("Read")
46        .with_style(parsing_style);
47    let prepare_bar = ProgressBar::new(0)
48        .with_message("Batch")
49        .with_style(prepare_style);
50
51    multi.add(reader_bar.clone());
52    multi.add(prepare_bar.clone());
53
54    while let Ok(event) = progress_rx.recv() {
55        tracing::debug!(?event);
56
57        match event {
58            Event::Reader(progress) => {
59                reader_bar.inc(progress.bytes as u64);
60                prepare_bar.inc_length(progress.statement_count as u64);
61                if progress.finished && verbosity > 1 {
62                    multi.println(format!(
63                        " ✅ Finished reading file {}",
64                        progress.filename.display()
65                    ))?;
66                }
67                state.update_reader_state(progress);
68            }
69            Event::Prepare(progress) => {
70                prepare_bar.inc(progress.statement_count as u64);
71                if verbosity > 1 {
72                    if let Some(filename) = progress
73                        .filename
74                        .file_name()
75                        .and_then(std::ffi::OsStr::to_str)
76                    {
77                        multi.println(format!(" ✅ Created batch {}", filename))?;
78                    }
79                }
80                state.update_prepare_state(progress);
81            }
82            Event::Publish(_) => unreachable!(),
83        }
84    }
85
86    reader_bar.finish();
87    prepare_bar.finish();
88
89    Ok(())
90}
91
92pub fn run_publish(
93    verbosity: u8,
94    mut state: PublishState,
95    progress_rx: Receiver<Event>,
96) -> Result<()> {
97    let parsing_style =
98        ProgressStyle::with_template("{msg:10} [{bar:40}] {binary_bytes} / {binary_total_bytes}")
99            .unwrap()
100            .progress_chars("##-");
101
102    let prepare_style =
103        ProgressStyle::with_template("{msg:10} [{bar:40}] {human_pos} / {human_len} statements")
104            .unwrap()
105            .progress_chars("##-");
106
107    let upload_style =
108        ProgressStyle::with_template("{msg:10} [{bar:40}] {human_pos} / {human_len} batches")
109            .unwrap()
110            .progress_chars("##-");
111
112    let multi = MultiProgress::new();
113    if verbosity < 1 {
114        // only show bars for `-v`
115        multi.set_draw_target(ProgressDrawTarget::hidden());
116    }
117
118    let reader_bar = multi.add(
119        ProgressBar::new(
120            state
121                .prepare
122                .as_ref()
123                .map(|state| state.total_bytes)
124                .unwrap_or_default() as u64,
125        )
126        .with_message("Read")
127        .with_style(parsing_style),
128    );
129    let prepare_bar = multi.add(
130        ProgressBar::new(0)
131            .with_message("Batch")
132            .with_style(prepare_style),
133    );
134    let upload_bar = multi.add(
135        ProgressBar::new(0)
136            .with_message("Upload")
137            .with_style(upload_style),
138    );
139
140    while let Ok(event) = progress_rx.recv() {
141        tracing::debug!(?event);
142
143        match event {
144            Event::Reader(progress) => {
145                reader_bar.inc(progress.bytes as u64);
146                prepare_bar.inc_length(progress.statement_count as u64);
147                if progress.finished && verbosity > 1 {
148                    multi.println(format!(
149                        " ✅ Finished reading file {}",
150                        progress.filename.display()
151                    ))?;
152                }
153                if let Some(ref mut state) = state.prepare {
154                    state.update_reader_state(progress);
155                }
156            }
157            Event::Prepare(progress) => {
158                prepare_bar.inc(progress.statement_count as u64);
159                upload_bar.inc_length(1);
160                if verbosity > 1 {
161                    if let Some(filename) = progress
162                        .filename
163                        .file_name()
164                        .and_then(std::ffi::OsStr::to_str)
165                    {
166                        multi.println(format!(" ✅ Created batch {}", filename))?;
167                    }
168                }
169                if let Some(ref mut state) = state.prepare {
170                    state.update_prepare_state(progress);
171                }
172            }
173            Event::Publish(progress) => {
174                upload_bar.inc(1);
175                if verbosity > 1 {
176                    if let Some(filename) = progress
177                        .filename
178                        .file_name()
179                        .and_then(std::ffi::OsStr::to_str)
180                    {
181                        multi.println(format!(" ✅ Uploaded batch {}", filename))?;
182                    }
183                }
184                state.update_publish_state(progress);
185            }
186        }
187    }
188
189    reader_bar.finish();
190    prepare_bar.finish();
191    upload_bar.finish();
192
193    Ok(())
194}