asimov_dataset_cli/ui/
mod.rs1use 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 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 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}