prospero_core/
postgres_store.rs

1//! sqlx-backed Postgres [`Store`] — the clustered-tier event backend.
2//!
3//! Mirrors [`crate::sqlite_store::SqliteStore`] with Postgres dialect (`$N`
4//! placeholders, `BIGSERIAL`/`BIGINT`). Runs the same `testkit` conformance
5//! batteries, gated on `DATABASE_URL` (skipped when unset). See spec §3/§4.
6
7use async_trait::async_trait;
8use sqlx::Row;
9use sqlx::postgres::PgPool;
10
11use crate::error::{CoreError, Result};
12use crate::event::FleetEvent;
13use crate::store::{Store, map_append_error};
14
15const SCHEMA: &str = "CREATE TABLE IF NOT EXISTS events (\
16    global_ordinal BIGSERIAL PRIMARY KEY,\
17    stream_key TEXT NOT NULL,\
18    seq        BIGINT NOT NULL,\
19    ts         TEXT NOT NULL,\
20    repo       TEXT NOT NULL,\
21    agent_id   TEXT NOT NULL,\
22    kind       TEXT NOT NULL,\
23    actor      TEXT,\
24    on_behalf_of TEXT,\
25    UNIQUE(stream_key, seq)\
26)";
27
28/// sqlx/Postgres-backed durable event store (clustered tier).
29pub struct PostgresStore {
30    pool: PgPool,
31}
32
33impl PostgresStore {
34    /// Connect to Postgres at `url` and ensure the schema exists.
35    pub async fn connect(url: &str) -> Result<Self> {
36        let pool = crate::pg::connect(url).await?;
37        crate::pg::ensure_schema(&pool, SCHEMA, "events table").await?;
38        // `ALTER TABLE ... ADD COLUMN IF NOT EXISTS` takes an ACCESS EXCLUSIVE
39        // lock even when the column already exists (the common case on every
40        // boot after the first), so probe first and only run the ALTER when
41        // it's actually missing. Mirrors SqliteStore::open's probe-then-ALTER.
42        for column in ["actor", "on_behalf_of"] {
43            let present: bool = sqlx::query_scalar(
44                "SELECT EXISTS (SELECT 1 FROM information_schema.columns \
45                 WHERE table_name = 'events' AND column_name = $1 \
46                 AND table_schema = current_schema())",
47            )
48            .bind(column)
49            .fetch_one(&pool)
50            .await
51            .map_err(|e| CoreError::Store(format!("inspecting postgres schema: {e}")))?;
52            if !present {
53                crate::pg::ensure_schema(
54                    &pool,
55                    &format!("ALTER TABLE events ADD COLUMN IF NOT EXISTS {column} TEXT"),
56                    "events column",
57                )
58                .await?;
59            }
60        }
61        Ok(Self { pool })
62    }
63
64    /// Truncate all events. Test-only (resets between conformance batteries).
65    #[cfg(any(test, feature = "testkit"))]
66    pub async fn reset_for_tests(&self) -> Result<()> {
67        sqlx::query("TRUNCATE events")
68            .execute(&self.pool)
69            .await
70            .map_err(|e| CoreError::Store(format!("reset: {e}")))?;
71        Ok(())
72    }
73}
74
75#[async_trait]
76impl Store for PostgresStore {
77    async fn append(&self, event: &FleetEvent) -> Result<()> {
78        let kind = serde_json::to_string(&event.kind)?;
79        sqlx::query(
80            "INSERT INTO events (stream_key, seq, ts, repo, agent_id, kind, actor, on_behalf_of) \
81             VALUES ($1, $2, $3, $4, $5, $6, $7, $8)",
82        )
83        .bind(event.stream_key())
84        .bind(event.seq as i64)
85        .bind(&event.ts)
86        .bind(&event.repo)
87        .bind(&event.agent_id)
88        .bind(kind)
89        .bind(&event.actor)
90        .bind(&event.on_behalf_of)
91        .execute(&self.pool)
92        .await
93        .map_err(map_append_error)?;
94        Ok(())
95    }
96
97    async fn replay(&self, stream_key: &str, from_seq: u64) -> Result<Vec<FleetEvent>> {
98        let rows = sqlx::query(
99            "SELECT seq, ts, repo, agent_id, kind, actor, on_behalf_of FROM events \
100             WHERE stream_key = $1 AND seq >= $2 ORDER BY seq",
101        )
102        .bind(stream_key)
103        .bind(from_seq as i64)
104        .fetch_all(&self.pool)
105        .await
106        .map_err(|e| CoreError::Store(format!("replay: {e}")))?;
107
108        let mut events = Vec::with_capacity(rows.len());
109        for row in rows {
110            let decode = |e: sqlx::Error| CoreError::Store(format!("replay decode: {e}"));
111            let seq: i64 = row.try_get("seq").map_err(decode)?;
112            let ts: String = row.try_get("ts").map_err(decode)?;
113            let repo: String = row.try_get("repo").map_err(decode)?;
114            let agent_id: String = row.try_get("agent_id").map_err(decode)?;
115            let kind_json: String = row.try_get("kind").map_err(decode)?;
116            let actor: Option<String> = row.try_get("actor").map_err(decode)?;
117            let on_behalf_of: Option<String> = row.try_get("on_behalf_of").map_err(decode)?;
118            events.push(FleetEvent {
119                seq: seq as u64,
120                ts,
121                repo,
122                agent_id,
123                kind: serde_json::from_str(&kind_json)?,
124                actor,
125                on_behalf_of,
126            });
127        }
128        Ok(events)
129    }
130
131    /// The cursor is `global_ordinal` — the `BIGSERIAL` primary key, i.e.
132    /// durable insertion order (#219).
133    async fn replay_fleet(
134        &self,
135        after_cursor: u64,
136        limit: usize,
137    ) -> Result<Vec<crate::store::CursoredEvent>> {
138        let rows = sqlx::query(
139            "SELECT global_ordinal, seq, ts, repo, agent_id, kind, actor, on_behalf_of FROM events \
140             WHERE global_ordinal > $1 ORDER BY global_ordinal LIMIT $2",
141        )
142        .bind(after_cursor as i64)
143        .bind(limit as i64)
144        .fetch_all(&self.pool)
145        .await
146        .map_err(|e| CoreError::Store(format!("replay_fleet: {e}")))?;
147
148        let mut out = Vec::with_capacity(rows.len());
149        for row in rows {
150            let decode = |e: sqlx::Error| CoreError::Store(format!("replay_fleet decode: {e}"));
151            let cursor: i64 = row.try_get("global_ordinal").map_err(decode)?;
152            let seq: i64 = row.try_get("seq").map_err(decode)?;
153            let ts: String = row.try_get("ts").map_err(decode)?;
154            let repo: String = row.try_get("repo").map_err(decode)?;
155            let agent_id: String = row.try_get("agent_id").map_err(decode)?;
156            let kind_json: String = row.try_get("kind").map_err(decode)?;
157            let actor: Option<String> = row.try_get("actor").map_err(decode)?;
158            let on_behalf_of: Option<String> = row.try_get("on_behalf_of").map_err(decode)?;
159            out.push(crate::store::CursoredEvent {
160                cursor: cursor as u64,
161                event: FleetEvent {
162                    seq: seq as u64,
163                    ts,
164                    repo,
165                    agent_id,
166                    kind: serde_json::from_str(&kind_json)?,
167                    actor,
168                    on_behalf_of,
169                },
170            });
171        }
172        Ok(out)
173    }
174
175    async fn latest_fleet_cursor(&self) -> Result<u64> {
176        let row = sqlx::query("SELECT COALESCE(MAX(global_ordinal), 0) AS head FROM events")
177            .fetch_one(&self.pool)
178            .await
179            .map_err(|e| CoreError::Store(format!("latest_fleet_cursor: {e}")))?;
180        let head: i64 = row
181            .try_get("head")
182            .map_err(|e| CoreError::Store(format!("latest_fleet_cursor decode: {e}")))?;
183        Ok(head as u64)
184    }
185
186    async fn high_water(&self, stream_key: &str) -> Result<u64> {
187        let row =
188            sqlx::query("SELECT COALESCE(MAX(seq), 0) AS hw FROM events WHERE stream_key = $1")
189                .bind(stream_key)
190                .fetch_one(&self.pool)
191                .await
192                .map_err(|e| CoreError::Store(format!("high_water: {e}")))?;
193        let hw: i64 = row
194            .try_get("hw")
195            .map_err(|e| CoreError::Store(format!("high_water decode: {e}")))?;
196        Ok(hw as u64)
197    }
198
199    async fn writable(&self) -> bool {
200        // Non-destructive write probe: insert a sentinel row inside a
201        // transaction we always roll back. Exercises the same write path as
202        // `append` (detecting a read-only / full store) without persisting
203        // anything and without DDL. `seq = -1` cannot collide with a real
204        // event (seq is u64) and the rollback ensures it never lands.
205        let Ok(mut tx) = self.pool.begin().await else {
206            return false;
207        };
208        let ok = sqlx::query(
209            "INSERT INTO events (stream_key, seq, ts, repo, agent_id, kind) \
210             VALUES ('__writable_probe__', -1, '', '', '', 'null')",
211        )
212        .execute(&mut *tx)
213        .await
214        .is_ok();
215        let _ = tx.rollback().await;
216        ok
217    }
218
219    async fn prune(&self, before_ts: &str) -> Result<u64> {
220        let res = sqlx::query("DELETE FROM events WHERE ts < $1")
221            .bind(before_ts)
222            .execute(&self.pool)
223            .await
224            .map_err(|e| CoreError::Store(format!("prune: {e}")))?;
225        Ok(res.rows_affected())
226    }
227
228    /// The sqlite query's Postgres twin: same grouping, same filter, same shape.
229    /// `kind` is a TEXT column (not `jsonb`), so each read casts before
230    /// extracting. The explicit `::bigint` casts on the counts are load-bearing —
231    /// Postgres widens `SUM` over an integer to `numeric`, which sqlx will not
232    /// decode as `i64`.
233    async fn usage(&self, since: &str, until: &str) -> Result<Vec<crate::store::UsageRow>> {
234        let rows = sqlx::query(
235            "SELECT repo AS workspace, substr(ts, 1, 10) AS day, \
236                COALESCE(SUM(CASE WHEN kind::jsonb->>'kind' = 'agent_finished' \
237                    THEN (kind::jsonb->>'cost_usd')::double precision END), 0.0)::double precision \
238                    AS cost_usd, \
239                COALESCE(SUM(CASE WHEN kind::jsonb->>'kind' = 'agent_finished' \
240                    THEN (kind::jsonb->>'turns')::bigint END), 0)::bigint AS turns, \
241                COALESCE(SUM(CASE WHEN kind::jsonb->>'to' = 'done' THEN 1 END), 0)::bigint AS done, \
242                COALESCE(SUM(CASE WHEN kind::jsonb->>'to' = 'failed' THEN 1 END), 0)::bigint \
243                    AS failed, \
244                COALESCE(SUM(CASE WHEN kind::jsonb->>'to' = 'killed' THEN 1 END), 0)::bigint \
245                    AS killed, \
246                COALESCE(SUM(CASE WHEN kind::jsonb->>'to' = 'crashed' THEN 1 END), 0)::bigint \
247                    AS crashed, \
248                COALESCE(SUM(CASE WHEN kind::jsonb->>'kind' = 'agent_timed_out' THEN 1 END), \
249                    0)::bigint AS timed_out \
250             FROM events \
251             WHERE ts >= $1 AND ts < $2 AND ( \
252                kind::jsonb->>'kind' = 'agent_finished' OR ( \
253                    kind::jsonb->>'kind' = 'status_changed' \
254                    AND kind::jsonb->>'to' IN ('done', 'failed', 'killed', 'crashed')) \
255                OR kind::jsonb->>'kind' = 'agent_timed_out') \
256             GROUP BY repo, substr(ts, 1, 10) \
257             ORDER BY repo, day",
258        )
259        .bind(since)
260        .bind(until)
261        .fetch_all(&self.pool)
262        .await
263        .map_err(|e| CoreError::Store(format!("usage: {e}")))?;
264
265        let decode = |e: sqlx::Error| CoreError::Store(format!("usage decode: {e}"));
266        let mut out = Vec::with_capacity(rows.len());
267        for row in rows {
268            out.push(crate::store::UsageRow {
269                workspace: row.try_get("workspace").map_err(decode)?,
270                day: row.try_get("day").map_err(decode)?,
271                cost_usd: row.try_get::<f64, _>("cost_usd").map_err(decode)?,
272                turns: row.try_get::<i64, _>("turns").map_err(decode)? as u64,
273                done: row.try_get::<i64, _>("done").map_err(decode)? as u64,
274                failed: row.try_get::<i64, _>("failed").map_err(decode)? as u64,
275                killed: row.try_get::<i64, _>("killed").map_err(decode)? as u64,
276                crashed: row.try_get::<i64, _>("crashed").map_err(decode)? as u64,
277                timed_out: row.try_get::<i64, _>("timed_out").map_err(decode)? as u64,
278            });
279        }
280        Ok(out)
281    }
282}
283
284#[cfg(test)]
285mod tests {
286    use super::*;
287
288    async fn connect_or_skip() -> Option<PostgresStore> {
289        let url = std::env::var("DATABASE_URL").ok()?;
290        let store = PostgresStore::connect(&url).await.unwrap();
291        store.reset_for_tests().await.unwrap();
292        Some(store)
293    }
294
295    #[tokio::test]
296    async fn postgres_store_satisfies_conformance() {
297        let Some(store) = connect_or_skip().await else {
298            eprintln!("SKIP postgres_store_satisfies_conformance: DATABASE_URL unset");
299            return;
300        };
301        crate::testkit::store_conformance(&store).await;
302        store.reset_for_tests().await.unwrap();
303        crate::testkit::store_prune_conformance(&store).await;
304        store.reset_for_tests().await.unwrap();
305        crate::testkit::store_usage_conformance(&store).await;
306        store.reset_for_tests().await.unwrap();
307        crate::testkit::store_fleet_replay_conformance(&store).await;
308    }
309}
prospero_core/postgres_store.rs

prospero_core/
postgres_store.rs
