Skip to main content

mlua_swarm/store/run/
sqlite.rs

1//! `SqliteRunStore` — SQLite-backed [`RunStore`] using [`rusqlite-isle`].
2//!
3//! The `Connection` is confined to a dedicated OS thread by `AsyncIsle`;
4//! every call is a typed closure dispatched over a bounded channel.
5//! `step_entries` and `result_ref` are stored as JSON blobs — the former is
6//! a pure trace/observability artifact (not queried relationally), the
7//! latter is caller-defined payload shape. `append_step_entry` runs as a
8//! read-modify-write inside a single transaction so concurrent appenders
9//! don't clobber each other's entries.
10//!
11//! ## Schema
12//!
13//! ```sql
14//! CREATE TABLE IF NOT EXISTS runs (
15//!   id                 TEXT PRIMARY KEY,
16//!   task_id            TEXT NOT NULL,
17//!   status             TEXT NOT NULL,      -- JSON-encoded `RunStatus`
18//!   step_entries_json  TEXT NOT NULL,      -- JSON-encoded `Vec<StepEntry>`
19//!   operator_sid       TEXT,
20//!   result_ref_json    TEXT,               -- JSON-encoded `serde_json::Value`, NULL when unset
21//!   created_at         INTEGER NOT NULL,
22//!   updated_at         INTEGER NOT NULL
23//! );
24//! CREATE INDEX IF NOT EXISTS ix_runs_task_id ON runs(task_id, created_at);
25//! ```
26
27use super::{RunId, RunRecord, RunStatus, RunStore, RunStoreError, StepEntry, TaskId};
28use async_trait::async_trait;
29use rusqlite::{params, OptionalExtension};
30use rusqlite_isle::{AsyncIsle, AsyncIsleDriver, IsleError};
31use std::path::Path;
32
33const SCHEMA_SQL: &str = "\
34CREATE TABLE IF NOT EXISTS runs (\
35  id                 TEXT PRIMARY KEY, \
36  task_id            TEXT NOT NULL, \
37  status             TEXT NOT NULL, \
38  step_entries_json  TEXT NOT NULL, \
39  operator_sid       TEXT, \
40  result_ref_json    TEXT, \
41  created_at         INTEGER NOT NULL, \
42  updated_at         INTEGER NOT NULL\
43);\
44CREATE INDEX IF NOT EXISTS ix_runs_task_id ON runs(task_id, created_at);\
45";
46
47/// SQLite-backed persistent [`RunStore`].
48///
49/// Open with [`SqliteRunStore::open`] (file path) or
50/// [`SqliteRunStore::open_in_memory`] (tests). Both return the store plus
51/// an [`AsyncIsleDriver`] the caller must `shutdown().await` when done —
52/// dropping the driver without a shutdown call leaves the SQLite thread
53/// as-is until the process exits.
54pub struct SqliteRunStore {
55    isle: AsyncIsle,
56}
57
58impl SqliteRunStore {
59    /// Open (or create) a SQLite database file and run the schema
60    /// migrations.
61    pub async fn open(path: impl AsRef<Path>) -> Result<(Self, AsyncIsleDriver), RunStoreError> {
62        let (isle, driver) = AsyncIsle::spawn(path.as_ref().to_path_buf(), |conn| {
63            conn.execute_batch(SCHEMA_SQL)
64        })
65        .await
66        .map_err(map_isle_err)?;
67        Ok((Self { isle }, driver))
68    }
69
70    /// Open an ephemeral in-memory database (tests, doctests).
71    pub async fn open_in_memory() -> Result<(Self, AsyncIsleDriver), RunStoreError> {
72        let (isle, driver) = AsyncIsle::open_in_memory(|conn| conn.execute_batch(SCHEMA_SQL))
73            .await
74            .map_err(map_isle_err)?;
75        Ok((Self { isle }, driver))
76    }
77}
78
79fn map_isle_err(e: IsleError) -> RunStoreError {
80    RunStoreError::Other(format!("sqlite: {e}"))
81}
82
83/// One `runs` SELECT row in column order: id, task_id, status,
84/// step_entries_json, operator_sid, result_ref_json, created_at,
85/// updated_at.
86type RunRow = (
87    String,
88    String,
89    String,
90    String,
91    Option<String>,
92    Option<String>,
93    i64,
94    i64,
95);
96
97fn row_to_record(row: RunRow) -> Result<RunRecord, RunStoreError> {
98    let (
99        id,
100        task_id,
101        status_json,
102        step_entries_json,
103        operator_sid,
104        result_ref_json,
105        created_at,
106        updated_at,
107    ) = row;
108    let status: RunStatus = serde_json::from_str(&status_json)
109        .map_err(|e| RunStoreError::Other(format!("decode status: {e}")))?;
110    let step_entries: Vec<StepEntry> = serde_json::from_str(&step_entries_json)
111        .map_err(|e| RunStoreError::Other(format!("decode step_entries: {e}")))?;
112    let result_ref: Option<serde_json::Value> = match result_ref_json {
113        Some(text) => Some(
114            serde_json::from_str(&text)
115                .map_err(|e| RunStoreError::Other(format!("decode result_ref: {e}")))?,
116        ),
117        None => None,
118    };
119    // Ids were minted by us before landing in the table; a prefix mismatch
120    // here means the row predates the issue #13 prefix reconciliation or
121    // the file was written by something else — fail loud either way.
122    let id = RunId::parse(id).map_err(|e| RunStoreError::Other(format!("decode id: {e}")))?;
123    let task_id =
124        TaskId::parse(task_id).map_err(|e| RunStoreError::Other(format!("decode task_id: {e}")))?;
125    Ok(RunRecord {
126        id,
127        task_id,
128        status,
129        step_entries,
130        operator_sid,
131        result_ref,
132        created_at: created_at as u64,
133        updated_at: updated_at as u64,
134    })
135}
136
137#[async_trait]
138impl RunStore for SqliteRunStore {
139    fn name(&self) -> &str {
140        "sqlite"
141    }
142
143    async fn create(&self, record: RunRecord) -> Result<(), RunStoreError> {
144        let id = record.id.to_string();
145        let id_for_conflict = record.id.clone();
146        let task_id = record.task_id.to_string();
147        let status_json = serde_json::to_string(&record.status)
148            .map_err(|e| RunStoreError::Other(format!("encode status: {e}")))?;
149        let step_entries_json = serde_json::to_string(&record.step_entries)
150            .map_err(|e| RunStoreError::Other(format!("encode step_entries: {e}")))?;
151        let operator_sid = record.operator_sid.clone();
152        let result_ref_json = record
153            .result_ref
154            .as_ref()
155            .map(serde_json::to_string)
156            .transpose()
157            .map_err(|e| RunStoreError::Other(format!("encode result_ref: {e}")))?;
158        let created_at = record.created_at as i64;
159        let updated_at = record.updated_at as i64;
160
161        self.isle
162            .call(move |conn| {
163                let tx = conn.transaction()?;
164                let exists: i64 = tx.query_row(
165                    "SELECT COUNT(*) FROM runs WHERE id = ?1",
166                    params![id],
167                    |row| row.get(0),
168                )?;
169                if exists > 0 {
170                    return Err(rusqlite::Error::SqliteFailure(
171                        rusqlite::ffi::Error::new(rusqlite::ffi::SQLITE_CONSTRAINT),
172                        Some(format!("__mlua_swarm_duplicate:{id}")),
173                    ));
174                }
175                tx.execute(
176                    "INSERT INTO runs (id, task_id, status, step_entries_json, operator_sid, \
177                     result_ref_json, created_at, updated_at) \
178                     VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
179                    params![
180                        id,
181                        task_id,
182                        status_json,
183                        step_entries_json,
184                        operator_sid,
185                        result_ref_json,
186                        created_at,
187                        updated_at,
188                    ],
189                )?;
190                tx.commit()?;
191                Ok(())
192            })
193            .await
194            .map_err(|e| match &e {
195                IsleError::Sqlite(rusqlite::Error::SqliteFailure(_, Some(msg)))
196                    if msg.starts_with("__mlua_swarm_duplicate:") =>
197                {
198                    RunStoreError::Duplicate(id_for_conflict.clone())
199                }
200                _ => map_isle_err(e),
201            })
202    }
203
204    async fn get(&self, id: &RunId) -> Result<RunRecord, RunStoreError> {
205        let id_str = id.to_string();
206        let id_for_notfound = id.clone();
207        let row = self
208            .isle
209            .call(move |conn| {
210                conn.query_row(
211                    "SELECT id, task_id, status, step_entries_json, operator_sid, \
212                     result_ref_json, created_at, updated_at FROM runs WHERE id = ?1",
213                    params![id_str],
214                    |row| {
215                        Ok((
216                            row.get::<_, String>(0)?,
217                            row.get::<_, String>(1)?,
218                            row.get::<_, String>(2)?,
219                            row.get::<_, String>(3)?,
220                            row.get::<_, Option<String>>(4)?,
221                            row.get::<_, Option<String>>(5)?,
222                            row.get::<_, i64>(6)?,
223                            row.get::<_, i64>(7)?,
224                        ))
225                    },
226                )
227                .optional()
228            })
229            .await
230            .map_err(map_isle_err)?;
231        match row {
232            Some(row) => row_to_record(row),
233            None => Err(RunStoreError::NotFound(id_for_notfound)),
234        }
235    }
236
237    async fn list_by_task(&self, task_id: &TaskId) -> Result<Vec<RunRecord>, RunStoreError> {
238        let task_id_str = task_id.to_string();
239        let rows = self
240            .isle
241            .call(move |conn| {
242                let mut stmt = conn.prepare(
243                    "SELECT id, task_id, status, step_entries_json, operator_sid, \
244                     result_ref_json, created_at, updated_at FROM runs \
245                     WHERE task_id = ?1 ORDER BY created_at ASC",
246                )?;
247                let iter = stmt.query_map(params![task_id_str], |row| {
248                    Ok((
249                        row.get::<_, String>(0)?,
250                        row.get::<_, String>(1)?,
251                        row.get::<_, String>(2)?,
252                        row.get::<_, String>(3)?,
253                        row.get::<_, Option<String>>(4)?,
254                        row.get::<_, Option<String>>(5)?,
255                        row.get::<_, i64>(6)?,
256                        row.get::<_, i64>(7)?,
257                    ))
258                })?;
259                let mut out = Vec::new();
260                for r in iter {
261                    out.push(r?);
262                }
263                Ok(out)
264            })
265            .await
266            .map_err(map_isle_err)?;
267        rows.into_iter().map(row_to_record).collect()
268    }
269
270    async fn append_step_entry(&self, id: &RunId, entry: StepEntry) -> Result<(), RunStoreError> {
271        let id_str = id.to_string();
272        let id_for_notfound = id.clone();
273        let updated_at = crate::types::now_unix() as i64;
274
275        let updated = self
276            .isle
277            .call(move |conn| {
278                let tx = conn.transaction()?;
279                let existing: Option<String> = tx
280                    .query_row(
281                        "SELECT step_entries_json FROM runs WHERE id = ?1",
282                        params![id_str],
283                        |row| row.get(0),
284                    )
285                    .optional()?;
286                let Some(existing_json) = existing else {
287                    return Ok(false);
288                };
289                let mut entries: Vec<StepEntry> = serde_json::from_str(&existing_json)
290                    .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?;
291                entries.push(entry);
292                let new_json = serde_json::to_string(&entries)
293                    .map_err(|e| rusqlite::Error::ToSqlConversionFailure(Box::new(e)))?;
294                tx.execute(
295                    "UPDATE runs SET step_entries_json = ?1, updated_at = ?2 WHERE id = ?3",
296                    params![new_json, updated_at, id_str],
297                )?;
298                tx.commit()?;
299                Ok(true)
300            })
301            .await
302            .map_err(map_isle_err)?;
303
304        if updated {
305            Ok(())
306        } else {
307            Err(RunStoreError::NotFound(id_for_notfound))
308        }
309    }
310
311    async fn update_status(&self, id: &RunId, status: RunStatus) -> Result<(), RunStoreError> {
312        let id_str = id.to_string();
313        let id_for_notfound = id.clone();
314        let status_json = serde_json::to_string(&status)
315            .map_err(|e| RunStoreError::Other(format!("encode status: {e}")))?;
316        let updated_at = crate::types::now_unix() as i64;
317        let n = self
318            .isle
319            .call(move |conn| {
320                conn.execute(
321                    "UPDATE runs SET status = ?1, updated_at = ?2 WHERE id = ?3",
322                    params![status_json, updated_at, id_str],
323                )
324            })
325            .await
326            .map_err(map_isle_err)?;
327        if n == 0 {
328            Err(RunStoreError::NotFound(id_for_notfound))
329        } else {
330            Ok(())
331        }
332    }
333
334    async fn set_result(
335        &self,
336        id: &RunId,
337        result_ref: serde_json::Value,
338    ) -> Result<(), RunStoreError> {
339        let id_str = id.to_string();
340        let id_for_notfound = id.clone();
341        let result_ref_json = serde_json::to_string(&result_ref)
342            .map_err(|e| RunStoreError::Other(format!("encode result_ref: {e}")))?;
343        let updated_at = crate::types::now_unix() as i64;
344        let n = self
345            .isle
346            .call(move |conn| {
347                conn.execute(
348                    "UPDATE runs SET result_ref_json = ?1, updated_at = ?2 WHERE id = ?3",
349                    params![result_ref_json, updated_at, id_str],
350                )
351            })
352            .await
353            .map_err(map_isle_err)?;
354        if n == 0 {
355            Err(RunStoreError::NotFound(id_for_notfound))
356        } else {
357            Ok(())
358        }
359    }
360}
361
362// ──────────────────────────────────────────────────────────────────────────
363// tests
364// ──────────────────────────────────────────────────────────────────────────
365
366#[cfg(test)]
367mod tests {
368    use super::*;
369    use serde_json::json;
370
371    fn mk(id: &str, task_id: &str, created_at: u64) -> RunRecord {
372        RunRecord {
373            id: RunId::parse(id).unwrap(),
374            task_id: TaskId::parse(task_id).unwrap(),
375            status: RunStatus::Pending,
376            step_entries: vec![],
377            operator_sid: None,
378            result_ref: None,
379            created_at,
380            updated_at: created_at,
381        }
382    }
383
384    #[tokio::test]
385    async fn create_then_get() {
386        let (s, driver) = SqliteRunStore::open_in_memory().await.unwrap();
387        s.create(mk("R-1", "T-1", 100)).await.unwrap();
388        let got = s.get(&RunId::parse("R-1").unwrap()).await.unwrap();
389        assert_eq!(got.task_id, TaskId::parse("T-1").unwrap());
390        assert_eq!(got.status, RunStatus::Pending);
391        assert!(got.step_entries.is_empty());
392        assert_eq!(got.result_ref, None);
393        drop(s);
394        driver.shutdown().await.unwrap();
395    }
396
397    #[tokio::test]
398    async fn duplicate_create_rejected() {
399        let (s, driver) = SqliteRunStore::open_in_memory().await.unwrap();
400        s.create(mk("R-1", "T-1", 100)).await.unwrap();
401        let err = s.create(mk("R-1", "T-1", 200)).await.unwrap_err();
402        assert!(matches!(err, RunStoreError::Duplicate(_)), "got: {err:?}");
403        drop(s);
404        driver.shutdown().await.unwrap();
405    }
406
407    #[tokio::test]
408    async fn get_missing_returns_not_found() {
409        let (s, driver) = SqliteRunStore::open_in_memory().await.unwrap();
410        let err = s.get(&RunId::parse("R-nope").unwrap()).await.unwrap_err();
411        assert!(matches!(err, RunStoreError::NotFound(_)));
412        drop(s);
413        driver.shutdown().await.unwrap();
414    }
415
416    #[tokio::test]
417    async fn list_by_task_filters_and_orders_ascending() {
418        let (s, driver) = SqliteRunStore::open_in_memory().await.unwrap();
419        s.create(mk("R-1", "T-1", 300)).await.unwrap();
420        s.create(mk("R-2", "T-2", 50)).await.unwrap();
421        s.create(mk("R-3", "T-1", 100)).await.unwrap();
422        let list = s
423            .list_by_task(&TaskId::parse("T-1").unwrap())
424            .await
425            .unwrap();
426        let ids: Vec<_> = list.iter().map(|r| r.id.to_string()).collect();
427        assert_eq!(ids, vec!["R-3", "R-1"]);
428        drop(s);
429        driver.shutdown().await.unwrap();
430    }
431
432    #[tokio::test]
433    async fn append_step_entry_accumulates_in_order() {
434        let (s, driver) = SqliteRunStore::open_in_memory().await.unwrap();
435        s.create(mk("R-1", "T-1", 100)).await.unwrap();
436        s.append_step_entry(
437            &RunId::parse("R-1").unwrap(),
438            StepEntry {
439                step_id: crate::types::StepId::parse("ST-1").unwrap(),
440                step_ref: Some("step-a".into()),
441                status: Some("dispatched".into()),
442                at: 101,
443            },
444        )
445        .await
446        .unwrap();
447        s.append_step_entry(
448            &RunId::parse("R-1").unwrap(),
449            StepEntry {
450                step_id: crate::types::StepId::parse("ST-2").unwrap(),
451                step_ref: Some("step-b".into()),
452                status: Some("passed".into()),
453                at: 102,
454            },
455        )
456        .await
457        .unwrap();
458        let got = s.get(&RunId::parse("R-1").unwrap()).await.unwrap();
459        assert_eq!(got.step_entries.len(), 2);
460        assert_eq!(got.step_entries[0].step_ref, Some("step-a".into()));
461        assert_eq!(got.step_entries[1].step_ref, Some("step-b".into()));
462        drop(s);
463        driver.shutdown().await.unwrap();
464    }
465
466    #[tokio::test]
467    async fn append_step_entry_unknown_run_fails() {
468        let (s, driver) = SqliteRunStore::open_in_memory().await.unwrap();
469        let err = s
470            .append_step_entry(
471                &RunId::parse("R-nope").unwrap(),
472                StepEntry {
473                    step_id: crate::types::StepId::parse("ST-1").unwrap(),
474                    step_ref: None,
475                    status: None,
476                    at: 1,
477                },
478            )
479            .await
480            .unwrap_err();
481        assert!(matches!(err, RunStoreError::NotFound(_)));
482        drop(s);
483        driver.shutdown().await.unwrap();
484    }
485
486    #[tokio::test]
487    async fn update_status_persists() {
488        let (s, driver) = SqliteRunStore::open_in_memory().await.unwrap();
489        s.create(mk("R-1", "T-1", 100)).await.unwrap();
490        s.update_status(&RunId::parse("R-1").unwrap(), RunStatus::Done)
491            .await
492            .unwrap();
493        let got = s.get(&RunId::parse("R-1").unwrap()).await.unwrap();
494        assert_eq!(got.status, RunStatus::Done);
495        drop(s);
496        driver.shutdown().await.unwrap();
497    }
498
499    #[tokio::test]
500    async fn set_result_persists() {
501        let (s, driver) = SqliteRunStore::open_in_memory().await.unwrap();
502        s.create(mk("R-1", "T-1", 100)).await.unwrap();
503        s.set_result(&RunId::parse("R-1").unwrap(), json!({"ok": true}))
504            .await
505            .unwrap();
506        let got = s.get(&RunId::parse("R-1").unwrap()).await.unwrap();
507        assert_eq!(got.result_ref, Some(json!({"ok": true})));
508        drop(s);
509        driver.shutdown().await.unwrap();
510    }
511
512    #[tokio::test]
513    async fn persists_across_reopen() {
514        let dir = tempfile::tempdir().unwrap();
515        let path = dir.path().join("runs.db");
516
517        {
518            let (s, driver) = SqliteRunStore::open(&path).await.unwrap();
519            s.create(mk("R-keep", "T-keep", 42)).await.unwrap();
520            s.append_step_entry(
521                &RunId::parse("R-keep").unwrap(),
522                StepEntry {
523                    step_id: crate::types::StepId::parse("ST-1").unwrap(),
524                    step_ref: Some("step-a".into()),
525                    status: Some("dispatched".into()),
526                    at: 43,
527                },
528            )
529            .await
530            .unwrap();
531            drop(s);
532            driver.shutdown().await.unwrap();
533        }
534
535        let (s, driver) = SqliteRunStore::open(&path).await.unwrap();
536        let got = s.get(&RunId::parse("R-keep").unwrap()).await.unwrap();
537        assert_eq!(got.task_id, TaskId::parse("T-keep").unwrap());
538        assert_eq!(got.step_entries.len(), 1);
539        assert_eq!(got.step_entries[0].step_ref, Some("step-a".into()));
540        drop(s);
541        driver.shutdown().await.unwrap();
542    }
543}