@kromatv/sdk 0.1.40-canary.0

package/rust/kroma-db/src/pipeline/query.rs

@@ -0,0 +1,255 @@
+//! Ledger read / status queries: per-stage tallies, per-subject status rollups,
+//! and the lean row-mappers that back the pipeline elements list. All read-only.
+
+use std::collections::HashMap;
+
+use anyhow::Result;
+use rusqlite::params;
+
+use crate::pool::Pool;
+use kroma_domain::{PipelineTaskView, StageStat};
+
+mod elements;
+
+#[cfg(test)]
+mod test_support;
+
+pub use elements::*;
+
+/// Per-stage status tally `(pending, running, done, failed, blocked)`.
+pub fn counts(pool: &Pool, stage: &str) -> Result<(i64, i64, i64, i64, i64)> {
+    let conn = pool.get()?;
+    let mut stmt =
+        conn.prepare("SELECT status, COUNT(*) FROM pipeline_tasks WHERE stage=?1 GROUP BY status")?;
+    let rows = stmt.query_map(params![stage], |r| {
+        Ok((r.get::<_, String>(0)?, r.get::<_, i64>(1)?))
+    })?;
+    let mut c = [0i64; 5]; // pending, running, done, failed, blocked
+    for row in rows {
+        let (st, n) = row?;
+        match st.as_str() {
+            "pending" => c[0] = n,
+            "running" => c[1] = n,
+            "done" => c[2] = n,
+            "failed" => c[3] = n,
+            "blocked" => c[4] = n,
+            _ => {}
+        }
+    }
+    Ok((c[0], c[1], c[2], c[3], c[4]))
+}
+
+/// The `StageStat` for one stage (counts + identity), for the API + WS event.
+pub fn stage_stat(pool: &Pool, stage: &str, key: &str, subject_kind: &str) -> Result<StageStat> {
+    let (pending, running, done, failed, blocked) = counts(pool, stage)?;
+    Ok(StageStat {
+        stage: stage.to_string(),
+        key: key.to_string(),
+        subject_kind: subject_kind.to_string(),
+        pending,
+        running,
+        done,
+        failed,
+        blocked,
+    })
+}
+
+/// Every task of one stage as `subject_id -> (status, error)`. Bulk map for the
+/// pipeline elements list (overlays the ledger's running/failed/pending states,
+/// with the failure message, onto the cheap artifact signals).
+pub fn stage_statuses(
+    pool: &Pool,
+    stage: &str,
+) -> Result<HashMap<String, (String, Option<String>)>> {
+    let conn = pool.get()?;
+    let mut stmt =
+        conn.prepare("SELECT subject_id, status, error FROM pipeline_tasks WHERE stage=?1")?;
+    let rows = stmt.query_map(params![stage], |r| {
+        Ok((
+            r.get::<_, String>(0)?,
+            (r.get::<_, String>(1)?, r.get::<_, Option<String>>(2)?),
+        ))
+    })?;
+    Ok(rows.collect::<rusqlite::Result<HashMap<_, _>>>()?)
+}
+
+/// The ledger status of one task, or `None` if no task exists for it yet.
+pub fn task_status(pool: &Pool, stage: &str, subject_id: &str) -> Result<Option<String>> {
+    let conn = pool.get()?;
+    let mut stmt =
+        conn.prepare("SELECT status FROM pipeline_tasks WHERE stage=?1 AND subject_id=?2")?;
+    let mut rows = stmt.query_map(params![stage, subject_id], |r| r.get::<_, String>(0))?;
+    match rows.next() {
+        Some(r) => Ok(Some(r?)),
+        None => Ok(None),
+    }
+}
+
+/// The "worst" ledger status across several subjects of one stage, by severity
+/// (`failed` > `running` > `pending` > `done`), or `None` if none has a task.
+/// Used to roll a show's per-episode/per-file tasks up to one treatment state.
+/// ONE query over the whole subject set (was N+1: a connection+query per id),
+/// backed by the `(stage, subject_id)` index.
+pub fn worst_status(pool: &Pool, stage: &str, subject_ids: &[String]) -> Result<Option<String>> {
+    if subject_ids.is_empty() {
+        return Ok(None);
+    }
+    let rank = |s: &str| match s {
+        "failed" => 4,
+        "running" => 3,
+        "pending" => 2,
+        "done" => 1,
+        _ => 0,
+    };
+    let conn = pool.get()?;
+    let placeholders = vec!["?"; subject_ids.len()].join(",");
+    let sql = format!(
+        "SELECT subject_id, status FROM pipeline_tasks \
+         WHERE stage=? AND subject_id IN ({placeholders})"
+    );
+    let mut stmt = conn.prepare(&sql)?;
+    let mut binds: Vec<&str> = Vec::with_capacity(subject_ids.len() + 1);
+    binds.push(stage);
+    binds.extend(subject_ids.iter().map(String::as_str));
+    let rows = stmt.query_map(rusqlite::params_from_iter(binds), |r| r.get::<_, String>(1))?;
+    let mut worst: Option<String> = None;
+    for st in rows {
+        let st = st?;
+        if worst.as_deref().is_none_or(|w| rank(&st) > rank(w)) {
+            worst = Some(st);
+        }
+    }
+    Ok(worst)
+}
+
+/// Failed tasks for a stage's drill-down (newest failure first). `title` is left
+/// as the raw id here; the API layer resolves it against the catalog.
+pub fn failed_tasks(pool: &Pool, stage: &str, limit: usize) -> Result<Vec<PipelineTaskView>> {
+    let conn = pool.get()?;
+    let mut stmt = conn.prepare(
+        "SELECT subject_kind, subject_id, status, attempts, error, finished_at \
+         FROM pipeline_tasks WHERE stage=?1 AND status='failed' \
+         ORDER BY finished_at DESC LIMIT ?2",
+    )?;
+    let rows = stmt.query_map(params![stage, limit as i64], |r| {
+        let subject_id: String = r.get(1)?;
+        Ok(PipelineTaskView {
+            stage: stage.to_string(),
+            subject_kind: r.get(0)?,
+            title: subject_id.clone(),
+            subject_id,
+            status: r.get(2)?,
+            attempts: r.get(3)?,
+            error: r.get(4)?,
+            finished_at: r.get(5)?,
+        })
+    })?;
+    Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::pipeline::query::test_support::*;
+
+    #[test]
+    fn stage_stat_statuses_and_worst() {
+        let p = pool();
+        {
+            let conn = p.get().unwrap();
+            task(&conn, "st", "f1", "done", None, Some(100));
+            task(&conn, "st", "f2", "pending", None, None);
+            task(&conn, "st", "f3", "failed", Some("boom"), Some(500));
+            task(&conn, "st", "f4", "running", None, None);
+        }
+        let stat = stage_stat(&p, "st", "the-key", "file").unwrap();
+        assert_eq!(
+            (
+                stat.pending,
+                stat.running,
+                stat.done,
+                stat.failed,
+                stat.blocked
+            ),
+            (1, 1, 1, 1, 0)
+        );
+        assert_eq!(stat.key, "the-key");
+
+        let map = stage_statuses(&p, "st").unwrap();
+        assert_eq!(map["f3"], ("failed".to_string(), Some("boom".to_string())));
+        assert_eq!(map["f1"].0, "done");
+
+        assert_eq!(
+            task_status(&p, "st", "f3").unwrap().as_deref(),
+            Some("failed")
+        );
+        assert!(task_status(&p, "st", "ghost").unwrap().is_none());
+
+        // worst_status ranks failed > running > pending > done.
+        assert_eq!(
+            worst_status(&p, "st", &["f1".into(), "f2".into(), "f3".into()])
+                .unwrap()
+                .as_deref(),
+            Some("failed")
+        );
+        assert_eq!(
+            worst_status(&p, "st", &["f1".into()]).unwrap().as_deref(),
+            Some("done")
+        );
+        assert!(worst_status(&p, "st", &[]).unwrap().is_none());
+        assert!(worst_status(&p, "st", &["ghost".into()]).unwrap().is_none());
+    }
+
+    #[test]
+    fn a_blocked_task_is_counted_and_ranks_below_every_named_status() {
+        let p = pool();
+        {
+            let conn = p.get().unwrap();
+            task(&conn, "st", "f1", "blocked", None, None);
+            task(&conn, "st", "f2", "blocked", None, None);
+            task(&conn, "st", "f3", "done", None, Some(1));
+            task(&conn, "st", "f4", "quarantined", None, None);
+        }
+        let stat = stage_stat(&p, "st", "k", "file").unwrap();
+        assert_eq!(
+            (
+                stat.pending,
+                stat.running,
+                stat.done,
+                stat.failed,
+                stat.blocked
+            ),
+            (0, 0, 1, 0, 2)
+        );
+
+        assert_eq!(
+            worst_status(&p, "st", &["f1".into(), "f3".into()])
+                .unwrap()
+                .as_deref(),
+            Some("done")
+        );
+        assert_eq!(
+            worst_status(&p, "st", &["f4".into()]).unwrap().as_deref(),
+            Some("quarantined"),
+            "the only row is still the worst one"
+        );
+    }
+
+    #[test]
+    fn failed_tasks_newest_first() {
+        let p = pool();
+        {
+            let conn = p.get().unwrap();
+            task(&conn, "st", "old", "failed", Some("e1"), Some(100));
+            task(&conn, "st", "new", "failed", Some("e2"), Some(900));
+            task(&conn, "st", "ok", "done", None, Some(500));
+        }
+        let failed = failed_tasks(&p, "st", 10).unwrap();
+        // Only failed rows, newest finished_at first.
+        assert_eq!(failed.len(), 2);
+        assert_eq!(failed[0].subject_id, "new");
+        assert_eq!(failed[0].title, "new"); // title defaults to the raw id
+        assert_eq!(failed[0].error.as_deref(), Some("e2"));
+        assert_eq!(failed[1].subject_id, "old");
+    }
+}

package/rust/kroma-db/src/pipeline/tests.rs

@@ -0,0 +1,319 @@
+use super::ops::{retry_backoff_ms, Subject, MAX_ATTEMPTS};
+use super::*;
+use crate::testing::TempPool;
+use crate::Pool;
+
+fn pool() -> TempPool {
+    crate::testing::temp_pool("pipe")
+}
+
+// `(pending, running, done, failed, blocked)` for the test stage.
+fn c(p: &Pool) -> (i64, i64, i64, i64, i64) {
+    counts(p, "s").unwrap()
+}
+
+fn subj(pairs: &[(&str, &str)]) -> Vec<Subject> {
+    pairs
+        .iter()
+        .map(|(id, sig)| (id.to_string(), sig.to_string()))
+        .collect()
+}
+
+#[test]
+fn reconcile_is_incremental_and_idempotent() {
+    let p = pool();
+    let s = subj(&[("a", "v1"), ("b", "v1")]);
+
+    reconcile(&p, "s", "item", &s, 1).unwrap();
+    assert_eq!(c(&p), (2, 0, 0, 0, 0));
+
+    let batch = claim_batch(&p, "s", 10, 2).unwrap();
+    assert_eq!(batch.len(), 2);
+    let ok: Vec<TaskResult> = batch
+        .iter()
+        .map(|(id, _)| TaskResult {
+            id: id.clone(),
+            error: None,
+            duration_ms: 1,
+        })
+        .collect();
+    finish_batch(&p, "s", &ok, 3).unwrap();
+    assert_eq!(c(&p), (0, 0, 2, 0, 0));
+
+    reconcile(&p, "s", "item", &s, 4).unwrap();
+    assert_eq!(c(&p), (0, 0, 2, 0, 0));
+    assert!(claim_batch(&p, "s", 10, 5).unwrap().is_empty());
+
+    reconcile(&p, "s", "item", &subj(&[("a", "v2"), ("b", "v1")]), 6).unwrap();
+    assert_eq!(c(&p), (1, 0, 1, 0, 0));
+    let re = claim_batch(&p, "s", 10, 7).unwrap();
+    assert_eq!(re, vec![("a".to_string(), "v2".to_string())]);
+
+    reconcile(&p, "s", "item", &subj(&[("b", "v1")]), 8).unwrap();
+    let all: i64 = {
+        let conn = p.get().unwrap();
+        conn.query_row(
+            "SELECT COUNT(*) FROM pipeline_tasks WHERE subject_id='a'",
+            [],
+            |r| r.get(0),
+        )
+        .unwrap()
+    };
+    assert_eq!(all, 0);
+}
+
+#[test]
+fn a_task_a_worker_already_claimed_is_left_alone_when_its_input_changes() {
+    let p = pool();
+    reconcile(&p, "s", "item", &subj(&[("a", "v1")]), 1).unwrap();
+    assert_eq!(
+        claim_batch(&p, "s", 10, 2).unwrap(),
+        vec![("a".to_string(), "v1".to_string())]
+    );
+    assert_eq!(c(&p), (0, 1, 0, 0, 0));
+
+    reconcile(&p, "s", "item", &subj(&[("a", "v2")]), 3).unwrap();
+    assert_eq!(c(&p), (0, 1, 0, 0, 0));
+    let sig: Option<String> = p
+        .get()
+        .unwrap()
+        .query_row(
+            "SELECT input_sig FROM pipeline_tasks WHERE stage='s' AND subject_id='a'",
+            [],
+            |r| r.get(0),
+        )
+        .unwrap();
+    assert_eq!(sig.as_deref(), Some("v1"));
+}
+
+#[test]
+fn failures_retry_up_to_max_then_stick() {
+    let p = pool();
+    reconcile(&p, "s", "item", &subj(&[("a", "v1")]), 1).unwrap();
+
+    for i in 0..MAX_ATTEMPTS {
+        let batch = claim_batch(&p, "s", 10, 10 + i).unwrap();
+        assert_eq!(
+            batch.len(),
+            1,
+            "attempt {i} should have a pending task to claim"
+        );
+        let failed_at = 20 + i;
+        finish_batch(
+            &p,
+            "s",
+            &[TaskResult {
+                id: "a".into(),
+                error: Some("boom".into()),
+                duration_ms: 1,
+            }],
+            failed_at,
+        )
+        .unwrap();
+        reconcile(
+            &p,
+            "s",
+            "item",
+            &subj(&[("a", "v1")]),
+            failed_at + retry_backoff_ms(i + 1),
+        )
+        .unwrap();
+    }
+    assert_eq!(c(&p), (0, 0, 0, 1, 0));
+    assert!(claim_batch(&p, "s", 10, 99).unwrap().is_empty());
+
+    assert_eq!(retry(&p, "s", Some("a")).unwrap(), 1);
+    assert_eq!(c(&p), (1, 0, 0, 0, 0));
+}
+
+#[test]
+fn auto_retry_waits_for_the_backoff_window() {
+    let p = pool();
+    reconcile(&p, "s", "item", &subj(&[("a", "v1")]), 1).unwrap();
+    claim_batch(&p, "s", 10, 2).unwrap();
+    finish_batch(
+        &p,
+        "s",
+        &[TaskResult {
+            id: "a".into(),
+            error: Some("boom".into()),
+            duration_ms: 1,
+        }],
+        3,
+    )
+    .unwrap();
+
+    reconcile(&p, "s", "item", &subj(&[("a", "v1")]), 4).unwrap();
+    assert_eq!(c(&p), (0, 0, 0, 1, 0));
+
+    reconcile(
+        &p,
+        "s",
+        "item",
+        &subj(&[("a", "v1")]),
+        3 + retry_backoff_ms(1),
+    )
+    .unwrap();
+    assert_eq!(c(&p), (1, 0, 0, 0, 0));
+}
+
+#[test]
+fn changed_signature_requeues_despite_backoff() {
+    let p = pool();
+    reconcile(&p, "s", "item", &subj(&[("a", "v1")]), 1).unwrap();
+    claim_batch(&p, "s", 10, 2).unwrap();
+    finish_batch(
+        &p,
+        "s",
+        &[TaskResult {
+            id: "a".into(),
+            error: Some("boom".into()),
+            duration_ms: 1,
+        }],
+        3,
+    )
+    .unwrap();
+
+    reconcile(&p, "s", "item", &subj(&[("a", "v2")]), 4).unwrap();
+    assert_eq!(c(&p), (1, 0, 0, 0, 0));
+}
+
+#[test]
+fn manual_retry_ignores_the_backoff_window() {
+    let p = pool();
+    reconcile(&p, "s", "item", &subj(&[("a", "v1")]), 1).unwrap();
+    claim_batch(&p, "s", 10, 2).unwrap();
+    finish_batch(
+        &p,
+        "s",
+        &[TaskResult {
+            id: "a".into(),
+            error: Some("boom".into()),
+            duration_ms: 1,
+        }],
+        3,
+    )
+    .unwrap();
+
+    assert_eq!(retry(&p, "s", Some("a")).unwrap(), 1);
+    assert_eq!(c(&p), (1, 0, 0, 0, 0));
+    let batch = claim_batch(&p, "s", 10, 4).unwrap();
+    assert_eq!(batch.len(), 1);
+}
+
+#[test]
+fn manual_retry_jumps_the_queue() {
+    let p = pool();
+    reconcile(&p, "s", "item", &subj(&[("routine", "v1")]), 1).unwrap();
+    {
+        let conn = p.get().unwrap();
+        conn.execute(
+            "INSERT INTO pipeline_tasks(stage,subject_kind,subject_id,status,attempts,priority,enqueued_at,updated_at) \
+             VALUES ('s','item','boom','failed',1,0,0,0)",
+            [],
+        )
+        .unwrap();
+    }
+    assert_eq!(retry(&p, "s", Some("boom")).unwrap(), 1);
+    let batch = claim_batch(&p, "s", 1, 5).unwrap();
+    assert_eq!(
+        batch.iter().map(|(id, _)| id.as_str()).collect::<Vec<_>>(),
+        vec!["boom"]
+    );
+}
+
+#[test]
+fn enqueue_null_sig_is_backfilled_not_rerun() {
+    let p = pool();
+    // `enqueue` inserts with input_sig = NULL.
+    enqueue(&p, "s", "item", "a", 100, 1).unwrap();
+    let batch = claim_batch(&p, "s", 10, 2).unwrap();
+    assert_eq!(batch.len(), 1);
+    finish_batch(
+        &p,
+        "s",
+        &[TaskResult {
+            id: "a".into(),
+            error: None,
+            duration_ms: 1,
+        }],
+        3,
+    )
+    .unwrap();
+    assert_eq!(c(&p), (0, 0, 1, 0, 0));
+
+    reconcile(&p, "s", "item", &subj(&[("a", "v1")]), 4).unwrap();
+    assert_eq!(c(&p), (0, 0, 1, 0, 0));
+    assert!(claim_batch(&p, "s", 10, 5).unwrap().is_empty());
+
+    reconcile(&p, "s", "item", &subj(&[("a", "v1")]), 6).unwrap();
+    assert_eq!(c(&p), (0, 0, 1, 0, 0));
+    reconcile(&p, "s", "item", &subj(&[("a", "v2")]), 7).unwrap();
+    assert_eq!(c(&p), (1, 0, 0, 0, 0));
+}
+
+#[test]
+fn reset_running_recovers_stranded_tasks() {
+    let p = pool();
+    reconcile(&p, "s", "item", &subj(&[("a", "v1"), ("b", "v1")]), 1).unwrap();
+    claim_batch(&p, "s", 10, 2).unwrap();
+    assert_eq!(c(&p), (0, 2, 0, 0, 0));
+    assert_eq!(reset_running(&p, None).unwrap(), 2);
+    assert_eq!(c(&p), (2, 0, 0, 0, 0));
+}
+
+#[test]
+fn unreadable_signature_never_requeues_or_deletes() {
+    let p = pool();
+    reconcile(&p, "s", "item", &subj(&[("a", "v1")]), 1).unwrap();
+    let batch = claim_batch(&p, "s", 10, 2).unwrap();
+    let ok: Vec<TaskResult> = batch
+        .iter()
+        .map(|(id, _)| TaskResult {
+            id: id.clone(),
+            error: None,
+            duration_ms: 1,
+        })
+        .collect();
+    finish_batch(&p, "s", &ok, 3).unwrap();
+    assert_eq!(c(&p), (0, 0, 1, 0, 0));
+
+    reconcile(&p, "s", "item", &subj(&[("a", UNREADABLE_SIG)]), 4).unwrap();
+    assert_eq!(c(&p), (0, 0, 1, 0, 0));
+    assert!(claim_batch(&p, "s", 10, 5).unwrap().is_empty());
+
+    // The stored sig was never overwritten by the sentinel, so no recompute.
+    reconcile(&p, "s", "item", &subj(&[("a", "v1")]), 6).unwrap();
+    assert_eq!(c(&p), (0, 0, 1, 0, 0));
+    assert!(claim_batch(&p, "s", 10, 7).unwrap().is_empty());
+}
+
+#[test]
+fn requeue_stage_rebuilds_done_tasks_after_cache_wipe() {
+    let p = pool();
+    reconcile(&p, "s", "item", &subj(&[("a", "v1"), ("b", "v1")]), 1).unwrap();
+    let batch = claim_batch(&p, "s", 10, 2).unwrap();
+    let ok: Vec<TaskResult> = batch
+        .iter()
+        .map(|(id, _)| TaskResult {
+            id: id.clone(),
+            error: None,
+            duration_ms: 1,
+        })
+        .collect();
+    finish_batch(&p, "s", &ok, 3).unwrap();
+    assert_eq!(c(&p), (0, 0, 2, 0, 0));
+
+    assert_eq!(requeue_stage(&p, "s", 4).unwrap(), 2);
+    assert_eq!(c(&p), (2, 0, 0, 0, 0));
+    assert_eq!(claim_batch(&p, "s", 10, 5).unwrap().len(), 2);
+}
+
+#[test]
+fn unreadable_brand_new_subject_is_deferred_not_inserted() {
+    let p = pool();
+    reconcile(&p, "s", "item", &subj(&[("a", UNREADABLE_SIG)]), 1).unwrap();
+    assert_eq!(c(&p), (0, 0, 0, 0, 0));
+    reconcile(&p, "s", "item", &subj(&[("a", "v1")]), 2).unwrap();
+    assert_eq!(c(&p), (1, 0, 0, 0, 0));
+}

package/rust/kroma-db/src/pipeline.rs

@@ -0,0 +1,26 @@
+//! Persistence for the per-element processing ledger (`pipeline_tasks`). See
+//! `crate::services::pipeline`. One row per `(stage, subject)`; all writes are
+//! batched into single transactions so the many stage workers never contend on
+//! SQLite's single writer (the dispatcher owns every write here). Timestamps are
+//! epoch milliseconds.
+//!
+//! The module splits into two concerns, both re-exported flat here so the public
+//! `crate::pipeline::<item>` paths resolve unchanged:
+//! - [`ops`]: ledger write / lifecycle ops (reconcile, enqueue, claim, finish,
+//!   reset, retry, reprocess) plus their shared consts/types.
+//! - [`query`]: read-only status queries and the lean elements-list row-mappers.
+
+mod ops;
+mod query;
+
+#[cfg(test)]
+mod tests;
+
+pub use ops::{
+    claim_batch, enqueue, finish_batch, reconcile, reprocess, requeue_stage, reset_running, retry,
+    retry_backoff_ms, TaskResult, UNREADABLE_SIG,
+};
+pub use query::{
+    counts, failed_tasks, item_titles, raw_items, raw_shows, show_titles, stage_stat,
+    stage_statuses, task_status, worst_status, RawItem,
+};