mlua_swarm/store/task/
sqlite.rs

1//! `SqliteTaskStore` — SQLite-backed [`TaskStore`] 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//! `blueprint_ref` / `input_ctx` / `task_input_spec` are stored as JSON
6//! blobs (all three are already `serde_json::Value` on [`TaskRecord`]), so
7//! schema evolution of the caller-side selector / spec shape never
8//! requires a migration here.
9//!
10//! ## Schema
11//!
12//! ```sql
13//! CREATE TABLE IF NOT EXISTS tasks (
14//!   id                    TEXT PRIMARY KEY,
15//!   goal                  TEXT NOT NULL,
16//!   blueprint_ref_json    TEXT NOT NULL,
17//!   input_ctx_json        TEXT NOT NULL,
18//!   status                TEXT NOT NULL,      -- JSON-encoded `TaskRecordStatus`
19//!   created_at            INTEGER NOT NULL,
20//!   updated_at            INTEGER NOT NULL,
21//!   task_input_spec_json  TEXT                 -- nullable; JSON-encoded `TaskInputSpec`
22//! );
23//! CREATE INDEX IF NOT EXISTS ix_tasks_created_at ON tasks(created_at);
24//! ```
25//!
26//! Issue #19 ST4: `task_input_spec_json` is a nullable column added after
27//! the original schema shipped. A fresh `open`/`open_in_memory` gets it
28//! straight from `CREATE TABLE IF NOT EXISTS` above; a database file
29//! created before ST4 already has a `tasks` table without it, so
30//! [`ensure_task_input_spec_column`] runs an `ALTER TABLE ... ADD COLUMN`
31//! backfill guarded by a `PRAGMA table_info` existence check (idempotent —
32//! safe to run against both pre-ST4 and post-ST4 files).
33
34use super::{TaskId, TaskRecord, TaskRecordStatus, TaskStore, TaskStoreError};
35use async_trait::async_trait;
36use rusqlite::{params, OptionalExtension};
37use rusqlite_isle::{AsyncIsle, AsyncIsleDriver, IsleError};
38use std::path::Path;
39
40const SCHEMA_SQL: &str = "\
41CREATE TABLE IF NOT EXISTS tasks (\
42  id                    TEXT PRIMARY KEY, \
43  goal                  TEXT NOT NULL, \
44  blueprint_ref_json    TEXT NOT NULL, \
45  input_ctx_json        TEXT NOT NULL, \
46  status                TEXT NOT NULL, \
47  created_at            INTEGER NOT NULL, \
48  updated_at            INTEGER NOT NULL, \
49  task_input_spec_json  TEXT\
50);\
51CREATE INDEX IF NOT EXISTS ix_tasks_created_at ON tasks(created_at);\
52";
53
54/// Backward-compat migration for pre-ST4 database files: add the
55/// `task_input_spec_json` column if the `tasks` table predates it. A
56/// no-op against files already created with the current [`SCHEMA_SQL`]
57/// (or already migrated once).
58fn ensure_task_input_spec_column(conn: &rusqlite::Connection) -> rusqlite::Result<()> {
59    let has_column = conn
60        .prepare("PRAGMA table_info(tasks)")?
61        .query_map([], |row| row.get::<_, String>(1))?
62        .filter_map(|r| r.ok())
63        .any(|name| name == "task_input_spec_json");
64    if !has_column {
65        conn.execute_batch("ALTER TABLE tasks ADD COLUMN task_input_spec_json TEXT")?;
66    }
67    Ok(())
68}
69
70/// SQLite-backed persistent [`TaskStore`].
71///
72/// Open with [`SqliteTaskStore::open`] (file path) or
73/// [`SqliteTaskStore::open_in_memory`] (tests). Both return the store plus
74/// an [`AsyncIsleDriver`] the caller must `shutdown().await` when done —
75/// dropping the driver without a shutdown call leaves the SQLite thread
76/// as-is until the process exits.
77pub struct SqliteTaskStore {
78    isle: AsyncIsle,
79}
80
81impl SqliteTaskStore {
82    /// Open (or create) a SQLite database file and run the schema
83    /// migrations.
84    pub async fn open(path: impl AsRef<Path>) -> Result<(Self, AsyncIsleDriver), TaskStoreError> {
85        let (isle, driver) = AsyncIsle::spawn(path.as_ref().to_path_buf(), |conn| {
86            conn.execute_batch(SCHEMA_SQL)?;
87            ensure_task_input_spec_column(conn)
88        })
89        .await
90        .map_err(map_isle_err)?;
91        Ok((Self { isle }, driver))
92    }
93
94    /// Open an ephemeral in-memory database (tests, doctests).
95    pub async fn open_in_memory() -> Result<(Self, AsyncIsleDriver), TaskStoreError> {
96        let (isle, driver) = AsyncIsle::open_in_memory(|conn| {
97            conn.execute_batch(SCHEMA_SQL)?;
98            ensure_task_input_spec_column(conn)
99        })
100        .await
101        .map_err(map_isle_err)?;
102        Ok((Self { isle }, driver))
103    }
104}
105
106fn map_isle_err(e: IsleError) -> TaskStoreError {
107    TaskStoreError::Other(format!("sqlite: {e}"))
108}
109
110/// One `tasks` SELECT row in column order: id, goal, blueprint_ref_json,
111/// input_ctx_json, status, created_at, updated_at, task_input_spec_json.
112/// The last column is nullable (issue #19 ST4) — `None` for every row
113/// created before that column existed, or by a caller with no Task-level
114/// fields to snapshot.
115type TaskRow = (
116    String,
117    String,
118    String,
119    String,
120    String,
121    i64,
122    i64,
123    Option<String>,
124);
125
126fn row_to_record(row: TaskRow) -> Result<TaskRecord, TaskStoreError> {
127    let (
128        id,
129        goal,
130        blueprint_ref_json,
131        input_ctx_json,
132        status_json,
133        created_at,
134        updated_at,
135        task_input_spec_json,
136    ) = row;
137    let blueprint_ref: serde_json::Value = serde_json::from_str(&blueprint_ref_json)
138        .map_err(|e| TaskStoreError::Other(format!("decode blueprint_ref: {e}")))?;
139    let input_ctx: serde_json::Value = serde_json::from_str(&input_ctx_json)
140        .map_err(|e| TaskStoreError::Other(format!("decode input_ctx: {e}")))?;
141    let status: TaskRecordStatus = serde_json::from_str(&status_json)
142        .map_err(|e| TaskStoreError::Other(format!("decode status: {e}")))?;
143    let task_input_spec: Option<serde_json::Value> = task_input_spec_json
144        .map(|s| serde_json::from_str(&s))
145        .transpose()
146        .map_err(|e| TaskStoreError::Other(format!("decode task_input_spec: {e}")))?;
147    // Ids were minted by us before landing in the table; a prefix mismatch
148    // here means the row predates the issue #13 prefix reconciliation or
149    // the file was written by something else — fail loud either way.
150    let id = TaskId::parse(id).map_err(|e| TaskStoreError::Other(format!("decode id: {e}")))?;
151    Ok(TaskRecord {
152        id,
153        goal,
154        blueprint_ref,
155        input_ctx,
156        task_input_spec,
157        status,
158        created_at: created_at as u64,
159        updated_at: updated_at as u64,
160    })
161}
162
163#[async_trait]
164impl TaskStore for SqliteTaskStore {
165    fn name(&self) -> &str {
166        "sqlite"
167    }
168
169    async fn create(&self, record: TaskRecord) -> Result<(), TaskStoreError> {
170        let id = record.id.to_string();
171        let id_for_conflict = record.id.clone();
172        let goal = record.goal.clone();
173        let blueprint_ref_json = serde_json::to_string(&record.blueprint_ref)
174            .map_err(|e| TaskStoreError::Other(format!("encode blueprint_ref: {e}")))?;
175        let input_ctx_json = serde_json::to_string(&record.input_ctx)
176            .map_err(|e| TaskStoreError::Other(format!("encode input_ctx: {e}")))?;
177        let status_json = serde_json::to_string(&record.status)
178            .map_err(|e| TaskStoreError::Other(format!("encode status: {e}")))?;
179        let created_at = record.created_at as i64;
180        let updated_at = record.updated_at as i64;
181        let task_input_spec_json = record
182            .task_input_spec
183            .as_ref()
184            .map(serde_json::to_string)
185            .transpose()
186            .map_err(|e| TaskStoreError::Other(format!("encode task_input_spec: {e}")))?;
187
188        self.isle
189            .call(move |conn| {
190                let tx = conn.transaction()?;
191                let exists: i64 = tx.query_row(
192                    "SELECT COUNT(*) FROM tasks WHERE id = ?1",
193                    params![id],
194                    |row| row.get(0),
195                )?;
196                if exists > 0 {
197                    return Err(rusqlite::Error::SqliteFailure(
198                        rusqlite::ffi::Error::new(rusqlite::ffi::SQLITE_CONSTRAINT),
199                        Some(format!("__mlua_swarm_duplicate:{id}")),
200                    ));
201                }
202                tx.execute(
203                    "INSERT INTO tasks (id, goal, blueprint_ref_json, input_ctx_json, status, \
204                     created_at, updated_at, task_input_spec_json) VALUES \
205                     (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
206                    params![
207                        id,
208                        goal,
209                        blueprint_ref_json,
210                        input_ctx_json,
211                        status_json,
212                        created_at,
213                        updated_at,
214                        task_input_spec_json,
215                    ],
216                )?;
217                tx.commit()?;
218                Ok(())
219            })
220            .await
221            .map_err(|e| match &e {
222                IsleError::Sqlite(rusqlite::Error::SqliteFailure(_, Some(msg)))
223                    if msg.starts_with("__mlua_swarm_duplicate:") =>
224                {
225                    TaskStoreError::Duplicate(id_for_conflict.clone())
226                }
227                _ => map_isle_err(e),
228            })
229    }
230
231    async fn get(&self, id: &TaskId) -> Result<TaskRecord, TaskStoreError> {
232        let id_str = id.to_string();
233        let id_for_notfound = id.clone();
234        let row = self
235            .isle
236            .call(move |conn| {
237                conn.query_row(
238                    "SELECT id, goal, blueprint_ref_json, input_ctx_json, status, created_at, \
239                     updated_at, task_input_spec_json FROM tasks WHERE id = ?1",
240                    params![id_str],
241                    |row| {
242                        Ok((
243                            row.get::<_, String>(0)?,
244                            row.get::<_, String>(1)?,
245                            row.get::<_, String>(2)?,
246                            row.get::<_, String>(3)?,
247                            row.get::<_, String>(4)?,
248                            row.get::<_, i64>(5)?,
249                            row.get::<_, i64>(6)?,
250                            row.get::<_, Option<String>>(7)?,
251                        ))
252                    },
253                )
254                .optional()
255            })
256            .await
257            .map_err(map_isle_err)?;
258        match row {
259            Some(row) => row_to_record(row),
260            None => Err(TaskStoreError::NotFound(id_for_notfound)),
261        }
262    }
263
264    async fn list(&self) -> Result<Vec<TaskRecord>, TaskStoreError> {
265        let rows = self
266            .isle
267            .call(|conn| {
268                let mut stmt = conn.prepare(
269                    "SELECT id, goal, blueprint_ref_json, input_ctx_json, status, created_at, \
270                     updated_at, task_input_spec_json FROM tasks ORDER BY created_at DESC",
271                )?;
272                let iter = stmt.query_map([], |row| {
273                    Ok((
274                        row.get::<_, String>(0)?,
275                        row.get::<_, String>(1)?,
276                        row.get::<_, String>(2)?,
277                        row.get::<_, String>(3)?,
278                        row.get::<_, String>(4)?,
279                        row.get::<_, i64>(5)?,
280                        row.get::<_, i64>(6)?,
281                        row.get::<_, Option<String>>(7)?,
282                    ))
283                })?;
284                let mut out = Vec::new();
285                for r in iter {
286                    out.push(r?);
287                }
288                Ok(out)
289            })
290            .await
291            .map_err(map_isle_err)?;
292        rows.into_iter().map(row_to_record).collect()
293    }
294
295    async fn update_status(
296        &self,
297        id: &TaskId,
298        status: TaskRecordStatus,
299    ) -> Result<(), TaskStoreError> {
300        let id_str = id.to_string();
301        let id_for_notfound = id.clone();
302        let status_json = serde_json::to_string(&status)
303            .map_err(|e| TaskStoreError::Other(format!("encode status: {e}")))?;
304        let updated_at = crate::types::now_unix() as i64;
305        let n = self
306            .isle
307            .call(move |conn| {
308                conn.execute(
309                    "UPDATE tasks SET status = ?1, updated_at = ?2 WHERE id = ?3",
310                    params![status_json, updated_at, id_str],
311                )
312            })
313            .await
314            .map_err(map_isle_err)?;
315        if n == 0 {
316            Err(TaskStoreError::NotFound(id_for_notfound))
317        } else {
318            Ok(())
319        }
320    }
321}
322
323// ──────────────────────────────────────────────────────────────────────────
324// tests
325// ──────────────────────────────────────────────────────────────────────────
326
327#[cfg(test)]
328mod tests {
329    use super::*;
330    use serde_json::json;
331
332    fn mk(id: &str, created_at: u64) -> TaskRecord {
333        TaskRecord {
334            id: TaskId::parse(id).unwrap(),
335            goal: format!("goal for {id}"),
336            blueprint_ref: json!({"id": "bp-1"}),
337            input_ctx: json!({"k": "v"}),
338            task_input_spec: None,
339            status: TaskRecordStatus::Pending,
340            created_at,
341            updated_at: created_at,
342        }
343    }
344
345    #[tokio::test]
346    async fn create_then_get() {
347        let (s, driver) = SqliteTaskStore::open_in_memory().await.unwrap();
348        s.create(mk("T-1", 100)).await.unwrap();
349        let got = s.get(&TaskId::parse("T-1").unwrap()).await.unwrap();
350        assert_eq!(got.goal, "goal for T-1");
351        assert_eq!(got.status, TaskRecordStatus::Pending);
352        assert_eq!(got.blueprint_ref, json!({"id": "bp-1"}));
353        drop(s);
354        driver.shutdown().await.unwrap();
355    }
356
357    #[tokio::test]
358    async fn duplicate_create_rejected() {
359        let (s, driver) = SqliteTaskStore::open_in_memory().await.unwrap();
360        s.create(mk("T-1", 100)).await.unwrap();
361        let err = s.create(mk("T-1", 200)).await.unwrap_err();
362        assert!(matches!(err, TaskStoreError::Duplicate(_)), "got: {err:?}");
363        drop(s);
364        driver.shutdown().await.unwrap();
365    }
366
367    #[tokio::test]
368    async fn get_missing_returns_not_found() {
369        let (s, driver) = SqliteTaskStore::open_in_memory().await.unwrap();
370        let err = s.get(&TaskId::parse("T-nope").unwrap()).await.unwrap_err();
371        assert!(matches!(err, TaskStoreError::NotFound(_)));
372        drop(s);
373        driver.shutdown().await.unwrap();
374    }
375
376    #[tokio::test]
377    async fn list_returns_newest_first() {
378        let (s, driver) = SqliteTaskStore::open_in_memory().await.unwrap();
379        s.create(mk("T-1", 100)).await.unwrap();
380        s.create(mk("T-2", 300)).await.unwrap();
381        s.create(mk("T-3", 200)).await.unwrap();
382        let list = s.list().await.unwrap();
383        let ids: Vec<_> = list.iter().map(|r| r.id.to_string()).collect();
384        assert_eq!(ids, vec!["T-2", "T-3", "T-1"]);
385        drop(s);
386        driver.shutdown().await.unwrap();
387    }
388
389    #[tokio::test]
390    async fn update_status_persists() {
391        let (s, driver) = SqliteTaskStore::open_in_memory().await.unwrap();
392        s.create(mk("T-1", 100)).await.unwrap();
393        s.update_status(&TaskId::parse("T-1").unwrap(), TaskRecordStatus::Failed)
394            .await
395            .unwrap();
396        let got = s.get(&TaskId::parse("T-1").unwrap()).await.unwrap();
397        assert_eq!(got.status, TaskRecordStatus::Failed);
398        assert!(got.updated_at >= got.created_at);
399        drop(s);
400        driver.shutdown().await.unwrap();
401    }
402
403    #[tokio::test]
404    async fn update_status_unknown_fails() {
405        let (s, driver) = SqliteTaskStore::open_in_memory().await.unwrap();
406        let err = s
407            .update_status(&TaskId::parse("T-nope").unwrap(), TaskRecordStatus::Done)
408            .await
409            .unwrap_err();
410        assert!(matches!(err, TaskStoreError::NotFound(_)));
411        drop(s);
412        driver.shutdown().await.unwrap();
413    }
414
415    #[tokio::test]
416    async fn persists_across_reopen() {
417        let dir = tempfile::tempdir().unwrap();
418        let path = dir.path().join("tasks.db");
419
420        {
421            let (s, driver) = SqliteTaskStore::open(&path).await.unwrap();
422            s.create(mk("T-keep", 42)).await.unwrap();
423            drop(s);
424            driver.shutdown().await.unwrap();
425        }
426
427        let (s, driver) = SqliteTaskStore::open(&path).await.unwrap();
428        let got = s.get(&TaskId::parse("T-keep").unwrap()).await.unwrap();
429        assert_eq!(got.goal, "goal for T-keep");
430        assert_eq!(got.created_at, 42);
431        drop(s);
432        driver.shutdown().await.unwrap();
433    }
434
435    // ──────────────────────────────────────────────────────────────────
436    // issue #19 ST4: `task_input_spec` (nullable column + migration)
437    // ──────────────────────────────────────────────────────────────────
438
439    #[tokio::test]
440    async fn create_and_get_round_trip_task_input_spec() {
441        let (s, driver) = SqliteTaskStore::open_in_memory().await.unwrap();
442        let mut record = mk("T-1", 100);
443        record.task_input_spec = Some(json!({"project_root": "/repo", "work_dir": "/repo/work"}));
444        s.create(record).await.unwrap();
445        let got = s.get(&TaskId::parse("T-1").unwrap()).await.unwrap();
446        assert_eq!(
447            got.task_input_spec,
448            Some(json!({"project_root": "/repo", "work_dir": "/repo/work"}))
449        );
450        drop(s);
451        driver.shutdown().await.unwrap();
452    }
453
454    #[tokio::test]
455    async fn create_and_get_task_input_spec_absent_is_none() {
456        let (s, driver) = SqliteTaskStore::open_in_memory().await.unwrap();
457        s.create(mk("T-1", 100)).await.unwrap();
458        let got = s.get(&TaskId::parse("T-1").unwrap()).await.unwrap();
459        assert_eq!(got.task_input_spec, None);
460        drop(s);
461        driver.shutdown().await.unwrap();
462    }
463
464    #[tokio::test]
465    async fn list_carries_task_input_spec_per_row() {
466        let (s, driver) = SqliteTaskStore::open_in_memory().await.unwrap();
467        s.create(mk("T-1", 100)).await.unwrap();
468        let mut with_spec = mk("T-2", 200);
469        with_spec.task_input_spec = Some(json!({"work_dir": "/repo/work"}));
470        s.create(with_spec).await.unwrap();
471        let list = s.list().await.unwrap();
472        let t1 = list.iter().find(|r| r.id.to_string() == "T-1").unwrap();
473        let t2 = list.iter().find(|r| r.id.to_string() == "T-2").unwrap();
474        assert_eq!(t1.task_input_spec, None);
475        assert_eq!(t2.task_input_spec, Some(json!({"work_dir": "/repo/work"})));
476        drop(s);
477        driver.shutdown().await.unwrap();
478    }
479
480    #[tokio::test]
481    async fn open_migrates_pre_st4_schema_missing_task_input_spec_column() {
482        // Backward compat: a database file created before this column
483        // existed must still open cleanly, and its pre-existing rows must
484        // decode with `task_input_spec: None` (rather than failing to
485        // decode, or the migration wiping the table).
486        let dir = tempfile::tempdir().unwrap();
487        let path = dir.path().join("legacy.db");
488
489        // Seed a pre-ST4 `tasks` table (no `task_input_spec_json` column)
490        // plus one row, bypassing `SqliteTaskStore` entirely.
491        {
492            let conn = rusqlite::Connection::open(&path).unwrap();
493            conn.execute_batch(
494                "CREATE TABLE tasks (\
495                   id                 TEXT PRIMARY KEY, \
496                   goal               TEXT NOT NULL, \
497                   blueprint_ref_json TEXT NOT NULL, \
498                   input_ctx_json     TEXT NOT NULL, \
499                   status             TEXT NOT NULL, \
500                   created_at         INTEGER NOT NULL, \
501                   updated_at         INTEGER NOT NULL\
502                 );",
503            )
504            .unwrap();
505            conn.execute(
506                "INSERT INTO tasks (id, goal, blueprint_ref_json, input_ctx_json, status, \
507                 created_at, updated_at) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
508                params![
509                    "T-legacy",
510                    "pre-st4 goal",
511                    "{\"id\":\"bp-1\"}",
512                    "{\"k\":\"v\"}",
513                    "\"pending\"",
514                    100i64,
515                    100i64,
516                ],
517            )
518            .unwrap();
519        }
520
521        // `SqliteTaskStore::open` must migrate the column in and still
522        // read the pre-existing row untouched.
523        let (s, driver) = SqliteTaskStore::open(&path).await.unwrap();
524        let got = s.get(&TaskId::parse("T-legacy").unwrap()).await.unwrap();
525        assert_eq!(got.goal, "pre-st4 goal");
526        assert_eq!(got.task_input_spec, None);
527
528        // New rows created after the migration can use the column normally.
529        let mut fresh = mk("T-fresh", 200);
530        fresh.task_input_spec = Some(json!({"work_dir": "/repo/work"}));
531        s.create(fresh).await.unwrap();
532        let got_fresh = s.get(&TaskId::parse("T-fresh").unwrap()).await.unwrap();
533        assert_eq!(
534            got_fresh.task_input_spec,
535            Some(json!({"work_dir": "/repo/work"}))
536        );
537
538        drop(s);
539        driver.shutdown().await.unwrap();
540    }
541}
mlua_swarm/store/task/sqlite.rs

mlua_swarm/store/task/
sqlite.rs
