1use std::time::Duration;
2
3use chrono::{DateTime, Utc};
4use sbol_db_core::{DomainError, JobId};
5use serde_json::Value;
6use sqlx::Row;
7use uuid::Uuid;
8
9use crate::repo::db_err;
10use crate::PgPool;
11
12use sbol_db_storage::{
13 EnqueueOutcome, JobAttempt, JobLogRecord, JobStatus, ListJobsFilter, NewJob, OldestQueuedAge,
14 QueueDepthRow, SbolJob, DEFAULT_QUEUE,
15};
16
17#[derive(Clone)]
18pub struct JobRepository {
19 pool: PgPool,
20}
21
22impl JobRepository {
23 pub fn new(pool: PgPool) -> Self {
24 Self { pool }
25 }
26
27 pub fn pool(&self) -> &PgPool {
28 &self.pool
29 }
30
31 pub async fn enqueue(&self, input: NewJob) -> Result<EnqueueOutcome, DomainError> {
35 if let Some(key) = &input.idempotency_key {
36 if let Some(existing) = self.find_by_idempotency(&input.kind, key).await? {
37 return Ok(EnqueueOutcome::AlreadyExists(existing));
38 }
39 }
40
41 let queue = input
42 .queue
43 .clone()
44 .unwrap_or_else(|| DEFAULT_QUEUE.to_owned());
45 let priority = input.priority.unwrap_or(0);
46 let max_attempts = input.max_attempts.unwrap_or(5);
47
48 let row = sqlx::query(
49 r#"
50 INSERT INTO sbol_jobs (
51 kind, queue, priority, payload,
52 idempotency_key, max_attempts, available_at,
53 parent_job_id, correlation_id
54 ) VALUES (
55 $1, $2, $3, $4,
56 $5, $6, COALESCE($7, now()),
57 $8, $9
58 )
59 RETURNING id, kind, status::text AS status, priority, queue, payload,
60 result, error, idempotency_key, attempts, max_attempts,
61 available_at, leased_by, lease_expires_at,
62 parent_job_id, correlation_id,
63 created_at, started_at, finished_at
64 "#,
65 )
66 .bind(&input.kind)
67 .bind(&queue)
68 .bind(priority)
69 .bind(&input.payload)
70 .bind(input.idempotency_key.as_deref())
71 .bind(max_attempts)
72 .bind(input.available_at)
73 .bind(input.parent_job_id.map(|j| j.as_uuid()))
74 .bind(input.correlation_id)
75 .fetch_one(&self.pool)
76 .await
77 .map_err(db_err)?;
78
79 Ok(EnqueueOutcome::Inserted(row_to_job(row)?))
80 }
81
82 pub async fn append_log(
83 &self,
84 job_id: JobId,
85 attempt_no: Option<i32>,
86 level: &str,
87 message: &str,
88 fields: Value,
89 ) -> Result<JobLogRecord, DomainError> {
90 let level = normalize_log_level(level);
91 let row = sqlx::query(
92 r#"
93 INSERT INTO sbol_job_logs (job_id, attempt_no, level, message, fields)
94 VALUES ($1, $2, $3, $4, $5)
95 RETURNING id, job_id, attempt_no, level, message, fields, created_at
96 "#,
97 )
98 .bind(job_id.as_uuid())
99 .bind(attempt_no)
100 .bind(level)
101 .bind(message)
102 .bind(fields)
103 .fetch_one(&self.pool)
104 .await
105 .map_err(db_err)?;
106 row_to_job_log(row)
107 }
108
109 pub async fn list_logs(
110 &self,
111 id: JobId,
112 after_id: Option<i64>,
113 limit: u32,
114 ) -> Result<Vec<JobLogRecord>, DomainError> {
115 let limit = limit.clamp(1, 1000) as i64;
116 let after_id = after_id.unwrap_or(0);
117 let rows = sqlx::query(
118 r#"
119 SELECT id, job_id, attempt_no, level, message, fields, created_at
120 FROM sbol_job_logs
121 WHERE job_id = $1 AND id > $2
122 ORDER BY id ASC
123 LIMIT $3
124 "#,
125 )
126 .bind(id.as_uuid())
127 .bind(after_id)
128 .bind(limit)
129 .fetch_all(&self.pool)
130 .await
131 .map_err(db_err)?;
132 rows.into_iter().map(row_to_job_log).collect()
133 }
134
135 async fn find_by_idempotency(
136 &self,
137 kind: &str,
138 key: &str,
139 ) -> Result<Option<SbolJob>, DomainError> {
140 let row = sqlx::query(
141 r#"
142 SELECT id, kind, status::text AS status, priority, queue, payload,
143 result, error, idempotency_key, attempts, max_attempts,
144 available_at, leased_by, lease_expires_at,
145 parent_job_id, correlation_id,
146 created_at, started_at, finished_at
147 FROM sbol_jobs
148 WHERE kind = $1
149 AND idempotency_key = $2
150 AND status IN ('queued', 'running', 'succeeded')
151 ORDER BY created_at DESC
152 LIMIT 1
153 "#,
154 )
155 .bind(kind)
156 .bind(key)
157 .fetch_optional(&self.pool)
158 .await
159 .map_err(db_err)?;
160 row.map(row_to_job).transpose()
161 }
162
163 pub async fn dequeue(
167 &self,
168 queues: &[String],
169 worker_id: &str,
170 lease: Duration,
171 ) -> Result<Option<SbolJob>, DomainError> {
172 let lease_interval = format!("{} seconds", lease.as_secs_f64());
173 let row = sqlx::query(
174 r#"
175 WITH next AS (
176 SELECT id
177 FROM sbol_jobs
178 WHERE status = 'queued'
179 AND queue = ANY($1)
180 AND available_at <= now()
181 ORDER BY priority DESC, available_at
182 LIMIT 1
183 FOR UPDATE SKIP LOCKED
184 )
185 UPDATE sbol_jobs j
186 SET status = 'running',
187 leased_by = $2,
188 lease_expires_at = now() + $3::interval,
189 attempts = j.attempts + 1,
190 started_at = COALESCE(j.started_at, now())
191 FROM next
192 WHERE j.id = next.id
193 RETURNING j.id, j.kind, j.status::text AS status, j.priority, j.queue, j.payload,
194 j.result, j.error, j.idempotency_key, j.attempts, j.max_attempts,
195 j.available_at, j.leased_by, j.lease_expires_at,
196 j.parent_job_id, j.correlation_id,
197 j.created_at, j.started_at, j.finished_at
198 "#,
199 )
200 .bind(queues)
201 .bind(worker_id)
202 .bind(lease_interval)
203 .fetch_optional(&self.pool)
204 .await
205 .map_err(db_err)?;
206
207 let job = row.map(row_to_job).transpose()?;
208 if let Some(j) = &job {
209 self.record_attempt_started(j.id, j.attempts, worker_id)
210 .await?;
211 }
212 Ok(job)
213 }
214
215 async fn record_attempt_started(
216 &self,
217 job_id: JobId,
218 attempt_no: i32,
219 worker_id: &str,
220 ) -> Result<(), DomainError> {
221 sqlx::query(
222 r#"
223 INSERT INTO sbol_job_attempts (job_id, attempt_no, worker_id, status)
224 VALUES ($1, $2, $3, 'running')
225 ON CONFLICT (job_id, attempt_no) DO NOTHING
226 "#,
227 )
228 .bind(job_id.as_uuid())
229 .bind(attempt_no)
230 .bind(worker_id)
231 .execute(&self.pool)
232 .await
233 .map_err(db_err)?;
234 Ok(())
235 }
236
237 pub async fn renew_lease(
241 &self,
242 job_id: JobId,
243 worker_id: &str,
244 lease: Duration,
245 ) -> Result<bool, DomainError> {
246 let lease_interval = format!("{} seconds", lease.as_secs_f64());
247 let res = sqlx::query(
248 r#"
249 UPDATE sbol_jobs
250 SET lease_expires_at = now() + $3::interval
251 WHERE id = $1
252 AND leased_by = $2
253 AND status = 'running'
254 "#,
255 )
256 .bind(job_id.as_uuid())
257 .bind(worker_id)
258 .bind(lease_interval)
259 .execute(&self.pool)
260 .await
261 .map_err(db_err)?;
262 Ok(res.rows_affected() == 1)
263 }
264
265 pub async fn mark_succeeded(
266 &self,
267 job_id: JobId,
268 worker_id: &str,
269 result: Option<Value>,
270 ) -> Result<(), DomainError> {
271 let mut tx = self.pool.begin().await.map_err(db_err)?;
272 let row = sqlx::query(
273 r#"
274 UPDATE sbol_jobs
275 SET status = 'succeeded',
276 result = $3,
277 error = NULL,
278 leased_by = NULL,
279 lease_expires_at = NULL,
280 finished_at = now()
281 WHERE id = $1
282 AND leased_by = $2
283 AND status = 'running'
284 RETURNING attempts
285 "#,
286 )
287 .bind(job_id.as_uuid())
288 .bind(worker_id)
289 .bind(result)
290 .fetch_optional(&mut *tx)
291 .await
292 .map_err(db_err)?;
293
294 if let Some(r) = row {
295 let attempts: i32 = r.try_get("attempts").map_err(db_err)?;
296 finalize_attempt(&mut tx, job_id, attempts, JobStatus::Succeeded, None).await?;
297 }
298 tx.commit().await.map_err(db_err)?;
299 Ok(())
300 }
301
302 pub async fn mark_failed(
306 &self,
307 job_id: JobId,
308 worker_id: &str,
309 error: &str,
310 ) -> Result<JobStatus, DomainError> {
311 let mut tx = self.pool.begin().await.map_err(db_err)?;
312 let row = sqlx::query(
313 r#"
314 SELECT attempts, max_attempts
315 FROM sbol_jobs
316 WHERE id = $1 AND leased_by = $2 AND status = 'running'
317 FOR UPDATE
318 "#,
319 )
320 .bind(job_id.as_uuid())
321 .bind(worker_id)
322 .fetch_optional(&mut *tx)
323 .await
324 .map_err(db_err)?;
325
326 let Some(row) = row else {
327 tx.commit().await.map_err(db_err)?;
328 return Err(DomainError::Database(
329 "job lease lost before failure could be recorded".to_owned(),
330 ));
331 };
332 let attempts: i32 = row.try_get("attempts").map_err(db_err)?;
333 let max_attempts: i32 = row.try_get("max_attempts").map_err(db_err)?;
334
335 let next_status = if attempts >= max_attempts {
336 JobStatus::Dead
337 } else {
338 JobStatus::Queued
339 };
340
341 match next_status {
342 JobStatus::Queued => {
343 let backoff_secs = backoff_seconds(attempts);
344 sqlx::query(
345 r#"
346 UPDATE sbol_jobs
347 SET status = 'queued',
348 leased_by = NULL,
349 lease_expires_at = NULL,
350 available_at = now() + ($2 || ' seconds')::interval,
351 error = $3
352 WHERE id = $1
353 "#,
354 )
355 .bind(job_id.as_uuid())
356 .bind(backoff_secs.to_string())
357 .bind(error)
358 .execute(&mut *tx)
359 .await
360 .map_err(db_err)?;
361 }
362 JobStatus::Dead => {
363 sqlx::query(
364 r#"
365 UPDATE sbol_jobs
366 SET status = 'dead',
367 leased_by = NULL,
368 lease_expires_at = NULL,
369 error = $2,
370 finished_at = now()
371 WHERE id = $1
372 "#,
373 )
374 .bind(job_id.as_uuid())
375 .bind(error)
376 .execute(&mut *tx)
377 .await
378 .map_err(db_err)?;
379 }
380 _ => unreachable!("next_status is queued or dead"),
381 }
382 let _ = worker_id;
383
384 let attempt_status = match next_status {
385 JobStatus::Dead => JobStatus::Dead,
386 _ => JobStatus::Failed,
387 };
388 finalize_attempt(&mut tx, job_id, attempts, attempt_status, Some(error)).await?;
389 tx.commit().await.map_err(db_err)?;
390 Ok(next_status)
391 }
392
393 pub async fn reap_expired_leases(&self) -> Result<u64, DomainError> {
397 let res = sqlx::query(
398 r#"
399 UPDATE sbol_jobs
400 SET status = 'queued',
401 leased_by = NULL,
402 lease_expires_at = NULL,
403 available_at = now(),
404 error = COALESCE(error, 'lease expired')
405 WHERE status = 'running'
406 AND lease_expires_at < now()
407 "#,
408 )
409 .execute(&self.pool)
410 .await
411 .map_err(db_err)?;
412 Ok(res.rows_affected())
413 }
414
415 pub async fn get(&self, id: JobId) -> Result<Option<SbolJob>, DomainError> {
416 let row = sqlx::query(
417 r#"
418 SELECT id, kind, status::text AS status, priority, queue, payload,
419 result, error, idempotency_key, attempts, max_attempts,
420 available_at, leased_by, lease_expires_at,
421 parent_job_id, correlation_id,
422 created_at, started_at, finished_at
423 FROM sbol_jobs
424 WHERE id = $1
425 "#,
426 )
427 .bind(id.as_uuid())
428 .fetch_optional(&self.pool)
429 .await
430 .map_err(db_err)?;
431 row.map(row_to_job).transpose()
432 }
433
434 pub async fn list_attempts(&self, id: JobId) -> Result<Vec<JobAttempt>, DomainError> {
439 let rows = sqlx::query(
440 r#"
441 SELECT id, job_id, attempt_no, worker_id, started_at,
442 finished_at, status::text AS status, error
443 FROM sbol_job_attempts
444 WHERE job_id = $1
445 ORDER BY attempt_no DESC
446 "#,
447 )
448 .bind(id.as_uuid())
449 .fetch_all(&self.pool)
450 .await
451 .map_err(db_err)?;
452 rows.into_iter()
453 .map(|row| {
454 Ok(JobAttempt {
455 id: row.try_get("id").map_err(db_err)?,
456 job_id: JobId(row.try_get("job_id").map_err(db_err)?),
457 attempt_no: row.try_get("attempt_no").map_err(db_err)?,
458 worker_id: row.try_get("worker_id").map_err(db_err)?,
459 started_at: row.try_get("started_at").map_err(db_err)?,
460 finished_at: row.try_get("finished_at").map_err(db_err)?,
461 status: JobStatus::from_db_str(
462 row.try_get::<String, _>("status").map_err(db_err)?.as_str(),
463 )?,
464 error: row.try_get("error").map_err(db_err)?,
465 })
466 })
467 .collect()
468 }
469
470 pub async fn list(&self, filter: &ListJobsFilter) -> Result<Vec<SbolJob>, DomainError> {
471 let limit = filter.limit.clamp(1, 1000) as i64;
472 let rows = sqlx::query(
473 r#"
474 SELECT id, kind, status::text AS status, priority, queue, payload,
475 result, error, idempotency_key, attempts, max_attempts,
476 available_at, leased_by, lease_expires_at,
477 parent_job_id, correlation_id,
478 created_at, started_at, finished_at
479 FROM sbol_jobs
480 WHERE ($1::text IS NULL OR kind = $1)
481 AND ($2::text IS NULL OR status::text = $2)
482 AND ($3::text IS NULL OR queue = $3)
483 AND ($4::uuid IS NULL OR correlation_id = $4)
484 AND ($5::timestamptz IS NULL OR created_at >= $5)
485 ORDER BY created_at DESC, id DESC
486 LIMIT $6
487 "#,
488 )
489 .bind(filter.kind.as_deref())
490 .bind(filter.status.map(|s| s.as_db_str()))
491 .bind(filter.queue.as_deref())
492 .bind(filter.correlation_id)
493 .bind(filter.since)
494 .bind(limit)
495 .fetch_all(&self.pool)
496 .await
497 .map_err(db_err)?;
498 rows.into_iter().map(row_to_job).collect()
499 }
500
501 pub async fn cancel(&self, id: JobId) -> Result<bool, DomainError> {
505 let res = sqlx::query(
506 r#"
507 UPDATE sbol_jobs
508 SET status = 'cancelled',
509 leased_by = NULL,
510 lease_expires_at = NULL,
511 finished_at = COALESCE(finished_at, now())
512 WHERE id = $1
513 AND status IN ('queued', 'running')
514 "#,
515 )
516 .bind(id.as_uuid())
517 .execute(&self.pool)
518 .await
519 .map_err(db_err)?;
520 Ok(res.rows_affected() == 1)
521 }
522
523 pub async fn current_status(&self, id: JobId) -> Result<Option<JobStatus>, DomainError> {
526 let row = sqlx::query("SELECT status::text AS status FROM sbol_jobs WHERE id = $1")
527 .bind(id.as_uuid())
528 .fetch_optional(&self.pool)
529 .await
530 .map_err(db_err)?;
531 row.map(|r| {
532 let s: String = r.try_get("status").map_err(db_err)?;
533 JobStatus::from_db_str(&s)
534 })
535 .transpose()
536 }
537
538 pub async fn queue_depth_snapshot(&self) -> Result<Vec<QueueDepthRow>, DomainError> {
543 let rows = sqlx::query(
544 r#"
545 SELECT status::text AS status, queue, count(*)::bigint AS n
546 FROM sbol_jobs
547 WHERE status IN ('queued', 'running', 'failed', 'dead')
548 GROUP BY status, queue
549 "#,
550 )
551 .fetch_all(&self.pool)
552 .await
553 .map_err(db_err)?;
554 rows.into_iter()
555 .map(|r| {
556 let status_str: String = r.try_get("status").map_err(db_err)?;
557 let queue: String = r.try_get("queue").map_err(db_err)?;
558 let n: i64 = r.try_get("n").map_err(db_err)?;
559 Ok(QueueDepthRow {
560 status: JobStatus::from_db_str(&status_str)?,
561 queue,
562 count: n,
563 })
564 })
565 .collect()
566 }
567
568 pub async fn oldest_queued_age(&self) -> Result<Vec<OldestQueuedAge>, DomainError> {
572 let rows = sqlx::query(
573 r#"
574 SELECT queue, EXTRACT(EPOCH FROM (now() - MIN(available_at)))::float8 AS age_secs
575 FROM sbol_jobs
576 WHERE status = 'queued' AND available_at <= now()
577 GROUP BY queue
578 "#,
579 )
580 .fetch_all(&self.pool)
581 .await
582 .map_err(db_err)?;
583 rows.into_iter()
584 .map(|r| {
585 let queue: String = r.try_get("queue").map_err(db_err)?;
586 let age_secs: f64 = r.try_get("age_secs").map_err(db_err)?;
587 Ok(OldestQueuedAge { queue, age_secs })
588 })
589 .collect()
590 }
591
592 pub async fn recent_failure_count(&self, within_hours: i64) -> Result<i64, DomainError> {
593 let row = sqlx::query(
594 r#"
595 SELECT COUNT(*)::bigint AS n
596 FROM sbol_jobs
597 WHERE status IN ('failed', 'dead')
598 AND finished_at >= now() - make_interval(hours => $1)
599 "#,
600 )
601 .bind(within_hours.max(0) as i32)
602 .fetch_one(&self.pool)
603 .await
604 .map_err(db_err)?;
605 row.try_get::<i64, _>("n").map_err(db_err)
606 }
607}
608
609async fn finalize_attempt(
610 tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
611 job_id: JobId,
612 attempt_no: i32,
613 status: JobStatus,
614 error: Option<&str>,
615) -> Result<(), DomainError> {
616 sqlx::query(
617 r#"
618 UPDATE sbol_job_attempts
619 SET finished_at = now(),
620 status = $3::sbol_job_status,
621 error = $4
622 WHERE job_id = $1 AND attempt_no = $2
623 "#,
624 )
625 .bind(job_id.as_uuid())
626 .bind(attempt_no)
627 .bind(status.as_db_str())
628 .bind(error)
629 .execute(&mut **tx)
630 .await
631 .map_err(db_err)?;
632 Ok(())
633}
634
635fn backoff_seconds(attempts: i32) -> i64 {
636 let attempts = attempts.max(1);
637 let base: i64 = 60;
638 let cap: i64 = 60 * 60;
639 let shift = attempts.saturating_sub(1).min(6) as u32;
640 base.saturating_mul(1i64 << shift).min(cap)
641}
642
643fn normalize_log_level(level: &str) -> &'static str {
644 match level {
645 "debug" => "debug",
646 "warn" => "warn",
647 "error" => "error",
648 _ => "info",
649 }
650}
651
652fn row_to_job(row: sqlx::postgres::PgRow) -> Result<SbolJob, DomainError> {
653 let id: Uuid = row.try_get("id").map_err(db_err)?;
654 let kind: String = row.try_get("kind").map_err(db_err)?;
655 let status: String = row.try_get("status").map_err(db_err)?;
656 let priority: i16 = row.try_get("priority").map_err(db_err)?;
657 let queue: String = row.try_get("queue").map_err(db_err)?;
658 let payload: Value = row.try_get("payload").map_err(db_err)?;
659 let result: Option<Value> = row.try_get("result").map_err(db_err)?;
660 let error: Option<String> = row.try_get("error").map_err(db_err)?;
661 let idempotency_key: Option<String> = row.try_get("idempotency_key").map_err(db_err)?;
662 let attempts: i32 = row.try_get("attempts").map_err(db_err)?;
663 let max_attempts: i32 = row.try_get("max_attempts").map_err(db_err)?;
664 let available_at: DateTime<Utc> = row.try_get("available_at").map_err(db_err)?;
665 let leased_by: Option<String> = row.try_get("leased_by").map_err(db_err)?;
666 let lease_expires_at: Option<DateTime<Utc>> =
667 row.try_get("lease_expires_at").map_err(db_err)?;
668 let parent_job_id: Option<Uuid> = row.try_get("parent_job_id").map_err(db_err)?;
669 let correlation_id: Option<Uuid> = row.try_get("correlation_id").map_err(db_err)?;
670 let created_at: DateTime<Utc> = row.try_get("created_at").map_err(db_err)?;
671 let started_at: Option<DateTime<Utc>> = row.try_get("started_at").map_err(db_err)?;
672 let finished_at: Option<DateTime<Utc>> = row.try_get("finished_at").map_err(db_err)?;
673
674 Ok(SbolJob {
675 id: JobId(id),
676 kind,
677 status: JobStatus::from_db_str(&status)?,
678 priority,
679 queue,
680 payload,
681 result,
682 error,
683 idempotency_key,
684 attempts,
685 max_attempts,
686 available_at,
687 leased_by,
688 lease_expires_at,
689 parent_job_id: parent_job_id.map(JobId),
690 correlation_id,
691 created_at,
692 started_at,
693 finished_at,
694 })
695}
696
697fn row_to_job_log(row: sqlx::postgres::PgRow) -> Result<JobLogRecord, DomainError> {
698 let id: i64 = row.try_get("id").map_err(db_err)?;
699 let job_id: Uuid = row.try_get("job_id").map_err(db_err)?;
700 let attempt_no: Option<i32> = row.try_get("attempt_no").map_err(db_err)?;
701 let level: String = row.try_get("level").map_err(db_err)?;
702 let message: String = row.try_get("message").map_err(db_err)?;
703 let fields: Value = row.try_get("fields").map_err(db_err)?;
704 let created_at: DateTime<Utc> = row.try_get("created_at").map_err(db_err)?;
705 Ok(JobLogRecord {
706 id,
707 job_id: JobId(job_id),
708 attempt_no,
709 level,
710 message,
711 fields,
712 created_at,
713 })
714}