use super::lifecycle::{
self, Guard, LockedWorkflow, ParentState, Target, TransitionResult, TransitionStore,
};
use super::row::{
self, ARCHIVED_JOB_FIELDS, ENCRYPTION_COLUMNS, INSERT_JOB_COLUMNS, JOB_SELECT_FIELDS,
result_storage_str,
};
use super::{
DatabaseQueue, DeadJobSummary, JobOutcome, JobTransition, QueueStats, RecordedOutcome,
};
use crate::{
Result,
batch::{BatchProgress, BatchStatus, PartialFailureMode},
job::{Job, JobId, JobStatus},
priority::JobPriority,
rate_limit::ThrottleConfig,
workflow::{FailurePolicy, WorkflowProgress, WorkflowStatus},
};
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use sqlx::{FromRow, Postgres, Row};
use std::collections::HashMap;
#[derive(FromRow, Clone)]
pub(crate) struct JobRow {
pub id: uuid::Uuid,
pub queue_name: String,
pub payload: serde_json::Value,
pub status: String,
pub priority: i32,
pub attempts: i32,
pub max_attempts: i32,
pub timeout_seconds: Option<i32>,
pub created_at: DateTime<Utc>,
pub scheduled_at: DateTime<Utc>,
pub started_at: Option<DateTime<Utc>>,
pub completed_at: Option<DateTime<Utc>>,
pub failed_at: Option<DateTime<Utc>>,
pub timed_out_at: Option<DateTime<Utc>>,
pub error_message: Option<String>,
pub cron_schedule: Option<String>,
pub next_run_at: Option<DateTime<Utc>>,
pub recurring: bool,
pub timezone: Option<String>,
pub batch_id: Option<uuid::Uuid>,
pub result_data: Option<serde_json::Value>,
pub result_stored_at: Option<DateTime<Utc>>,
pub result_expires_at: Option<DateTime<Utc>>,
pub result_storage_type: Option<String>,
pub result_ttl_seconds: Option<i64>,
pub result_max_size_bytes: Option<i64>,
pub depends_on: Vec<uuid::Uuid>,
pub dependents: Vec<uuid::Uuid>,
pub dependency_status: Option<String>,
pub workflow_id: Option<uuid::Uuid>,
pub workflow_name: Option<String>,
pub trace_id: Option<String>,
pub correlation_id: Option<String>,
pub parent_span_id: Option<String>,
pub span_context: Option<String>,
pub retry_strategy: Option<serde_json::Value>,
pub is_encrypted: bool,
#[cfg(feature = "encryption")]
pub encryption_key_id: Option<String>,
#[cfg(feature = "encryption")]
pub encryption_algorithm: Option<String>,
#[cfg(feature = "encryption")]
pub encrypted_payload: Option<Vec<u8>>,
#[cfg(feature = "encryption")]
pub encryption_nonce: Option<Vec<u8>>,
#[cfg(feature = "encryption")]
pub encryption_tag: Option<Vec<u8>>,
#[cfg(feature = "encryption")]
pub encryption_metadata: Option<serde_json::Value>,
#[cfg(feature = "encryption")]
pub payload_hash: Option<String>,
pub pii_fields: Option<Vec<String>>,
#[cfg(feature = "encryption")]
pub retention_policy: Option<String>,
#[cfg(feature = "encryption")]
pub retention_delete_at: Option<DateTime<Utc>>,
#[cfg(feature = "encryption")]
pub encrypted_at: Option<DateTime<Utc>>,
}
async fn copy_encryption_columns(
conn: &mut sqlx::PgConnection,
from: &str,
to: &str,
ids: &[uuid::Uuid],
) -> Result<()> {
if ids.is_empty() {
return Ok(());
}
let assignments = ENCRYPTION_COLUMNS
.iter()
.map(|column| format!("{column} = src.{column}"))
.collect::<Vec<_>>()
.join(", ");
sqlx::query(&format!(
"UPDATE {to} dst SET {assignments} FROM {from} src \
WHERE dst.id = src.id AND src.id = ANY($1)"
))
.bind(ids)
.execute(&mut *conn)
.await?;
Ok(())
}
fn archived_job_from_row(row: &sqlx::postgres::PgRow) -> Result<Job> {
let payload = crate::archive::decode_archived_payload(
&row.try_get::<Vec<u8>, _>("payload")?,
row.try_get("payload_compressed")?,
)?;
Ok(Job {
id: row.try_get("id")?,
queue_name: row.try_get("queue_name")?,
payload,
status: JobStatus::Archived,
priority: row
.try_get::<String, _>("priority")?
.parse()
.unwrap_or(JobPriority::Normal),
attempts: row.try_get("attempts")?,
max_attempts: row.try_get("max_attempts")?,
created_at: row.try_get("created_at")?,
scheduled_at: row.try_get("scheduled_at")?,
started_at: row.try_get("started_at")?,
completed_at: row.try_get("completed_at")?,
failed_at: row.try_get("failed_at")?,
timed_out_at: row.try_get("timed_out_at")?,
timeout: row
.try_get::<Option<i32>, _>("timeout_seconds")?
.map(|s| super::db_seconds(s, "timeout_seconds"))
.transpose()?,
error_message: row.try_get("error_message")?,
cron_schedule: row.try_get("cron_schedule")?,
next_run_at: row.try_get("next_run_at")?,
recurring: row.try_get("recurring")?,
timezone: row.try_get("timezone")?,
batch_id: row.try_get("batch_id")?,
result_config: row
.try_get::<Option<serde_json::Value>, _>("result_config")?
.map(serde_json::from_value)
.transpose()?
.unwrap_or_default(),
result_data: row.try_get("result")?,
result_stored_at: None,
result_expires_at: row.try_get("result_ttl")?,
retry_strategy: row
.try_get::<Option<String>, _>("retry_strategy")?
.map(|s| serde_json::from_str(&s))
.transpose()?,
depends_on: row
.try_get::<Option<Vec<uuid::Uuid>>, _>("depends_on")?
.unwrap_or_default(),
dependents: Vec::new(),
dependency_status: crate::archive::parse_archived_dependency_status(
row.try_get::<Option<String>, _>("dependency_status")?
.as_deref(),
),
workflow_id: None,
workflow_name: None,
trace_id: row.try_get("trace_id")?,
correlation_id: row.try_get("correlation_id")?,
parent_span_id: row.try_get("parent_span_id")?,
span_context: row.try_get("span_context")?,
#[cfg(feature = "encryption")]
encryption_config: None,
pii_fields: row
.try_get::<Option<Vec<String>>, _>("pii_fields")?
.unwrap_or_default(),
#[cfg(feature = "encryption")]
retention_policy: None,
is_encrypted: row.try_get("is_encrypted")?,
#[cfg(feature = "encryption")]
encrypted_payload: None,
})
}
impl JobRow {
pub fn into_job(mut self) -> Result<Job> {
let pii_fields = self.pii_fields.take().unwrap_or_default();
#[cfg(feature = "encryption")]
let (encryption_config, retention_policy, encrypted_payload) =
self.encryption_row(&pii_fields).decode()?;
Ok(Job {
id: self.id,
queue_name: self.queue_name,
payload: self.payload,
status: row::job_status_from_row(&self.status)?,
priority: JobPriority::from_i32(self.priority).unwrap_or(JobPriority::Normal),
attempts: self.attempts,
max_attempts: self.max_attempts,
created_at: self.created_at,
scheduled_at: self.scheduled_at,
started_at: self.started_at,
completed_at: self.completed_at,
failed_at: self.failed_at,
timed_out_at: self.timed_out_at,
timeout: self
.timeout_seconds
.map(|s| super::db_seconds(s, "timeout_seconds"))
.transpose()?,
error_message: self.error_message,
cron_schedule: self.cron_schedule,
next_run_at: self.next_run_at,
recurring: self.recurring,
timezone: self.timezone,
batch_id: self.batch_id,
result_config: crate::job::ResultConfig {
storage: row::result_storage_from_db(self.result_storage_type.as_deref()),
ttl: self
.result_ttl_seconds
.map(|s| super::db_seconds(s, "result_ttl_seconds"))
.transpose()?,
max_size_bytes: self
.result_max_size_bytes
.map(|b| super::db_int(b, "result_max_size_bytes"))
.transpose()?,
},
result_data: self.result_data,
result_stored_at: self.result_stored_at,
result_expires_at: self.result_expires_at,
retry_strategy: super::retry_strategy_from_json(self.retry_strategy)?,
depends_on: self.depends_on,
dependents: self.dependents,
dependency_status: row::dependency_status_from_row(self.dependency_status.as_deref())?,
workflow_id: self.workflow_id,
workflow_name: self.workflow_name,
trace_id: self.trace_id,
correlation_id: self.correlation_id,
parent_span_id: self.parent_span_id,
span_context: self.span_context,
#[cfg(feature = "encryption")]
encryption_config,
pii_fields,
#[cfg(feature = "encryption")]
retention_policy,
is_encrypted: self.is_encrypted,
#[cfg(feature = "encryption")]
encrypted_payload,
})
}
#[cfg(feature = "encryption")]
fn encryption_row<'a>(&'a self, pii_fields: &'a [String]) -> row::EncryptionRow<'a> {
row::EncryptionRow {
is_encrypted: self.is_encrypted,
created_at: self.created_at,
encryption_key_id: self.encryption_key_id.as_deref(),
encryption_algorithm: self.encryption_algorithm.as_deref(),
encrypted_payload: self.encrypted_payload.as_deref(),
encryption_nonce: self.encryption_nonce.as_deref(),
encryption_tag: self.encryption_tag.as_deref(),
encryption_metadata: self.encryption_metadata.as_ref(),
payload_hash: self.payload_hash.as_deref(),
pii_fields,
retention_policy: self.retention_policy.as_deref(),
retention_delete_at: self.retention_delete_at,
encrypted_at: self.encrypted_at,
}
}
}
impl crate::queue::JobQueue<Postgres> {
async fn seal_for_insert(&self, jobs: &mut [Job]) -> Result<()> {
self.seal_jobs(jobs).await?;
for queue_name in self.encryption.guarded_queues(jobs) {
let holds_encrypted_jobs = sqlx::query(
"SELECT 1 FROM hammerwork_jobs WHERE queue_name = $1 AND is_encrypted = TRUE \
LIMIT 1",
)
.bind(&queue_name)
.fetch_optional(&self.pool)
.await?
.is_some();
if holds_encrypted_jobs {
return Err(
super::payload_encryption::PayloadEncryption::plaintext_refused(&queue_name),
);
}
}
Ok(())
}
}
#[async_trait]
impl DatabaseQueue for crate::queue::JobQueue<Postgres> {
type Database = Postgres;
async fn enqueue(&self, mut job: Job) -> Result<JobId> {
self.seal_for_insert(std::slice::from_mut(&mut job)).await?;
if job.depends_on.is_empty() {
let mut conn = self.pool.acquire().await?;
insert_jobs(&mut conn, std::slice::from_ref(&job)).await?;
return Ok(job.id);
}
super::retry_on_conflict(|| async {
let mut tx = self.pool.begin().await?;
let mut jobs = vec![job.clone()];
let result = insert_dependent_jobs(&mut tx, &mut jobs).await;
super::end_transaction(tx, result).await
})
.await?;
Ok(job.id)
}
async fn dequeue(&self, queue_name: &str) -> Result<Option<Job>> {
self.dequeue_leased(queue_name, None, super::DEFAULT_LEASE_DURATION)
.await
}
async fn dequeue_with_priority_weights(
&self,
queue_name: &str,
weights: &crate::priority::PriorityWeights,
) -> Result<Option<Job>> {
self.dequeue_leased(queue_name, Some(weights), super::DEFAULT_LEASE_DURATION)
.await
}
async fn dequeue_leased(
&self,
queue_name: &str,
weights: Option<&crate::priority::PriorityWeights>,
lease: std::time::Duration,
) -> Result<Option<Job>> {
let lease = super::lease_interval(lease);
let weights = match weights {
Some(weights) if !weights.is_strict() => weights,
_ => {
let row = sqlx::query_as::<_, JobRow>(&claim_sql(false))
.bind(queue_name)
.bind(lease)
.fetch_optional(&self.pool)
.await?;
return row.map(JobRow::into_job).transpose();
}
};
let levels: Vec<i32> =
sqlx::query_scalar(&super::runnable_priorities_sql(RUNNABLE_IN_QUEUE_SQL))
.bind(queue_name)
.fetch_all(&self.pool)
.await?;
let mut candidates: Vec<JobPriority> = levels
.into_iter()
.filter_map(|level| JobPriority::from_i32(level).ok())
.collect();
let mut attempt = 0;
while let Some(priority) = super::pick_weighted_priority(
&candidates,
weights,
super::weighted_seed(queue_name, attempt),
) {
let claimed = sqlx::query_as::<_, JobRow>(&claim_sql(true))
.bind(queue_name)
.bind(lease)
.bind(priority.as_i32())
.fetch_optional(&self.pool)
.await?;
if let Some(row) = claimed {
return row.into_job().map(Some);
}
candidates.retain(|candidate| *candidate != priority);
attempt += 1;
}
Ok(None)
}
async fn dequeue_batch_leased(
&self,
queue_name: &str,
weights: Option<&crate::priority::PriorityWeights>,
lease: std::time::Duration,
max_jobs: usize,
) -> Result<Vec<Job>> {
let max_jobs = max_jobs.min(super::MAX_DEQUEUE_BATCH_SIZE);
if max_jobs == 0 {
return Ok(Vec::new());
}
let lease = super::lease_interval(lease);
let weights = match weights {
Some(weights) if !weights.is_strict() => weights,
_ => {
let rows = sqlx::query_as::<_, JobRow>(&claim_batch_sql(false))
.bind(queue_name)
.bind(lease)
.bind(max_jobs as i64)
.fetch_all(&self.pool)
.await?;
return Ok(claimed_jobs(rows));
}
};
let levels: Vec<i32> =
sqlx::query_scalar(&super::runnable_priorities_sql(RUNNABLE_IN_QUEUE_SQL))
.bind(queue_name)
.fetch_all(&self.pool)
.await?;
let candidates: Vec<JobPriority> = levels
.into_iter()
.filter_map(|level| JobPriority::from_i32(level).ok())
.collect();
super::fill_weighted_batch(
queue_name,
candidates,
weights,
max_jobs,
|priority, room| async move {
let rows = sqlx::query_as::<_, JobRow>(&claim_batch_sql(true))
.bind(queue_name)
.bind(lease)
.bind(room as i64)
.bind(priority.as_i32())
.fetch_all(&self.pool)
.await?;
Ok(claimed_jobs(rows))
},
)
.await
}
async fn release_job_run(&self, run: &Job) -> Result<bool> {
let started = super::run_started_range(run);
let result = sqlx::query(
"UPDATE hammerwork_jobs SET status = $1, started_at = NULL, \
attempts = GREATEST(attempts - 1, 0), last_heartbeat_at = NULL, \
lease_expires_at = NULL WHERE id = $2 AND status = $3 AND attempts = $4 \
AND ($5::timestamptz IS NULL OR (started_at > $5 AND started_at < $6))",
)
.bind(JobStatus::Pending)
.bind(run.id)
.bind(JobStatus::Running)
.bind(run.attempts)
.bind(started.map(|(from, _)| from))
.bind(started.map(|(_, to)| to))
.execute(&self.pool)
.await?;
Ok(result.rows_affected() > 0)
}
async fn complete_job(&self, job_id: JobId) -> Result<()> {
lifecycle::manual_transition(self, job_id, JobTransition::Complete, Target::Completed).await
}
async fn fail_job(&self, job_id: JobId, error_message: &str) -> Result<()> {
lifecycle::manual_transition(
self,
job_id,
JobTransition::Fail,
Target::Failed(error_message.to_string()),
)
.await
}
async fn retry_job(&self, job_id: JobId, retry_at: DateTime<Utc>) -> Result<()> {
lifecycle::manual_transition(
self,
job_id,
JobTransition::Retry,
Target::Retry {
retry_at,
error: None,
timed_out: false,
},
)
.await
}
async fn finish_job_run(
&self,
run: &Job,
outcome: JobOutcome,
) -> Result<Option<RecordedOutcome>> {
lifecycle::finish_run(self, run, outcome).await
}
async fn complete_job_run_with_children(
&self,
run: &Job,
mut children: Vec<Job>,
) -> Result<Option<RecordedOutcome>> {
self.seal_jobs(&mut children).await?;
super::retry_on_conflict(|| async {
let mut tx = self.pool.begin().await?;
let result = async {
let transition = lifecycle::transition_in_tx::<Self>(
&mut tx,
run.id,
Guard::Run(run),
&Target::Completed,
true,
)
.await?;
let Some(mut recorded) = transition.into_run() else {
return Ok(None);
};
let mut jobs = children.clone();
insert_dependent_jobs(&mut tx, &mut jobs).await?;
recorded.spawned = jobs.iter().map(|job| job.id).collect();
Ok(Some(recorded))
}
.await;
super::end_transaction(tx, result).await
})
.await
}
async fn get_job(&self, job_id: JobId) -> Result<Option<Job>> {
let row = sqlx::query_as::<_, JobRow>(&format!(
"SELECT {} FROM hammerwork_jobs WHERE id = $1",
JOB_SELECT_FIELDS
))
.bind(job_id)
.fetch_optional(&self.pool)
.await?;
if let Some(job_row) = row {
return Ok(Some(job_row.into_job()?));
}
let archived = sqlx::query(&format!(
"SELECT {ARCHIVED_JOB_FIELDS} FROM hammerwork_jobs_archive WHERE id = $1"
))
.bind(job_id)
.fetch_optional(&self.pool)
.await?;
archived.as_ref().map(archived_job_from_row).transpose()
}
async fn delete_job(&self, job_id: JobId) -> Result<()> {
sqlx::query("DELETE FROM hammerwork_jobs WHERE id = $1")
.bind(job_id)
.execute(&self.pool)
.await?;
self.memory_results.remove(job_id);
Ok(())
}
async fn enqueue_batch(&self, batch: crate::batch::JobBatch) -> Result<crate::batch::BatchId> {
batch.validate()?;
let mut jobs: Vec<Job> = batch
.jobs
.iter()
.cloned()
.map(|mut job| {
job.batch_id = Some(batch.id);
job
})
.collect();
self.seal_for_insert(&mut jobs).await?;
let mut tx = self.pool.begin().await?;
sqlx::query(
r#"
INSERT INTO hammerwork_batches
(id, batch_name, total_jobs, completed_jobs, failed_jobs, pending_jobs, status, failure_mode, created_at, metadata)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
"#
)
.bind(batch.id)
.bind(&batch.name)
.bind(batch.jobs.len() as i32)
.bind(0i32) .bind(0i32) .bind(batch.jobs.len() as i32) .bind(BatchStatus::Pending)
.bind(serde_json::to_string(&batch.failure_mode)?)
.bind(batch.created_at)
.bind(serde_json::to_value(&batch.metadata)?)
.execute(&mut *tx)
.await?;
insert_dependent_jobs(&mut tx, &mut jobs).await?;
tx.commit().await?;
Ok(batch.id)
}
async fn get_batch_status(
&self,
batch_id: crate::batch::BatchId,
) -> Result<crate::batch::BatchResult> {
use crate::batch::BatchResult;
use std::collections::HashMap;
let batch_row = sqlx::query(
"SELECT batch_name, total_jobs, completed_jobs, failed_jobs, pending_jobs, status, failure_mode, created_at, completed_at, error_summary, metadata FROM hammerwork_batches WHERE id = $1"
)
.bind(batch_id)
.fetch_optional(&self.pool)
.await?;
let batch_row = batch_row.ok_or_else(|| crate::HammerworkError::JobNotFound {
id: batch_id.to_string(),
})?;
let total_jobs: u32 =
super::db_int(batch_row.try_get::<i32, _>("total_jobs")?, "total_jobs")?;
let failure_mode =
crate::batch::parse_failure_mode(&batch_row.try_get::<String, _>("failure_mode")?);
let status_counts: Vec<(String, i64)> = sqlx::query_as(
"SELECT status, COUNT(*) FROM (
SELECT status FROM hammerwork_jobs WHERE batch_id = $1
UNION ALL
SELECT status FROM hammerwork_jobs_archive WHERE batch_id = $1
) batch_jobs GROUP BY status",
)
.bind(batch_id)
.fetch_all(&self.pool)
.await?;
let progress = crate::batch::BatchProgress::from_status_counts(
status_counts,
total_jobs,
&failure_mode,
);
let created_at: DateTime<Utc> = batch_row.try_get("created_at")?;
let completed_at: Option<DateTime<Utc>> = batch_row.try_get("completed_at")?;
let error_summary: Option<String> = batch_row.try_get("error_summary")?;
let job_errors: Vec<(uuid::Uuid, String)> = sqlx::query_as(
"SELECT id, error_message FROM hammerwork_jobs WHERE batch_id = $1 AND error_message IS NOT NULL"
)
.bind(batch_id)
.fetch_all(&self.pool)
.await?;
let job_errors_map: HashMap<uuid::Uuid, String> = job_errors.into_iter().collect();
Ok(BatchResult {
batch_id,
total_jobs,
completed_jobs: progress.completed,
failed_jobs: progress.failed,
pending_jobs: progress.pending,
status: progress.status,
created_at,
completed_at,
error_summary,
job_errors: job_errors_map,
})
}
async fn get_batch_jobs(&self, batch_id: crate::batch::BatchId) -> Result<Vec<Job>> {
let rows = sqlx::query_as::<_, JobRow>(&format!(
"SELECT {} FROM hammerwork_jobs WHERE batch_id = $1 ORDER BY created_at ASC",
JOB_SELECT_FIELDS
))
.bind(batch_id)
.fetch_all(&self.pool)
.await?;
rows.into_iter().map(|row| row.into_job()).collect()
}
async fn delete_batch(&self, batch_id: crate::batch::BatchId) -> Result<()> {
let mut tx = self.pool.begin().await?;
sqlx::query("DELETE FROM hammerwork_jobs WHERE batch_id = $1")
.bind(batch_id)
.execute(&mut *tx)
.await?;
sqlx::query("DELETE FROM hammerwork_batches WHERE id = $1")
.bind(batch_id)
.execute(&mut *tx)
.await?;
tx.commit().await?;
Ok(())
}
async fn mark_job_dead(&self, job_id: JobId, error_message: &str) -> Result<()> {
lifecycle::manual_transition(
self,
job_id,
JobTransition::MarkDead,
Target::Dead(error_message.to_string()),
)
.await
}
async fn mark_job_timed_out(&self, job_id: JobId, error_message: &str) -> Result<()> {
lifecycle::manual_transition(
self,
job_id,
JobTransition::MarkTimedOut,
Target::TimedOut(error_message.to_string()),
)
.await
}
async fn get_dead_jobs(&self, limit: Option<u32>, offset: Option<u32>) -> Result<Vec<Job>> {
let limit = limit.unwrap_or(100) as i64;
let offset = offset.unwrap_or(0) as i64;
let rows = sqlx::query_as::<_, JobRow>(&format!(
"SELECT {JOB_SELECT_FIELDS} FROM hammerwork_jobs WHERE status = $1 ORDER BY failed_at DESC LIMIT $2 OFFSET $3"
))
.bind(JobStatus::Dead)
.bind(limit)
.bind(offset)
.fetch_all(&self.pool)
.await?;
Ok(super::decodable_rows(
"dead jobs",
rows.into_iter()
.map(|row| (row.id.to_string(), row.into_job())),
))
}
async fn get_dead_jobs_by_queue(
&self,
queue_name: &str,
limit: Option<u32>,
offset: Option<u32>,
) -> Result<Vec<Job>> {
let limit = limit.unwrap_or(100) as i64;
let offset = offset.unwrap_or(0) as i64;
let rows = sqlx::query_as::<_, JobRow>(&format!(
"SELECT {JOB_SELECT_FIELDS} FROM hammerwork_jobs WHERE status = $1 AND queue_name = $2 ORDER BY failed_at DESC LIMIT $3 OFFSET $4"
))
.bind(JobStatus::Dead)
.bind(queue_name)
.bind(limit)
.bind(offset)
.fetch_all(&self.pool)
.await?;
Ok(super::decodable_rows(
"dead jobs",
rows.into_iter()
.map(|row| (row.id.to_string(), row.into_job())),
))
}
async fn retry_dead_job(&self, job_id: JobId) -> Result<()> {
let transition = JobTransition::RetryDead;
let result = sqlx::query(&format!(
"UPDATE hammerwork_jobs SET status = $1, attempts = 0, scheduled_at = now(), \
started_at = NULL, failed_at = NULL, timed_out_at = NULL, \
last_heartbeat_at = NULL, lease_expires_at = NULL \
WHERE id = $2 AND status IN ({})",
transition.sql_status_list()
))
.bind(JobStatus::Pending)
.bind(job_id)
.execute(&self.pool)
.await?;
if result.rows_affected() == 0 {
return Err(self.rejected_transition(job_id, transition).await?);
}
Ok(())
}
async fn purge_dead_jobs(&self, older_than: DateTime<Utc>) -> Result<u64> {
let result =
sqlx::query("DELETE FROM hammerwork_jobs WHERE status = $1 AND failed_at < $2")
.bind(JobStatus::Dead)
.bind(older_than)
.execute(&self.pool)
.await?;
Ok(result.rows_affected())
}
async fn get_dead_job_summary(&self) -> Result<DeadJobSummary> {
let total_dead_jobs: (i64,) =
sqlx::query_as("SELECT COUNT(*) FROM hammerwork_jobs WHERE status = $1")
.bind(JobStatus::Dead)
.fetch_one(&self.pool)
.await?;
let dead_jobs_by_queue_rows: Vec<(String, i64)> = sqlx::query_as(
"SELECT queue_name, COUNT(*) FROM hammerwork_jobs WHERE status = $1 GROUP BY queue_name"
)
.bind(JobStatus::Dead)
.fetch_all(&self.pool)
.await?;
let timestamps: Vec<(Option<DateTime<Utc>>, Option<DateTime<Utc>>)> = sqlx::query_as(
"SELECT MIN(failed_at), MAX(failed_at) FROM hammerwork_jobs WHERE status = $1 AND failed_at IS NOT NULL"
)
.bind(JobStatus::Dead)
.fetch_all(&self.pool)
.await?;
let error_patterns_rows: Vec<(Option<String>, i64)> = sqlx::query_as(
"SELECT error_message, COUNT(*) FROM hammerwork_jobs WHERE status = $1 AND error_message IS NOT NULL GROUP BY error_message ORDER BY COUNT(*) DESC LIMIT 20"
)
.bind(JobStatus::Dead)
.fetch_all(&self.pool)
.await?;
let dead_jobs_by_queue = super::count_map(dead_jobs_by_queue_rows, "dead job count")?;
let error_patterns = super::count_map(
error_patterns_rows
.into_iter()
.filter_map(|(error, count)| error.map(|e| (e, count)))
.collect(),
"error count",
)?;
let (oldest_dead_job, newest_dead_job) = timestamps
.first()
.map(|(oldest, newest)| (*oldest, *newest))
.unwrap_or((None, None));
Ok(DeadJobSummary {
total_dead_jobs: super::db_int(total_dead_jobs.0, "dead job count")?,
dead_jobs_by_queue,
oldest_dead_job,
newest_dead_job,
error_patterns,
})
}
async fn get_queue_stats(&self, queue_name: &str) -> Result<QueueStats> {
use crate::stats::JobStatistics;
use std::time::Duration;
let status_counts: Vec<(String, i64)> = sqlx::query_as(
"SELECT status, COUNT(*) FROM hammerwork_jobs WHERE queue_name = $1 GROUP BY status",
)
.bind(queue_name)
.fetch_all(&self.pool)
.await?;
let counts = super::count_map(status_counts, "job count")?;
let pending_count = counts.get("Pending").copied().unwrap_or(0);
let running_count = counts.get("Running").copied().unwrap_or(0);
let dead_count = counts.get("Dead").copied().unwrap_or(0);
let timed_out_count = counts.get("TimedOut").copied().unwrap_or(0);
let completed_count = counts.get("Completed").copied().unwrap_or(0);
let statistics = JobStatistics {
total_processed: completed_count + dead_count,
completed: completed_count,
failed: counts.get("Failed").copied().unwrap_or(0),
dead: dead_count,
timed_out: timed_out_count,
running: running_count,
time_window: Duration::from_secs(3600), calculated_at: Utc::now(),
..Default::default()
};
Ok(QueueStats {
queue_name: queue_name.to_string(),
pending_count,
running_count,
dead_count,
timed_out_count,
completed_count,
statistics,
})
}
async fn get_all_queue_stats(&self) -> Result<Vec<QueueStats>> {
let queue_names: Vec<(String,)> =
sqlx::query_as("SELECT DISTINCT queue_name FROM hammerwork_jobs")
.fetch_all(&self.pool)
.await?;
let mut results = Vec::new();
for (queue_name,) in queue_names {
let stats = self.get_queue_stats(&queue_name).await?;
results.push(stats);
}
Ok(results)
}
async fn get_job_counts_by_status(
&self,
queue_name: &str,
) -> Result<std::collections::HashMap<String, u64>> {
let status_counts: Vec<(String, i64)> = sqlx::query_as(
"SELECT status, COUNT(*) FROM hammerwork_jobs WHERE queue_name = $1 GROUP BY status",
)
.bind(queue_name)
.fetch_all(&self.pool)
.await?;
super::count_map(status_counts, "job count")
}
async fn get_priority_stats(&self, queue_name: &str) -> Result<crate::priority::PriorityStats> {
let priority_counts: Vec<(i32, i64)> = sqlx::query_as(
"SELECT priority, COUNT(*) FROM hammerwork_jobs WHERE queue_name = $1 GROUP BY priority",
)
.bind(queue_name)
.fetch_all(&self.pool)
.await?;
let mut stats = crate::priority::PriorityStats::new();
for (priority_num, count) in priority_counts {
let priority = match priority_num {
0 => crate::priority::JobPriority::Background,
1 => crate::priority::JobPriority::Low,
2 => crate::priority::JobPriority::Normal,
3 => crate::priority::JobPriority::High,
4 => crate::priority::JobPriority::Critical,
_ => crate::priority::JobPriority::Normal, };
*stats.job_counts.entry(priority).or_insert(0) = super::db_int(count, "job count")?;
}
let processing_times: Vec<(i32, f64)> = sqlx::query_as(
r#"SELECT priority,
CAST(EXTRACT(EPOCH FROM (completed_at - started_at)) * 1000
AS DOUBLE PRECISION) AS processing_ms
FROM hammerwork_jobs
WHERE queue_name = $1 AND started_at IS NOT NULL AND completed_at IS NOT NULL"#,
)
.bind(queue_name)
.fetch_all(&self.pool)
.await?;
let mut priority_times: std::collections::HashMap<crate::priority::JobPriority, Vec<f64>> =
std::collections::HashMap::new();
for (priority_num, processing_ms) in processing_times {
let priority = match priority_num {
0 => crate::priority::JobPriority::Background,
1 => crate::priority::JobPriority::Low,
2 => crate::priority::JobPriority::Normal,
3 => crate::priority::JobPriority::High,
4 => crate::priority::JobPriority::Critical,
_ => crate::priority::JobPriority::Normal,
};
priority_times
.entry(priority)
.or_default()
.push(processing_ms);
}
for (priority, times) in priority_times {
if !times.is_empty() {
let avg = times.iter().sum::<f64>() / times.len() as f64;
stats.avg_processing_times.insert(priority, avg);
}
}
stats.calculate_distribution();
Ok(stats)
}
async fn get_processing_times(
&self,
queue_name: &str,
since: DateTime<Utc>,
) -> Result<Vec<i64>> {
let times: Vec<(Option<i64>,)> = sqlx::query_as(
r#"
SELECT CAST(ROUND(EXTRACT(EPOCH FROM (completed_at - started_at)) * 1000)
AS BIGINT) AS processing_time_ms
FROM hammerwork_jobs
WHERE queue_name = $1
AND started_at IS NOT NULL
AND completed_at IS NOT NULL
AND completed_at >= $2
ORDER BY completed_at DESC
LIMIT 1000
"#,
)
.bind(queue_name)
.bind(since)
.fetch_all(&self.pool)
.await?;
Ok(times.into_iter().filter_map(|(time,)| time).collect())
}
async fn get_error_frequencies(
&self,
queue_name: Option<&str>,
since: DateTime<Utc>,
) -> Result<std::collections::HashMap<String, u64>> {
let query = if queue_name.is_some() {
"SELECT error_message, COUNT(*) FROM hammerwork_jobs WHERE queue_name = $1 AND error_message IS NOT NULL AND failed_at >= $2 GROUP BY error_message ORDER BY COUNT(*) DESC LIMIT 50"
} else {
"SELECT error_message, COUNT(*) FROM hammerwork_jobs WHERE error_message IS NOT NULL AND failed_at >= $1 GROUP BY error_message ORDER BY COUNT(*) DESC LIMIT 50"
};
let error_frequencies: Vec<(String, i64)> = if let Some(queue) = queue_name {
sqlx::query_as(query)
.bind(queue)
.bind(since)
.fetch_all(&self.pool)
.await?
} else {
sqlx::query_as(query)
.bind(since)
.fetch_all(&self.pool)
.await?
};
super::count_map(error_frequencies, "error count")
}
async fn get_jobs_completed_in_range(
&self,
queue_name: Option<&str>,
start_time: DateTime<Utc>,
end_time: DateTime<Utc>,
limit: Option<u32>,
) -> Result<Vec<Job>> {
let limit_clause = if let Some(limit) = limit {
format!("LIMIT {}", limit)
} else {
"".to_string()
};
let query = if queue_name.is_some() {
format!(
r#"
SELECT {}
FROM hammerwork_jobs
WHERE queue_name = $1
AND status = 'Completed'
AND completed_at >= $2
AND completed_at < $3
ORDER BY completed_at DESC
{}
"#,
JOB_SELECT_FIELDS, limit_clause
)
} else {
format!(
r#"
SELECT {}
FROM hammerwork_jobs
WHERE status = 'Completed'
AND completed_at >= $1
AND completed_at < $2
ORDER BY completed_at DESC
{}
"#,
JOB_SELECT_FIELDS, limit_clause
)
};
let rows: Vec<JobRow> = if let Some(queue) = queue_name {
sqlx::query_as(&query)
.bind(queue)
.bind(start_time)
.bind(end_time)
.fetch_all(&self.pool)
.await?
} else {
sqlx::query_as(&query)
.bind(start_time)
.bind(end_time)
.fetch_all(&self.pool)
.await?
};
rows.into_iter().map(|row| row.into_job()).collect()
}
async fn enqueue_cron_job(&self, job: Job) -> Result<JobId> {
self.enqueue(super::prepare_cron_job(job)?).await
}
async fn get_due_cron_jobs(&self, queue_name: Option<&str>) -> Result<Vec<Job>> {
let query = if queue_name.is_some() {
format!(
r#"
SELECT {}
FROM hammerwork_jobs
WHERE recurring = TRUE
AND queue_name = $1
AND (next_run_at IS NULL OR next_run_at <= now())
AND status = $2
ORDER BY next_run_at ASC
"#,
JOB_SELECT_FIELDS
)
} else {
format!(
r#"
SELECT {}
FROM hammerwork_jobs
WHERE recurring = TRUE
AND (next_run_at IS NULL OR next_run_at <= now())
AND status = $1
ORDER BY next_run_at ASC
"#,
JOB_SELECT_FIELDS
)
};
let rows = if let Some(queue) = queue_name {
sqlx::query_as::<_, JobRow>(&query)
.bind(queue)
.bind(JobStatus::Pending)
.fetch_all(&self.pool)
.await?
} else {
sqlx::query_as::<_, JobRow>(&query)
.bind(JobStatus::Pending)
.fetch_all(&self.pool)
.await?
};
rows.into_iter().map(|row| row.into_job()).collect()
}
async fn reschedule_cron_job(&self, job_id: JobId, next_run_at: DateTime<Utc>) -> Result<()> {
let transition = JobTransition::RescheduleCron;
let result = sqlx::query(&format!(
r#"
UPDATE hammerwork_jobs
SET status = $1,
scheduled_at = $2,
next_run_at = $2,
attempts = 0,
started_at = NULL,
completed_at = NULL,
failed_at = NULL,
timed_out_at = NULL,
error_message = NULL,
last_heartbeat_at = NULL,
lease_expires_at = NULL
WHERE id = $3 AND recurring = TRUE AND status IN ({})
"#,
transition.sql_status_list()
))
.bind(JobStatus::Pending)
.bind(next_run_at)
.bind(job_id)
.execute(&self.pool)
.await?;
if result.rows_affected() == 0 {
return Err(self.rejected_transition(job_id, transition).await?);
}
Ok(())
}
async fn get_recurring_jobs(&self, queue_name: &str) -> Result<Vec<Job>> {
let rows = sqlx::query_as::<_, JobRow>(
&format!("SELECT {} FROM hammerwork_jobs WHERE queue_name = $1 AND recurring = TRUE ORDER BY next_run_at ASC", JOB_SELECT_FIELDS)
)
.bind(queue_name)
.fetch_all(&self.pool)
.await?;
Ok(super::decodable_rows(
"recurring jobs",
rows.into_iter()
.map(|row| (row.id.to_string(), row.into_job())),
))
}
async fn disable_recurring_job(&self, job_id: JobId) -> Result<()> {
let result = sqlx::query(
"UPDATE hammerwork_jobs SET recurring = FALSE, next_run_at = NULL \
WHERE id = $1 AND cron_schedule IS NOT NULL",
)
.bind(job_id)
.execute(&self.pool)
.await?;
if result.rows_affected() == 0 {
return super::disable_recurring_job_not_applied(self.get_job(job_id).await?, job_id);
}
Ok(())
}
async fn enable_recurring_job(&self, job_id: JobId) -> Result<()> {
super::retry_on_conflict(|| async {
let mut tx = self.pool.begin().await?;
let result = async {
let row = sqlx::query_as::<_, JobRow>(&format!(
"SELECT {JOB_SELECT_FIELDS} FROM hammerwork_jobs WHERE id = $1 FOR UPDATE"
))
.bind(job_id)
.fetch_optional(&mut *tx)
.await?;
let Some(row) = row else {
return Err(crate::HammerworkError::JobNotFound {
id: job_id.to_string(),
});
};
let job = row.into_job()?;
let now = db_now(&mut tx).await?;
match super::enabled_recurring_next_run(&job, now)? {
None => {
sqlx::query("UPDATE hammerwork_jobs SET recurring = TRUE WHERE id = $1")
.bind(job_id)
.execute(&mut *tx)
.await?;
}
Some(next_run_at) => {
sqlx::query(
"UPDATE hammerwork_jobs SET recurring = TRUE, status = 'Pending', \
scheduled_at = $2, next_run_at = $2, attempts = 0, \
started_at = NULL, completed_at = NULL, failed_at = NULL, \
timed_out_at = NULL, error_message = NULL, \
last_heartbeat_at = NULL, lease_expires_at = NULL WHERE id = $1",
)
.bind(job_id)
.bind(next_run_at)
.execute(&mut *tx)
.await?;
}
}
Ok(())
}
.await;
super::end_transaction(tx, result).await
})
.await
}
async fn set_throttle_config(&self, queue_name: &str, config: ThrottleConfig) -> Result<()> {
self.set_throttle(queue_name, config).await
}
async fn get_throttle_config(&self, queue_name: &str) -> Result<Option<ThrottleConfig>> {
Ok(self.get_throttle(queue_name).await)
}
async fn remove_throttle_config(&self, queue_name: &str) -> Result<()> {
self.remove_throttle(queue_name).await
}
async fn get_all_throttle_configs(&self) -> Result<HashMap<String, ThrottleConfig>> {
Ok(self.get_all_throttles().await)
}
async fn get_queue_depth(&self, queue_name: &str) -> Result<u64> {
let count: i64 = sqlx::query_scalar(
"SELECT COUNT(*) FROM hammerwork_jobs WHERE queue_name = $1 AND status = 'Pending'",
)
.bind(queue_name)
.fetch_one(&self.pool)
.await?;
super::db_int(count, "queue depth")
}
async fn store_job_result(
&self,
job_id: JobId,
result_data: serde_json::Value,
expires_at: Option<DateTime<Utc>>,
) -> Result<()> {
sqlx::query(
"UPDATE hammerwork_jobs SET result_data = $1, result_stored_at = now(), result_expires_at = $2 WHERE id = $3"
)
.bind(result_data)
.bind(expires_at)
.bind(job_id)
.execute(&self.pool)
.await?;
Ok(())
}
async fn get_job_result(&self, job_id: JobId) -> Result<Option<serde_json::Value>> {
if let Some(result) = self.memory_results.get(job_id) {
return Ok(Some(result));
}
let row = sqlx::query(
"SELECT result_data FROM hammerwork_jobs WHERE id = $1 AND result_data IS NOT NULL AND (result_expires_at IS NULL OR result_expires_at > now())"
)
.bind(job_id)
.fetch_optional(&self.pool)
.await?;
match row {
Some(row) => Ok(row.try_get("result_data")?),
None => Ok(None),
}
}
async fn delete_job_result(&self, job_id: JobId) -> Result<()> {
self.memory_results.remove(job_id);
sqlx::query(
"UPDATE hammerwork_jobs SET result_data = NULL, result_stored_at = NULL, result_expires_at = NULL WHERE id = $1"
)
.bind(job_id)
.execute(&self.pool)
.await?;
Ok(())
}
async fn cleanup_expired_results(&self) -> Result<u64> {
let result = sqlx::query(
"UPDATE hammerwork_jobs SET result_data = NULL, result_stored_at = NULL, result_expires_at = NULL WHERE result_expires_at IS NOT NULL AND result_expires_at <= now()"
)
.execute(&self.pool)
.await?;
Ok(result.rows_affected() + self.memory_results.cleanup_expired())
}
async fn enqueue_workflow(
&self,
workflow: crate::workflow::JobGroup,
) -> Result<crate::workflow::WorkflowId> {
workflow.validate()?;
let mut jobs = workflow.jobs.clone();
self.seal_for_insert(&mut jobs).await?;
let mut tx = self.pool.begin().await?;
sqlx::query(
r#"
INSERT INTO hammerwork_workflows
(id, name, status, created_at, completed_at, failed_at, total_jobs, completed_jobs, failed_jobs, failure_policy, metadata)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
"#
)
.bind(workflow.id)
.bind(&workflow.name)
.bind(workflow.status.as_str())
.bind(workflow.created_at)
.bind(workflow.completed_at)
.bind(workflow.failed_at)
.bind(workflow.total_jobs as i32)
.bind(workflow.completed_jobs as i32)
.bind(workflow.failed_jobs as i32)
.bind(workflow.failure_policy.as_str())
.bind(&workflow.metadata)
.execute(&mut *tx)
.await?;
insert_dependent_jobs(&mut tx, &mut jobs).await?;
tx.commit().await?;
Ok(workflow.id)
}
async fn get_workflow_status(
&self,
workflow_id: crate::workflow::WorkflowId,
) -> Result<Option<crate::workflow::JobGroup>> {
use crate::workflow::{FailurePolicy, JobGroup, WorkflowStatus};
let workflow_row = sqlx::query(
r#"
SELECT id, name, status, created_at, completed_at, failed_at, total_jobs, completed_jobs, failed_jobs, failure_policy, metadata
FROM hammerwork_workflows
WHERE id = $1
"#
)
.bind(workflow_id)
.fetch_optional(&self.pool)
.await?;
if let Some(row) = workflow_row {
let jobs = self.get_workflow_jobs(workflow_id).await?;
let mut dependencies = std::collections::HashMap::new();
for job in &jobs {
if !job.depends_on.is_empty() {
dependencies.insert(job.id, job.depends_on.clone());
}
}
let workflow = JobGroup {
id: row.try_get("id")?,
name: row.try_get("name")?,
status: WorkflowStatus::parse_from_db(row.try_get("status")?)?,
created_at: row.try_get("created_at")?,
completed_at: row.try_get("completed_at")?,
failed_at: row.try_get("failed_at")?,
failure_policy: FailurePolicy::parse_from_db(row.try_get("failure_policy")?)?,
jobs,
dependencies,
total_jobs: super::db_int(row.try_get::<i32, _>("total_jobs")?, "total_jobs")?,
completed_jobs: super::db_int(
row.try_get::<i32, _>("completed_jobs")?,
"completed_jobs",
)?,
failed_jobs: super::db_int(row.try_get::<i32, _>("failed_jobs")?, "failed_jobs")?,
metadata: row.try_get("metadata")?,
};
Ok(Some(workflow))
} else {
Ok(None)
}
}
async fn resolve_job_dependencies(&self, completed_job_id: JobId) -> Result<Vec<JobId>> {
super::retry_on_conflict(|| async {
let mut tx = self.pool.begin().await?;
let result = Self::resolve_dependents(&mut tx, completed_job_id).await;
super::end_transaction(tx, result).await
})
.await
}
async fn get_ready_jobs(&self, queue_name: &str, limit: u32) -> Result<Vec<Job>> {
let rows = sqlx::query_as::<_, JobRow>(&format!(
r#"
SELECT {}
FROM hammerwork_jobs
WHERE queue_name = $1
AND status = 'Pending'
AND dependency_status IN ('none', 'satisfied')
AND scheduled_at <= now()
ORDER BY priority DESC, scheduled_at ASC
LIMIT $2
"#,
JOB_SELECT_FIELDS
))
.bind(queue_name)
.bind(i64::from(limit))
.fetch_all(&self.pool)
.await?;
Ok(super::decodable_rows(
"ready jobs",
rows.into_iter()
.map(|row| (row.id.to_string(), row.into_job())),
))
}
async fn fail_job_dependencies(&self, failed_job_id: JobId) -> Result<Vec<JobId>> {
super::retry_on_conflict(|| async {
let mut tx = self.pool.begin().await?;
let result = Self::fail_job_dependencies_in_tx(&mut tx, failed_job_id).await;
super::end_transaction(tx, result).await
})
.await
}
async fn get_workflow_jobs(
&self,
workflow_id: crate::workflow::WorkflowId,
) -> Result<Vec<Job>> {
let rows = sqlx::query_as::<_, JobRow>(&format!(
r#"
SELECT {}
FROM hammerwork_jobs
WHERE workflow_id = $1
ORDER BY created_at ASC
"#,
JOB_SELECT_FIELDS
))
.bind(workflow_id)
.fetch_all(&self.pool)
.await?;
rows.into_iter().map(|row| row.into_job()).collect()
}
async fn cancel_workflow(&self, workflow_id: crate::workflow::WorkflowId) -> Result<()> {
super::retry_on_conflict(|| async {
let mut tx = self.pool.begin().await?;
let result = Self::cancel_workflow_in_tx(&mut tx, workflow_id).await;
super::end_transaction(tx, result).await
})
.await
}
async fn archive_jobs(
&self,
queue_name: Option<&str>,
policy: &crate::archive::ArchivalPolicy,
config: &crate::archive::ArchivalConfig,
reason: crate::archive::ArchivalReason,
archived_by: Option<&str>,
) -> Result<crate::archive::ArchivalStats> {
super::retry_on_conflict(|| {
self.archive_jobs_once(queue_name, policy, config, reason.clone(), archived_by)
})
.await
}
async fn restore_archived_job(&self, job_id: JobId) -> Result<Job> {
super::retry_on_conflict(|| self.restore_archived_job_once(job_id)).await
}
async fn list_archived_jobs(
&self,
queue_name: Option<&str>,
limit: Option<u32>,
offset: Option<u32>,
) -> Result<Vec<crate::archive::ArchivedJob>> {
use crate::archive::ArchivedJob;
let mut query = "SELECT id, queue_name, status, created_at, archived_at, archival_reason, original_payload_size, payload_compressed, archived_by FROM hammerwork_jobs_archive".to_string();
let mut conditions = Vec::new();
if let Some(_queue) = queue_name {
conditions.push("queue_name = $1".to_string());
}
if !conditions.is_empty() {
query.push_str(" WHERE ");
query.push_str(&conditions.join(" AND "));
}
query.push_str(" ORDER BY archived_at DESC");
if let Some(_limit) = limit {
let param_num = if queue_name.is_some() { 2 } else { 1 };
query.push_str(&format!(" LIMIT ${}", param_num));
}
if let Some(_offset) = offset {
let param_num = if queue_name.is_some() && limit.is_some() {
3
} else if queue_name.is_some() || limit.is_some() {
2
} else {
1
};
query.push_str(&format!(" OFFSET ${}", param_num));
}
let mut sql_query = sqlx::query(&query);
if let Some(queue) = queue_name {
sql_query = sql_query.bind(queue);
}
if let Some(limit) = limit {
sql_query = sql_query.bind(limit as i64);
}
if let Some(offset) = offset {
sql_query = sql_query.bind(offset as i64);
}
let rows = sql_query.fetch_all(&self.pool).await?;
let decode = |row: &sqlx::postgres::PgRow| -> Result<ArchivedJob> {
Ok(ArchivedJob {
id: row.try_get("id")?,
queue_name: row.try_get("queue_name")?,
status: crate::archive::parse_archived_status(&row.try_get::<String, _>("status")?),
created_at: row.try_get("created_at")?,
archived_at: row.try_get("archived_at")?,
archival_reason: crate::archive::ArchivalReason::parse_from_db(
&row.try_get::<String, _>("archival_reason")?,
)
.unwrap_or_default(),
original_payload_size: row
.try_get::<Option<i32>, _>("original_payload_size")?
.map(|s| super::db_int(s, "original_payload_size"))
.transpose()?,
payload_compressed: row.try_get("payload_compressed")?,
archived_by: row.try_get("archived_by")?,
})
};
Ok(super::decodable_rows(
"archived jobs",
rows.iter().map(|row| {
let id = row
.try_get::<uuid::Uuid, _>("id")
.map(|id| id.to_string())
.unwrap_or_default();
(id, decode(row))
}),
))
}
async fn purge_archived_jobs(&self, older_than: DateTime<Utc>) -> Result<u64> {
let result = sqlx::query("DELETE FROM hammerwork_jobs_archive WHERE archived_at <= $1")
.bind(older_than)
.execute(&self.pool)
.await?;
Ok(result.rows_affected())
}
async fn get_archival_stats(
&self,
queue_name: Option<&str>,
) -> Result<crate::archive::ArchivalStats> {
use crate::archive::ArchivalStats;
let mut base_query = "SELECT
COUNT(*) as job_count,
COALESCE(SUM(original_payload_size), 0) as total_original_size,
COALESCE(SUM(LENGTH(payload)), 0) as total_compressed_size,
MAX(archived_at) as last_archived_at
FROM hammerwork_jobs_archive"
.to_string();
if let Some(_queue) = queue_name {
base_query.push_str(" WHERE queue_name = $1");
}
let mut query = sqlx::query(&base_query);
if let Some(queue) = queue_name {
query = query.bind(queue);
}
let row = query.fetch_one(&self.pool).await?;
let job_count: i64 = row.try_get("job_count")?;
let total_original_size: i64 = row.try_get("total_original_size")?;
let total_compressed_size: i64 = row.try_get("total_compressed_size")?;
let last_archived_at: Option<DateTime<Utc>> = row.try_get("last_archived_at")?;
let compression_ratio = if total_compressed_size > 0 {
total_original_size as f64 / total_compressed_size as f64
} else {
1.0
};
Ok(ArchivalStats {
jobs_archived: super::db_int(job_count, "archived job count")?,
jobs_purged: 0, bytes_archived: super::db_int(total_compressed_size, "archived bytes")?,
bytes_purged: 0,
compression_ratio,
operation_duration: std::time::Duration::from_secs(0),
last_run_at: last_archived_at.unwrap_or(Utc::now()),
})
}
async fn pause_queue(&self, queue_name: &str, paused_by: Option<&str>) -> Result<()> {
sqlx::query(
r#"
INSERT INTO hammerwork_queue_pause (queue_name, paused_by, paused_at, created_at, updated_at)
VALUES ($1, $2, NOW(), NOW(), NOW())
ON CONFLICT (queue_name)
DO UPDATE SET
paused_by = EXCLUDED.paused_by,
paused_at = NOW(),
updated_at = NOW()
"#,
)
.bind(queue_name)
.bind(paused_by)
.execute(&self.pool)
.await?;
Ok(())
}
async fn resume_queue(&self, queue_name: &str, _resumed_by: Option<&str>) -> Result<()> {
sqlx::query("DELETE FROM hammerwork_queue_pause WHERE queue_name = $1")
.bind(queue_name)
.execute(&self.pool)
.await?;
Ok(())
}
async fn is_queue_paused(&self, queue_name: &str) -> Result<bool> {
let row = sqlx::query("SELECT 1 FROM hammerwork_queue_pause WHERE queue_name = $1")
.bind(queue_name)
.fetch_optional(&self.pool)
.await?;
Ok(row.is_some())
}
async fn get_queue_pause_info(
&self,
queue_name: &str,
) -> Result<Option<super::QueuePauseInfo>> {
let row = sqlx::query(
"SELECT queue_name, paused_at, paused_by, reason FROM hammerwork_queue_pause WHERE queue_name = $1"
)
.bind(queue_name)
.fetch_optional(&self.pool)
.await?;
row.as_ref().map(pause_info_from_row).transpose()
}
async fn get_paused_queues(&self) -> Result<Vec<super::QueuePauseInfo>> {
let rows = sqlx::query(
"SELECT queue_name, paused_at, paused_by, reason FROM hammerwork_queue_pause ORDER BY paused_at DESC"
)
.fetch_all(&self.pool)
.await?;
rows.iter().map(pause_info_from_row).collect()
}
async fn heartbeat_job(&self, run: &Job, lease: std::time::Duration) -> Result<bool> {
let started = super::run_started_range(run);
let result = sqlx::query(
"UPDATE hammerwork_jobs SET last_heartbeat_at = now(), \
lease_expires_at = now() + $1 WHERE id = $2 AND status = $3 AND attempts = $4 \
AND ($5::timestamptz IS NULL OR (started_at > $5 AND started_at < $6))",
)
.bind(super::lease_interval(lease))
.bind(run.id)
.bind(JobStatus::Running)
.bind(run.attempts)
.bind(started.map(|(from, _)| from))
.bind(started.map(|(_, to)| to))
.execute(&self.pool)
.await?;
Ok(result.rows_affected() > 0)
}
async fn start_claimed_run(&self, run: &mut Job, lease: std::time::Duration) -> Result<bool> {
let started = super::run_started_range(run);
let restarted: Option<DateTime<Utc>> = sqlx::query_scalar(
"UPDATE hammerwork_jobs SET started_at = now(), last_heartbeat_at = now(), \
lease_expires_at = now() + $1 WHERE id = $2 AND status = $3 AND attempts = $4 \
AND ($5::timestamptz IS NULL OR (started_at > $5 AND started_at < $6)) \
RETURNING started_at",
)
.bind(super::lease_interval(lease))
.bind(run.id)
.bind(JobStatus::Running)
.bind(run.attempts)
.bind(started.map(|(from, _)| from))
.bind(started.map(|(_, to)| to))
.fetch_optional(&self.pool)
.await?;
match restarted {
Some(started_at) => {
run.started_at = Some(started_at);
Ok(true)
}
None => Ok(false),
}
}
async fn requeue_stale_jobs(
&self,
older_than: std::time::Duration,
) -> Result<super::StaleJobRecovery> {
let recovery =
super::retry_on_conflict(|| self.requeue_stale_jobs_once(older_than)).await?;
if !recovery.is_empty() {
tracing::warn!(
requeued = recovery.requeued.len(),
dead = recovery.dead.len(),
"Reclaimed stale Running jobs whose lease expired"
);
}
Ok(recovery)
}
async fn purge_expired_encrypted_jobs(&self) -> Result<super::EncryptedJobPurge> {
let mut tx = self.pool.begin().await?;
let now = db_now(&mut tx).await?;
let jobs = sqlx::query(&format!(
"DELETE FROM hammerwork_jobs WHERE is_encrypted = true \
AND retention_delete_at IS NOT NULL AND retention_delete_at <= $1 \
AND status IN ({})",
super::FINISHED_STATUS_SQL
))
.bind(now)
.execute(&mut *tx)
.await?
.rows_affected();
let archived_jobs = sqlx::query(
"DELETE FROM hammerwork_jobs_archive WHERE is_encrypted = true \
AND retention_delete_at IS NOT NULL AND retention_delete_at <= $1",
)
.bind(now)
.execute(&mut *tx)
.await?
.rows_affected();
tx.commit().await?;
let purge = super::EncryptedJobPurge {
jobs,
archived_jobs,
};
if purge.total() > 0 {
tracing::info!(
jobs = purge.jobs,
archived_jobs = purge.archived_jobs,
"Deleted encrypted jobs whose retention period ended"
);
}
Ok(purge)
}
}
impl crate::queue::JobQueue<Postgres> {
async fn requeue_stale_jobs_once(
&self,
older_than: std::time::Duration,
) -> Result<super::StaleJobRecovery> {
let mut tx = self.pool.begin().await?;
let result = Self::requeue_stale_jobs_in_tx(&mut tx, older_than).await;
super::end_transaction(tx, result).await
}
async fn requeue_stale_jobs_in_tx(
conn: &mut sqlx::PgConnection,
older_than: std::time::Duration,
) -> Result<super::StaleJobRecovery> {
let now = db_now(conn).await?;
let cutoff = super::saturating_sub_from(now, older_than);
let rows = sqlx::query(
r#"
WITH stale AS (
SELECT id FROM hammerwork_jobs
WHERE status = $1
AND (
(last_heartbeat_at IS NOT NULL
AND last_heartbeat_at >= started_at
AND lease_expires_at < $2)
OR ((last_heartbeat_at IS NULL OR lease_expires_at IS NULL
OR last_heartbeat_at < started_at)
AND started_at < $3)
)
FOR UPDATE SKIP LOCKED
)
UPDATE hammerwork_jobs AS j
SET status = CASE WHEN j.attempts >= j.max_attempts THEN $4 ELSE $5 END,
failed_at = CASE WHEN j.attempts >= j.max_attempts THEN $2 ELSE j.failed_at END,
scheduled_at = CASE WHEN j.attempts >= j.max_attempts
THEN j.scheduled_at ELSE $2 END,
started_at = CASE WHEN j.attempts >= j.max_attempts
THEN j.started_at ELSE NULL END,
error_message = $6,
last_heartbeat_at = NULL,
lease_expires_at = NULL
FROM stale
WHERE j.id = stale.id AND j.status = $1
RETURNING j.id, j.status
"#,
)
.bind(JobStatus::Running)
.bind(now)
.bind(cutoff)
.bind(JobStatus::Dead)
.bind(JobStatus::Pending)
.bind(super::STALE_JOB_ERROR_MESSAGE)
.fetch_all(&mut *conn)
.await?;
let mut recovery = super::StaleJobRecovery::default();
for row in rows {
let id: uuid::Uuid = row.try_get("id")?;
let status: String = row.try_get("status")?;
if status == "Dead" {
recovery.dead.push(id);
} else {
recovery.requeued.push(id);
}
}
for id in &recovery.dead {
let row = sqlx::query_as::<_, JobRow>(&format!(
"SELECT {JOB_SELECT_FIELDS} FROM hammerwork_jobs WHERE id = $1"
))
.bind(id)
.fetch_one(&mut *conn)
.await?;
let job = match row.into_job() {
Ok(job) => job,
Err(error) => {
super::warn_undecodable_stale_job(*id, &error);
continue;
}
};
lifecycle::after_terminal::<Self>(conn, &job, JobStatus::Dead, now, true).await?;
}
Ok(recovery)
}
}
impl crate::queue::JobQueue<Postgres> {
async fn enqueue_with_tx(
&self,
tx: &mut sqlx::Transaction<'_, Postgres>,
job: Job,
) -> Result<JobId> {
insert_jobs(tx, std::slice::from_ref(&job)).await?;
Ok(job.id)
}
}
const INSERT_CHUNK_ROWS: usize = 1000;
async fn insert_jobs(conn: &mut sqlx::PgConnection, jobs: &[Job]) -> Result<()> {
let strategies = jobs
.iter()
.map(super::retry_strategy_json)
.collect::<Result<Vec<_>>>()?;
let encryption = jobs
.iter()
.map(super::EncryptionColumns::for_job)
.collect::<Result<Vec<_>>>()?;
for ((chunk, strategies), encryption) in jobs
.chunks(INSERT_CHUNK_ROWS)
.zip(strategies.chunks(INSERT_CHUNK_ROWS))
.zip(encryption.chunks(INSERT_CHUNK_ROWS))
{
let mut builder = sqlx::QueryBuilder::<Postgres>::new(format!(
"INSERT INTO hammerwork_jobs ({INSERT_JOB_COLUMNS}) "
));
let rows = chunk.iter().zip(strategies).zip(encryption);
builder.push_values(rows, |mut row, ((job, strategy), encryption)| {
row.push_bind(job.id)
.push_bind(&job.queue_name)
.push_bind(&job.payload)
.push_bind(job.status)
.push_bind(job.priority.as_i32())
.push_bind(job.attempts)
.push_bind(job.max_attempts)
.push_bind(super::timeout_db_seconds(job.timeout))
.push_bind(job.created_at)
.push_bind(job.scheduled_at)
.push_bind(job.started_at)
.push_bind(job.completed_at)
.push_bind(job.failed_at)
.push_bind(job.timed_out_at)
.push_bind(&job.error_message)
.push_bind(&job.cron_schedule)
.push_bind(job.next_run_at)
.push_bind(job.recurring)
.push_bind(&job.timezone)
.push_bind(job.batch_id)
.push_bind(result_storage_str(&job.result_config.storage))
.push_bind(
job.result_config
.ttl
.map(|d| super::saturating_i64(d.as_secs())),
)
.push_bind(
job.result_config
.max_size_bytes
.map(|s| super::saturating_i64(s as u64)),
)
.push_bind(strategy)
.push_bind(&job.depends_on)
.push_bind(&job.dependents)
.push_bind(job.dependency_status.as_str())
.push_bind(job.workflow_id)
.push_bind(&job.workflow_name)
.push_bind(&job.trace_id)
.push_bind(&job.correlation_id)
.push_bind(&job.parent_span_id)
.push_bind(&job.span_context)
.push_bind(encryption.is_encrypted)
.push_bind(&encryption.key_id)
.push_bind(encryption.algorithm)
.push_bind(&encryption.ciphertext)
.push_bind(&encryption.nonce)
.push_bind(&encryption.tag)
.push_bind(&encryption.metadata)
.push_bind(&encryption.payload_hash)
.push_bind(super::pii_fields_column(job))
.push_bind(encryption.retention_policy)
.push_bind(encryption.retention_delete_at)
.push_bind(encryption.encrypted_at);
});
builder.build().execute(&mut *conn).await?;
}
Ok(())
}
async fn db_now(conn: &mut sqlx::PgConnection) -> Result<DateTime<Utc>> {
Ok(sqlx::query_scalar("SELECT now()")
.fetch_one(&mut *conn)
.await?)
}
fn pause_info_from_row(row: &sqlx::postgres::PgRow) -> Result<super::QueuePauseInfo> {
Ok(super::QueuePauseInfo {
queue_name: row.try_get("queue_name")?,
paused_at: row.try_get("paused_at")?,
paused_by: row.try_get("paused_by")?,
reason: row.try_get("reason")?,
})
}
const RUNNABLE_IN_QUEUE_SQL: &str = "j.queue_name = $1 AND j.status = 'Pending' \
AND j.scheduled_at <= now() AND j.dependency_status IN ('none', 'satisfied') \
AND (j.recurring OR j.cron_schedule IS NULL) \
AND NOT EXISTS (SELECT 1 FROM hammerwork_queue_pause p WHERE p.queue_name = $1)";
fn claim_sql(by_priority: bool) -> String {
format!(
"UPDATE hammerwork_jobs SET status = 'Running', started_at = now(), \
attempts = attempts + 1, last_heartbeat_at = now(), lease_expires_at = now() + $2 \
WHERE id = (SELECT j.id FROM hammerwork_jobs j WHERE {RUNNABLE_IN_QUEUE_SQL}{} \
ORDER BY j.priority DESC, j.scheduled_at ASC LIMIT 1 FOR UPDATE SKIP LOCKED) \
RETURNING {JOB_SELECT_FIELDS}",
if by_priority {
" AND j.priority = $3"
} else {
""
}
)
}
fn claim_batch_sql(by_priority: bool) -> String {
format!(
"UPDATE hammerwork_jobs SET status = 'Running', started_at = now(), \
attempts = attempts + 1, last_heartbeat_at = now(), lease_expires_at = now() + $2 \
WHERE id IN (SELECT j.id FROM hammerwork_jobs j WHERE {RUNNABLE_IN_QUEUE_SQL}{} \
ORDER BY j.priority DESC, j.scheduled_at ASC LIMIT $3 FOR UPDATE SKIP LOCKED) \
RETURNING {JOB_SELECT_FIELDS}",
if by_priority {
" AND j.priority = $4"
} else {
""
}
)
}
fn claimed_jobs(rows: Vec<JobRow>) -> Vec<Job> {
let mut jobs = super::decodable_rows(
"claimed jobs",
rows.into_iter()
.map(|row| (row.id.to_string(), row.into_job())),
);
super::sort_in_claim_order(&mut jobs);
jobs
}
async fn insert_dependent_jobs(conn: &mut sqlx::PgConnection, jobs: &mut [Job]) -> Result<()> {
let parents = lifecycle::external_dependencies(jobs);
if parents.is_empty() {
return insert_jobs(conn, jobs).await;
}
let mut states = HashMap::new();
let stored: Vec<(uuid::Uuid, String, Option<String>)> = sqlx::query_as(
"SELECT j.id, j.status, w.failure_policy FROM hammerwork_jobs j \
LEFT JOIN hammerwork_workflows w ON w.id = j.workflow_id \
WHERE j.id = ANY($1) ORDER BY j.id FOR SHARE OF j",
)
.bind(&parents)
.fetch_all(&mut *conn)
.await?;
for (id, status, policy) in stored {
states.insert(id, ParentState::from_db(&status, policy.as_deref()));
}
let archived: Vec<(uuid::Uuid, String)> =
sqlx::query_as("SELECT id, status FROM hammerwork_jobs_archive WHERE id = ANY($1)")
.bind(&parents)
.fetch_all(&mut *conn)
.await?;
for (id, status) in archived {
states
.entry(id)
.or_insert_with(|| ParentState::from_archived(&status));
}
let now = db_now(conn).await?;
let failed = lifecycle::settle_new_dependents(jobs, &states, now);
insert_jobs(conn, jobs).await?;
let mut workflows: Vec<uuid::Uuid> = jobs
.iter()
.filter(|job| failed.contains(&job.id))
.filter_map(|job| job.workflow_id)
.collect();
workflows.sort_unstable();
workflows.dedup();
for id in workflows {
if let Some(workflow) = crate::queue::JobQueue::<Postgres>::lock_workflow(conn, id).await? {
crate::queue::JobQueue::<Postgres>::refresh_workflow(conn, &workflow, now).await?;
}
}
Ok(())
}
impl TransitionStore for crate::queue::JobQueue<Postgres> {
type Conn = sqlx::PgConnection;
async fn lock_job(conn: &mut sqlx::PgConnection, job_id: JobId) -> Result<Option<Job>> {
let row = sqlx::query_as::<_, JobRow>(&format!(
"SELECT {JOB_SELECT_FIELDS} FROM hammerwork_jobs WHERE id = $1 FOR UPDATE"
))
.bind(job_id)
.fetch_optional(&mut *conn)
.await?;
row.map(JobRow::into_job).transpose()
}
async fn db_now(conn: &mut sqlx::PgConnection) -> Result<DateTime<Utc>> {
db_now(conn).await
}
async fn transition_job(
&self,
job_id: JobId,
guard: Guard<'_>,
target: &Target,
reschedule_recurring: bool,
) -> Result<TransitionResult> {
super::retry_on_conflict(|| {
self.transition_job_once(job_id, guard, target, reschedule_recurring)
})
.await
}
async fn write_target(
conn: &mut sqlx::PgConnection,
job_id: JobId,
target: &Target,
now: DateTime<Utc>,
) -> Result<()> {
let query = match target {
Target::Completed => sqlx::query(
"UPDATE hammerwork_jobs SET status = 'Completed', completed_at = $2, \
last_heartbeat_at = NULL, lease_expires_at = NULL WHERE id = $1",
)
.bind(job_id)
.bind(now),
Target::Failed(error) | Target::Dead(error) => sqlx::query(
"UPDATE hammerwork_jobs SET status = $2, error_message = $3, failed_at = $4, \
last_heartbeat_at = NULL, lease_expires_at = NULL WHERE id = $1",
)
.bind(job_id)
.bind(target.terminal_status())
.bind(error)
.bind(now),
Target::TimedOut(error) => sqlx::query(
"UPDATE hammerwork_jobs SET status = 'TimedOut', error_message = $2, \
timed_out_at = $3, last_heartbeat_at = NULL, lease_expires_at = NULL \
WHERE id = $1",
)
.bind(job_id)
.bind(error)
.bind(now),
Target::Retry {
retry_at,
error,
timed_out,
} => sqlx::query(
"UPDATE hammerwork_jobs SET status = 'Pending', scheduled_at = $2, \
started_at = NULL, error_message = COALESCE($3, error_message), \
timed_out_at = COALESCE($4, timed_out_at), \
last_heartbeat_at = NULL, lease_expires_at = NULL WHERE id = $1",
)
.bind(job_id)
.bind(retry_at)
.bind(error)
.bind(timed_out.then_some(now)),
};
query.execute(&mut *conn).await?;
Ok(())
}
async fn reschedule_after_run(
conn: &mut sqlx::PgConnection,
job_id: JobId,
next_run_at: DateTime<Utc>,
run_status: JobStatus,
) -> Result<()> {
sqlx::query(
r#"
UPDATE hammerwork_jobs
SET status = 'Pending',
scheduled_at = $2,
next_run_at = $2,
attempts = 0,
started_at = NULL,
completed_at = NULL,
failed_at = CASE WHEN $3 THEN failed_at ELSE NULL END,
timed_out_at = CASE WHEN $3 THEN timed_out_at ELSE NULL END,
error_message = CASE WHEN $3 THEN error_message ELSE NULL END,
last_heartbeat_at = NULL,
lease_expires_at = NULL
WHERE id = $1
"#,
)
.bind(job_id)
.bind(next_run_at)
.bind(run_status != JobStatus::Completed)
.execute(&mut *conn)
.await?;
Ok(())
}
async fn lock_workflow(
conn: &mut sqlx::PgConnection,
workflow_id: uuid::Uuid,
) -> Result<Option<LockedWorkflow>> {
let row = sqlx::query(
"SELECT failure_policy, status, total_jobs FROM hammerwork_workflows \
WHERE id = $1 FOR UPDATE",
)
.bind(workflow_id)
.fetch_optional(&mut *conn)
.await?;
let Some(row) = row else {
return Ok(None);
};
Ok(Some(LockedWorkflow {
id: workflow_id,
policy: FailurePolicy::parse_from_db(row.try_get("failure_policy")?)?,
status: WorkflowStatus::parse_from_db(row.try_get("status")?)?,
total_jobs: usize::try_from(row.try_get::<i32, _>("total_jobs")?).unwrap_or(0),
}))
}
async fn resolve_dependents(
conn: &mut sqlx::PgConnection,
completed_job_id: JobId,
) -> Result<Vec<JobId>> {
let dependents: Vec<uuid::Uuid> = sqlx::query_scalar(
r#"
SELECT id FROM hammerwork_jobs
WHERE depends_on @> ARRAY[$1]::uuid[]
AND dependency_status = 'waiting'
AND status = 'Pending'
ORDER BY id
FOR UPDATE
"#,
)
.bind(completed_job_id)
.fetch_all(&mut *conn)
.await?;
if dependents.is_empty() {
return Ok(Vec::new());
}
Ok(sqlx::query_scalar(
r#"
UPDATE hammerwork_jobs j
SET dependency_status = 'satisfied'
WHERE j.id = ANY($1)
AND NOT EXISTS (
SELECT 1 FROM unnest(j.depends_on) AS dependency(id)
WHERE NOT EXISTS (
SELECT 1 FROM hammerwork_jobs d
WHERE d.id = dependency.id AND d.status = 'Completed')
AND NOT EXISTS (
SELECT 1 FROM hammerwork_jobs_archive a
WHERE a.id = dependency.id AND a.status = 'Completed')
)
RETURNING j.id
"#,
)
.bind(&dependents)
.fetch_all(&mut *conn)
.await?)
}
async fn fail_dependents(
conn: &mut sqlx::PgConnection,
failed_job_id: JobId,
now: DateTime<Utc>,
) -> Result<Vec<JobId>> {
let mut failed = Vec::new();
let mut to_check = vec![failed_job_id];
while let Some(current) = to_check.pop() {
let ids: Vec<uuid::Uuid> = sqlx::query_scalar(
r#"
UPDATE hammerwork_jobs
SET dependency_status = 'failed',
status = 'Failed',
failed_at = $2,
error_message = $3
WHERE depends_on @> ARRAY[$1]::uuid[]
AND dependency_status IN ('waiting', 'satisfied')
AND status = 'Pending'
RETURNING id
"#,
)
.bind(current)
.bind(now)
.bind(lifecycle::dependency_failed_message(current))
.fetch_all(&mut *conn)
.await?;
for id in ids {
if !failed.contains(&id) {
failed.push(id);
to_check.push(id);
}
}
}
Ok(failed)
}
async fn fail_workflow_pending(
conn: &mut sqlx::PgConnection,
workflow_id: uuid::Uuid,
failed_job_id: JobId,
now: DateTime<Utc>,
) -> Result<Vec<JobId>> {
Ok(sqlx::query_scalar(
r#"
UPDATE hammerwork_jobs
SET status = 'Failed',
failed_at = $3,
error_message = $4,
dependency_status = CASE WHEN dependency_status = 'waiting'
THEN 'failed' ELSE dependency_status END
WHERE workflow_id = $1
AND id <> $2
AND status IN ('Pending', 'Retrying')
RETURNING id
"#,
)
.bind(workflow_id)
.bind(failed_job_id)
.bind(now)
.bind(lifecycle::workflow_failed_message(failed_job_id))
.fetch_all(&mut *conn)
.await?)
}
async fn refresh_workflow(
conn: &mut sqlx::PgConnection,
workflow: &LockedWorkflow,
now: DateTime<Utc>,
) -> Result<()> {
let counts: Vec<(String, i64)> = sqlx::query_as(
"SELECT status, COUNT(*) FROM hammerwork_jobs WHERE workflow_id = $1 GROUP BY status",
)
.bind(workflow.id)
.fetch_all(&mut *conn)
.await?;
let progress = WorkflowProgress::from_status_counts(counts);
let status = progress.status(workflow.total_jobs, &workflow.policy, &workflow.status);
sqlx::query(
r#"
UPDATE hammerwork_workflows
SET completed_jobs = $2,
failed_jobs = $3,
status = $4,
completed_at = COALESCE(completed_at, $5),
failed_at = COALESCE(failed_at, $6)
WHERE id = $1
"#,
)
.bind(workflow.id)
.bind(progress.completed as i32)
.bind(progress.failed as i32)
.bind(status.as_str())
.bind((status == WorkflowStatus::Completed).then_some(now))
.bind((status == WorkflowStatus::Failed).then_some(now))
.execute(&mut *conn)
.await?;
Ok(())
}
async fn apply_batch_outcome(
conn: &mut sqlx::PgConnection,
batch_id: uuid::Uuid,
job_id: JobId,
status: JobStatus,
now: DateTime<Utc>,
) -> Result<Vec<JobId>> {
let row = sqlx::query(
"SELECT total_jobs, failure_mode FROM hammerwork_batches WHERE id = $1 FOR UPDATE",
)
.bind(batch_id)
.fetch_optional(&mut *conn)
.await?;
let Some(row) = row else {
return Ok(Vec::new());
};
let total_jobs: i32 = row.try_get("total_jobs")?;
let mode = crate::batch::parse_failure_mode(&row.try_get::<String, _>("failure_mode")?);
let fail_fast = status != JobStatus::Completed && mode == PartialFailureMode::FailFast;
let cancelled: Vec<JobId> = if fail_fast {
sqlx::query_scalar(
r#"
UPDATE hammerwork_jobs
SET status = 'Failed', failed_at = $2, error_message = $3
WHERE batch_id = $1 AND status IN ('Pending', 'Retrying')
RETURNING id
"#,
)
.bind(batch_id)
.bind(now)
.bind(lifecycle::batch_failed_message(job_id))
.fetch_all(&mut *conn)
.await?
} else {
Vec::new()
};
let unfinished: bool = sqlx::query_scalar(&format!(
"SELECT EXISTS (SELECT 1 FROM hammerwork_jobs WHERE batch_id = $1 \
AND status IN ({}))",
lifecycle::UNFINISHED_STATUS_SQL
))
.bind(batch_id)
.fetch_one(&mut *conn)
.await?;
if unfinished && !fail_fast {
sqlx::query("UPDATE hammerwork_batches SET status = $2 WHERE id = $1 AND status = $3")
.bind(batch_id)
.bind(BatchStatus::Processing)
.bind(BatchStatus::Pending)
.execute(&mut *conn)
.await?;
return Ok(cancelled);
}
let counts: Vec<(String, i64)> = sqlx::query_as(
"SELECT status, COUNT(*) FROM (
SELECT status FROM hammerwork_jobs WHERE batch_id = $1
UNION ALL
SELECT status FROM hammerwork_jobs_archive WHERE batch_id = $1
) batch_jobs GROUP BY status",
)
.bind(batch_id)
.fetch_all(&mut *conn)
.await?;
let progress = BatchProgress::from_status_counts(
counts,
super::db_int(total_jobs, "total_jobs")?,
&mode,
);
sqlx::query(
r#"
UPDATE hammerwork_batches
SET completed_jobs = $2,
failed_jobs = $3,
pending_jobs = $4,
status = $5,
completed_at = COALESCE(completed_at, $6),
error_summary = COALESCE($7, error_summary)
WHERE id = $1
"#,
)
.bind(batch_id)
.bind(progress.completed as i32)
.bind(progress.failed as i32)
.bind(progress.pending as i32)
.bind(progress.status)
.bind((progress.pending == 0).then_some(now))
.bind(
(progress.failed > 0)
.then(|| format!("{} of {} jobs failed", progress.failed, total_jobs)),
)
.execute(&mut *conn)
.await?;
Ok(cancelled)
}
}
impl crate::queue::JobQueue<Postgres> {
async fn transition_job_once(
&self,
job_id: JobId,
guard: Guard<'_>,
target: &Target,
reschedule_recurring: bool,
) -> Result<TransitionResult> {
let mut tx = self.pool.begin().await?;
let result = lifecycle::transition_in_tx::<Self>(
&mut tx,
job_id,
guard,
target,
reschedule_recurring,
)
.await;
super::end_transaction(tx, result).await
}
async fn fail_job_dependencies_in_tx(
conn: &mut sqlx::PgConnection,
failed_job_id: JobId,
) -> Result<Vec<JobId>> {
let workflow_id: Option<Option<uuid::Uuid>> =
sqlx::query_scalar("SELECT workflow_id FROM hammerwork_jobs WHERE id = $1 FOR UPDATE")
.bind(failed_job_id)
.fetch_optional(&mut *conn)
.await?;
lifecycle::fail_job_dependencies::<Self>(conn, failed_job_id, workflow_id.flatten()).await
}
async fn cancel_workflow_in_tx(
conn: &mut sqlx::PgConnection,
workflow_id: crate::workflow::WorkflowId,
) -> Result<()> {
Self::lock_workflow(conn, workflow_id).await?;
let now = db_now(conn).await?;
sqlx::query(&format!(
r#"
UPDATE hammerwork_jobs
SET status = 'Failed',
failed_at = $2,
error_message = $3,
dependency_status = CASE WHEN dependency_status = 'waiting'
THEN 'failed' ELSE dependency_status END,
last_heartbeat_at = NULL,
lease_expires_at = NULL
WHERE workflow_id = $1
AND status IN ({})
"#,
lifecycle::UNFINISHED_STATUS_SQL
))
.bind(workflow_id)
.bind(now)
.bind(lifecycle::WORKFLOW_CANCELLED_MESSAGE)
.execute(&mut *conn)
.await?;
sqlx::query(
r#"
UPDATE hammerwork_workflows
SET status = 'cancelled',
failed_at = COALESCE(failed_at, $2)
WHERE id = $1
"#,
)
.bind(workflow_id)
.bind(now)
.execute(&mut *conn)
.await?;
Ok(())
}
async fn rejected_transition(
&self,
job_id: JobId,
transition: JobTransition,
) -> Result<crate::HammerworkError> {
let status: Option<String> =
sqlx::query_scalar("SELECT status FROM hammerwork_jobs WHERE id = $1")
.bind(job_id)
.fetch_optional(&self.pool)
.await?;
Ok(lifecycle::rejection_error(
job_id,
transition,
status.as_deref(),
))
}
}
impl crate::queue::JobQueue<Postgres> {
async fn archive_jobs_once(
&self,
queue_name: Option<&str>,
policy: &crate::archive::ArchivalPolicy,
config: &crate::archive::ArchivalConfig,
reason: crate::archive::ArchivalReason,
archived_by: Option<&str>,
) -> Result<crate::archive::ArchivalStats> {
use crate::archive::ArchivalStats;
if !policy.enabled {
return Ok(ArchivalStats::default());
}
let start_time = Utc::now();
let mut jobs_archived = 0u64;
let mut bytes_archived = 0u64;
let mut total_compression_ratio = 0.0;
let now = Utc::now();
let mut thresholds = Vec::new();
let mut status_conditions = Vec::new();
let mut param = if queue_name.is_some() { 2 } else { 1 };
for (after, status, column) in crate::archive::archival_candidates(policy) {
if let Some(threshold) =
after.and_then(|after| crate::archive::age_threshold(now, after))
{
status_conditions.push(format!(
"(status = '{status}' AND {column} IS NOT NULL AND {column} <= ${param})"
));
thresholds.push(threshold);
param += 1;
}
}
if status_conditions.is_empty() || policy.batch_size == 0 {
return Ok(ArchivalStats::default());
}
let select = |limit: usize| {
format!(
"SELECT {} FROM hammerwork_jobs WHERE archived_at IS NULL{} AND ({}) \
AND NOT (id = ANY(${param})) \
ORDER BY created_at ASC LIMIT {limit} FOR UPDATE SKIP LOCKED",
JOB_SELECT_FIELDS,
if queue_name.is_some() {
" AND queue_name = $1"
} else {
""
},
status_conditions.join(" OR "),
)
};
let mut tx = self.pool.begin().await?;
let mut jobs_to_archive = Vec::new();
let mut examined: Vec<uuid::Uuid> = Vec::new();
while jobs_to_archive.len() < policy.batch_size {
let wanted = policy.batch_size - jobs_to_archive.len();
let select = select(wanted);
let mut query = sqlx::query_as::<_, JobRow>(&select);
if let Some(queue) = queue_name {
query = query.bind(queue);
}
for threshold in &thresholds {
query = query.bind(*threshold);
}
let rows = query.bind(&examined).fetch_all(&mut *tx).await?;
let fetched = rows.len();
let mut skipped = false;
for row in rows {
let id = row.id;
examined.push(id);
match row.into_job() {
Ok(job) => jobs_to_archive.push(job),
Err(error) => {
tracing::warn!(
job_id = %id,
error = %error,
"Skipping a job that cannot be decoded while archiving"
);
skipped = true;
}
}
}
if !skipped || fetched < wanted {
break;
}
}
let archived_at = Utc::now();
let mut archived_ids = Vec::with_capacity(jobs_to_archive.len());
let mut encrypted_ids = Vec::new();
for job in jobs_to_archive {
let (final_payload, is_compressed, original_size) =
crate::archive::encode_archived_payload(&job.payload, policy, config)?;
total_compression_ratio += if is_compressed {
original_size as f64 / final_payload.len() as f64
} else {
1.0
};
sqlx::query(
r#"
INSERT INTO hammerwork_jobs_archive (
id, queue_name, payload, payload_compressed, original_payload_size,
status, priority, attempts, max_attempts, created_at, scheduled_at,
started_at, completed_at, failed_at, timed_out_at, archived_at,
error_message, result, result_ttl, retry_strategy, timeout_seconds,
priority_weight, cron_schedule, next_run_at, recurring, timezone,
batch_id, depends_on, dependency_status, result_config,
trace_id, correlation_id, parent_span_id, span_context,
archival_reason, archived_by, original_table
) VALUES (
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16,
$17, $18, $19, $20, $21, $22, $23, $24, $25, $26, $27, $28, $29, $30,
$31, $32, $33, $34, $35, $36, $37
)
"#,
)
.bind(job.id)
.bind(&job.queue_name)
.bind(&final_payload)
.bind(is_compressed)
.bind(original_size as i32)
.bind(job.status)
.bind(job.priority.to_string())
.bind(job.attempts)
.bind(job.max_attempts)
.bind(job.created_at)
.bind(job.scheduled_at)
.bind(job.started_at)
.bind(job.completed_at)
.bind(job.failed_at)
.bind(job.timed_out_at)
.bind(archived_at)
.bind(job.error_message)
.bind(job.result_data)
.bind(job.result_expires_at)
.bind(
job.retry_strategy
.map(|rs| serde_json::to_string(&rs))
.transpose()?,
)
.bind(super::timeout_db_seconds(job.timeout))
.bind(job.priority.weight() as i32)
.bind(job.cron_schedule)
.bind(job.next_run_at)
.bind(job.recurring)
.bind(job.timezone)
.bind(job.batch_id)
.bind(&job.depends_on)
.bind(job.dependency_status.as_str())
.bind(serde_json::to_value(&job.result_config)?)
.bind(job.trace_id)
.bind(job.correlation_id)
.bind(job.parent_span_id)
.bind(job.span_context)
.bind(reason.as_str())
.bind(archived_by)
.bind("hammerwork_jobs")
.execute(&mut *tx)
.await?;
archived_ids.push(job.id);
if job.is_encrypted || !job.pii_fields.is_empty() {
encrypted_ids.push(job.id);
}
jobs_archived += 1;
bytes_archived += final_payload.len() as u64;
}
copy_encryption_columns(
&mut tx,
"hammerwork_jobs",
"hammerwork_jobs_archive",
&encrypted_ids,
)
.await?;
if !archived_ids.is_empty() {
sqlx::query("DELETE FROM hammerwork_jobs WHERE id = ANY($1)")
.bind(&archived_ids)
.execute(&mut *tx)
.await?;
}
tx.commit().await?;
let operation_duration = Utc::now() - start_time;
let compression_ratio = if jobs_archived > 0 && total_compression_ratio > 0.0 {
total_compression_ratio / jobs_archived as f64
} else {
1.0
};
Ok(ArchivalStats {
jobs_archived,
jobs_purged: 0,
bytes_archived,
bytes_purged: 0,
compression_ratio,
operation_duration: operation_duration
.to_std()
.unwrap_or(std::time::Duration::from_secs(0)),
last_run_at: Utc::now(),
})
}
async fn restore_archived_job_once(&self, job_id: JobId) -> Result<Job> {
let mut tx = self.pool.begin().await?;
let archived_row = sqlx::query(&format!(
"SELECT {ARCHIVED_JOB_FIELDS} FROM hammerwork_jobs_archive WHERE id = $1 FOR UPDATE"
))
.bind(job_id)
.fetch_optional(&mut *tx)
.await?
.ok_or_else(|| crate::HammerworkError::JobNotFound {
id: job_id.to_string(),
})?;
let job = crate::archive::reset_for_restore(archived_job_from_row(&archived_row)?);
sqlx::query("DELETE FROM hammerwork_jobs WHERE id = $1 AND archived_at IS NOT NULL")
.bind(job_id)
.execute(&mut *tx)
.await?;
let mut row = job.clone();
row.is_encrypted = false;
self.enqueue_with_tx(&mut tx, row).await?;
if job.is_encrypted || !job.pii_fields.is_empty() {
copy_encryption_columns(
&mut tx,
"hammerwork_jobs_archive",
"hammerwork_jobs",
&[job_id],
)
.await?;
}
sqlx::query("DELETE FROM hammerwork_jobs_archive WHERE id = $1")
.bind(job_id)
.execute(&mut *tx)
.await?;
tx.commit().await?;
Ok(job)
}
}