a3s_flow/worker/

postgres.rs

use std::fmt;

use a3s_orm::{
    sql_query, Executor, FromRow, PostgresDialect, PostgresError, PostgresExecutor, PostgresRow,
    PostgresTransactionError, Query, SqlQuery,
};
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use uuid::Uuid;

use crate::error::{FlowError, Result};
use crate::store::{migrate_postgres_flow, verify_postgres_flow};

pub use super::task::PostgresDeadLetteredTask;
use super::{timestamp_nanos_saturating, FlowTask, FlowTaskLease, FlowTaskQueue};

/// A3S ORM-backed PostgreSQL task queue for shared workers.
///
/// Pending and inflight tasks live in one table and are scoped by `queue_name`.
/// Leasing uses an atomic `FOR UPDATE SKIP LOCKED` CTE, so multiple workers can
/// lease concurrently without taking the same task. Heartbeats rotate fencing
/// tokens; stale acknowledgements cannot delete a task owned by another worker.
#[derive(Clone)]
pub struct PostgresFlowTaskQueue {
    executor: PostgresExecutor,
    queue_name: String,
}

impl fmt::Debug for PostgresFlowTaskQueue {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("PostgresFlowTaskQueue")
            .field("queue_name", &self.queue_name)
            .finish_non_exhaustive()
    }
}

impl PostgresFlowTaskQueue {
    /// Connects to PostgreSQL, migrates the schema, and uses the default queue.
    pub async fn connect(database_url: impl AsRef<str>) -> Result<Self> {
        Self::connect_with_queue(database_url, "default").await
    }

    /// Connects to PostgreSQL, verifies a separately migrated schema, and uses
    /// the default queue without acquiring DDL authority.
    pub async fn connect_verified(database_url: impl AsRef<str>) -> Result<Self> {
        Self::connect_verified_with_queue(database_url, "default").await
    }

    /// Connect with the ORM's bounded non-TLS pool and run Flow migrations.
    ///
    /// Production hosts that require TLS or custom pool controls should create
    /// a configured [`PostgresExecutor`] and call
    /// [`Self::from_executor_with_queue`].
    pub async fn connect_with_queue(
        database_url: impl AsRef<str>,
        queue_name: impl AsRef<str>,
    ) -> Result<Self> {
        let executor = PostgresExecutor::connect_no_tls(database_url.as_ref(), 5)
            .map_err(postgres_queue_driver_error)?;
        Self::from_executor_with_queue(executor, queue_name).await
    }

    /// Connect with the ORM's bounded non-TLS pool and verify a separately
    /// migrated Flow schema for the named queue.
    pub async fn connect_verified_with_queue(
        database_url: impl AsRef<str>,
        queue_name: impl AsRef<str>,
    ) -> Result<Self> {
        let executor = PostgresExecutor::connect_no_tls(database_url.as_ref(), 5)
            .map_err(postgres_queue_driver_error)?;
        Self::from_executor_verified_with_queue(executor, queue_name).await
    }

    /// Uses a configured executor, migrates the schema, and uses the default queue.
    pub async fn from_executor(executor: PostgresExecutor) -> Result<Self> {
        Self::from_executor_with_queue(executor, "default").await
    }

    /// Uses a configured executor, verifies the Flow schema without mutation,
    /// and uses the default queue.
    pub async fn from_executor_verified(executor: PostgresExecutor) -> Result<Self> {
        Self::from_executor_verified_with_queue(executor, "default").await
    }

    /// Uses a configured executor and migrates the named queue schema.
    pub async fn from_executor_with_queue(
        executor: PostgresExecutor,
        queue_name: impl AsRef<str>,
    ) -> Result<Self> {
        let queue_name = validated_queue_name(queue_name.as_ref())?;
        migrate_postgres_flow(&executor).await?;
        Ok(Self::new(executor, queue_name))
    }

    /// Uses a configured executor and verifies the complete Flow schema for a
    /// named queue without applying migrations.
    pub async fn from_executor_verified_with_queue(
        executor: PostgresExecutor,
        queue_name: impl AsRef<str>,
    ) -> Result<Self> {
        let queue_name = validated_queue_name(queue_name.as_ref())?;
        verify_postgres_flow(&executor).await?;
        Ok(Self::new(executor, queue_name))
    }

    fn new(executor: PostgresExecutor, queue_name: String) -> Self {
        Self {
            executor,
            queue_name,
        }
    }

    /// Returns the configured A3S ORM executor.
    pub fn executor(&self) -> &PostgresExecutor {
        &self.executor
    }

    /// Returns the logical queue name used to scope rows.
    pub fn queue_name(&self) -> &str {
        &self.queue_name
    }

    /// Returns the number of currently leased tasks.
    pub async fn inflight_len(&self) -> Result<usize> {
        self.count_by_status("inflight").await
    }

    /// Returns the number of dead-lettered tasks.
    pub async fn dead_letter_len(&self) -> Result<usize> {
        let count = fetch_one_query(
            &self.executor,
            sql_query::<i64>(
                "SELECT COUNT(*)::BIGINT FROM flow_task_dead_letters WHERE queue_name = ",
            )
            .bind(self.queue_name.clone()),
        )
        .await?;
        postgres_count_to_usize(count)
    }

    /// Loads dead-lettered tasks in durable insertion order.
    pub async fn dead_lettered_tasks(&self) -> Result<Vec<PostgresDeadLetteredTask>> {
        let rows = fetch_all_query(
            &self.executor,
            sql_query::<(String, String, String, i64)>(
                "SELECT lease_id, task_json, reason, dead_lettered_at_nanos \
                 FROM flow_task_dead_letters WHERE queue_name = ",
            )
            .bind(self.queue_name.clone())
            .append(" ORDER BY dead_lettered_at_nanos ASC, dead_letter_id ASC"),
        )
        .await?;
        rows.into_iter().map(dead_letter_row).collect()
    }

    /// Returns leases at or before `cutoff` to pending dispatch.
    pub async fn requeue_inflight_older_than(&self, cutoff: DateTime<Utc>) -> Result<usize> {
        let rows = execute_query(
            &self.executor,
            sql_query::<()>(
                "UPDATE flow_tasks SET status = 'pending', lease_id = NULL, \
                 leased_at_nanos = NULL, updated_at_nanos = ",
            )
            .bind(timestamp_nanos_saturating(Utc::now()))
            .append(" WHERE queue_name = ")
            .bind(self.queue_name.clone())
            .append(" AND status = 'inflight' AND leased_at_nanos <= ")
            .bind(timestamp_nanos_saturating(cutoff)),
        )
        .await?;
        postgres_rows_affected_to_usize(rows)
    }

    /// Atomically moves leases at or before `cutoff` to the dead-letter table.
    pub async fn dead_letter_inflight_older_than(
        &self,
        cutoff: DateTime<Utc>,
        reason: impl Into<String>,
    ) -> Result<usize> {
        let queue_name = self.queue_name.clone();
        let reason = reason.into();
        let cutoff = timestamp_nanos_saturating(cutoff);
        let result = self
            .executor
            .transaction(|transaction| {
                Box::pin(async move {
                    let rows = fetch_all_query(
                        transaction,
                        sql_query::<(String, Option<String>, String, Option<i64>)>(
                            "SELECT task_id, lease_id, task_json, leased_at_nanos \
                             FROM flow_tasks WHERE queue_name = ",
                        )
                        .bind(queue_name.clone())
                        .append(" AND status = 'inflight' AND leased_at_nanos <= ")
                        .bind(cutoff)
                        .append(
                            " ORDER BY leased_at_nanos ASC, task_id ASC \
                             FOR UPDATE SKIP LOCKED",
                        ),
                    )
                    .await?;

                    let dead_lettered_at = timestamp_nanos_saturating(Utc::now());
                    for (task_id, lease_id, task_json, leased_at) in &rows {
                        let lease_id = lease_id.as_ref().ok_or_else(|| {
                            FlowError::Store(format!(
                                "inflight PostgreSQL task {task_id} has no lease"
                            ))
                        })?;
                        execute_query(
                            transaction,
                            sql_query::<()>(
                                "INSERT INTO flow_task_dead_letters (queue_name, \
                                 dead_letter_id, lease_id, task_json, reason, \
                                 dead_lettered_at_nanos, leased_at_nanos) VALUES (",
                            )
                            .bind(queue_name.clone())
                            .append(", ")
                            .bind(Uuid::new_v4().to_string())
                            .append(", ")
                            .bind(lease_id.clone())
                            .append(", ")
                            .bind(task_json.clone())
                            .append(", ")
                            .bind(reason.clone())
                            .append(", ")
                            .bind(dead_lettered_at)
                            .append(", ")
                            .bind(*leased_at)
                            .append(")"),
                        )
                        .await?;
                        execute_query(
                            transaction,
                            sql_query::<()>("DELETE FROM flow_tasks WHERE queue_name = ")
                                .bind(queue_name.clone())
                                .append(" AND task_id = ")
                                .bind(task_id.clone()),
                        )
                        .await?;
                    }
                    Ok(rows.len())
                })
            })
            .await;
        map_postgres_queue_transaction(result)
    }

    async fn count_by_status(&self, status: &str) -> Result<usize> {
        let count = fetch_one_query(
            &self.executor,
            sql_query::<i64>("SELECT COUNT(*)::BIGINT FROM flow_tasks WHERE queue_name = ")
                .bind(self.queue_name.clone())
                .append(" AND status = ")
                .bind(status),
        )
        .await?;
        postgres_count_to_usize(count)
    }
}

#[async_trait]
impl FlowTaskQueue for PostgresFlowTaskQueue {
    async fn enqueue(&self, task: FlowTask) -> Result<()> {
        let now = timestamp_nanos_saturating(Utc::now());
        execute_query(
            &self.executor,
            sql_query::<()>(
                "INSERT INTO flow_tasks (queue_name, task_id, task_json, status, \
                 enqueued_at_nanos, updated_at_nanos) VALUES (",
            )
            .bind(self.queue_name.clone())
            .append(", ")
            .bind(Uuid::new_v4().to_string())
            .append(", ")
            .bind(serde_json::to_string(&task)?)
            .append(", 'pending', ")
            .bind(now)
            .append(", ")
            .bind(now)
            .append(")"),
        )
        .await?;
        Ok(())
    }

    async fn lease(&self) -> Result<Option<FlowTaskLease>> {
        let lease_id = Uuid::new_v4().to_string();
        let now = timestamp_nanos_saturating(Utc::now());
        let row = fetch_optional_query(
            &self.executor,
            sql_query::<(String, String)>(
                "WITH next_task AS (SELECT task_id FROM flow_tasks \
                 WHERE queue_name = ",
            )
            .bind(self.queue_name.clone())
            .append(
                " AND status = 'pending' ORDER BY enqueued_at_nanos ASC, task_id ASC \
                 FOR UPDATE SKIP LOCKED LIMIT 1) UPDATE flow_tasks \
                 SET status = 'inflight', lease_id = ",
            )
            .bind(lease_id)
            .append(", leased_at_nanos = ")
            .bind(now)
            .append(", updated_at_nanos = ")
            .bind(now)
            .append(" FROM next_task WHERE flow_tasks.queue_name = ")
            .bind(self.queue_name.clone())
            .append(
                " AND flow_tasks.task_id = next_task.task_id \
                 RETURNING flow_tasks.lease_id, flow_tasks.task_json",
            ),
        )
        .await?;
        row.map(|(lease_id, task_json)| {
            Ok(FlowTaskLease {
                lease_id,
                task: serde_json::from_str(&task_json)?,
            })
        })
        .transpose()
    }

    async fn heartbeat(&self, lease_id: &str) -> Result<String> {
        let renewed_lease_id = Uuid::new_v4().to_string();
        let now = timestamp_nanos_saturating(Utc::now());
        let rows = execute_query(
            &self.executor,
            sql_query::<()>("UPDATE flow_tasks SET lease_id = ")
                .bind(renewed_lease_id.clone())
                .append(", leased_at_nanos = ")
                .bind(now)
                .append(", updated_at_nanos = ")
                .bind(now)
                .append(" WHERE queue_name = ")
                .bind(self.queue_name.clone())
                .append(" AND status = 'inflight' AND lease_id = ")
                .bind(lease_id),
        )
        .await?;
        if rows == 0 {
            return Err(FlowError::LeaseLost(lease_id.to_string()));
        }
        Ok(renewed_lease_id)
    }

    async fn ack(&self, lease_id: &str) -> Result<()> {
        let rows = execute_query(
            &self.executor,
            sql_query::<()>("DELETE FROM flow_tasks WHERE queue_name = ")
                .bind(self.queue_name.clone())
                .append(" AND status = 'inflight' AND lease_id = ")
                .bind(lease_id),
        )
        .await?;
        if rows == 0 {
            return Err(FlowError::LeaseLost(lease_id.to_string()));
        }
        Ok(())
    }

    async fn requeue_inflight(&self) -> Result<usize> {
        let rows = execute_query(
            &self.executor,
            sql_query::<()>(
                "UPDATE flow_tasks SET status = 'pending', lease_id = NULL, \
                 leased_at_nanos = NULL, updated_at_nanos = ",
            )
            .bind(timestamp_nanos_saturating(Utc::now()))
            .append(" WHERE queue_name = ")
            .bind(self.queue_name.clone())
            .append(" AND status = 'inflight'"),
        )
        .await?;
        postgres_rows_affected_to_usize(rows)
    }

    async fn len(&self) -> Result<usize> {
        self.count_by_status("pending").await
    }
}

async fn execute_query<E>(executor: &E, query: SqlQuery<()>) -> Result<u64>
where
    E: Executor<Row = PostgresRow, Error = PostgresError>,
{
    let query = query
        .compile(&PostgresDialect)
        .map_err(postgres_queue_query_error)?;
    Ok(executor
        .execute(&query)
        .await
        .map_err(postgres_queue_driver_error)?
        .rows_affected)
}

async fn fetch_all_query<T, E>(executor: &E, query: SqlQuery<T>) -> Result<Vec<T>>
where
    T: FromRow + Send,
    E: Executor<Row = PostgresRow, Error = PostgresError>,
{
    let query = query
        .compile(&PostgresDialect)
        .map_err(postgres_queue_query_error)?;
    executor
        .fetch_all(&query)
        .await
        .map_err(postgres_queue_driver_error)?
        .rows
        .iter()
        .map(T::from_row)
        .collect::<std::result::Result<Vec<_>, _>>()
        .map_err(postgres_queue_decode_error)
}

async fn fetch_optional_query<T, E>(executor: &E, query: SqlQuery<T>) -> Result<Option<T>>
where
    T: FromRow + Send,
    E: Executor<Row = PostgresRow, Error = PostgresError>,
{
    let mut rows = fetch_all_query(executor, query).await?;
    match rows.len() {
        0 => Ok(None),
        1 => Ok(rows.pop()),
        actual => Err(FlowError::Store(format!(
            "PostgreSQL Flow query returned {actual} rows where at most one was expected"
        ))),
    }
}

async fn fetch_one_query<T, E>(executor: &E, query: SqlQuery<T>) -> Result<T>
where
    T: FromRow + Send,
    E: Executor<Row = PostgresRow, Error = PostgresError>,
{
    fetch_optional_query(executor, query)
        .await?
        .ok_or_else(|| FlowError::Store("PostgreSQL Flow query returned no rows".to_string()))
}

fn dead_letter_row(
    (lease_id, task_json, reason, dead_lettered_at): (String, String, String, i64),
) -> Result<PostgresDeadLetteredTask> {
    Ok(PostgresDeadLetteredTask {
        lease_id,
        task: serde_json::from_str(&task_json)?,
        reason,
        dead_lettered_at: nanos_to_datetime(dead_lettered_at)?,
    })
}

fn map_postgres_queue_transaction<T>(
    result: std::result::Result<T, PostgresTransactionError<FlowError>>,
) -> Result<T> {
    match result {
        Ok(value) => Ok(value),
        Err(PostgresTransactionError::Operation(error)) => Err(error),
        Err(error) => Err(FlowError::Store(format!(
            "PostgreSQL Flow task transaction failed: {error}"
        ))),
    }
}

fn postgres_count_to_usize(count: i64) -> Result<usize> {
    usize::try_from(count).map_err(|error| {
        FlowError::Store(format!(
            "invalid PostgreSQL Flow task count {count}: {error}"
        ))
    })
}

fn postgres_rows_affected_to_usize(rows: u64) -> Result<usize> {
    usize::try_from(rows).map_err(|error| {
        FlowError::Store(format!(
            "PostgreSQL Flow affected row count {rows} exceeds usize range: {error}"
        ))
    })
}

fn validated_queue_name(queue_name: &str) -> Result<String> {
    let queue_name = queue_name.trim();
    if queue_name.is_empty() {
        return Err(FlowError::Store(
            "PostgreSQL task queue name cannot be empty".to_string(),
        ));
    }
    Ok(queue_name.to_string())
}

fn nanos_to_datetime(nanos: i64) -> Result<DateTime<Utc>> {
    let seconds = nanos.div_euclid(1_000_000_000);
    let subsecond_nanos = nanos.rem_euclid(1_000_000_000) as u32;
    DateTime::from_timestamp(seconds, subsecond_nanos)
        .ok_or_else(|| FlowError::Store(format!("invalid PostgreSQL Flow task timestamp {nanos}")))
}

fn postgres_queue_query_error(error: a3s_orm::Error) -> FlowError {
    FlowError::Store(format!("PostgreSQL Flow task query build failed: {error}"))
}

fn postgres_queue_driver_error(error: PostgresError) -> FlowError {
    FlowError::Store(format!("PostgreSQL Flow task storage failed: {error}"))
}

fn postgres_queue_decode_error(error: a3s_orm::DecodeError) -> FlowError {
    FlowError::Store(format!("PostgreSQL Flow task row decoding failed: {error}"))
}