1use crate::dlq::DlqMetadata;
2use crate::error::AwaError;
3use crate::job::{JobRow, JobState};
4use crate::queue_storage::QueueStorage;
5use chrono::{DateTime, Duration, Utc};
6use serde::{Deserialize, Serialize};
7use sqlx::types::Json;
8use sqlx::Acquire;
9use sqlx::PgExecutor;
10use sqlx::PgPool;
11use std::cmp::max;
12use std::collections::HashMap;
13use std::fmt;
14use std::str::FromStr;
15use uuid::Uuid;
16
17#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
18pub struct JobTimelineEvent {
19 pub timestamp: DateTime<Utc>,
20 pub label: String,
21 pub state: Option<JobState>,
22 pub detail: Option<String>,
23 pub is_error: bool,
24}
25
26#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
27pub struct JobDumpSummary {
28 pub original_priority: i16,
29 pub can_retry: bool,
30 pub can_cancel: bool,
31 pub error_count: usize,
32 pub latest_error: Option<serde_json::Value>,
33}
34
35#[derive(Debug, Clone, Serialize, Deserialize)]
36pub struct JobDump {
37 pub job: JobRow,
38 pub summary: JobDumpSummary,
39 pub timeline: Vec<JobTimelineEvent>,
40 #[serde(default, skip_serializing_if = "Option::is_none")]
41 pub dlq: Option<DlqMetadata>,
42}
43
44#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
45#[serde(rename_all = "snake_case")]
46pub enum RunDumpSource {
47 CurrentJobRow,
48 ErrorHistory,
49}
50
51#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
52pub struct CallbackDump {
53 pub callback_id: Option<Uuid>,
54 pub callback_timeout_at: Option<DateTime<Utc>>,
55 pub callback_filter: Option<String>,
56 pub callback_on_complete: Option<String>,
57 pub callback_on_fail: Option<String>,
58 pub callback_transform: Option<String>,
59}
60
61#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
62pub struct RunDump {
63 pub job_id: i64,
64 pub kind: String,
65 pub queue: String,
66 pub selected_attempt: i16,
67 pub current_attempt: i16,
68 pub current_run_lease: i64,
69 pub selected_run_lease: Option<i64>,
70 pub source: RunDumpSource,
71 pub state: JobState,
72 pub started_at: Option<DateTime<Utc>>,
73 pub finished_at: Option<DateTime<Utc>>,
74 pub error: Option<String>,
75 pub terminal: Option<bool>,
76 pub progress: Option<serde_json::Value>,
77 pub metadata: Option<serde_json::Value>,
78 pub callback: Option<CallbackDump>,
79 pub raw_error_entry: Option<serde_json::Value>,
80 pub notes: Vec<String>,
81}
82
83#[derive(Debug, Clone)]
84struct ErrorEntry {
85 attempt: Option<i16>,
86 error: Option<String>,
87 at: Option<DateTime<Utc>>,
88 terminal: bool,
89 raw: serde_json::Value,
90}
91
92fn parse_error_entry(value: &serde_json::Value) -> Option<ErrorEntry> {
93 let obj = value.as_object()?;
94 let attempt = obj
95 .get("attempt")
96 .and_then(|v| v.as_i64())
97 .and_then(|v| i16::try_from(v).ok());
98 let error = obj
99 .get("error")
100 .and_then(|v| v.as_str())
101 .map(ToOwned::to_owned);
102 let at = obj
103 .get("at")
104 .and_then(|v| v.as_str())
105 .and_then(|v| chrono::DateTime::parse_from_rfc3339(v).ok())
106 .map(|v| v.with_timezone(&Utc));
107 let terminal = obj
108 .get("terminal")
109 .and_then(|v| v.as_bool())
110 .unwrap_or(false);
111 Some(ErrorEntry {
112 attempt,
113 error,
114 at,
115 terminal,
116 raw: value.clone(),
117 })
118}
119
120fn original_priority(job: &JobRow) -> i16 {
121 job.metadata
122 .get("_awa_original_priority")
123 .and_then(|v| v.as_i64())
124 .and_then(|v| i16::try_from(v).ok())
125 .unwrap_or(job.priority)
126}
127
128fn build_job_timeline(job: &JobRow) -> Vec<JobTimelineEvent> {
129 let mut events = Vec::new();
130 events.push(JobTimelineEvent {
131 timestamp: job.created_at,
132 label: "Created".to_string(),
133 state: None,
134 detail: None,
135 is_error: false,
136 });
137
138 if let Some(errors) = &job.errors {
139 for entry in errors.iter().filter_map(parse_error_entry) {
140 if let Some(timestamp) = entry.at {
141 let label = match entry.attempt {
142 Some(attempt) => format!("Attempt {attempt} failed"),
143 None => "Error".to_string(),
144 };
145 events.push(JobTimelineEvent {
146 timestamp,
147 label,
148 state: Some(JobState::Failed),
149 detail: entry.error,
150 is_error: true,
151 });
152 }
153 }
154 }
155
156 if let Some(attempted_at) = job.attempted_at {
157 let label = if job.state == JobState::WaitingExternal {
158 format!("Attempt {} — waiting for callback", job.attempt)
159 } else if job.state.is_terminal() {
160 format!("Attempt {} started", job.attempt)
161 } else {
162 format!("Attempt {} running", job.attempt)
163 };
164 let state = if job.state.is_terminal() {
165 Some(JobState::Running)
166 } else {
167 Some(job.state)
168 };
169 events.push(JobTimelineEvent {
170 timestamp: attempted_at,
171 label,
172 state,
173 detail: None,
174 is_error: false,
175 });
176 }
177
178 if let Some(finalized_at) = job.finalized_at {
179 let label = match job.state {
180 JobState::Completed => format!("Completed (attempt {})", job.attempt),
181 JobState::Failed => format!(
182 "Failed after {} attempt{}",
183 job.attempt,
184 if job.attempt == 1 { "" } else { "s" }
185 ),
186 JobState::Cancelled => "Cancelled".to_string(),
187 _ => "Finalized".to_string(),
188 };
189 events.push(JobTimelineEvent {
190 timestamp: finalized_at,
191 label,
192 state: Some(job.state),
193 detail: None,
194 is_error: false,
195 });
196 }
197
198 events.sort_by_key(|event| event.timestamp);
199 events
200}
201
202fn build_job_dump(job: JobRow, dlq: Option<DlqMetadata>) -> JobDump {
203 let latest_error = job
204 .errors
205 .as_ref()
206 .and_then(|errors| errors.last())
207 .cloned();
208 let summary = JobDumpSummary {
209 original_priority: original_priority(&job),
210 can_retry: dlq.is_none()
211 && matches!(
212 job.state,
213 JobState::Failed | JobState::Cancelled | JobState::WaitingExternal
214 ),
215 can_cancel: dlq.is_none() && !job.state.is_terminal(),
216 error_count: job.errors.as_ref().map(|errors| errors.len()).unwrap_or(0),
217 latest_error,
218 };
219 let timeline = build_job_timeline(&job);
220 JobDump {
221 job,
222 summary,
223 timeline,
224 dlq,
225 }
226}
227
228fn callback_dump(job: &JobRow) -> Option<CallbackDump> {
229 let callback = CallbackDump {
230 callback_id: job.callback_id,
231 callback_timeout_at: job.callback_timeout_at,
232 callback_filter: job.callback_filter.clone(),
233 callback_on_complete: job.callback_on_complete.clone(),
234 callback_on_fail: job.callback_on_fail.clone(),
235 callback_transform: job.callback_transform.clone(),
236 };
237 if callback.callback_id.is_none()
238 && callback.callback_timeout_at.is_none()
239 && callback.callback_filter.is_none()
240 && callback.callback_on_complete.is_none()
241 && callback.callback_on_fail.is_none()
242 && callback.callback_transform.is_none()
243 {
244 None
245 } else {
246 Some(callback)
247 }
248}
249
250async fn active_queue_storage(pool: &PgPool) -> Result<Option<QueueStorage>, AwaError> {
251 QueueStorage::active_schema(pool)
252 .await?
253 .map(QueueStorage::from_existing_schema)
254 .transpose()
255}
256
257async fn active_queue_storage_in_tx(
264 tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
265) -> Result<Option<QueueStorage>, AwaError> {
266 QueueStorage::active_schema_in_tx(tx)
267 .await?
268 .map(QueueStorage::from_existing_schema)
269 .transpose()
270}
271
272fn queue_storage_current_jobs_cte(schema: &str) -> String {
273 format!(
274 r#"
275 WITH current_available AS (
276 SELECT
277 ready.job_id,
278 ready.kind,
279 ready.queue,
280 'available'::awa.job_state AS state,
281 ready.created_at,
282 ready.run_at,
283 NULL::timestamptz AS finalized_at
284 FROM {schema}.ready_entries AS ready
285 JOIN {schema}.queue_claim_heads AS claims
286 ON claims.queue = ready.queue
287 AND claims.priority = ready.priority
288 AND claims.enqueue_shard = ready.enqueue_shard
289 WHERE ready.lane_seq >= {schema}.sequence_next_value(claims.seq_name)
290 AND NOT EXISTS (
291 SELECT 1 FROM {schema}.ready_tombstones AS tomb
292 WHERE tomb.ready_slot = ready.ready_slot
293 AND tomb.ready_generation = ready.ready_generation
294 AND tomb.queue = ready.queue
295 AND tomb.priority = ready.priority
296 AND tomb.enqueue_shard = ready.enqueue_shard
297 AND tomb.lane_seq = ready.lane_seq
298 )
299 ),
300 current_jobs AS (
301 SELECT job_id, kind, queue, state, created_at, run_at, finalized_at
302 FROM current_available
303 UNION ALL
304 SELECT job_id, kind, queue, state, created_at, run_at, finalized_at
305 FROM {schema}.deferred_jobs
306 UNION ALL
307 SELECT
308 leases.job_id,
309 ready.kind,
310 leases.queue,
311 leases.state,
312 ready.created_at,
313 ready.run_at,
314 NULL::timestamptz AS finalized_at
315 FROM {schema}.leases AS leases
316 JOIN {schema}.ready_entries AS ready
317 ON ready.ready_slot = leases.ready_slot
318 AND ready.ready_generation = leases.ready_generation
319 AND ready.queue = leases.queue
320 AND ready.priority = leases.priority
321 AND ready.enqueue_shard = leases.enqueue_shard
322 AND ready.lane_seq = leases.lane_seq
323 UNION ALL
324 SELECT job_id, kind, queue, state, created_at, run_at, finalized_at
325 FROM {schema}.terminal_jobs
326 UNION ALL
327 SELECT
328 job_id,
329 kind,
330 queue,
331 'failed'::awa.job_state AS state,
332 created_at,
333 run_at,
334 finalized_at
335 FROM {schema}.dlq_entries
336 )
337 "#
338 )
339}
340
341async fn notify_cancellation_tx<'a>(
342 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
343 job_id: i64,
344 run_lease: i64,
345) -> Result<(), AwaError> {
346 let payload = serde_json::json!({ "job_id": job_id, "run_lease": run_lease }).to_string();
347 sqlx::query("SELECT pg_notify('awa:cancel', $1)")
348 .bind(payload)
349 .execute(tx.as_mut())
350 .await?;
351 Ok(())
352}
353
354async fn list_queue_storage_jobs(
355 store: &QueueStorage,
356 pool: &PgPool,
357 filter: &ListJobsFilter,
358) -> Result<Vec<JobRow>, AwaError> {
359 let limit = filter.limit.unwrap_or(100).clamp(1, 1000);
360 let candidate_limit = if filter.tag.is_some() {
361 limit.saturating_mul(10).min(5000)
362 } else {
363 limit.saturating_mul(3).min(2000)
364 };
365
366 let sql = format!(
367 "{} \
368 SELECT job_id \
369 FROM current_jobs \
370 WHERE ($1::awa.job_state IS NULL OR state = $1) \
371 AND ($2::text IS NULL OR kind = $2) \
372 AND ($3::text IS NULL OR queue = $3) \
373 AND ($4::bigint IS NULL OR job_id < $4) \
374 ORDER BY job_id DESC \
375 LIMIT $5",
376 queue_storage_current_jobs_cte(store.schema())
377 );
378
379 let mut jobs = Vec::new();
380 let mut cursor = filter.before_id;
381
382 loop {
383 let ids: Vec<i64> = sqlx::query_scalar(&sql)
384 .bind(filter.state)
385 .bind(&filter.kind)
386 .bind(&filter.queue)
387 .bind(cursor)
388 .bind(candidate_limit)
389 .fetch_all(pool)
390 .await?;
391 if ids.is_empty() {
392 break;
393 }
394
395 for job_id in &ids {
396 let Some(job) = store.load_job(pool, *job_id).await? else {
397 continue;
398 };
399
400 if let Some(state) = filter.state {
401 if job.state != state {
402 continue;
403 }
404 }
405 if let Some(kind) = &filter.kind {
406 if &job.kind != kind {
407 continue;
408 }
409 }
410 if let Some(queue) = &filter.queue {
411 if &job.queue != queue {
412 continue;
413 }
414 }
415 if let Some(tag) = &filter.tag {
416 if !job.tags.iter().any(|job_tag| job_tag == tag) {
417 continue;
418 }
419 }
420
421 jobs.push(job);
422 if jobs.len() as i64 >= limit {
423 break;
424 }
425 }
426
427 if jobs.len() as i64 >= limit || ids.len() < candidate_limit as usize {
428 break;
429 }
430 cursor = ids.last().copied();
431 }
432
433 jobs.sort_by_key(|job| std::cmp::Reverse(job.id));
434 Ok(jobs)
435}
436
437pub async fn retry(pool: &PgPool, job_id: i64) -> Result<Option<JobRow>, AwaError> {
439 if let Some(store) = active_queue_storage(pool).await? {
440 return store
441 .retry_job(pool, job_id)
442 .await?
443 .ok_or(AwaError::JobNotFound { id: job_id })
444 .map(Some);
445 }
446
447 sqlx::query_as::<_, JobRow>(
448 r#"
449 UPDATE awa.jobs
450 SET state = 'available', attempt = 0, run_at = now(),
451 finalized_at = NULL, heartbeat_at = NULL, deadline_at = NULL,
452 callback_id = NULL, callback_timeout_at = NULL,
453 callback_filter = NULL, callback_on_complete = NULL,
454 callback_on_fail = NULL, callback_transform = NULL
455 WHERE id = $1 AND state IN ('failed', 'cancelled', 'waiting_external')
456 RETURNING *
457 "#,
458 )
459 .bind(job_id)
460 .fetch_optional(pool)
461 .await?
462 .ok_or(AwaError::JobNotFound { id: job_id })
463 .map(Some)
464}
465
466pub async fn cancel(pool: &PgPool, job_id: i64) -> Result<Option<JobRow>, AwaError> {
468 if let Some(store) = active_queue_storage(pool).await? {
469 return store
470 .cancel_job(pool, job_id)
471 .await?
472 .ok_or(AwaError::JobNotFound { id: job_id })
473 .map(Some);
474 }
475
476 let mut tx = pool.begin().await?;
481
482 let prior_state: Option<JobState> =
483 sqlx::query_scalar::<_, JobState>("SELECT state FROM awa.jobs WHERE id = $1 FOR UPDATE")
484 .bind(job_id)
485 .fetch_optional(tx.as_mut())
486 .await?;
487
488 let Some(prior_state) = prior_state else {
489 tx.rollback().await.ok();
490 return Err(AwaError::JobNotFound { id: job_id });
491 };
492
493 let job: Option<JobRow> = sqlx::query_as::<_, JobRow>(
494 r#"
495 UPDATE awa.jobs
496 SET state = 'cancelled', finalized_at = now(),
497 callback_id = NULL, callback_timeout_at = NULL,
498 callback_filter = NULL, callback_on_complete = NULL,
499 callback_on_fail = NULL, callback_transform = NULL
500 WHERE id = $1 AND state NOT IN ('completed', 'failed', 'cancelled')
501 RETURNING *
502 "#,
503 )
504 .bind(job_id)
505 .fetch_optional(tx.as_mut())
506 .await?;
507
508 let Some(job) = job else {
509 tx.rollback().await.ok();
510 return Err(AwaError::JobNotFound { id: job_id });
511 };
512
513 if matches!(prior_state, JobState::Running | JobState::WaitingExternal) {
514 notify_cancellation_tx(&mut tx, job.id, job.run_lease).await?;
515 }
516
517 tx.commit().await?;
518 Ok(Some(job))
519}
520
521pub async fn cancel_tx(
541 conn: &mut sqlx::PgConnection,
542 job_id: i64,
543) -> Result<Option<JobRow>, AwaError> {
544 let mut tx = conn.begin().await?;
545 match cancel_in_tx(&mut tx, job_id).await {
546 Ok(row) => {
547 tx.commit().await?;
548 Ok(row)
549 }
550 Err(err) => {
551 tx.rollback().await.ok();
552 Err(err)
553 }
554 }
555}
556
557async fn cancel_in_tx<'a>(
560 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
561 job_id: i64,
562) -> Result<Option<JobRow>, AwaError> {
563 if let Some(store) = active_queue_storage_in_tx(tx).await? {
564 return store
565 .cancel_job_in_tx(tx, job_id)
566 .await?
567 .ok_or(AwaError::JobNotFound { id: job_id })
568 .map(Some);
569 }
570
571 let prior_state: Option<JobState> =
572 sqlx::query_scalar::<_, JobState>("SELECT state FROM awa.jobs WHERE id = $1 FOR UPDATE")
573 .bind(job_id)
574 .fetch_optional(tx.as_mut())
575 .await?;
576
577 let Some(prior_state) = prior_state else {
578 return Err(AwaError::JobNotFound { id: job_id });
579 };
580
581 let job: Option<JobRow> = sqlx::query_as::<_, JobRow>(
582 r#"
583 UPDATE awa.jobs
584 SET state = 'cancelled', finalized_at = now(),
585 callback_id = NULL, callback_timeout_at = NULL,
586 callback_filter = NULL, callback_on_complete = NULL,
587 callback_on_fail = NULL, callback_transform = NULL
588 WHERE id = $1 AND state NOT IN ('completed', 'failed', 'cancelled')
589 RETURNING *
590 "#,
591 )
592 .bind(job_id)
593 .fetch_optional(tx.as_mut())
594 .await?;
595
596 let Some(job) = job else {
597 return Err(AwaError::JobNotFound { id: job_id });
598 };
599
600 if matches!(prior_state, JobState::Running | JobState::WaitingExternal) {
601 notify_cancellation_tx(tx, job.id, job.run_lease).await?;
602 }
603
604 Ok(Some(job))
605}
606
607fn unique_key_candidate_sql(schema: &str) -> String {
611 format!(
612 r#"
613 WITH current_available AS (
614 SELECT ready.job_id, ready.unique_key
615 FROM {schema}.ready_entries AS ready
616 JOIN {schema}.queue_claim_heads AS claims
617 ON claims.queue = ready.queue
618 AND claims.priority = ready.priority
619 AND claims.enqueue_shard = ready.enqueue_shard
620 WHERE ready.lane_seq >= {schema}.sequence_next_value(claims.seq_name)
621 AND NOT EXISTS (
622 SELECT 1 FROM {schema}.ready_tombstones AS tomb
623 WHERE tomb.ready_slot = ready.ready_slot
624 AND tomb.ready_generation = ready.ready_generation
625 AND tomb.queue = ready.queue
626 AND tomb.priority = ready.priority
627 AND tomb.enqueue_shard = ready.enqueue_shard
628 AND tomb.lane_seq = ready.lane_seq
629 )
630 ),
631 candidates AS (
632 SELECT job_id
633 FROM current_available
634 WHERE unique_key = $1
635 UNION ALL
636 SELECT job_id
637 FROM {schema}.deferred_jobs
638 WHERE unique_key = $1
639 UNION ALL
640 -- A running job's unique_key is stored on its ready_entries row;
641 -- reach it through the lease's lane identity.
642 SELECT leases.job_id
643 FROM {schema}.leases AS leases
644 JOIN {schema}.ready_entries AS ready
645 ON ready.ready_slot = leases.ready_slot
646 AND ready.ready_generation = leases.ready_generation
647 AND ready.queue = leases.queue
648 AND ready.priority = leases.priority
649 AND ready.enqueue_shard = leases.enqueue_shard
650 AND ready.lane_seq = leases.lane_seq
651 WHERE ready.unique_key = $1
652 )
653 SELECT job_id
654 FROM candidates
655 ORDER BY job_id ASC
656 LIMIT 1
657 "#,
658 schema = schema
659 )
660}
661
662const UNIQUE_KEY_CANDIDATE_SQL_CANONICAL: &str = r#"
664 WITH candidates AS (
665 SELECT id FROM awa.jobs_hot
666 WHERE unique_key = $1 AND state NOT IN ('completed', 'failed', 'cancelled')
667 UNION ALL
668 SELECT id FROM awa.scheduled_jobs
669 WHERE unique_key = $1 AND state NOT IN ('completed', 'failed', 'cancelled')
670 ORDER BY id ASC
671 LIMIT 1
672 )
673 SELECT id
674 FROM candidates
675 "#;
676
677pub async fn cancel_by_unique_key(
711 pool: &PgPool,
712 kind: &str,
713 queue: Option<&str>,
714 args: Option<&serde_json::Value>,
715 period_bucket: Option<i64>,
716) -> Result<Option<JobRow>, AwaError> {
717 let unique_key = crate::unique::compute_unique_key(kind, queue, args, period_bucket);
718
719 if let Some(store) = active_queue_storage(pool).await? {
720 let sql = unique_key_candidate_sql(store.schema());
721 let candidate: Option<i64> = sqlx::query_scalar(&sql)
722 .bind(&unique_key)
723 .fetch_optional(pool)
724 .await?;
725
726 return match candidate {
727 Some(job_id) => cancel(pool, job_id).await,
728 None => Ok(None),
729 };
730 }
731
732 let candidate: Option<i64> = sqlx::query_scalar(UNIQUE_KEY_CANDIDATE_SQL_CANONICAL)
736 .bind(&unique_key)
737 .fetch_optional(pool)
738 .await?;
739
740 match candidate {
741 Some(job_id) => match cancel(pool, job_id).await {
742 Ok(row) => Ok(row),
743 Err(AwaError::JobNotFound { .. }) => Ok(None),
744 Err(err) => Err(err),
745 },
746 None => Ok(None),
747 }
748}
749
750pub async fn cancel_by_unique_key_tx(
758 conn: &mut sqlx::PgConnection,
759 kind: &str,
760 queue: Option<&str>,
761 args: Option<&serde_json::Value>,
762 period_bucket: Option<i64>,
763) -> Result<Option<JobRow>, AwaError> {
764 let mut tx = conn.begin().await?;
765 match cancel_by_unique_key_in_tx(&mut tx, kind, queue, args, period_bucket).await {
766 Ok(row) => {
767 tx.commit().await?;
768 Ok(row)
769 }
770 Err(err) => {
771 tx.rollback().await.ok();
772 Err(err)
773 }
774 }
775}
776
777async fn cancel_by_unique_key_in_tx<'a>(
778 tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
779 kind: &str,
780 queue: Option<&str>,
781 args: Option<&serde_json::Value>,
782 period_bucket: Option<i64>,
783) -> Result<Option<JobRow>, AwaError> {
784 let unique_key = crate::unique::compute_unique_key(kind, queue, args, period_bucket);
785
786 if let Some(store) = active_queue_storage_in_tx(tx).await? {
787 let sql = unique_key_candidate_sql(store.schema());
788 let candidate: Option<i64> = sqlx::query_scalar(&sql)
789 .bind(&unique_key)
790 .fetch_optional(tx.as_mut())
791 .await?;
792
793 return match candidate {
794 Some(job_id) => cancel_in_tx(tx, job_id).await,
795 None => Ok(None),
796 };
797 }
798
799 let candidate: Option<i64> = sqlx::query_scalar(UNIQUE_KEY_CANDIDATE_SQL_CANONICAL)
800 .bind(&unique_key)
801 .fetch_optional(tx.as_mut())
802 .await?;
803
804 match candidate {
805 Some(job_id) => match cancel_in_tx(tx, job_id).await {
806 Ok(row) => Ok(row),
807 Err(AwaError::JobNotFound { .. }) => Ok(None),
808 Err(err) => Err(err),
809 },
810 None => Ok(None),
811 }
812}
813
814#[derive(Debug, Clone, Serialize, Deserialize)]
817pub struct RetryFailedOutcome {
818 pub retried: Vec<JobRow>,
820 pub matched: u64,
824 pub pruned_failed_count: Option<u64>,
831}
832
833pub async fn retry_failed_by_kind(
835 pool: &PgPool,
836 kind: &str,
837) -> Result<RetryFailedOutcome, AwaError> {
838 if let Some(store) = active_queue_storage(pool).await? {
839 let sql = format!(
840 r#"
841 SELECT DISTINCT job_id
842 FROM {schema}.done_entries
843 WHERE kind = $1
844 AND state = 'failed'
845 ORDER BY job_id ASC
846 "#,
847 schema = store.schema()
848 );
849 let ids: Vec<i64> = sqlx::query_scalar(&sql).bind(kind).fetch_all(pool).await?;
850 let (retried, matched) = store.retry_jobs_by_ids(pool, &ids).await?;
851 return Ok(RetryFailedOutcome {
852 retried,
853 matched,
854 pruned_failed_count: None,
855 });
856 }
857
858 let rows = sqlx::query_as::<_, JobRow>(
859 r#"
860 UPDATE awa.jobs
861 SET state = 'available', attempt = 0, run_at = now(),
862 finalized_at = NULL, heartbeat_at = NULL, deadline_at = NULL
863 WHERE kind = $1 AND state = 'failed'
864 RETURNING *
865 "#,
866 )
867 .bind(kind)
868 .fetch_all(pool)
869 .await?;
870
871 Ok(RetryFailedOutcome {
872 matched: rows.len() as u64,
873 retried: rows,
874 pruned_failed_count: None,
875 })
876}
877
878pub async fn retry_failed_by_queue(
880 pool: &PgPool,
881 queue: &str,
882) -> Result<RetryFailedOutcome, AwaError> {
883 if let Some(store) = active_queue_storage(pool).await? {
884 let sql = format!(
885 r#"
886 SELECT DISTINCT job_id
887 FROM {schema}.done_entries
888 WHERE queue = $1
889 AND state = 'failed'
890 ORDER BY job_id ASC
891 "#,
892 schema = store.schema()
893 );
894 let ids: Vec<i64> = sqlx::query_scalar(&sql).bind(queue).fetch_all(pool).await?;
895 let (retried, matched) = store.retry_jobs_by_ids(pool, &ids).await?;
896 let pruned_failed_count = store.pruned_failed_count_for_queue(pool, queue).await?;
897 return Ok(RetryFailedOutcome {
898 retried,
899 matched,
900 pruned_failed_count: Some(pruned_failed_count),
901 });
902 }
903
904 let rows = sqlx::query_as::<_, JobRow>(
905 r#"
906 UPDATE awa.jobs
907 SET state = 'available', attempt = 0, run_at = now(),
908 finalized_at = NULL, heartbeat_at = NULL, deadline_at = NULL
909 WHERE queue = $1 AND state = 'failed'
910 RETURNING *
911 "#,
912 )
913 .bind(queue)
914 .fetch_all(pool)
915 .await?;
916
917 Ok(RetryFailedOutcome {
918 matched: rows.len() as u64,
919 retried: rows,
920 pruned_failed_count: None,
921 })
922}
923
924pub async fn discard_failed(pool: &PgPool, kind: &str) -> Result<u64, AwaError> {
926 if let Some(store) = active_queue_storage(pool).await? {
927 return store.discard_failed_by_kind(pool, kind).await;
928 }
929
930 let result = sqlx::query("DELETE FROM awa.jobs WHERE kind = $1 AND state = 'failed'")
931 .bind(kind)
932 .execute(pool)
933 .await?;
934
935 Ok(result.rows_affected())
936}
937
938pub async fn pause_queue<'e, E>(
940 executor: E,
941 queue: &str,
942 paused_by: Option<&str>,
943) -> Result<(), AwaError>
944where
945 E: PgExecutor<'e>,
946{
947 sqlx::query(
948 r#"
949 INSERT INTO awa.queue_meta (queue, paused, paused_at, paused_by)
950 VALUES ($1, TRUE, now(), $2)
951 ON CONFLICT (queue) DO UPDATE SET paused = TRUE, paused_at = now(), paused_by = $2
952 "#,
953 )
954 .bind(queue)
955 .bind(paused_by)
956 .execute(executor)
957 .await?;
958
959 Ok(())
960}
961
962pub async fn resume_queue<'e, E>(executor: E, queue: &str) -> Result<(), AwaError>
964where
965 E: PgExecutor<'e>,
966{
967 sqlx::query("UPDATE awa.queue_meta SET paused = FALSE WHERE queue = $1")
968 .bind(queue)
969 .execute(executor)
970 .await?;
971
972 Ok(())
973}
974
975pub async fn drain_queue(pool: &PgPool, queue: &str) -> Result<u64, AwaError> {
977 if let Some(store) = active_queue_storage(pool).await? {
978 let sql = format!(
979 "{} \
980 SELECT job_id \
981 FROM current_jobs \
982 WHERE queue = $1 \
983 AND state IN ('available', 'scheduled', 'retryable', 'waiting_external') \
984 ORDER BY job_id ASC",
985 queue_storage_current_jobs_cte(store.schema())
986 );
987 let ids: Vec<i64> = sqlx::query_scalar(&sql).bind(queue).fetch_all(pool).await?;
988 return store
989 .cancel_jobs_by_ids(pool, &ids)
990 .await
991 .map(|rows| rows.len() as u64);
992 }
993
994 let result = sqlx::query(
995 r#"
996 UPDATE awa.jobs
997 SET state = 'cancelled', finalized_at = now(),
998 callback_id = NULL, callback_timeout_at = NULL,
999 callback_filter = NULL, callback_on_complete = NULL,
1000 callback_on_fail = NULL, callback_transform = NULL
1001 WHERE queue = $1 AND state IN ('available', 'scheduled', 'retryable', 'waiting_external')
1002 "#,
1003 )
1004 .bind(queue)
1005 .execute(pool)
1006 .await?;
1007
1008 Ok(result.rows_affected())
1009}
1010
1011#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
1013pub struct QueueDescriptor {
1014 pub display_name: Option<String>,
1015 pub description: Option<String>,
1016 pub owner: Option<String>,
1017 pub docs_url: Option<String>,
1018 pub tags: Vec<String>,
1019 pub extra: serde_json::Value,
1020}
1021
1022impl Default for QueueDescriptor {
1023 fn default() -> Self {
1024 Self {
1025 display_name: None,
1026 description: None,
1027 owner: None,
1028 docs_url: None,
1029 tags: Vec::new(),
1030 extra: serde_json::json!({}),
1031 }
1032 }
1033}
1034
1035#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
1036pub struct NamedQueueDescriptor {
1037 pub queue: String,
1038 #[serde(flatten)]
1039 pub descriptor: QueueDescriptor,
1040}
1041
1042#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
1043pub struct JobKindDescriptor {
1044 pub display_name: Option<String>,
1045 pub description: Option<String>,
1046 pub owner: Option<String>,
1047 pub docs_url: Option<String>,
1048 pub tags: Vec<String>,
1049 pub extra: serde_json::Value,
1050}
1051
1052impl Default for JobKindDescriptor {
1053 fn default() -> Self {
1054 Self {
1055 display_name: None,
1056 description: None,
1057 owner: None,
1058 docs_url: None,
1059 tags: Vec::new(),
1060 extra: serde_json::json!({}),
1061 }
1062 }
1063}
1064
1065#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
1066pub struct NamedJobKindDescriptor {
1067 pub kind: String,
1068 #[serde(flatten)]
1069 pub descriptor: JobKindDescriptor,
1070}
1071
1072const DESCRIPTOR_PARAMS_PER_ROW: u32 = 9;
1075
1076const DESCRIPTOR_BATCH_SIZE: usize = 5000;
1079
1080fn build_descriptor_upsert(table: &str, pk: &str, count: usize) -> String {
1084 let mut query = format!(
1085 "INSERT INTO {table} (\
1086 {pk}, display_name, description, owner, docs_url, tags, extra, \
1087 descriptor_hash, sync_interval_ms, created_at, updated_at, last_seen_at\
1088 ) VALUES "
1089 );
1090 let mut param_index = 1u32;
1091 for i in 0..count {
1092 if i > 0 {
1093 query.push_str(", ");
1094 }
1095 query.push_str(&format!(
1096 "(${}, ${}, ${}, ${}, ${}, ${}, ${}, ${}, ${}, now(), now(), now())",
1097 param_index,
1098 param_index + 1,
1099 param_index + 2,
1100 param_index + 3,
1101 param_index + 4,
1102 param_index + 5,
1103 param_index + 6,
1104 param_index + 7,
1105 param_index + 8,
1106 ));
1107 param_index += DESCRIPTOR_PARAMS_PER_ROW;
1108 }
1109 query.push_str(&format!(
1113 " ON CONFLICT ({pk}) DO UPDATE SET \
1114 display_name = EXCLUDED.display_name, \
1115 description = EXCLUDED.description, \
1116 owner = EXCLUDED.owner, \
1117 docs_url = EXCLUDED.docs_url, \
1118 tags = EXCLUDED.tags, \
1119 extra = EXCLUDED.extra, \
1120 descriptor_hash = EXCLUDED.descriptor_hash, \
1121 sync_interval_ms = EXCLUDED.sync_interval_ms, \
1122 updated_at = CASE \
1123 WHEN {table}.descriptor_hash IS DISTINCT FROM EXCLUDED.descriptor_hash \
1124 THEN now() ELSE {table}.updated_at \
1125 END, \
1126 last_seen_at = now()"
1127 ));
1128 query
1129}
1130
1131pub async fn sync_queue_descriptors(
1141 pool: &PgPool,
1142 descriptors: &[NamedQueueDescriptor],
1143 sync_interval: std::time::Duration,
1144) -> Result<(), AwaError> {
1145 if descriptors.is_empty() {
1146 return Ok(());
1147 }
1148 let sync_interval_ms = sync_interval.as_millis() as i64;
1149
1150 for chunk in descriptors.chunks(DESCRIPTOR_BATCH_SIZE) {
1151 let hashes: Vec<String> = chunk
1152 .iter()
1153 .map(|named| named.descriptor.descriptor_hash())
1154 .collect();
1155 let sql = build_descriptor_upsert("awa.queue_descriptors", "queue", chunk.len());
1156 let mut query = sqlx::query(&sql);
1157 for (named, hash) in chunk.iter().zip(hashes.iter()) {
1158 query = query
1159 .bind(&named.queue)
1160 .bind(named.descriptor.display_name.as_deref())
1161 .bind(named.descriptor.description.as_deref())
1162 .bind(named.descriptor.owner.as_deref())
1163 .bind(named.descriptor.docs_url.as_deref())
1164 .bind(&named.descriptor.tags)
1165 .bind(&named.descriptor.extra)
1166 .bind(hash.as_str())
1167 .bind(sync_interval_ms);
1168 }
1169 query.execute(pool).await?;
1170 }
1171
1172 Ok(())
1173}
1174
1175pub async fn sync_job_kind_descriptors(
1178 pool: &PgPool,
1179 descriptors: &[NamedJobKindDescriptor],
1180 sync_interval: std::time::Duration,
1181) -> Result<(), AwaError> {
1182 if descriptors.is_empty() {
1183 return Ok(());
1184 }
1185 let sync_interval_ms = sync_interval.as_millis() as i64;
1186
1187 for chunk in descriptors.chunks(DESCRIPTOR_BATCH_SIZE) {
1188 let hashes: Vec<String> = chunk
1189 .iter()
1190 .map(|named| named.descriptor.descriptor_hash())
1191 .collect();
1192 let sql = build_descriptor_upsert("awa.job_kind_descriptors", "kind", chunk.len());
1193 let mut query = sqlx::query(&sql);
1194 for (named, hash) in chunk.iter().zip(hashes.iter()) {
1195 query = query
1196 .bind(&named.kind)
1197 .bind(named.descriptor.display_name.as_deref())
1198 .bind(named.descriptor.description.as_deref())
1199 .bind(named.descriptor.owner.as_deref())
1200 .bind(named.descriptor.docs_url.as_deref())
1201 .bind(&named.descriptor.tags)
1202 .bind(&named.descriptor.extra)
1203 .bind(hash.as_str())
1204 .bind(sync_interval_ms);
1205 }
1206 query.execute(pool).await?;
1207 }
1208
1209 Ok(())
1210}
1211
1212impl QueueDescriptor {
1213 pub fn new() -> Self {
1214 Self::default()
1215 }
1216
1217 pub fn display_name(mut self, display_name: impl Into<String>) -> Self {
1218 self.display_name = Some(display_name.into());
1219 self
1220 }
1221
1222 pub fn description(mut self, description: impl Into<String>) -> Self {
1223 self.description = Some(description.into());
1224 self
1225 }
1226
1227 pub fn owner(mut self, owner: impl Into<String>) -> Self {
1228 self.owner = Some(owner.into());
1229 self
1230 }
1231
1232 pub fn docs_url(mut self, docs_url: impl Into<String>) -> Self {
1233 self.docs_url = Some(docs_url.into());
1234 self
1235 }
1236
1237 pub fn tags<I, S>(mut self, tags: I) -> Self
1238 where
1239 I: IntoIterator<Item = S>,
1240 S: Into<String>,
1241 {
1242 self.tags = tags.into_iter().map(Into::into).collect();
1243 self
1244 }
1245
1246 pub fn tag(mut self, tag: impl Into<String>) -> Self {
1247 self.tags.push(tag.into());
1248 self
1249 }
1250
1251 pub fn extra(mut self, extra: serde_json::Value) -> Self {
1252 self.extra = extra;
1253 self
1254 }
1255
1256 pub fn descriptor_hash(&self) -> String {
1257 let payload = serde_json::json!({
1258 "display_name": self.display_name,
1259 "description": self.description,
1260 "owner": self.owner,
1261 "docs_url": self.docs_url,
1262 "tags": self.tags,
1263 "extra": canonicalize_json(&self.extra),
1264 });
1265 let encoded = serde_json::to_vec(&payload)
1266 .expect("queue descriptor JSON serialization should not fail");
1267 hex::encode(blake3::hash(&encoded).as_bytes())
1268 }
1269}
1270
1271impl JobKindDescriptor {
1272 pub fn new() -> Self {
1273 Self::default()
1274 }
1275
1276 pub fn display_name(mut self, display_name: impl Into<String>) -> Self {
1277 self.display_name = Some(display_name.into());
1278 self
1279 }
1280
1281 pub fn description(mut self, description: impl Into<String>) -> Self {
1282 self.description = Some(description.into());
1283 self
1284 }
1285
1286 pub fn owner(mut self, owner: impl Into<String>) -> Self {
1287 self.owner = Some(owner.into());
1288 self
1289 }
1290
1291 pub fn docs_url(mut self, docs_url: impl Into<String>) -> Self {
1292 self.docs_url = Some(docs_url.into());
1293 self
1294 }
1295
1296 pub fn tags<I, S>(mut self, tags: I) -> Self
1297 where
1298 I: IntoIterator<Item = S>,
1299 S: Into<String>,
1300 {
1301 self.tags = tags.into_iter().map(Into::into).collect();
1302 self
1303 }
1304
1305 pub fn tag(mut self, tag: impl Into<String>) -> Self {
1306 self.tags.push(tag.into());
1307 self
1308 }
1309
1310 pub fn extra(mut self, extra: serde_json::Value) -> Self {
1311 self.extra = extra;
1312 self
1313 }
1314
1315 pub fn descriptor_hash(&self) -> String {
1316 let payload = serde_json::json!({
1317 "display_name": self.display_name,
1318 "description": self.description,
1319 "owner": self.owner,
1320 "docs_url": self.docs_url,
1321 "tags": self.tags,
1322 "extra": canonicalize_json(&self.extra),
1323 });
1324 let encoded = serde_json::to_vec(&payload)
1325 .expect("job kind descriptor JSON serialization should not fail");
1326 hex::encode(blake3::hash(&encoded).as_bytes())
1327 }
1328}
1329
1330fn canonicalize_json(value: &serde_json::Value) -> serde_json::Value {
1331 match value {
1332 serde_json::Value::Object(map) => {
1333 let mut keys: Vec<_> = map.keys().cloned().collect();
1334 keys.sort();
1335 let mut out = serde_json::Map::with_capacity(map.len());
1336 for key in keys {
1337 if let Some(child) = map.get(&key) {
1338 out.insert(key, canonicalize_json(child));
1339 }
1340 }
1341 serde_json::Value::Object(out)
1342 }
1343 serde_json::Value::Array(values) => {
1344 serde_json::Value::Array(values.iter().map(canonicalize_json).collect())
1345 }
1346 _ => value.clone(),
1347 }
1348}
1349
1350#[derive(Debug, Clone, Serialize, sqlx::FromRow)]
1351pub struct QueueOverview {
1352 pub queue: String,
1353 pub display_name: Option<String>,
1354 pub description: Option<String>,
1355 pub owner: Option<String>,
1356 pub docs_url: Option<String>,
1357 pub tags: Vec<String>,
1358 pub extra: serde_json::Value,
1359 pub descriptor_last_seen_at: Option<DateTime<Utc>>,
1360 pub descriptor_stale: bool,
1361 pub descriptor_mismatch: bool,
1362 pub total_queued: i64,
1364 pub scheduled: i64,
1365 pub available: i64,
1366 pub retryable: i64,
1367 pub running: i64,
1368 pub failed: i64,
1369 pub waiting_external: i64,
1370 pub completed_last_hour: i64,
1371 pub lag_seconds: Option<f64>,
1372 pub paused: bool,
1373}
1374
1375#[derive(Debug, Clone, Serialize, sqlx::FromRow)]
1376pub struct JobKindOverview {
1377 pub kind: String,
1378 pub display_name: Option<String>,
1379 pub description: Option<String>,
1380 pub owner: Option<String>,
1381 pub docs_url: Option<String>,
1382 pub tags: Vec<String>,
1383 pub extra: serde_json::Value,
1384 pub descriptor_last_seen_at: Option<DateTime<Utc>>,
1385 pub descriptor_stale: bool,
1386 pub descriptor_mismatch: bool,
1387 pub job_count: i64,
1388 pub queue_count: i64,
1389 pub completed_last_hour: i64,
1390}
1391
1392#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
1394pub struct RateLimitSnapshot {
1395 pub max_rate: f64,
1396 pub burst: u32,
1397}
1398
1399#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
1401#[serde(rename_all = "snake_case")]
1402pub enum QueueRuntimeMode {
1403 HardReserved,
1404 Weighted,
1405}
1406
1407#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
1409pub struct QueueRuntimeConfigSnapshot {
1410 pub mode: QueueRuntimeMode,
1411 pub max_workers: Option<u32>,
1412 pub min_workers: Option<u32>,
1413 pub weight: Option<u32>,
1414 pub global_max_workers: Option<u32>,
1415 pub poll_interval_ms: u64,
1416 pub deadline_duration_secs: u64,
1417 pub priority_aging_interval_secs: u64,
1418 #[serde(default)]
1419 pub claimers: Option<u16>,
1420 #[serde(default)]
1421 pub claim_batch_size: Option<usize>,
1422 #[serde(default)]
1423 pub dlq_enabled: Option<bool>,
1424 pub rate_limit: Option<RateLimitSnapshot>,
1425}
1426
1427#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
1429pub struct QueueRuntimeSnapshot {
1430 pub queue: String,
1431 pub in_flight: u32,
1432 pub overflow_held: Option<u32>,
1433 pub config: QueueRuntimeConfigSnapshot,
1434}
1435
1436#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
1437#[serde(rename_all = "snake_case")]
1438pub enum StorageCapability {
1439 Canonical,
1440 CanonicalDrainOnly,
1441 QueueStorage,
1442}
1443
1444impl StorageCapability {
1445 pub fn as_str(self) -> &'static str {
1446 match self {
1447 Self::Canonical => "canonical",
1448 Self::CanonicalDrainOnly => "canonical_drain_only",
1449 Self::QueueStorage => "queue_storage",
1450 }
1451 }
1452}
1453
1454impl fmt::Display for StorageCapability {
1455 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1456 f.write_str(self.as_str())
1457 }
1458}
1459
1460impl FromStr for StorageCapability {
1461 type Err = String;
1462
1463 fn from_str(value: &str) -> Result<Self, Self::Err> {
1464 match value {
1465 "canonical" => Ok(Self::Canonical),
1466 "canonical_drain_only" => Ok(Self::CanonicalDrainOnly),
1467 "queue_storage" => Ok(Self::QueueStorage),
1468 _ => Err(value.to_string()),
1469 }
1470 }
1471}
1472
1473#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq, Hash)]
1481#[serde(rename_all = "snake_case")]
1482pub enum TransitionRole {
1483 #[default]
1489 Auto,
1490 CanonicalDrain,
1494 QueueStorageTarget,
1498}
1499
1500impl TransitionRole {
1501 pub fn as_str(self) -> &'static str {
1502 match self {
1503 Self::Auto => "auto",
1504 Self::CanonicalDrain => "canonical_drain",
1505 Self::QueueStorageTarget => "queue_storage_target",
1506 }
1507 }
1508}
1509
1510impl fmt::Display for TransitionRole {
1511 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1512 f.write_str(self.as_str())
1513 }
1514}
1515
1516impl FromStr for TransitionRole {
1517 type Err = String;
1518
1519 fn from_str(value: &str) -> Result<Self, Self::Err> {
1520 match value {
1521 "auto" => Ok(Self::Auto),
1522 "canonical_drain" => Ok(Self::CanonicalDrain),
1523 "queue_storage_target" => Ok(Self::QueueStorageTarget),
1524 _ => Err(value.to_string()),
1525 }
1526 }
1527}
1528
1529#[derive(Debug, Clone, Serialize, Deserialize)]
1531pub struct RuntimeSnapshotInput {
1532 pub instance_id: Uuid,
1533 pub hostname: Option<String>,
1534 pub pid: i32,
1535 pub version: String,
1536 pub storage_capability: StorageCapability,
1537 #[serde(default)]
1538 pub transition_role: TransitionRole,
1539 pub started_at: DateTime<Utc>,
1540 pub snapshot_interval_ms: i64,
1541 pub healthy: bool,
1542 pub postgres_connected: bool,
1543 pub poll_loop_alive: bool,
1544 pub heartbeat_alive: bool,
1545 pub maintenance_alive: bool,
1546 pub shutting_down: bool,
1547 pub leader: bool,
1548 pub global_max_workers: Option<u32>,
1549 pub queues: Vec<QueueRuntimeSnapshot>,
1550 pub queue_descriptor_hashes: HashMap<String, String>,
1551 pub job_kind_descriptor_hashes: HashMap<String, String>,
1552}
1553
1554#[derive(Debug, Clone, Serialize, Deserialize)]
1556pub struct RuntimeInstance {
1557 pub instance_id: Uuid,
1558 pub hostname: Option<String>,
1559 pub pid: i32,
1560 pub version: String,
1561 pub storage_capability: StorageCapability,
1562 #[serde(default)]
1563 pub transition_role: TransitionRole,
1564 pub started_at: DateTime<Utc>,
1565 pub last_seen_at: DateTime<Utc>,
1566 pub snapshot_interval_ms: i64,
1567 pub stale: bool,
1568 pub healthy: bool,
1569 pub postgres_connected: bool,
1570 pub poll_loop_alive: bool,
1571 pub heartbeat_alive: bool,
1572 pub maintenance_alive: bool,
1573 pub shutting_down: bool,
1574 pub leader: bool,
1575 pub global_max_workers: Option<u32>,
1576 pub queues: Vec<QueueRuntimeSnapshot>,
1577}
1578
1579impl RuntimeInstance {
1580 fn stale_cutoff(interval_ms: i64) -> Duration {
1581 let interval_ms = max(interval_ms, 1_000);
1582 Duration::milliseconds(max(interval_ms.saturating_mul(3), 30_000))
1583 }
1584
1585 fn from_db_row(row: RuntimeInstanceRow, now: DateTime<Utc>) -> Result<Self, AwaError> {
1586 let stale = row.last_seen_at + Self::stale_cutoff(row.snapshot_interval_ms) < now;
1587 let storage_capability =
1588 StorageCapability::from_str(&row.storage_capability).map_err(|value| {
1589 AwaError::Validation(format!(
1590 "invalid storage capability in runtime_instances: {value}"
1591 ))
1592 })?;
1593 let transition_role = TransitionRole::from_str(&row.transition_role).map_err(|value| {
1594 AwaError::Validation(format!(
1595 "invalid transition_role in runtime_instances: {value}"
1596 ))
1597 })?;
1598 Ok(Self {
1599 instance_id: row.instance_id,
1600 hostname: row.hostname,
1601 pid: row.pid,
1602 version: row.version,
1603 storage_capability,
1604 transition_role,
1605 started_at: row.started_at,
1606 last_seen_at: row.last_seen_at,
1607 snapshot_interval_ms: row.snapshot_interval_ms,
1608 stale,
1609 healthy: row.healthy,
1610 postgres_connected: row.postgres_connected,
1611 poll_loop_alive: row.poll_loop_alive,
1612 heartbeat_alive: row.heartbeat_alive,
1613 maintenance_alive: row.maintenance_alive,
1614 shutting_down: row.shutting_down,
1615 leader: row.leader,
1616 global_max_workers: row.global_max_workers.map(|v| v as u32),
1617 queues: row.queues.0,
1618 })
1619 }
1620}
1621
1622#[derive(Debug, Clone, Serialize, Deserialize)]
1624pub struct RuntimeOverview {
1625 pub total_instances: usize,
1626 pub live_instances: usize,
1627 pub stale_instances: usize,
1628 pub healthy_instances: usize,
1629 pub leader_instances: usize,
1630 pub instances: Vec<RuntimeInstance>,
1631}
1632
1633#[derive(Debug, Clone, Serialize, Deserialize)]
1635pub struct QueueRuntimeSummary {
1636 pub queue: String,
1637 pub instance_count: usize,
1638 pub live_instances: usize,
1639 pub stale_instances: usize,
1640 pub healthy_instances: usize,
1641 pub total_in_flight: u64,
1642 pub overflow_held_total: Option<u64>,
1643 pub config_mismatch: bool,
1644 pub config: Option<QueueRuntimeConfigSnapshot>,
1645}
1646
1647fn queue_runtime_configs_match(
1648 left: &QueueRuntimeConfigSnapshot,
1649 right: &QueueRuntimeConfigSnapshot,
1650) -> bool {
1651 left.mode == right.mode
1652 && left.max_workers == right.max_workers
1653 && left.min_workers == right.min_workers
1654 && left.weight == right.weight
1655 && left.global_max_workers == right.global_max_workers
1656 && left.poll_interval_ms == right.poll_interval_ms
1657 && left.deadline_duration_secs == right.deadline_duration_secs
1658 && left.priority_aging_interval_secs == right.priority_aging_interval_secs
1659 && left.rate_limit == right.rate_limit
1660 && (left.dlq_enabled.is_none()
1661 || right.dlq_enabled.is_none()
1662 || left.dlq_enabled == right.dlq_enabled)
1663}
1664
1665fn representative_queue_runtime_config(
1666 configs: &[QueueRuntimeConfigSnapshot],
1667) -> Option<QueueRuntimeConfigSnapshot> {
1668 let first = configs.first()?;
1669 if first.dlq_enabled.is_some() {
1670 return Some(first.clone());
1671 }
1672
1673 configs
1674 .iter()
1675 .find(|candidate| {
1676 candidate.dlq_enabled.is_some() && queue_runtime_configs_match(first, candidate)
1677 })
1678 .cloned()
1679 .or_else(|| Some(first.clone()))
1680}
1681
1682#[derive(Debug, sqlx::FromRow)]
1683struct RuntimeInstanceRow {
1684 instance_id: Uuid,
1685 hostname: Option<String>,
1686 pid: i32,
1687 version: String,
1688 storage_capability: String,
1689 transition_role: String,
1690 started_at: DateTime<Utc>,
1691 last_seen_at: DateTime<Utc>,
1692 snapshot_interval_ms: i64,
1693 healthy: bool,
1694 postgres_connected: bool,
1695 poll_loop_alive: bool,
1696 heartbeat_alive: bool,
1697 maintenance_alive: bool,
1698 shutting_down: bool,
1699 leader: bool,
1700 global_max_workers: Option<i32>,
1701 queues: Json<Vec<QueueRuntimeSnapshot>>,
1702}
1703
1704pub async fn upsert_runtime_snapshot<'e, E>(
1706 executor: E,
1707 snapshot: &RuntimeSnapshotInput,
1708) -> Result<(), AwaError>
1709where
1710 E: PgExecutor<'e>,
1711{
1712 sqlx::query(
1713 r#"
1714 INSERT INTO awa.runtime_instances (
1715 instance_id,
1716 hostname,
1717 pid,
1718 version,
1719 storage_capability,
1720 transition_role,
1721 started_at,
1722 last_seen_at,
1723 snapshot_interval_ms,
1724 healthy,
1725 postgres_connected,
1726 poll_loop_alive,
1727 heartbeat_alive,
1728 maintenance_alive,
1729 shutting_down,
1730 leader,
1731 global_max_workers,
1732 queues,
1733 queue_descriptor_hashes,
1734 job_kind_descriptor_hashes
1735 )
1736 VALUES (
1737 $1, $2, $3, $4, $5, $6, $7, now(), $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19
1738 )
1739 ON CONFLICT (instance_id) DO UPDATE SET
1740 hostname = EXCLUDED.hostname,
1741 pid = EXCLUDED.pid,
1742 version = EXCLUDED.version,
1743 storage_capability = EXCLUDED.storage_capability,
1744 transition_role = EXCLUDED.transition_role,
1745 started_at = EXCLUDED.started_at,
1746 last_seen_at = now(),
1747 snapshot_interval_ms = EXCLUDED.snapshot_interval_ms,
1748 healthy = EXCLUDED.healthy,
1749 postgres_connected = EXCLUDED.postgres_connected,
1750 poll_loop_alive = EXCLUDED.poll_loop_alive,
1751 heartbeat_alive = EXCLUDED.heartbeat_alive,
1752 maintenance_alive = EXCLUDED.maintenance_alive,
1753 shutting_down = EXCLUDED.shutting_down,
1754 leader = EXCLUDED.leader,
1755 global_max_workers = EXCLUDED.global_max_workers,
1756 queues = EXCLUDED.queues,
1757 queue_descriptor_hashes = EXCLUDED.queue_descriptor_hashes,
1758 job_kind_descriptor_hashes = EXCLUDED.job_kind_descriptor_hashes
1759 "#,
1760 )
1761 .bind(snapshot.instance_id)
1762 .bind(snapshot.hostname.as_deref())
1763 .bind(snapshot.pid)
1764 .bind(&snapshot.version)
1765 .bind(snapshot.storage_capability.as_str())
1766 .bind(snapshot.transition_role.as_str())
1767 .bind(snapshot.started_at)
1768 .bind(snapshot.snapshot_interval_ms)
1769 .bind(snapshot.healthy)
1770 .bind(snapshot.postgres_connected)
1771 .bind(snapshot.poll_loop_alive)
1772 .bind(snapshot.heartbeat_alive)
1773 .bind(snapshot.maintenance_alive)
1774 .bind(snapshot.shutting_down)
1775 .bind(snapshot.leader)
1776 .bind(snapshot.global_max_workers.map(|v| v as i32))
1777 .bind(Json(&snapshot.queues))
1778 .bind(Json(&snapshot.queue_descriptor_hashes))
1779 .bind(Json(&snapshot.job_kind_descriptor_hashes))
1780 .execute(executor)
1781 .await?;
1782
1783 Ok(())
1784}
1785
1786pub async fn cleanup_runtime_snapshots<'e, E>(
1788 executor: E,
1789 max_age: Duration,
1790) -> Result<u64, AwaError>
1791where
1792 E: PgExecutor<'e>,
1793{
1794 let seconds = max(max_age.num_seconds(), 1);
1795 let result = sqlx::query(
1796 "DELETE FROM awa.runtime_instances WHERE last_seen_at < now() - make_interval(secs => $1)",
1797 )
1798 .bind(seconds)
1799 .execute(executor)
1800 .await?;
1801
1802 Ok(result.rows_affected())
1803}
1804
1805pub async fn cleanup_stale_descriptors<'e, E>(
1817 executor: E,
1818 table: &str,
1819 max_age: Duration,
1820) -> Result<u64, AwaError>
1821where
1822 E: PgExecutor<'e>,
1823{
1824 if !matches!(table, "awa.queue_descriptors" | "awa.job_kind_descriptors") {
1825 return Err(AwaError::Validation(format!(
1826 "cleanup_stale_descriptors: unknown table {table:?}"
1827 )));
1828 }
1829 let seconds = max(max_age.num_seconds(), 1);
1830 let sql = format!("DELETE FROM {table} WHERE last_seen_at < now() - make_interval(secs => $1)");
1833 let result = sqlx::query(&sql).bind(seconds).execute(executor).await?;
1834 Ok(result.rows_affected())
1835}
1836
1837pub async fn list_runtime_instances<'e, E>(executor: E) -> Result<Vec<RuntimeInstance>, AwaError>
1839where
1840 E: PgExecutor<'e>,
1841{
1842 let rows = sqlx::query_as::<_, RuntimeInstanceRow>(
1843 r#"
1844 SELECT
1845 instance_id,
1846 hostname,
1847 pid,
1848 version,
1849 storage_capability,
1850 transition_role,
1851 started_at,
1852 last_seen_at,
1853 snapshot_interval_ms,
1854 healthy,
1855 postgres_connected,
1856 poll_loop_alive,
1857 heartbeat_alive,
1858 maintenance_alive,
1859 shutting_down,
1860 leader,
1861 global_max_workers,
1862 queues
1863 FROM awa.runtime_instances
1864 ORDER BY leader DESC, last_seen_at DESC, started_at DESC
1865 "#,
1866 )
1867 .fetch_all(executor)
1868 .await?;
1869
1870 let now = Utc::now();
1871 rows.into_iter()
1872 .map(|row| RuntimeInstance::from_db_row(row, now))
1873 .collect()
1874}
1875
1876pub async fn runtime_overview<'e, E>(executor: E) -> Result<RuntimeOverview, AwaError>
1878where
1879 E: PgExecutor<'e>,
1880{
1881 let instances = list_runtime_instances(executor).await?;
1882 let total_instances = instances.len();
1883 let stale_instances = instances.iter().filter(|i| i.stale).count();
1884 let live_instances = total_instances.saturating_sub(stale_instances);
1885 let healthy_instances = instances.iter().filter(|i| !i.stale && i.healthy).count();
1886 let leader_instances = instances.iter().filter(|i| !i.stale && i.leader).count();
1887
1888 Ok(RuntimeOverview {
1889 total_instances,
1890 live_instances,
1891 stale_instances,
1892 healthy_instances,
1893 leader_instances,
1894 instances,
1895 })
1896}
1897
1898pub async fn queue_runtime_summary<'e, E>(executor: E) -> Result<Vec<QueueRuntimeSummary>, AwaError>
1900where
1901 E: PgExecutor<'e>,
1902{
1903 let instances = list_runtime_instances(executor).await?;
1904 let mut by_queue: HashMap<String, Vec<(bool, bool, QueueRuntimeSnapshot)>> = HashMap::new();
1905
1906 for instance in instances {
1907 let is_live = !instance.stale;
1908 let is_healthy = is_live && instance.healthy;
1909 for queue in instance.queues {
1910 by_queue
1911 .entry(queue.queue.clone())
1912 .or_default()
1913 .push((is_live, is_healthy, queue));
1914 }
1915 }
1916
1917 let mut summaries: Vec<_> = by_queue
1918 .into_iter()
1919 .map(|(queue, entries)| {
1920 let instance_count = entries.len();
1921 let live_instances = entries.iter().filter(|(live, _, _)| *live).count();
1922 let stale_instances = instance_count.saturating_sub(live_instances);
1923 let healthy_instances = entries.iter().filter(|(_, healthy, _)| *healthy).count();
1924 let total_in_flight = entries
1925 .iter()
1926 .filter(|(live, _, _)| *live)
1927 .map(|(_, _, queue)| u64::from(queue.in_flight))
1928 .sum();
1929
1930 let overflow_total: u64 = entries
1931 .iter()
1932 .filter(|(live, _, _)| *live)
1933 .filter_map(|(_, _, queue)| queue.overflow_held.map(u64::from))
1934 .sum();
1935
1936 let live_configs: Vec<_> = entries
1937 .iter()
1938 .filter(|(live, _, _)| *live)
1939 .map(|(_, _, queue)| queue.config.clone())
1940 .collect();
1941 let config_candidates = if live_configs.is_empty() {
1942 entries
1943 .iter()
1944 .map(|(_, _, queue)| queue.config.clone())
1945 .collect::<Vec<_>>()
1946 } else {
1947 live_configs
1948 };
1949 let config = representative_queue_runtime_config(&config_candidates);
1950 let config_mismatch = config_candidates.iter().skip(1).any(|candidate| {
1951 config
1952 .as_ref()
1953 .is_some_and(|cfg| !queue_runtime_configs_match(cfg, candidate))
1954 });
1955
1956 QueueRuntimeSummary {
1957 queue,
1958 instance_count,
1959 live_instances,
1960 stale_instances,
1961 healthy_instances,
1962 total_in_flight,
1963 overflow_held_total: config
1964 .as_ref()
1965 .filter(|cfg| cfg.mode == QueueRuntimeMode::Weighted)
1966 .map(|_| overflow_total),
1967 config_mismatch,
1968 config,
1969 }
1970 })
1971 .collect();
1972
1973 summaries.sort_by(|a, b| a.queue.cmp(&b.queue));
1974 Ok(summaries)
1975}
1976
1977pub async fn queue_overviews(pool: &PgPool) -> Result<Vec<QueueOverview>, AwaError> {
1990 if let Some(store) = active_queue_storage(pool).await? {
1991 let sql = format!(
1992 r#"
1993 {current_jobs_cte},
1994 all_queues AS (
1995 SELECT queue FROM current_jobs
1996 UNION
1997 SELECT queue FROM awa.queue_descriptors
1998 UNION
1999 SELECT queue FROM awa.queue_meta
2000 UNION
2001 SELECT DISTINCT descriptor.key AS queue
2002 FROM awa.runtime_instances runtime
2003 CROSS JOIN LATERAL jsonb_each_text(runtime.queue_descriptor_hashes) AS descriptor(key, value)
2004 ),
2005 live_queue_descriptor_variants AS (
2006 SELECT
2007 descriptor.key AS queue,
2008 count(DISTINCT descriptor.value)::bigint AS descriptor_variant_count
2009 FROM awa.runtime_instances runtime
2010 CROSS JOIN LATERAL jsonb_each_text(runtime.queue_descriptor_hashes) AS descriptor(key, value)
2011 WHERE runtime.last_seen_at + make_interval(
2012 secs => GREATEST(((GREATEST(runtime.snapshot_interval_ms, 1000) / 1000) * 3)::int, 30)
2013 ) >= now()
2014 GROUP BY descriptor.key
2015 ),
2016 queue_counts AS (
2017 SELECT
2018 queue,
2019 count(*) FILTER (WHERE state = 'scheduled')::bigint AS scheduled,
2020 count(*) FILTER (WHERE state = 'available')::bigint AS available,
2021 count(*) FILTER (WHERE state = 'retryable')::bigint AS retryable,
2022 count(*) FILTER (WHERE state = 'running')::bigint AS running,
2023 count(*) FILTER (WHERE state = 'failed')::bigint AS failed,
2024 count(*) FILTER (WHERE state = 'waiting_external')::bigint AS waiting_external
2025 FROM current_jobs
2026 GROUP BY queue
2027 ),
2028 available_lag AS (
2029 SELECT
2030 queue,
2031 EXTRACT(EPOCH FROM (clock_timestamp() - min(run_at)))::float8 AS lag_seconds
2032 FROM current_jobs
2033 WHERE state = 'available'
2034 GROUP BY queue
2035 ),
2036 completed_recent AS (
2037 SELECT
2038 queue,
2039 count(*)::bigint AS completed_last_hour
2040 FROM current_jobs
2041 WHERE state = 'completed'
2042 AND finalized_at > clock_timestamp() - interval '1 hour'
2043 GROUP BY queue
2044 )
2045 SELECT
2046 q.queue,
2047 qd.display_name,
2048 qd.description,
2049 qd.owner,
2050 qd.docs_url,
2051 COALESCE(qd.tags, ARRAY[]::text[]) AS tags,
2052 COALESCE(qd.extra, '{{}}'::jsonb) AS extra,
2053 qd.last_seen_at AS descriptor_last_seen_at,
2054 CASE
2055 WHEN qd.last_seen_at IS NULL THEN FALSE
2056 ELSE qd.last_seen_at + make_interval(
2057 secs => GREATEST(((COALESCE(qd.sync_interval_ms, 10000) / 1000) * 3)::int, 30)
2058 ) < now()
2059 END AS descriptor_stale,
2060 COALESCE(qdv.descriptor_variant_count, 0) > 1 AS descriptor_mismatch,
2061 COALESCE(qc.scheduled + qc.available + qc.running + qc.retryable + qc.waiting_external, 0) AS total_queued,
2062 COALESCE(qc.scheduled, 0) AS scheduled,
2063 COALESCE(qc.available, 0) AS available,
2064 COALESCE(qc.retryable, 0) AS retryable,
2065 COALESCE(qc.running, 0) AS running,
2066 COALESCE(qc.failed, 0) AS failed,
2067 COALESCE(qc.waiting_external, 0) AS waiting_external,
2068 COALESCE(cr.completed_last_hour, 0) AS completed_last_hour,
2069 al.lag_seconds,
2070 COALESCE(qm.paused, FALSE) AS paused
2071 FROM all_queues q
2072 LEFT JOIN queue_counts qc ON qc.queue = q.queue
2073 LEFT JOIN awa.queue_descriptors qd ON qd.queue = q.queue
2074 LEFT JOIN live_queue_descriptor_variants qdv ON qdv.queue = q.queue
2075 LEFT JOIN available_lag al ON al.queue = q.queue
2076 LEFT JOIN completed_recent cr ON cr.queue = q.queue
2077 LEFT JOIN awa.queue_meta qm ON qm.queue = q.queue
2078 ORDER BY q.queue
2079 "#,
2080 current_jobs_cte = queue_storage_current_jobs_cte(store.schema())
2081 );
2082
2083 let rows = sqlx::query_as::<_, QueueOverview>(&sql)
2084 .fetch_all(pool)
2085 .await?;
2086 return Ok(rows);
2087 }
2088
2089 let rows = sqlx::query_as::<_, QueueOverview>(
2090 r#"
2091 WITH all_queues AS (
2092 -- Union every source of queue-name knowledge so /queues never
2093 -- hides a queue that exists *somewhere* in the system:
2094 -- - queue_state_counts: has observed jobs
2095 -- - queue_descriptors: declared by worker code
2096 -- - queue_meta: pause state recorded by an operator
2097 -- - runtime_instances: registered by a currently-running worker
2098 SELECT queue FROM awa.queue_state_counts
2099 UNION
2100 SELECT queue FROM awa.queue_descriptors
2101 UNION
2102 SELECT queue FROM awa.queue_meta
2103 UNION
2104 SELECT DISTINCT descriptor.key AS queue
2105 FROM awa.runtime_instances runtime
2106 CROSS JOIN LATERAL jsonb_each_text(runtime.queue_descriptor_hashes) AS descriptor(key, value)
2107 ),
2108 live_queue_descriptor_variants AS (
2109 SELECT
2110 descriptor.key AS queue,
2111 count(DISTINCT descriptor.value)::bigint AS descriptor_variant_count
2112 FROM awa.runtime_instances runtime
2113 CROSS JOIN LATERAL jsonb_each_text(runtime.queue_descriptor_hashes) AS descriptor(key, value)
2114 WHERE runtime.last_seen_at + make_interval(
2115 secs => GREATEST(((GREATEST(runtime.snapshot_interval_ms, 1000) / 1000) * 3)::int, 30)
2116 ) >= now()
2117 GROUP BY descriptor.key
2118 ),
2119 available_lag AS (
2120 SELECT
2121 queue,
2122 EXTRACT(EPOCH FROM (now() - min(run_at)))::float8 AS lag_seconds
2123 FROM awa.jobs_hot
2124 WHERE state = 'available'
2125 GROUP BY queue
2126 ),
2127 completed_recent AS (
2128 SELECT
2129 queue,
2130 count(*)::bigint AS completed_last_hour
2131 FROM awa.jobs_hot
2132 WHERE state = 'completed'
2133 AND finalized_at > now() - interval '1 hour'
2134 GROUP BY queue
2135 )
2136 SELECT
2137 q.queue,
2138 qd.display_name,
2139 qd.description,
2140 qd.owner,
2141 qd.docs_url,
2142 COALESCE(qd.tags, ARRAY[]::text[]) AS tags,
2143 COALESCE(qd.extra, '{}'::jsonb) AS extra,
2144 qd.last_seen_at AS descriptor_last_seen_at,
2145 CASE
2146 WHEN qd.last_seen_at IS NULL THEN FALSE
2147 ELSE qd.last_seen_at + make_interval(
2148 secs => GREATEST(((COALESCE(qd.sync_interval_ms, 10000) / 1000) * 3)::int, 30)
2149 ) < now()
2150 END AS descriptor_stale,
2151 COALESCE(qdv.descriptor_variant_count, 0) > 1 AS descriptor_mismatch,
2152 COALESCE(qs.scheduled + qs.available + qs.running + qs.retryable + qs.waiting_external, 0) AS total_queued,
2153 COALESCE(qs.scheduled, 0) AS scheduled,
2154 COALESCE(qs.available, 0) AS available,
2155 COALESCE(qs.retryable, 0) AS retryable,
2156 COALESCE(qs.running, 0) AS running,
2157 COALESCE(qs.failed, 0) AS failed,
2158 COALESCE(qs.waiting_external, 0) AS waiting_external,
2159 COALESCE(cr.completed_last_hour, 0) AS completed_last_hour,
2160 al.lag_seconds,
2161 COALESCE(qm.paused, FALSE) AS paused
2162 FROM all_queues q
2163 LEFT JOIN awa.queue_state_counts qs ON qs.queue = q.queue
2164 LEFT JOIN awa.queue_descriptors qd ON qd.queue = q.queue
2165 LEFT JOIN live_queue_descriptor_variants qdv ON qdv.queue = q.queue
2166 LEFT JOIN available_lag al ON al.queue = q.queue
2167 LEFT JOIN completed_recent cr ON cr.queue = q.queue
2168 LEFT JOIN awa.queue_meta qm ON qm.queue = q.queue
2169 ORDER BY q.queue
2170 "#,
2171 )
2172 .fetch_all(pool)
2173 .await?;
2174
2175 Ok(rows)
2176}
2177
2178pub async fn queue_overview(pool: &PgPool, queue: &str) -> Result<Option<QueueOverview>, AwaError> {
2180 let rows = queue_overviews(pool).await?;
2181 Ok(rows.into_iter().find(|row| row.queue == queue))
2182}
2183
2184pub async fn queue_descriptors_for_names<'e, E>(
2185 executor: E,
2186 queues: &[String],
2187) -> Result<HashMap<String, QueueDescriptor>, AwaError>
2188where
2189 E: PgExecutor<'e>,
2190{
2191 if queues.is_empty() {
2192 return Ok(HashMap::new());
2193 }
2194
2195 let rows = sqlx::query_as::<
2196 _,
2197 (
2198 String,
2199 Option<String>,
2200 Option<String>,
2201 Option<String>,
2202 Option<String>,
2203 Vec<String>,
2204 serde_json::Value,
2205 ),
2206 >(
2207 r#"
2208 SELECT
2209 queue,
2210 display_name,
2211 description,
2212 owner,
2213 docs_url,
2214 tags,
2215 extra
2216 FROM awa.queue_descriptors
2217 WHERE queue = ANY($1)
2218 "#,
2219 )
2220 .bind(queues)
2221 .fetch_all(executor)
2222 .await?;
2223
2224 Ok(rows
2225 .into_iter()
2226 .map(
2227 |(queue, display_name, description, owner, docs_url, tags, extra)| {
2228 (
2229 queue,
2230 QueueDescriptor {
2231 display_name,
2232 description,
2233 owner,
2234 docs_url,
2235 tags,
2236 extra,
2237 },
2238 )
2239 },
2240 )
2241 .collect())
2242}
2243
2244pub async fn job_kind_descriptors_for_names<'e, E>(
2245 executor: E,
2246 kinds: &[String],
2247) -> Result<HashMap<String, JobKindDescriptor>, AwaError>
2248where
2249 E: PgExecutor<'e>,
2250{
2251 if kinds.is_empty() {
2252 return Ok(HashMap::new());
2253 }
2254
2255 let rows = sqlx::query_as::<
2256 _,
2257 (
2258 String,
2259 Option<String>,
2260 Option<String>,
2261 Option<String>,
2262 Option<String>,
2263 Vec<String>,
2264 serde_json::Value,
2265 ),
2266 >(
2267 r#"
2268 SELECT
2269 kind,
2270 display_name,
2271 description,
2272 owner,
2273 docs_url,
2274 tags,
2275 extra
2276 FROM awa.job_kind_descriptors
2277 WHERE kind = ANY($1)
2278 "#,
2279 )
2280 .bind(kinds)
2281 .fetch_all(executor)
2282 .await?;
2283
2284 Ok(rows
2285 .into_iter()
2286 .map(
2287 |(kind, display_name, description, owner, docs_url, tags, extra)| {
2288 (
2289 kind,
2290 JobKindDescriptor {
2291 display_name,
2292 description,
2293 owner,
2294 docs_url,
2295 tags,
2296 extra,
2297 },
2298 )
2299 },
2300 )
2301 .collect())
2302}
2303
2304#[derive(Debug, Clone, Default, Serialize)]
2306pub struct ListJobsFilter {
2307 pub state: Option<JobState>,
2308 pub kind: Option<String>,
2309 pub queue: Option<String>,
2310 pub tag: Option<String>,
2311 pub before_id: Option<i64>,
2312 pub limit: Option<i64>,
2313}
2314
2315pub async fn list_jobs(pool: &PgPool, filter: &ListJobsFilter) -> Result<Vec<JobRow>, AwaError> {
2317 if let Some(store) = active_queue_storage(pool).await? {
2318 return list_queue_storage_jobs(&store, pool, filter).await;
2319 }
2320
2321 let limit = filter.limit.unwrap_or(100);
2322
2323 let rows = sqlx::query_as::<_, JobRow>(
2324 r#"
2325 SELECT * FROM awa.jobs
2326 WHERE ($1::awa.job_state IS NULL OR state = $1)
2327 AND ($2::text IS NULL OR kind = $2)
2328 AND ($3::text IS NULL OR queue = $3)
2329 AND ($4::text IS NULL OR tags @> ARRAY[$4]::text[])
2330 AND ($5::bigint IS NULL OR id < $5)
2331 ORDER BY id DESC
2332 LIMIT $6
2333 "#,
2334 )
2335 .bind(filter.state)
2336 .bind(&filter.kind)
2337 .bind(&filter.queue)
2338 .bind(&filter.tag)
2339 .bind(filter.before_id)
2340 .bind(limit)
2341 .fetch_all(pool)
2342 .await?;
2343
2344 Ok(rows)
2345}
2346
2347pub async fn get_job(pool: &PgPool, job_id: i64) -> Result<JobRow, AwaError> {
2349 if let Some(store) = active_queue_storage(pool).await? {
2350 let row = store.load_job(pool, job_id).await?;
2351 return row.ok_or(AwaError::JobNotFound { id: job_id });
2352 }
2353
2354 let row = sqlx::query_as::<_, JobRow>("SELECT * FROM awa.jobs WHERE id = $1")
2355 .bind(job_id)
2356 .fetch_optional(pool)
2357 .await?;
2358
2359 row.ok_or(AwaError::JobNotFound { id: job_id })
2360}
2361
2362pub async fn get_job_with_source(
2365 pool: &PgPool,
2366 job_id: i64,
2367) -> Result<(JobRow, Option<DlqMetadata>), AwaError> {
2368 let job = get_job(pool, job_id).await?;
2369 let dlq = if active_queue_storage(pool).await?.is_some() {
2370 crate::dlq::get_dlq_job(pool, job_id)
2371 .await?
2372 .map(|row| row.metadata())
2373 } else {
2374 None
2375 };
2376 Ok((job, dlq))
2377}
2378
2379pub async fn dump_job(pool: &PgPool, job_id: i64) -> Result<JobDump, AwaError> {
2381 let (job, dlq) = get_job_with_source(pool, job_id).await?;
2382 Ok(build_job_dump(job, dlq))
2383}
2384
2385pub async fn dump_run(
2391 pool: &PgPool,
2392 job_id: i64,
2393 attempt: Option<i16>,
2394) -> Result<RunDump, AwaError> {
2395 let job = get_job(pool, job_id).await?;
2396 let selected_attempt = attempt.unwrap_or(job.attempt);
2397
2398 if selected_attempt < 0 {
2399 return Err(AwaError::Validation("attempt must be >= 0".to_string()));
2400 }
2401
2402 if job.attempt == 0 && selected_attempt == 0 {
2403 return Ok(RunDump {
2404 job_id: job.id,
2405 kind: job.kind.clone(),
2406 queue: job.queue.clone(),
2407 selected_attempt,
2408 current_attempt: job.attempt,
2409 current_run_lease: job.run_lease,
2410 selected_run_lease: Some(job.run_lease),
2411 source: RunDumpSource::CurrentJobRow,
2412 state: job.state,
2413 started_at: job.attempted_at,
2414 finished_at: job.finalized_at,
2415 error: None,
2416 terminal: None,
2417 progress: job.progress.clone(),
2418 metadata: Some(job.metadata.clone()),
2419 callback: callback_dump(&job),
2420 raw_error_entry: None,
2421 notes: vec!["Job has not been claimed yet; attempt 0 is the pre-run snapshot.".into()],
2422 });
2423 }
2424
2425 if selected_attempt == job.attempt {
2426 let latest_error = job
2427 .errors
2428 .as_ref()
2429 .into_iter()
2430 .flatten()
2431 .filter_map(parse_error_entry)
2432 .find(|entry| entry.attempt == Some(selected_attempt));
2433 return Ok(RunDump {
2434 job_id: job.id,
2435 kind: job.kind.clone(),
2436 queue: job.queue.clone(),
2437 selected_attempt,
2438 current_attempt: job.attempt,
2439 current_run_lease: job.run_lease,
2440 selected_run_lease: Some(job.run_lease),
2441 source: RunDumpSource::CurrentJobRow,
2442 state: job.state,
2443 started_at: job.attempted_at,
2444 finished_at: if job.state.is_terminal() {
2445 job.finalized_at
2446 } else {
2447 None
2448 },
2449 error: latest_error.as_ref().and_then(|entry| entry.error.clone()),
2450 terminal: latest_error.as_ref().map(|entry| entry.terminal),
2451 progress: job.progress.clone(),
2452 metadata: Some(job.metadata.clone()),
2453 callback: callback_dump(&job),
2454 raw_error_entry: latest_error.map(|entry| entry.raw),
2455 notes: vec![],
2456 });
2457 }
2458
2459 if selected_attempt == 0 {
2460 return Err(AwaError::Validation(format!(
2461 "attempt {selected_attempt} is not available for job {job_id}; current attempt is {}",
2462 job.attempt
2463 )));
2464 }
2465
2466 if job.attempt > 0 && selected_attempt > job.attempt {
2467 return Err(AwaError::Validation(format!(
2468 "attempt {selected_attempt} is not available for job {job_id}; current attempt is {}",
2469 job.attempt
2470 )));
2471 }
2472
2473 let historical = job
2474 .errors
2475 .as_ref()
2476 .into_iter()
2477 .flatten()
2478 .filter_map(parse_error_entry)
2479 .find(|entry| entry.attempt == Some(selected_attempt))
2480 .ok_or_else(|| {
2481 AwaError::Validation(format!(
2482 "attempt {selected_attempt} is not present in the recorded error history for job {job_id}"
2483 ))
2484 })?;
2485
2486 Ok(RunDump {
2487 job_id: job.id,
2488 kind: job.kind.clone(),
2489 queue: job.queue.clone(),
2490 selected_attempt,
2491 current_attempt: job.attempt,
2492 current_run_lease: job.run_lease,
2493 selected_run_lease: None,
2494 source: RunDumpSource::ErrorHistory,
2495 state: if historical.terminal {
2496 JobState::Failed
2497 } else {
2498 JobState::Retryable
2499 },
2500 started_at: None,
2501 finished_at: historical.at,
2502 error: historical.error,
2503 terminal: Some(historical.terminal),
2504 progress: None,
2505 metadata: None,
2506 callback: None,
2507 raw_error_entry: Some(historical.raw),
2508 notes: vec![
2509 "Historical attempts are reconstructed from errors[] because Awa does not persist a standalone runs table.".into(),
2510 "Only the current attempt has live progress, callback, and metadata fields.".into(),
2511 ],
2512 })
2513}
2514
2515pub async fn state_counts(pool: &PgPool) -> Result<HashMap<JobState, i64>, AwaError> {
2519 if let Some(store) = active_queue_storage(pool).await? {
2520 let sql = format!(
2521 r#"
2522 SELECT
2523 COALESCE((SELECT count(*)::bigint FROM {schema}.deferred_jobs WHERE state = 'scheduled'), 0) AS scheduled,
2524 COALESCE((
2525 SELECT count(*)::bigint
2526 FROM {schema}.ready_entries AS ready
2527 JOIN {schema}.queue_claim_heads AS claims
2528 ON claims.queue = ready.queue
2529 AND claims.priority = ready.priority
2530 AND claims.enqueue_shard = ready.enqueue_shard
2531 WHERE ready.lane_seq >= {schema}.sequence_next_value(claims.seq_name)
2532 AND NOT EXISTS (
2533 SELECT 1 FROM {schema}.ready_tombstones AS tomb
2534 WHERE tomb.queue = ready.queue
2535 AND tomb.priority = ready.priority
2536 AND tomb.enqueue_shard = ready.enqueue_shard
2537 AND tomb.lane_seq = ready.lane_seq
2538 AND tomb.ready_slot = ready.ready_slot
2539 AND tomb.ready_generation = ready.ready_generation
2540 )
2541 ), 0) AS available,
2542 COALESCE((SELECT count(*)::bigint FROM {schema}.leases WHERE state = 'running'), 0) AS running,
2543 COALESCE((SELECT count(*)::bigint FROM {schema}.terminal_jobs WHERE state = 'completed'), 0) AS completed,
2544 COALESCE((SELECT count(*)::bigint FROM {schema}.deferred_jobs WHERE state = 'retryable'), 0) AS retryable,
2545 COALESCE((SELECT count(*)::bigint FROM {schema}.done_entries WHERE state = 'failed'), 0)
2546 + COALESCE((SELECT count(*)::bigint FROM {schema}.dlq_entries), 0) AS failed,
2547 COALESCE((SELECT count(*)::bigint FROM {schema}.done_entries WHERE state = 'cancelled'), 0) AS cancelled,
2548 COALESCE((SELECT count(*)::bigint FROM {schema}.leases WHERE state = 'waiting_external'), 0) AS waiting_external
2549 "#,
2550 schema = store.schema()
2551 );
2552
2553 let (
2554 scheduled,
2555 available,
2556 running,
2557 completed,
2558 retryable,
2559 failed,
2560 cancelled,
2561 waiting_external,
2562 ): (i64, i64, i64, i64, i64, i64, i64, i64) = sqlx::query_as(&sql).fetch_one(pool).await?;
2563
2564 return Ok(HashMap::from([
2565 (JobState::Scheduled, scheduled),
2566 (JobState::Available, available),
2567 (JobState::Running, running),
2568 (JobState::Completed, completed),
2569 (JobState::Retryable, retryable),
2570 (JobState::Failed, failed),
2571 (JobState::Cancelled, cancelled),
2572 (JobState::WaitingExternal, waiting_external),
2573 ]));
2574 }
2575
2576 let rows = sqlx::query_as::<_, (JobState, i64)>(
2579 r#"
2580 SELECT v.state, v.total FROM (
2581 SELECT
2582 COALESCE(sum(scheduled), 0)::bigint AS scheduled,
2583 COALESCE(sum(available), 0)::bigint AS available,
2584 COALESCE(sum(running), 0)::bigint AS running,
2585 COALESCE(sum(completed), 0)::bigint AS completed,
2586 COALESCE(sum(retryable), 0)::bigint AS retryable,
2587 COALESCE(sum(failed), 0)::bigint AS failed,
2588 COALESCE(sum(cancelled), 0)::bigint AS cancelled,
2589 COALESCE(sum(waiting_external), 0)::bigint AS waiting_external
2590 FROM awa.queue_state_counts
2591 ) s,
2592 LATERAL (VALUES
2593 ('scheduled'::awa.job_state, s.scheduled),
2594 ('available'::awa.job_state, s.available),
2595 ('running'::awa.job_state, s.running),
2596 ('completed'::awa.job_state, s.completed),
2597 ('retryable'::awa.job_state, s.retryable),
2598 ('failed'::awa.job_state, s.failed),
2599 ('cancelled'::awa.job_state, s.cancelled),
2600 ('waiting_external'::awa.job_state, s.waiting_external)
2601 ) AS v(state, total)
2602 "#,
2603 )
2604 .fetch_all(pool)
2605 .await?;
2606
2607 Ok(rows.into_iter().collect())
2608}
2609
2610pub async fn job_kind_overviews(pool: &PgPool) -> Result<Vec<JobKindOverview>, AwaError> {
2612 if let Some(store) = active_queue_storage(pool).await? {
2613 let sql = format!(
2614 r#"
2615 {current_jobs_cte},
2616 all_kinds AS (
2617 SELECT kind FROM current_jobs
2618 UNION
2619 SELECT kind FROM awa.job_kind_descriptors
2620 UNION
2621 SELECT DISTINCT descriptor.key AS kind
2622 FROM awa.runtime_instances runtime
2623 CROSS JOIN LATERAL jsonb_each_text(runtime.job_kind_descriptor_hashes) AS descriptor(key, value)
2624 ),
2625 live_kind_descriptor_variants AS (
2626 SELECT
2627 descriptor.key AS kind,
2628 count(DISTINCT descriptor.value)::bigint AS descriptor_variant_count
2629 FROM awa.runtime_instances runtime
2630 CROSS JOIN LATERAL jsonb_each_text(runtime.job_kind_descriptor_hashes) AS descriptor(key, value)
2631 WHERE runtime.last_seen_at + make_interval(
2632 secs => GREATEST(((GREATEST(runtime.snapshot_interval_ms, 1000) / 1000) * 3)::int, 30)
2633 ) >= now()
2634 GROUP BY descriptor.key
2635 ),
2636 kind_counts AS (
2637 SELECT
2638 kind,
2639 count(*)::bigint AS job_count,
2640 count(DISTINCT queue)::bigint AS queue_count
2641 FROM current_jobs
2642 GROUP BY kind
2643 ),
2644 completed_recent AS (
2645 SELECT
2646 kind,
2647 count(*)::bigint AS completed_last_hour
2648 FROM current_jobs
2649 WHERE state = 'completed'
2650 AND finalized_at > clock_timestamp() - interval '1 hour'
2651 GROUP BY kind
2652 )
2653 SELECT
2654 k.kind,
2655 kd.display_name,
2656 kd.description,
2657 kd.owner,
2658 kd.docs_url,
2659 COALESCE(kd.tags, ARRAY[]::text[]) AS tags,
2660 COALESCE(kd.extra, '{{}}'::jsonb) AS extra,
2661 kd.last_seen_at AS descriptor_last_seen_at,
2662 CASE
2663 WHEN kd.last_seen_at IS NULL THEN FALSE
2664 ELSE kd.last_seen_at + make_interval(
2665 secs => GREATEST(((COALESCE(kd.sync_interval_ms, 10000) / 1000) * 3)::int, 30)
2666 ) < now()
2667 END AS descriptor_stale,
2668 COALESCE(kdv.descriptor_variant_count, 0) > 1 AS descriptor_mismatch,
2669 COALESCE(kc.job_count, 0) AS job_count,
2670 COALESCE(kc.queue_count, 0) AS queue_count,
2671 COALESCE(cr.completed_last_hour, 0) AS completed_last_hour
2672 FROM all_kinds k
2673 LEFT JOIN kind_counts kc ON kc.kind = k.kind
2674 LEFT JOIN awa.job_kind_descriptors kd ON kd.kind = k.kind
2675 LEFT JOIN live_kind_descriptor_variants kdv ON kdv.kind = k.kind
2676 LEFT JOIN completed_recent cr ON cr.kind = k.kind
2677 ORDER BY k.kind
2678 "#,
2679 current_jobs_cte = queue_storage_current_jobs_cte(store.schema())
2680 );
2681
2682 let rows = sqlx::query_as::<_, JobKindOverview>(&sql)
2683 .fetch_all(pool)
2684 .await?;
2685 return Ok(rows);
2686 }
2687
2688 let rows = sqlx::query_as::<_, JobKindOverview>(
2689 r#"
2690 WITH all_kinds AS (
2691 -- Union every source of kind-name knowledge so /kinds never
2692 -- hides a kind that exists *somewhere* in the system:
2693 -- - job_kind_catalog: has observed jobs
2694 -- - job_kind_descriptors: declared by worker code
2695 -- - runtime_instances: reported by a currently-running worker
2696 SELECT kind FROM awa.job_kind_catalog WHERE ref_count > 0
2697 UNION
2698 SELECT kind FROM awa.job_kind_descriptors
2699 UNION
2700 SELECT DISTINCT descriptor.key AS kind
2701 FROM awa.runtime_instances runtime
2702 CROSS JOIN LATERAL jsonb_each_text(runtime.job_kind_descriptor_hashes) AS descriptor(key, value)
2703 ),
2704 live_kind_descriptor_variants AS (
2705 SELECT
2706 descriptor.key AS kind,
2707 count(DISTINCT descriptor.value)::bigint AS descriptor_variant_count
2708 FROM awa.runtime_instances runtime
2709 CROSS JOIN LATERAL jsonb_each_text(runtime.job_kind_descriptor_hashes) AS descriptor(key, value)
2710 WHERE runtime.last_seen_at + make_interval(
2711 secs => GREATEST(((GREATEST(runtime.snapshot_interval_ms, 1000) / 1000) * 3)::int, 30)
2712 ) >= now()
2713 GROUP BY descriptor.key
2714 ),
2715 completed_recent AS (
2716 SELECT
2717 kind,
2718 count(*)::bigint AS completed_last_hour
2719 FROM awa.jobs_hot
2720 WHERE state = 'completed'
2721 AND finalized_at > now() - interval '1 hour'
2722 GROUP BY kind
2723 ),
2724 queue_counts AS (
2725 -- Restrict the per-kind queue fan-out to `jobs_hot` rather than
2726 -- the full `awa.jobs` view. jobs_hot is bounded by retention
2727 -- (default 24h completed / 72h failed) so the scan cost is
2728 -- tied to in-flight volume, not historical volume. The
2729 -- semantic this produces — "queues this kind is currently
2730 -- active on" — is the one admin surfaces actually care about;
2731 -- a kind that hasn't enqueued in 3 months shouldn't be
2732 -- counted as still spanning N queues.
2733 SELECT
2734 kind,
2735 count(DISTINCT queue)::bigint AS queue_count
2736 FROM awa.jobs_hot
2737 GROUP BY kind
2738 )
2739 SELECT
2740 k.kind,
2741 kd.display_name,
2742 kd.description,
2743 kd.owner,
2744 kd.docs_url,
2745 COALESCE(kd.tags, ARRAY[]::text[]) AS tags,
2746 COALESCE(kd.extra, '{}'::jsonb) AS extra,
2747 kd.last_seen_at AS descriptor_last_seen_at,
2748 CASE
2749 WHEN kd.last_seen_at IS NULL THEN FALSE
2750 ELSE kd.last_seen_at + make_interval(
2751 secs => GREATEST(((COALESCE(kd.sync_interval_ms, 10000) / 1000) * 3)::int, 30)
2752 ) < now()
2753 END AS descriptor_stale,
2754 COALESCE(kdv.descriptor_variant_count, 0) > 1 AS descriptor_mismatch,
2755 COALESCE(kc.ref_count, 0) AS job_count,
2756 COALESCE(qc.queue_count, 0) AS queue_count,
2757 COALESCE(cr.completed_last_hour, 0) AS completed_last_hour
2758 FROM all_kinds k
2759 LEFT JOIN awa.job_kind_catalog kc ON kc.kind = k.kind
2760 LEFT JOIN awa.job_kind_descriptors kd ON kd.kind = k.kind
2761 LEFT JOIN live_kind_descriptor_variants kdv ON kdv.kind = k.kind
2762 LEFT JOIN queue_counts qc ON qc.kind = k.kind
2763 LEFT JOIN completed_recent cr ON cr.kind = k.kind
2764 ORDER BY k.kind
2765 "#,
2766 )
2767 .fetch_all(pool)
2768 .await?;
2769
2770 Ok(rows)
2771}
2772
2773pub async fn job_kind_overview(
2774 pool: &PgPool,
2775 kind: &str,
2776) -> Result<Option<JobKindOverview>, AwaError> {
2777 let rows = job_kind_overviews(pool).await?;
2778 Ok(rows.into_iter().find(|row| row.kind == kind))
2779}
2780
2781pub async fn distinct_kinds(pool: &PgPool) -> Result<Vec<String>, AwaError> {
2785 if let Some(store) = active_queue_storage(pool).await? {
2786 let sql = format!(
2787 "{} \
2788 SELECT DISTINCT kind \
2789 FROM current_jobs \
2790 ORDER BY kind",
2791 queue_storage_current_jobs_cte(store.schema())
2792 );
2793 return sqlx::query_scalar(&sql)
2794 .fetch_all(pool)
2795 .await
2796 .map_err(AwaError::from);
2797 }
2798
2799 let rows = sqlx::query_scalar::<_, String>(
2800 "SELECT kind FROM awa.job_kind_catalog WHERE ref_count > 0 ORDER BY kind",
2801 )
2802 .fetch_all(pool)
2803 .await?;
2804
2805 Ok(rows)
2806}
2807
2808pub async fn distinct_queues(pool: &PgPool) -> Result<Vec<String>, AwaError> {
2812 if let Some(store) = active_queue_storage(pool).await? {
2813 let sql = format!(
2814 "{} \
2815 SELECT DISTINCT queue \
2816 FROM current_jobs \
2817 ORDER BY queue",
2818 queue_storage_current_jobs_cte(store.schema())
2819 );
2820 return sqlx::query_scalar(&sql)
2821 .fetch_all(pool)
2822 .await
2823 .map_err(AwaError::from);
2824 }
2825
2826 let rows = sqlx::query_scalar::<_, String>(
2827 "SELECT queue FROM awa.job_queue_catalog WHERE ref_count > 0 ORDER BY queue",
2828 )
2829 .fetch_all(pool)
2830 .await?;
2831
2832 Ok(rows)
2833}
2834
2835pub async fn recompute_dirty_admin_metadata(pool: &PgPool) -> Result<i32, AwaError> {
2841 if active_queue_storage(pool).await?.is_some() {
2842 return Ok(0);
2843 }
2844 let count: i32 = sqlx::query_scalar("SELECT awa.recompute_dirty_admin_metadata(100)")
2845 .fetch_one(pool)
2846 .await?;
2847 Ok(count)
2848}
2849
2850pub async fn flush_dirty_admin_metadata(pool: &PgPool) -> Result<i32, AwaError> {
2857 if active_queue_storage(pool).await?.is_some() {
2858 return Ok(0);
2859 }
2860 let mut total = 0i32;
2861 loop {
2862 let count: i32 = sqlx::query_scalar("SELECT awa.recompute_dirty_admin_metadata(100)")
2863 .fetch_one(pool)
2864 .await?;
2865 total += count;
2866 if count == 0 {
2867 break;
2868 }
2869 }
2870 Ok(total)
2871}
2872
2873pub async fn refresh_admin_metadata(pool: &PgPool) -> Result<(), AwaError> {
2879 if active_queue_storage(pool).await?.is_some() {
2880 return Ok(());
2881 }
2882 sqlx::query("SELECT awa.refresh_admin_metadata()")
2883 .execute(pool)
2884 .await?;
2885 Ok(())
2886}
2887
2888pub async fn bulk_retry(pool: &PgPool, ids: &[i64]) -> Result<Vec<JobRow>, AwaError> {
2890 if let Some(store) = active_queue_storage(pool).await? {
2891 let (rows, _attempted) = store.retry_jobs_by_ids(pool, ids).await?;
2892 return Ok(rows);
2893 }
2894
2895 let rows = sqlx::query_as::<_, JobRow>(
2896 r#"
2897 UPDATE awa.jobs
2898 SET state = 'available', attempt = 0, run_at = now(),
2899 finalized_at = NULL, heartbeat_at = NULL, deadline_at = NULL,
2900 callback_id = NULL, callback_timeout_at = NULL,
2901 callback_filter = NULL, callback_on_complete = NULL,
2902 callback_on_fail = NULL, callback_transform = NULL
2903 WHERE id = ANY($1) AND state IN ('failed', 'cancelled', 'waiting_external')
2904 RETURNING *
2905 "#,
2906 )
2907 .bind(ids)
2908 .fetch_all(pool)
2909 .await?;
2910
2911 Ok(rows)
2912}
2913
2914pub async fn bulk_cancel(pool: &PgPool, ids: &[i64]) -> Result<Vec<JobRow>, AwaError> {
2916 if let Some(store) = active_queue_storage(pool).await? {
2917 return store.cancel_jobs_by_ids(pool, ids).await;
2918 }
2919
2920 let rows = sqlx::query_as::<_, JobRow>(
2921 r#"
2922 UPDATE awa.jobs
2923 SET state = 'cancelled', finalized_at = now(),
2924 callback_id = NULL, callback_timeout_at = NULL,
2925 callback_filter = NULL, callback_on_complete = NULL,
2926 callback_on_fail = NULL, callback_transform = NULL
2927 WHERE id = ANY($1) AND state NOT IN ('completed', 'failed', 'cancelled')
2928 RETURNING *
2929 "#,
2930 )
2931 .bind(ids)
2932 .fetch_all(pool)
2933 .await?;
2934
2935 Ok(rows)
2936}
2937
2938#[derive(Debug, Clone, Serialize)]
2940pub struct StateTimeseriesBucket {
2941 pub bucket: chrono::DateTime<chrono::Utc>,
2942 pub state: JobState,
2943 pub count: i64,
2944}
2945
2946pub async fn state_timeseries(
2948 pool: &PgPool,
2949 minutes: i32,
2950) -> Result<Vec<StateTimeseriesBucket>, AwaError> {
2951 if let Some(store) = active_queue_storage(pool).await? {
2952 let sql = format!(
2953 "{} \
2954 SELECT \
2955 date_trunc('minute', created_at) AS bucket, \
2956 state, \
2957 count(*) AS count \
2958 FROM current_jobs \
2959 WHERE created_at >= clock_timestamp() - make_interval(mins => $1) \
2960 GROUP BY bucket, state \
2961 ORDER BY bucket",
2962 queue_storage_current_jobs_cte(store.schema())
2963 );
2964 let rows = sqlx::query_as::<_, (chrono::DateTime<chrono::Utc>, JobState, i64)>(&sql)
2965 .bind(minutes)
2966 .fetch_all(pool)
2967 .await?;
2968
2969 return Ok(rows
2970 .into_iter()
2971 .map(|(bucket, state, count)| StateTimeseriesBucket {
2972 bucket,
2973 state,
2974 count,
2975 })
2976 .collect());
2977 }
2978
2979 let rows = sqlx::query_as::<_, (chrono::DateTime<chrono::Utc>, JobState, i64)>(
2980 r#"
2981 SELECT
2982 date_trunc('minute', created_at) AS bucket,
2983 state,
2984 count(*) AS count
2985 FROM awa.jobs
2986 WHERE created_at >= now() - make_interval(mins => $1)
2987 GROUP BY bucket, state
2988 ORDER BY bucket
2989 "#,
2990 )
2991 .bind(minutes)
2992 .fetch_all(pool)
2993 .await?;
2994
2995 Ok(rows
2996 .into_iter()
2997 .map(|(bucket, state, count)| StateTimeseriesBucket {
2998 bucket,
2999 state,
3000 count,
3001 })
3002 .collect())
3003}
3004
3005pub async fn register_callback(
3014 pool: &PgPool,
3015 job_id: i64,
3016 run_lease: i64,
3017 timeout: std::time::Duration,
3018) -> Result<Uuid, AwaError> {
3019 if let Some(store) = active_queue_storage(pool).await? {
3020 return store
3021 .register_callback(pool, job_id, run_lease, timeout)
3022 .await;
3023 }
3024
3025 let callback_id = Uuid::new_v4();
3026 let timeout_secs = timeout.as_secs_f64();
3027 let result = sqlx::query(
3028 r#"UPDATE awa.jobs
3029 SET callback_id = $2,
3030 callback_timeout_at = now() + make_interval(secs => $3),
3031 callback_filter = NULL,
3032 callback_on_complete = NULL,
3033 callback_on_fail = NULL,
3034 callback_transform = NULL
3035 WHERE id = $1 AND state = 'running' AND run_lease = $4"#,
3036 )
3037 .bind(job_id)
3038 .bind(callback_id)
3039 .bind(timeout_secs)
3040 .bind(run_lease)
3041 .execute(pool)
3042 .await?;
3043 if result.rows_affected() == 0 {
3044 return Err(AwaError::Validation("job is not in running state".into()));
3045 }
3046 Ok(callback_id)
3047}
3048
3049pub async fn complete_external(
3060 pool: &PgPool,
3061 callback_id: Uuid,
3062 payload: Option<serde_json::Value>,
3063 run_lease: Option<i64>,
3064) -> Result<JobRow, AwaError> {
3065 let mut tx = pool.begin().await?;
3066 let row = complete_external_in_tx(&mut tx, callback_id, payload, run_lease).await?;
3067 tx.commit().await?;
3068 Ok(row)
3069}
3070
3071pub async fn complete_external_in_tx(
3077 tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
3078 callback_id: Uuid,
3079 payload: Option<serde_json::Value>,
3080 run_lease: Option<i64>,
3081) -> Result<JobRow, AwaError> {
3082 complete_external_in_tx_inner(tx, callback_id, payload, run_lease, false).await
3083}
3084
3085pub async fn resume_external(
3091 pool: &PgPool,
3092 callback_id: Uuid,
3093 payload: Option<serde_json::Value>,
3094 run_lease: Option<i64>,
3095) -> Result<JobRow, AwaError> {
3096 let mut tx = pool.begin().await?;
3097 let row = resume_external_in_tx(&mut tx, callback_id, payload, run_lease).await?;
3098 tx.commit().await?;
3099 Ok(row)
3100}
3101
3102pub async fn resume_external_in_tx(
3110 tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
3111 callback_id: Uuid,
3112 payload: Option<serde_json::Value>,
3113 run_lease: Option<i64>,
3114) -> Result<JobRow, AwaError> {
3115 complete_external_in_tx_inner(tx, callback_id, payload, run_lease, true).await
3116}
3117
3118async fn complete_external_in_tx_inner(
3119 tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
3120 callback_id: Uuid,
3121 payload: Option<serde_json::Value>,
3122 run_lease: Option<i64>,
3123 resume: bool,
3124) -> Result<JobRow, AwaError> {
3125 if let Some(store) = active_queue_storage_in_tx(tx).await? {
3126 return store
3127 .complete_external_in_tx(tx, callback_id, payload, run_lease, resume)
3128 .await;
3129 }
3130
3131 let row = if resume {
3132 let payload_json = payload.unwrap_or(serde_json::Value::Null);
3135 sqlx::query_as::<_, JobRow>(
3136 r#"
3137 UPDATE awa.jobs
3138 SET state = 'running',
3139 callback_id = NULL,
3140 callback_timeout_at = NULL,
3141 callback_filter = NULL,
3142 callback_on_complete = NULL,
3143 callback_on_fail = NULL,
3144 callback_transform = NULL,
3145 heartbeat_at = now(),
3146 metadata = metadata || jsonb_build_object('_awa_callback_result', $3::jsonb)
3147 WHERE callback_id = $1 AND state IN ('waiting_external', 'running')
3148 AND ($2::bigint IS NULL OR run_lease = $2)
3149 RETURNING *
3150 "#,
3151 )
3152 .bind(callback_id)
3153 .bind(run_lease)
3154 .bind(&payload_json)
3155 .fetch_optional(tx.as_mut())
3156 .await?
3157 } else {
3158 sqlx::query_as::<_, JobRow>(
3160 r#"
3161 UPDATE awa.jobs
3162 SET state = 'completed',
3163 finalized_at = now(),
3164 callback_id = NULL,
3165 callback_timeout_at = NULL,
3166 callback_filter = NULL,
3167 callback_on_complete = NULL,
3168 callback_on_fail = NULL,
3169 callback_transform = NULL,
3170 heartbeat_at = NULL,
3171 deadline_at = NULL,
3172 progress = NULL
3173 WHERE callback_id = $1 AND state IN ('waiting_external', 'running')
3174 AND ($2::bigint IS NULL OR run_lease = $2)
3175 RETURNING *
3176 "#,
3177 )
3178 .bind(callback_id)
3179 .bind(run_lease)
3180 .fetch_optional(tx.as_mut())
3181 .await?
3182 };
3183
3184 row.ok_or(AwaError::CallbackNotFound {
3185 callback_id: callback_id.to_string(),
3186 })
3187}
3188
3189pub async fn fail_external(
3193 pool: &PgPool,
3194 callback_id: Uuid,
3195 error: &str,
3196 run_lease: Option<i64>,
3197) -> Result<JobRow, AwaError> {
3198 let mut tx = pool.begin().await?;
3199 let row = fail_external_in_tx(&mut tx, callback_id, error, run_lease).await?;
3200 tx.commit().await?;
3201 Ok(row)
3202}
3203
3204pub async fn fail_external_in_tx(
3206 tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
3207 callback_id: Uuid,
3208 error: &str,
3209 run_lease: Option<i64>,
3210) -> Result<JobRow, AwaError> {
3211 if let Some(store) = active_queue_storage_in_tx(tx).await? {
3212 return store
3213 .fail_external_with_error_entry_in_tx(
3214 tx,
3215 callback_id,
3216 serde_json::json!({ "error": error }),
3217 run_lease,
3218 )
3219 .await;
3220 }
3221
3222 let row = sqlx::query_as::<_, JobRow>(
3223 r#"
3224 UPDATE awa.jobs
3225 SET state = 'failed',
3226 finalized_at = now(),
3227 callback_id = NULL,
3228 callback_timeout_at = NULL,
3229 callback_filter = NULL,
3230 callback_on_complete = NULL,
3231 callback_on_fail = NULL,
3232 callback_transform = NULL,
3233 heartbeat_at = NULL,
3234 deadline_at = NULL,
3235 errors = errors || jsonb_build_object(
3236 'error', $2::text,
3237 'attempt', attempt,
3238 'at', now()
3239 )::jsonb
3240 WHERE callback_id = $1 AND state IN ('waiting_external', 'running')
3241 AND ($3::bigint IS NULL OR run_lease = $3)
3242 RETURNING *
3243 "#,
3244 )
3245 .bind(callback_id)
3246 .bind(error)
3247 .bind(run_lease)
3248 .fetch_optional(tx.as_mut())
3249 .await?;
3250
3251 row.ok_or(AwaError::CallbackNotFound {
3252 callback_id: callback_id.to_string(),
3253 })
3254}
3255
3256pub async fn retry_external(
3266 pool: &PgPool,
3267 callback_id: Uuid,
3268 run_lease: Option<i64>,
3269) -> Result<JobRow, AwaError> {
3270 let mut tx = pool.begin().await?;
3271 let row = retry_external_in_tx(&mut tx, callback_id, run_lease).await?;
3272 tx.commit().await?;
3273 Ok(row)
3274}
3275
3276pub async fn retry_external_in_tx(
3278 tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
3279 callback_id: Uuid,
3280 run_lease: Option<i64>,
3281) -> Result<JobRow, AwaError> {
3282 if let Some(store) = active_queue_storage_in_tx(tx).await? {
3283 return store.retry_external_in_tx(tx, callback_id, run_lease).await;
3284 }
3285
3286 let row = sqlx::query_as::<_, JobRow>(
3287 r#"
3288 UPDATE awa.jobs
3289 SET state = 'available',
3290 attempt = 0,
3291 run_at = now(),
3292 finalized_at = NULL,
3293 callback_id = NULL,
3294 callback_timeout_at = NULL,
3295 callback_filter = NULL,
3296 callback_on_complete = NULL,
3297 callback_on_fail = NULL,
3298 callback_transform = NULL,
3299 heartbeat_at = NULL,
3300 deadline_at = NULL
3301 WHERE callback_id = $1 AND state = 'waiting_external'
3302 AND ($2::bigint IS NULL OR run_lease = $2)
3303 RETURNING *
3304 "#,
3305 )
3306 .bind(callback_id)
3307 .bind(run_lease)
3308 .fetch_optional(tx.as_mut())
3309 .await?;
3310
3311 row.ok_or(AwaError::CallbackNotFound {
3312 callback_id: callback_id.to_string(),
3313 })
3314}
3315
3316pub async fn heartbeat_callback(
3325 pool: &PgPool,
3326 callback_id: Uuid,
3327 timeout: std::time::Duration,
3328) -> Result<JobRow, AwaError> {
3329 if let Some(store) = active_queue_storage(pool).await? {
3330 return store.heartbeat_callback(pool, callback_id, timeout).await;
3331 }
3332
3333 let timeout_secs = timeout.as_secs_f64();
3334 let row = sqlx::query_as::<_, JobRow>(
3335 r#"
3336 UPDATE awa.jobs
3337 SET callback_timeout_at = now() + make_interval(secs => $2)
3338 WHERE callback_id = $1 AND state = 'waiting_external'
3339 RETURNING *
3340 "#,
3341 )
3342 .bind(callback_id)
3343 .bind(timeout_secs)
3344 .fetch_optional(pool)
3345 .await?;
3346
3347 row.ok_or(AwaError::CallbackNotFound {
3348 callback_id: callback_id.to_string(),
3349 })
3350}
3351
3352pub async fn cancel_callback(pool: &PgPool, job_id: i64, run_lease: i64) -> Result<bool, AwaError> {
3358 if let Some(store) = active_queue_storage(pool).await? {
3359 return store.cancel_callback(pool, job_id, run_lease).await;
3360 }
3361
3362 let result = sqlx::query(
3363 r#"
3364 UPDATE awa.jobs
3365 SET callback_id = NULL,
3366 callback_timeout_at = NULL,
3367 callback_filter = NULL,
3368 callback_on_complete = NULL,
3369 callback_on_fail = NULL,
3370 callback_transform = NULL
3371 WHERE id = $1 AND callback_id IS NOT NULL AND state = 'running' AND run_lease = $2
3372 "#,
3373 )
3374 .bind(job_id)
3375 .bind(run_lease)
3376 .execute(pool)
3377 .await?;
3378
3379 Ok(result.rows_affected() > 0)
3380}
3381
3382#[derive(Debug)]
3390pub enum CallbackPollResult {
3391 Resolved(serde_json::Value),
3393 Pending,
3395 Stale {
3397 token: Uuid,
3398 current: Uuid,
3399 state: JobState,
3400 },
3401 UnexpectedState { token: Uuid, state: JobState },
3403 NotFound,
3405}
3406
3407pub async fn enter_callback_wait(
3412 pool: &PgPool,
3413 job_id: i64,
3414 run_lease: i64,
3415 callback_id: Uuid,
3416) -> Result<bool, AwaError> {
3417 if let Some(store) = active_queue_storage(pool).await? {
3418 return store
3419 .enter_callback_wait(pool, job_id, run_lease, callback_id)
3420 .await;
3421 }
3422
3423 let result = sqlx::query(
3424 r#"
3425 UPDATE awa.jobs
3426 SET state = 'waiting_external',
3427 heartbeat_at = NULL,
3428 deadline_at = NULL
3429 WHERE id = $1 AND state = 'running' AND run_lease = $2 AND callback_id = $3
3430 "#,
3431 )
3432 .bind(job_id)
3433 .bind(run_lease)
3434 .bind(callback_id)
3435 .execute(pool)
3436 .await?;
3437
3438 Ok(result.rows_affected() > 0)
3439}
3440
3441pub async fn check_callback_state(
3447 pool: &PgPool,
3448 job_id: i64,
3449 callback_id: Uuid,
3450) -> Result<CallbackPollResult, AwaError> {
3451 if let Some(store) = active_queue_storage(pool).await? {
3452 return store.check_callback_state(pool, job_id, callback_id).await;
3453 }
3454
3455 let row: Option<(JobState, Option<Uuid>, serde_json::Value)> =
3456 sqlx::query_as("SELECT state, callback_id, metadata FROM awa.jobs WHERE id = $1")
3457 .bind(job_id)
3458 .fetch_optional(pool)
3459 .await?;
3460
3461 match row {
3462 Some((JobState::Running, None, metadata))
3463 if metadata.get("_awa_callback_result").is_some() =>
3464 {
3465 let payload = take_callback_payload(pool, job_id, metadata).await?;
3466 Ok(CallbackPollResult::Resolved(payload))
3467 }
3468 Some((state, Some(current_callback_id), _)) if current_callback_id != callback_id => {
3469 Ok(CallbackPollResult::Stale {
3470 token: callback_id,
3471 current: current_callback_id,
3472 state,
3473 })
3474 }
3475 Some((JobState::WaitingExternal, Some(current), _)) if current == callback_id => {
3476 Ok(CallbackPollResult::Pending)
3477 }
3478 Some((state, _, _)) => Ok(CallbackPollResult::UnexpectedState {
3479 token: callback_id,
3480 state,
3481 }),
3482 None => Ok(CallbackPollResult::NotFound),
3483 }
3484}
3485
3486pub async fn take_callback_payload(
3488 pool: &PgPool,
3489 job_id: i64,
3490 metadata: serde_json::Value,
3491) -> Result<serde_json::Value, AwaError> {
3492 let payload = metadata
3493 .get("_awa_callback_result")
3494 .cloned()
3495 .unwrap_or(serde_json::Value::Null);
3496
3497 sqlx::query("UPDATE awa.jobs SET metadata = metadata - '_awa_callback_result' WHERE id = $1")
3498 .bind(job_id)
3499 .execute(pool)
3500 .await?;
3501
3502 Ok(payload)
3503}
3504
3505#[derive(Debug, Clone, Default)]
3512pub struct CallbackConfig {
3513 pub filter: Option<String>,
3515 pub on_complete: Option<String>,
3517 pub on_fail: Option<String>,
3519 pub transform: Option<String>,
3521}
3522
3523impl CallbackConfig {
3524 pub fn is_empty(&self) -> bool {
3526 self.filter.is_none()
3527 && self.on_complete.is_none()
3528 && self.on_fail.is_none()
3529 && self.transform.is_none()
3530 }
3531}
3532
3533#[derive(Debug, Clone, Copy, PartialEq, Eq)]
3536pub enum DefaultAction {
3537 Complete,
3538 Fail,
3539 Ignore,
3540}
3541
3542#[derive(Debug)]
3544pub enum ResolveOutcome {
3545 Completed {
3546 payload: Option<serde_json::Value>,
3547 job: JobRow,
3548 },
3549 Failed {
3550 job: JobRow,
3551 },
3552 Ignored {
3553 reason: String,
3554 },
3555}
3556
3557impl ResolveOutcome {
3558 pub fn is_completed(&self) -> bool {
3559 matches!(self, ResolveOutcome::Completed { .. })
3560 }
3561 pub fn is_failed(&self) -> bool {
3562 matches!(self, ResolveOutcome::Failed { .. })
3563 }
3564 pub fn is_ignored(&self) -> bool {
3565 matches!(self, ResolveOutcome::Ignored { .. })
3566 }
3567}
3568
3569pub async fn register_callback_with_config(
3578 pool: &PgPool,
3579 job_id: i64,
3580 run_lease: i64,
3581 timeout: std::time::Duration,
3582 config: &CallbackConfig,
3583) -> Result<Uuid, AwaError> {
3584 #[cfg(feature = "cel")]
3586 {
3587 for (name, expr) in [
3588 ("filter", &config.filter),
3589 ("on_complete", &config.on_complete),
3590 ("on_fail", &config.on_fail),
3591 ("transform", &config.transform),
3592 ] {
3593 if let Some(src) = expr {
3594 let program = cel::Program::compile(src).map_err(|e| {
3595 AwaError::Validation(format!("invalid CEL expression for {name}: {e}"))
3596 })?;
3597
3598 let refs = program.references();
3602 let bad_vars: Vec<&str> = refs
3603 .variables()
3604 .into_iter()
3605 .filter(|v| *v != "payload")
3606 .collect();
3607 if !bad_vars.is_empty() {
3608 return Err(AwaError::Validation(format!(
3609 "CEL expression for {name} references undeclared variable(s): {}; \
3610 only 'payload' is available",
3611 bad_vars.join(", ")
3612 )));
3613 }
3614 }
3615 }
3616 }
3617
3618 #[cfg(not(feature = "cel"))]
3619 {
3620 if !config.is_empty() {
3621 return Err(AwaError::Validation(
3622 "CEL expressions require the 'cel' feature".into(),
3623 ));
3624 }
3625 }
3626
3627 if let Some(store) = active_queue_storage(pool).await? {
3628 return store
3629 .register_callback_with_config(pool, job_id, run_lease, timeout, config)
3630 .await;
3631 }
3632
3633 let callback_id = Uuid::new_v4();
3634 let timeout_secs = timeout.as_secs_f64();
3635
3636 let result = sqlx::query(
3637 r#"UPDATE awa.jobs
3638 SET callback_id = $2,
3639 callback_timeout_at = now() + make_interval(secs => $3),
3640 callback_filter = $4,
3641 callback_on_complete = $5,
3642 callback_on_fail = $6,
3643 callback_transform = $7
3644 WHERE id = $1 AND state = 'running' AND run_lease = $8"#,
3645 )
3646 .bind(job_id)
3647 .bind(callback_id)
3648 .bind(timeout_secs)
3649 .bind(&config.filter)
3650 .bind(&config.on_complete)
3651 .bind(&config.on_fail)
3652 .bind(&config.transform)
3653 .bind(run_lease)
3654 .execute(pool)
3655 .await?;
3656
3657 if result.rows_affected() == 0 {
3658 return Err(AwaError::Validation("job is not in running state".into()));
3659 }
3660 Ok(callback_id)
3661}
3662
3663enum ResolveAction {
3665 Complete(Option<serde_json::Value>),
3666 Fail {
3667 error: String,
3668 expression: Option<String>,
3669 },
3670 Ignore(String),
3671}
3672
3673pub async fn resolve_callback(
3679 pool: &PgPool,
3680 callback_id: Uuid,
3681 payload: Option<serde_json::Value>,
3682 default_action: DefaultAction,
3683 run_lease: Option<i64>,
3684) -> Result<ResolveOutcome, AwaError> {
3685 let mut tx = pool.begin().await?;
3686 let outcome =
3687 resolve_callback_in_tx(&mut tx, callback_id, payload, default_action, run_lease).await?;
3688 tx.commit().await?;
3689 Ok(outcome)
3690}
3691
3692pub async fn resolve_callback_in_tx(
3706 tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
3707 callback_id: Uuid,
3708 payload: Option<serde_json::Value>,
3709 default_action: DefaultAction,
3710 run_lease: Option<i64>,
3711) -> Result<ResolveOutcome, AwaError> {
3712 if let Some(store) = active_queue_storage_in_tx(tx).await? {
3713 let job = store
3714 .callback_job_in_tx(tx, callback_id, run_lease, true)
3715 .await?
3716 .ok_or(AwaError::CallbackNotFound {
3717 callback_id: callback_id.to_string(),
3718 })?;
3719
3720 let action = evaluate_or_default(&job, &payload, default_action)?;
3721
3722 return match action {
3723 ResolveAction::Complete(transformed_payload) => {
3724 let completed_job = store
3725 .complete_external_in_tx(tx, callback_id, None, run_lease, false)
3726 .await?;
3727 Ok(ResolveOutcome::Completed {
3728 payload: transformed_payload,
3729 job: completed_job,
3730 })
3731 }
3732 ResolveAction::Fail { error, expression } => {
3733 let mut error_json = serde_json::json!({
3734 "error": error,
3735 "attempt": job.attempt,
3736 "at": chrono::Utc::now().to_rfc3339(),
3737 });
3738 if let Some(expr) = expression {
3739 error_json["expression"] = serde_json::Value::String(expr);
3740 }
3741
3742 let failed_job = store
3743 .fail_external_with_error_entry_in_tx(tx, callback_id, error_json, run_lease)
3744 .await?;
3745 Ok(ResolveOutcome::Failed { job: failed_job })
3746 }
3747 ResolveAction::Ignore(reason) => Ok(ResolveOutcome::Ignored { reason }),
3748 };
3749 }
3750
3751 let job = sqlx::query_as::<_, JobRow>(
3760 "SELECT * FROM awa.jobs_hot WHERE callback_id = $1
3761 AND state IN ('waiting_external', 'running')
3762 AND ($2::bigint IS NULL OR run_lease = $2)
3763 FOR UPDATE",
3764 )
3765 .bind(callback_id)
3766 .bind(run_lease)
3767 .fetch_optional(tx.as_mut())
3768 .await?
3769 .ok_or(AwaError::CallbackNotFound {
3770 callback_id: callback_id.to_string(),
3771 })?;
3772
3773 let action = evaluate_or_default(&job, &payload, default_action)?;
3774
3775 match action {
3776 ResolveAction::Complete(transformed_payload) => {
3777 let completed_job = sqlx::query_as::<_, JobRow>(
3778 r#"
3779 UPDATE awa.jobs
3780 SET state = 'completed',
3781 finalized_at = now(),
3782 callback_id = NULL,
3783 callback_timeout_at = NULL,
3784 callback_filter = NULL,
3785 callback_on_complete = NULL,
3786 callback_on_fail = NULL,
3787 callback_transform = NULL,
3788 heartbeat_at = NULL,
3789 deadline_at = NULL,
3790 progress = NULL
3791 WHERE id = $1
3792 RETURNING *
3793 "#,
3794 )
3795 .bind(job.id)
3796 .fetch_one(tx.as_mut())
3797 .await?;
3798
3799 Ok(ResolveOutcome::Completed {
3800 payload: transformed_payload,
3801 job: completed_job,
3802 })
3803 }
3804 ResolveAction::Fail { error, expression } => {
3805 let mut error_json = serde_json::json!({
3806 "error": error,
3807 "attempt": job.attempt,
3808 "at": chrono::Utc::now().to_rfc3339(),
3809 });
3810 if let Some(expr) = expression {
3811 error_json["expression"] = serde_json::Value::String(expr);
3812 }
3813
3814 let failed_job = sqlx::query_as::<_, JobRow>(
3815 r#"
3816 UPDATE awa.jobs
3817 SET state = 'failed',
3818 finalized_at = now(),
3819 callback_id = NULL,
3820 callback_timeout_at = NULL,
3821 callback_filter = NULL,
3822 callback_on_complete = NULL,
3823 callback_on_fail = NULL,
3824 callback_transform = NULL,
3825 heartbeat_at = NULL,
3826 deadline_at = NULL,
3827 errors = errors || $2::jsonb
3828 WHERE id = $1
3829 RETURNING *
3830 "#,
3831 )
3832 .bind(job.id)
3833 .bind(error_json)
3834 .fetch_one(tx.as_mut())
3835 .await?;
3836
3837 Ok(ResolveOutcome::Failed { job: failed_job })
3838 }
3839 ResolveAction::Ignore(reason) => {
3840 Ok(ResolveOutcome::Ignored { reason })
3843 }
3844 }
3845}
3846
3847fn evaluate_or_default(
3849 job: &JobRow,
3850 payload: &Option<serde_json::Value>,
3851 default_action: DefaultAction,
3852) -> Result<ResolveAction, AwaError> {
3853 let has_expressions = job.callback_filter.is_some()
3854 || job.callback_on_complete.is_some()
3855 || job.callback_on_fail.is_some()
3856 || job.callback_transform.is_some();
3857
3858 if !has_expressions {
3859 return Ok(apply_default(default_action, payload));
3860 }
3861
3862 #[cfg(feature = "cel")]
3863 {
3864 Ok(evaluate_cel(job, payload, default_action))
3865 }
3866
3867 #[cfg(not(feature = "cel"))]
3868 {
3869 let _ = (payload, default_action);
3872 Err(AwaError::Validation(
3873 "CEL expressions present but 'cel' feature is not enabled".into(),
3874 ))
3875 }
3876}
3877
3878fn apply_default(
3879 default_action: DefaultAction,
3880 payload: &Option<serde_json::Value>,
3881) -> ResolveAction {
3882 match default_action {
3883 DefaultAction::Complete => ResolveAction::Complete(payload.clone()),
3884 DefaultAction::Fail => ResolveAction::Fail {
3885 error: "callback failed: default action".to_string(),
3886 expression: None,
3887 },
3888 DefaultAction::Ignore => {
3889 ResolveAction::Ignore("no expressions configured, default is ignore".to_string())
3890 }
3891 }
3892}
3893
3894#[cfg(feature = "cel")]
3895fn evaluate_cel(
3896 job: &JobRow,
3897 payload: &Option<serde_json::Value>,
3898 default_action: DefaultAction,
3899) -> ResolveAction {
3900 let payload_value = payload.as_ref().cloned().unwrap_or(serde_json::Value::Null);
3901
3902 if let Some(filter_expr) = &job.callback_filter {
3904 match eval_bool(filter_expr, &payload_value, job.id, "filter") {
3905 Ok(true) => {} Ok(false) => {
3907 return ResolveAction::Ignore("filter expression returned false".to_string());
3908 }
3909 Err(_) => {
3910 }
3912 }
3913 }
3914
3915 if let Some(on_fail_expr) = &job.callback_on_fail {
3917 match eval_bool(on_fail_expr, &payload_value, job.id, "on_fail") {
3918 Ok(true) => {
3919 return ResolveAction::Fail {
3920 error: "callback failed: on_fail expression matched".to_string(),
3921 expression: Some(on_fail_expr.clone()),
3922 };
3923 }
3924 Ok(false) => {} Err(_) => {
3926 }
3928 }
3929 }
3930
3931 if let Some(on_complete_expr) = &job.callback_on_complete {
3933 match eval_bool(on_complete_expr, &payload_value, job.id, "on_complete") {
3934 Ok(true) => {
3935 let transformed = apply_transform(job, &payload_value);
3937 return ResolveAction::Complete(Some(transformed));
3938 }
3939 Ok(false) => {} Err(_) => {
3941 }
3943 }
3944 }
3945
3946 apply_default(default_action, payload)
3948}
3949
3950#[cfg(feature = "cel")]
3951fn eval_bool(
3952 expression: &str,
3953 payload_value: &serde_json::Value,
3954 job_id: i64,
3955 expression_name: &str,
3956) -> Result<bool, ()> {
3957 let program = match cel::Program::compile(expression) {
3958 Ok(p) => p,
3959 Err(e) => {
3960 tracing::warn!(
3961 job_id,
3962 expression_name,
3963 expression,
3964 error = %e,
3965 "CEL compilation error during evaluation"
3966 );
3967 return Err(());
3968 }
3969 };
3970
3971 let mut context = cel::Context::default();
3972 if let Err(e) = context.add_variable("payload", payload_value.clone()) {
3973 tracing::warn!(
3974 job_id,
3975 expression_name,
3976 error = %e,
3977 "Failed to add payload variable to CEL context"
3978 );
3979 return Err(());
3980 }
3981
3982 match program.execute(&context) {
3983 Ok(cel::Value::Bool(b)) => Ok(b),
3984 Ok(other) => {
3985 tracing::warn!(
3986 job_id,
3987 expression_name,
3988 expression,
3989 result_type = ?other.type_of(),
3990 "CEL expression returned non-bool"
3991 );
3992 Err(())
3993 }
3994 Err(e) => {
3995 tracing::warn!(
3996 job_id,
3997 expression_name,
3998 expression,
3999 error = %e,
4000 "CEL execution error"
4001 );
4002 Err(())
4003 }
4004 }
4005}
4006
4007#[cfg(feature = "cel")]
4008fn apply_transform(job: &JobRow, payload_value: &serde_json::Value) -> serde_json::Value {
4009 let transform_expr = match &job.callback_transform {
4010 Some(expr) => expr,
4011 None => return payload_value.clone(),
4012 };
4013
4014 let program = match cel::Program::compile(transform_expr) {
4015 Ok(p) => p,
4016 Err(e) => {
4017 tracing::warn!(
4018 job_id = job.id,
4019 expression = transform_expr,
4020 error = %e,
4021 "CEL transform compilation error, using original payload"
4022 );
4023 return payload_value.clone();
4024 }
4025 };
4026
4027 let mut context = cel::Context::default();
4028 if let Err(e) = context.add_variable("payload", payload_value.clone()) {
4029 tracing::warn!(
4030 job_id = job.id,
4031 error = %e,
4032 "Failed to add payload variable for transform"
4033 );
4034 return payload_value.clone();
4035 }
4036
4037 match program.execute(&context) {
4038 Ok(value) => match value.json() {
4039 Ok(json) => json,
4040 Err(e) => {
4041 tracing::warn!(
4042 job_id = job.id,
4043 expression = transform_expr,
4044 error = %e,
4045 "CEL transform result could not be converted to JSON, using original payload"
4046 );
4047 payload_value.clone()
4048 }
4049 },
4050 Err(e) => {
4051 tracing::warn!(
4052 job_id = job.id,
4053 expression = transform_expr,
4054 error = %e,
4055 "CEL transform execution error, using original payload"
4056 );
4057 payload_value.clone()
4058 }
4059 }
4060}