later 0.0.49

use super::{
    JobIndexPage, JobIndexRow, JobIndexWaitStats, RangeItem, RangeOrder, RangePage, Storage,
    StorageOperation,
};
use crate::UtcDateTime;
use sqlx::{postgres::PgPoolOptions, PgPool, Row};

const METRIC_BUCKET_MS: i64 = 60_000;

fn minute_bucket(at: UtcDateTime) -> i64 {
    at.timestamp_millis().div_euclid(METRIC_BUCKET_MS)
}

fn row_to_job_index(row: &sqlx::postgres::PgRow) -> Result<JobIndexRow, sqlx::Error> {
    Ok(JobIndexRow {
        job_id: row.try_get("job_id")?,
        payload_type: row.try_get("payload_type")?,
        stage: row.try_get("stage")?,
        stage_date: row.try_get("stage_date")?,
        revision: row.try_get("revision")?,
        created_at: row.try_get("created_at")?,
        wait_ms: row.try_get("wait_ms")?,
        wait_mode: row.try_get("wait_mode")?,
        topic: row.try_get("topic")?,
        partition: row.try_get("partition_id")?,
        sequence: row.try_get("sequence")?,
        parent_job_id: row.try_get("parent_job_id")?,
        date_expire: row.try_get("date_expire")?,
    })
}

#[derive(Clone)]
/// Durable job state stored in PostgreSQL.
///
/// Clones share the same SQLx pool. Later applies its embedded migrations when
/// this type is constructed.
pub struct Postgres {
    pool: PgPool,
}

/// The table Later's own migration history lives in - deliberately not
/// sqlx's default `_sqlx_migrations`. See the identical constant/function in
/// `storage::sqlite` for the full reasoning: `Postgres::from_pool` is meant
/// to be usable against a database an application already owns and
/// migrates with sqlx itself, and sqlx's migrator assumes exclusive
/// ownership of whatever table it's pointed at - there is no supported way
/// to give it a different table name instead. Tracking Later's own history
/// separately here avoids that collision entirely.
const MIGRATIONS_TABLE: &str = "later_schema_migrations";

/// Applies every `.sql` file under `./migrations` that isn't already
/// recorded in [`MIGRATIONS_TABLE`], each in its own transaction alongside
/// the row that records it - see `storage::sqlite::run_migrations`, which
/// this mirrors.
async fn run_migrations(pool: &PgPool) -> anyhow::Result<()> {
    sqlx::query(&format!(
        "CREATE TABLE IF NOT EXISTS {MIGRATIONS_TABLE} ( \
            version BIGINT PRIMARY KEY, \
            description TEXT NOT NULL, \
            checksum BYTEA NOT NULL, \
            applied_at TIMESTAMPTZ NOT NULL DEFAULT now() \
        )"
    ))
    .execute(pool)
    .await?;

    let migrator = sqlx::migrate!("./migrations");
    for migration in migrator.iter() {
        if migration.migration_type.is_down_migration() {
            continue;
        }

        let existing: Option<Vec<u8>> = sqlx::query_scalar(&format!(
            "SELECT checksum FROM {MIGRATIONS_TABLE} WHERE version = $1"
        ))
        .bind(migration.version)
        .fetch_optional(pool)
        .await?;

        match existing {
            Some(checksum) if checksum == migration.checksum.as_ref() => continue,
            Some(_) => {
                anyhow::bail!(
                    "Later migration {} ({}) has changed since it was applied - \
                     migrations must never be edited after release",
                    migration.version,
                    migration.description,
                );
            }
            None => {}
        }

        let mut tx = pool.begin().await?;
        sqlx::raw_sql(&migration.sql).execute(&mut *tx).await?;
        sqlx::query(&format!(
            "INSERT INTO {MIGRATIONS_TABLE} (version, description, checksum) VALUES ($1, $2, $3)"
        ))
        .bind(migration.version)
        .bind(migration.description.as_ref())
        .bind(migration.checksum.as_ref())
        .execute(&mut *tx)
        .await?;
        tx.commit().await?;
    }

    Ok(())
}

impl Postgres {
    /// Connects with a new pool and applies Later's PostgreSQL migrations.
    pub async fn new(connection_string: &str) -> anyhow::Result<Self> {
        let pool = PgPoolOptions::new()
            .max_connections(20)
            .connect(connection_string)
            .await?;
        run_migrations(&pool).await?;
        Ok(Self { pool })
    }

    /// Uses an application-owned pool after applying Later's migrations.
    ///
    /// Safe to call against a database that already has its own, unrelated
    /// migration history - see [`MIGRATIONS_TABLE`].
    pub async fn from_pool(pool: PgPool) -> anyhow::Result<Self> {
        run_migrations(&pool).await?;
        Ok(Self { pool })
    }

    /// Returns the SQLx pool used by this storage instance.
    pub fn pool(&self) -> &PgPool {
        &self.pool
    }
}

/// The job-index upsert, usable on the pool or inside a transaction.
async fn upsert_job_index<'e, E>(
    executor: E,
    namespace: &str,
    row: &JobIndexRow,
) -> anyhow::Result<()>
where
    E: sqlx::Executor<'e, Database = sqlx::Postgres>,
{
    sqlx::query(
        "INSERT INTO later_jobs_index \
               (namespace, job_id, payload_type, stage, stage_date, revision, created_at, \
                wait_ms, wait_mode, topic, partition_id, sequence, parent_job_id, date_expire) \
             VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14) \
             ON CONFLICT (namespace, job_id) DO UPDATE SET \
               payload_type = excluded.payload_type, \
               stage = excluded.stage, \
               stage_date = excluded.stage_date, \
               revision = excluded.revision, \
               wait_ms = COALESCE(excluded.wait_ms, later_jobs_index.wait_ms), \
               wait_mode = COALESCE(excluded.wait_mode, later_jobs_index.wait_mode), \
               topic = excluded.topic, \
               partition_id = excluded.partition_id, \
               sequence = excluded.sequence, \
               parent_job_id = COALESCE(excluded.parent_job_id, later_jobs_index.parent_job_id), \
               date_expire = excluded.date_expire \
             WHERE excluded.revision >= later_jobs_index.revision",
    )
    .bind(namespace)
    .bind(&row.job_id)
    .bind(&row.payload_type)
    .bind(&row.stage)
    .bind(row.stage_date)
    .bind(row.revision)
    .bind(row.created_at)
    .bind(row.wait_ms)
    .bind(&row.wait_mode)
    .bind(&row.topic)
    .bind(row.partition)
    .bind(row.sequence)
    .bind(&row.parent_job_id)
    .bind(row.date_expire)
    .execute(executor)
    .await?;
    Ok(())
}

#[async_trait::async_trait]
impl Storage for Postgres {
    async fn get(&self, key: &str) -> anyhow::Result<Option<Vec<u8>>> {
        sqlx::query(
            "SELECT value FROM later_storage \
             WHERE key = $1 AND (date_expire IS NULL OR date_expire > NOW())",
        )
        .bind(key)
        .fetch_optional(&self.pool)
        .await?
        .map(|row| row.try_get("value"))
        .transpose()
        .map_err(anyhow::Error::from)
    }

    async fn apply(&self, operations: Vec<StorageOperation>) -> anyhow::Result<()> {
        let mut transaction = self.pool.begin().await?;
        crate::backend::apply_postgres_operations(&mut transaction, operations).await?;
        transaction.commit().await?;
        Ok(())
    }

    async fn set_if_absent(&self, key: &str, value: &[u8]) -> anyhow::Result<bool> {
        let row: Option<(i32,)> = sqlx::query_as(
            "INSERT INTO later_storage (key, value, update_count, date_created) \
             VALUES ($1, $2, 0, NOW()) \
             ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value, \
               date_updated = NOW(), update_count = later_storage.update_count + 1, \
               date_expire = NULL \
             WHERE later_storage.date_expire IS NOT NULL AND later_storage.date_expire <= NOW() \
             RETURNING 1",
        )
        .bind(key)
        .bind(value)
        .fetch_optional(&self.pool)
        .await?;
        Ok(row.is_some())
    }

    async fn range_page(
        &self,
        key: &str,
        cursor: Option<i64>,
        limit: usize,
        order: RangeOrder,
    ) -> anyhow::Result<RangePage> {
        let limit = i64::try_from(limit)?;
        let rows = match (order, cursor) {
            (RangeOrder::OldestFirst, Some(cursor)) => {
                sqlx::query(
                    "SELECT sequence, value FROM later_storage_range \
                     WHERE range_key = $1 AND sequence > $2 \
                       AND (date_expire IS NULL OR date_expire > NOW()) \
                     ORDER BY sequence ASC LIMIT $3",
                )
                .bind(key)
                .bind(cursor)
                .bind(limit)
                .fetch_all(&self.pool)
                .await?
            }
            (RangeOrder::OldestFirst, None) => {
                sqlx::query(
                    "SELECT sequence, value FROM later_storage_range \
                     WHERE range_key = $1 AND (date_expire IS NULL OR date_expire > NOW()) \
                     ORDER BY sequence ASC LIMIT $2",
                )
                .bind(key)
                .bind(limit)
                .fetch_all(&self.pool)
                .await?
            }
            (RangeOrder::NewestFirst, Some(cursor)) => {
                sqlx::query(
                    "SELECT sequence, value FROM later_storage_range \
                     WHERE range_key = $1 AND sequence < $2 \
                       AND (date_expire IS NULL OR date_expire > NOW()) \
                     ORDER BY sequence DESC LIMIT $3",
                )
                .bind(key)
                .bind(cursor)
                .bind(limit)
                .fetch_all(&self.pool)
                .await?
            }
            (RangeOrder::NewestFirst, None) => {
                sqlx::query(
                    "SELECT sequence, value FROM later_storage_range \
                     WHERE range_key = $1 AND (date_expire IS NULL OR date_expire > NOW()) \
                     ORDER BY sequence DESC LIMIT $2",
                )
                .bind(key)
                .bind(limit)
                .fetch_all(&self.pool)
                .await?
            }
        };
        let items = rows
            .into_iter()
            .map(|row| {
                Ok(RangeItem {
                    cursor: row.try_get("sequence")?,
                    value: row.try_get("value")?,
                })
            })
            .collect::<Result<Vec<_>, sqlx::Error>>()?;
        let next_cursor = (items.len() == usize::try_from(limit)?)
            .then(|| items.last().map(|item| item.cursor))
            .flatten();
        Ok(RangePage { items, next_cursor })
    }

    async fn range_count(&self, key: &str) -> anyhow::Result<usize> {
        // A plain read - see the identical comment and rationale on
        // `storage::sqlite::Sqlite::range_count` (this used to also run an
        // opportunistic bounded delete here). `sweep_expired` (run
        // periodically via `PollExpiredStorage`) is solely responsible for
        // physically deleting expired rows now.
        let count = sqlx::query(
            "SELECT COUNT(*) AS item_count FROM later_storage_range \
             WHERE range_key = $1 AND (date_expire IS NULL OR date_expire > NOW())",
        )
        .bind(key)
        .fetch_one(&self.pool)
        .await?
        .try_get::<i64, _>("item_count")?;
        Ok(usize::try_from(count)?)
    }

    async fn total_db_size_bytes(&self) -> anyhow::Result<Option<u64>> {
        let size: i64 = sqlx::query_scalar("SELECT pg_database_size(current_database())")
            .fetch_one(&self.pool)
            .await?;
        Ok(Some(u64::try_from(size)?))
    }

    async fn sweep_expired(&self, limit: usize) -> anyhow::Result<usize> {
        let limit = i64::try_from(limit).unwrap_or(i64::MAX);
        let storage_removed = sqlx::query(
            "DELETE FROM later_storage WHERE key IN ( \
               SELECT key FROM later_storage \
               WHERE date_expire IS NOT NULL AND date_expire <= NOW() LIMIT $1 \
             )",
        )
        .bind(limit)
        .execute(&self.pool)
        .await?
        .rows_affected();
        // Shares `limit`'s remaining budget with the table above rather
        // than getting its own full `limit` - the contract is "at most
        // `limit` rows total per call", not per table.
        let remaining = limit.saturating_sub(i64::try_from(storage_removed).unwrap_or(i64::MAX));
        let range_removed = if remaining > 0 {
            sqlx::query(
                "DELETE FROM later_storage_range WHERE sequence IN ( \
                   SELECT sequence FROM later_storage_range \
                   WHERE date_expire IS NOT NULL AND date_expire <= NOW() LIMIT $1 \
                 )",
            )
            .bind(remaining)
            .execute(&self.pool)
            .await?
            .rows_affected()
        } else {
            0
        };
        Ok(usize::try_from(storage_removed + range_removed)?)
    }

    async fn job_index_upsert(&self, namespace: &str, row: JobIndexRow) -> anyhow::Result<()> {
        upsert_job_index(&self.pool, namespace, &row).await
    }

    async fn apply_with_job_index(
        &self,
        operations: Vec<StorageOperation>,
        namespace: &str,
        change: crate::storage::JobIndexChange,
    ) -> anyhow::Result<()> {
        // One transaction instead of two: the job write and its dashboard
        // row commit together, saving a commit per job transition.
        let mut transaction = self.pool.begin().await?;
        crate::backend::apply_postgres_operations(&mut transaction, operations).await?;
        match change {
            crate::storage::JobIndexChange::Upsert(row) => {
                upsert_job_index(&mut *transaction, namespace, &row).await?
            }
            crate::storage::JobIndexChange::Remove(job_id) => {
                sqlx::query("DELETE FROM later_jobs_index WHERE namespace = $1 AND job_id = $2")
                    .bind(namespace)
                    .bind(job_id)
                    .execute(&mut *transaction)
                    .await?;
            }
        }
        transaction.commit().await?;
        Ok(())
    }

    async fn job_index_remove(&self, namespace: &str, job_id: &str) -> anyhow::Result<()> {
        sqlx::query("DELETE FROM later_jobs_index WHERE namespace = $1 AND job_id = $2")
            .bind(namespace)
            .bind(job_id)
            .execute(&self.pool)
            .await?;
        Ok(())
    }

    async fn queue_sample_record(
        &self,
        namespace: &str,
        at: UtcDateTime,
        queued: i64,
        partitioned: i64,
    ) -> anyhow::Result<()> {
        let bucket = minute_bucket(at);
        let mut tx = self.pool.begin().await?;
        sqlx::query(
            "INSERT INTO later_queue_samples (namespace, minute_bucket, queued, partitioned) \
             VALUES ($1, $2, $3, $4) \
             ON CONFLICT (namespace, minute_bucket) \
               DO UPDATE SET queued = excluded.queued, partitioned = excluded.partitioned",
        )
        .bind(namespace)
        .bind(bucket)
        .bind(queued)
        .bind(partitioned)
        .execute(&mut *tx)
        .await?;
        sqlx::query("DELETE FROM later_queue_samples WHERE namespace = $1 AND minute_bucket < $2")
            .bind(namespace)
            .bind(bucket - 7 * 24 * 60)
            .execute(&mut *tx)
            .await?;
        tx.commit().await?;
        Ok(())
    }

    async fn queue_samples_since(
        &self,
        namespace: &str,
        since: UtcDateTime,
    ) -> anyhow::Result<Vec<crate::storage::QueueSample>> {
        use chrono::TimeZone;
        let rows = sqlx::query(
            "SELECT minute_bucket, queued, partitioned FROM later_queue_samples \
             WHERE namespace = $1 AND minute_bucket >= $2 ORDER BY minute_bucket",
        )
        .bind(namespace)
        .bind(minute_bucket(since))
        .fetch_all(&self.pool)
        .await?;
        rows.iter()
            .map(|row| {
                let bucket: i64 = row.try_get("minute_bucket")?;
                Ok(crate::storage::QueueSample {
                    at: chrono::Utc
                        .timestamp_millis_opt(bucket * METRIC_BUCKET_MS)
                        .single()
                        .ok_or_else(|| anyhow::anyhow!("queue sample time out of range"))?,
                    queued: row.try_get("queued")?,
                    partitioned: row.try_get("partitioned")?,
                })
            })
            .collect()
    }

    async fn purge_unused_dashboard_rows(
        &self,
        namespace: &str,
        limit: usize,
    ) -> anyhow::Result<usize> {
        let mut budget = i64::try_from(limit)?;
        let mut removed = 0usize;
        // Keys earlier dashboards kept in `later_storage_range` that nothing
        // reads any more; the delayed / requeued polling keys, the worker
        // heartbeat and continuation keys are live and not listed.
        for suffix in [
            "all-jobs",
            "stage-enqueued-jobs",
            "stage-running-jobs",
            "stage-waiting-jobs",
        ] {
            if budget == 0 {
                break;
            }
            let done = sqlx::query(
                "DELETE FROM later_storage_range WHERE sequence IN ( \
                   SELECT sequence FROM later_storage_range WHERE range_key = $1 LIMIT $2)",
            )
            .bind(format!("{namespace}-{suffix}"))
            .bind(budget)
            .execute(&self.pool)
            .await?
            .rows_affected();
            budget -= i64::try_from(done)?;
            removed += usize::try_from(done)?;
        }
        for prefix in ["stats-job-in-", "stats-partition-jobs-"] {
            if budget == 0 {
                break;
            }
            let from = format!("{namespace}-{prefix}");
            let mut to = from.clone();
            to.pop();
            to.push('.');
            let done = sqlx::query(
                "DELETE FROM later_storage_range WHERE sequence IN ( \
                   SELECT sequence FROM later_storage_range \
                   WHERE range_key >= $1 AND range_key < $2 LIMIT $3)",
            )
            .bind(from)
            .bind(to)
            .bind(budget)
            .execute(&self.pool)
            .await?
            .rows_affected();
            budget -= i64::try_from(done)?;
            removed += usize::try_from(done)?;
        }
        Ok(removed)
    }

    async fn job_index_stage_counts(
        &self,
        namespace: &str,
    ) -> anyhow::Result<std::collections::HashMap<String, usize>> {
        // Live stages are counter reads (see `later_stage_counts`); finished
        // stages are bounded by retention and keep the expiry filter.
        let rows = sqlx::query(
            "SELECT stage, count AS item_count FROM later_stage_counts \
             WHERE namespace = $1 AND count > 0 \
               AND stage IN ('delayed', 'waiting', 'enqueued', 'running', 'requeued') \
             UNION ALL \
             SELECT stage, COUNT(*) AS item_count FROM later_jobs_index \
             WHERE namespace = $1 AND stage IN ('success', 'failed') \
               AND (date_expire IS NULL OR date_expire > NOW()) \
             GROUP BY stage",
        )
        .bind(namespace)
        .fetch_all(&self.pool)
        .await?;
        rows.into_iter()
            .map(|row| {
                let stage: String = row.try_get("stage")?;
                let count: i64 = row.try_get("item_count")?;
                Ok((stage, usize::try_from(count)?))
            })
            .collect()
    }

    async fn job_index_reconcile_stage_counts(&self, namespace: &str) -> anyhow::Result<()> {
        // One transaction, so the recount and the replacement see the same
        // rows and no writer can slip a change between them.
        let mut tx = self.pool.begin().await?;
        sqlx::query("DELETE FROM later_stage_counts WHERE namespace = $1")
            .bind(namespace)
            .execute(&mut *tx)
            .await?;
        sqlx::query(
            "INSERT INTO later_stage_counts (namespace, stage, count) \
             SELECT namespace, stage, COUNT(*) FROM later_jobs_index \
             WHERE namespace = $1 \
               AND stage IN ('delayed', 'waiting', 'enqueued', 'running', 'requeued') \
             GROUP BY namespace, stage",
        )
        .bind(namespace)
        .execute(&mut *tx)
        .await?;
        tx.commit().await?;
        Ok(())
    }

    async fn job_index_list_by_stage(
        &self,
        namespace: &str,
        stage: &str,
        cursor: Option<i64>,
        limit: usize,
    ) -> anyhow::Result<JobIndexPage> {
        use chrono::TimeZone;
        let limit_i64 = i64::try_from(limit)?;
        let rows = match cursor {
            Some(cursor) => {
                let cursor = chrono::Utc.timestamp_millis_opt(cursor).unwrap();
                sqlx::query(
                    "SELECT * FROM later_jobs_index \
                     WHERE namespace = $1 AND stage = $2 AND stage_date < $3 \
                       AND (date_expire IS NULL OR date_expire > NOW()) \
                     ORDER BY stage_date DESC, job_id DESC LIMIT $4",
                )
                .bind(namespace)
                .bind(stage)
                .bind(cursor)
                .bind(limit_i64)
                .fetch_all(&self.pool)
                .await?
            }
            None => {
                sqlx::query(
                    "SELECT * FROM later_jobs_index \
                     WHERE namespace = $1 AND stage = $2 \
                       AND (date_expire IS NULL OR date_expire > NOW()) \
                     ORDER BY stage_date DESC, job_id DESC LIMIT $3",
                )
                .bind(namespace)
                .bind(stage)
                .bind(limit_i64)
                .fetch_all(&self.pool)
                .await?
            }
        };
        let items = rows
            .iter()
            .map(row_to_job_index)
            .collect::<Result<Vec<_>, _>>()?;
        let next_cursor = (items.len() == limit)
            .then(|| items.last().map(|item| item.stage_date.timestamp_millis()))
            .flatten();
        Ok(JobIndexPage { items, next_cursor })
    }

    async fn job_index_list_by_partition(
        &self,
        namespace: &str,
        topic: &str,
        partition: u32,
        cursor: Option<i64>,
        limit: usize,
    ) -> anyhow::Result<JobIndexPage> {
        let limit_i64 = i64::try_from(limit)?;
        let partition_id = i64::from(partition);
        let rows = match cursor {
            Some(cursor) => {
                sqlx::query(
                    "SELECT * FROM later_jobs_index \
                     WHERE namespace = $1 AND topic = $2 AND partition_id = $3 \
                       AND sequence < $4 \
                     ORDER BY sequence DESC LIMIT $5",
                )
                .bind(namespace)
                .bind(topic)
                .bind(partition_id)
                .bind(cursor)
                .bind(limit_i64)
                .fetch_all(&self.pool)
                .await?
            }
            None => {
                sqlx::query(
                    "SELECT * FROM later_jobs_index \
                     WHERE namespace = $1 AND topic = $2 AND partition_id = $3 \
                     ORDER BY sequence DESC LIMIT $4",
                )
                .bind(namespace)
                .bind(topic)
                .bind(partition_id)
                .bind(limit_i64)
                .fetch_all(&self.pool)
                .await?
            }
        };
        let items = rows
            .iter()
            .map(row_to_job_index)
            .collect::<Result<Vec<_>, _>>()?;
        let next_cursor = (items.len() == limit)
            .then(|| items.last().and_then(|item| item.sequence))
            .flatten();
        Ok(JobIndexPage { items, next_cursor })
    }

    async fn job_index_partition_neighbor(
        &self,
        namespace: &str,
        topic: &str,
        partition: u32,
        sequence: i64,
        older: bool,
    ) -> anyhow::Result<Option<JobIndexRow>> {
        let target = if older { sequence - 1 } else { sequence + 1 };
        let row = sqlx::query(
            "SELECT * FROM later_jobs_index \
             WHERE namespace = $1 AND topic = $2 AND partition_id = $3 AND sequence = $4",
        )
        .bind(namespace)
        .bind(topic)
        .bind(i64::from(partition))
        .bind(target)
        .fetch_optional(&self.pool)
        .await?;
        Ok(row.as_ref().map(row_to_job_index).transpose()?)
    }

    async fn job_index_list_continuations(
        &self,
        namespace: &str,
        parent_job_id: &str,
        limit: usize,
    ) -> anyhow::Result<Vec<JobIndexRow>> {
        let rows = sqlx::query(
            "SELECT * FROM later_jobs_index \
             WHERE namespace = $1 AND parent_job_id = $2 \
             ORDER BY stage_date ASC LIMIT $3",
        )
        .bind(namespace)
        .bind(parent_job_id)
        .bind(i64::try_from(limit)?)
        .fetch_all(&self.pool)
        .await?;
        rows.iter()
            .map(row_to_job_index)
            .collect::<Result<_, _>>()
            .map_err(anyhow::Error::from)
    }

    async fn job_index_truncate_partition(
        &self,
        namespace: &str,
        topic: &str,
        partition: u32,
        keep_last: usize,
    ) -> anyhow::Result<usize> {
        let partition_id = i64::from(partition);
        let keep_last = i64::try_from(keep_last)?;
        let result = sqlx::query(
            "DELETE FROM later_jobs_index \
             WHERE namespace = $1 AND topic = $2 AND partition_id = $3 \
               AND job_id NOT IN ( \
                 SELECT job_id FROM later_jobs_index \
                 WHERE namespace = $1 AND topic = $2 AND partition_id = $3 \
                 ORDER BY sequence DESC LIMIT $4 \
               )",
        )
        .bind(namespace)
        .bind(topic)
        .bind(partition_id)
        .bind(keep_last)
        .execute(&self.pool)
        .await?;
        Ok(usize::try_from(result.rows_affected())?)
    }

    async fn job_index_record_transition(
        &self,
        namespace: &str,
        stage: &str,
        at: UtcDateTime,
    ) -> anyhow::Result<()> {
        sqlx::query(
            "INSERT INTO later_job_transition_metrics (namespace, stage, minute_bucket, count) \
             VALUES ($1, $2, $3, 1) \
             ON CONFLICT (namespace, stage, minute_bucket) \
               DO UPDATE SET count = later_job_transition_metrics.count + 1",
        )
        .bind(namespace)
        .bind(stage)
        .bind(minute_bucket(at))
        .execute(&self.pool)
        .await?;
        Ok(())
    }

    async fn job_index_record_metrics_batch(
        &self,
        namespace: &str,
        batch: &crate::storage::JobIndexMetricsBatch,
    ) -> anyhow::Result<()> {
        let mut tx = self.pool.begin().await?;
        for (stage, at, count) in &batch.transitions {
            sqlx::query(
                "INSERT INTO later_job_transition_metrics (namespace, stage, minute_bucket, count) \
                 VALUES ($1, $2, $3, $4) \
                 ON CONFLICT (namespace, stage, minute_bucket) \
                   DO UPDATE SET count = later_job_transition_metrics.count + $4",
            )
            .bind(namespace)
            .bind(stage)
            .bind(minute_bucket(*at))
            .bind(count)
            .execute(&mut *tx)
            .await?;
        }
        for (mode, at, count, sum_ms, max_ms) in &batch.waits {
            sqlx::query(
                "INSERT INTO later_job_wait_metrics \
                   (namespace, mode, minute_bucket, count, sum_ms, max_ms) \
                 VALUES ($1, $2, $3, $4, $5, $6) \
                 ON CONFLICT (namespace, mode, minute_bucket) DO UPDATE SET \
                   count = later_job_wait_metrics.count + $4, \
                   sum_ms = later_job_wait_metrics.sum_ms + $5, \
                   max_ms = GREATEST(later_job_wait_metrics.max_ms, $6)",
            )
            .bind(namespace)
            .bind(mode)
            .bind(minute_bucket(*at))
            .bind(count)
            .bind(sum_ms)
            .bind(max_ms)
            .execute(&mut *tx)
            .await?;
        }
        tx.commit().await?;
        Ok(())
    }

    async fn job_index_recent_transition_count(
        &self,
        namespace: &str,
        stage: &str,
        now: UtcDateTime,
    ) -> anyhow::Result<usize> {
        let count: i64 = sqlx::query_scalar(
            "SELECT COALESCE(SUM(count), 0) FROM later_job_transition_metrics \
             WHERE namespace = $1 AND stage = $2 AND minute_bucket >= $3",
        )
        .bind(namespace)
        .bind(stage)
        .bind(minute_bucket(now) - 1)
        .fetch_one(&self.pool)
        .await?;
        Ok(usize::try_from(count)?)
    }

    async fn job_index_record_wait(
        &self,
        namespace: &str,
        mode: &str,
        wait_ms: i64,
        at: UtcDateTime,
    ) -> anyhow::Result<()> {
        sqlx::query(
            "INSERT INTO later_job_wait_metrics \
               (namespace, mode, minute_bucket, count, sum_ms, max_ms) \
             VALUES ($1, $2, $3, 1, $4, $4) \
             ON CONFLICT (namespace, mode, minute_bucket) DO UPDATE SET \
               count = later_job_wait_metrics.count + 1, \
               sum_ms = later_job_wait_metrics.sum_ms + $4, \
               max_ms = GREATEST(later_job_wait_metrics.max_ms, $4)",
        )
        .bind(namespace)
        .bind(mode)
        .bind(minute_bucket(at))
        .bind(wait_ms)
        .execute(&self.pool)
        .await?;
        Ok(())
    }

    async fn job_index_recent_wait_stats(
        &self,
        namespace: &str,
        mode: &str,
        now: UtcDateTime,
    ) -> anyhow::Result<JobIndexWaitStats> {
        let row = sqlx::query(
            "SELECT COALESCE(SUM(count), 0) AS c, COALESCE(SUM(sum_ms), 0) AS s, \
                    COALESCE(MAX(max_ms), 0) AS m \
             FROM later_job_wait_metrics \
             WHERE namespace = $1 AND mode = $2 AND minute_bucket >= $3",
        )
        .bind(namespace)
        .bind(mode)
        .bind(minute_bucket(now) - 1)
        .fetch_one(&self.pool)
        .await?;
        let count: i64 = row.try_get("c")?;
        let sum_ms: i64 = row.try_get("s")?;
        let max_ms: i64 = row.try_get("m")?;
        let count = usize::try_from(count)?;
        Ok(JobIndexWaitStats {
            count,
            avg_ms: if count == 0 {
                0
            } else {
                sum_ms / i64::try_from(count)?
            },
            max_ms,
        })
    }

    async fn job_index_sweep_expired(
        &self,
        namespace: &str,
        now: UtcDateTime,
        limit: usize,
    ) -> anyhow::Result<usize> {
        let limit_i64 = i64::try_from(limit)?;
        let result = sqlx::query(
            "DELETE FROM later_jobs_index WHERE job_id IN ( \
               SELECT job_id FROM later_jobs_index \
               WHERE namespace = $1 AND date_expire IS NOT NULL AND date_expire <= NOW() \
               LIMIT $2 \
             ) AND namespace = $1",
        )
        .bind(namespace)
        .bind(limit_i64)
        .execute(&self.pool)
        .await?;
        let stale_bucket = minute_bucket(now) - 5;
        sqlx::query(
            "DELETE FROM later_job_transition_metrics \
             WHERE namespace = $1 AND minute_bucket < $2",
        )
        .bind(namespace)
        .bind(stale_bucket)
        .execute(&self.pool)
        .await?;
        sqlx::query(
            "DELETE FROM later_job_wait_metrics WHERE namespace = $1 AND minute_bucket < $2",
        )
        .bind(namespace)
        .bind(stale_bucket)
        .execute(&self.pool)
        .await?;
        Ok(usize::try_from(result.rows_affected())?)
    }
}