rivet/state/
cdc_snapshot_store.rs

1use crate::error::Result;
2
3use super::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        // ON CONFLICT DO UPDATE upsert (SQLite ≥ 3.24 + Postgres) — the codebase's
21        // canonical upsert, replacing the old SQLite `INSERT OR REPLACE` arm.
22        self.execute(
23            "INSERT INTO cdc_snapshot (export_name, table_name, run_id, completed_at)
24             VALUES (?1, ?2, ?3, ?4)
25             ON CONFLICT (export_name, table_name) DO UPDATE SET
26                 run_id       = excluded.run_id,
27                 completed_at = excluded.completed_at",
28            &[
29                export_name.into(),
30                table_name.into(),
31                run_id.into(),
32                now.into(),
33            ],
34        )?;
35        Ok(())
36    }
37
38    /// Whether `(export_name, table_name)`'s initial snapshot has completed per
39    /// the state DB — the authoritative, GCS-independent signal.
40    pub fn snapshot_done(&self, export_name: &str, table_name: &str) -> Result<bool> {
41        Ok(self
42            .query_opt(
43                "SELECT COUNT(*) FROM cdc_snapshot WHERE export_name = ?1 AND table_name = ?2",
44                &[export_name.into(), table_name.into()],
45                |r| r.i64(0),
46            )?
47            .unwrap_or(0)
48            > 0)
49    }
50}
51
52#[cfg(test)]
53mod tests {
54    use super::*;
55
56    #[test]
57    fn snapshot_done_is_false_until_marked_then_true() {
58        let s = StateStore::open_in_memory().unwrap();
59        assert!(!s.snapshot_done("customers", "customers").unwrap());
60        s.mark_snapshot_done("customers", "customers", "run_1")
61            .unwrap();
62        assert!(s.snapshot_done("customers", "customers").unwrap());
63        // Scoped per (export, table): a sibling table is still un-snapshotted.
64        assert!(!s.snapshot_done("customers", "orders").unwrap());
65    }
66
67    #[test]
68    fn mark_snapshot_done_is_idempotent() {
69        let s = StateStore::open_in_memory().unwrap();
70        s.mark_snapshot_done("e", "t", "run_1").unwrap();
71        s.mark_snapshot_done("e", "t", "run_2").unwrap(); // replay/re-record
72        assert!(s.snapshot_done("e", "t").unwrap());
73    }
74}
rivet/state/cdc_snapshot_store.rs

rivet/state/
cdc_snapshot_store.rs
