Skip to main content

orbit_runtime/store/
sqlite.rs

1//! SQLite backend (default; a single file, WAL mode with `synchronous=FULL`).
2
3use super::{
4    ExecutionSummary, FireRecord, PutDefinition, Result, Rollover, ScheduleRow, WorkerRecord,
5    SCHEMA,
6};
7use crate::model::{Event, ExecutionState, HistoryEvent};
8use orbit_ir::WorkflowIr;
9use rusqlite::{params, Connection, OptionalExtension};
10use std::path::Path;
11
12pub struct SqliteStore {
13    conn: Connection,
14}
15
16impl SqliteStore {
17    pub fn open(path: &Path) -> Result<SqliteStore> {
18        let conn = Connection::open(path)?;
19        Self::init(conn)
20    }
21
22    pub fn in_memory() -> Result<SqliteStore> {
23        Self::init(Connection::open_in_memory()?)
24    }
25
26    fn init(conn: Connection) -> Result<SqliteStore> {
27        // WAL + FULL sync: a committed event survives power loss.
28        conn.pragma_update(None, "journal_mode", "WAL")?;
29        conn.pragma_update(None, "synchronous", "FULL")?;
30        conn.pragma_update(None, "foreign_keys", "ON")?;
31        conn.busy_timeout(std::time::Duration::from_secs(5))?;
32        conn.execute_batch(&SCHEMA.replace("BIGINT", "INTEGER"))?;
33        if conn.query_row("SELECT COUNT(*) FROM schema_version", [], |r| {
34            r.get::<_, i64>(0)
35        })? == 0
36        {
37            conn.execute("INSERT INTO schema_version (version) VALUES (1)", [])?;
38        }
39        Ok(SqliteStore { conn })
40    }
41
42    // ------------------------------------------------------------ definitions
43
44    pub fn put_definition(&self, ir: &WorkflowIr, now: u64) -> Result<PutDefinition> {
45        let existing: Option<String> = self
46            .conn
47            .query_row(
48                "SELECT definition_hash FROM workflow_definitions WHERE name = ?1 AND version = ?2",
49                params![ir.name, ir.version],
50                |r| r.get(0),
51            )
52            .optional()?;
53        match existing {
54            Some(h) if h == ir.definition_hash => Ok(PutDefinition::Unchanged),
55            Some(h) => Ok(PutDefinition::Conflict { existing_hash: h }),
56            None => {
57                let text = serde_json::to_string(ir).expect("IR serialises");
58                self.conn.execute(
59                    "INSERT INTO workflow_definitions (name, version, definition_hash, ir, created_at_ms) VALUES (?1, ?2, ?3, ?4, ?5)",
60                    params![ir.name, ir.version, ir.definition_hash, text, now as i64],
61                )?;
62                Ok(PutDefinition::Created)
63            }
64        }
65    }
66
67    pub fn definitions(&self) -> Result<Vec<WorkflowIr>> {
68        let mut stmt = self
69            .conn
70            .prepare("SELECT ir FROM workflow_definitions ORDER BY name, created_at_ms")?;
71        let rows = stmt.query_map([], |r| r.get::<_, String>(0))?;
72        let mut out = vec![];
73        for text in rows {
74            match serde_json::from_str::<WorkflowIr>(&text?) {
75                Ok(ir) => out.push(ir),
76                Err(e) => tracing::error!(error = %e, "skipping unreadable workflow definition"),
77            }
78        }
79        Ok(out)
80    }
81
82    // ------------------------------------------------------------- executions
83
84    pub fn execution_exists(&self, id: &str) -> Result<bool> {
85        Ok(self
86            .conn
87            .query_row(
88                "SELECT 1 FROM workflow_executions WHERE execution_id = ?1",
89                params![id],
90                |_| Ok(()),
91            )
92            .optional()?
93            .is_some())
94    }
95
96    /// Atomically append events and update the execution projection.
97    pub fn append(&mut self, state: &ExecutionState, events: &[HistoryEvent]) -> Result<()> {
98        let tx = self.conn.transaction()?;
99        write_events(&tx, state, events)?;
100        Ok(tx.commit()?)
101    }
102
103    pub fn append_rollover(
104        &mut self,
105        state: &ExecutionState,
106        events: &[HistoryEvent],
107        rollover: Rollover<'_>,
108    ) -> Result<()> {
109        let tx = self.conn.transaction()?;
110        let run = rollover.run as i64;
111        {
112            let mut ins = tx.prepare_cached(
113                "INSERT INTO execution_run_events (execution_id, run, seq, event_type, step, payload, timestamp_ms) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
114            )?;
115            for h in rollover.events {
116                let payload = serde_json::to_string(&h.event).expect("event serialises");
117                ins.execute(params![
118                    state.id,
119                    run,
120                    h.seq as i64,
121                    h.event.name(),
122                    h.event.step(),
123                    payload,
124                    h.timestamp_ms as i64
125                ])?;
126            }
127        }
128        tx.execute(
129            "INSERT INTO execution_run_events (execution_id, run, seq, event_type, step, payload, timestamp_ms)
130             SELECT execution_id, ?2, seq, event_type, step, payload, timestamp_ms FROM execution_events WHERE execution_id = ?1",
131            params![state.id, run],
132        )?;
133        tx.execute(
134            "DELETE FROM execution_events WHERE execution_id = ?1",
135            params![state.id],
136        )?;
137        write_events(&tx, state, events)?;
138        Ok(tx.commit()?)
139    }
140
141    pub fn run_history(&self, id: &str, run: u32) -> Result<Vec<HistoryEvent>> {
142        let mut stmt = self.conn.prepare_cached(
143            "SELECT seq, timestamp_ms, payload FROM execution_run_events WHERE execution_id = ?1 AND run = ?2 ORDER BY seq",
144        )?;
145        let rows = stmt.query_map(params![id, run as i64], |r| {
146            Ok((
147                r.get::<_, i64>(0)?,
148                r.get::<_, i64>(1)?,
149                r.get::<_, String>(2)?,
150            ))
151        })?;
152        events_from(rows)
153    }
154
155    pub fn execution_status(&self, id: &str) -> Result<Option<String>> {
156        Ok(self
157            .conn
158            .query_row(
159                "SELECT status FROM workflow_executions WHERE execution_id = ?1",
160                params![id],
161                |r| r.get(0),
162            )
163            .optional()?)
164    }
165}
166
167/// Write events and the projection inside an open transaction.
168fn write_events(
169    tx: &rusqlite::Transaction<'_>,
170    state: &ExecutionState,
171    events: &[HistoryEvent],
172) -> Result<()> {
173    {
174        let mut ins = tx.prepare_cached(
175            "INSERT INTO execution_events (execution_id, seq, event_type, step, payload, timestamp_ms) VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
176        )?;
177        for h in events {
178            let payload = serde_json::to_string(&h.event).expect("event serialises");
179            ins.execute(params![
180                state.id,
181                h.seq as i64,
182                h.event.name(),
183                h.event.step(),
184                payload,
185                h.timestamp_ms as i64
186            ])?;
187        }
188    }
189    {
190        tx.execute(
191            "INSERT INTO workflow_executions (execution_id, workflow, version, definition_hash, status, created_at_ms, updated_at_ms, last_seq)
192             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)
193             ON CONFLICT(execution_id) DO UPDATE SET status = excluded.status, updated_at_ms = excluded.updated_at_ms, last_seq = excluded.last_seq",
194            params![
195                state.id,
196                state.workflow,
197                state.version,
198                state.definition_hash,
199                state.status.as_str(),
200                state.created_at_ms as i64,
201                state.updated_at_ms as i64,
202                state.last_seq as i64
203            ],
204        )?;
205    }
206    Ok(())
207}
208
209fn events_from(
210    rows: impl Iterator<Item = rusqlite::Result<(i64, i64, String)>>,
211) -> Result<Vec<HistoryEvent>> {
212    let mut out = vec![];
213    for row in rows {
214        let (seq, ts, payload) = row?;
215        let event: Event = serde_json::from_str(&payload).map_err(|e| {
216            rusqlite::Error::FromSqlConversionFailure(2, rusqlite::types::Type::Text, Box::new(e))
217        })?;
218        out.push(HistoryEvent {
219            seq: seq as u64,
220            timestamp_ms: ts as u64,
221            event,
222        });
223    }
224    Ok(out)
225}
226
227const SCHEDULE_COLUMNS: &str = "schedule_id, spec, paused, next_fire_ms, due_ms, last_fire_ms, last_execution_id, fired_count, skipped_count, last_error, created_at_ms, updated_at_ms";
228
229fn schedule_row(r: &rusqlite::Row<'_>) -> rusqlite::Result<ScheduleRow> {
230    let opt = |v: Option<i64>| v.map(|x| x as u64);
231    Ok(ScheduleRow {
232        schedule_id: r.get(0)?,
233        spec: r.get(1)?,
234        paused: r.get::<_, i64>(2)? != 0,
235        next_fire_ms: opt(r.get(3)?),
236        due_ms: opt(r.get(4)?),
237        last_fire_ms: opt(r.get(5)?),
238        last_execution_id: r.get(6)?,
239        fired: r.get::<_, i64>(7)? as u64,
240        skipped: r.get::<_, i64>(8)? as u64,
241        last_error: r.get(9)?,
242        created_at_ms: r.get::<_, i64>(10)? as u64,
243        updated_at_ms: r.get::<_, i64>(11)? as u64,
244    })
245}
246
247fn opt_i64(v: Option<u64>) -> Option<i64> {
248    v.map(|x| x as i64)
249}
250
251impl SqliteStore {
252    pub fn insert_schedule(&self, row: &ScheduleRow) -> Result<bool> {
253        let n = self.conn.execute(
254            &format!(
255                "INSERT INTO orbit_schedules ({SCHEDULE_COLUMNS}) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12)
256                 ON CONFLICT (schedule_id) DO NOTHING"
257            ),
258            params![
259                row.schedule_id,
260                row.spec,
261                row.paused as i64,
262                opt_i64(row.next_fire_ms),
263                opt_i64(row.due_ms),
264                opt_i64(row.last_fire_ms),
265                row.last_execution_id,
266                row.fired as i64,
267                row.skipped as i64,
268                row.last_error,
269                row.created_at_ms as i64,
270                row.updated_at_ms as i64
271            ],
272        )?;
273        Ok(n == 1)
274    }
275
276    pub fn update_schedule(&self, row: &ScheduleRow) -> Result<bool> {
277        let n = self.conn.execute(
278            "UPDATE orbit_schedules SET spec = ?2, paused = ?3, next_fire_ms = ?4, due_ms = ?5, last_error = ?6,
279             claimed_by = NULL, claimed_until_ms = NULL, updated_at_ms = ?7 WHERE schedule_id = ?1",
280            params![
281                row.schedule_id,
282                row.spec,
283                row.paused as i64,
284                opt_i64(row.next_fire_ms),
285                opt_i64(row.due_ms),
286                row.last_error,
287                row.updated_at_ms as i64
288            ],
289        )?;
290        Ok(n == 1)
291    }
292
293    pub fn schedule(&self, id: &str) -> Result<Option<ScheduleRow>> {
294        Ok(self
295            .conn
296            .query_row(
297                &format!("SELECT {SCHEDULE_COLUMNS} FROM orbit_schedules WHERE schedule_id = ?1"),
298                params![id],
299                schedule_row,
300            )
301            .optional()?)
302    }
303
304    pub fn schedules(&self) -> Result<Vec<ScheduleRow>> {
305        let mut stmt = self.conn.prepare(&format!(
306            "SELECT {SCHEDULE_COLUMNS} FROM orbit_schedules ORDER BY schedule_id"
307        ))?;
308        let rows = stmt.query_map([], schedule_row)?;
309        Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
310    }
311
312    pub fn delete_schedule(&self, id: &str) -> Result<bool> {
313        Ok(self.conn.execute(
314            "DELETE FROM orbit_schedules WHERE schedule_id = ?1",
315            params![id],
316        )? == 1)
317    }
318
319    pub fn due_schedules(&self, now: u64) -> Result<Vec<ScheduleRow>> {
320        let mut stmt = self.conn.prepare_cached(&format!(
321            "SELECT {SCHEDULE_COLUMNS} FROM orbit_schedules
322             WHERE paused = 0 AND due_ms IS NOT NULL AND due_ms <= ?1
323             AND (claimed_until_ms IS NULL OR claimed_until_ms <= ?1)
324             ORDER BY due_ms LIMIT 100"
325        ))?;
326        let rows = stmt.query_map(params![now as i64], schedule_row)?;
327        Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
328    }
329
330    pub fn claim_schedule(
331        &self,
332        id: &str,
333        next_fire_ms: u64,
334        node: &str,
335        now: u64,
336        until: u64,
337    ) -> Result<bool> {
338        Ok(self.conn.execute(
339            "UPDATE orbit_schedules SET claimed_by = ?3, claimed_until_ms = ?5
340             WHERE schedule_id = ?1 AND next_fire_ms = ?2 AND paused = 0
341             AND (claimed_until_ms IS NULL OR claimed_until_ms <= ?4)",
342            params![id, next_fire_ms as i64, node, now as i64, until as i64],
343        )? == 1)
344    }
345
346    pub fn finish_schedule_fire(
347        &self,
348        id: &str,
349        next_fire_ms: u64,
350        node: &str,
351        now: u64,
352        rec: &FireRecord,
353    ) -> Result<bool> {
354        Ok(self.conn.execute(
355            "UPDATE orbit_schedules SET next_fire_ms = ?4, due_ms = ?5,
356             last_fire_ms = COALESCE(?6, last_fire_ms), last_execution_id = COALESCE(?7, last_execution_id),
357             fired_count = fired_count + ?8, skipped_count = skipped_count + ?9, last_error = ?10,
358             claimed_by = NULL, claimed_until_ms = NULL, updated_at_ms = ?11
359             WHERE schedule_id = ?1 AND next_fire_ms = ?2 AND claimed_by = ?3",
360            params![
361                id,
362                next_fire_ms as i64,
363                node,
364                opt_i64(rec.next_fire_ms),
365                opt_i64(rec.due_ms),
366                opt_i64(rec.last_fire_ms),
367                rec.last_execution_id,
368                rec.fired as i64,
369                rec.skipped as i64,
370                rec.last_error,
371                now as i64
372            ],
373        )? == 1)
374    }
375
376    pub fn release_schedule_claim(
377        &self,
378        id: &str,
379        node: &str,
380        retry_at: u64,
381        error: &str,
382        now: u64,
383    ) -> Result<()> {
384        self.conn.execute(
385            "UPDATE orbit_schedules SET claimed_by = NULL, claimed_until_ms = ?3, last_error = ?4, updated_at_ms = ?5
386             WHERE schedule_id = ?1 AND claimed_by = ?2",
387            params![id, node, retry_at as i64, error, now as i64],
388        )?;
389        Ok(())
390    }
391
392    pub fn release_schedule_claims(&self) -> Result<u64> {
393        Ok(self.conn.execute(
394            "UPDATE orbit_schedules SET claimed_by = NULL, claimed_until_ms = NULL WHERE claimed_by IS NOT NULL",
395            [],
396        )? as u64)
397    }
398
399    pub fn record_schedule_trigger(&self, id: &str, execution_id: &str, now: u64) -> Result<()> {
400        self.conn.execute(
401            "UPDATE orbit_schedules SET last_execution_id = ?2, fired_count = fired_count + 1, updated_at_ms = ?3
402             WHERE schedule_id = ?1",
403            params![id, execution_id, now as i64],
404        )?;
405        Ok(())
406    }
407}
408
409impl SqliteStore {
410    pub fn history(&self, id: &str) -> Result<Vec<HistoryEvent>> {
411        let mut stmt = self.conn.prepare_cached(
412            "SELECT seq, timestamp_ms, payload FROM execution_events WHERE execution_id = ?1 ORDER BY seq",
413        )?;
414        let rows = stmt.query_map(params![id], |r| {
415            Ok((
416                r.get::<_, i64>(0)?,
417                r.get::<_, i64>(1)?,
418                r.get::<_, String>(2)?,
419            ))
420        })?;
421        events_from(rows)
422    }
423
424    /// Ids of executions that have not reached a terminal status.
425    pub fn active_execution_ids(&self) -> Result<Vec<String>> {
426        let mut stmt = self.conn.prepare(
427            "SELECT execution_id FROM workflow_executions WHERE status IN ('RUNNING', 'COMPENSATING') ORDER BY created_at_ms",
428        )?;
429        let rows = stmt.query_map([], |r| r.get(0))?;
430        Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
431    }
432
433    pub fn list_executions(
434        &self,
435        status: Option<&str>,
436        workflow: Option<&str>,
437        limit: u32,
438    ) -> Result<Vec<ExecutionSummary>> {
439        let mut stmt = self.conn.prepare(
440            "SELECT execution_id, workflow, version, status, created_at_ms, updated_at_ms FROM workflow_executions
441             WHERE (?1 IS NULL OR status = ?1) AND (?2 IS NULL OR workflow = ?2)
442             ORDER BY created_at_ms DESC LIMIT ?3",
443        )?;
444        let rows = stmt.query_map(params![status, workflow, limit], |r| {
445            Ok(ExecutionSummary {
446                execution_id: r.get(0)?,
447                workflow: r.get(1)?,
448                version: r.get(2)?,
449                status: r.get(3)?,
450                created_at_ms: r.get::<_, i64>(4)? as u64,
451                updated_at_ms: r.get::<_, i64>(5)? as u64,
452            })
453        })?;
454        Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
455    }
456
457    pub fn count_by_status(&self) -> Result<Vec<(String, u64)>> {
458        let mut stmt = self
459            .conn
460            .prepare("SELECT status, COUNT(*) FROM workflow_executions GROUP BY status")?;
461        let rows = stmt.query_map([], |r| Ok((r.get(0)?, r.get::<_, i64>(1)? as u64)))?;
462        Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
463    }
464
465    // ---------------------------------------------------------------- workers
466
467    pub fn upsert_worker(&self, w: &WorkerRecord) -> Result<()> {
468        self.conn.execute(
469            "INSERT INTO worker_registry (worker_id, runtime, sdk, contracts, registered_at_ms, last_seen_ms) VALUES (?1, ?2, ?3, ?4, ?5, ?6)
470             ON CONFLICT(worker_id) DO UPDATE SET runtime = excluded.runtime, sdk = excluded.sdk, contracts = excluded.contracts,
471             registered_at_ms = excluded.registered_at_ms, last_seen_ms = excluded.last_seen_ms",
472            params![w.worker_id, w.runtime, w.sdk, w.contracts.to_string(), w.registered_at_ms as i64, w.last_seen_ms as i64],
473        )?;
474        Ok(())
475    }
476
477    pub fn touch_worker(&self, worker_id: &str, now: u64) -> Result<()> {
478        self.conn.execute(
479            "UPDATE worker_registry SET last_seen_ms = ?2 WHERE worker_id = ?1",
480            params![worker_id, now as i64],
481        )?;
482        Ok(())
483    }
484
485    pub fn workers(&self) -> Result<Vec<WorkerRecord>> {
486        let mut stmt = self.conn.prepare(
487            "SELECT worker_id, runtime, sdk, contracts, registered_at_ms, last_seen_ms FROM worker_registry ORDER BY worker_id",
488        )?;
489        let rows = stmt.query_map([], |r| {
490            Ok(WorkerRecord {
491                worker_id: r.get(0)?,
492                runtime: r.get(1)?,
493                sdk: r.get(2)?,
494                contracts: serde_json::from_str(&r.get::<_, String>(3)?).unwrap_or_default(),
495                registered_at_ms: r.get::<_, i64>(4)? as u64,
496                last_seen_ms: r.get::<_, i64>(5)? as u64,
497            })
498        })?;
499        Ok(rows.collect::<std::result::Result<Vec<_>, _>>()?)
500    }
501}