use rusqlite::Connection;
use std::path::Path;
use std::sync::Mutex;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LedgerJobStatus {
Running,
Completed,
Failed,
Cancelled,
}
impl LedgerJobStatus {
fn from_str(s: &str) -> Self {
match s {
"completed" => Self::Completed,
"failed" => Self::Failed,
"cancelled" => Self::Cancelled,
_ => Self::Running,
}
}
}
pub struct Ledger {
conn: Mutex<Connection>,
job_dir: std::path::PathBuf,
}
impl std::fmt::Debug for Ledger {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Ledger")
.field("job_dir", &self.job_dir)
.finish_non_exhaustive()
}
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
pub struct Escalation {
pub node: String,
pub item: Option<usize>,
pub reason: String,
pub input: Option<serde_json::Value>,
pub drained_at: Option<String>,
pub at: String,
}
#[derive(Debug, thiserror::Error)]
pub enum LedgerError {
#[error(transparent)]
Sqlite(#[from] rusqlite::Error),
#[error(
"this job's ledger predates per-item fan-out checkpoints and cannot be resumed \
— re-submit the job to start a fresh one"
)]
StaleSchema,
}
impl Ledger {
pub fn open(path: &Path, graph_fingerprint: &str) -> Result<Self, LedgerError> {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent).ok();
}
let conn = Connection::open(path)?;
conn.busy_timeout(std::time::Duration::from_secs(5))?;
let existing_columns: Vec<String> = {
let mut stmt = conn.prepare("PRAGMA table_info(checkpoints)")?;
let rows = stmt.query_map([], |r| r.get::<_, String>(1))?;
rows.collect::<rusqlite::Result<_>>()?
};
if !existing_columns.is_empty() && !existing_columns.iter().any(|c| c == "item_index") {
return Err(LedgerError::StaleSchema);
}
for (column, ddl) in [
(
"input_json",
"ALTER TABLE checkpoints ADD COLUMN input_json TEXT",
),
(
"drained_at",
"ALTER TABLE checkpoints ADD COLUMN drained_at TEXT",
),
] {
if !existing_columns.is_empty() && !existing_columns.iter().any(|c| c == column) {
conn.execute(ddl, [])?;
}
}
conn.execute_batch(
"CREATE TABLE IF NOT EXISTS checkpoints (
node_name TEXT NOT NULL,
item_index INTEGER NOT NULL DEFAULT -1,
status TEXT NOT NULL,
output_json TEXT,
error_text TEXT,
completed_at TEXT NOT NULL,
input_json TEXT,
drained_at TEXT,
PRIMARY KEY (node_name, item_index)
);
CREATE TABLE IF NOT EXISTS job_status (status TEXT NOT NULL, graph_fingerprint TEXT NOT NULL);
CREATE TABLE IF NOT EXISTS fanout_manifests (
node_name TEXT PRIMARY KEY,
digest TEXT NOT NULL,
item_count INTEGER NOT NULL
);",
)?;
let count: i64 = conn.query_row("SELECT COUNT(*) FROM job_status", [], |r| r.get(0))?;
if count == 0 {
conn.execute(
"INSERT INTO job_status (status, graph_fingerprint) VALUES ('running', ?1)",
[graph_fingerprint],
)?;
}
Ok(Self {
conn: Mutex::new(conn),
job_dir: path
.parent()
.unwrap_or_else(|| Path::new("."))
.to_path_buf(),
})
}
pub fn job_dir(&self) -> &Path {
&self.job_dir
}
pub fn graph_fingerprint(&self) -> rusqlite::Result<String> {
self.lock()
.query_row("SELECT graph_fingerprint FROM job_status", [], |r| r.get(0))
}
pub fn get_completed(&self, node_name: &str) -> rusqlite::Result<Option<serde_json::Value>> {
let result: Option<(String, Option<String>)> = match self.lock().query_row(
"SELECT status, output_json FROM checkpoints WHERE node_name = ?1 AND item_index = -1",
[node_name],
|r| Ok((r.get(0)?, r.get(1)?)),
) {
Ok(row) => Some(row),
Err(rusqlite::Error::QueryReturnedNoRows) => None,
Err(e) => return Err(e),
};
match result {
Some((status, Some(json))) if status == "completed" => Ok(Some(
serde_json::from_str(&json).expect("ledger never stores invalid JSON"),
)),
_ => Ok(None),
}
}
pub fn is_skipped(&self, node_name: &str) -> rusqlite::Result<bool> {
let status: Option<String> = match self.lock().query_row(
"SELECT status FROM checkpoints WHERE node_name = ?1 AND item_index = -1",
[node_name],
|r| r.get(0),
) {
Ok(status) => Some(status),
Err(rusqlite::Error::QueryReturnedNoRows) => None,
Err(e) => return Err(e),
};
Ok(status.as_deref() == Some("skipped"))
}
pub fn write_completed(
&self,
node_name: &str,
output: &serde_json::Value,
) -> rusqlite::Result<()> {
self.lock().execute(
"INSERT OR REPLACE INTO checkpoints
(node_name, item_index, status, output_json, error_text, completed_at)
VALUES (?1, -1, 'completed', ?2, NULL, ?3)",
rusqlite::params![node_name, output.to_string(), now_marker()],
)?;
Ok(())
}
pub fn write_skipped(&self, node_name: &str) -> rusqlite::Result<()> {
self.lock().execute(
"INSERT OR REPLACE INTO checkpoints
(node_name, item_index, status, output_json, error_text, completed_at)
VALUES (?1, -1, 'skipped', NULL, NULL, ?2)",
rusqlite::params![node_name, now_marker()],
)?;
Ok(())
}
pub fn write_item_completed(
&self,
node_name: &str,
item_index: usize,
output: &serde_json::Value,
input: Option<&serde_json::Value>,
) -> rusqlite::Result<()> {
self.lock().execute(
"INSERT OR REPLACE INTO checkpoints
(node_name, item_index, status, output_json, error_text, completed_at, input_json)
VALUES (?1, ?2, 'completed', ?3, NULL, ?4, ?5)",
rusqlite::params![
node_name,
item_index as i64,
output.to_string(),
now_marker(),
input.map(|v| v.to_string())
],
)?;
Ok(())
}
pub fn write_item_failed(
&self,
node_name: &str,
item_index: usize,
error: &str,
input: Option<&serde_json::Value>,
) -> rusqlite::Result<()> {
self.lock().execute(
"INSERT OR REPLACE INTO checkpoints
(node_name, item_index, status, output_json, error_text, completed_at, input_json)
VALUES (?1, ?2, 'failed', NULL, ?3, ?4, ?5)",
rusqlite::params![
node_name,
item_index as i64,
error,
now_marker(),
input.map(|v| v.to_string())
],
)?;
Ok(())
}
pub fn get_item_completed(
&self,
node_name: &str,
item_index: usize,
) -> rusqlite::Result<Option<serde_json::Value>> {
let row: Option<(String, Option<String>)> = match self.lock().query_row(
"SELECT status, output_json FROM checkpoints
WHERE node_name = ?1 AND item_index = ?2",
rusqlite::params![node_name, item_index as i64],
|r| Ok((r.get(0)?, r.get(1)?)),
) {
Ok(row) => Some(row),
Err(rusqlite::Error::QueryReturnedNoRows) => None,
Err(e) => return Err(e),
};
Ok(match row {
Some((status, Some(json))) if status == "completed" => {
Some(serde_json::from_str(&json).expect("ledger never stores invalid JSON"))
}
_ => None,
})
}
pub fn item_concluded(&self, node_name: &str, item_index: usize) -> rusqlite::Result<bool> {
let count: i64 = self.lock().query_row(
"SELECT COUNT(*) FROM checkpoints WHERE node_name = ?1 AND item_index = ?2",
rusqlite::params![node_name, item_index as i64],
|r| r.get(0),
)?;
Ok(count > 0)
}
#[allow(clippy::type_complexity)]
pub fn concluded_items(
&self,
node_name: &str,
) -> rusqlite::Result<Vec<(usize, Option<serde_json::Value>, Option<String>)>> {
let conn = self.lock();
let mut stmt = conn.prepare(
"SELECT item_index, status, output_json, error_text FROM checkpoints
WHERE node_name = ?1 AND item_index >= 0 ORDER BY item_index",
)?;
let rows = stmt.query_map([node_name], |r| {
let index: i64 = r.get(0)?;
let status: String = r.get(1)?;
let output: Option<String> = r.get(2)?;
let error: Option<String> = r.get(3)?;
Ok((
index as usize,
if status == "completed" {
output.map(|j| {
serde_json::from_str(&j).expect("ledger never stores invalid JSON")
})
} else {
None
},
if status == "completed" { None } else { error },
))
})?;
rows.collect()
}
pub fn concluded_rows(&self, node_name: &str) -> rusqlite::Result<Vec<ConcludedRow>> {
let conn = self.lock();
let mut stmt = conn.prepare(
"SELECT item_index, status, output_json, error_text, completed_at, input_json
FROM checkpoints
WHERE node_name = ?1 AND item_index >= 0 ORDER BY item_index",
)?;
let rows = stmt.query_map([node_name], |r| {
let status: String = r.get(1)?;
let output: Option<String> = r.get(2)?;
Ok(ConcludedRow {
item: r.get::<_, i64>(0)?,
output: if status == "completed" {
output.map(|j| {
serde_json::from_str(&j).expect("ledger never stores invalid JSON")
})
} else {
None
},
error: if status == "completed" {
None
} else {
r.get(3)?
},
status,
concluded_at: r.get(4)?,
source_input: r.get(5)?,
})
})?;
rows.collect()
}
pub fn write_escalated(
&self,
node_name: &str,
item_index: Option<usize>,
reason: &str,
input: Option<&serde_json::Value>,
) -> rusqlite::Result<()> {
self.lock().execute(
"INSERT OR REPLACE INTO checkpoints
(node_name, item_index, status, output_json, error_text, completed_at, input_json)
VALUES (?1, ?2, 'escalated', NULL, ?3, ?4, ?5)",
rusqlite::params![
node_name,
item_index.map(|i| i as i64).unwrap_or(-1),
reason,
now_marker(),
input.map(|v| v.to_string())
],
)?;
Ok(())
}
pub fn open_read_only(path: &Path) -> Result<Self, LedgerError> {
let conn = Connection::open_with_flags(
path,
rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY | rusqlite::OpenFlags::SQLITE_OPEN_URI,
)?;
conn.busy_timeout(std::time::Duration::from_secs(5))?;
Ok(Self {
conn: Mutex::new(conn),
job_dir: path
.parent()
.unwrap_or_else(|| Path::new("."))
.to_path_buf(),
})
}
fn has_drain_columns(&self) -> bool {
let conn = self.lock();
let Ok(mut stmt) = conn.prepare("PRAGMA table_info(checkpoints)") else {
return false;
};
let Ok(rows) = stmt.query_map([], |r| r.get::<_, String>(1)) else {
return false;
};
let columns: Vec<String> = rows.filter_map(Result::ok).collect();
columns.iter().any(|c| c == "input_json") && columns.iter().any(|c| c == "drained_at")
}
pub fn escalations(&self) -> rusqlite::Result<Vec<Escalation>> {
self.query_escalations(true)
}
pub fn all_escalations(&self) -> rusqlite::Result<Vec<Escalation>> {
self.query_escalations(false)
}
fn query_escalations(&self, outstanding_only: bool) -> rusqlite::Result<Vec<Escalation>> {
if !self.has_drain_columns() {
let conn = self.lock();
let mut stmt = conn.prepare(
"SELECT node_name, item_index, error_text, completed_at FROM checkpoints
WHERE status = 'escalated' ORDER BY node_name, item_index",
)?;
let rows = stmt.query_map([], |r| {
let index: i64 = r.get(1)?;
Ok(Escalation {
node: r.get(0)?,
item: (index >= 0).then_some(index as usize),
reason: r.get::<_, Option<String>>(2)?.unwrap_or_default(),
at: r.get(3)?,
input: None,
drained_at: None,
})
})?;
return rows.collect();
}
let conn = self.lock();
let mut stmt = conn.prepare(
"SELECT node_name, item_index, error_text, completed_at, input_json, drained_at
FROM checkpoints
WHERE status = 'escalated' AND (?1 = 0 OR drained_at IS NULL)
ORDER BY node_name, item_index",
)?;
let rows = stmt.query_map([i64::from(outstanding_only)], |r| {
let index: i64 = r.get(1)?;
let input: Option<String> = r.get(4)?;
Ok(Escalation {
node: r.get(0)?,
item: (index >= 0).then_some(index as usize),
reason: r.get::<_, Option<String>>(2)?.unwrap_or_default(),
at: r.get(3)?,
input: input.and_then(|j| serde_json::from_str(&j).ok()),
drained_at: r.get(5)?,
})
})?;
rows.collect()
}
pub fn mark_drained(&self, node_name: &str, item_index: Option<usize>) -> rusqlite::Result<()> {
self.lock().execute(
"UPDATE checkpoints SET drained_at = ?3
WHERE node_name = ?1 AND item_index = ?2 AND status = 'escalated'",
rusqlite::params![
node_name,
item_index.map(|i| i as i64).unwrap_or(-1),
now_marker()
],
)?;
Ok(())
}
pub fn check_or_record_manifest(
&self,
node_name: &str,
digest: &str,
item_count: usize,
) -> rusqlite::Result<Result<(), String>> {
let conn = self.lock();
let existing: Option<String> = match conn.query_row(
"SELECT digest FROM fanout_manifests WHERE node_name = ?1",
[node_name],
|r| r.get(0),
) {
Ok(d) => Some(d),
Err(rusqlite::Error::QueryReturnedNoRows) => None,
Err(e) => return Err(e),
};
match existing {
Some(previous) if previous != digest => Ok(Err(previous)),
Some(_) => Ok(Ok(())),
None => {
conn.execute(
"INSERT INTO fanout_manifests (node_name, digest, item_count)
VALUES (?1, ?2, ?3)",
rusqlite::params![node_name, digest, item_count as i64],
)?;
Ok(Ok(()))
}
}
}
pub fn job_status(&self) -> rusqlite::Result<LedgerJobStatus> {
let s: String = self
.lock()
.query_row("SELECT status FROM job_status", [], |r| r.get(0))?;
Ok(LedgerJobStatus::from_str(&s))
}
pub fn finish(&self, status: &str) -> rusqlite::Result<()> {
self.lock()
.execute("UPDATE job_status SET status = ?1", [status])?;
Ok(())
}
fn lock(&self) -> std::sync::MutexGuard<'_, Connection> {
self.conn.lock().expect("ledger connection mutex poisoned")
}
}
pub fn jobs_root() -> Option<std::path::PathBuf> {
if let Ok(dir) = std::env::var("CUTTLEFISH_JOBS_HOME") {
return Some(std::path::PathBuf::from(dir));
}
crate::catalog::cuttlefish_home().map(|h| h.join("jobs"))
}
fn now_marker() -> String {
crate::catalog::now_rfc3339()
}
#[derive(Debug, Clone)]
pub struct ConcludedRow {
pub item: i64,
pub status: String,
pub output: Option<serde_json::Value>,
pub error: Option<String>,
pub concluded_at: String,
pub source_input: Option<String>,
}