use crate::error::Result;
use crate::store::dialect::SqlDialect;
use crate::store::query::{QueryBuilder, QueryParam};
use crate::store::sqlite::dialect::SqliteDialect;
use crate::store::sqlite::{format_sqlite_timestamp, parse_sqlite_timestamp};
use crate::store::tables::DialectWorkerTable;
use crate::types::{WorkerRecord, WorkerStatus};
use async_trait::async_trait;
use chrono::Utc;
use sqlx::{Row, SqlitePool};
use std::str::FromStr;
const INSERT_WORKER: &str = r#"
INSERT INTO pgqrs_workers (name, queue_id, started_at, heartbeat_at, status)
VALUES ($1, $2, $3, $4, $5)
RETURNING id;
"#;
const GET_WORKER_BY_ID: &str = r#"
SELECT id, name, queue_id, started_at, heartbeat_at, shutdown_at, status
FROM pgqrs_workers
WHERE id = $1;
"#;
const LIST_ALL_WORKERS: &str = r#"
SELECT id, name, queue_id, started_at, heartbeat_at, shutdown_at, status
FROM pgqrs_workers
ORDER BY started_at DESC;
"#;
const LIST_WORKERS_BY_QUEUE: &str = r#"
SELECT id, name, queue_id, started_at, heartbeat_at, shutdown_at, status
FROM pgqrs_workers
WHERE queue_id = $1
ORDER BY started_at DESC;
"#;
const DELETE_WORKER_BY_ID: &str = r#"
DELETE FROM pgqrs_workers
WHERE id = $1;
"#;
const INSERT_EPHEMERAL_WORKER: &str = r#"
INSERT INTO pgqrs_workers (name, queue_id, status)
VALUES ($1, $2, 'ready')
RETURNING id, name, queue_id, started_at, heartbeat_at, shutdown_at, status;
"#;
#[derive(Debug, Clone)]
pub struct SqliteWorkerTable {
pool: SqlitePool,
}
impl SqliteWorkerTable {
pub fn new(pool: SqlitePool) -> Self {
Self { pool }
}
fn map_row(row: sqlx::sqlite::SqliteRow) -> Result<WorkerRecord> {
let id: i64 = row.try_get("id")?;
let name: String = row.try_get("name")?;
let queue_id: Option<i64> = row.try_get("queue_id")?;
let started_at_str: String = row.try_get("started_at")?;
let started_at = parse_sqlite_timestamp(&started_at_str)?;
let heartbeat_at_str: String = row.try_get("heartbeat_at")?;
let heartbeat_at = parse_sqlite_timestamp(&heartbeat_at_str)?;
let shutdown_at_str: Option<String> = row.try_get("shutdown_at")?;
let shutdown_at = match shutdown_at_str {
Some(s) => Some(parse_sqlite_timestamp(&s)?),
None => None,
};
let status_str: String = row.try_get("status")?;
let status = WorkerStatus::from_str(&status_str)
.map_err(|e| crate::error::Error::Internal { message: e })?;
Ok(WorkerRecord {
id,
name,
queue_id,
started_at,
heartbeat_at,
shutdown_at,
status,
})
}
pub async fn get_status(&self, worker_id: i64) -> Result<WorkerStatus> {
let status_str: String =
sqlx::query_scalar("SELECT status FROM pgqrs_workers WHERE id = $1")
.bind(worker_id)
.fetch_one(&self.pool)
.await
.map_err(|e| crate::error::Error::QueryFailed {
query: "GET_WORKER_STATUS".into(),
source: Box::new(e),
context: format!("Failed to get worker {} status", worker_id),
})?;
WorkerStatus::from_str(&status_str)
.map_err(|e| crate::error::Error::Internal { message: e })
}
pub async fn is_healthy(&self, worker_id: i64, max_age: chrono::Duration) -> Result<bool> {
let threshold = Utc::now() - max_age;
let threshold_str = format_sqlite_timestamp(&threshold);
let is_healthy: bool =
sqlx::query_scalar("SELECT heartbeat_at > $2 FROM pgqrs_workers WHERE id = $1")
.bind(worker_id)
.bind(threshold_str)
.fetch_one(&self.pool)
.await
.map_err(|e| crate::error::Error::QueryFailed {
query: "CHECK_WORKER_HEALTH".into(),
source: Box::new(e),
context: format!("Failed to check health for worker {}", worker_id),
})?;
Ok(is_healthy)
}
}
#[async_trait]
impl crate::store::WorkerTable for SqliteWorkerTable {
async fn insert(&self, data: crate::types::NewWorkerRecord) -> Result<WorkerRecord> {
let now = Utc::now();
let now_str = format_sqlite_timestamp(&now);
let status_str = WorkerStatus::Ready.to_string();
let id: i64 = sqlx::query_scalar(INSERT_WORKER)
.bind(&data.name)
.bind(data.queue_id)
.bind(&now_str)
.bind(&now_str)
.bind(&status_str)
.fetch_one(&self.pool)
.await
.map_err(|e| crate::error::Error::QueryFailed {
query: "INSERT_WORKER".into(),
source: Box::new(e),
context: format!("Failed to insert worker {}", data.name),
})?;
Ok(WorkerRecord {
id,
name: data.name,
queue_id: data.queue_id,
started_at: now,
heartbeat_at: now,
shutdown_at: None,
status: WorkerStatus::Ready,
})
}
async fn get(&self, id: i64) -> Result<WorkerRecord> {
let row = sqlx::query(GET_WORKER_BY_ID)
.bind(id)
.fetch_one(&self.pool)
.await
.map_err(|e| crate::error::Error::QueryFailed {
query: format!("GET_WORKER_BY_ID ({})", id),
source: Box::new(e),
context: format!("Failed to get worker {}", id),
})?;
Self::map_row(row)
}
async fn list(&self) -> Result<Vec<WorkerRecord>> {
let rows = sqlx::query(LIST_ALL_WORKERS)
.fetch_all(&self.pool)
.await
.map_err(|e| crate::error::Error::QueryFailed {
query: "LIST_ALL_WORKERS".into(),
source: Box::new(e),
context: "Failed to list all workers".into(),
})?;
let mut workers = Vec::with_capacity(rows.len());
for row in rows {
workers.push(Self::map_row(row)?);
}
Ok(workers)
}
async fn count(&self) -> Result<i64> {
let count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM pgqrs_workers")
.fetch_one(&self.pool)
.await
.map_err(|e| crate::error::Error::QueryFailed {
query: "COUNT_WORKERS".into(),
source: Box::new(e),
context: "Failed to count workers".into(),
})?;
Ok(count)
}
async fn delete(&self, id: i64) -> Result<u64> {
let result = sqlx::query(DELETE_WORKER_BY_ID)
.bind(id)
.execute(&self.pool)
.await
.map_err(|e| crate::error::Error::QueryFailed {
query: format!("DELETE_WORKER_BY_ID ({})", id),
source: Box::new(e),
context: format!("Failed to delete worker {}", id),
})?;
Ok(result.rows_affected())
}
async fn filter_by_fk(&self, queue_id: i64) -> Result<Vec<WorkerRecord>> {
let rows = sqlx::query(LIST_WORKERS_BY_QUEUE)
.bind(queue_id)
.fetch_all(&self.pool)
.await
.map_err(|e| crate::error::Error::QueryFailed {
query: format!("LIST_WORKERS_BY_QUEUE ({})", queue_id),
source: Box::new(e),
context: format!("Failed to filter workers by queue ID {}", queue_id),
})?;
let mut workers = Vec::with_capacity(rows.len());
for row in rows {
workers.push(Self::map_row(row)?);
}
Ok(workers)
}
async fn count_by_fk(&self, queue_id: i64) -> Result<i64> {
let count: i64 =
sqlx::query_scalar("SELECT COUNT(*) FROM pgqrs_workers WHERE queue_id = $1")
.bind(queue_id)
.fetch_one(&self.pool)
.await
.map_err(|e| crate::error::Error::QueryFailed {
query: "COUNT_WORKERS_BY_QUEUE".into(),
source: Box::new(e),
context: format!("Failed to count workers for queue {}", queue_id),
})?;
Ok(count)
}
async fn mark_stopped(&self, id: i64) -> Result<()> {
sqlx::query(SqliteDialect::WORKER.mark_stopped)
.bind(id)
.execute(&self.pool)
.await
.map_err(|e| crate::error::Error::QueryFailed {
query: "MARK_WORKER_STOPPED".into(),
source: Box::new(e),
context: format!("Failed to mark worker {} as stopped", id),
})?;
Ok(())
}
async fn count_for_queue(
&self,
queue_id: i64,
state: crate::types::WorkerStatus,
) -> Result<i64> {
let state_str = state.to_string();
let count: i64 = sqlx::query_scalar(
"SELECT COUNT(*) FROM pgqrs_workers WHERE queue_id = $1 AND status = $2",
)
.bind(queue_id)
.bind(state_str)
.fetch_one(&self.pool)
.await
.map_err(|e| crate::error::Error::QueryFailed {
query: format!("COUNT_WORKERS_BY_STATE (queue_id={})", queue_id),
source: Box::new(e),
context: format!("Failed to count workers for queue {}", queue_id),
})?;
Ok(count)
}
async fn count_zombies_for_queue(
&self,
queue_id: i64,
older_than: chrono::Duration,
) -> Result<i64> {
let threshold = Utc::now() - older_than;
let threshold_str = format_sqlite_timestamp(&threshold);
let count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM pgqrs_workers WHERE queue_id = $1 AND status IN ('ready', 'polling', 'suspended', 'interrupted') AND heartbeat_at < $2")
.bind(queue_id)
.bind(threshold_str)
.fetch_one(&self.pool)
.await
.map_err(|e| crate::error::Error::QueryFailed {
query: format!("COUNT_ZOMBIE_WORKERS (queue_id={})", queue_id),
source: Box::new(e),
context: format!("Failed to count zombie workers for queue {}", queue_id),
})?;
Ok(count)
}
async fn list_for_queue(
&self,
queue_id: i64,
state: crate::types::WorkerStatus,
) -> Result<Vec<WorkerRecord>> {
let state_str = state.to_string();
let rows = sqlx::query("SELECT id, name, queue_id, started_at, heartbeat_at, shutdown_at, status FROM pgqrs_workers WHERE queue_id = $1 AND status = $2 ORDER BY started_at DESC")
.bind(queue_id)
.bind(state_str)
.fetch_all(&self.pool)
.await
.map_err(|e| crate::error::Error::QueryFailed {
query: format!("LIST_WORKERS_BY_QUEUE_AND_STATE (queue_id={})", queue_id),
source: Box::new(e),
context: format!("Failed to list workers for queue {}", queue_id),
})?;
let mut workers = Vec::with_capacity(rows.len());
for row in rows {
workers.push(Self::map_row(row)?);
}
Ok(workers)
}
async fn list_zombies_for_queue(
&self,
queue_id: i64,
older_than: chrono::Duration,
) -> Result<Vec<WorkerRecord>> {
let threshold = Utc::now() - older_than;
let threshold_str = format_sqlite_timestamp(&threshold);
let rows = sqlx::query("SELECT id, name, queue_id, started_at, heartbeat_at, shutdown_at, status FROM pgqrs_workers WHERE queue_id = $1 AND status IN ('ready', 'polling', 'suspended', 'interrupted') AND heartbeat_at < $2 ORDER BY heartbeat_at ASC")
.bind(queue_id)
.bind(threshold_str)
.fetch_all(&self.pool)
.await
.map_err(|e| crate::error::Error::QueryFailed {
query: format!("LIST_ZOMBIE_WORKERS (queue_id={})", queue_id),
source: Box::new(e),
context: format!("Failed to list zombie workers for queue {}", queue_id),
})?;
let mut workers = Vec::with_capacity(rows.len());
for row in rows {
workers.push(Self::map_row(row)?);
}
Ok(workers)
}
async fn register(&self, queue_id: Option<i64>, name: &str) -> Result<WorkerRecord> {
let existing = sqlx::query("SELECT id, name, queue_id, started_at, heartbeat_at, shutdown_at, status FROM pgqrs_workers WHERE name = $1")
.bind(name)
.fetch_optional(&self.pool)
.await
.map_err(|e| crate::error::Error::QueryFailed {
query: format!("FIND_WORKER_BY_NAME ({})", name),
source: Box::new(e),
context: format!("Failed to find worker {}", name),
})?;
if let Some(row) = existing {
let worker = Self::map_row(row)?;
match worker.status {
WorkerStatus::Stopped => {
let now = Utc::now();
let now_str = format_sqlite_timestamp(&now);
let row = sqlx::query("UPDATE pgqrs_workers SET status = 'ready', queue_id = $2, started_at = $3, heartbeat_at = $3, shutdown_at = NULL WHERE id = $1 RETURNING id, name, queue_id, started_at, heartbeat_at, shutdown_at, status")
.bind(worker.id)
.bind(queue_id)
.bind(now_str)
.fetch_one(&self.pool)
.await
.map_err(|e| crate::error::Error::QueryFailed {
query: format!("RESET_WORKER_TO_READY ({})", worker.id),
source: Box::new(e),
context: format!("Failed to reset worker {}", name),
})?;
Self::map_row(row)
}
WorkerStatus::Ready => Err(crate::error::Error::ValidationFailed {
reason: format!(
"Worker {} is already active. Cannot register duplicate.",
name
),
}),
WorkerStatus::Suspended => Err(crate::error::Error::ValidationFailed {
reason: format!("Worker {} is suspended. Use resume() to reactivate.", name),
}),
WorkerStatus::Polling | WorkerStatus::Interrupted => {
Err(crate::error::Error::ValidationFailed {
reason: format!(
"Worker {} is already active. Cannot register duplicate.",
name
),
})
}
}
} else {
self.insert(crate::types::NewWorkerRecord {
name: name.to_string(),
queue_id,
})
.await
}
}
async fn register_ephemeral(&self, queue_id: Option<i64>) -> Result<WorkerRecord> {
let name = format!("__ephemeral__{}", uuid::Uuid::new_v4());
let row = sqlx::query(INSERT_EPHEMERAL_WORKER)
.bind(&name)
.bind(queue_id)
.fetch_one(&self.pool)
.await
.map_err(|e| crate::error::Error::QueryFailed {
query: "INSERT_EPHEMERAL_WORKER".into(),
source: Box::new(e),
context: "Failed to create ephemeral worker".into(),
})?;
Self::map_row(row)
}
async fn get_status(&self, id: i64) -> Result<WorkerStatus> {
self.get_status(id).await
}
async fn suspend(&self, id: i64) -> Result<()> {
<Self as DialectWorkerTable>::dialect_suspend(self, id).await
}
async fn resume(&self, id: i64) -> Result<()> {
<Self as DialectWorkerTable>::dialect_resume(self, id).await
}
async fn complete_poll(&self, id: i64) -> Result<()> {
<Self as DialectWorkerTable>::dialect_complete_poll(self, id).await
}
async fn shutdown(&self, id: i64) -> Result<()> {
<Self as DialectWorkerTable>::dialect_shutdown(self, id).await
}
async fn poll(&self, id: i64) -> Result<()> {
<Self as DialectWorkerTable>::dialect_poll(self, id).await
}
async fn interrupt(&self, id: i64) -> Result<()> {
<Self as DialectWorkerTable>::dialect_interrupt(self, id).await
}
async fn heartbeat(&self, id: i64) -> Result<()> {
<Self as DialectWorkerTable>::dialect_heartbeat(self, id).await
}
async fn is_healthy(&self, id: i64, max_age: chrono::Duration) -> Result<bool> {
self.is_healthy(id, max_age).await
}
}
#[async_trait]
impl DialectWorkerTable for SqliteWorkerTable {
type Dialect = SqliteDialect;
async fn execute_worker_update(&self, query: QueryBuilder) -> Result<u64> {
let mut builder = sqlx::query(query.sql());
for param in query.params() {
builder = match param {
QueryParam::I64(value) => builder.bind(*value),
QueryParam::I32(value) => builder.bind(*value),
QueryParam::String(value) => builder.bind(value),
QueryParam::Json(value) => builder.bind(value.to_string()),
QueryParam::DateTime(value) => {
builder.bind(value.map(|dt| format_sqlite_timestamp(&dt)))
}
};
}
let result =
builder
.execute(&self.pool)
.await
.map_err(|e| crate::error::Error::QueryFailed {
query: "SQLITE_WORKER_UPDATE".into(),
source: Box::new(e),
context: "Failed to execute sqlite worker update".into(),
})?;
Ok(result.rows_affected())
}
async fn query_worker_status(&self, worker_id: i64) -> Result<WorkerStatus> {
self.get_status(worker_id).await
}
}