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 UNIQUE(stream_key, seq)\
30)";
31
32pub struct SqliteStore {
34 pool: SqlitePool,
35}
36
37impl SqliteStore {
38 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 .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 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 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}