Skip to main content

prospero_core/
postgres_config_store.rs

1//! sqlx-backed Postgres [`ConfigStore`] — the clustered-tier config backend.
2//!
3//! Mirrors [`crate::config_store::SqliteConfigStore`] with Postgres dialect.
4//! Runs the same conformance battery, gated on `DATABASE_URL`.
5
6use async_trait::async_trait;
7use sqlx::Row;
8use sqlx::postgres::PgPool;
9
10use crate::automation::{AutomationRun, RUN_HISTORY_LIMIT, StoredAutomation};
11use crate::config_store::ConfigStore;
12use crate::error::{CoreError, Result};
13use crate::registry::RegisteredWorkspace;
14
15const SCHEMA: &str = "CREATE TABLE IF NOT EXISTS repos (\
16    name   TEXT PRIMARY KEY,\
17    root   TEXT NOT NULL,\
18    config TEXT NOT NULL\
19)";
20
21/// Mirrors the sqlite schema; see `config_store::AUTOMATIONS_SCHEMA` for why
22/// `last_fired_at` is a column rather than a field inside `spec`.
23const AUTOMATIONS_SCHEMA: &str = "CREATE TABLE IF NOT EXISTS automations (\
24    id             TEXT PRIMARY KEY,\
25    spec           TEXT NOT NULL,\
26    webhook_secret TEXT,\
27    last_fired_at  TEXT\
28)";
29
30const RUNS_SCHEMA: &str = "CREATE TABLE IF NOT EXISTS automation_runs (\
31    automation_id TEXT NOT NULL,\
32    fired_at      TEXT NOT NULL,\
33    source        TEXT NOT NULL,\
34    agent_id      TEXT,\
35    error         TEXT\
36)";
37
38const RUNS_INDEX: &str = "CREATE INDEX IF NOT EXISTS automation_runs_by_time ON automation_runs (automation_id, fired_at DESC)";
39
40/// sqlx/Postgres-backed config store (clustered tier).
41pub struct PostgresConfigStore {
42    pool: PgPool,
43}
44
45impl PostgresConfigStore {
46    /// Connect to Postgres at `url` and ensure the schema exists.
47    pub async fn connect(url: &str) -> Result<Self> {
48        let pool = crate::pg::connect(url).await?;
49        crate::pg::ensure_schema(&pool, SCHEMA, "repos table").await?;
50        crate::pg::ensure_schema(&pool, AUTOMATIONS_SCHEMA, "automations table").await?;
51        crate::pg::ensure_schema(&pool, RUNS_SCHEMA, "automation_runs table").await?;
52        crate::pg::ensure_schema(&pool, RUNS_INDEX, "automation_runs index").await?;
53        Ok(Self { pool })
54    }
55
56    /// Truncate all repos. Test-only.
57    #[cfg(any(test, feature = "testkit"))]
58    pub async fn reset_for_tests(&self) -> Result<()> {
59        sqlx::query("TRUNCATE repos, automations, automation_runs")
60            .execute(&self.pool)
61            .await
62            .map_err(|e| CoreError::Store(format!("reset: {e}")))?;
63        Ok(())
64    }
65}
66
67#[async_trait]
68impl ConfigStore for PostgresConfigStore {
69    async fn list_repos(&self) -> Result<Vec<RegisteredWorkspace>> {
70        let rows = sqlx::query("SELECT name, root, config FROM repos ORDER BY name")
71            .fetch_all(&self.pool)
72            .await
73            .map_err(|e| CoreError::Store(format!("list_repos: {e}")))?;
74        let mut repos = Vec::with_capacity(rows.len());
75        for row in rows {
76            let decode = |e: sqlx::Error| CoreError::Store(format!("list_repos decode: {e}"));
77            let name: String = row.try_get("name").map_err(decode)?;
78            let root: String = row.try_get("root").map_err(decode)?;
79            let config_json: String = row.try_get("config").map_err(decode)?;
80            repos.push(RegisteredWorkspace {
81                name,
82                root: root.into(),
83                config: serde_json::from_str(&config_json)?,
84            });
85        }
86        Ok(repos)
87    }
88
89    async fn upsert_repo(&self, repo: &RegisteredWorkspace) -> Result<()> {
90        let config = serde_json::to_string(&repo.config)?;
91        let root = repo
92            .root
93            .to_str()
94            .ok_or_else(|| CoreError::Store(format!("non-UTF8 repo root path: {:?}", repo.root)))?;
95        sqlx::query(
96            "INSERT INTO repos (name, root, config) VALUES ($1, $2, $3) \
97             ON CONFLICT (name) DO UPDATE SET root = excluded.root, config = excluded.config",
98        )
99        .bind(&repo.name)
100        .bind(root)
101        .bind(config)
102        .execute(&self.pool)
103        .await
104        .map_err(|e| CoreError::Store(format!("upsert_repo: {e}")))?;
105        Ok(())
106    }
107
108    async fn delete_repo(&self, name: &str) -> Result<bool> {
109        let res = sqlx::query("DELETE FROM repos WHERE name = $1")
110            .bind(name)
111            .execute(&self.pool)
112            .await
113            .map_err(|e| CoreError::Store(format!("delete_repo: {e}")))?;
114        Ok(res.rows_affected() > 0)
115    }
116
117    async fn list_automations(&self) -> Result<Vec<StoredAutomation>> {
118        let rows =
119            sqlx::query("SELECT spec, webhook_secret, last_fired_at FROM automations ORDER BY id")
120                .fetch_all(&self.pool)
121                .await
122                .map_err(|e| CoreError::Store(format!("list_automations: {e}")))?;
123        rows.iter()
124            .map(crate::config_store::decode_automation)
125            .collect()
126    }
127
128    async fn upsert_automation(&self, automation: &StoredAutomation) -> Result<()> {
129        let spec = crate::config_store::encode_spec(automation)?;
130        sqlx::query(
131            "INSERT INTO automations (id, spec, webhook_secret, last_fired_at) \
132             VALUES ($1, $2, $3, NULL) \
133             ON CONFLICT (id) DO UPDATE SET spec = excluded.spec, \
134             webhook_secret = excluded.webhook_secret",
135        )
136        .bind(&automation.automation.id)
137        .bind(spec)
138        .bind(&automation.webhook_secret)
139        .execute(&self.pool)
140        .await
141        .map_err(|e| CoreError::Store(format!("upsert_automation: {e}")))?;
142        Ok(())
143    }
144
145    async fn delete_automation(&self, id: &str) -> Result<bool> {
146        sqlx::query("DELETE FROM automation_runs WHERE automation_id = $1")
147            .bind(id)
148            .execute(&self.pool)
149            .await
150            .map_err(|e| CoreError::Store(format!("delete_automation runs: {e}")))?;
151        let res = sqlx::query("DELETE FROM automations WHERE id = $1")
152            .bind(id)
153            .execute(&self.pool)
154            .await
155            .map_err(|e| CoreError::Store(format!("delete_automation: {e}")))?;
156        Ok(res.rows_affected() > 0)
157    }
158
159    async fn claim_automation_fire(&self, id: &str, fire_at: &str) -> Result<bool> {
160        let res = sqlx::query(
161            "UPDATE automations SET last_fired_at = $1 \
162             WHERE id = $2 AND (last_fired_at IS NULL OR last_fired_at < $1)",
163        )
164        .bind(fire_at)
165        .bind(id)
166        .execute(&self.pool)
167        .await
168        .map_err(|e| CoreError::Store(format!("claim_automation_fire: {e}")))?;
169        Ok(res.rows_affected() == 1)
170    }
171
172    async fn record_run(&self, run: &AutomationRun) -> Result<()> {
173        sqlx::query(
174            "INSERT INTO automation_runs (automation_id, fired_at, source, agent_id, error) \
175             VALUES ($1, $2, $3, $4, $5)",
176        )
177        .bind(&run.automation_id)
178        .bind(&run.fired_at)
179        .bind(crate::config_store::encode_source(run.source)?)
180        .bind(&run.agent_id)
181        .bind(&run.error)
182        .execute(&self.pool)
183        .await
184        .map_err(|e| CoreError::Store(format!("record_run: {e}")))?;
185
186        // `ctid` is Postgres' physical row identifier — the equivalent of
187        // sqlite's `rowid` here, and the reason the trim can address individual
188        // rows even when two runs share a `fired_at`.
189        sqlx::query(
190            "DELETE FROM automation_runs WHERE automation_id = $1 AND ctid NOT IN \
191             (SELECT ctid FROM automation_runs WHERE automation_id = $1 \
192              ORDER BY fired_at DESC LIMIT $2)",
193        )
194        .bind(&run.automation_id)
195        .bind(RUN_HISTORY_LIMIT as i64)
196        .execute(&self.pool)
197        .await
198        .map_err(|e| CoreError::Store(format!("trim runs: {e}")))?;
199        Ok(())
200    }
201
202    async fn list_runs(&self, automation_id: &str, limit: usize) -> Result<Vec<AutomationRun>> {
203        let rows = sqlx::query(
204            "SELECT automation_id, fired_at, source, agent_id, error FROM automation_runs \
205             WHERE automation_id = $1 ORDER BY fired_at DESC LIMIT $2",
206        )
207        .bind(automation_id)
208        .bind(limit as i64)
209        .fetch_all(&self.pool)
210        .await
211        .map_err(|e| CoreError::Store(format!("list_runs: {e}")))?;
212        rows.iter().map(crate::config_store::decode_run).collect()
213    }
214}
215
216#[cfg(test)]
217mod tests {
218    use super::*;
219
220    /// Serialize the Postgres-gated config-store tests against each other.
221    /// They share one database and each one starts by TRUNCATE-ing it, so run
222    /// in parallel a sibling's reset deletes the rows this test just wrote —
223    /// and, worse, a sibling claiming the same tick makes the exactly-once
224    /// assertion fail for a reason that has nothing to do with the code under
225    /// test. Distinct ids would not help: the reset is what collides. Held
226    /// across awaits, so it must be a `tokio` mutex; taken right after the
227    /// `DATABASE_URL` guard, since an unset-DB skip never contends. Mirrors
228    /// `distributed_bus`'s `BUS_TEST_SERIAL`.
229    static CONFIG_TEST_SERIAL: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
230
231    #[tokio::test]
232    async fn postgres_config_store_satisfies_conformance() {
233        let Ok(url) = std::env::var("DATABASE_URL") else {
234            eprintln!("SKIP postgres_config_store_satisfies_conformance: DATABASE_URL unset");
235            return;
236        };
237        let _serial = CONFIG_TEST_SERIAL.lock().await;
238        let store = PostgresConfigStore::connect(&url).await.unwrap();
239        store.reset_for_tests().await.unwrap();
240        crate::testkit::config_store_conformance(&store).await;
241    }
242
243    #[tokio::test]
244    async fn postgres_config_store_satisfies_automation_conformance() {
245        let Ok(url) = std::env::var("DATABASE_URL") else {
246            eprintln!(
247                "SKIP postgres_config_store_satisfies_automation_conformance: DATABASE_URL unset"
248            );
249            return;
250        };
251        let _serial = CONFIG_TEST_SERIAL.lock().await;
252        let store = PostgresConfigStore::connect(&url).await.unwrap();
253        store.reset_for_tests().await.unwrap();
254        crate::testkit::automation_store_conformance(&store).await;
255    }
256
257    #[tokio::test]
258    async fn postgres_config_store_claims_a_tick_once() {
259        let Ok(url) = std::env::var("DATABASE_URL") else {
260            eprintln!("SKIP postgres_config_store_claims_a_tick_once: DATABASE_URL unset");
261            return;
262        };
263        let _serial = CONFIG_TEST_SERIAL.lock().await;
264        let store = PostgresConfigStore::connect(&url).await.unwrap();
265        store.reset_for_tests().await.unwrap();
266        crate::testkit::automation_claim_conformance(&store).await;
267    }
268}