1use 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
18const 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
34pub struct SqliteStore {
36 pool: SqlitePool,
37}
38
39impl SqliteStore {
40 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 .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 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 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 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 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 {
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 drop(store);
374 let store = SqliteStore::open(dir.path()).await.unwrap();
375 crate::testkit::store_conformance(&store).await;
376 }
377}