1use std::path::{Path, PathBuf};
11use std::sync::Arc;
12use std::sync::atomic::{AtomicU8, Ordering};
13use std::sync::mpsc::Sender;
14use std::time::Duration;
15
16use dre_protocol::host::LogSink;
17use serde_json::{Map as JsonMap, Value as Json};
18
19use crate::run::{BindingOutcome, BindingPlan, Level, RunTargets, ShownMessage, Ui};
20
21#[derive(Debug, Clone)]
23pub enum RunEvent {
24 Targets(RunTargets),
26 Plan(usize),
28 BindingStart {
29 report: String,
30 set: Option<String>,
31 },
32 BindingLog(PathBuf),
34 BindingVars {
36 schedule: Option<String>,
37 schedule_vars: Option<JsonMap<String, Json>>,
38 vars: JsonMap<String, Json>,
39 },
40 Step {
42 level: Level,
43 verb: String,
44 detail: String,
45 elapsed: Option<Duration>,
46 },
47 Warn(String),
48 BindingEnd(BindingOutcome),
49 Compiled(BindingPlan),
51 Message(ShownMessage),
53 PluginLog(String, String),
55 Sql(String, String),
57 Running {
60 label: String,
61 running: bool,
62 },
63 Rows(crate::run::RowProgress),
65 Block(Vec<RunEvent>),
68}
69
70#[derive(Clone)]
73pub struct Events {
74 tx: Sender<RunEvent>,
75 held: Option<std::sync::Arc<std::sync::Mutex<Vec<RunEvent>>>>,
77}
78
79impl Events {
80 pub fn new(sender: Sender<RunEvent>) -> Events {
81 Events {
82 tx: sender,
83 held: None,
84 }
85 }
86
87 pub fn held(&self) -> Events {
89 Events {
90 tx: self.tx.clone(),
91 held: Some(Default::default()),
92 }
93 }
94
95 pub fn emit(&self, e: RunEvent) {
97 match &self.held {
98 Some(h) => h.lock().unwrap().push(e),
99 None => {
100 let _ = self.tx.send(e);
101 }
102 }
103 }
104
105 pub fn emit_now(&self, e: RunEvent) {
107 let _ = self.tx.send(e);
108 }
109
110 pub fn flush(&self) {
112 if let Some(h) = &self.held {
113 let events = std::mem::take(&mut *h.lock().unwrap());
114 if !events.is_empty() {
115 let _ = self.tx.send(RunEvent::Block(events));
116 }
117 }
118 }
119
120 pub fn step(&self, level: Level, verb: &str, detail: &str, elapsed: Option<Duration>) {
121 self.emit(RunEvent::Step {
122 level,
123 verb: verb.to_string(),
124 detail: detail.to_string(),
125 elapsed,
126 });
127 }
128
129 pub fn warn(&self, msg: &str) {
130 self.emit(RunEvent::Warn(msg.to_string()));
131 }
132
133 pub fn rows(&self, progress: crate::run::RowProgress) {
135 self.emit_now(RunEvent::Rows(progress));
136 }
137
138 pub fn compiled(&self, plan: &BindingPlan) {
139 self.emit(RunEvent::Compiled(plan.clone()));
140 }
141
142 pub fn message(&self, message: &ShownMessage) {
143 self.emit(RunEvent::Message(message.clone()));
144 }
145
146 pub fn binding_vars(
147 &self,
148 schedule: Option<&str>,
149 schedule_vars: Option<&JsonMap<String, Json>>,
150 vars: &JsonMap<String, Json>,
151 ) {
152 self.emit(RunEvent::BindingVars {
153 schedule: schedule.map(str::to_string),
154 schedule_vars: schedule_vars.cloned(),
155 vars: vars.clone(),
156 });
157 }
158
159 pub fn plugin_log(&self) -> LogSink {
161 let events = self.clone();
162 Arc::new(move |plugin: &str, line: &str| events.emit(RunEvent::PluginLog(plugin.into(), line.into())))
163 }
164
165 pub fn sql_log(&self) -> LogSink {
167 let events = self.clone();
168 Arc::new(move |label: &str, sql: &str| events.emit(RunEvent::Sql(label.into(), sql.into())))
169 }
170}
171
172pub fn dispatch(ui: &mut dyn Ui, e: RunEvent, plugin_log: &LogSink, sql_log: &LogSink) {
175 match e {
176 RunEvent::Targets(t) => ui.targets(&t),
177 RunEvent::Plan(n) => ui.plan(n),
178 RunEvent::BindingStart { report, set } => ui.binding_start(&report, set.as_deref()),
179 RunEvent::BindingLog(path) => ui.binding_log(&path),
180 RunEvent::BindingVars {
181 schedule,
182 schedule_vars,
183 vars,
184 } => ui.binding_vars(schedule.as_deref(), schedule_vars.as_ref(), &vars),
185 RunEvent::Step {
186 level,
187 verb,
188 detail,
189 elapsed,
190 } => ui.step(level, &verb, &detail, elapsed),
191 RunEvent::Warn(msg) => ui.warn(&msg),
192 RunEvent::BindingEnd(o) => ui.binding_end(&o),
193 RunEvent::Compiled(p) => ui.compiled(&p),
194 RunEvent::Message(m) => ui.message(&m),
195 RunEvent::PluginLog(plugin, line) => plugin_log(&plugin, &line),
196 RunEvent::Sql(label, sql) => sql_log(&label, &sql),
197 RunEvent::Running { label, running } => ui.running(&label, running),
198 RunEvent::Rows(p) => ui.rows(&p),
199 RunEvent::Block(events) => {
200 for e in events {
201 dispatch(ui, e, plugin_log, sql_log);
202 }
203 }
204 }
205}
206
207#[derive(Debug, Clone, Copy, PartialEq, Eq)]
209pub enum CancelReason {
210 Interrupt,
212 Terminate,
214 Timeout,
216}
217
218impl CancelReason {
219 pub fn exit_code(self) -> u8 {
222 match self {
223 CancelReason::Interrupt => 130,
224 CancelReason::Terminate => 143,
225 CancelReason::Timeout => 124,
226 }
227 }
228
229 fn code(self) -> u8 {
230 match self {
231 CancelReason::Interrupt => 1,
232 CancelReason::Terminate => 2,
233 CancelReason::Timeout => 3,
234 }
235 }
236}
237
238#[derive(Debug, Clone, Default)]
240pub struct CancelToken(Arc<AtomicU8>);
241
242impl CancelToken {
243 pub fn new() -> CancelToken {
244 CancelToken::default()
245 }
246
247 pub fn cancel(&self) {
249 self.cancel_with(CancelReason::Interrupt);
250 }
251
252 pub fn cancel_with(&self, reason: CancelReason) {
253 let _ = self
254 .0
255 .compare_exchange(0, reason.code(), Ordering::SeqCst, Ordering::SeqCst);
256 }
257
258 pub fn is_cancelled(&self) -> bool {
259 self.reason().is_some()
260 }
261
262 pub fn reason(&self) -> Option<CancelReason> {
263 match self.0.load(Ordering::SeqCst) {
264 1 => Some(CancelReason::Interrupt),
265 2 => Some(CancelReason::Terminate),
266 3 => Some(CancelReason::Timeout),
267 _ => None,
268 }
269 }
270}
271
272#[derive(Debug, Clone)]
276pub struct RunStore {
277 target: PathBuf,
278 claims: Option<std::sync::Arc<std::sync::Mutex<std::collections::HashMap<String, String>>>>,
281}
282
283impl RunStore {
284 pub fn new(target: &Path) -> RunStore {
285 RunStore {
286 target: target.to_path_buf(),
287 claims: None,
288 }
289 }
290
291 pub fn concurrent(target: &Path) -> RunStore {
293 RunStore {
294 target: target.to_path_buf(),
295 claims: Some(Default::default()),
296 }
297 }
298
299 pub fn claim(&self, profile: &str, path: &str, binding: &str) -> Option<String> {
301 let mut claims = self.claims.as_ref()?.lock().unwrap();
302 let owner = claims
303 .entry(format!("{profile}\u{0}{path}"))
304 .or_insert_with(|| binding.to_string());
305 (owner != binding).then(|| owner.clone())
306 }
307
308 fn binding(&self, kind: &str, report: &str, binding: &str) -> PathBuf {
309 self.target.join(kind).join(report).join(binding)
310 }
311
312 pub fn compiled_dir(&self, report: &str, binding: &str) -> PathBuf {
314 self.binding("compiled", report, binding)
315 }
316
317 pub fn run_dir(&self, report: &str, binding: &str) -> PathBuf {
319 self.binding("run", report, binding)
320 }
321
322 pub fn schema_dir(&self, report: &str, binding: &str) -> PathBuf {
324 self.binding("schema", report, binding)
325 }
326
327 pub fn write_results(
329 &self,
330 run_dir: &Path,
331 results: &crate::run_results::RunResults,
332 ) -> std::io::Result<()> {
333 std::fs::create_dir_all(run_dir)?;
334 crate::fs::write_atomic(
335 &run_dir.join("run_results.json"),
336 (crate::secrets::to_json_pretty(results)? + "\n").as_bytes(),
337 )
338 }
339}
340
341#[cfg(test)]
342mod tests {
343 use super::*;
344
345 #[test]
346 fn a_cancel_is_seen_by_every_clone() {
347 let a = CancelToken::new();
348 let b = a.clone();
349 assert!(!b.is_cancelled());
350 a.cancel_with(CancelReason::Terminate);
351 assert!(b.is_cancelled());
352 b.cancel();
354 assert_eq!(a.reason(), Some(CancelReason::Terminate));
355 assert_eq!(CancelReason::Terminate.exit_code(), 143);
356 }
357
358 #[test]
359 fn the_store_lays_out_a_bindings_folders() {
360 let s = RunStore::new(Path::new("/t"));
361 assert_eq!(s.run_dir("daily", "default"), Path::new("/t/run/daily/default"));
362 assert_eq!(s.compiled_dir("daily", "eu"), Path::new("/t/compiled/daily/eu"));
363 assert_eq!(s.schema_dir("daily", "eu"), Path::new("/t/schema/daily/eu"));
364 }
365
366 #[test]
367 fn events_arrive_in_order_and_log_sinks_feed_the_stream() {
368 let (tx, rx) = std::sync::mpsc::channel();
369 let events = Events::new(tx);
370 events.warn("one");
371 events.sql_log()("q.sql:1", "select 1");
372 events.plugin_log()("duckdb", "hello");
373 drop(events);
374 let got: Vec<String> = rx.iter().map(|e| format!("{e:?}")).collect();
375 assert_eq!(got.len(), 3);
376 assert!(got[0].starts_with("Warn") && got[1].starts_with("Sql") && got[2].starts_with("PluginLog"));
377 }
378}