sbol_db_sqlite/repo/
job.rs

1//! The async job queue over SQLite.
2//!
3//! Dequeue is a single atomic `UPDATE ... RETURNING` whose subquery picks the
4//! next ready job; SQLite serializes writers, so this needs no `SKIP LOCKED`.
5//! Timestamps are bound as `DateTime<Utc>` throughout so RFC3339 TEXT
6//! comparisons (`available_at <= now`) order correctly.
7
8use std::time::Duration;
9
10use async_trait::async_trait;
11use chrono::{DateTime, Utc};
12use sbol_db_core::{DomainError, JobId};
13use serde_json::Value;
14use sqlx::{QueryBuilder, Row, Sqlite, SqlitePool};
15
16use sbol_db_storage::{
17    EnqueueOutcome, JobAttempt, JobLogRecord, JobQueue, JobStatus, ListJobsFilter, NewJob,
18    OldestQueuedAge, QueueDepthRow, SbolJob, DEFAULT_QUEUE,
19};
20
21use crate::pool::db_err;
22
23const JOB_COLS: &str = "id, kind, status, priority, queue, payload, result, error, \
24    idempotency_key, attempts, max_attempts, available_at, leased_by, lease_expires_at, \
25    parent_job_id, correlation_id, created_at, started_at, finished_at";
26
27#[derive(Clone)]
28pub struct SqliteJobRepository {
29    pool: SqlitePool,
30}
31
32impl SqliteJobRepository {
33    pub fn new(pool: SqlitePool) -> Self {
34        Self { pool }
35    }
36
37    async fn find_by_idempotency(
38        &self,
39        kind: &str,
40        key: &str,
41    ) -> Result<Option<SbolJob>, DomainError> {
42        let row = sqlx::query(&format!(
43            "SELECT {JOB_COLS} FROM sbol_jobs \
44             WHERE kind = ? AND idempotency_key = ? \
45               AND status IN ('queued', 'running', 'succeeded') \
46             ORDER BY created_at DESC LIMIT 1"
47        ))
48        .bind(kind)
49        .bind(key)
50        .fetch_optional(&self.pool)
51        .await
52        .map_err(db_err)?;
53        row.map(row_to_job).transpose()
54    }
55
56    async fn record_attempt_started(
57        &self,
58        job_id: JobId,
59        attempt_no: i32,
60        worker_id: &str,
61    ) -> Result<(), DomainError> {
62        sqlx::query(
63            "INSERT OR IGNORE INTO sbol_job_attempts \
64             (job_id, attempt_no, worker_id, status, started_at) \
65             VALUES (?, ?, ?, 'running', ?)",
66        )
67        .bind(job_id.as_uuid().to_string())
68        .bind(attempt_no)
69        .bind(worker_id)
70        .bind(Utc::now())
71        .execute(&self.pool)
72        .await
73        .map_err(db_err)?;
74        Ok(())
75    }
76}
77
78#[async_trait]
79impl JobQueue for SqliteJobRepository {
80    async fn enqueue(&self, input: NewJob) -> Result<EnqueueOutcome, DomainError> {
81        if let Some(key) = &input.idempotency_key {
82            if let Some(existing) = self.find_by_idempotency(&input.kind, key).await? {
83                return Ok(EnqueueOutcome::AlreadyExists(existing));
84            }
85        }
86
87        let now = Utc::now();
88        let available_at = input.available_at.unwrap_or(now);
89        let queue = input.queue.unwrap_or_else(|| DEFAULT_QUEUE.to_owned());
90        let priority = input.priority.unwrap_or(0) as i64;
91        let max_attempts = input.max_attempts.unwrap_or(5) as i64;
92        let payload = serde_json::to_string(&input.payload).map_err(db_err)?;
93
94        let row = sqlx::query(&format!(
95            r#"
96            INSERT INTO sbol_jobs (
97                id, kind, status, priority, queue, payload, idempotency_key,
98                attempts, max_attempts, available_at, parent_job_id,
99                correlation_id, created_at
100            ) VALUES (?, ?, 'queued', ?, ?, ?, ?, 0, ?, ?, ?, ?, ?)
101            RETURNING {JOB_COLS}
102            "#
103        ))
104        .bind(uuid::Uuid::new_v4().to_string())
105        .bind(&input.kind)
106        .bind(priority)
107        .bind(&queue)
108        .bind(payload)
109        .bind(input.idempotency_key.as_deref())
110        .bind(max_attempts)
111        .bind(available_at)
112        .bind(input.parent_job_id.map(|j| j.as_uuid().to_string()))
113        .bind(input.correlation_id.map(|c| c.to_string()))
114        .bind(now)
115        .fetch_one(&self.pool)
116        .await
117        .map_err(db_err)?;
118        Ok(EnqueueOutcome::Inserted(row_to_job(row)?))
119    }
120
121    async fn append_log(
122        &self,
123        job_id: JobId,
124        attempt_no: Option<i32>,
125        level: &str,
126        message: &str,
127        fields: Value,
128    ) -> Result<JobLogRecord, DomainError> {
129        let fields_s = serde_json::to_string(&fields).map_err(db_err)?;
130        let row = sqlx::query(
131            "INSERT INTO sbol_job_logs (job_id, attempt_no, level, message, fields, created_at) \
132             VALUES (?, ?, ?, ?, ?, ?) \
133             RETURNING id, job_id, attempt_no, level, message, fields, created_at",
134        )
135        .bind(job_id.as_uuid().to_string())
136        .bind(attempt_no)
137        .bind(normalize_log_level(level))
138        .bind(message)
139        .bind(fields_s)
140        .bind(Utc::now())
141        .fetch_one(&self.pool)
142        .await
143        .map_err(db_err)?;
144        row_to_job_log(row)
145    }
146
147    async fn list_logs(
148        &self,
149        id: JobId,
150        after_id: Option<i64>,
151        limit: u32,
152    ) -> Result<Vec<JobLogRecord>, DomainError> {
153        let rows = sqlx::query(
154            "SELECT id, job_id, attempt_no, level, message, fields, created_at \
155             FROM sbol_job_logs WHERE job_id = ? AND id > ? ORDER BY id ASC LIMIT ?",
156        )
157        .bind(id.as_uuid().to_string())
158        .bind(after_id.unwrap_or(0))
159        .bind(limit.clamp(1, 1000) as i64)
160        .fetch_all(&self.pool)
161        .await
162        .map_err(db_err)?;
163        rows.into_iter().map(row_to_job_log).collect()
164    }
165
166    async fn dequeue(
167        &self,
168        queues: &[String],
169        worker_id: &str,
170        lease: Duration,
171    ) -> Result<Option<SbolJob>, DomainError> {
172        if queues.is_empty() {
173            return Ok(None);
174        }
175        let now = Utc::now();
176        let lease_expires = now
177            + chrono::Duration::from_std(lease).unwrap_or_else(|_| chrono::Duration::seconds(60));
178
179        let mut qb: QueryBuilder<Sqlite> =
180            QueryBuilder::new("UPDATE sbol_jobs SET status = 'running', leased_by = ");
181        qb.push_bind(worker_id.to_owned());
182        qb.push(", lease_expires_at = ").push_bind(lease_expires);
183        qb.push(", attempts = attempts + 1, started_at = COALESCE(started_at, ")
184            .push_bind(now)
185            .push(")");
186        qb.push(
187            " WHERE id = (SELECT id FROM sbol_jobs WHERE status = 'queued' AND available_at <= ",
188        )
189        .push_bind(now);
190        qb.push(" AND queue IN (");
191        {
192            let mut sep = qb.separated(", ");
193            for q in queues {
194                sep.push_bind(q.clone());
195            }
196        }
197        qb.push(") ORDER BY priority DESC, available_at ASC LIMIT 1) RETURNING ");
198        qb.push(JOB_COLS);
199
200        let row = qb
201            .build()
202            .fetch_optional(&self.pool)
203            .await
204            .map_err(db_err)?;
205        let job = row.map(row_to_job).transpose()?;
206        if let Some(j) = &job {
207            self.record_attempt_started(j.id, j.attempts, worker_id)
208                .await?;
209        }
210        Ok(job)
211    }
212
213    async fn renew_lease(
214        &self,
215        job_id: JobId,
216        worker_id: &str,
217        lease: Duration,
218    ) -> Result<bool, DomainError> {
219        let lease_expires = Utc::now()
220            + chrono::Duration::from_std(lease).unwrap_or_else(|_| chrono::Duration::seconds(60));
221        let res = sqlx::query(
222            "UPDATE sbol_jobs SET lease_expires_at = ? \
223             WHERE id = ? AND leased_by = ? AND status = 'running'",
224        )
225        .bind(lease_expires)
226        .bind(job_id.as_uuid().to_string())
227        .bind(worker_id)
228        .execute(&self.pool)
229        .await
230        .map_err(db_err)?;
231        Ok(res.rows_affected() == 1)
232    }
233
234    async fn mark_succeeded(
235        &self,
236        job_id: JobId,
237        worker_id: &str,
238        result: Option<Value>,
239    ) -> Result<(), DomainError> {
240        let result_s = result
241            .map(|v| serde_json::to_string(&v))
242            .transpose()
243            .map_err(db_err)?;
244        let mut tx = self.pool.begin().await.map_err(db_err)?;
245        let row = sqlx::query(
246            "UPDATE sbol_jobs SET status = 'succeeded', result = ?, error = NULL, \
247             leased_by = NULL, lease_expires_at = NULL, finished_at = ? \
248             WHERE id = ? AND leased_by = ? AND status = 'running' RETURNING attempts",
249        )
250        .bind(result_s)
251        .bind(Utc::now())
252        .bind(job_id.as_uuid().to_string())
253        .bind(worker_id)
254        .fetch_optional(&mut *tx)
255        .await
256        .map_err(db_err)?;
257
258        if let Some(r) = row {
259            let attempts: i64 = r.try_get("attempts").map_err(db_err)?;
260            finalize_attempt(&mut tx, job_id, attempts as i32, JobStatus::Succeeded, None).await?;
261        }
262        tx.commit().await.map_err(db_err)?;
263        Ok(())
264    }
265
266    async fn mark_failed(
267        &self,
268        job_id: JobId,
269        worker_id: &str,
270        error: &str,
271    ) -> Result<JobStatus, DomainError> {
272        let mut tx = self.pool.begin().await.map_err(db_err)?;
273        let row = sqlx::query(
274            "SELECT attempts, max_attempts FROM sbol_jobs \
275             WHERE id = ? AND leased_by = ? AND status = 'running'",
276        )
277        .bind(job_id.as_uuid().to_string())
278        .bind(worker_id)
279        .fetch_optional(&mut *tx)
280        .await
281        .map_err(db_err)?;
282
283        let Some(row) = row else {
284            tx.commit().await.map_err(db_err)?;
285            return Err(DomainError::Database(
286                "job lease lost before failure could be recorded".to_owned(),
287            ));
288        };
289        let attempts = row.try_get::<i64, _>("attempts").map_err(db_err)? as i32;
290        let max_attempts = row.try_get::<i64, _>("max_attempts").map_err(db_err)? as i32;
291
292        let next_status = if attempts >= max_attempts {
293            JobStatus::Dead
294        } else {
295            JobStatus::Queued
296        };
297
298        match next_status {
299            JobStatus::Queued => {
300                let available_at =
301                    Utc::now() + chrono::Duration::seconds(backoff_seconds(attempts));
302                sqlx::query(
303                    "UPDATE sbol_jobs SET status = 'queued', leased_by = NULL, \
304                     lease_expires_at = NULL, available_at = ?, error = ? WHERE id = ?",
305                )
306                .bind(available_at)
307                .bind(error)
308                .bind(job_id.as_uuid().to_string())
309                .execute(&mut *tx)
310                .await
311                .map_err(db_err)?;
312            }
313            JobStatus::Dead => {
314                sqlx::query(
315                    "UPDATE sbol_jobs SET status = 'dead', leased_by = NULL, \
316                     lease_expires_at = NULL, error = ?, finished_at = ? WHERE id = ?",
317                )
318                .bind(error)
319                .bind(Utc::now())
320                .bind(job_id.as_uuid().to_string())
321                .execute(&mut *tx)
322                .await
323                .map_err(db_err)?;
324            }
325            _ => unreachable!("next_status is queued or dead"),
326        }
327
328        let attempt_status = match next_status {
329            JobStatus::Dead => JobStatus::Dead,
330            _ => JobStatus::Failed,
331        };
332        finalize_attempt(&mut tx, job_id, attempts, attempt_status, Some(error)).await?;
333        tx.commit().await.map_err(db_err)?;
334        Ok(next_status)
335    }
336
337    async fn reap_expired_leases(&self) -> Result<u64, DomainError> {
338        let res = sqlx::query(
339            "UPDATE sbol_jobs SET status = 'queued', leased_by = NULL, \
340             lease_expires_at = NULL, available_at = ?, \
341             error = COALESCE(error, 'lease expired') \
342             WHERE status = 'running' AND lease_expires_at < ?",
343        )
344        .bind(Utc::now())
345        .bind(Utc::now())
346        .execute(&self.pool)
347        .await
348        .map_err(db_err)?;
349        Ok(res.rows_affected())
350    }
351
352    async fn get(&self, id: JobId) -> Result<Option<SbolJob>, DomainError> {
353        let row = sqlx::query(&format!("SELECT {JOB_COLS} FROM sbol_jobs WHERE id = ?"))
354            .bind(id.as_uuid().to_string())
355            .fetch_optional(&self.pool)
356            .await
357            .map_err(db_err)?;
358        row.map(row_to_job).transpose()
359    }
360
361    async fn list_attempts(&self, id: JobId) -> Result<Vec<JobAttempt>, DomainError> {
362        let rows = sqlx::query(
363            "SELECT id, job_id, attempt_no, worker_id, started_at, finished_at, status, error \
364             FROM sbol_job_attempts WHERE job_id = ? ORDER BY attempt_no DESC",
365        )
366        .bind(id.as_uuid().to_string())
367        .fetch_all(&self.pool)
368        .await
369        .map_err(db_err)?;
370        rows.into_iter()
371            .map(|row| {
372                Ok(JobAttempt {
373                    id: row.try_get("id").map_err(db_err)?,
374                    job_id: JobId(parse_uuid(row.try_get("job_id").map_err(db_err)?)?),
375                    attempt_no: row.try_get::<i64, _>("attempt_no").map_err(db_err)? as i32,
376                    worker_id: row.try_get("worker_id").map_err(db_err)?,
377                    started_at: row.try_get("started_at").map_err(db_err)?,
378                    finished_at: row.try_get("finished_at").map_err(db_err)?,
379                    status: JobStatus::from_db_str(
380                        row.try_get::<String, _>("status").map_err(db_err)?.as_str(),
381                    )?,
382                    error: row.try_get("error").map_err(db_err)?,
383                })
384            })
385            .collect()
386    }
387
388    async fn list(&self, filter: &ListJobsFilter) -> Result<Vec<SbolJob>, DomainError> {
389        let mut qb: QueryBuilder<Sqlite> =
390            QueryBuilder::new(format!("SELECT {JOB_COLS} FROM sbol_jobs WHERE 1=1"));
391        if let Some(kind) = &filter.kind {
392            qb.push(" AND kind = ").push_bind(kind.clone());
393        }
394        if let Some(status) = filter.status {
395            qb.push(" AND status = ").push_bind(status.as_db_str());
396        }
397        if let Some(queue) = &filter.queue {
398            qb.push(" AND queue = ").push_bind(queue.clone());
399        }
400        if let Some(correlation_id) = filter.correlation_id {
401            qb.push(" AND correlation_id = ")
402                .push_bind(correlation_id.to_string());
403        }
404        if let Some(since) = filter.since {
405            qb.push(" AND created_at >= ").push_bind(since);
406        }
407        qb.push(" ORDER BY created_at DESC, id DESC LIMIT ")
408            .push_bind(filter.limit.clamp(1, 1000) as i64);
409        let rows = qb.build().fetch_all(&self.pool).await.map_err(db_err)?;
410        rows.into_iter().map(row_to_job).collect()
411    }
412
413    async fn cancel(&self, id: JobId) -> Result<bool, DomainError> {
414        let res = sqlx::query(
415            "UPDATE sbol_jobs SET status = 'cancelled', leased_by = NULL, \
416             lease_expires_at = NULL, finished_at = COALESCE(finished_at, ?) \
417             WHERE id = ? AND status IN ('queued', 'running')",
418        )
419        .bind(Utc::now())
420        .bind(id.as_uuid().to_string())
421        .execute(&self.pool)
422        .await
423        .map_err(db_err)?;
424        Ok(res.rows_affected() == 1)
425    }
426
427    async fn current_status(&self, id: JobId) -> Result<Option<JobStatus>, DomainError> {
428        let row = sqlx::query("SELECT status FROM sbol_jobs WHERE id = ?")
429            .bind(id.as_uuid().to_string())
430            .fetch_optional(&self.pool)
431            .await
432            .map_err(db_err)?;
433        row.map(|r| JobStatus::from_db_str(&r.try_get::<String, _>("status").map_err(db_err)?))
434            .transpose()
435    }
436
437    async fn queue_depth_snapshot(&self) -> Result<Vec<QueueDepthRow>, DomainError> {
438        let rows = sqlx::query(
439            "SELECT status, queue, count(*) AS n FROM sbol_jobs \
440             WHERE status IN ('queued', 'running', 'failed', 'dead') \
441             GROUP BY status, queue",
442        )
443        .fetch_all(&self.pool)
444        .await
445        .map_err(db_err)?;
446        rows.into_iter()
447            .map(|r| {
448                Ok(QueueDepthRow {
449                    status: JobStatus::from_db_str(
450                        r.try_get::<String, _>("status").map_err(db_err)?.as_str(),
451                    )?,
452                    queue: r.try_get("queue").map_err(db_err)?,
453                    count: r.try_get("n").map_err(db_err)?,
454                })
455            })
456            .collect()
457    }
458
459    async fn oldest_queued_age(&self) -> Result<Vec<OldestQueuedAge>, DomainError> {
460        let now = Utc::now();
461        let rows = sqlx::query(
462            "SELECT queue, MIN(available_at) AS oldest FROM sbol_jobs \
463             WHERE status = 'queued' AND available_at <= ? GROUP BY queue",
464        )
465        .bind(now)
466        .fetch_all(&self.pool)
467        .await
468        .map_err(db_err)?;
469        rows.into_iter()
470            .map(|r| {
471                let oldest: DateTime<Utc> = r.try_get("oldest").map_err(db_err)?;
472                Ok(OldestQueuedAge {
473                    queue: r.try_get("queue").map_err(db_err)?,
474                    age_secs: (now - oldest).num_milliseconds() as f64 / 1000.0,
475                })
476            })
477            .collect()
478    }
479}
480
481async fn finalize_attempt(
482    tx: &mut sqlx::Transaction<'_, Sqlite>,
483    job_id: JobId,
484    attempt_no: i32,
485    status: JobStatus,
486    error: Option<&str>,
487) -> Result<(), DomainError> {
488    sqlx::query(
489        "UPDATE sbol_job_attempts SET finished_at = ?, status = ?, error = ? \
490         WHERE job_id = ? AND attempt_no = ?",
491    )
492    .bind(Utc::now())
493    .bind(status.as_db_str())
494    .bind(error)
495    .bind(job_id.as_uuid().to_string())
496    .bind(attempt_no)
497    .execute(&mut **tx)
498    .await
499    .map_err(db_err)?;
500    Ok(())
501}
502
503fn backoff_seconds(attempts: i32) -> i64 {
504    let attempts = attempts.max(1);
505    let base: i64 = 60;
506    let cap: i64 = 60 * 60;
507    let shift = attempts.saturating_sub(1).min(6) as u32;
508    base.saturating_mul(1i64 << shift).min(cap)
509}
510
511fn normalize_log_level(level: &str) -> &'static str {
512    match level {
513        "debug" => "debug",
514        "warn" => "warn",
515        "error" => "error",
516        _ => "info",
517    }
518}
519
520fn parse_uuid(s: String) -> Result<uuid::Uuid, DomainError> {
521    uuid::Uuid::parse_str(&s).map_err(db_err)
522}
523
524fn row_to_job(row: sqlx::sqlite::SqliteRow) -> Result<SbolJob, DomainError> {
525    let payload: String = row.try_get("payload").map_err(db_err)?;
526    let result: Option<String> = row.try_get("result").map_err(db_err)?;
527    let parent_job_id: Option<String> = row.try_get("parent_job_id").map_err(db_err)?;
528    let correlation_id: Option<String> = row.try_get("correlation_id").map_err(db_err)?;
529
530    Ok(SbolJob {
531        id: JobId(parse_uuid(row.try_get("id").map_err(db_err)?)?),
532        kind: row.try_get("kind").map_err(db_err)?,
533        status: JobStatus::from_db_str(
534            row.try_get::<String, _>("status").map_err(db_err)?.as_str(),
535        )?,
536        priority: row.try_get::<i64, _>("priority").map_err(db_err)? as i16,
537        queue: row.try_get("queue").map_err(db_err)?,
538        payload: serde_json::from_str(&payload).map_err(db_err)?,
539        result: result
540            .map(|s| serde_json::from_str(&s))
541            .transpose()
542            .map_err(db_err)?,
543        error: row.try_get("error").map_err(db_err)?,
544        idempotency_key: row.try_get("idempotency_key").map_err(db_err)?,
545        attempts: row.try_get::<i64, _>("attempts").map_err(db_err)? as i32,
546        max_attempts: row.try_get::<i64, _>("max_attempts").map_err(db_err)? as i32,
547        available_at: row.try_get("available_at").map_err(db_err)?,
548        leased_by: row.try_get("leased_by").map_err(db_err)?,
549        lease_expires_at: row.try_get("lease_expires_at").map_err(db_err)?,
550        parent_job_id: parent_job_id.map(parse_uuid).transpose()?.map(JobId),
551        correlation_id: correlation_id.map(parse_uuid).transpose()?,
552        created_at: row.try_get("created_at").map_err(db_err)?,
553        started_at: row.try_get("started_at").map_err(db_err)?,
554        finished_at: row.try_get("finished_at").map_err(db_err)?,
555    })
556}
557
558fn row_to_job_log(row: sqlx::sqlite::SqliteRow) -> Result<JobLogRecord, DomainError> {
559    let fields: String = row.try_get("fields").map_err(db_err)?;
560    Ok(JobLogRecord {
561        id: row.try_get("id").map_err(db_err)?,
562        job_id: JobId(parse_uuid(row.try_get("job_id").map_err(db_err)?)?),
563        attempt_no: row
564            .try_get::<Option<i64>, _>("attempt_no")
565            .map_err(db_err)?
566            .map(|n| n as i32),
567        level: row.try_get("level").map_err(db_err)?,
568        message: row.try_get("message").map_err(db_err)?,
569        fields: serde_json::from_str(&fields).map_err(db_err)?,
570        created_at: row.try_get("created_at").map_err(db_err)?,
571    })
572}
sbol_db_sqlite/repo/job.rs

sbol_db_sqlite/repo/
job.rs
