alembic_engine/
state.rs

1//! local uid -> backend id state store.
2
3use crate::types::BackendId;
4use alembic_core::{TypeName, Uid};
5use anyhow::{anyhow, Context, Result};
6use serde::{Deserialize, Serialize};
7use std::collections::BTreeMap;
8use std::fs;
9use std::path::{Path, PathBuf};
10use std::sync::Arc;
11use tokio::sync::Mutex;
12use tokio_postgres::Client;
13
14/// on-disk state schema.
15#[derive(Debug, Default, Clone, Serialize, Deserialize)]
16pub struct StateData {
17    #[serde(default)]
18    pub mappings: BTreeMap<TypeName, BTreeMap<Uid, BackendId>>,
19}
20
21/// TLS configuration for postgres state backend connections.
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23pub enum PostgresTlsMode {
24    Disable,
25    Require,
26}
27
28/// trait for pluggable state backends.
29#[async_trait::async_trait]
30pub trait StateBackend: Send + Sync + std::fmt::Debug {
31    async fn load(&mut self) -> Result<StateData>;
32    async fn save(&mut self, data: &StateData) -> Result<()>;
33}
34
35/// state store wrapper with load/save helpers.
36#[derive(Debug, Clone)]
37pub struct StateStore {
38    backend: Option<Arc<Mutex<dyn StateBackend>>>,
39    data: StateData,
40    journal_dir: Option<PathBuf>,
41}
42
43impl StateStore {
44    /// create a new state store with an optional backend.
45    pub fn new(backend: Option<Arc<Mutex<dyn StateBackend>>>, data: StateData) -> Self {
46        Self {
47            backend,
48            data,
49            journal_dir: None,
50        }
51    }
52
53    /// set the directory used to persist apply journals for resumable applies.
54    pub fn with_journal_dir(mut self, dir: PathBuf) -> Self {
55        self.journal_dir = Some(dir);
56        self
57    }
58
59    /// directory used to persist apply journals, if configured.
60    pub fn journal_dir(&self) -> Option<&Path> {
61        self.journal_dir.as_deref()
62    }
63
64    /// load state from a file path.
65    pub fn load(path: impl AsRef<Path>) -> Result<Self> {
66        let path = path.as_ref().to_path_buf();
67        // always create a backend so we can save to the same path later.
68        // the backend's load() method handles missing files gracefully.
69        let backend: Option<Arc<Mutex<dyn StateBackend>>> =
70            Some(Arc::new(Mutex::new(LocalBackend { path: path.clone() }))
71                as Arc<Mutex<dyn StateBackend>>);
72        let data = if path.exists() {
73            let raw = fs::read_to_string(&path)
74                .with_context(|| format!("read state: {}", path.display()))?;
75            serde_json::from_str::<StateData>(&raw)
76                .with_context(|| format!("parse state: {}", path.display()))?
77        } else {
78            StateData::default()
79        };
80        Ok(Self::new(backend, data))
81    }
82
83    /// load state from a postgres backend.
84    pub async fn load_postgres(
85        url: impl Into<String>,
86        key: impl Into<String>,
87        tls_mode: PostgresTlsMode,
88    ) -> Result<Self> {
89        let mut postgres_backend = PostgresBackend {
90            url: url.into(),
91            key: key.into(),
92            tls_mode,
93            loaded_version: None,
94            table_ensured: false,
95        };
96        let data = postgres_backend.load().await?;
97        let backend: Arc<Mutex<dyn StateBackend>> = Arc::new(Mutex::new(postgres_backend));
98        Ok(Self::new(Some(backend), data))
99    }
100
101    /// load state from the configured backend.
102    pub async fn load_async(&mut self) -> Result<()> {
103        if let Some(backend) = &self.backend {
104            self.data = backend.lock().await.load().await?;
105        }
106        Ok(())
107    }
108
109    /// persist state to the configured backend.
110    pub async fn save_async(&self) -> Result<()> {
111        if let Some(backend) = &self.backend {
112            backend.lock().await.save(&self.data).await?;
113        }
114        Ok(())
115    }
116
117    /// lookup a backend id by type + uid.
118    pub fn backend_id(&self, type_name: TypeName, uid: Uid) -> Option<BackendId> {
119        self.data
120            .mappings
121            .get(&type_name)
122            .and_then(|map| map.get(&uid).cloned())
123    }
124
125    /// set a backend id mapping.
126    pub fn set_backend_id(&mut self, type_name: TypeName, uid: Uid, backend_id: BackendId) {
127        self.data
128            .mappings
129            .entry(type_name)
130            .or_default()
131            .insert(uid, backend_id);
132    }
133
134    /// remove a backend id mapping.
135    pub fn remove_backend_id(&mut self, type_name: TypeName, uid: Uid) {
136        if let Some(type_map) = self.data.mappings.get_mut(&type_name) {
137            type_map.remove(&uid);
138        }
139    }
140
141    /// return all mappings for external use.
142    pub fn all_mappings(&self) -> &BTreeMap<TypeName, BTreeMap<Uid, BackendId>> {
143        &self.data.mappings
144    }
145}
146
147impl From<&StateStore> for StateData {
148    fn from(store: &StateStore) -> Self {
149        Self {
150            mappings: store.all_mappings().clone(),
151        }
152    }
153}
154
155#[derive(Debug)]
156struct LocalBackend {
157    path: PathBuf,
158}
159
160#[async_trait::async_trait]
161impl StateBackend for LocalBackend {
162    async fn load(&mut self) -> Result<StateData> {
163        if self.path.exists() {
164            let raw = fs::read_to_string(&self.path)
165                .with_context(|| format!("read state: {}", self.path.display()))?;
166            let data = serde_json::from_str::<StateData>(&raw)
167                .with_context(|| format!("parse state: {}", self.path.display()))?;
168            Ok(data)
169        } else {
170            Ok(StateData::default())
171        }
172    }
173
174    async fn save(&mut self, data: &StateData) -> Result<()> {
175        if let Some(parent) = self.path.parent() {
176            fs::create_dir_all(parent)
177                .with_context(|| format!("create state dir: {}", parent.display()))?;
178        }
179        let raw = serde_json::to_string_pretty(data)?;
180        let tmp = self.path.with_extension("json.tmp");
181        fs::write(&tmp, &raw).with_context(|| format!("write state tmp: {}", tmp.display()))?;
182        fs::rename(&tmp, &self.path)
183            .with_context(|| format!("write state: {}", self.path.display()))?;
184        Ok(())
185    }
186}
187
188#[derive(Debug)]
189struct PostgresBackend {
190    url: String,
191    key: String,
192    tls_mode: PostgresTlsMode,
193    loaded_version: Option<i32>,
194    table_ensured: bool,
195}
196
197#[async_trait::async_trait]
198impl StateBackend for PostgresBackend {
199    async fn load(&mut self) -> Result<StateData> {
200        let client = self.connect().await?;
201
202        let row = client
203            .query_opt(
204                "SELECT payload::text, loaded_version FROM alembic_state WHERE state_key = $1",
205                &[&self.key],
206            )
207            .await
208            .with_context(|| "load postgres state payload")?;
209
210        let Some(row) = row else {
211            self.loaded_version = Some(0);
212            return Ok(StateData::default());
213        };
214
215        self.loaded_version = Some(row.try_get::<_, i32>("loaded_version")?);
216
217        let raw: String = row
218            .try_get(0)
219            .with_context(|| "decode postgres state payload")?;
220        serde_json::from_str::<StateData>(&raw).with_context(|| "parse postgres state payload")
221    }
222
223    async fn save(&mut self, data: &StateData) -> Result<()> {
224        let client = self.connect().await?;
225        let Some(loaded_version) = self.loaded_version else {
226            return Err(anyhow!("must load state before saving"));
227        };
228
229        let payload = serde_json::to_string(data)?;
230        let rows_modified = client
231            .execute(
232                "INSERT INTO alembic_state (state_key, payload, updated_at)
233                 VALUES ($1, CAST($2 AS TEXT)::jsonb, NOW())
234                 ON CONFLICT (state_key)
235                 DO UPDATE SET payload = EXCLUDED.payload, updated_at = NOW(), loaded_version = $3+1
236                 WHERE alembic_state.loaded_version = $3",
237                &[&self.key, &payload, &loaded_version],
238            )
239            .await
240            .with_context(|| "save postgres state payload")?;
241
242        if rows_modified == 0 {
243            return Err(anyhow!(
244                "failed to save postgres state for key '{}': optimistic lock failed (expected loaded_version={})",
245                self.key,
246                loaded_version
247            ));
248        }
249
250        self.loaded_version = Some(loaded_version + 1);
251
252        Ok(())
253    }
254}
255
256impl PostgresBackend {
257    async fn ensure_table(&mut self, client: &Client) -> Result<()> {
258        if self.table_ensured {
259            return Ok(());
260        }
261
262        client
263            .execute(
264                "CREATE TABLE IF NOT EXISTS alembic_state (\
265                state_key TEXT PRIMARY KEY, \
266                payload JSONB NOT NULL, \
267                updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), \
268                loaded_version INTEGER NOT NULL DEFAULT 1\
269                )",
270                &[],
271            )
272            .await
273            .with_context(|| "ensure postgres alembic_state table exists")?;
274
275        // if the table already existed, CREATE TABLE IF NOT EXISTS won't add new columns.
276        client
277            .execute(
278                "ALTER TABLE alembic_state ADD COLUMN IF NOT EXISTS loaded_version INTEGER NOT NULL DEFAULT 1",
279                &[],
280            )
281            .await
282            .with_context(|| "ensure postgres alembic_state.loaded_version column exists")?;
283
284        self.table_ensured = true;
285        Ok(())
286    }
287
288    async fn connect(&mut self) -> Result<tokio_postgres::Client> {
289        match self.tls_mode {
290            PostgresTlsMode::Disable => {
291                let (client, connection) =
292                    tokio_postgres::connect(&self.url, tokio_postgres::NoTls)
293                        .await
294                        .with_context(|| "connect postgres state backend")?;
295                tokio::spawn(async move {
296                    if let Err(err) = connection.await {
297                        tracing::warn!("postgres state backend connection error: {err}");
298                    }
299                });
300                self.ensure_table(&client).await?;
301                Ok(client)
302            }
303            PostgresTlsMode::Require => {
304                let connector = native_tls::TlsConnector::builder()
305                    .build()
306                    .with_context(|| "build postgres TLS connector")?;
307                let connector = postgres_native_tls::MakeTlsConnector::new(connector);
308                let (client, connection) = tokio_postgres::connect(&self.url, connector)
309                    .await
310                    .with_context(|| "connect postgres state backend")?;
311                tokio::spawn(async move {
312                    if let Err(err) = connection.await {
313                        tracing::warn!("postgres state backend connection error: {err}");
314                    }
315                });
316                self.ensure_table(&client).await?;
317                Ok(client)
318            }
319        }
320    }
321}
322
323#[cfg(test)]
324mod tests {
325    use super::*;
326    use tempfile::TempDir;
327
328    fn t(s: &str) -> TypeName {
329        TypeName::new(s)
330    }
331
332    fn uid(n: u128) -> Uid {
333        Uid::from_u128(n)
334    }
335
336    #[test]
337    fn backend_id_returns_none_for_missing_type() {
338        let store = StateStore::new(None, StateData::default());
339        assert_eq!(store.backend_id(t("site"), uid(1)), None);
340    }
341
342    #[test]
343    fn backend_id_returns_none_for_missing_uid() {
344        let mut data = StateData::default();
345        data.mappings
346            .entry(t("site"))
347            .or_default()
348            .insert(uid(1), BackendId::Int(42));
349        let store = StateStore::new(None, data);
350        assert_eq!(store.backend_id(t("site"), uid(2)), None);
351    }
352
353    #[test]
354    fn backend_id_returns_value_for_existing_mapping() {
355        let mut data = StateData::default();
356        data.mappings
357            .entry(t("site"))
358            .or_default()
359            .insert(uid(1), BackendId::Int(42));
360        let store = StateStore::new(None, data);
361        assert_eq!(
362            store.backend_id(t("site"), uid(1)),
363            Some(BackendId::Int(42))
364        );
365    }
366
367    #[test]
368    fn set_backend_id_creates_mapping() {
369        let mut store = StateStore::new(None, StateData::default());
370        store.set_backend_id(t("site"), uid(1), BackendId::Int(42));
371        assert_eq!(
372            store.backend_id(t("site"), uid(1)),
373            Some(BackendId::Int(42))
374        );
375    }
376
377    #[test]
378    fn set_backend_id_overwrites_existing() {
379        let mut data = StateData::default();
380        data.mappings
381            .entry(t("site"))
382            .or_default()
383            .insert(uid(1), BackendId::Int(42));
384        let mut store = StateStore::new(None, data);
385        store.set_backend_id(t("site"), uid(1), BackendId::Int(99));
386        assert_eq!(
387            store.backend_id(t("site"), uid(1)),
388            Some(BackendId::Int(99))
389        );
390    }
391
392    #[test]
393    fn remove_backend_id_removes_mapping() {
394        let mut data = StateData::default();
395        data.mappings
396            .entry(t("site"))
397            .or_default()
398            .insert(uid(1), BackendId::Int(42));
399        let mut store = StateStore::new(None, data);
400        store.remove_backend_id(t("site"), uid(1));
401        assert_eq!(store.backend_id(t("site"), uid(1)), None);
402    }
403
404    #[test]
405    fn remove_backend_id_noop_for_missing() {
406        let mut store = StateStore::new(None, StateData::default());
407        store.remove_backend_id(t("site"), uid(1));
408        // should not panic
409    }
410
411    #[test]
412    fn journal_dir_defaults_to_none_and_round_trips() {
413        let store = StateStore::new(None, StateData::default());
414        assert_eq!(store.journal_dir(), None);
415
416        let dir = PathBuf::from("/tmp/alembic-journal");
417        let store = store.with_journal_dir(dir.clone());
418        assert_eq!(store.journal_dir(), Some(dir.as_path()));
419    }
420
421    #[tokio::test]
422    async fn local_backend_load_missing_returns_empty() {
423        let dir = TempDir::new().unwrap();
424        let mut backend = LocalBackend {
425            path: dir.path().join("nope.json"),
426        };
427        let data = backend.load().await.unwrap();
428        assert!(data.mappings.is_empty());
429    }
430
431    #[tokio::test]
432    async fn local_backend_save_load_round_trip() {
433        let dir = TempDir::new().unwrap();
434        let path = dir.path().join("sub").join("state.json");
435        let mut backend = LocalBackend { path: path.clone() };
436
437        let mut data = StateData::default();
438        data.mappings
439            .entry(t("site"))
440            .or_default()
441            .insert(uid(10), BackendId::String("site-001".into()));
442
443        backend.save(&data).await.unwrap();
444        assert!(path.exists());
445
446        let loaded = backend.load().await.unwrap();
447        assert_eq!(
448            loaded.mappings[&t("site")][&uid(10)],
449            BackendId::String("site-001".into())
450        );
451    }
452
453    #[tokio::test]
454    async fn store_save_without_backend_is_noop() {
455        let store = StateStore::new(None, StateData::default());
456        store.save_async().await.unwrap();
457    }
458
459    #[tokio::test]
460    async fn store_load_async_without_backend_is_noop() {
461        let mut store = StateStore::new(None, StateData::default());
462        store.set_backend_id(t("x"), uid(1), BackendId::Int(1));
463        store.load_async().await.unwrap();
464        assert_eq!(store.backend_id(t("x"), uid(1)), Some(BackendId::Int(1)));
465    }
466}
alembic_engine/state.rs

alembic_engine/
state.rs
