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` are stored as JSON blobs (both are already
6//! `serde_json::Value`), so schema evolution of the caller-side selector
7//! shape never requires a migration here.
8//!
9//! ## Schema
10//!
11//! ```sql
12//! CREATE TABLE IF NOT EXISTS tasks (
13//!   id                 TEXT PRIMARY KEY,
14//!   goal               TEXT NOT NULL,
15//!   blueprint_ref_json TEXT NOT NULL,
16//!   input_ctx_json     TEXT NOT NULL,
17//!   status             TEXT NOT NULL,      -- JSON-encoded `TaskRecordStatus`
18//!   created_at         INTEGER NOT NULL,
19//!   updated_at         INTEGER NOT NULL
20//! );
21//! CREATE INDEX IF NOT EXISTS ix_tasks_created_at ON tasks(created_at);
22//! ```
23
24use super::{TaskId, TaskRecord, TaskRecordStatus, TaskStore, TaskStoreError};
25use async_trait::async_trait;
26use rusqlite::{params, OptionalExtension};
27use rusqlite_isle::{AsyncIsle, AsyncIsleDriver, IsleError};
28use std::path::Path;
29
30const SCHEMA_SQL: &str = "\
31CREATE TABLE IF NOT EXISTS tasks (\
32  id                 TEXT PRIMARY KEY, \
33  goal               TEXT NOT NULL, \
34  blueprint_ref_json TEXT NOT NULL, \
35  input_ctx_json     TEXT NOT NULL, \
36  status             TEXT NOT NULL, \
37  created_at         INTEGER NOT NULL, \
38  updated_at         INTEGER NOT NULL\
39);\
40CREATE INDEX IF NOT EXISTS ix_tasks_created_at ON tasks(created_at);\
41";
42
43/// SQLite-backed persistent [`TaskStore`].
44///
45/// Open with [`SqliteTaskStore::open`] (file path) or
46/// [`SqliteTaskStore::open_in_memory`] (tests). Both return the store plus
47/// an [`AsyncIsleDriver`] the caller must `shutdown().await` when done —
48/// dropping the driver without a shutdown call leaves the SQLite thread
49/// as-is until the process exits.
50pub struct SqliteTaskStore {
51    isle: AsyncIsle,
52}
53
54impl SqliteTaskStore {
55    /// Open (or create) a SQLite database file and run the schema
56    /// migrations.
57    pub async fn open(path: impl AsRef<Path>) -> Result<(Self, AsyncIsleDriver), TaskStoreError> {
58        let (isle, driver) = AsyncIsle::spawn(path.as_ref().to_path_buf(), |conn| {
59            conn.execute_batch(SCHEMA_SQL)
60        })
61        .await
62        .map_err(map_isle_err)?;
63        Ok((Self { isle }, driver))
64    }
65
66    /// Open an ephemeral in-memory database (tests, doctests).
67    pub async fn open_in_memory() -> Result<(Self, AsyncIsleDriver), TaskStoreError> {
68        let (isle, driver) = AsyncIsle::open_in_memory(|conn| conn.execute_batch(SCHEMA_SQL))
69            .await
70            .map_err(map_isle_err)?;
71        Ok((Self { isle }, driver))
72    }
73}
74
75fn map_isle_err(e: IsleError) -> TaskStoreError {
76    TaskStoreError::Other(format!("sqlite: {e}"))
77}
78
79/// One `tasks` SELECT row in column order: id, goal, blueprint_ref_json,
80/// input_ctx_json, status, created_at, updated_at.
81type TaskRow = (String, String, String, String, String, i64, i64);
82
83fn row_to_record(row: TaskRow) -> Result<TaskRecord, TaskStoreError> {
84    let (id, goal, blueprint_ref_json, input_ctx_json, status_json, created_at, updated_at) = row;
85    let blueprint_ref: serde_json::Value = serde_json::from_str(&blueprint_ref_json)
86        .map_err(|e| TaskStoreError::Other(format!("decode blueprint_ref: {e}")))?;
87    let input_ctx: serde_json::Value = serde_json::from_str(&input_ctx_json)
88        .map_err(|e| TaskStoreError::Other(format!("decode input_ctx: {e}")))?;
89    let status: TaskRecordStatus = serde_json::from_str(&status_json)
90        .map_err(|e| TaskStoreError::Other(format!("decode status: {e}")))?;
91    // Ids were minted by us before landing in the table; a prefix mismatch
92    // here means the row predates the issue #13 prefix reconciliation or
93    // the file was written by something else — fail loud either way.
94    let id = TaskId::parse(id).map_err(|e| TaskStoreError::Other(format!("decode id: {e}")))?;
95    Ok(TaskRecord {
96        id,
97        goal,
98        blueprint_ref,
99        input_ctx,
100        status,
101        created_at: created_at as u64,
102        updated_at: updated_at as u64,
103    })
104}
105
106#[async_trait]
107impl TaskStore for SqliteTaskStore {
108    fn name(&self) -> &str {
109        "sqlite"
110    }
111
112    async fn create(&self, record: TaskRecord) -> Result<(), TaskStoreError> {
113        let id = record.id.to_string();
114        let id_for_conflict = record.id.clone();
115        let goal = record.goal.clone();
116        let blueprint_ref_json = serde_json::to_string(&record.blueprint_ref)
117            .map_err(|e| TaskStoreError::Other(format!("encode blueprint_ref: {e}")))?;
118        let input_ctx_json = serde_json::to_string(&record.input_ctx)
119            .map_err(|e| TaskStoreError::Other(format!("encode input_ctx: {e}")))?;
120        let status_json = serde_json::to_string(&record.status)
121            .map_err(|e| TaskStoreError::Other(format!("encode status: {e}")))?;
122        let created_at = record.created_at as i64;
123        let updated_at = record.updated_at as i64;
124
125        self.isle
126            .call(move |conn| {
127                let tx = conn.transaction()?;
128                let exists: i64 = tx.query_row(
129                    "SELECT COUNT(*) FROM tasks WHERE id = ?1",
130                    params![id],
131                    |row| row.get(0),
132                )?;
133                if exists > 0 {
134                    return Err(rusqlite::Error::SqliteFailure(
135                        rusqlite::ffi::Error::new(rusqlite::ffi::SQLITE_CONSTRAINT),
136                        Some(format!("__mlua_swarm_duplicate:{id}")),
137                    ));
138                }
139                tx.execute(
140                    "INSERT INTO tasks (id, goal, blueprint_ref_json, input_ctx_json, status, \
141                     created_at, updated_at) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
142                    params![
143                        id,
144                        goal,
145                        blueprint_ref_json,
146                        input_ctx_json,
147                        status_json,
148                        created_at,
149                        updated_at,
150                    ],
151                )?;
152                tx.commit()?;
153                Ok(())
154            })
155            .await
156            .map_err(|e| match &e {
157                IsleError::Sqlite(rusqlite::Error::SqliteFailure(_, Some(msg)))
158                    if msg.starts_with("__mlua_swarm_duplicate:") =>
159                {
160                    TaskStoreError::Duplicate(id_for_conflict.clone())
161                }
162                _ => map_isle_err(e),
163            })
164    }
165
166    async fn get(&self, id: &TaskId) -> Result<TaskRecord, TaskStoreError> {
167        let id_str = id.to_string();
168        let id_for_notfound = id.clone();
169        let row = self
170            .isle
171            .call(move |conn| {
172                conn.query_row(
173                    "SELECT id, goal, blueprint_ref_json, input_ctx_json, status, created_at, \
174                     updated_at FROM tasks WHERE id = ?1",
175                    params![id_str],
176                    |row| {
177                        Ok((
178                            row.get::<_, String>(0)?,
179                            row.get::<_, String>(1)?,
180                            row.get::<_, String>(2)?,
181                            row.get::<_, String>(3)?,
182                            row.get::<_, String>(4)?,
183                            row.get::<_, i64>(5)?,
184                            row.get::<_, i64>(6)?,
185                        ))
186                    },
187                )
188                .optional()
189            })
190            .await
191            .map_err(map_isle_err)?;
192        match row {
193            Some(row) => row_to_record(row),
194            None => Err(TaskStoreError::NotFound(id_for_notfound)),
195        }
196    }
197
198    async fn list(&self) -> Result<Vec<TaskRecord>, TaskStoreError> {
199        let rows = self
200            .isle
201            .call(|conn| {
202                let mut stmt = conn.prepare(
203                    "SELECT id, goal, blueprint_ref_json, input_ctx_json, status, created_at, \
204                     updated_at FROM tasks ORDER BY created_at DESC",
205                )?;
206                let iter = stmt.query_map([], |row| {
207                    Ok((
208                        row.get::<_, String>(0)?,
209                        row.get::<_, String>(1)?,
210                        row.get::<_, String>(2)?,
211                        row.get::<_, String>(3)?,
212                        row.get::<_, String>(4)?,
213                        row.get::<_, i64>(5)?,
214                        row.get::<_, i64>(6)?,
215                    ))
216                })?;
217                let mut out = Vec::new();
218                for r in iter {
219                    out.push(r?);
220                }
221                Ok(out)
222            })
223            .await
224            .map_err(map_isle_err)?;
225        rows.into_iter().map(row_to_record).collect()
226    }
227
228    async fn update_status(
229        &self,
230        id: &TaskId,
231        status: TaskRecordStatus,
232    ) -> Result<(), TaskStoreError> {
233        let id_str = id.to_string();
234        let id_for_notfound = id.clone();
235        let status_json = serde_json::to_string(&status)
236            .map_err(|e| TaskStoreError::Other(format!("encode status: {e}")))?;
237        let updated_at = crate::types::now_unix() as i64;
238        let n = self
239            .isle
240            .call(move |conn| {
241                conn.execute(
242                    "UPDATE tasks SET status = ?1, updated_at = ?2 WHERE id = ?3",
243                    params![status_json, updated_at, id_str],
244                )
245            })
246            .await
247            .map_err(map_isle_err)?;
248        if n == 0 {
249            Err(TaskStoreError::NotFound(id_for_notfound))
250        } else {
251            Ok(())
252        }
253    }
254}
255
256// ──────────────────────────────────────────────────────────────────────────
257// tests
258// ──────────────────────────────────────────────────────────────────────────
259
260#[cfg(test)]
261mod tests {
262    use super::*;
263    use serde_json::json;
264
265    fn mk(id: &str, created_at: u64) -> TaskRecord {
266        TaskRecord {
267            id: TaskId::parse(id).unwrap(),
268            goal: format!("goal for {id}"),
269            blueprint_ref: json!({"id": "bp-1"}),
270            input_ctx: json!({"k": "v"}),
271            status: TaskRecordStatus::Pending,
272            created_at,
273            updated_at: created_at,
274        }
275    }
276
277    #[tokio::test]
278    async fn create_then_get() {
279        let (s, driver) = SqliteTaskStore::open_in_memory().await.unwrap();
280        s.create(mk("T-1", 100)).await.unwrap();
281        let got = s.get(&TaskId::parse("T-1").unwrap()).await.unwrap();
282        assert_eq!(got.goal, "goal for T-1");
283        assert_eq!(got.status, TaskRecordStatus::Pending);
284        assert_eq!(got.blueprint_ref, json!({"id": "bp-1"}));
285        drop(s);
286        driver.shutdown().await.unwrap();
287    }
288
289    #[tokio::test]
290    async fn duplicate_create_rejected() {
291        let (s, driver) = SqliteTaskStore::open_in_memory().await.unwrap();
292        s.create(mk("T-1", 100)).await.unwrap();
293        let err = s.create(mk("T-1", 200)).await.unwrap_err();
294        assert!(matches!(err, TaskStoreError::Duplicate(_)), "got: {err:?}");
295        drop(s);
296        driver.shutdown().await.unwrap();
297    }
298
299    #[tokio::test]
300    async fn get_missing_returns_not_found() {
301        let (s, driver) = SqliteTaskStore::open_in_memory().await.unwrap();
302        let err = s.get(&TaskId::parse("T-nope").unwrap()).await.unwrap_err();
303        assert!(matches!(err, TaskStoreError::NotFound(_)));
304        drop(s);
305        driver.shutdown().await.unwrap();
306    }
307
308    #[tokio::test]
309    async fn list_returns_newest_first() {
310        let (s, driver) = SqliteTaskStore::open_in_memory().await.unwrap();
311        s.create(mk("T-1", 100)).await.unwrap();
312        s.create(mk("T-2", 300)).await.unwrap();
313        s.create(mk("T-3", 200)).await.unwrap();
314        let list = s.list().await.unwrap();
315        let ids: Vec<_> = list.iter().map(|r| r.id.to_string()).collect();
316        assert_eq!(ids, vec!["T-2", "T-3", "T-1"]);
317        drop(s);
318        driver.shutdown().await.unwrap();
319    }
320
321    #[tokio::test]
322    async fn update_status_persists() {
323        let (s, driver) = SqliteTaskStore::open_in_memory().await.unwrap();
324        s.create(mk("T-1", 100)).await.unwrap();
325        s.update_status(&TaskId::parse("T-1").unwrap(), TaskRecordStatus::Failed)
326            .await
327            .unwrap();
328        let got = s.get(&TaskId::parse("T-1").unwrap()).await.unwrap();
329        assert_eq!(got.status, TaskRecordStatus::Failed);
330        assert!(got.updated_at >= got.created_at);
331        drop(s);
332        driver.shutdown().await.unwrap();
333    }
334
335    #[tokio::test]
336    async fn update_status_unknown_fails() {
337        let (s, driver) = SqliteTaskStore::open_in_memory().await.unwrap();
338        let err = s
339            .update_status(&TaskId::parse("T-nope").unwrap(), TaskRecordStatus::Done)
340            .await
341            .unwrap_err();
342        assert!(matches!(err, TaskStoreError::NotFound(_)));
343        drop(s);
344        driver.shutdown().await.unwrap();
345    }
346
347    #[tokio::test]
348    async fn persists_across_reopen() {
349        let dir = tempfile::tempdir().unwrap();
350        let path = dir.path().join("tasks.db");
351
352        {
353            let (s, driver) = SqliteTaskStore::open(&path).await.unwrap();
354            s.create(mk("T-keep", 42)).await.unwrap();
355            drop(s);
356            driver.shutdown().await.unwrap();
357        }
358
359        let (s, driver) = SqliteTaskStore::open(&path).await.unwrap();
360        let got = s.get(&TaskId::parse("T-keep").unwrap()).await.unwrap();
361        assert_eq!(got.goal, "goal for T-keep");
362        assert_eq!(got.created_at, 42);
363        drop(s);
364        driver.shutdown().await.unwrap();
365    }
366}
mlua_swarm/store/task/sqlite.rs

mlua_swarm/store/task/
sqlite.rs
