prospero_core/
postgres_config_store.rs1use 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
21const 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
40pub struct PostgresConfigStore {
42 pool: PgPool,
43}
44
45impl PostgresConfigStore {
46 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 #[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 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 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}