344 lines
12 KiB
Rust
344 lines
12 KiB
Rust
//! Enhance jobs (#192): one row per ComfyUI run submitted for a photo.
|
|
//!
|
|
//! The row is the only state: the poller advances it from ComfyUI's
|
|
//! `/queue` and `/history`, the endpoints read it, approve appends to
|
|
//! `saved`. `pending` / `saved` hold small JSON documents owned by
|
|
//! `crate::enhance`; this layer treats them as opaque text.
|
|
|
|
use std::sync::{Arc, Mutex};
|
|
|
|
use diesel::prelude::*;
|
|
use diesel::sql_query;
|
|
use diesel::sql_types::{BigInt, Integer, Nullable, Text};
|
|
use diesel::sqlite::SqliteConnection;
|
|
|
|
use crate::database::{DbError, DbErrorKind, connect};
|
|
use crate::otel::trace_db_call;
|
|
|
|
pub const STATUS_QUEUED: &str = "queued";
|
|
pub const STATUS_RUNNING: &str = "running";
|
|
pub const STATUS_READY: &str = "ready";
|
|
pub const STATUS_APPROVED: &str = "approved";
|
|
pub const STATUS_FAILED: &str = "failed";
|
|
pub const STATUS_EXPIRED: &str = "expired";
|
|
pub const STATUS_UNAVAILABLE: &str = "unavailable";
|
|
|
|
#[derive(Debug, Clone, PartialEq, Eq, QueryableByName)]
|
|
pub struct EnhanceJobRow {
|
|
#[diesel(sql_type = Text)]
|
|
pub id: String,
|
|
#[diesel(sql_type = Text)]
|
|
pub original_hash: String,
|
|
#[diesel(sql_type = Integer)]
|
|
pub library_id: i32,
|
|
#[diesel(sql_type = Text)]
|
|
pub rel_path: String,
|
|
#[diesel(sql_type = Text)]
|
|
pub template: String,
|
|
#[diesel(sql_type = Nullable<Text>)]
|
|
pub prompt: Option<String>,
|
|
#[diesel(sql_type = Nullable<Text>)]
|
|
pub reference_hash: Option<String>,
|
|
#[diesel(sql_type = Nullable<Text>)]
|
|
pub reference_path: Option<String>,
|
|
#[diesel(sql_type = Text)]
|
|
pub comfy_prompt_id: String,
|
|
#[diesel(sql_type = Text)]
|
|
pub status: String,
|
|
#[diesel(sql_type = Nullable<Text>)]
|
|
pub error: Option<String>,
|
|
#[diesel(sql_type = Nullable<Integer>)]
|
|
pub queue_position: Option<i32>,
|
|
#[diesel(sql_type = Text)]
|
|
pub pending: String,
|
|
#[diesel(sql_type = Text)]
|
|
pub saved: String,
|
|
#[diesel(sql_type = Integer)]
|
|
pub discarded: i32,
|
|
#[diesel(sql_type = BigInt)]
|
|
pub created_at: i64,
|
|
#[diesel(sql_type = BigInt)]
|
|
pub updated_at: i64,
|
|
}
|
|
|
|
const COLUMNS: &str = "id, original_hash, library_id, rel_path, template, prompt, reference_hash, \
|
|
reference_path, comfy_prompt_id, status, error, queue_position, pending, saved, \
|
|
discarded, created_at, updated_at";
|
|
|
|
/// What the poller / approve changes on an existing job.
|
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
|
pub struct JobUpdate {
|
|
pub status: String,
|
|
pub error: Option<String>,
|
|
pub queue_position: Option<i32>,
|
|
/// `Some` replaces the `pending` document.
|
|
pub pending: Option<String>,
|
|
/// `Some` replaces the `saved` document.
|
|
pub saved: Option<String>,
|
|
}
|
|
|
|
impl JobUpdate {
|
|
pub fn status(status: &str) -> Self {
|
|
JobUpdate {
|
|
status: status.to_string(),
|
|
error: None,
|
|
queue_position: None,
|
|
pending: None,
|
|
saved: None,
|
|
}
|
|
}
|
|
}
|
|
|
|
pub trait EnhanceJobDao: Sync + Send {
|
|
fn insert(&mut self, job: &EnhanceJobRow) -> Result<(), DbError>;
|
|
fn get(&mut self, id: &str) -> Result<Option<EnhanceJobRow>, DbError>;
|
|
/// Newest first.
|
|
fn list_for_hash(
|
|
&mut self,
|
|
original_hash: &str,
|
|
limit: i64,
|
|
) -> Result<Vec<EnhanceJobRow>, DbError>;
|
|
/// Jobs the poller still has to advance (`queued` / `running`).
|
|
fn list_active(&mut self) -> Result<Vec<EnhanceJobRow>, DbError>;
|
|
fn update(&mut self, id: &str, update: &JobUpdate) -> Result<(), DbError>;
|
|
/// Hide a job from `list_for_hash`.
|
|
fn discard(&mut self, id: &str) -> Result<(), DbError>;
|
|
/// Delete jobs last touched before `cutoff` (unix seconds) that are not
|
|
/// active. Returns the number removed.
|
|
fn prune(&mut self, cutoff: i64) -> Result<usize, DbError>;
|
|
}
|
|
|
|
pub struct SqliteEnhanceJobDao {
|
|
connection: Arc<Mutex<SqliteConnection>>,
|
|
}
|
|
|
|
impl Default for SqliteEnhanceJobDao {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
impl SqliteEnhanceJobDao {
|
|
pub fn new() -> Self {
|
|
SqliteEnhanceJobDao {
|
|
connection: Arc::new(Mutex::new(connect())),
|
|
}
|
|
}
|
|
|
|
/// Test-only: share a connection with other DAOs.
|
|
#[cfg(test)]
|
|
pub fn from_shared(connection: Arc<Mutex<SqliteConnection>>) -> Self {
|
|
SqliteEnhanceJobDao { connection }
|
|
}
|
|
}
|
|
|
|
impl EnhanceJobDao for SqliteEnhanceJobDao {
|
|
fn insert(&mut self, job: &EnhanceJobRow) -> Result<(), DbError> {
|
|
trace_db_call("insert", "enhance_job_insert", |_span| {
|
|
let mut conn = self.connection.lock().expect("Unable to get EnhanceJobDao");
|
|
sql_query(format!(
|
|
"INSERT INTO enhance_jobs ({COLUMNS}) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)"
|
|
))
|
|
.bind::<Text, _>(&job.id)
|
|
.bind::<Text, _>(&job.original_hash)
|
|
.bind::<Integer, _>(job.library_id)
|
|
.bind::<Text, _>(&job.rel_path)
|
|
.bind::<Text, _>(&job.template)
|
|
.bind::<Nullable<Text>, _>(&job.prompt)
|
|
.bind::<Nullable<Text>, _>(&job.reference_hash)
|
|
.bind::<Nullable<Text>, _>(&job.reference_path)
|
|
.bind::<Text, _>(&job.comfy_prompt_id)
|
|
.bind::<Text, _>(&job.status)
|
|
.bind::<Nullable<Text>, _>(&job.error)
|
|
.bind::<Nullable<Integer>, _>(job.queue_position)
|
|
.bind::<Text, _>(&job.pending)
|
|
.bind::<Text, _>(&job.saved)
|
|
.bind::<Integer, _>(job.discarded)
|
|
.bind::<BigInt, _>(job.created_at)
|
|
.bind::<BigInt, _>(job.updated_at)
|
|
.execute(&mut *conn)
|
|
.map(|_| ())
|
|
.map_err(|e| anyhow::anyhow!("Insert error: {}", e))
|
|
})
|
|
.map_err(|e| DbError::log(DbErrorKind::InsertError, e))
|
|
}
|
|
|
|
fn get(&mut self, id: &str) -> Result<Option<EnhanceJobRow>, DbError> {
|
|
trace_db_call("query", "enhance_job_get", |_span| {
|
|
let mut conn = self.connection.lock().expect("Unable to get EnhanceJobDao");
|
|
sql_query(format!("SELECT {COLUMNS} FROM enhance_jobs WHERE id = ?"))
|
|
.bind::<Text, _>(id)
|
|
.load::<EnhanceJobRow>(&mut *conn)
|
|
.map(|rows| rows.into_iter().next())
|
|
.map_err(|e| anyhow::anyhow!("Query error: {}", e))
|
|
})
|
|
.map_err(|e| DbError::log(DbErrorKind::QueryError, e))
|
|
}
|
|
|
|
fn list_for_hash(
|
|
&mut self,
|
|
original_hash: &str,
|
|
limit: i64,
|
|
) -> Result<Vec<EnhanceJobRow>, DbError> {
|
|
trace_db_call("query", "enhance_job_list_for_hash", |_span| {
|
|
let mut conn = self.connection.lock().expect("Unable to get EnhanceJobDao");
|
|
sql_query(format!(
|
|
"SELECT {COLUMNS} FROM enhance_jobs WHERE original_hash = ? AND discarded = 0 \
|
|
ORDER BY created_at DESC, rowid DESC LIMIT ?"
|
|
))
|
|
.bind::<Text, _>(original_hash)
|
|
.bind::<BigInt, _>(limit)
|
|
.load::<EnhanceJobRow>(&mut *conn)
|
|
.map_err(|e| anyhow::anyhow!("Query error: {}", e))
|
|
})
|
|
.map_err(|e| DbError::log(DbErrorKind::QueryError, e))
|
|
}
|
|
|
|
fn list_active(&mut self) -> Result<Vec<EnhanceJobRow>, DbError> {
|
|
trace_db_call("query", "enhance_job_list_active", |_span| {
|
|
let mut conn = self.connection.lock().expect("Unable to get EnhanceJobDao");
|
|
sql_query(format!(
|
|
"SELECT {COLUMNS} FROM enhance_jobs WHERE status IN ('queued', 'running') \
|
|
ORDER BY created_at, rowid"
|
|
))
|
|
.load::<EnhanceJobRow>(&mut *conn)
|
|
.map_err(|e| anyhow::anyhow!("Query error: {}", e))
|
|
})
|
|
.map_err(|e| DbError::log(DbErrorKind::QueryError, e))
|
|
}
|
|
|
|
fn update(&mut self, id: &str, u: &JobUpdate) -> Result<(), DbError> {
|
|
trace_db_call("update", "enhance_job_update", |_span| {
|
|
let mut conn = self.connection.lock().expect("Unable to get EnhanceJobDao");
|
|
sql_query(
|
|
"UPDATE enhance_jobs SET status = ?, error = ?, queue_position = ?, \
|
|
pending = COALESCE(?, pending), saved = COALESCE(?, saved), updated_at = ? \
|
|
WHERE id = ?",
|
|
)
|
|
.bind::<Text, _>(&u.status)
|
|
.bind::<Nullable<Text>, _>(&u.error)
|
|
.bind::<Nullable<Integer>, _>(u.queue_position)
|
|
.bind::<Nullable<Text>, _>(&u.pending)
|
|
.bind::<Nullable<Text>, _>(&u.saved)
|
|
.bind::<BigInt, _>(chrono::Utc::now().timestamp())
|
|
.bind::<Text, _>(id)
|
|
.execute(&mut *conn)
|
|
.map(|_| ())
|
|
.map_err(|e| anyhow::anyhow!("Update error: {}", e))
|
|
})
|
|
.map_err(|e| DbError::log(DbErrorKind::UpdateError, e))
|
|
}
|
|
|
|
fn discard(&mut self, id: &str) -> Result<(), DbError> {
|
|
trace_db_call("update", "enhance_job_discard", |_span| {
|
|
let mut conn = self.connection.lock().expect("Unable to get EnhanceJobDao");
|
|
sql_query("UPDATE enhance_jobs SET discarded = 1, updated_at = ? WHERE id = ?")
|
|
.bind::<BigInt, _>(chrono::Utc::now().timestamp())
|
|
.bind::<Text, _>(id)
|
|
.execute(&mut *conn)
|
|
.map(|_| ())
|
|
.map_err(|e| anyhow::anyhow!("Update error: {}", e))
|
|
})
|
|
.map_err(|e| DbError::log(DbErrorKind::UpdateError, e))
|
|
}
|
|
|
|
fn prune(&mut self, cutoff: i64) -> Result<usize, DbError> {
|
|
trace_db_call("delete", "enhance_job_prune", |_span| {
|
|
let mut conn = self.connection.lock().expect("Unable to get EnhanceJobDao");
|
|
sql_query(
|
|
"DELETE FROM enhance_jobs WHERE updated_at < ? \
|
|
AND status NOT IN ('queued', 'running')",
|
|
)
|
|
.bind::<BigInt, _>(cutoff)
|
|
.execute(&mut *conn)
|
|
.map_err(|e| anyhow::anyhow!("Delete error: {}", e))
|
|
})
|
|
.map_err(|e| DbError::log(DbErrorKind::UpdateError, e))
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::database::test::in_memory_db_connection;
|
|
|
|
fn dao() -> SqliteEnhanceJobDao {
|
|
SqliteEnhanceJobDao::from_shared(Arc::new(Mutex::new(in_memory_db_connection())))
|
|
}
|
|
|
|
fn job(id: &str, hash: &str, status: &str, created: i64) -> EnhanceJobRow {
|
|
EnhanceJobRow {
|
|
id: id.into(),
|
|
original_hash: hash.into(),
|
|
library_id: 1,
|
|
rel_path: "a.jpg".into(),
|
|
template: "restore".into(),
|
|
prompt: Some("warmer".into()),
|
|
reference_hash: None,
|
|
reference_path: None,
|
|
comfy_prompt_id: format!("p-{id}"),
|
|
status: status.into(),
|
|
error: None,
|
|
queue_position: None,
|
|
pending: "[]".into(),
|
|
saved: "[]".into(),
|
|
discarded: 0,
|
|
created_at: created,
|
|
updated_at: created,
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn insert_get_list_update_prune() {
|
|
let mut d = dao();
|
|
d.insert(&job("a", "h1", "queued", 100)).unwrap();
|
|
d.insert(&job("b", "h1", "ready", 200)).unwrap();
|
|
d.insert(&job("c", "h2", "running", 300)).unwrap();
|
|
|
|
assert_eq!(d.get("a").unwrap().unwrap().comfy_prompt_id, "p-a");
|
|
assert!(d.get("zzz").unwrap().is_none());
|
|
|
|
let ids: Vec<String> = d
|
|
.list_for_hash("h1", 10)
|
|
.unwrap()
|
|
.into_iter()
|
|
.map(|j| j.id)
|
|
.collect();
|
|
assert_eq!(ids, ["b", "a"], "newest first");
|
|
assert_eq!(d.list_for_hash("h1", 1).unwrap().len(), 1);
|
|
|
|
let active: Vec<String> = d.list_active().unwrap().into_iter().map(|j| j.id).collect();
|
|
assert_eq!(active, ["a", "c"]);
|
|
|
|
d.update(
|
|
"a",
|
|
&JobUpdate {
|
|
status: "ready".into(),
|
|
error: None,
|
|
queue_position: None,
|
|
pending: Some("[{\"id\":\"x\"}]".into()),
|
|
saved: None,
|
|
},
|
|
)
|
|
.unwrap();
|
|
let a = d.get("a").unwrap().unwrap();
|
|
assert_eq!(a.status, "ready");
|
|
assert_eq!(a.pending, "[{\"id\":\"x\"}]");
|
|
assert_eq!(a.saved, "[]", "unset columns are left alone");
|
|
assert!(a.updated_at >= 100);
|
|
|
|
d.discard("b").unwrap();
|
|
assert_eq!(
|
|
d.list_for_hash("h1", 10).unwrap().len(),
|
|
1,
|
|
"discarded hidden"
|
|
);
|
|
|
|
// Prune drops only old, finished jobs (updated_at was just bumped
|
|
// for `a`; `c` is active).
|
|
assert_eq!(d.prune(i64::MAX).unwrap(), 2);
|
|
assert!(d.get("c").unwrap().is_some());
|
|
assert!(d.get("b").unwrap().is_none());
|
|
}
|
|
}
|