Skip to main content

cortex/
backend.rs

1// Copyright 2015-2025 Deyan Ginev. See the LICENSE
2// file at the top-level directory of this distribution.
3//
4// Licensed under the MIT license <LICENSE-MIT or http://opensource.org/licenses/MIT>.
5// This file may not be copied, modified, or distributed
6// except according to those terms.
7
8//! All aggregate operations over the CorTeX PostgresQL store are accessed through the connection of
9//! a `Backend` object.
10
11mod 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};
24// `pub` (not `pub(crate)`): the `cortex diff --tasks` subcommand calls it directly, giving the
25// per-task changed-tasks drill a third (CLI) surface alongside the agent `/api/runs/<c>/<s>/tasks`
26// and the web screen.
27pub 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;
31// `pub` (not `pub(crate)`): the `cortex` CLI's `diff` subcommand calls it directly, so the run-diff
32// summary has a third (CLI) surface alongside the agent `/api/runs/<c>/<s>/diff` and the web
33// screen.
34pub 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
56/// The production database postgresql address, from the runtime [`crate::config`] configuration
57pub fn default_db_address() -> &'static str { &config().database.url }
58/// The test database postgresql address, from the runtime [`crate::config`] configuration
59pub fn test_db_address() -> &'static str { &config().database.test_url }
60
61/// Provides an interface to the Postgres task store
62pub struct Backend {
63  /// The Diesel PgConnection object
64  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
76/// Constructs a new Task store representation from a Postgres DB address.
77///
78/// A transient connect failure — a momentary `too many clients` / a DB restarting — is retried a
79/// few times with linear backoff before giving up, so a blip doesn't become an instant fatal panic
80/// (the robustness mandate: degrade through transients, fail fast only when the DB is truly down).
81pub 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
99/// A pool of PostgreSQL connections (Diesel + r2d2).
100pub type DbPool = Pool<ConnectionManager<PgConnection>>;
101/// A connection checked out from a [`DbPool`]; dereferences to a `PgConnection`.
102pub type PooledConn = PooledConnection<ConnectionManager<PgConnection>>;
103
104/// Builds a lazily-initialized connection pool: connections are established on first checkout, so
105/// this never blocks or fails at startup even if the database is momentarily unavailable.
106pub 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
111/// The configured database URL, managed as Rocket state so background jobs open their own
112/// connection against the same database as the request pool (notably the test database in
113/// integration tests).
114pub struct DatabaseUrl(pub String);
115/// Constructs the default Backend struct for testing
116pub fn testdb() -> Backend {
117  Backend {
118    connection: connection_at(test_db_address()),
119  }
120}
121/// Constructs a Backend at a given address
122pub fn from_address(address: &str) -> Backend {
123  Backend {
124    connection: connection_at(address),
125  }
126}
127
128/// Options container for relevant fields in requesting a `(corpus, service)` rerun
129pub struct RerunOptions<'a> {
130  /// corpus to rerun
131  pub corpus: &'a Corpus,
132  /// service to rerun
133  pub service: &'a Service,
134  /// optionally, severity level filter
135  pub severity_opt: Option<String>,
136  /// optionally, category level filter
137  pub category_opt: Option<String>,
138  /// optionally, what level filter
139  pub what_opt: Option<String>,
140  /// optionally, owner of the rerun (default is "admin")
141  pub owner_opt: Option<String>,
142  /// optionally, description of the rerun (default is "rerun")
143  pub description_opt: Option<String>,
144}
145
146/// Instance methods
147impl Backend {
148  /// Insert a vector of new `NewTask` tasks into the Task store
149  /// For example, on import, or when a new service is activated on a corpus
150  pub fn mark_imported(&mut self, imported_tasks: &[NewTask]) -> Result<usize, Error> {
151    mark::mark_imported(&mut self.connection, imported_tasks)
152  }
153  /// Insert a vector of `TaskReport` reports into the Task store, also marking their tasks as
154  /// completed with the correct status code.
155  pub fn mark_done(&mut self, reports: &[TaskReport]) -> Result<(), Error> {
156    mark::mark_done(&mut self.connection, reports)
157  }
158  /// Given a complex selector, of a `Corpus`, `Service`, and the optional `severity`, `category`
159  /// and `what` mark all matching tasks to be rerun
160  pub fn mark_rerun(&mut self, options: RerunOptions) -> Result<(), Error> {
161    mark::mark_rerun(&mut self.connection, options)
162  }
163
164  /// **Pause** a `(corpus, service)` run: block every in-progress task (`status >= 0`) so the
165  /// dispatcher stops leasing them. Returns the number paused. The CLI/agent twin of the report
166  /// screen's "Pause run". Reversible with [`Backend::resume_run`].
167  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  /// **Resume** a paused `(corpus, service)` run: return every Blocked task (`status < -5`) to TODO
172  /// so the dispatcher picks them up. Returns the number resumed. The inverse of
173  /// [`Backend::pause_run`].
174  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  /// While not changing any status information for Tasks, add a new historical run bookmark
179  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  /// Save the current historical tasks for reference
190  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  /// Generic delete method, uses primary "id" field
199  pub fn delete<Model: CortexDeletable>(&mut self, object: &Model) -> Result<usize, Error> {
200    object.delete_by(&mut self.connection, "id")
201  }
202  /// Delete all entries matching the "field" value of a given object
203  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  /// Generic addition method, attempting to insert in the DB a Task store datum
211  /// applicable for any struct implementing the `CortexORM` trait
212  /// (for example `Corpus`, `Service`, `Task`)
213  pub fn add<Model: CortexInsertable>(&mut self, object: &Model) -> Result<usize, Error> {
214    object.create(&mut self.connection)
215  }
216
217  /// Fetches no more than `limit` queued tasks for a given `Service`
218  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  /// Globally resets any "in progress" tasks back to "queued".
222  /// Particularly useful for dispatcher restarts, when all "in progress" tasks need to be
223  /// invalidated
224  pub fn clear_limbo_tasks(&mut self) -> Result<usize, Error> {
225    tasks_aggregate::clear_limbo_tasks(&mut self.connection)
226  }
227
228  /// Like [`Self::clear_limbo_tasks`], but preserves the given in-flight task ids (the live
229  /// `progress_queue` on a ventilator restart) so tasks a worker is actively processing are not
230  /// reset to `TODO` and double-dispatched. See `tasks_aggregate::clear_limbo_tasks_except`.
231  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  /// Activates an existing service on the given corpus.
236  /// if the service has previously been registered, this call will `RESET` the service into a mint
237  /// state also removing any related log messages. The new run is attributed to `owner` with
238  /// `description` (the UI/API thread the actor; the CLI passes a default). Takes the resolved
239  /// `Corpus` (not a path): a sandbox shares its parent's path, so a path lookup would hit the
240  /// parent.
241  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  /// Extends an existing service on the given corpus.
252  /// if the service has previously been registered, this call will ignore existing entries and
253  /// simply add newly encountered ones. Takes the resolved `Corpus` (not a path) for the same
254  /// sandbox-shares-parent-path reason as `register_service`.
255  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  /// Permanently destroys a service by name — its definition plus all of its tasks + `log_*` rows
260  /// across every corpus, in one transaction (orphan-free + crash-consistent; closes R-6). Refuses
261  /// the magic `init`/`import` services. See [`services_aggregate::destroy_service_by_name`].
262  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  /// Returns a vector of currently available corpora in the Task store
267  pub fn corpora(&mut self) -> Vec<Corpus> { corpora_aggregate::list_corpora(&mut self.connection) }
268
269  /// Returns a vector of tasks for a given Corpus, Service and status
270  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  /// Returns a vector of task entry paths for a given Corpus, Service and status
279  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  /// Given a complex selector, of a `Corpus`, `Service`, and the optional `severity`, `category`
289  /// and `what`, Provide a progress report at the chosen granularity
290  pub fn task_report(&mut self, options: TaskReportOptions) -> Vec<HashMap<String, String>> {
291    reports::task_report(&mut self.connection, options)
292  }
293  /// Provides a progress report, grouped by severity, for a given `Corpus` and `Service` pair
294  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  /// Recomputes the `report_summary` rollup (Arm 14 #6); call on the run-completion path so the
299  /// cheap category/what report reads stay fresh.
300  pub fn refresh_report_summary(&mut self) -> Result<(), Error> {
301    rollup::refresh_report_summary(&mut self.connection)
302  }
303  /// Drop the cached report grains for one `(corpus, service)` scope so the next report view
304  /// repopulates from current data. Called per touched scope on the run-completion path — scoped, a
305  /// cheap keyed `DELETE`, never the global all-corpora scan that used to stall conversions.
306  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  /// Run-completion-on-drain: close the `(corpus, service)` pair's open historical run iff its work
310  /// is exhausted (every task terminal). Delegates to
311  /// [`HistoricalRun::complete_if_drained`](crate::models::HistoricalRun::complete_if_drained);
312  /// returns `true` iff a run was closed. Called per touched scope from the dispatcher's finalize
313  /// loop when the queue idles, so a finished run is closed at once rather than at the next rerun.
314  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      // The run just closed on drain — its per-task outcomes are settled. Snapshot them into
326      // `historical_tasks` as the BASELINE for the NEXT run's live run-diff: captured once, here,
327      // at the natural completion point (in the finalize thread, off the request path). Best-effort
328      // — a snapshot failure must never undo the (critical) run close, so log and carry on.
329      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  /// Category-grain report for `(corpus, service, severity)`, read from the `report_summary`
341  /// rollup, windowed to `[offset, offset + limit)` (ordered by descending task count).
342  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  /// `what`-grain drill-down for `(corpus, service, severity, category)`, read from the rollup,
361  /// windowed to `[offset, offset + limit)` (ordered by descending task count).
362  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  /// The per-severity grand totals `(distinct tasks, total messages)` for `(corpus, service,
383  /// severity)`, read from the rollup — the denominators the category-grain report shows. `(0, 0)`
384  /// when the severity has no logged messages.
385  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  /// The category-grain totals `(distinct tasks, total messages)` for `(corpus, service, severity,
396  /// category)` — the denominators the `what` drill-down shows. `(0, 0)` when the category is
397  /// empty.
398  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}