Skip to main content

runledger_postgres/jobs/
admin.rs

1use runledger_core::jobs::{JobStatus, JobType, WorkflowStepStatus};
2use serde_json::Value;
3use sqlx::types::Uuid;
4
5use crate::{DbPool, DbTx, Error, Result};
6
7use super::errors::{
8    invalid_job_state_error, job_not_found_error, validate_page_limit, validate_pagination,
9    workflow_requeue_not_supported_error,
10};
11use super::row_decode::{
12    parse_job_event_type, parse_job_stage, parse_job_status, parse_job_type_name,
13};
14use super::types::{JobEventRecord, JobListFilter, JobMetricsRecord, JobQueueRecord};
15use super::workflows::on_terminal;
16
17const JOB_PAYLOAD_UUID_ARRAY_FIELD_UPDATE_LOCK_TIMEOUT: &str = "1s";
18const JOB_PAYLOAD_UUID_ARRAY_FIELD_UPDATE_LOCK_TIMEOUT_MS: i64 = 1_000;
19
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21#[must_use = "callers must inspect Updated/NotFound/Rejected"]
22#[non_exhaustive]
23pub enum JobPayloadUuidArrayFieldUpdate {
24    Updated,
25    NotFound,
26    Rejected {
27        reason: JobPayloadUuidArrayFieldUpdateRejection,
28    },
29}
30
31#[derive(Debug, Clone, Copy, PartialEq, Eq)]
32#[non_exhaustive]
33pub enum JobPayloadUuidArrayFieldUpdateRejection {
34    WorkflowManaged,
35    IdempotentRequestSnapshot,
36    NotPendingOrClaimed,
37}
38
39async fn rollback_and_classify_missing_job_mutation(
40    tx: DbTx<'_>,
41    pool: &DbPool,
42    organization_id: Option<Uuid>,
43    job_id: Uuid,
44) -> Result<Error> {
45    if let Err(error) = tx.rollback().await {
46        tracing::warn!(error = %error, "failed to rollback missing job mutation transaction");
47    }
48    let exists = get_job_by_id(pool, organization_id, job_id).await?;
49    Ok(if exists.is_none() {
50        job_not_found_error()
51    } else {
52        invalid_job_state_error()
53    })
54}
55
56async fn workflow_managed_job_exists_tx(
57    tx: &mut DbTx<'_>,
58    job_id: Uuid,
59    organization_id: Option<Uuid>,
60) -> Result<bool> {
61    let exists: bool = sqlx::query_scalar!(
62        "SELECT EXISTS (
63            SELECT 1
64            FROM job_queue jq
65            WHERE jq.id = $1
66              AND jq.workflow_step_id IS NOT NULL
67              AND ($2::uuid IS NULL OR jq.organization_id = $2)
68         ) AS \"exists!\"",
69        job_id,
70        organization_id,
71    )
72    .fetch_one(&mut **tx)
73    .await
74    .map_err(|error| {
75        Error::from_query_sqlx_with_context("requeue workflow-managed job check", error)
76    })?;
77
78    Ok(exists)
79}
80
81#[derive(sqlx::FromRow)]
82struct JobQueueRecordRow {
83    id: Uuid,
84    job_type: String,
85    organization_id: Option<Uuid>,
86    payload: Value,
87    status: String,
88    priority: i32,
89    run_number: i32,
90    attempt: i32,
91    max_attempts: i32,
92    timeout_seconds: i32,
93    next_run_at: chrono::DateTime<chrono::Utc>,
94    lease_expires_at: Option<chrono::DateTime<chrono::Utc>>,
95    last_heartbeat_at: Option<chrono::DateTime<chrono::Utc>>,
96    worker_id: Option<String>,
97    started_at: Option<chrono::DateTime<chrono::Utc>>,
98    finished_at: Option<chrono::DateTime<chrono::Utc>>,
99    stage: String,
100    progress_done: Option<i64>,
101    progress_total: Option<i64>,
102    progress_pct: Option<f64>,
103    checkpoint: Option<Value>,
104    output: Option<Value>,
105    idempotency_key: Option<String>,
106    status_reason: Option<String>,
107    last_error_code: Option<String>,
108    last_error_message: Option<String>,
109    created_at: chrono::DateTime<chrono::Utc>,
110    updated_at: chrono::DateTime<chrono::Utc>,
111}
112
113#[derive(sqlx::FromRow)]
114struct JobPayloadUuidArrayFieldUpdateCandidate {
115    status: String,
116    worker_id: Option<String>,
117    lease_expires_at: Option<chrono::DateTime<chrono::Utc>>,
118    workflow_step_id: Option<Uuid>,
119    idempotency_key: Option<String>,
120    enqueue_request: Option<Value>,
121}
122
123fn job_queue_record_from_row(row: JobQueueRecordRow) -> Result<JobQueueRecord> {
124    Ok(JobQueueRecord {
125        id: row.id,
126        job_type: parse_job_type_name(row.job_type)?,
127        organization_id: row.organization_id,
128        payload: row.payload,
129        status: parse_job_status(row.status)?,
130        priority: row.priority,
131        run_number: row.run_number,
132        attempt: row.attempt,
133        max_attempts: row.max_attempts,
134        timeout_seconds: row.timeout_seconds,
135        next_run_at: row.next_run_at,
136        lease_expires_at: row.lease_expires_at,
137        last_heartbeat_at: row.last_heartbeat_at,
138        worker_id: row.worker_id,
139        started_at: row.started_at,
140        finished_at: row.finished_at,
141        stage: parse_job_stage(row.stage)?,
142        progress_done: row.progress_done,
143        progress_total: row.progress_total,
144        progress_pct: row.progress_pct,
145        checkpoint: row.checkpoint,
146        output: row.output,
147        idempotency_key: row.idempotency_key,
148        status_reason: row.status_reason,
149        last_error_code: row.last_error_code,
150        last_error_message: row.last_error_message,
151        created_at: row.created_at,
152        updated_at: row.updated_at,
153    })
154}
155
156pub async fn list_jobs(pool: &DbPool, filter: &JobListFilter<'_>) -> Result<Vec<JobQueueRecord>> {
157    validate_pagination(filter.limit, filter.offset)?;
158
159    let status_filter = filter.status.map(JobStatus::as_db_value);
160
161    let rows = sqlx::query!(
162        "SELECT
163            id,
164            job_type,
165            organization_id,
166            payload,
167            status::text AS \"status!\",
168            priority,
169            run_number,
170            attempt,
171            max_attempts,
172            timeout_seconds,
173            next_run_at,
174            lease_expires_at,
175            last_heartbeat_at,
176            worker_id,
177            started_at,
178            finished_at,
179            stage,
180            progress_done,
181            progress_total,
182            progress_pct::float8 AS progress_pct,
183            checkpoint,
184            output,
185            idempotency_key,
186            status_reason,
187            last_error_code,
188            last_error_message,
189            created_at,
190            updated_at
191         FROM job_queue
192         WHERE ($1::uuid IS NULL OR organization_id = $1)
193           AND ($2::text::job_status IS NULL OR status = $2::text::job_status)
194           AND ($3::text IS NULL OR job_type ILIKE '%' || $3 || '%')
195         ORDER BY created_at DESC, id DESC
196         LIMIT $4
197         OFFSET $5",
198        filter.organization_id,
199        status_filter,
200        filter.job_type,
201        filter.limit,
202        filter.offset,
203    )
204    .fetch_all(pool)
205    .await
206    .map_err(|error| Error::from_query_sqlx_with_context("list jobs", error))?;
207
208    rows.into_iter()
209        .map(|row| {
210            Ok(JobQueueRecord {
211                id: row.id,
212                job_type: parse_job_type_name(row.job_type)?,
213                organization_id: row.organization_id,
214                payload: row.payload,
215                status: parse_job_status(row.status)?,
216                priority: row.priority,
217                run_number: row.run_number,
218                attempt: row.attempt,
219                max_attempts: row.max_attempts,
220                timeout_seconds: row.timeout_seconds,
221                next_run_at: row.next_run_at,
222                lease_expires_at: row.lease_expires_at,
223                last_heartbeat_at: row.last_heartbeat_at,
224                worker_id: row.worker_id,
225                started_at: row.started_at,
226                finished_at: row.finished_at,
227                stage: parse_job_stage(row.stage)?,
228                progress_done: row.progress_done,
229                progress_total: row.progress_total,
230                progress_pct: row.progress_pct,
231                checkpoint: row.checkpoint,
232                output: row.output,
233                idempotency_key: row.idempotency_key,
234                status_reason: row.status_reason,
235                last_error_code: row.last_error_code,
236                last_error_message: row.last_error_message,
237                created_at: row.created_at,
238                updated_at: row.updated_at,
239            })
240        })
241        .collect::<Result<Vec<_>>>()
242}
243
244pub async fn get_job_by_id(
245    pool: &DbPool,
246    organization_id: Option<Uuid>,
247    job_id: Uuid,
248) -> Result<Option<JobQueueRecord>> {
249    let row = sqlx::query!(
250        "SELECT
251            id,
252            job_type,
253            organization_id,
254            payload,
255            status::text AS \"status!\",
256            priority,
257            run_number,
258            attempt,
259            max_attempts,
260            timeout_seconds,
261            next_run_at,
262            lease_expires_at,
263            last_heartbeat_at,
264            worker_id,
265            started_at,
266            finished_at,
267            stage,
268            progress_done,
269            progress_total,
270            progress_pct::float8 AS progress_pct,
271            checkpoint,
272            output,
273            idempotency_key,
274            status_reason,
275            last_error_code,
276            last_error_message,
277            created_at,
278            updated_at
279         FROM job_queue
280         WHERE id = $1
281           AND ($2::uuid IS NULL OR organization_id = $2)
282         LIMIT 1",
283        job_id,
284        organization_id,
285    )
286    .fetch_optional(pool)
287    .await
288    .map_err(|error| Error::from_query_sqlx_with_context("get job by id", error))?;
289
290    row.map(|row| {
291        Ok(JobQueueRecord {
292            id: row.id,
293            job_type: parse_job_type_name(row.job_type)?,
294            organization_id: row.organization_id,
295            payload: row.payload,
296            status: parse_job_status(row.status)?,
297            priority: row.priority,
298            run_number: row.run_number,
299            attempt: row.attempt,
300            max_attempts: row.max_attempts,
301            timeout_seconds: row.timeout_seconds,
302            next_run_at: row.next_run_at,
303            lease_expires_at: row.lease_expires_at,
304            last_heartbeat_at: row.last_heartbeat_at,
305            worker_id: row.worker_id,
306            started_at: row.started_at,
307            finished_at: row.finished_at,
308            stage: parse_job_stage(row.stage)?,
309            progress_done: row.progress_done,
310            progress_total: row.progress_total,
311            progress_pct: row.progress_pct,
312            checkpoint: row.checkpoint,
313            output: row.output,
314            idempotency_key: row.idempotency_key,
315            status_reason: row.status_reason,
316            last_error_code: row.last_error_code,
317            last_error_message: row.last_error_message,
318            created_at: row.created_at,
319            updated_at: row.updated_at,
320        })
321    })
322    .transpose()
323}
324
325pub async fn get_job_payload_by_idempotency_key(
326    pool: &DbPool,
327    organization_id: Uuid,
328    job_type: JobType<'_>,
329    idempotency_key: &str,
330) -> Result<Option<(Uuid, serde_json::Value)>> {
331    let row = sqlx::query!(
332        "SELECT id, payload
333         FROM job_queue
334         WHERE organization_id = $1
335           AND job_type = $2
336           AND idempotency_key = $3
337         LIMIT 1",
338        organization_id,
339        job_type as _,
340        idempotency_key,
341    )
342    .fetch_optional(pool)
343    .await
344    .map_err(|error| {
345        Error::from_query_sqlx_with_context("get job payload by idempotency key", error)
346    })?;
347
348    Ok(row.map(|row| (row.id, row.payload)))
349}
350
351pub async fn get_latest_job_payload_for_run(
352    pool: &DbPool,
353    organization_id: Uuid,
354    job_type: JobType<'_>,
355    run_id: Uuid,
356) -> Result<Option<(Uuid, serde_json::Value)>> {
357    let run_id_text = run_id.to_string();
358    let row = sqlx::query!(
359        "SELECT id, payload
360         FROM job_queue
361         WHERE organization_id = $1
362           AND job_type = $2
363           AND payload->>'run_id' = $3
364         ORDER BY created_at DESC, id DESC
365         LIMIT 1",
366        organization_id,
367        job_type as _,
368        run_id_text,
369    )
370    .fetch_optional(pool)
371    .await
372    .map_err(|error| {
373        Error::from_query_sqlx_with_context("get latest job payload for run", error)
374    })?;
375
376    Ok(row.map(|row| (row.id, row.payload)))
377}
378
379/// Updates one UUID-array payload field on a direct, unclaimed pending job.
380///
381/// Returns a classified rejection when the row is already claimed or terminal,
382/// belongs to a workflow step, or has an idempotency request snapshot that this
383/// API cannot keep consistent.
384pub async fn update_job_payload_uuid_array_field(
385    pool: &DbPool,
386    organization_id: Uuid,
387    job_id: Uuid,
388    job_type: JobType<'_>,
389    payload_field: &str,
390    values: &[Uuid],
391) -> Result<JobPayloadUuidArrayFieldUpdate> {
392    let mut tx = pool.begin().await.map_err(|error| {
393        Error::from_query_sqlx_with_context(
394            "begin job payload uuid array update transaction",
395            error,
396        )
397    })?;
398
399    let previous_lock_timeout =
400        cap_job_payload_uuid_array_field_update_lock_timeout_tx(&mut tx).await?;
401
402    let row_result = sqlx::query_as::<_, JobPayloadUuidArrayFieldUpdateCandidate>(
403        "SELECT
404             status::text AS status,
405             worker_id,
406             lease_expires_at,
407             workflow_step_id,
408             idempotency_key,
409             enqueue_request
410           FROM job_queue
411           WHERE id = $1
412             AND organization_id = $2
413             AND job_type = $3
414           FOR UPDATE",
415    )
416    .bind(job_id)
417    .bind(organization_id)
418    .bind(job_type)
419    .fetch_optional(&mut *tx)
420    .await;
421
422    let row = match row_result {
423        Ok(row) => {
424            set_local_lock_timeout_tx(
425                &mut tx,
426                &previous_lock_timeout,
427                "restore job payload uuid array update lock timeout",
428            )
429            .await?;
430            row
431        }
432        Err(error) => {
433            return Err(Error::from_query_sqlx_with_context(
434                "classify job payload uuid array update",
435                error,
436            ));
437        }
438    };
439
440    let Some(row) = row else {
441        tx.commit().await.map_err(|error| {
442            Error::from_query_sqlx_with_context(
443                "commit job payload uuid array update transaction",
444                error,
445            )
446        })?;
447        return Ok(JobPayloadUuidArrayFieldUpdate::NotFound);
448    };
449
450    // Order matters: workflow-managed jobs can also carry request snapshots, so
451    // return the ownership rejection before the snapshot-consistency rejection.
452    let rejection = if row.workflow_step_id.is_some() {
453        Some(JobPayloadUuidArrayFieldUpdateRejection::WorkflowManaged)
454    } else if row.idempotency_key.is_some() || row.enqueue_request.is_some() {
455        Some(JobPayloadUuidArrayFieldUpdateRejection::IdempotentRequestSnapshot)
456    } else if row.status != JobStatus::Pending.as_db_value()
457        || row.worker_id.is_some()
458        || row.lease_expires_at.is_some()
459    {
460        Some(JobPayloadUuidArrayFieldUpdateRejection::NotPendingOrClaimed)
461    } else {
462        None
463    };
464
465    if let Some(reason) = rejection {
466        tx.commit().await.map_err(|error| {
467            Error::from_query_sqlx_with_context(
468                "commit job payload uuid array update transaction",
469                error,
470            )
471        })?;
472        return Ok(JobPayloadUuidArrayFieldUpdate::Rejected { reason });
473    }
474
475    sqlx::query!(
476        "UPDATE job_queue
477         SET
478             payload = jsonb_set(
479                 payload,
480                 ARRAY[$4::text],
481                 to_jsonb($5::uuid[]),
482                 true
483             ),
484             updated_at = now()
485         WHERE id = $1
486           AND organization_id = $2
487           AND job_type = $3",
488        job_id,
489        organization_id,
490        job_type as _,
491        payload_field,
492        values,
493    )
494    .execute(&mut *tx)
495    .await
496    .map_err(|error| {
497        Error::from_query_sqlx_with_context("update job payload uuid array field", error)
498    })?;
499
500    tx.commit().await.map_err(|error| {
501        Error::from_query_sqlx_with_context(
502            "commit job payload uuid array update transaction",
503            error,
504        )
505    })?;
506    Ok(JobPayloadUuidArrayFieldUpdate::Updated)
507}
508
509async fn cap_job_payload_uuid_array_field_update_lock_timeout_tx(
510    tx: &mut DbTx<'_>,
511) -> Result<String> {
512    sqlx::query_scalar::<_, String>(
513        "WITH previous AS MATERIALIZED (
514             SELECT
515                current_setting('lock_timeout') AS lock_timeout,
516                setting::bigint AS lock_timeout_ms
517             FROM pg_settings
518             WHERE name = 'lock_timeout'
519         )
520         SELECT previous.lock_timeout
521         FROM previous,
522              LATERAL (
523                SELECT set_config(
524                    'lock_timeout',
525                    CASE
526                        WHEN previous.lock_timeout_ms = 0 THEN $1
527                        WHEN previous.lock_timeout_ms <= $2 THEN previous.lock_timeout
528                        ELSE $1
529                    END,
530                    true
531                )
532              ) AS applied",
533    )
534    .bind(JOB_PAYLOAD_UUID_ARRAY_FIELD_UPDATE_LOCK_TIMEOUT)
535    .bind(JOB_PAYLOAD_UUID_ARRAY_FIELD_UPDATE_LOCK_TIMEOUT_MS)
536    .fetch_one(&mut **tx)
537    .await
538    .map_err(|error| {
539        Error::from_query_sqlx_with_context("set job payload uuid array update lock timeout", error)
540    })
541}
542
543async fn set_local_lock_timeout_tx(
544    tx: &mut DbTx<'_>,
545    lock_timeout: &str,
546    context: &'static str,
547) -> Result<()> {
548    sqlx::query_scalar::<_, String>("SELECT set_config('lock_timeout', $1, true)")
549        .bind(lock_timeout)
550        .fetch_one(&mut **tx)
551        .await
552        .map_err(|error| Error::from_query_sqlx_with_context(context, error))?;
553
554    Ok(())
555}
556
557pub async fn list_job_events(
558    pool: &DbPool,
559    organization_id: Option<Uuid>,
560    job_id: Uuid,
561    limit: i64,
562    after_id: Option<i64>,
563) -> Result<Vec<JobEventRecord>> {
564    validate_page_limit(limit)?;
565
566    let rows = sqlx::query!(
567        "SELECT
568            je.id,
569            je.job_id,
570            je.run_number,
571            je.attempt,
572            je.event_type::text AS \"event_type!\",
573            je.stage,
574            je.progress_done,
575            je.progress_total,
576            je.payload,
577            je.occurred_at
578         FROM job_events je
579         JOIN job_queue jq ON jq.id = je.job_id
580         WHERE je.job_id = $1
581           AND ($2::uuid IS NULL OR jq.organization_id = $2)
582           AND ($3::bigint IS NULL OR je.id > $3)
583         ORDER BY je.id ASC
584         LIMIT $4",
585        job_id,
586        organization_id,
587        after_id,
588        limit,
589    )
590    .fetch_all(pool)
591    .await
592    .map_err(|error| Error::from_query_sqlx_with_context("list job events", error))?;
593
594    rows.into_iter()
595        .map(|row| {
596            Ok(JobEventRecord {
597                id: row.id,
598                job_id: row.job_id,
599                run_number: row.run_number,
600                attempt: row.attempt,
601                event_type: parse_job_event_type(row.event_type)?,
602                stage: row.stage.map(parse_job_stage).transpose()?,
603                progress_done: row.progress_done,
604                progress_total: row.progress_total,
605                payload: row.payload,
606                occurred_at: row.occurred_at,
607            })
608        })
609        .collect::<Result<Vec<_>>>()
610}
611
612pub async fn get_job_metrics(
613    pool: &DbPool,
614    organization_id: Option<Uuid>,
615    job_type: Option<&str>,
616) -> Result<Vec<JobMetricsRecord>> {
617    let rows = sqlx::query!(
618        "SELECT
619            jd.job_type AS \"job_type!\",
620            COALESCE(SUM(jmr.pending_count), 0)::bigint AS \"pending_count!\",
621            COALESCE(SUM(jmr.leased_count), 0)::bigint AS \"leased_count!\",
622            COALESCE(SUM(jmr.stale_leases), 0)::bigint AS \"stale_leases!\",
623            COALESCE(SUM(jmr.succeeded_24h), 0)::bigint AS \"succeeded_24h!\",
624            COALESCE(SUM(jmr.retryable_24h), 0)::bigint AS \"retryable_24h!\",
625            COALESCE(SUM(jmr.terminal_24h), 0)::bigint AS \"terminal_24h!\",
626            COALESCE(SUM(jmr.panicked_24h), 0)::bigint AS \"panicked_24h!\",
627            COALESCE(SUM(jmr.timeout_24h), 0)::bigint AS \"timeout_24h!\",
628            COALESCE(SUM(jmr.dead_lettered_24h), 0)::bigint AS \"dead_lettered_24h!\",
629            AVG(jmr.p50_duration_ms_24h) AS p50_duration_ms_24h,
630            AVG(jmr.p95_duration_ms_24h) AS p95_duration_ms_24h
631         FROM job_definitions jd
632         LEFT JOIN job_metrics_rollup jmr
633           ON jmr.job_type = jd.job_type
634          AND ($1::uuid IS NULL OR jmr.organization_id = $1)
635         WHERE ($2::text IS NULL OR jd.job_type = $2)
636         GROUP BY jd.job_type
637         ORDER BY jd.job_type ASC",
638        organization_id,
639        job_type,
640    )
641    .fetch_all(pool)
642    .await
643    .map_err(|error| Error::from_query_sqlx_with_context("get job metrics", error))?;
644
645    rows.into_iter()
646        .map(|row| {
647            Ok(JobMetricsRecord {
648                job_type: parse_job_type_name(row.job_type)?,
649                pending_count: row.pending_count,
650                leased_count: row.leased_count,
651                stale_leases: row.stale_leases,
652                succeeded_24h: row.succeeded_24h,
653                retryable_24h: row.retryable_24h,
654                terminal_24h: row.terminal_24h,
655                panicked_24h: row.panicked_24h,
656                timeout_24h: row.timeout_24h,
657                dead_lettered_24h: row.dead_lettered_24h,
658                p50_duration_ms_24h: row.p50_duration_ms_24h,
659                p95_duration_ms_24h: row.p95_duration_ms_24h,
660            })
661        })
662        .collect::<Result<Vec<_>>>()
663}
664
665pub async fn cancel_job(
666    pool: &DbPool,
667    organization_id: Option<Uuid>,
668    job_id: Uuid,
669    reason: Option<&str>,
670) -> Result<JobQueueRecord> {
671    let mut tx = pool
672        .begin()
673        .await
674        .map_err(|error| Error::ConnectionError(error.to_string()))?;
675    let Some(record) = cancel_job_tx(&mut tx, organization_id, job_id, reason).await? else {
676        return Err(
677            rollback_and_classify_missing_job_mutation(tx, pool, organization_id, job_id).await?,
678        );
679    };
680
681    tx.commit()
682        .await
683        .map_err(|error| Error::ConnectionError(error.to_string()))?;
684
685    Ok(record)
686}
687
688pub(crate) async fn cancel_job_tx(
689    tx: &mut DbTx<'_>,
690    organization_id: Option<Uuid>,
691    job_id: Uuid,
692    reason: Option<&str>,
693) -> Result<Option<JobQueueRecord>> {
694    let row = sqlx::query_as!(
695        JobQueueRecordRow,
696        "UPDATE job_queue
697         SET status = 'CANCELED',
698             lease_expires_at = NULL,
699             last_heartbeat_at = NULL,
700             worker_id = NULL,
701             finished_at = now(),
702             output = NULL,
703             status_reason = COALESCE($3, 'CANCELED'),
704             updated_at = now()
705         WHERE id = $1
706           AND ($2::uuid IS NULL OR organization_id = $2)
707           AND status IN ('PENDING', 'LEASED')
708         RETURNING
709            id,
710            job_type,
711            organization_id,
712            payload,
713            status::text AS \"status!\",
714            priority,
715            run_number,
716            attempt,
717            max_attempts,
718            timeout_seconds,
719            next_run_at,
720            lease_expires_at,
721            last_heartbeat_at,
722            worker_id,
723            started_at,
724            finished_at,
725            stage,
726            progress_done,
727            progress_total,
728            progress_pct::float8 AS progress_pct,
729            checkpoint,
730            output,
731            idempotency_key,
732            status_reason,
733            last_error_code,
734            last_error_message,
735            created_at,
736            updated_at",
737        job_id,
738        organization_id,
739        reason,
740    )
741    .fetch_optional(&mut **tx)
742    .await
743    .map_err(|error| Error::from_query_sqlx_with_context("cancel job", error))?;
744
745    let Some(row) = row else {
746        return Ok(None);
747    };
748
749    let record = job_queue_record_from_row(row)?;
750
751    sqlx::query!(
752        "UPDATE job_attempts
753         SET finished_at = now()
754         WHERE job_id = $1
755           AND run_number = $2
756           AND attempt = $3
757           AND finished_at IS NULL",
758        record.id,
759        record.run_number,
760        record.attempt,
761    )
762    .execute(&mut **tx)
763    .await
764    .map_err(|error| Error::from_query_sqlx_with_context("close canceled attempt", error))?;
765
766    let event_attempt = (record.attempt > 0).then_some(record.attempt);
767    sqlx::query!(
768        "INSERT INTO job_events (
769            job_id,
770            run_number,
771            attempt,
772            event_type,
773            payload
774         )
775         VALUES (
776            $1,
777            $2,
778            $3,
779            'CANCELED',
780            jsonb_build_object('reason', $4::text)
781         )",
782        record.id,
783        record.run_number,
784        event_attempt,
785        record.status_reason.as_deref(),
786    )
787    .execute(&mut **tx)
788    .await
789    .map_err(|error| Error::from_query_sqlx_with_context("insert canceled event", error))?;
790
791    on_terminal(
792        tx,
793        record.id,
794        WorkflowStepStatus::Canceled,
795        record.status_reason.as_deref(),
796        None,
797        None,
798        None,
799    )
800    .await?;
801
802    Ok(Some(record))
803}
804
805fn ensure_workflow_requeue_rejection_rollback_succeeded(
806    rollback_result: std::result::Result<(), sqlx::Error>,
807) -> Result<()> {
808    rollback_result.map_err(|error| Error::ConnectionError(error.to_string()))
809}
810
811pub async fn requeue_job(
812    pool: &DbPool,
813    organization_id: Option<Uuid>,
814    job_id: Uuid,
815    reason: Option<&str>,
816) -> Result<JobQueueRecord> {
817    let mut tx = pool
818        .begin()
819        .await
820        .map_err(|error| Error::ConnectionError(error.to_string()))?;
821
822    let workflow_managed = workflow_managed_job_exists_tx(&mut tx, job_id, organization_id).await?;
823    if workflow_managed {
824        ensure_workflow_requeue_rejection_rollback_succeeded(tx.rollback().await)?;
825        return Err(workflow_requeue_not_supported_error());
826    }
827
828    let previous_run = sqlx::query!(
829        "SELECT run_number, attempt
830         FROM job_queue
831         WHERE id = $1
832           AND ($2::uuid IS NULL OR organization_id = $2)
833           AND status IN ('DEAD_LETTERED', 'CANCELED', 'SUCCEEDED')
834         FOR UPDATE",
835        job_id,
836        organization_id,
837    )
838    .fetch_optional(&mut *tx)
839    .await
840    .map_err(|error| Error::from_query_sqlx_with_context("requeue job prefetch attempt", error))?;
841
842    let Some(previous_run) = previous_run else {
843        return Err(
844            rollback_and_classify_missing_job_mutation(tx, pool, organization_id, job_id).await?,
845        );
846    };
847    let previous_run_number: i32 = previous_run.run_number;
848    let previous_attempt: i32 = previous_run.attempt;
849
850    let row = sqlx::query_as!(
851        JobQueueRecordRow,
852        "UPDATE job_queue
853         SET status = 'PENDING',
854             stage = 'queued',
855             progress_done = NULL,
856             progress_total = NULL,
857             checkpoint = NULL,
858             output = NULL,
859             run_number = run_number + 1,
860             attempt = 0,
861             lease_expires_at = NULL,
862             last_heartbeat_at = NULL,
863             worker_id = NULL,
864             next_run_at = now(),
865             started_at = NULL,
866             finished_at = NULL,
867             status_reason = COALESCE($3, 'REQUEUED'),
868             last_error_code = NULL,
869             last_error_message = NULL,
870             updated_at = now()
871         WHERE id = $1
872           AND ($2::uuid IS NULL OR organization_id = $2)
873           AND status IN ('DEAD_LETTERED', 'CANCELED', 'SUCCEEDED')
874         RETURNING
875            id,
876            job_type,
877            organization_id,
878            payload,
879            status::text AS \"status!\",
880            priority,
881            run_number,
882            attempt,
883            max_attempts,
884            timeout_seconds,
885            next_run_at,
886            lease_expires_at,
887            last_heartbeat_at,
888            worker_id,
889            started_at,
890            finished_at,
891            stage,
892            progress_done,
893            progress_total,
894            progress_pct::float8 AS progress_pct,
895            checkpoint,
896            output,
897            idempotency_key,
898            status_reason,
899            last_error_code,
900            last_error_message,
901            created_at,
902            updated_at",
903        job_id,
904        organization_id,
905        reason,
906    )
907    .fetch_optional(&mut *tx)
908    .await
909    .map_err(|error| Error::from_query_sqlx_with_context("requeue job", error))?;
910
911    let Some(row) = row else {
912        return Err(
913            rollback_and_classify_missing_job_mutation(tx, pool, organization_id, job_id).await?,
914        );
915    };
916
917    let record = job_queue_record_from_row(row)?;
918
919    let event_attempt = (previous_attempt > 0).then_some(previous_attempt);
920    sqlx::query!(
921        "INSERT INTO job_events (
922            job_id,
923            run_number,
924            attempt,
925            event_type,
926            payload
927         )
928         VALUES (
929            $1,
930            $2,
931            $3,
932            'REQUEUED',
933            jsonb_build_object('reason', $4::text)
934         )",
935        record.id,
936        previous_run_number,
937        event_attempt,
938        record.status_reason.as_deref(),
939    )
940    .execute(&mut *tx)
941    .await
942    .map_err(|error| Error::from_query_sqlx_with_context("insert requeued event", error))?;
943
944    tx.commit()
945        .await
946        .map_err(|error| Error::ConnectionError(error.to_string()))?;
947
948    Ok(record)
949}
950
951#[cfg(test)]
952mod tests {
953    use super::ensure_workflow_requeue_rejection_rollback_succeeded;
954    use crate::Error;
955
956    #[test]
957    fn workflow_requeue_rejection_allows_validation_error_after_successful_rollback() {
958        let result = ensure_workflow_requeue_rejection_rollback_succeeded(Ok(()));
959        assert!(result.is_ok());
960    }
961
962    #[test]
963    fn workflow_requeue_rejection_returns_internal_error_when_rollback_fails() {
964        let result =
965            ensure_workflow_requeue_rejection_rollback_succeeded(Err(sqlx::Error::PoolTimedOut));
966        match result {
967            Err(Error::ConnectionError(message)) => assert!(!message.is_empty()),
968            other => panic!("expected connection error, got {other:?}"),
969        }
970    }
971}