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    UNIQUE(stream_key, seq)\
24)";
25
26/// sqlx/Postgres-backed durable event store (clustered tier).
27pub struct PostgresStore {
28    pool: PgPool,
29}
30
31impl PostgresStore {
32    /// Connect to Postgres at `url` and ensure the schema exists.
33    pub async fn connect(url: &str) -> Result<Self> {
34        let pool = crate::pg::connect(url).await?;
35        crate::pg::ensure_schema(&pool, SCHEMA, "events table").await?;
36        Ok(Self { pool })
37    }
38
39    /// Truncate all events. Test-only (resets between conformance batteries).
40    #[cfg(any(test, feature = "testkit"))]
41    pub async fn reset_for_tests(&self) -> Result<()> {
42        sqlx::query("TRUNCATE events")
43            .execute(&self.pool)
44            .await
45            .map_err(|e| CoreError::Store(format!("reset: {e}")))?;
46        Ok(())
47    }
48}
49
50#[async_trait]
51impl Store for PostgresStore {
52    async fn append(&self, event: &FleetEvent) -> Result<()> {
53        let kind = serde_json::to_string(&event.kind)?;
54        sqlx::query(
55            "INSERT INTO events (stream_key, seq, ts, repo, agent_id, kind) \
56             VALUES ($1, $2, $3, $4, $5, $6)",
57        )
58        .bind(event.stream_key())
59        .bind(event.seq as i64)
60        .bind(&event.ts)
61        .bind(&event.repo)
62        .bind(&event.agent_id)
63        .bind(kind)
64        .execute(&self.pool)
65        .await
66        .map_err(map_append_error)?;
67        Ok(())
68    }
69
70    async fn replay(&self, stream_key: &str, from_seq: u64) -> Result<Vec<FleetEvent>> {
71        let rows = sqlx::query(
72            "SELECT seq, ts, repo, agent_id, kind FROM events \
73             WHERE stream_key = $1 AND seq >= $2 ORDER BY seq",
74        )
75        .bind(stream_key)
76        .bind(from_seq as i64)
77        .fetch_all(&self.pool)
78        .await
79        .map_err(|e| CoreError::Store(format!("replay: {e}")))?;
80
81        let mut events = Vec::with_capacity(rows.len());
82        for row in rows {
83            let decode = |e: sqlx::Error| CoreError::Store(format!("replay decode: {e}"));
84            let seq: i64 = row.try_get("seq").map_err(decode)?;
85            let ts: String = row.try_get("ts").map_err(decode)?;
86            let repo: String = row.try_get("repo").map_err(decode)?;
87            let agent_id: String = row.try_get("agent_id").map_err(decode)?;
88            let kind_json: String = row.try_get("kind").map_err(decode)?;
89            events.push(FleetEvent {
90                seq: seq as u64,
91                ts,
92                repo,
93                agent_id,
94                kind: serde_json::from_str(&kind_json)?,
95            });
96        }
97        Ok(events)
98    }
99
100    async fn high_water(&self, stream_key: &str) -> Result<u64> {
101        let row =
102            sqlx::query("SELECT COALESCE(MAX(seq), 0) AS hw FROM events WHERE stream_key = $1")
103                .bind(stream_key)
104                .fetch_one(&self.pool)
105                .await
106                .map_err(|e| CoreError::Store(format!("high_water: {e}")))?;
107        let hw: i64 = row
108            .try_get("hw")
109            .map_err(|e| CoreError::Store(format!("high_water decode: {e}")))?;
110        Ok(hw as u64)
111    }
112
113    async fn writable(&self) -> bool {
114        // Non-destructive write probe: insert a sentinel row inside a
115        // transaction we always roll back. Exercises the same write path as
116        // `append` (detecting a read-only / full store) without persisting
117        // anything and without DDL. `seq = -1` cannot collide with a real
118        // event (seq is u64) and the rollback ensures it never lands.
119        let Ok(mut tx) = self.pool.begin().await else {
120            return false;
121        };
122        let ok = sqlx::query(
123            "INSERT INTO events (stream_key, seq, ts, repo, agent_id, kind) \
124             VALUES ('__writable_probe__', -1, '', '', '', 'null')",
125        )
126        .execute(&mut *tx)
127        .await
128        .is_ok();
129        let _ = tx.rollback().await;
130        ok
131    }
132
133    async fn prune(&self, before_ts: &str) -> Result<u64> {
134        let res = sqlx::query("DELETE FROM events WHERE ts < $1")
135            .bind(before_ts)
136            .execute(&self.pool)
137            .await
138            .map_err(|e| CoreError::Store(format!("prune: {e}")))?;
139        Ok(res.rows_affected())
140    }
141
142    /// The sqlite query's Postgres twin: same grouping, same filter, same shape.
143    /// `kind` is a TEXT column (not `jsonb`), so each read casts before
144    /// extracting. The explicit `::bigint` casts on the counts are load-bearing —
145    /// Postgres widens `SUM` over an integer to `numeric`, which sqlx will not
146    /// decode as `i64`.
147    async fn usage(&self, since: &str, until: &str) -> Result<Vec<crate::store::UsageRow>> {
148        let rows = sqlx::query(
149            "SELECT repo AS workspace, substr(ts, 1, 10) AS day, \
150                COALESCE(SUM(CASE WHEN kind::jsonb->>'kind' = 'agent_finished' \
151                    THEN (kind::jsonb->>'cost_usd')::double precision END), 0.0)::double precision \
152                    AS cost_usd, \
153                COALESCE(SUM(CASE WHEN kind::jsonb->>'kind' = 'agent_finished' \
154                    THEN (kind::jsonb->>'turns')::bigint END), 0)::bigint AS turns, \
155                COALESCE(SUM(CASE WHEN kind::jsonb->>'to' = 'done' THEN 1 END), 0)::bigint AS done, \
156                COALESCE(SUM(CASE WHEN kind::jsonb->>'to' = 'failed' THEN 1 END), 0)::bigint \
157                    AS failed, \
158                COALESCE(SUM(CASE WHEN kind::jsonb->>'to' = 'killed' THEN 1 END), 0)::bigint \
159                    AS killed, \
160                COALESCE(SUM(CASE WHEN kind::jsonb->>'to' = 'crashed' THEN 1 END), 0)::bigint \
161                    AS crashed \
162             FROM events \
163             WHERE ts >= $1 AND ts < $2 AND ( \
164                kind::jsonb->>'kind' = 'agent_finished' OR ( \
165                    kind::jsonb->>'kind' = 'status_changed' \
166                    AND kind::jsonb->>'to' IN ('done', 'failed', 'killed', 'crashed'))) \
167             GROUP BY repo, substr(ts, 1, 10) \
168             ORDER BY repo, day",
169        )
170        .bind(since)
171        .bind(until)
172        .fetch_all(&self.pool)
173        .await
174        .map_err(|e| CoreError::Store(format!("usage: {e}")))?;
175
176        let decode = |e: sqlx::Error| CoreError::Store(format!("usage decode: {e}"));
177        let mut out = Vec::with_capacity(rows.len());
178        for row in rows {
179            out.push(crate::store::UsageRow {
180                workspace: row.try_get("workspace").map_err(decode)?,
181                day: row.try_get("day").map_err(decode)?,
182                cost_usd: row.try_get::<f64, _>("cost_usd").map_err(decode)?,
183                turns: row.try_get::<i64, _>("turns").map_err(decode)? as u64,
184                done: row.try_get::<i64, _>("done").map_err(decode)? as u64,
185                failed: row.try_get::<i64, _>("failed").map_err(decode)? as u64,
186                killed: row.try_get::<i64, _>("killed").map_err(decode)? as u64,
187                crashed: row.try_get::<i64, _>("crashed").map_err(decode)? as u64,
188            });
189        }
190        Ok(out)
191    }
192}
193
194#[cfg(test)]
195mod tests {
196    use super::*;
197
198    async fn connect_or_skip() -> Option<PostgresStore> {
199        let url = std::env::var("DATABASE_URL").ok()?;
200        let store = PostgresStore::connect(&url).await.unwrap();
201        store.reset_for_tests().await.unwrap();
202        Some(store)
203    }
204
205    #[tokio::test]
206    async fn postgres_store_satisfies_conformance() {
207        let Some(store) = connect_or_skip().await else {
208            eprintln!("SKIP postgres_store_satisfies_conformance: DATABASE_URL unset");
209            return;
210        };
211        crate::testkit::store_conformance(&store).await;
212        store.reset_for_tests().await.unwrap();
213        crate::testkit::store_prune_conformance(&store).await;
214        store.reset_for_tests().await.unwrap();
215        crate::testkit::store_usage_conformance(&store).await;
216    }
217}
prospero_core/postgres_store.rs

prospero_core/
postgres_store.rs
