rivet/state/
cdc_snapshot_store.rs

1use crate::error::Result;
2
3use super::{StateConn, StateStore};
4
5impl StateStore {
6    /// Record that `cdc.initial: snapshot`'s backfill for `(export_name,
7    /// table_name)` completed, on run `run_id`. Idempotent (upsert on the
8    /// `(export_name, table_name)` key), so a retried run never double-inserts.
9    ///
10    /// This is the durable, cleanup-proof twin of the GCS `snapshot/_SUCCESS`
11    /// marker: once here, `cleanup_source: true` may wipe the bucket without the
12    /// next run mistaking the table for un-snapshotted and re-snapshotting it.
13    pub fn mark_snapshot_done(
14        &self,
15        export_name: &str,
16        table_name: &str,
17        run_id: &str,
18    ) -> Result<()> {
19        let now = chrono::Utc::now().to_rfc3339();
20        match &self.conn {
21            StateConn::Sqlite(c) => {
22                c.execute(
23                    "INSERT OR REPLACE INTO cdc_snapshot
24                       (export_name, table_name, run_id, completed_at)
25                     VALUES (?1, ?2, ?3, ?4)",
26                    rusqlite::params![export_name, table_name, run_id, now],
27                )?;
28            }
29            StateConn::Postgres(client) => {
30                let mut c = client.borrow_mut();
31                c.execute(
32                    "INSERT INTO cdc_snapshot
33                       (export_name, table_name, run_id, completed_at)
34                     VALUES ($1, $2, $3, $4)
35                     ON CONFLICT (export_name, table_name) DO UPDATE SET
36                         run_id       = excluded.run_id,
37                         completed_at = excluded.completed_at",
38                    &[&export_name, &table_name, &run_id, &now],
39                )?;
40            }
41        }
42        Ok(())
43    }
44
45    /// Whether `(export_name, table_name)`'s initial snapshot has completed per
46    /// the state DB — the authoritative, GCS-independent signal.
47    pub fn snapshot_done(&self, export_name: &str, table_name: &str) -> Result<bool> {
48        match &self.conn {
49            StateConn::Sqlite(c) => {
50                let n: i64 = c.query_row(
51                    "SELECT COUNT(*) FROM cdc_snapshot
52                     WHERE export_name = ?1 AND table_name = ?2",
53                    rusqlite::params![export_name, table_name],
54                    |r| r.get(0),
55                )?;
56                Ok(n > 0)
57            }
58            StateConn::Postgres(client) => {
59                let mut c = client.borrow_mut();
60                let row = c.query_one(
61                    "SELECT COUNT(*) FROM cdc_snapshot
62                     WHERE export_name = $1 AND table_name = $2",
63                    &[&export_name, &table_name],
64                )?;
65                Ok(row.get::<_, i64>(0) > 0)
66            }
67        }
68    }
69}
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74
75    #[test]
76    fn snapshot_done_is_false_until_marked_then_true() {
77        let s = StateStore::open_in_memory().unwrap();
78        assert!(!s.snapshot_done("customers", "customers").unwrap());
79        s.mark_snapshot_done("customers", "customers", "run_1")
80            .unwrap();
81        assert!(s.snapshot_done("customers", "customers").unwrap());
82        // Scoped per (export, table): a sibling table is still un-snapshotted.
83        assert!(!s.snapshot_done("customers", "orders").unwrap());
84    }
85
86    #[test]
87    fn mark_snapshot_done_is_idempotent() {
88        let s = StateStore::open_in_memory().unwrap();
89        s.mark_snapshot_done("e", "t", "run_1").unwrap();
90        s.mark_snapshot_done("e", "t", "run_2").unwrap(); // replay/re-record
91        assert!(s.snapshot_done("e", "t").unwrap());
92    }
93}
rivet/state/cdc_snapshot_store.rs

rivet/state/
cdc_snapshot_store.rs
