Skip to main content

rivet/state/
journal_store.rs

1use crate::error::Result;
2use crate::journal::RunJournal;
3
4use super::StateStore;
5
6impl StateStore {
7    /// Persist a completed `RunJournal` to the state DB.
8    ///
9    /// Called once per export run, after `RunCompleted` has been recorded.
10    /// Overwrites any existing row for the same `run_id` (idempotent on retry).
11    pub fn store_journal(&self, journal: &RunJournal) -> Result<()> {
12        let json = serde_json::to_string(journal)?;
13        let now = chrono::Utc::now().to_rfc3339();
14        self.execute(
15            "INSERT INTO run_journal (run_id, export_name, finished_at, journal_json)
16             VALUES (?1, ?2, ?3, ?4)
17             ON CONFLICT (run_id) DO UPDATE SET
18                 export_name  = excluded.export_name,
19                 finished_at  = excluded.finished_at,
20                 journal_json = excluded.journal_json",
21            &[
22                journal.run_id.as_str().into(),
23                journal.export_name.as_str().into(),
24                now.into(),
25                json.into(),
26            ],
27        )?;
28        Ok(())
29    }
30
31    /// Load a journal by `run_id`.  Returns `None` if the run is not found.
32    #[allow(dead_code)]
33    pub fn load_journal(&self, run_id: &str) -> Result<Option<RunJournal>> {
34        let json = self.query_opt(
35            "SELECT journal_json FROM run_journal WHERE run_id = ?1",
36            &[run_id.into()],
37            |r| r.text(0),
38        )?;
39        Ok(match json {
40            Some(s) => Some(serde_json::from_str(&s)?),
41            None => None,
42        })
43    }
44
45    /// Return the most recent `limit` journal entries for an export, newest first.
46    #[allow(dead_code)]
47    pub fn recent_journals(&self, export_name: &str, limit: usize) -> Result<Vec<RunJournal>> {
48        let jsons = self.query(
49            "SELECT journal_json FROM run_journal
50             WHERE export_name = ?1
51             ORDER BY finished_at DESC
52             LIMIT ?2",
53            &[export_name.into(), (limit as i64).into()],
54            |r| r.text(0),
55        )?;
56        jsons
57            .iter()
58            .map(|j| serde_json::from_str::<RunJournal>(j).map_err(Into::into))
59            .collect()
60    }
61}
62
63#[cfg(test)]
64mod tests {
65    use super::*;
66    use crate::journal::{RunEvent, RunJournal};
67
68    fn make_journal(run_id: &str, export: &str) -> RunJournal {
69        let mut j = RunJournal::new(run_id, export);
70        j.record(RunEvent::FileWritten {
71            file_name: "part0.parquet".into(),
72            rows: 1_000,
73            bytes: 65_536,
74            part_index: 0,
75        });
76        j.record(RunEvent::RunCompleted {
77            status: "success".into(),
78            error_message: None,
79            duration_ms: 420,
80        });
81        j
82    }
83
84    #[test]
85    fn store_and_load_roundtrip() {
86        let store = StateStore::open_in_memory().unwrap();
87        let j = make_journal("run_abc_001", "orders");
88        store.store_journal(&j).unwrap();
89
90        let loaded = store.load_journal("run_abc_001").unwrap().unwrap();
91        assert_eq!(loaded.run_id, "run_abc_001");
92        assert_eq!(loaded.export_name, "orders");
93        assert_eq!(loaded.entries.len(), 2);
94        assert!(matches!(
95            loaded.entries[0].event,
96            RunEvent::FileWritten { rows: 1_000, .. }
97        ));
98        assert!(matches!(
99            loaded.entries[1].event,
100            RunEvent::RunCompleted { ref status, .. } if status == "success"
101        ));
102    }
103
104    #[test]
105    fn load_missing_returns_none() {
106        let store = StateStore::open_in_memory().unwrap();
107        assert!(store.load_journal("nonexistent").unwrap().is_none());
108    }
109
110    #[test]
111    fn store_is_idempotent_on_same_run_id() {
112        let store = StateStore::open_in_memory().unwrap();
113        let j = make_journal("run_idem", "payments");
114        store.store_journal(&j).unwrap();
115        store.store_journal(&j).unwrap();
116
117        let loaded = store.load_journal("run_idem").unwrap().unwrap();
118        assert_eq!(loaded.entries.len(), 2);
119    }
120
121    #[test]
122    fn recent_journals_returns_newest_first() {
123        let store = StateStore::open_in_memory().unwrap();
124        for i in 1..=3_u32 {
125            std::thread::sleep(std::time::Duration::from_millis(2));
126            store
127                .store_journal(&make_journal(&format!("run_{i:03}"), "events"))
128                .unwrap();
129        }
130
131        let recent = store.recent_journals("events", 2).unwrap();
132        assert_eq!(recent.len(), 2);
133        assert_eq!(recent[0].run_id, "run_003");
134        assert_eq!(recent[1].run_id, "run_002");
135    }
136}