Skip to main content

prospero_core/
sqlite_store.rs

1//! sqlx-backed sqlite [`Store`] — the default `prosperod` backend.
2//!
3//! One `events` table; `global_ordinal` (the rowid) records durable insertion
4//! order for future fleet-wide queries, while consumers see the per-stream
5//! `seq`. The SQL uses sqlx's runtime query API (no compile-time `DATABASE_URL`)
6//! so the same statements port to Postgres in Phase 2.
7
8use std::path::Path;
9
10use async_trait::async_trait;
11use sqlx::Row;
12use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePool, SqlitePoolOptions};
13
14use crate::error::{CoreError, Result};
15use crate::event::FleetEvent;
16use crate::store::{Store, map_append_error};
17
18/// One-table schema. `global_ordinal` (rowid) is durable insertion order;
19/// `UNIQUE(stream_key, seq)` is the per-stream monotonicity backstop and also
20/// the index `replay` scans.
21const SCHEMA: &str = "CREATE TABLE IF NOT EXISTS events (\
22    global_ordinal INTEGER PRIMARY KEY AUTOINCREMENT,\
23    stream_key TEXT NOT NULL,\
24    seq        INTEGER NOT NULL,\
25    ts         TEXT NOT NULL,\
26    repo       TEXT NOT NULL,\
27    agent_id   TEXT NOT NULL,\
28    kind       TEXT NOT NULL,\
29    actor      TEXT,\
30    on_behalf_of TEXT,\
31    UNIQUE(stream_key, seq)\
32)";
33
34/// sqlx/sqlite-backed durable event store.
35pub struct SqliteStore {
36    pool: SqlitePool,
37}
38
39impl SqliteStore {
40    /// Open (creating it + parent dirs if missing) the store at `dir/events.db`.
41    pub async fn open(dir: &Path) -> Result<Self> {
42        std::fs::create_dir_all(dir)?;
43        let path = dir.join("events.db");
44        let opts = SqliteConnectOptions::new()
45            .filename(&path)
46            .create_if_missing(true)
47            .journal_mode(SqliteJournalMode::Wal)
48            // WAL lets readers run during a write; concurrent *writers* still
49            // serialize, so wait for the write lock instead of erroring with
50            // SQLITE_BUSY immediately.
51            .busy_timeout(std::time::Duration::from_secs(5));
52        let pool = SqlitePoolOptions::new()
53            .connect_with(opts)
54            .await
55            .map_err(|e| CoreError::Store(format!("opening sqlite store: {e}")))?;
56        sqlx::query(SCHEMA)
57            .execute(&pool)
58            .await
59            .map_err(|e| CoreError::Store(format!("initializing sqlite schema: {e}")))?;
60        // #2 / #251: databases created before a column existed lack it. SQLite
61        // has no `ADD COLUMN IF NOT EXISTS`, so probe the table first.
62        for column in ["actor", "on_behalf_of"] {
63            let present: bool = sqlx::query_scalar(
64                "SELECT COUNT(*) > 0 FROM pragma_table_info('events') WHERE name = ?",
65            )
66            .bind(column)
67            .fetch_one(&pool)
68            .await
69            .map_err(|e| CoreError::Store(format!("inspecting sqlite schema: {e}")))?;
70            if !present {
71                sqlx::query(&format!("ALTER TABLE events ADD COLUMN {column} TEXT"))
72                    .execute(&pool)
73                    .await
74                    .map_err(|e| CoreError::Store(format!("adding {column} column: {e}")))?;
75            }
76        }
77        Ok(Self { pool })
78    }
79}
80
81#[async_trait]
82impl Store for SqliteStore {
83    async fn append(&self, event: &FleetEvent) -> Result<()> {
84        let kind = serde_json::to_string(&event.kind)?;
85        sqlx::query(
86            "INSERT INTO events (stream_key, seq, ts, repo, agent_id, kind, actor, on_behalf_of) \
87             VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
88        )
89        .bind(event.stream_key())
90        .bind(event.seq as i64)
91        .bind(&event.ts)
92        .bind(&event.repo)
93        .bind(&event.agent_id)
94        .bind(kind)
95        .bind(&event.actor)
96        .bind(&event.on_behalf_of)
97        .execute(&self.pool)
98        .await
99        .map_err(map_append_error)?;
100        Ok(())
101    }
102
103    async fn replay(&self, stream_key: &str, from_seq: u64) -> Result<Vec<FleetEvent>> {
104        let rows = sqlx::query(
105            "SELECT seq, ts, repo, agent_id, kind, actor, on_behalf_of FROM events \
106             WHERE stream_key = ? AND seq >= ? ORDER BY seq",
107        )
108        .bind(stream_key)
109        .bind(from_seq as i64)
110        .fetch_all(&self.pool)
111        .await
112        .map_err(|e| CoreError::Store(format!("replay: {e}")))?;
113
114        let mut events = Vec::with_capacity(rows.len());
115        for row in rows {
116            let decode = |e: sqlx::Error| CoreError::Store(format!("replay decode: {e}"));
117            let seq: i64 = row.try_get("seq").map_err(decode)?;
118            let ts: String = row.try_get("ts").map_err(decode)?;
119            let repo: String = row.try_get("repo").map_err(decode)?;
120            let agent_id: String = row.try_get("agent_id").map_err(decode)?;
121            let kind_json: String = row.try_get("kind").map_err(decode)?;
122            let actor: Option<String> = row.try_get("actor").map_err(decode)?;
123            let on_behalf_of: Option<String> = row.try_get("on_behalf_of").map_err(decode)?;
124            events.push(FleetEvent {
125                seq: seq as u64,
126                ts,
127                repo,
128                agent_id,
129                kind: serde_json::from_str(&kind_json)?,
130                actor,
131                on_behalf_of,
132            });
133        }
134        Ok(events)
135    }
136
137    /// The cursor is `global_ordinal` — the rowid, i.e. durable insertion
138    /// order (#219).
139    async fn replay_fleet(
140        &self,
141        after_cursor: u64,
142        limit: usize,
143    ) -> Result<Vec<crate::store::CursoredEvent>> {
144        let rows = sqlx::query(
145            "SELECT global_ordinal, seq, ts, repo, agent_id, kind, actor, on_behalf_of FROM events \
146             WHERE global_ordinal > ? ORDER BY global_ordinal LIMIT ?",
147        )
148        .bind(after_cursor as i64)
149        .bind(limit as i64)
150        .fetch_all(&self.pool)
151        .await
152        .map_err(|e| CoreError::Store(format!("replay_fleet: {e}")))?;
153
154        let mut out = Vec::with_capacity(rows.len());
155        for row in rows {
156            let decode = |e: sqlx::Error| CoreError::Store(format!("replay_fleet decode: {e}"));
157            let cursor: i64 = row.try_get("global_ordinal").map_err(decode)?;
158            let seq: i64 = row.try_get("seq").map_err(decode)?;
159            let ts: String = row.try_get("ts").map_err(decode)?;
160            let repo: String = row.try_get("repo").map_err(decode)?;
161            let agent_id: String = row.try_get("agent_id").map_err(decode)?;
162            let kind_json: String = row.try_get("kind").map_err(decode)?;
163            let actor: Option<String> = row.try_get("actor").map_err(decode)?;
164            let on_behalf_of: Option<String> = row.try_get("on_behalf_of").map_err(decode)?;
165            out.push(crate::store::CursoredEvent {
166                cursor: cursor as u64,
167                event: FleetEvent {
168                    seq: seq as u64,
169                    ts,
170                    repo,
171                    agent_id,
172                    kind: serde_json::from_str(&kind_json)?,
173                    actor,
174                    on_behalf_of,
175                },
176            });
177        }
178        Ok(out)
179    }
180
181    async fn latest_fleet_cursor(&self) -> Result<u64> {
182        let row = sqlx::query("SELECT COALESCE(MAX(global_ordinal), 0) AS head FROM events")
183            .fetch_one(&self.pool)
184            .await
185            .map_err(|e| CoreError::Store(format!("latest_fleet_cursor: {e}")))?;
186        let head: i64 = row
187            .try_get("head")
188            .map_err(|e| CoreError::Store(format!("latest_fleet_cursor decode: {e}")))?;
189        Ok(head as u64)
190    }
191
192    async fn high_water(&self, stream_key: &str) -> Result<u64> {
193        let row =
194            sqlx::query("SELECT COALESCE(MAX(seq), 0) AS hw FROM events WHERE stream_key = ?")
195                .bind(stream_key)
196                .fetch_one(&self.pool)
197                .await
198                .map_err(|e| CoreError::Store(format!("high_water: {e}")))?;
199        let hw: i64 = row
200            .try_get("hw")
201            .map_err(|e| CoreError::Store(format!("high_water decode: {e}")))?;
202        Ok(hw as u64)
203    }
204
205    async fn writable(&self) -> bool {
206        // Non-destructive write probe: insert a sentinel row inside a
207        // transaction we always roll back. Exercises the same write path as
208        // `append` (detecting a read-only / full store) without persisting
209        // anything and without DDL. `seq = -1` cannot collide with a real
210        // event (seq is u64) and the rollback ensures it never lands.
211        let Ok(mut tx) = self.pool.begin().await else {
212            return false;
213        };
214        let ok = sqlx::query(
215            "INSERT INTO events (stream_key, seq, ts, repo, agent_id, kind) \
216             VALUES ('__writable_probe__', -1, '', '', '', 'null')",
217        )
218        .execute(&mut *tx)
219        .await
220        .is_ok();
221        let _ = tx.rollback().await;
222        ok
223    }
224
225    async fn prune(&self, before_ts: &str) -> Result<u64> {
226        let res = sqlx::query("DELETE FROM events WHERE ts < ?")
227            .bind(before_ts)
228            .execute(&self.pool)
229            .await
230            .map_err(|e| CoreError::Store(format!("prune: {e}")))?;
231        Ok(res.rows_affected())
232    }
233
234    /// Aggregated in SQL (JSON1 `json_extract` over the `kind` column) so a long
235    /// log is grouped by the database rather than replayed into memory. The
236    /// `WHERE` clause discards every other event kind before grouping, so
237    /// workspaces that only produced output never open a row.
238    async fn usage(&self, since: &str, until: &str) -> Result<Vec<crate::store::UsageRow>> {
239        let rows = sqlx::query(
240            "SELECT repo AS workspace, substr(ts, 1, 10) AS day, \
241                COALESCE(SUM(CASE WHEN json_extract(kind, '$.kind') = 'agent_finished' \
242                    THEN json_extract(kind, '$.cost_usd') END), 0.0) AS cost_usd, \
243                COALESCE(SUM(CASE WHEN json_extract(kind, '$.kind') = 'agent_finished' \
244                    THEN json_extract(kind, '$.turns') END), 0) AS turns, \
245                COALESCE(SUM(CASE WHEN json_extract(kind, '$.to') = 'done' THEN 1 END), 0) AS done, \
246                COALESCE(SUM(CASE WHEN json_extract(kind, '$.to') = 'failed' THEN 1 END), 0) AS failed, \
247                COALESCE(SUM(CASE WHEN json_extract(kind, '$.to') = 'killed' THEN 1 END), 0) AS killed, \
248                COALESCE(SUM(CASE WHEN json_extract(kind, '$.to') = 'crashed' THEN 1 END), 0) AS crashed, \
249                COALESCE(SUM(CASE WHEN json_extract(kind, '$.kind') = 'agent_timed_out' \
250                    THEN 1 END), 0) AS timed_out \
251             FROM events \
252             WHERE ts >= ? AND ts < ? AND ( \
253                json_extract(kind, '$.kind') = 'agent_finished' OR ( \
254                    json_extract(kind, '$.kind') = 'status_changed' \
255                    AND json_extract(kind, '$.to') IN ('done', 'failed', 'killed', 'crashed')) \
256                OR json_extract(kind, '$.kind') = 'agent_timed_out') \
257             GROUP BY repo, substr(ts, 1, 10) \
258             ORDER BY repo, day",
259        )
260        .bind(since)
261        .bind(until)
262        .fetch_all(&self.pool)
263        .await
264        .map_err(|e| CoreError::Store(format!("usage: {e}")))?;
265
266        let decode = |e: sqlx::Error| CoreError::Store(format!("usage decode: {e}"));
267        let mut out = Vec::with_capacity(rows.len());
268        for row in rows {
269            out.push(crate::store::UsageRow {
270                workspace: row.try_get("workspace").map_err(decode)?,
271                day: row.try_get("day").map_err(decode)?,
272                cost_usd: row.try_get::<f64, _>("cost_usd").map_err(decode)?,
273                turns: row.try_get::<i64, _>("turns").map_err(decode)? as u64,
274                done: row.try_get::<i64, _>("done").map_err(decode)? as u64,
275                failed: row.try_get::<i64, _>("failed").map_err(decode)? as u64,
276                killed: row.try_get::<i64, _>("killed").map_err(decode)? as u64,
277                crashed: row.try_get::<i64, _>("crashed").map_err(decode)? as u64,
278                timed_out: row.try_get::<i64, _>("timed_out").map_err(decode)? as u64,
279            });
280        }
281        Ok(out)
282    }
283}
284
285#[cfg(test)]
286mod tests {
287    use super::*;
288
289    fn ev(seq: u64, agent: &str) -> crate::event::FleetEvent {
290        crate::event::FleetEvent {
291            seq,
292            ts: "t".into(),
293            repo: "r".into(),
294            agent_id: agent.into(),
295            kind: crate::event::EventKind::AgentSpawned,
296            actor: None,
297            on_behalf_of: None,
298        }
299    }
300
301    #[tokio::test]
302    async fn sqlite_store_satisfies_conformance() {
303        let dir = tempfile::tempdir().unwrap();
304        let store = SqliteStore::open(dir.path()).await.unwrap();
305        crate::testkit::store_conformance(&store).await;
306    }
307
308    #[tokio::test]
309    async fn sqlite_store_satisfies_fleet_replay_conformance() {
310        let dir = tempfile::tempdir().unwrap();
311        let store = SqliteStore::open(dir.path()).await.unwrap();
312        crate::testkit::store_fleet_replay_conformance(&store).await;
313    }
314
315    #[tokio::test]
316    async fn reopen_resumes_per_stream_high_water() {
317        let dir = tempfile::tempdir().unwrap();
318        {
319            let s = SqliteStore::open(dir.path()).await.unwrap();
320            s.append(&ev(5, "a")).await.unwrap();
321        }
322        let s = SqliteStore::open(dir.path()).await.unwrap();
323        assert_eq!(s.high_water("a").await.unwrap(), 5);
324    }
325
326    #[tokio::test]
327    async fn duplicate_stream_seq_is_rejected() {
328        let dir = tempfile::tempdir().unwrap();
329        let s = SqliteStore::open(dir.path()).await.unwrap();
330        s.append(&ev(1, "a")).await.unwrap();
331        assert!(s.append(&ev(1, "a")).await.is_err());
332    }
333
334    #[tokio::test]
335    async fn sqlite_store_prunes_by_age() {
336        let dir = tempfile::tempdir().unwrap();
337        let store = SqliteStore::open(dir.path()).await.unwrap();
338        crate::testkit::store_prune_conformance(&store).await;
339    }
340
341    #[tokio::test]
342    async fn sqlite_store_aggregates_usage() {
343        let dir = tempfile::tempdir().unwrap();
344        let store = SqliteStore::open(dir.path()).await.unwrap();
345        crate::testkit::store_usage_conformance(&store).await;
346    }
347
348    #[tokio::test]
349    async fn opening_a_pre_actor_database_adds_the_column() {
350        let dir = tempfile::tempdir().unwrap();
351        // A database created by an older prosperod: no `actor` column.
352        {
353            let opts = sqlx::sqlite::SqliteConnectOptions::new()
354                .filename(dir.path().join("events.db"))
355                .create_if_missing(true);
356            let pool = sqlx::sqlite::SqlitePoolOptions::new()
357                .connect_with(opts)
358                .await
359                .unwrap();
360            sqlx::query(
361                "CREATE TABLE events (global_ordinal INTEGER PRIMARY KEY AUTOINCREMENT, \
362                 stream_key TEXT NOT NULL, seq INTEGER NOT NULL, ts TEXT NOT NULL, \
363                 repo TEXT NOT NULL, agent_id TEXT NOT NULL, kind TEXT NOT NULL, \
364                 UNIQUE(stream_key, seq))",
365            )
366            .execute(&pool)
367            .await
368            .unwrap();
369            pool.close().await;
370        }
371        let store = SqliteStore::open(dir.path()).await.unwrap();
372        // Opening twice must also be fine (idempotent migration).
373        drop(store);
374        let store = SqliteStore::open(dir.path()).await.unwrap();
375        crate::testkit::store_conformance(&store).await;
376    }
377}