awa_model/

migrations.rs

use crate::error::AwaError;
use sqlx::postgres::PgConnection;
use sqlx::PgPool;
use tracing::info;

/// Current schema version.
pub const CURRENT_VERSION: i32 = 40;

/// All migrations in order. SQL lives in `awa-model/migrations/*.sql`
/// for easy inspection by users who run their own migration tooling.
///
/// ## Migration policy
///
/// Migrations MUST be **additive only**:
/// - Add tables, columns (with defaults), indexes, functions
/// - Never drop columns, change types, or tighten constraints
///
/// This ensures running workers are not broken by a schema upgrade.
/// Representation changes require a release-specific compatibility and
/// authority-transition procedure.
const MIGRATIONS: &[(i32, &str, &[&str])] = &[
    (1, "Canonical schema with UI indexes", &[V1_UP]),
    (2, "Runtime observability snapshots", &[V2_UP]),
    (3, "Maintenance loop health in runtime snapshots", &[V3_UP]),
    (4, "Admin metadata cache tables", &[V4_UP]),
    (5, "Statement-level admin metadata triggers", &[V5_UP]),
    (
        6,
        "Dirty-key statement triggers for deadlock-free admin metadata",
        &[V6_UP],
    ),
    (
        7,
        "Backoff interval creation avoids scientific-notation parse failures",
        &[V7_UP],
    ),
    // v008 is reserved for the dead-letter-queue migration on a parallel
    // branch; leave the slot open so both PRs can land without renumbering.
    (9, "Queue and job-kind descriptor catalogs", &[V9_UP]),
    (
        10,
        "Storage transition metadata and canonical compat routing",
        &[V10_UP],
    ),
    (
        11,
        "Storage transition self-heal: NULL-safe engine resolution and singleton re-seed",
        &[V11_UP],
    ),
    (
        12,
        "Queue storage compatibility layer and active backend selection",
        &[V12_UP],
    ),
    (
        13,
        "Storage auto-finalize and queue-storage count maintenance",
        &[V13_UP],
    ),
    (
        14,
        "Storage transition role tracking and tightened mixed-transition gate",
        &[V14_UP],
    ),
    (15, "Cron missed-fire policy", &[V15_UP]),
    (
        16,
        "Drop redundant queue_lanes.available_count cache; reader derives from heads",
        &[V16_UP],
    ),
    (
        17,
        "Shard queue_enqueue_heads/queue_claim_heads/ready_entries by enqueue_shard",
        &[V17_UP],
    ),
    (
        18,
        "Thread ordering_key through insert_job_compat for queue-storage producers",
        &[V18_UP],
    ),
    (
        19,
        "Make queue-storage jobs compatibility view shard-aware",
        &[V19_UP],
    ),
    (
        20,
        "Derive active queue-storage schema from transition state",
        &[V20_UP],
    ),
    (
        21,
        "Shard-aware lane indexes on ready_entries/done_entries/leases",
        &[V21_UP],
    ),
    (
        22,
        "delete_job_compat decrements queue_terminal_live_counts for done_entries deletes",
        &[V22_UP],
    ),
    (
        23,
        "Install default awa queue-storage substrate via SQL helper",
        &[V23_UP],
    ),
    (
        24,
        "Lower fillfactor to 50 on leases and lease_claims partitions",
        &[V24_UP],
    ),
    (
        25,
        "Drop idx_<schema>_leases_<slot>_state_hb on all AWA substrates",
        &[V25_UP],
    ),
    (
        26,
        "Add paused_at + paused_by to cron_jobs for per-schedule pause",
        &[V26_UP],
    ),
    (
        27,
        "Move lane cursors to sequences and stripe terminal counters",
        &[V23_UP, V22_UP, V27_UP],
    ),
    (
        28,
        "Add ready_tombstones ledger and compatibility filters (#295)",
        &[V23_UP, V22_UP, V28_UP],
    ),
    (29, "Add durable batch operations control table", &[V29_UP]),
    (
        30,
        "Add terminal-count delta ledger and async rollup",
        &[V23_UP, V30_UP],
    ),
    (
        31,
        "Backfill failed done-entry metric indexes for queue storage",
        &[V31_UP],
    ),
    (
        32,
        "Add pruned_failed_count to queue_terminal_rollups for the failed terminal retention floor",
        &[V32_UP],
    ),
    (
        33,
        "Add per-slot receipt-rescue cursors for queue storage",
        &[V23_UP, V33_UP],
    ),
    (
        34,
        "Materialize receipt closures before terminal compatibility deletes",
        &[V34_UP],
    ),
    (
        35,
        "Add per-slot receipt deadline-rescue cursors for queue storage",
        &[V35_UP],
    ),
    (
        36,
        "Compact successful receipt completions into batch terminal history",
        &[V18_UP, V23_UP, V36_UP],
    ),
    (
        37,
        "Add ready_segments claim-routing ledger for queue storage",
        &[V18_UP, V23_UP, V37_UP],
    ),
    (
        38,
        "Add compact receipt claim batch ledger for queue storage",
        &[V23_UP, V38_UP],
    ),
    (
        39,
        "Refresh claim_ready_runtime to cache-free ready-segment routing",
        &[V23_UP, V39_UP],
    ),
    (
        40,
        "Allow queue-storage finalization with live canonical drain-only runtimes",
        &[V40_UP],
    ),
];

const V1_UP: &str = include_str!("../migrations/v001_canonical_schema.sql");
const V2_UP: &str = include_str!("../migrations/v002_runtime_instances.sql");
const V3_UP: &str = include_str!("../migrations/v003_maintenance_health.sql");
const V4_UP: &str = include_str!("../migrations/v004_admin_metadata.sql");
const V5_UP: &str = include_str!("../migrations/v005_admin_metadata_stmt_triggers.sql");
const V6_UP: &str = include_str!("../migrations/v006_remove_hot_table_triggers.sql");
const V7_UP: &str = include_str!("../migrations/v007_backoff_interval_fix.sql");
const V9_UP: &str = include_str!("../migrations/v009_descriptors.sql");
const V10_UP: &str = include_str!("../migrations/v010_storage_transition_prep.sql");
const V11_UP: &str = include_str!("../migrations/v011_storage_transition_self_heal.sql");
const V12_UP: &str = include_str!("../migrations/v012_queue_storage_compat.sql");
const V13_UP: &str = include_str!("../migrations/v013_storage_auto_finalize.sql");
const V14_UP: &str = include_str!("../migrations/v014_storage_transition_role.sql");
const V15_UP: &str = include_str!("../migrations/v015_cron_missed_fire_policy.sql");
const V16_UP: &str = include_str!("../migrations/v016_drop_queue_lanes_available_count.sql");
const V17_UP: &str = include_str!("../migrations/v017_shard_queue_enqueue_heads.sql");
const V18_UP: &str = include_str!("../migrations/v018_insert_job_compat_ordering_key.sql");
const V19_UP: &str = include_str!("../migrations/v019_queue_storage_jobs_compat_shard_joins.sql");
const V20_UP: &str = include_str!("../migrations/v020_active_queue_storage_schema_fallback.sql");
const V21_UP: &str = include_str!("../migrations/v021_shard_aware_lane_indexes.sql");
const V22_UP: &str = include_str!("../migrations/v022_delete_compat_terminal_counter.sql");
const V23_UP: &str = include_str!("../migrations/v023_install_queue_storage_substrate.sql");
const V24_UP: &str = include_str!("../migrations/v024_receipt_plane_fillfactor.sql");
const V25_UP: &str = include_str!("../migrations/v025_drop_leases_state_hb_index.sql");
const V26_UP: &str = include_str!("../migrations/v026_cron_jobs_pause.sql");
const V27_UP: &str = include_str!("../migrations/v027_sequence_lane_cursors.sql");
const V28_UP: &str = include_str!("../migrations/v028_ready_tombstones.sql");
const V29_UP: &str = include_str!("../migrations/v029_batch_operations.sql");
const V30_UP: &str = include_str!("../migrations/v030_terminal_count_deltas.sql");
const V31_UP: &str = include_str!("../migrations/v031_queue_storage_failed_done_indexes.sql");
const V32_UP: &str = include_str!("../migrations/v032_failed_terminal_retention.sql");
const V33_UP: &str = include_str!("../migrations/v033_receipt_rescue_cursors.sql");
const V34_UP: &str = include_str!("../migrations/v034_receipt_terminal_delete_closures.sql");
const V35_UP: &str = include_str!("../migrations/v035_receipt_deadline_rescue_cursors.sql");
const V36_UP: &str = include_str!("../migrations/v036_compact_receipt_completions.sql");
const V37_UP: &str = include_str!("../migrations/v037_ready_segments.sql");
const V38_UP: &str = include_str!("../migrations/v038_compact_claim_batches.sql");
const V39_UP: &str = include_str!("../migrations/v039_claim_head_cold_routing.sql");
const V40_UP: &str = include_str!("../migrations/v040_finalize_with_drain_runtimes.sql");

/// Old version numbers from pre-0.4 releases that used V3/V4/V5 numbering.
/// Also tolerates the unreleased inline-V6 branch numbering used during review.
/// Maps old max version → equivalent new version.
fn normalize_legacy_version(old_version: i32) -> i32 {
    match old_version {
        v if v >= 6 => 4, // legacy/unreleased V6 admin metadata = V4 (new)
        5 => 3,           // V5 (0.3.x) = V3 (new)
        4 => 2,           // V4 = V2 (new)
        3 => 1,           // V3 = V1 (new)
        _ => 0,           // Pre-canonical or fresh
    }
}

/// Run all pending migrations against the database.
///
/// Applies only migrations newer than the current schema version.
/// V1 bootstraps the canonical schema from scratch; V2+ are incremental
/// and use `IF NOT EXISTS` guards so they are safe to re-run.
///
/// by replacing the legacy `schema_version` rows with the new numbering.
///
/// Takes `&PgPool` for ergonomic use from Rust.
pub async fn run(pool: &PgPool) -> Result<(), AwaError> {
    let lock_key: i64 = 0x4157_415f_4d49_4752; // "AWA_MIGR"
    let mut conn = pool.acquire().await?;
    sqlx::query("SELECT pg_advisory_lock($1)")
        .bind(lock_key)
        .execute(&mut *conn)
        .await?;

    let result = run_inner(&mut conn).await;

    let _ = sqlx::query("SELECT pg_advisory_unlock($1)")
        .bind(lock_key)
        .execute(&mut *conn)
        .await;

    result
}

async fn run_inner(conn: &mut PgConnection) -> Result<(), AwaError> {
    let has_schema: bool =
        sqlx::query_scalar("SELECT EXISTS(SELECT 1 FROM pg_namespace WHERE nspname = 'awa')")
            .fetch_one(&mut *conn)
            .await?;

    let current = if has_schema {
        current_version_conn(conn).await?
    } else {
        0
    };

    if current > CURRENT_VERSION {
        info!(
            schema_version = current,
            supported_schema_version = CURRENT_VERSION,
            "Schema is newer but forward-compatible; no migrations to apply"
        );
    } else if !(has_schema && current == CURRENT_VERSION) {
        for &(version, description, steps) in MIGRATIONS {
            if version <= current {
                continue;
            }
            info!(version, description, "Applying migration");
            for step in steps {
                sqlx::raw_sql(step).execute(&mut *conn).await?;
            }
            info!(version, "Migration applied");
        }
    } else {
        info!(version = current, "Schema is up to date");
    }

    // Ensure the admin metadata cache is warm. Since v006 removed the
    // synchronous triggers on jobs_hot, the cache is only updated by the
    // maintenance leader. Refreshing here guarantees queue_stats() and
    // state_counts() return accurate data immediately after migrate().
    let has_refresh: bool = sqlx::query_scalar(
        "SELECT EXISTS(SELECT 1 FROM pg_proc WHERE proname = 'refresh_admin_metadata' AND pronamespace = (SELECT oid FROM pg_namespace WHERE nspname = 'awa'))",
    )
    .fetch_one(&mut *conn)
    .await?;
    if has_refresh {
        // Best-effort cache warmup. Uses a short statement timeout to avoid
        // blocking if a previous runtime's maintenance leader is still
        // holding the cache advisory lock during a slow shutdown.
        // Wrapped in an explicit transaction because SET LOCAL is only
        // effective inside a transaction block (not in autocommit mode).
        let _ = sqlx::raw_sql(
            "BEGIN; SET LOCAL statement_timeout = '5s'; SELECT awa.refresh_admin_metadata(); COMMIT;",
        )
        .execute(&mut *conn)
        .await;
    }

    Ok(())
}

/// Get the current schema version.
pub async fn current_version(pool: &PgPool) -> Result<i32, AwaError> {
    let mut conn = pool.acquire().await?;
    current_version_conn(&mut conn).await
}

async fn current_version_conn(conn: &mut PgConnection) -> Result<i32, AwaError> {
    let has_schema: bool =
        sqlx::query_scalar("SELECT EXISTS(SELECT 1 FROM pg_namespace WHERE nspname = 'awa')")
            .fetch_one(&mut *conn)
            .await?;

    if !has_schema {
        return Ok(0);
    }

    let has_table: bool = sqlx::query_scalar(
        "SELECT EXISTS(SELECT 1 FROM information_schema.tables WHERE table_schema = 'awa' AND table_name = 'schema_version')",
    )
    .fetch_one(&mut *conn)
    .await?;

    if !has_table {
        return Ok(0);
    }

    let version: Option<i32> = sqlx::query_scalar("SELECT MAX(version) FROM awa.schema_version")
        .fetch_one(&mut *conn)
        .await?;

    let raw_version = version.unwrap_or(0);

    if raw_version > CURRENT_VERSION {
        if forward_compatible_v043_columns(conn, raw_version).await? {
            return Ok(raw_version);
        }
        tracing::error!(
            schema_version = raw_version,
            supported_schema_version = CURRENT_VERSION,
            "This Awa binary is too old for the database schema; upgrade Awa before retrying"
        );
        return Err(AwaError::SchemaNotMigrated {
            expected: CURRENT_VERSION,
            found: raw_version,
        });
    }

    // If max version is within the current MIGRATIONS range and the expected
    // tables exist, this is a current install — skip legacy detection.
    if (1..=CURRENT_VERSION).contains(&raw_version) {
        // Quick check: does the schema match what we expect at this version?
        // If queue_state_counts exists, we're past v4 in the current numbering.
        let has_admin_tables: bool = sqlx::query_scalar(
            "SELECT EXISTS(SELECT 1 FROM information_schema.tables WHERE table_schema = 'awa' AND table_name = 'queue_state_counts')",
        )
        .fetch_one(&mut *conn)
        .await
        .unwrap_or(false);

        // Current v4+ has queue_state_counts. If we're at v4+ and have
        // the table, this is definitely a current install.
        if raw_version >= 4 && has_admin_tables {
            return Ok(raw_version);
        }
        // Current v1-v3 don't have queue_state_counts.
        if raw_version <= 3 {
            let has_runtime: bool = sqlx::query_scalar(
                "SELECT EXISTS(SELECT 1 FROM information_schema.tables WHERE table_schema = 'awa' AND table_name = 'runtime_instances')",
            )
            .fetch_one(&mut *conn)
            .await
            .unwrap_or(false);
            // v2+ has runtime_instances. If present, current install.
            if (raw_version >= 2 && has_runtime) || raw_version == 1 {
                return Ok(raw_version);
            }
        }
    }

    // Detect legacy version numbering from pre-0.4 releases.
    // Legacy installs used a different numbering scheme where v3-v6 mapped
    // to what is now v1-v4.
    let has_legacy_high: bool =
        sqlx::query_scalar("SELECT EXISTS(SELECT 1 FROM awa.schema_version WHERE version >= 6)")
            .fetch_one(&mut *conn)
            .await
            .unwrap_or(false);

    let has_admin_metadata: bool = sqlx::query_scalar(
        "SELECT EXISTS(SELECT 1 FROM information_schema.tables WHERE table_schema = 'awa' AND table_name = 'queue_state_counts')",
    )
    .fetch_one(&mut *conn)
    .await
    .unwrap_or(false);

    let is_legacy_v5_only = raw_version == 5 && !has_legacy_high && !has_admin_metadata;
    let is_legacy_v4_only = raw_version == 4 && !has_legacy_high && !has_admin_metadata;

    // Also detect a single legacy V3 row (0.3.0 with only canonical schema)
    // by checking if runtime_instances exists — if not, this is legacy V3.
    let is_legacy_v3_only = raw_version == 3
        && !has_legacy_high
        && {
            let has_runtime: bool = sqlx::query_scalar(
            "SELECT EXISTS(SELECT 1 FROM information_schema.tables WHERE table_schema = 'awa' AND table_name = 'runtime_instances')",
        )
        .fetch_one(&mut *conn)
        .await
        .unwrap_or(false);
            !has_runtime
        };

    if has_legacy_high || is_legacy_v5_only || is_legacy_v4_only || is_legacy_v3_only {
        let normalized = normalize_legacy_version(raw_version);
        info!(
            old_version = raw_version,
            new_version = normalized,
            "Normalizing legacy version numbering"
        );
        // Replace legacy rows so future calls return the new numbering.
        sqlx::query("DELETE FROM awa.schema_version WHERE version >= 3")
            .execute(&mut *conn)
            .await?;
        for &(v, desc, _) in MIGRATIONS {
            if v <= normalized {
                sqlx::query(
                    "INSERT INTO awa.schema_version (version, description) VALUES ($1, $2) ON CONFLICT (version) DO NOTHING",
                )
                .bind(v)
                .bind(desc)
                .execute(&mut *conn)
                .await?;
            }
        }
        return Ok(normalized);
    }

    Ok(raw_version)
}

/// The 0.6.2 rolling-upgrade stepping-stone recognizes exactly the additive
/// v043 shape, and only before any queue-storage schema has flipped to ledger
/// authority. This lets an already-rolled 0.6.2 fleet survive a migration-first
/// 0.7 rollout. Unknown future schemas and the post-flip boundary fail closed.
async fn forward_compatible_v043_columns(
    conn: &mut PgConnection,
    raw_version: i32,
) -> Result<bool, AwaError> {
    if raw_version != 43 {
        return Ok(false);
    }

    // quote_ident is computed by PostgreSQL before the identifier is used in
    // the dynamic query below. Queue-storage schema names are validated when
    // installed, but retaining quoting here keeps this probe safe for catalogs
    // created by external migration tooling too.
    let schemas: Vec<String> = sqlx::query_scalar(
        "SELECT quote_ident(n.nspname) \
         FROM pg_namespace AS n \
         JOIN pg_class AS c ON c.relnamespace = n.oid \
         WHERE c.relname = 'ring_cursor_authority' \
           AND c.relkind IN ('r', 'p') \
         ORDER BY n.nspname",
    )
    .fetch_all(&mut *conn)
    .await?;

    if schemas.is_empty() {
        return Ok(false);
    }

    for schema in schemas {
        let authority: Option<String> = sqlx::query_scalar(&format!(
            "SELECT authority FROM {schema}.ring_cursor_authority WHERE singleton"
        ))
        .fetch_optional(&mut *conn)
        .await?;
        if authority.as_deref() != Some("columns") {
            return Ok(false);
        }
    }

    Ok(true)
}

/// Get the raw SQL for all migrations (for extraction / external tooling).
pub fn migration_sql() -> Vec<(i32, &'static str, String)> {
    MIGRATIONS
        .iter()
        .map(|&(v, d, steps)| (v, d, steps.join("\n")))
        .collect()
}

/// Get migration SQL for a version range `(from, to]` — `from` is exclusive,
/// `to` is inclusive. Returns only migrations where `from < version <= to`.
pub fn migration_sql_range(from: i32, to: i32) -> Vec<(i32, &'static str, String)> {
    MIGRATIONS
        .iter()
        .filter(|&&(v, _, _)| v > from && v <= to)
        .map(|&(v, d, steps)| (v, d, steps.join("\n")))
        .collect()
}

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

    #[test]
    fn migration_sql_range_all() {
        let all = migration_sql_range(0, CURRENT_VERSION);
        assert_eq!(all.len(), MIGRATIONS.len());
        assert_eq!(all.first().unwrap().0, 1);
        assert_eq!(all.last().unwrap().0, CURRENT_VERSION);
    }

    #[test]
    fn migration_sql_range_subset() {
        let subset = migration_sql_range(2, CURRENT_VERSION);
        assert!(subset.iter().all(|(v, _, _)| *v > 2));
        let expected = MIGRATIONS.iter().filter(|&&(v, _, _)| v > 2).count();
        assert_eq!(subset.len(), expected);
    }

    #[test]
    fn migration_sql_range_single() {
        let single = migration_sql_range(2, 3);
        assert_eq!(single.len(), 1);
        assert_eq!(single[0].0, 3);
        assert!(!single[0].2.is_empty());
    }

    #[test]
    fn migration_sql_range_empty_when_equal() {
        let empty = migration_sql_range(CURRENT_VERSION, CURRENT_VERSION);
        assert!(empty.is_empty());
    }

    #[test]
    fn migration_sql_range_empty_when_inverted() {
        let empty = migration_sql_range(3, 1);
        assert!(empty.is_empty());
    }

    #[test]
    fn migration_sql_range_matches_full() {
        let full = migration_sql();
        let ranged = migration_sql_range(0, CURRENT_VERSION);
        assert_eq!(full.len(), ranged.len());
        for (f, r) in full.iter().zip(ranged.iter()) {
            assert_eq!(f.0, r.0);
            assert_eq!(f.1, r.1);
            assert_eq!(f.2, r.2);
        }
    }
}