1mod corpora_aggregate;
12mod export;
13mod mark;
14mod reports;
15mod rollup;
16mod sandbox;
17mod services_aggregate;
18mod tasks_aggregate;
19pub use export::{DatasetExportOutcome, GroupBy, export_html_dataset};
20pub(crate) use mark::{
21 mark_all_blocked, mark_blocked, mark_rerun, resume_all_blocked, resume_blocked,
22 save_historical_tasks,
23};
24pub use reports::list_task_diffs;
28pub(crate) use reports::live_run_diff;
29pub(crate) use reports::progress_report;
30pub(crate) use reports::report_uses_rollup;
31pub use reports::TaskReportOptions;
35pub use reports::summary_task_diffs;
36pub(crate) use reports::task_report;
37pub use reports::{DOCUMENT_MESSAGE_CAP, MessageCounts, task_messages};
38pub use rollup::ReportSummaryRow;
39pub(crate) use rollup::{
40 category_rollup, category_total, invalidate_all, invalidate_scope, populate_scope,
41 populate_scope_bounded, report_cache_computed_at, scope_cached, severity_total, what_rollup,
42};
43pub use sandbox::{SandboxOutcome, SandboxSelection, create_sandbox};
44
45use diesel::r2d2::{ConnectionManager, Pool, PooledConnection};
46use diesel::result::Error;
47use diesel::*;
48use std::collections::HashMap;
49use std::fmt;
50
51use crate::concerns::{CortexDeletable, CortexInsertable};
52use crate::config::config;
53use crate::helpers::{TaskReport, TaskStatus};
54use crate::models::{Corpus, NewTask, Service, Task};
55
56pub fn default_db_address() -> &'static str { &config().database.url }
58pub fn test_db_address() -> &'static str { &config().database.test_url }
60
61pub struct Backend {
63 pub connection: PgConnection,
65}
66impl fmt::Debug for Backend {
67 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { f.write_str("<Backend omitted>") }
68}
69impl Default for Backend {
70 fn default() -> Self {
71 let connection = connection_at(default_db_address());
72 Backend { connection }
73 }
74}
75
76pub fn connection_at(address: &str) -> PgConnection {
82 const MAX_ATTEMPTS: u32 = 5;
83 let mut attempt = 1;
84 loop {
85 match PgConnection::establish(address) {
86 Ok(connection) => return connection,
87 Err(e) if attempt < MAX_ATTEMPTS => {
88 eprintln!(
89 "-- connect to {address} failed (attempt {attempt}/{MAX_ATTEMPTS}): {e}; retrying"
90 );
91 std::thread::sleep(std::time::Duration::from_millis(100 * u64::from(attempt)));
92 attempt += 1;
93 },
94 Err(e) => panic!("Error connecting to {address} after {MAX_ATTEMPTS} attempts: {e}"),
95 }
96 }
97}
98
99pub type DbPool = Pool<ConnectionManager<PgConnection>>;
101pub type PooledConn = PooledConnection<ConnectionManager<PgConnection>>;
103
104pub fn build_pool(database_url: &str, max_size: u32) -> DbPool {
107 let manager = ConnectionManager::<PgConnection>::new(database_url);
108 Pool::builder().max_size(max_size).build_unchecked(manager)
109}
110
111pub struct DatabaseUrl(pub String);
115pub fn testdb() -> Backend {
117 Backend {
118 connection: connection_at(test_db_address()),
119 }
120}
121pub fn from_address(address: &str) -> Backend {
123 Backend {
124 connection: connection_at(address),
125 }
126}
127
128pub struct RerunOptions<'a> {
130 pub corpus: &'a Corpus,
132 pub service: &'a Service,
134 pub severity_opt: Option<String>,
136 pub category_opt: Option<String>,
138 pub what_opt: Option<String>,
140 pub owner_opt: Option<String>,
142 pub description_opt: Option<String>,
144}
145
146impl Backend {
148 pub fn mark_imported(&mut self, imported_tasks: &[NewTask]) -> Result<usize, Error> {
151 mark::mark_imported(&mut self.connection, imported_tasks)
152 }
153 pub fn mark_done(&mut self, reports: &[TaskReport]) -> Result<(), Error> {
156 mark::mark_done(&mut self.connection, reports)
157 }
158 pub fn mark_rerun(&mut self, options: RerunOptions) -> Result<(), Error> {
161 mark::mark_rerun(&mut self.connection, options)
162 }
163
164 pub fn pause_run(&mut self, corpus_id: i32, service_id: i32) -> Result<usize, Error> {
168 mark::mark_blocked(&mut self.connection, corpus_id, service_id)
169 }
170
171 pub fn resume_run(&mut self, corpus_id: i32, service_id: i32) -> Result<usize, Error> {
175 mark::resume_blocked(&mut self.connection, corpus_id, service_id)
176 }
177
178 pub fn mark_new_run(
180 &mut self,
181 corpus: &Corpus,
182 service: &Service,
183 owner: String,
184 description: String,
185 ) -> Result<(), Error> {
186 mark::mark_new_run(&mut self.connection, corpus, service, owner, description)
187 }
188
189 pub fn save_historical_tasks(
191 &mut self,
192 corpus: &Corpus,
193 service: &Service,
194 ) -> Result<usize, Error> {
195 mark::save_historical_tasks(&mut self.connection, corpus, service)
196 }
197
198 pub fn delete<Model: CortexDeletable>(&mut self, object: &Model) -> Result<usize, Error> {
200 object.delete_by(&mut self.connection, "id")
201 }
202 pub fn delete_by<Model: CortexDeletable>(
204 &mut self,
205 object: &Model,
206 field: &str,
207 ) -> Result<usize, Error> {
208 object.delete_by(&mut self.connection, field)
209 }
210 pub fn add<Model: CortexInsertable>(&mut self, object: &Model) -> Result<usize, Error> {
214 object.create(&mut self.connection)
215 }
216
217 pub fn fetch_tasks(&mut self, service: &Service, limit: usize) -> Result<Vec<Task>, Error> {
219 tasks_aggregate::fetch_tasks(&mut self.connection, service, limit)
220 }
221 pub fn clear_limbo_tasks(&mut self) -> Result<usize, Error> {
225 tasks_aggregate::clear_limbo_tasks(&mut self.connection)
226 }
227
228 pub fn clear_limbo_tasks_except(&mut self, in_flight: &[i64]) -> Result<usize, Error> {
232 tasks_aggregate::clear_limbo_tasks_except(&mut self.connection, in_flight)
233 }
234
235 pub fn register_service(
242 &mut self,
243 service: &Service,
244 corpus: &Corpus,
245 owner: String,
246 description: String,
247 ) -> Result<(), Error> {
248 services_aggregate::register_service(&mut self.connection, service, corpus, owner, description)
249 }
250
251 pub fn extend_service(&mut self, service: &Service, corpus: &Corpus) -> Result<(), Error> {
256 services_aggregate::extend_service(&mut self.connection, service, corpus)
257 }
258
259 pub fn destroy_service_by_name(&mut self, name: &str) -> Result<usize, Error> {
263 services_aggregate::destroy_service_by_name(&mut self.connection, name)
264 }
265
266 pub fn corpora(&mut self) -> Vec<Corpus> { corpora_aggregate::list_corpora(&mut self.connection) }
268
269 pub fn tasks(
271 &mut self,
272 corpus: &Corpus,
273 service: &Service,
274 task_status: &TaskStatus,
275 ) -> Vec<Task> {
276 reports::list_tasks(&mut self.connection, corpus, service, task_status)
277 }
278 pub fn entries(
280 &mut self,
281 corpus: &Corpus,
282 service: &Service,
283 task_status: &TaskStatus,
284 ) -> Vec<String> {
285 reports::list_entries(&mut self.connection, corpus, service, task_status)
286 }
287
288 pub fn task_report(&mut self, options: TaskReportOptions) -> Vec<HashMap<String, String>> {
291 reports::task_report(&mut self.connection, options)
292 }
293 pub fn progress_report(&mut self, corpus: &Corpus, service: &Service) -> HashMap<String, f64> {
295 reports::progress_report(&mut self.connection, corpus.id, service.id)
296 }
297
298 pub fn refresh_report_summary(&mut self) -> Result<(), Error> {
301 rollup::refresh_report_summary(&mut self.connection)
302 }
303 pub fn invalidate_report_cache(&mut self, corpus_id: i32, service_id: i32) -> Result<(), Error> {
307 rollup::invalidate_scope(&mut self.connection, corpus_id, service_id)
308 }
309 pub fn complete_run_if_drained(
315 &mut self,
316 corpus_id: i32,
317 service_id: i32,
318 ) -> Result<bool, Error> {
319 let closed = crate::models::HistoricalRun::complete_if_drained(
320 corpus_id,
321 service_id,
322 &mut self.connection,
323 )?;
324 if closed {
325 if let Err(e) = mark::snapshot_tasks(&mut self.connection, corpus_id, service_id) {
330 tracing::warn!(
331 corpus_id,
332 service_id,
333 error = ?e,
334 "run-completion-on-drain: baseline snapshot failed (non-fatal)"
335 );
336 }
337 }
338 Ok(closed)
339 }
340 pub fn category_rollup(
343 &mut self,
344 corpus: &Corpus,
345 service: &Service,
346 severity: &str,
347 limit: i64,
348 offset: i64,
349 ) -> Vec<ReportSummaryRow> {
350 rollup::category_rollup(
351 &mut self.connection,
352 corpus.id,
353 service.id,
354 severity,
355 limit,
356 offset,
357 )
358 .unwrap_or_default()
359 }
360 pub fn what_rollup(
363 &mut self,
364 corpus: &Corpus,
365 service: &Service,
366 severity: &str,
367 category: &str,
368 limit: i64,
369 offset: i64,
370 ) -> Vec<ReportSummaryRow> {
371 rollup::what_rollup(
372 &mut self.connection,
373 corpus.id,
374 service.id,
375 severity,
376 category,
377 limit,
378 offset,
379 )
380 .unwrap_or_default()
381 }
382 pub fn severity_totals(
386 &mut self,
387 corpus: &Corpus,
388 service: &Service,
389 severity: &str,
390 ) -> (i64, i64) {
391 rollup::severity_total(&mut self.connection, corpus.id, service.id, severity)
392 .unwrap_or_default()
393 .map_or((0, 0), |row| (row.task_count, row.message_count))
394 }
395 pub fn category_totals(
399 &mut self,
400 corpus: &Corpus,
401 service: &Service,
402 severity: &str,
403 category: &str,
404 ) -> (i64, i64) {
405 rollup::category_total(
406 &mut self.connection,
407 corpus.id,
408 service.id,
409 severity,
410 category,
411 )
412 .unwrap_or_default()
413 .map_or((0, 0), |row| (row.task_count, row.message_count))
414 }
415}