mlua_swarm/store/output/
sqlite.rs

1//! `SqliteOutputStore` — SQLite-backed [`OutputStore`].
2//!
3//! One row per emit (`OutputRecord`). `event` and `parent_refs` are stored
4//! as JSON blobs (both types already carry `Serialize + Deserialize`), so
5//! adding a new [`OutputEvent`] variant does not require a schema migration.
6//!
7//! Ordering guarantees:
8//!
9//! - `list_for_attempt` returns rows in insertion order (per the trait
10//!   contract) via an autoincrementing `seq` column.
11//! - `get_latest_by_name` picks the row with the largest `seq` for a given
12//!   `producer_agent`.
13//! - `get_latest_by_name_in_run` (GH #23 Layer 2) is the same query
14//!   additionally filtered to one `(task_id, attempt)` run — see
15//!   `ix_outputs_producer_run`, applied idempotently via `CREATE INDEX IF
16//!   NOT EXISTS` on every `open`/`open_in_memory` call, so no migration
17//!   step is needed for pre-existing databases.
18
19use super::{OutputEvent, OutputRecord, OutputRef, OutputStore, OutputStoreError};
20use async_trait::async_trait;
21use rusqlite::{params, OptionalExtension};
22use rusqlite_isle::{AsyncIsle, AsyncIsleDriver, IsleError};
23use std::path::Path;
24
25const SCHEMA_SQL: &str = "\
26CREATE TABLE IF NOT EXISTS outputs (\
27  id             TEXT PRIMARY KEY, \
28  task_id        TEXT NOT NULL, \
29  attempt        INTEGER NOT NULL, \
30  producer_agent TEXT NOT NULL, \
31  event_json     TEXT NOT NULL, \
32  parent_refs_json TEXT NOT NULL, \
33  seq            INTEGER NOT NULL\
34);\
35CREATE INDEX IF NOT EXISTS ix_outputs_attempt ON outputs(task_id, attempt, seq);\
36CREATE INDEX IF NOT EXISTS ix_outputs_producer ON outputs(producer_agent, seq);\
37CREATE INDEX IF NOT EXISTS ix_outputs_producer_run ON outputs(producer_agent, task_id, attempt, seq);\
38";
39
40/// SQLite-backed [`OutputStore`].
41pub struct SqliteOutputStore {
42    isle: AsyncIsle,
43}
44
45impl SqliteOutputStore {
46    /// Open (or create) a SQLite file and apply the schema.
47    pub async fn open(path: impl AsRef<Path>) -> Result<(Self, AsyncIsleDriver), OutputStoreError> {
48        let (isle, driver) = AsyncIsle::spawn(path.as_ref().to_path_buf(), |conn| {
49            conn.execute_batch(SCHEMA_SQL)
50        })
51        .await
52        .map_err(map_isle_err)?;
53        Ok((Self { isle }, driver))
54    }
55
56    /// Open an ephemeral in-memory database (tests).
57    pub async fn open_in_memory() -> Result<(Self, AsyncIsleDriver), OutputStoreError> {
58        let (isle, driver) = AsyncIsle::open_in_memory(|conn| conn.execute_batch(SCHEMA_SQL))
59            .await
60            .map_err(map_isle_err)?;
61        Ok((Self { isle }, driver))
62    }
63}
64
65fn map_isle_err(e: IsleError) -> OutputStoreError {
66    OutputStoreError::Internal(format!("sqlite: {e}"))
67}
68
69fn decode_record(
70    id: String,
71    task_id: String,
72    attempt: i64,
73    producer_agent: String,
74    event_json: String,
75    parent_refs_json: String,
76) -> Result<OutputRecord, OutputStoreError> {
77    let event: OutputEvent = serde_json::from_str(&event_json)
78        .map_err(|e| OutputStoreError::Internal(format!("decode event: {e}")))?;
79    let parent_refs: Vec<OutputRef> = serde_json::from_str(&parent_refs_json)
80        .map_err(|e| OutputStoreError::Internal(format!("decode parent_refs: {e}")))?;
81    Ok(OutputRecord {
82        id: OutputRef(id),
83        task_id,
84        attempt: attempt as u32,
85        producer_agent,
86        event,
87        parent_refs,
88    })
89}
90
91#[async_trait]
92impl OutputStore for SqliteOutputStore {
93    async fn append(
94        &self,
95        task_id: &str,
96        attempt: u32,
97        producer_agent: &str,
98        event: OutputEvent,
99        parent_refs: Vec<OutputRef>,
100    ) -> Result<OutputRef, OutputStoreError> {
101        let id = OutputRef::new();
102        let id_str = id.0.clone();
103        let task_id = task_id.to_string();
104        let attempt = attempt as i64;
105        let producer_agent = producer_agent.to_string();
106        let event_json = serde_json::to_string(&event)
107            .map_err(|e| OutputStoreError::Internal(format!("encode event: {e}")))?;
108        let parent_refs_json = serde_json::to_string(&parent_refs)
109            .map_err(|e| OutputStoreError::Internal(format!("encode parent_refs: {e}")))?;
110
111        self.isle
112            .call(move |conn| {
113                let tx = conn.transaction()?;
114                let seq: i64 =
115                    tx.query_row("SELECT COALESCE(MAX(seq), 0) + 1 FROM outputs", [], |row| {
116                        row.get(0)
117                    })?;
118                tx.execute(
119                    "INSERT INTO outputs (id, task_id, attempt, producer_agent, event_json, \
120                     parent_refs_json, seq) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
121                    params![
122                        id_str,
123                        task_id,
124                        attempt,
125                        producer_agent,
126                        event_json,
127                        parent_refs_json,
128                        seq,
129                    ],
130                )?;
131                tx.commit()?;
132                Ok(())
133            })
134            .await
135            .map_err(map_isle_err)?;
136        Ok(id)
137    }
138
139    async fn get(&self, id: &OutputRef) -> Result<OutputRecord, OutputStoreError> {
140        let id_str = id.0.clone();
141        let id_for_notfound = id.0.clone();
142        let row = self
143            .isle
144            .call(move |conn| {
145                conn.query_row(
146                    "SELECT id, task_id, attempt, producer_agent, event_json, parent_refs_json \
147                     FROM outputs WHERE id = ?1",
148                    params![id_str],
149                    |row| {
150                        Ok((
151                            row.get::<_, String>(0)?,
152                            row.get::<_, String>(1)?,
153                            row.get::<_, i64>(2)?,
154                            row.get::<_, String>(3)?,
155                            row.get::<_, String>(4)?,
156                            row.get::<_, String>(5)?,
157                        ))
158                    },
159                )
160                .optional()
161            })
162            .await
163            .map_err(map_isle_err)?;
164        match row {
165            Some((id, task_id, attempt, producer, event, parent)) => {
166                decode_record(id, task_id, attempt, producer, event, parent)
167            }
168            None => Err(OutputStoreError::NotFound(id_for_notfound)),
169        }
170    }
171
172    async fn get_latest_by_name(&self, name: &str) -> Result<OutputRecord, OutputStoreError> {
173        let name_str = name.to_string();
174        let name_for_notfound = name.to_string();
175        let row = self
176            .isle
177            .call(move |conn| {
178                conn.query_row(
179                    "SELECT id, task_id, attempt, producer_agent, event_json, parent_refs_json \
180                     FROM outputs WHERE producer_agent = ?1 ORDER BY seq DESC LIMIT 1",
181                    params![name_str],
182                    |row| {
183                        Ok((
184                            row.get::<_, String>(0)?,
185                            row.get::<_, String>(1)?,
186                            row.get::<_, i64>(2)?,
187                            row.get::<_, String>(3)?,
188                            row.get::<_, String>(4)?,
189                            row.get::<_, String>(5)?,
190                        ))
191                    },
192                )
193                .optional()
194            })
195            .await
196            .map_err(map_isle_err)?;
197        match row {
198            Some((id, task_id, attempt, producer, event, parent)) => {
199                decode_record(id, task_id, attempt, producer, event, parent)
200            }
201            None => Err(OutputStoreError::NotFound(name_for_notfound)),
202        }
203    }
204
205    async fn get_latest_by_name_in_run(
206        &self,
207        task_id: &str,
208        attempt: u32,
209        name: &str,
210    ) -> Result<OutputRecord, OutputStoreError> {
211        let name_str = name.to_string();
212        let task_id_str = task_id.to_string();
213        let task_id_for_notfound = task_id.to_string();
214        let name_for_notfound = name.to_string();
215        let attempt_i64 = attempt as i64;
216        let row = self
217            .isle
218            .call(move |conn| {
219                conn.query_row(
220                    "SELECT id, task_id, attempt, producer_agent, event_json, parent_refs_json \
221                     FROM outputs WHERE producer_agent = ?1 AND task_id = ?2 AND attempt = ?3 \
222                     ORDER BY seq DESC LIMIT 1",
223                    params![name_str, task_id_str, attempt_i64],
224                    |row| {
225                        Ok((
226                            row.get::<_, String>(0)?,
227                            row.get::<_, String>(1)?,
228                            row.get::<_, i64>(2)?,
229                            row.get::<_, String>(3)?,
230                            row.get::<_, String>(4)?,
231                            row.get::<_, String>(5)?,
232                        ))
233                    },
234                )
235                .optional()
236            })
237            .await
238            .map_err(map_isle_err)?;
239        match row {
240            Some((id, task_id, attempt, producer, event, parent)) => {
241                decode_record(id, task_id, attempt, producer, event, parent)
242            }
243            None => Err(OutputStoreError::NotFound(format!(
244                "{task_id_for_notfound}/{attempt}/{name_for_notfound}"
245            ))),
246        }
247    }
248
249    async fn list_for_attempt(
250        &self,
251        task_id: &str,
252        attempt: u32,
253    ) -> Result<Vec<OutputRecord>, OutputStoreError> {
254        let task_id = task_id.to_string();
255        let attempt = attempt as i64;
256        let rows = self
257            .isle
258            .call(move |conn| {
259                let mut stmt = conn.prepare(
260                    "SELECT id, task_id, attempt, producer_agent, event_json, parent_refs_json \
261                     FROM outputs WHERE task_id = ?1 AND attempt = ?2 ORDER BY seq ASC",
262                )?;
263                let iter = stmt.query_map(params![task_id, attempt], |row| {
264                    Ok((
265                        row.get::<_, String>(0)?,
266                        row.get::<_, String>(1)?,
267                        row.get::<_, i64>(2)?,
268                        row.get::<_, String>(3)?,
269                        row.get::<_, String>(4)?,
270                        row.get::<_, String>(5)?,
271                    ))
272                })?;
273                let mut out = Vec::new();
274                for r in iter {
275                    out.push(r?);
276                }
277                Ok(out)
278            })
279            .await
280            .map_err(map_isle_err)?;
281        rows.into_iter()
282            .map(|(id, task_id, attempt, producer, event, parent)| {
283                decode_record(id, task_id, attempt, producer, event, parent)
284            })
285            .collect()
286    }
287}
288
289#[cfg(test)]
290mod tests {
291    use super::*;
292    use crate::store::output::ContentRef;
293
294    fn mk_final(text: &str, ok: bool) -> OutputEvent {
295        OutputEvent::Final {
296            content: ContentRef::inline_text(text),
297            ok,
298        }
299    }
300
301    #[tokio::test]
302    async fn append_then_get_roundtrip() {
303        let (s, driver) = SqliteOutputStore::open_in_memory().await.unwrap();
304        let id = s
305            .append("task-1", 1, "producer-a", mk_final("hello", true), vec![])
306            .await
307            .unwrap();
308        let got = s.get(&id).await.unwrap();
309        assert_eq!(got.id, id);
310        assert_eq!(got.task_id, "task-1");
311        assert_eq!(got.attempt, 1);
312        assert_eq!(got.producer_agent, "producer-a");
313        match got.event {
314            OutputEvent::Final { ok, .. } => assert!(ok),
315            other => panic!("unexpected: {other:?}"),
316        }
317        drop(s);
318        driver.shutdown().await.unwrap();
319    }
320
321    #[tokio::test]
322    async fn get_not_found_returns_error() {
323        let (s, driver) = SqliteOutputStore::open_in_memory().await.unwrap();
324        let err = s.get(&OutputRef("missing".into())).await.unwrap_err();
325        assert!(matches!(err, OutputStoreError::NotFound(_)));
326        drop(s);
327        driver.shutdown().await.unwrap();
328    }
329
330    #[tokio::test]
331    async fn list_for_attempt_orders_by_insertion() {
332        let (s, driver) = SqliteOutputStore::open_in_memory().await.unwrap();
333        let a = s
334            .append("t", 1, "p1", mk_final("a", true), vec![])
335            .await
336            .unwrap();
337        let b = s
338            .append("t", 1, "p2", mk_final("b", true), vec![])
339            .await
340            .unwrap();
341        // Not part of the same attempt — must be skipped by the filter.
342        let _ = s
343            .append("t", 2, "p1", mk_final("other-attempt", true), vec![])
344            .await
345            .unwrap();
346        let c = s
347            .append("t", 1, "p3", mk_final("c", true), vec![])
348            .await
349            .unwrap();
350
351        let listed = s.list_for_attempt("t", 1).await.unwrap();
352        let ids: Vec<_> = listed.iter().map(|r| r.id.clone()).collect();
353        assert_eq!(ids, vec![a, b, c]);
354        drop(s);
355        driver.shutdown().await.unwrap();
356    }
357
358    #[tokio::test]
359    async fn get_latest_by_name_returns_newest_emit() {
360        let (s, driver) = SqliteOutputStore::open_in_memory().await.unwrap();
361        let _ = s
362            .append("t", 1, "same-producer", mk_final("v1", true), vec![])
363            .await
364            .unwrap();
365        let _ = s
366            .append(
367                "t",
368                1,
369                "other-producer",
370                mk_final("unrelated", true),
371                vec![],
372            )
373            .await
374            .unwrap();
375        let latest_id = s
376            .append("t", 2, "same-producer", mk_final("v2", true), vec![])
377            .await
378            .unwrap();
379        let got = s.get_latest_by_name("same-producer").await.unwrap();
380        assert_eq!(got.id, latest_id);
381        drop(s);
382        driver.shutdown().await.unwrap();
383    }
384
385    #[tokio::test]
386    async fn get_latest_by_name_unknown_returns_not_found() {
387        let (s, driver) = SqliteOutputStore::open_in_memory().await.unwrap();
388        let err = s.get_latest_by_name("nobody").await.unwrap_err();
389        assert!(matches!(err, OutputStoreError::NotFound(_)));
390        drop(s);
391        driver.shutdown().await.unwrap();
392    }
393
394    #[tokio::test]
395    async fn get_latest_by_name_in_run_does_not_cross_resolve_between_runs() {
396        let (s, driver) = SqliteOutputStore::open_in_memory().await.unwrap();
397        let id_t1 = s
398            .append("t1", 1, "same-producer", mk_final("run-1", true), vec![])
399            .await
400            .unwrap();
401        let id_t2 = s
402            .append("t2", 1, "same-producer", mk_final("run-2", true), vec![])
403            .await
404            .unwrap();
405
406        let got_t1 = s
407            .get_latest_by_name_in_run("t1", 1, "same-producer")
408            .await
409            .unwrap();
410        assert_eq!(got_t1.id, id_t1, "must not cross-resolve to t2's emit");
411
412        let got_t2 = s
413            .get_latest_by_name_in_run("t2", 1, "same-producer")
414            .await
415            .unwrap();
416        assert_eq!(got_t2.id, id_t2, "must not cross-resolve to t1's emit");
417        drop(s);
418        driver.shutdown().await.unwrap();
419    }
420
421    #[tokio::test]
422    async fn get_latest_by_name_in_run_returns_newest_within_run() {
423        let (s, driver) = SqliteOutputStore::open_in_memory().await.unwrap();
424        let _ = s
425            .append("t", 1, "p", mk_final("first", true), vec![])
426            .await
427            .unwrap();
428        let id2 = s
429            .append("t", 1, "p", mk_final("second", true), vec![])
430            .await
431            .unwrap();
432        let got = s.get_latest_by_name_in_run("t", 1, "p").await.unwrap();
433        assert_eq!(got.id, id2, "latest emit within the run wins");
434        drop(s);
435        driver.shutdown().await.unwrap();
436    }
437
438    #[tokio::test]
439    async fn get_latest_by_name_in_run_wrong_run_returns_not_found() {
440        let (s, driver) = SqliteOutputStore::open_in_memory().await.unwrap();
441        let _ = s
442            .append("t", 1, "same-producer", mk_final("x", true), vec![])
443            .await
444            .unwrap();
445        // Right name, wrong attempt — must not fall back to a different run.
446        let err = s
447            .get_latest_by_name_in_run("t", 2, "same-producer")
448            .await
449            .unwrap_err();
450        assert!(matches!(err, OutputStoreError::NotFound(_)));
451        drop(s);
452        driver.shutdown().await.unwrap();
453    }
454
455    #[tokio::test]
456    async fn parent_refs_are_persisted() {
457        let (s, driver) = SqliteOutputStore::open_in_memory().await.unwrap();
458        let a = s
459            .append("t", 1, "p", mk_final("parent", true), vec![])
460            .await
461            .unwrap();
462        let b = s
463            .append("t", 1, "p", mk_final("child", true), vec![a.clone()])
464            .await
465            .unwrap();
466        let got = s.get(&b).await.unwrap();
467        assert_eq!(got.parent_refs, vec![a]);
468        drop(s);
469        driver.shutdown().await.unwrap();
470    }
471
472    #[tokio::test]
473    async fn persists_across_reopen() {
474        let dir = tempfile::tempdir().unwrap();
475        let path = dir.path().join("outputs.db");
476        let id;
477        {
478            let (s, driver) = SqliteOutputStore::open(&path).await.unwrap();
479            id = s
480                .append("keep", 1, "p", mk_final("body", true), vec![])
481                .await
482                .unwrap();
483            drop(s);
484            driver.shutdown().await.unwrap();
485        }
486        let (s, driver) = SqliteOutputStore::open(&path).await.unwrap();
487        let got = s.get(&id).await.unwrap();
488        assert_eq!(got.task_id, "keep");
489        drop(s);
490        driver.shutdown().await.unwrap();
491    }
492}
mlua_swarm/store/output/sqlite.rs

mlua_swarm/store/output/
sqlite.rs
