1use 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
47pub struct SqliteRunStore {
55 isle: AsyncIsle,
56}
57
58impl SqliteRunStore {
59 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 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
83type 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 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#[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}