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    UNIQUE(stream_key, seq)\
30)";
31
32/// sqlx/sqlite-backed durable event store.
33pub struct SqliteStore {
34    pool: SqlitePool,
35}
36
37impl SqliteStore {
38    /// Open (creating it + parent dirs if missing) the store at `dir/events.db`.
39    pub async fn open(dir: &Path) -> Result<Self> {
40        std::fs::create_dir_all(dir)?;
41        let path = dir.join("events.db");
42        let opts = SqliteConnectOptions::new()
43            .filename(&path)
44            .create_if_missing(true)
45            .journal_mode(SqliteJournalMode::Wal)
46            // WAL lets readers run during a write; concurrent *writers* still
47            // serialize, so wait for the write lock instead of erroring with
48            // SQLITE_BUSY immediately.
49            .busy_timeout(std::time::Duration::from_secs(5));
50        let pool = SqlitePoolOptions::new()
51            .connect_with(opts)
52            .await
53            .map_err(|e| CoreError::Store(format!("opening sqlite store: {e}")))?;
54        sqlx::query(SCHEMA)
55            .execute(&pool)
56            .await
57            .map_err(|e| CoreError::Store(format!("initializing sqlite schema: {e}")))?;
58        Ok(Self { pool })
59    }
60}
61
62#[async_trait]
63impl Store for SqliteStore {
64    async fn append(&self, event: &FleetEvent) -> Result<()> {
65        let kind = serde_json::to_string(&event.kind)?;
66        sqlx::query(
67            "INSERT INTO events (stream_key, seq, ts, repo, agent_id, kind) \
68             VALUES (?, ?, ?, ?, ?, ?)",
69        )
70        .bind(event.stream_key())
71        .bind(event.seq as i64)
72        .bind(&event.ts)
73        .bind(&event.repo)
74        .bind(&event.agent_id)
75        .bind(kind)
76        .execute(&self.pool)
77        .await
78        .map_err(map_append_error)?;
79        Ok(())
80    }
81
82    async fn replay(&self, stream_key: &str, from_seq: u64) -> Result<Vec<FleetEvent>> {
83        let rows = sqlx::query(
84            "SELECT seq, ts, repo, agent_id, kind FROM events \
85             WHERE stream_key = ? AND seq >= ? ORDER BY seq",
86        )
87        .bind(stream_key)
88        .bind(from_seq as i64)
89        .fetch_all(&self.pool)
90        .await
91        .map_err(|e| CoreError::Store(format!("replay: {e}")))?;
92
93        let mut events = Vec::with_capacity(rows.len());
94        for row in rows {
95            let decode = |e: sqlx::Error| CoreError::Store(format!("replay decode: {e}"));
96            let seq: i64 = row.try_get("seq").map_err(decode)?;
97            let ts: String = row.try_get("ts").map_err(decode)?;
98            let repo: String = row.try_get("repo").map_err(decode)?;
99            let agent_id: String = row.try_get("agent_id").map_err(decode)?;
100            let kind_json: String = row.try_get("kind").map_err(decode)?;
101            events.push(FleetEvent {
102                seq: seq as u64,
103                ts,
104                repo,
105                agent_id,
106                kind: serde_json::from_str(&kind_json)?,
107            });
108        }
109        Ok(events)
110    }
111
112    async fn high_water(&self, stream_key: &str) -> Result<u64> {
113        let row =
114            sqlx::query("SELECT COALESCE(MAX(seq), 0) AS hw FROM events WHERE stream_key = ?")
115                .bind(stream_key)
116                .fetch_one(&self.pool)
117                .await
118                .map_err(|e| CoreError::Store(format!("high_water: {e}")))?;
119        let hw: i64 = row
120            .try_get("hw")
121            .map_err(|e| CoreError::Store(format!("high_water decode: {e}")))?;
122        Ok(hw as u64)
123    }
124
125    async fn writable(&self) -> bool {
126        // Non-destructive write probe: insert a sentinel row inside a
127        // transaction we always roll back. Exercises the same write path as
128        // `append` (detecting a read-only / full store) without persisting
129        // anything and without DDL. `seq = -1` cannot collide with a real
130        // event (seq is u64) and the rollback ensures it never lands.
131        let Ok(mut tx) = self.pool.begin().await else {
132            return false;
133        };
134        let ok = sqlx::query(
135            "INSERT INTO events (stream_key, seq, ts, repo, agent_id, kind) \
136             VALUES ('__writable_probe__', -1, '', '', '', 'null')",
137        )
138        .execute(&mut *tx)
139        .await
140        .is_ok();
141        let _ = tx.rollback().await;
142        ok
143    }
144
145    async fn prune(&self, before_ts: &str) -> Result<u64> {
146        let res = sqlx::query("DELETE FROM events WHERE ts < ?")
147            .bind(before_ts)
148            .execute(&self.pool)
149            .await
150            .map_err(|e| CoreError::Store(format!("prune: {e}")))?;
151        Ok(res.rows_affected())
152    }
153
154    /// Aggregated in SQL (JSON1 `json_extract` over the `kind` column) so a long
155    /// log is grouped by the database rather than replayed into memory. The
156    /// `WHERE` clause discards every other event kind before grouping, so
157    /// workspaces that only produced output never open a row.
158    async fn usage(&self, since: &str, until: &str) -> Result<Vec<crate::store::UsageRow>> {
159        let rows = sqlx::query(
160            "SELECT repo AS workspace, substr(ts, 1, 10) AS day, \
161                COALESCE(SUM(CASE WHEN json_extract(kind, '$.kind') = 'agent_finished' \
162                    THEN json_extract(kind, '$.cost_usd') END), 0.0) AS cost_usd, \
163                COALESCE(SUM(CASE WHEN json_extract(kind, '$.kind') = 'agent_finished' \
164                    THEN json_extract(kind, '$.turns') END), 0) AS turns, \
165                COALESCE(SUM(CASE WHEN json_extract(kind, '$.to') = 'done' THEN 1 END), 0) AS done, \
166                COALESCE(SUM(CASE WHEN json_extract(kind, '$.to') = 'failed' THEN 1 END), 0) AS failed, \
167                COALESCE(SUM(CASE WHEN json_extract(kind, '$.to') = 'killed' THEN 1 END), 0) AS killed, \
168                COALESCE(SUM(CASE WHEN json_extract(kind, '$.to') = 'crashed' THEN 1 END), 0) AS crashed \
169             FROM events \
170             WHERE ts >= ? AND ts < ? AND ( \
171                json_extract(kind, '$.kind') = 'agent_finished' OR ( \
172                    json_extract(kind, '$.kind') = 'status_changed' \
173                    AND json_extract(kind, '$.to') IN ('done', 'failed', 'killed', 'crashed'))) \
174             GROUP BY repo, substr(ts, 1, 10) \
175             ORDER BY repo, day",
176        )
177        .bind(since)
178        .bind(until)
179        .fetch_all(&self.pool)
180        .await
181        .map_err(|e| CoreError::Store(format!("usage: {e}")))?;
182
183        let decode = |e: sqlx::Error| CoreError::Store(format!("usage decode: {e}"));
184        let mut out = Vec::with_capacity(rows.len());
185        for row in rows {
186            out.push(crate::store::UsageRow {
187                workspace: row.try_get("workspace").map_err(decode)?,
188                day: row.try_get("day").map_err(decode)?,
189                cost_usd: row.try_get::<f64, _>("cost_usd").map_err(decode)?,
190                turns: row.try_get::<i64, _>("turns").map_err(decode)? as u64,
191                done: row.try_get::<i64, _>("done").map_err(decode)? as u64,
192                failed: row.try_get::<i64, _>("failed").map_err(decode)? as u64,
193                killed: row.try_get::<i64, _>("killed").map_err(decode)? as u64,
194                crashed: row.try_get::<i64, _>("crashed").map_err(decode)? as u64,
195            });
196        }
197        Ok(out)
198    }
199}
200
201#[cfg(test)]
202mod tests {
203    use super::*;
204
205    fn ev(seq: u64, agent: &str) -> crate::event::FleetEvent {
206        crate::event::FleetEvent {
207            seq,
208            ts: "t".into(),
209            repo: "r".into(),
210            agent_id: agent.into(),
211            kind: crate::event::EventKind::AgentSpawned,
212        }
213    }
214
215    #[tokio::test]
216    async fn sqlite_store_satisfies_conformance() {
217        let dir = tempfile::tempdir().unwrap();
218        let store = SqliteStore::open(dir.path()).await.unwrap();
219        crate::testkit::store_conformance(&store).await;
220    }
221
222    #[tokio::test]
223    async fn reopen_resumes_per_stream_high_water() {
224        let dir = tempfile::tempdir().unwrap();
225        {
226            let s = SqliteStore::open(dir.path()).await.unwrap();
227            s.append(&ev(5, "a")).await.unwrap();
228        }
229        let s = SqliteStore::open(dir.path()).await.unwrap();
230        assert_eq!(s.high_water("a").await.unwrap(), 5);
231    }
232
233    #[tokio::test]
234    async fn duplicate_stream_seq_is_rejected() {
235        let dir = tempfile::tempdir().unwrap();
236        let s = SqliteStore::open(dir.path()).await.unwrap();
237        s.append(&ev(1, "a")).await.unwrap();
238        assert!(s.append(&ev(1, "a")).await.is_err());
239    }
240
241    #[tokio::test]
242    async fn sqlite_store_prunes_by_age() {
243        let dir = tempfile::tempdir().unwrap();
244        let store = SqliteStore::open(dir.path()).await.unwrap();
245        crate::testkit::store_prune_conformance(&store).await;
246    }
247
248    #[tokio::test]
249    async fn sqlite_store_aggregates_usage() {
250        let dir = tempfile::tempdir().unwrap();
251        let store = SqliteStore::open(dir.path()).await.unwrap();
252        crate::testkit::store_usage_conformance(&store).await;
253    }
254}