prop/rust/prop_grid_rs/src/db.rs

//! `grid_tasks` claim/complete + `propagation_ready` NOTIFY.
//!
//! Contract:
//!   * Elixir seeds one row per `(run_time, forecast_hour)` with status=queued.
//!   * Rust claims one queued, non-zero-forecast-hour row via
//!     `SELECT ... FOR UPDATE SKIP LOCKED` and marks it running.
//!   * On success: status=done, NOTIFY propagation_ready '<iso valid_time>'.
//!   * On failure: status=failed, error=<string>. An hourly cron reseeds.

use chrono::{DateTime, NaiveDateTime, Utc};
use sqlx::postgres::{PgNotification, PgPool, PgPoolOptions};
use sqlx::Row;
use std::time::Duration;
use uuid::Uuid;

pub const NOTIFY_CHANNEL: &str = "propagation_ready";

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TaskKind {
    /// f01..f18 — fetch + decode + score the forecast-hour grid. This is
    /// what Rust has owned since Phase 2 cutover.
    Forecast,
    /// f00 enrichment — adds native-level duct merge, NEXRAD composite,
    /// and commercial-link degradation on top of the forecast pipeline,
    /// and writes the ProfilesFile consumed by /weather. Introduced in
    /// Phase 3 Stream A.
    Analysis,
}

impl TaskKind {
    fn from_str(s: &str) -> Self {
        match s {
            "analysis" => TaskKind::Analysis,
            // Default to Forecast so older producers (which wrote rows
            // before the kind column existed, defaulted to 'forecast'
            // by migration) keep working.
            _ => TaskKind::Forecast,
        }
    }
}

/// NWP source the task should fetch from. Set by the Elixir seeder
/// (`GridTaskEnqueuer.seed_with_analysis/2`); the worker dispatches the
/// fetch + decode + score pipeline on this enum.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TaskSource {
    /// HRRR — CONUS coverage, hourly cycles, NOAA AWS S3.
    Hrrr,
    /// HRDPS — Canadian extent (49-60°N, -141 to -52°W minus the HRRR
    /// overlap), 4×/day at 00/06/12/18Z, MSC Datamart.
    Hrdps,
}

impl TaskSource {
    fn from_str(s: &str) -> Self {
        match s {
            "hrdps" => TaskSource::Hrdps,
            // Default to Hrrr for the legacy rows that predated the
            // source column (and for any source string we don't
            // recognize — we'd rather keep doing the safe HRRR path
            // than refuse to claim).
            _ => TaskSource::Hrrr,
        }
    }

    pub fn as_str(self) -> &'static str {
        match self {
            TaskSource::Hrrr => "hrrr",
            TaskSource::Hrdps => "hrdps",
        }
    }
}

#[derive(Debug, Clone)]
pub struct ClaimedTask {
    pub id: Uuid,
    pub run_time: DateTime<Utc>,
    pub forecast_hour: i16,
    pub valid_time: DateTime<Utc>,
    pub attempt: i32,
    pub kind: TaskKind,
    pub source: TaskSource,
}

#[derive(Debug, thiserror::Error)]
pub enum DbError {
    #[error(transparent)]
    Sqlx(#[from] sqlx::Error),
}

pub async fn connect(url: &str, max_connections: u32) -> Result<PgPool, DbError> {
    let pool = PgPoolOptions::new()
        .max_connections(max_connections)
        .acquire_timeout(Duration::from_secs(10))
        .connect(url)
        .await?;
    Ok(pool)
}

/// Atomically claim one queued grid_tasks row with forecast_hour > 0.
/// Transitions it to status='running', bumps the attempt counter, and
/// stamps claimed_at. Returns `None` if the queue is empty.
pub async fn claim_next(pool: &PgPool) -> Result<Option<ClaimedTask>, DbError> {
    let mut tx = pool.begin().await?;
    // Forecast-only claim. Analysis rows (kind='analysis') go through
    // `claim_next_analysis` once the Rust-side f00 pipeline is wired;
    // keeping them out of the forecast lane means a half-deployed cluster
    // (Elixir still owns f00, Rust pulls forecasts) doesn't deadlock on
    // an unfulfillable analysis row.
    //
    // Order: newest run first, then smallest forecast_hour first. Users
    // loading /map always want the most recent cycle's near-hours before
    // anything else. A stale run whose tasks never drained (broken HRRR
    // publish, etc.) shouldn't block the current hour — the hourly cron
    // reseeds the new run and the prune cron clears rotted rows.
    let row = sqlx::query(
        r#"
        SELECT id, run_time, forecast_hour, valid_time, attempt, kind, source
        FROM grid_tasks
        WHERE status = 'queued' AND forecast_hour > 0 AND kind = 'forecast'
        ORDER BY run_time DESC, forecast_hour ASC
        FOR UPDATE SKIP LOCKED
        LIMIT 1
        "#,
    )
    .fetch_optional(&mut *tx)
    .await?;

    let Some(row) = row else {
        tx.commit().await?;
        return Ok(None);
    };

    let id: Uuid = row.try_get("id")?;
    // Elixir writes run_time/valid_time via Ecto `:utc_datetime`, which
    // lands in Postgres as TIMESTAMP WITHOUT TIME ZONE (values are UTC
    // by Ecto contract). Read as NaiveDateTime and reattach UTC so the
    // rest of the worker can keep its DateTime<Utc> signature.
    let run_time_naive: NaiveDateTime = row.try_get("run_time")?;
    let run_time = DateTime::<Utc>::from_naive_utc_and_offset(run_time_naive, Utc);
    let forecast_hour: i32 = row.try_get("forecast_hour")?;
    let forecast_hour = forecast_hour as i16;
    let valid_time_naive: NaiveDateTime = row.try_get("valid_time")?;
    let valid_time = DateTime::<Utc>::from_naive_utc_and_offset(valid_time_naive, Utc);
    let attempt: i32 = row.try_get("attempt")?;
    let kind_str: String = row.try_get("kind")?;
    let kind = TaskKind::from_str(&kind_str);
    let source_str: String = row.try_get("source")?;
    let source = TaskSource::from_str(&source_str);

    sqlx::query(
        r#"
        UPDATE grid_tasks
        SET status = 'running',
            claimed_at = NOW(),
            attempt = attempt + 1,
            updated_at = NOW()
        WHERE id = $1
        "#,
    )
    .bind(id)
    .execute(&mut *tx)
    .await?;

    tx.commit().await?;

    Ok(Some(ClaimedTask {
        id,
        run_time,
        forecast_hour,
        valid_time,
        attempt: attempt + 1,
        kind,
        source,
    }))
}

/// Claim one queued `kind='analysis'` row. Analysis tasks carry the
/// f00 enrichment (native-level duct merge, NEXRAD composite,
/// commercial-link degradation, ProfilesFile write) that Rust only
/// learned to do in Phase 3 Stream A. Kept on a separate lane from
/// `claim_next` so operators can gate the analysis path independently
/// of the forecast path during rollout.
///
/// Restricted to `source = 'hrrr'` for now — the HRDPS analysis path
/// (native-level duct + NEXRAD + commercial merges) hasn't been ported
/// yet, so HRDPS analysis rows would fail mid-pipeline. HRDPS analysis
/// stays Elixir-side / out of scope until that gap is closed.
pub async fn claim_next_analysis(pool: &PgPool) -> Result<Option<ClaimedTask>, DbError> {
    let mut tx = pool.begin().await?;
    let row = sqlx::query(
        r#"
        SELECT id, run_time, forecast_hour, valid_time, attempt, kind, source
        FROM grid_tasks
        WHERE status = 'queued' AND kind = 'analysis' AND source = 'hrrr'
        ORDER BY run_time DESC, forecast_hour ASC
        FOR UPDATE SKIP LOCKED
        LIMIT 1
        "#,
    )
    .fetch_optional(&mut *tx)
    .await?;

    let Some(row) = row else {
        tx.commit().await?;
        return Ok(None);
    };

    let id: Uuid = row.try_get("id")?;
    let run_time_naive: NaiveDateTime = row.try_get("run_time")?;
    let run_time = DateTime::<Utc>::from_naive_utc_and_offset(run_time_naive, Utc);
    let forecast_hour: i32 = row.try_get("forecast_hour")?;
    let forecast_hour = forecast_hour as i16;
    let valid_time_naive: NaiveDateTime = row.try_get("valid_time")?;
    let valid_time = DateTime::<Utc>::from_naive_utc_and_offset(valid_time_naive, Utc);
    let attempt: i32 = row.try_get("attempt")?;
    let kind_str: String = row.try_get("kind")?;
    let kind = TaskKind::from_str(&kind_str);
    let source_str: String = row.try_get("source")?;
    let source = TaskSource::from_str(&source_str);

    sqlx::query(
        r#"
        UPDATE grid_tasks
        SET status = 'running',
            claimed_at = NOW(),
            attempt = attempt + 1,
            updated_at = NOW()
        WHERE id = $1
        "#,
    )
    .bind(id)
    .execute(&mut *tx)
    .await?;

    tx.commit().await?;

    Ok(Some(ClaimedTask {
        id,
        run_time,
        forecast_hour,
        valid_time,
        attempt: attempt + 1,
        kind,
        source,
    }))
}

/// Mark a task done and emit NOTIFY propagation_ready '<iso valid_time>'.
/// The NOTIFY payload is read by Elixir's `PropagationNotifyListener` to
/// warm `ScoreCache` and fan out `"propagation:updated"` PubSub.
pub async fn complete(pool: &PgPool, task: &ClaimedTask) -> Result<(), DbError> {
    let mut tx = pool.begin().await?;
    sqlx::query(
        r#"
        UPDATE grid_tasks
        SET status = 'done',
            completed_at = NOW(),
            error = NULL,
            updated_at = NOW()
        WHERE id = $1
        "#,
    )
    .bind(task.id)
    .execute(&mut *tx)
    .await?;

    let payload = task.valid_time.format("%Y-%m-%dT%H:%M:%SZ").to_string();
    let notify = format!("NOTIFY {}, '{}'", NOTIFY_CHANNEL, payload);
    sqlx::query(&notify).execute(&mut *tx).await?;
    tx.commit().await?;
    Ok(())
}

pub async fn fail(pool: &PgPool, task_id: Uuid, error: &str) -> Result<(), DbError> {
    sqlx::query(
        r#"
        UPDATE grid_tasks
        SET status = 'failed',
            completed_at = NOW(),
            error = $2,
            updated_at = NOW()
        WHERE id = $1
        "#,
    )
    .bind(task_id)
    .bind(error)
    .execute(pool)
    .await?;
    Ok(())
}

/// Park the task back in the queue (e.g. SIGTERM received mid-work).
/// Resets status to 'queued' without incrementing attempt so a restart
/// doesn't chew through the retry budget.
pub async fn requeue(pool: &PgPool, task_id: Uuid) -> Result<(), DbError> {
    sqlx::query(
        r#"
        UPDATE grid_tasks
        SET status = 'queued',
            claimed_at = NULL,
            updated_at = NOW()
        WHERE id = $1 AND status = 'running'
        "#,
    )
    .bind(task_id)
    .execute(pool)
    .await?;
    Ok(())
}

// Imported only to keep the PgNotification type used from `sqlx::postgres`.
const _: fn() = || {
    fn assert_send<T: Send>() {}
    assert_send::<PgNotification>();
};

/// Per-QSO HRRR fetch task. One row per (valid_time, JSONB array of
/// points). Consumed by `bin/hrrr_point_worker`.
#[derive(Debug, Clone)]
pub struct HrrrFetchTask {
    pub id: Uuid,
    pub valid_time: DateTime<Utc>,
    pub points: Vec<(f64, f64)>,
    pub attempt: i32,
}

pub async fn claim_next_hrrr_task(pool: &PgPool) -> Result<Option<HrrrFetchTask>, DbError> {
    let mut tx = pool.begin().await?;
    let row = sqlx::query(
        r#"
        SELECT id, valid_time, points, attempt
        FROM hrrr_fetch_tasks
        WHERE status = 'queued'
        ORDER BY valid_time DESC
        FOR UPDATE SKIP LOCKED
        LIMIT 1
        "#,
    )
    .fetch_optional(&mut *tx)
    .await?;

    let Some(row) = row else {
        tx.commit().await?;
        return Ok(None);
    };

    let id: Uuid = row.try_get("id")?;
    let vt_naive: NaiveDateTime = row.try_get("valid_time")?;
    let valid_time = DateTime::<Utc>::from_naive_utc_and_offset(vt_naive, Utc);
    let attempt: i32 = row.try_get("attempt")?;
    let points_json: serde_json::Value = row.try_get("points")?;
    let points = points_from_json(&points_json);

    sqlx::query(
        r#"
        UPDATE hrrr_fetch_tasks
        SET status = 'running',
            claimed_at = NOW(),
            attempt = attempt + 1,
            updated_at = NOW()
        WHERE id = $1
        "#,
    )
    .bind(id)
    .execute(&mut *tx)
    .await?;

    tx.commit().await?;

    Ok(Some(HrrrFetchTask {
        id,
        valid_time,
        points,
        attempt: attempt + 1,
    }))
}

pub async fn complete_hrrr_task(pool: &PgPool, task_id: Uuid) -> Result<(), DbError> {
    sqlx::query(
        r#"
        UPDATE hrrr_fetch_tasks
        SET status = 'done',
            completed_at = NOW(),
            error = NULL,
            updated_at = NOW()
        WHERE id = $1
        "#,
    )
    .bind(task_id)
    .execute(pool)
    .await?;
    Ok(())
}

pub async fn fail_hrrr_task(pool: &PgPool, task_id: Uuid, error: &str) -> Result<(), DbError> {
    sqlx::query(
        r#"
        UPDATE hrrr_fetch_tasks
        SET status = 'failed',
            completed_at = NOW(),
            error = $2,
            updated_at = NOW()
        WHERE id = $1
        "#,
    )
    .bind(task_id)
    .bind(error)
    .execute(pool)
    .await?;
    Ok(())
}

fn points_from_json(v: &serde_json::Value) -> Vec<(f64, f64)> {
    let Some(arr) = v.as_array() else {
        return Vec::new();
    };
    arr.iter()
        .filter_map(|p| {
            let lat = p.get("lat")?.as_f64()?;
            let lon = p.get("lon")?.as_f64()?;
            Some((lat, lon))
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;

    /// Live-DB integration tests — gated behind the `PROP_GRID_RS_TEST_DB`
    /// env var so `cargo test` stays green without Postgres.
    #[tokio::test]
    async fn claim_and_complete_roundtrip() {
        let Ok(url) = std::env::var("PROP_GRID_RS_TEST_DB") else {
            eprintln!("skip: set PROP_GRID_RS_TEST_DB to run db tests");
            return;
        };
        let pool = connect(&url, 4).await.unwrap();

        // Seed a unique fh>0 row. run_time/valid_time bind as
        // NaiveDateTime to match Ecto's :utc_datetime column type.
        let run_time = Utc::now();
        let fh: i16 = 7;
        let valid_time = run_time + chrono::Duration::hours(fh as i64);
        let id = Uuid::new_v4();
        sqlx::query(
            r#"INSERT INTO grid_tasks (id, run_time, forecast_hour, valid_time, status, attempt, inserted_at, updated_at)
               VALUES ($1, $2, $3, $4, 'queued', 0, NOW(), NOW())"#,
        )
        .bind(id)
        .bind(run_time.naive_utc())
        .bind(fh as i32)
        .bind(valid_time.naive_utc())
        .execute(&pool)
        .await
        .unwrap();

        let claimed = claim_next(&pool).await.unwrap().expect("had a row");
        assert!(claimed.forecast_hour > 0);
        assert_eq!(claimed.attempt, 1);
        assert_eq!(claimed.kind, TaskKind::Forecast);
        // Default-DB rows land with source='hrrr' per the migration default.
        assert_eq!(claimed.source, TaskSource::Hrrr);

        complete(&pool, &claimed).await.unwrap();

        // Clean up.
        sqlx::query("DELETE FROM grid_tasks WHERE id = $1")
            .bind(claimed.id)
            .execute(&pool)
            .await
            .unwrap();
    }
}