PiCloud/crates/manager-core/tests/atomic_write.rs

//! Audit fix #6: a data write and the trigger fan-out it produces must commit
//! ATOMICALLY.
//!
//! Before this, a service wrote the row (commit #1), then asked the emitter to
//! resolve matching triggers and insert outbox rows (commit #2). If commit #2
//! failed, the row was permanently in the store with its trigger having never
//! fired — unobservable to the caller, unrecoverable by any retry, and only
//! logged. `PostgresKvWriter` runs both on ONE connection in ONE transaction, so
//! an emit failure rolls the write back and surfaces as an error.
//!
//! Fault injection: a Postgres BEFORE-INSERT trigger on `outbox` that raises,
//! scoped to THIS test's freshly-minted `app_id`, so it cannot affect any test
//! running in parallel against the same database. Skips when `DATABASE_URL` is
//! unset.

use picloud_manager_core::atomic_write::{KvWriter, PostgresKvWriter};
use picloud_shared::{AppId, ExecutionId, RequestId, ScriptId, SdkCallCx};
use sqlx::postgres::PgPoolOptions;
use sqlx::PgPool;
use uuid::Uuid;

async fn pool_or_skip() -> Option<PgPool> {
    let Ok(url) = std::env::var("DATABASE_URL") else {
        eprintln!("atomic_write: DATABASE_URL unset — skipping");
        return None;
    };
    let pool = PgPoolOptions::new()
        .max_connections(3)
        .connect(&url)
        .await
        .expect("connect");
    sqlx::migrate!("./migrations")
        .run(&pool)
        .await
        .expect("migrate");
    Some(pool)
}

fn cx(app_id: Uuid) -> SdkCallCx {
    SdkCallCx {
        app_id: AppId::from(app_id),
        script_id: ScriptId::new(),
        principal: None,
        execution_id: ExecutionId::new(),
        request_id: RequestId::new(),
        trigger_depth: 0,
        root_execution_id: ExecutionId::new(),
        is_dead_letter_handler: false,
        event: None,
    }
}

/// An app with one script and one enabled KV trigger watching `collection`.
struct Fixture {
    app: Uuid,
}

async fn setup(pool: &PgPool, collection: &str) -> Fixture {
    let uniq = Uuid::new_v4().simple().to_string();
    let (group,): (Uuid,) =
        sqlx::query_as("INSERT INTO groups (slug, name) VALUES ($1, $1) RETURNING id")
            .bind(format!("aw-grp-{uniq}"))
            .fetch_one(pool)
            .await
            .expect("group");
    let (app,): (Uuid,) =
        sqlx::query_as("INSERT INTO apps (slug, name, group_id) VALUES ($1, $1, $2) RETURNING id")
            .bind(format!("aw-app-{uniq}"))
            .bind(group)
            .fetch_one(pool)
            .await
            .expect("app");
    let (admin,): (Uuid,) = sqlx::query_as(
        "INSERT INTO admin_users (username, password_hash) VALUES ($1, 'x') RETURNING id",
    )
    .bind(format!("aw-admin-{uniq}"))
    .fetch_one(pool)
    .await
    .expect("admin");
    let (script,): (Uuid,) = sqlx::query_as(
        "INSERT INTO scripts (app_id, name, source) VALUES ($1, 'handler', '') RETURNING id",
    )
    .bind(app)
    .fetch_one(pool)
    .await
    .expect("script");
    let (trigger,): (Uuid,) = sqlx::query_as(
        "INSERT INTO triggers ( \
            app_id, script_id, kind, retry_max_attempts, retry_backoff, \
            retry_base_ms, registered_by_principal \
         ) VALUES ($1, $2, 'kv', 3, 'exponential', 1000, $3) RETURNING id",
    )
    .bind(app)
    .bind(script)
    .bind(admin)
    .fetch_one(pool)
    .await
    .expect("trigger");
    sqlx::query(
        "INSERT INTO kv_trigger_details (trigger_id, collection_glob, ops) \
         VALUES ($1, $2, '{}')",
    )
    .bind(trigger)
    .bind(collection)
    .execute(pool)
    .await
    .expect("kv_trigger_details");
    Fixture { app }
}

/// Make every outbox insert for `app` fail. Scoped by app_id so parallel tests
/// (which use their own fresh apps) are untouched.
async fn break_outbox_for(pool: &PgPool, app: Uuid, tag: &str) {
    sqlx::query(&format!(
        "CREATE FUNCTION break_outbox_{tag}() RETURNS TRIGGER AS $$ \
         BEGIN \
           IF NEW.app_id = '{app}'::uuid THEN \
             RAISE EXCEPTION 'injected outbox failure'; \
           END IF; \
           RETURN NEW; \
         END; $$ LANGUAGE plpgsql"
    ))
    .execute(pool)
    .await
    .expect("create fn");
    sqlx::query(&format!(
        "CREATE TRIGGER break_outbox_{tag} BEFORE INSERT ON outbox \
         FOR EACH ROW EXECUTE FUNCTION break_outbox_{tag}()"
    ))
    .execute(pool)
    .await
    .expect("create trigger");
}

async fn unbreak_outbox(pool: &PgPool, tag: &str) {
    sqlx::query(&format!(
        "DROP TRIGGER IF EXISTS break_outbox_{tag} ON outbox"
    ))
    .execute(pool)
    .await
    .expect("drop trigger");
    sqlx::query(&format!("DROP FUNCTION IF EXISTS break_outbox_{tag}()"))
        .execute(pool)
        .await
        .expect("drop fn");
}

async fn kv_count(pool: &PgPool, app: Uuid) -> i64 {
    let (n,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM kv_entries WHERE app_id = $1")
        .bind(app)
        .fetch_one(pool)
        .await
        .expect("count kv");
    n
}

async fn outbox_count(pool: &PgPool, app: Uuid) -> i64 {
    let (n,): (i64,) = sqlx::query_as("SELECT COUNT(*) FROM outbox WHERE app_id = $1")
        .bind(app)
        .fetch_one(pool)
        .await
        .expect("count outbox");
    n
}

/// `scripts` is RESTRICT on app delete (code is not data), so drop it first.
async fn cleanup(pool: &PgPool, app: Uuid) {
    sqlx::query("DELETE FROM scripts WHERE app_id = $1")
        .bind(app)
        .execute(pool)
        .await
        .expect("cleanup scripts");
    sqlx::query("DELETE FROM apps WHERE id = $1")
        .bind(app)
        .execute(pool)
        .await
        .expect("cleanup app");
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn a_failed_fan_out_rolls_the_kv_write_back() {
    let Some(pool) = pool_or_skip().await else {
        return;
    };
    let f = setup(&pool, "watched").await;
    let writer = PostgresKvWriter::new(pool.clone());
    let cx = cx(f.app);
    // A distinct, valid SQL identifier for this run's trigger/function names.
    let tag = Uuid::new_v4().simple().to_string();

    // --- Control: with the outbox healthy, the write lands AND fans out. -----
    writer
        .set(&cx, "watched", "k", serde_json::json!(1))
        .await
        .expect("healthy set succeeds");
    assert_eq!(kv_count(&pool, f.app).await, 1, "the row is written");
    assert_eq!(
        outbox_count(&pool, f.app).await,
        1,
        "the matching trigger fanned out"
    );

    // --- The fix: a fan-out failure must roll the write back. ---------------
    break_outbox_for(&pool, f.app, &tag).await;
    let err = writer
        .set(&cx, "watched", "k2", serde_json::json!(2))
        .await
        .expect_err("an outbox failure must surface as an error, not be swallowed");
    unbreak_outbox(&pool, &tag).await;

    assert!(
        format!("{err}").contains("event emit"),
        "the error should name the emit failure, got: {err}"
    );
    assert_eq!(
        kv_count(&pool, f.app).await,
        1,
        "k2 must NOT be in the store — the write is rolled back with its fan-out. \
         A committed row whose trigger never fired is exactly the durability hole \
         the transactional outbox closes."
    );
    assert_eq!(
        outbox_count(&pool, f.app).await,
        1,
        "and no outbox row for the rolled-back write"
    );

    // --- An unwatched collection still writes (no trigger → no fan-out). ----
    writer
        .set(&cx, "unwatched", "k", serde_json::json!(3))
        .await
        .expect("a write with no matching trigger needs no outbox row");
    assert_eq!(kv_count(&pool, f.app).await, 2);
    assert_eq!(outbox_count(&pool, f.app).await, 1, "still no new fan-out");

    cleanup(&pool, f.app).await;
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn a_failed_fan_out_rolls_a_delete_back() {
    let Some(pool) = pool_or_skip().await else {
        return;
    };
    let f = setup(&pool, "watched").await;
    let writer = PostgresKvWriter::new(pool.clone());
    let cx = cx(f.app);
    let tag = Uuid::new_v4().simple().to_string();

    writer
        .set(&cx, "watched", "doomed", serde_json::json!(1))
        .await
        .expect("seed");
    assert_eq!(kv_count(&pool, f.app).await, 1);

    // A delete that can't fan out must leave the row in place — otherwise the
    // key is gone and nothing downstream ever learns it was removed.
    break_outbox_for(&pool, f.app, &tag).await;
    let err = writer
        .delete(&cx, "watched", "doomed")
        .await
        .expect_err("outbox failure surfaces");
    unbreak_outbox(&pool, &tag).await;
    assert!(format!("{err}").contains("event emit"));
    assert_eq!(
        kv_count(&pool, f.app).await,
        1,
        "the delete rolled back — the key is still there"
    );

    cleanup(&pool, f.app).await;
}