Skip to main content

improv_store/
lib.rs

1//! One model store over either backend: an embedded Mentat file (a path) or
2//! a Postgres-hosted model (`postgres://…?model=<name>`).
3//!
4//! The GUI, TUI, CLI and server open models through [`Store::open`]; every
5//! method delegates to `improv_storage_mentat::ModelStore` or
6//! `improv_storage_pg::PgModelStore`, which share one surface.
7#![cfg_attr(coverage_nightly, feature(coverage_attribute))]
8#![cfg_attr(coverage_nightly, allow(unused_features))]
9
10use improv_core_model::{MeasureId, Model};
11use improv_storage_mentat::ModelStore;
12use improv_storage_pg::PgModelStore;
13
14/// Either backend's error. Callers match [`StoreError::Conflict`] for the
15/// "reload and retry" path; everything else is reported as text.
16#[derive(Debug, thiserror::Error)]
17pub enum StoreError {
18    #[error(transparent)]
19    Mentat(improv_storage_mentat::StoreError),
20    #[error(transparent)]
21    Pg(improv_storage_pg::StoreError),
22    /// Another writer committed since this handle loaded (either backend).
23    #[error("conflict: {0}")]
24    Conflict(String),
25}
26
27impl From<improv_storage_mentat::StoreError> for StoreError {
28    fn from(e: improv_storage_mentat::StoreError) -> Self {
29        match e {
30            improv_storage_mentat::StoreError::Conflict(m) => StoreError::Conflict(m),
31            e => StoreError::Mentat(e),
32        }
33    }
34}
35
36impl From<improv_storage_pg::StoreError> for StoreError {
37    fn from(e: improv_storage_pg::StoreError) -> Self {
38        match e {
39            improv_storage_pg::StoreError::Conflict(m) => StoreError::Conflict(m),
40            e => StoreError::Pg(e),
41        }
42    }
43}
44
45pub type Result<T> = std::result::Result<T, StoreError>;
46
47/// A model store: Mentat file or Postgres model.
48pub enum Store {
49    Mentat(ModelStore),
50    Pg(PgModelStore),
51}
52
53/// True when `target` names a Postgres-hosted model rather than a file.
54pub fn is_postgres(target: &str) -> bool {
55    target.starts_with("postgres://") || target.starts_with("postgresql://")
56}
57
58macro_rules! delegate {
59    ($self:ident, $s:ident => $e:expr) => {
60        match $self {
61            Store::Mentat($s) => Ok($e?),
62            Store::Pg($s) => Ok($e?),
63        }
64    };
65}
66
67impl Store {
68    /// Open `target`: a `postgres://` / `postgresql://` URL with `?model=`
69    /// opens that Postgres model; anything else is a Mentat file path (`""`
70    /// is in-memory).
71    pub fn open(target: &str) -> Result<Self> {
72        Ok(if is_postgres(target) {
73            Store::Pg(PgModelStore::open(target)?)
74        } else {
75            Store::Mentat(ModelStore::open(target)?)
76        })
77    }
78
79    pub fn load_model(&mut self) -> Result<Model> {
80        delegate!(self, s => s.load_model())
81    }
82
83    pub fn load_partial(&mut self, measure_ids: &[MeasureId]) -> Result<Model> {
84        delegate!(self, s => s.load_partial(measure_ids))
85    }
86
87    pub fn save_model(&mut self, model: &Model) -> Result<()> {
88        delegate!(self, s => s.save_model(model))
89    }
90
91    pub fn reload_model(&mut self) -> Result<Model> {
92        delegate!(self, s => s.reload_model())
93    }
94
95    pub fn tx_high_water(&mut self) -> Result<i64> {
96        delegate!(self, s => s.tx_high_water())
97    }
98
99    pub fn loaded_tx(&self) -> Option<i64> {
100        match self {
101            Store::Mentat(s) => s.loaded_tx(),
102            Store::Pg(s) => s.loaded_tx(),
103        }
104    }
105
106    pub fn set_loaded_tx(&mut self, tx: Option<i64>) {
107        match self {
108            Store::Mentat(s) => s.set_loaded_tx(tx),
109            Store::Pg(s) => s.set_loaded_tx(tx),
110        }
111    }
112
113    /// Everything a later handle needs to save safely over what this one
114    /// loaded or saved: the optimistic-concurrency token and (Postgres) the
115    /// snapshot a save diffs against, so it writes only what changed.
116    /// Callers that reopen the store per save carry this across handles.
117    pub fn save_token(&self) -> SaveToken {
118        match self {
119            Store::Mentat(s) => SaveToken {
120                tx: s.loaded_tx(),
121                pg: None,
122            },
123            Store::Pg(s) => SaveToken {
124                tx: s.loaded_tx(),
125                pg: s.snapshot().cloned(),
126            },
127        }
128    }
129
130    /// Adopt a token from [`Store::save_token`] before saving.
131    pub fn set_save_token(&mut self, token: &SaveToken) {
132        match self {
133            Store::Mentat(s) => s.set_loaded_tx(token.tx),
134            Store::Pg(s) => {
135                s.set_loaded_tx(token.tx);
136                s.set_snapshot(token.pg.clone());
137            }
138        }
139    }
140}
141
142/// What a store handle last loaded or saved, carried to the next handle (see
143/// [`Store::save_token`]). `Default` is "nothing loaded": the save writes the
144/// whole model unchecked.
145#[derive(Debug, Clone, Default)]
146pub struct SaveToken {
147    tx: Option<i64>,
148    pg: Option<improv_storage_pg::PgSnapshot>,
149}
150
151impl SaveToken {
152    /// The optimistic-concurrency high-water this token carries.
153    pub fn tx(&self) -> Option<i64> {
154        self.tx
155    }
156}
157
158#[cfg(test)]
159#[cfg_attr(coverage_nightly, coverage(off))]
160mod tests {
161    use super::*;
162    use improv_core_model::{CategoryId, Coordinate, ItemId, Measure, MeasureKind, Name, Value};
163
164    fn model() -> Model {
165        let mut m = Model::new();
166        m.add_category(CategoryId(1), "Time");
167        m.add_item(ItemId(10), CategoryId(1), "Q1");
168        m.add_measure(Measure {
169            id: MeasureId(100),
170            name: Name("Price".into()),
171            value_type: improv_core_model::ValueType::Number,
172            categories: vec![CategoryId(1)],
173            kind: MeasureKind::Input,
174            description: None,
175        });
176        m.set_input(
177            MeasureId(100),
178            Coordinate::from_pairs([(CategoryId(1), ItemId(10))]),
179            Value::Number(3.0),
180        );
181        m
182    }
183
184    /// Every delegating method, plus Conflict mapping, on one open store.
185    fn exercise(target: &str) {
186        let m = model();
187        let mut a = Store::open(target).unwrap();
188        let mut b = Store::open(target).unwrap();
189        a.save_model(&m).unwrap();
190        assert_eq!(a.load_model().unwrap(), m);
191        assert_eq!(a.reload_model().unwrap(), m);
192        assert_eq!(a.load_partial(&[MeasureId(100)]).unwrap(), m);
193        assert_eq!(a.loaded_tx(), Some(a.tx_high_water().unwrap()));
194        b.set_loaded_tx(Some(-1));
195        let e = b.save_model(&m).unwrap_err();
196        assert!(matches!(e, StoreError::Conflict(_)), "{e}");
197        assert!(e.to_string().starts_with("conflict:"));
198    }
199
200    #[test]
201    fn a_path_opens_mentat() {
202        let dir = std::env::temp_dir().join(format!("improv_store_{}", std::process::id()));
203        std::fs::create_dir_all(&dir).unwrap();
204        let db = dir.join("m.db");
205        let db = db.to_str().unwrap();
206        assert!(matches!(Store::open(db).unwrap(), Store::Mentat(_)));
207        exercise(db);
208        // A non-conflict backend error keeps its backend.
209        let mut s = Store::open("").unwrap();
210        let mut bad = model();
211        bad.measures.get_mut(&MeasureId(100)).unwrap().categories =
212            (0..6000).map(CategoryId).collect();
213        let e = s.save_model(&bad).unwrap_err();
214        assert!(matches!(e, StoreError::Mentat(_)), "{e}");
215        std::fs::remove_dir_all(&dir).unwrap();
216    }
217
218    #[test]
219    fn postgres_urls_dispatch_to_pg_and_need_a_model() {
220        assert!(is_postgres("postgres://u@h/d") && is_postgres("postgresql://u@h/d"));
221        assert!(!is_postgres("/tmp/postgres.db") && !is_postgres(""));
222        for bad in ["postgres://u@localhost/d", "postgresql://u@localhost/d?x=1"] {
223            let e = Store::open(bad).err().unwrap();
224            assert!(matches!(e, StoreError::Pg(_)), "{e}");
225            assert!(e.to_string().contains("?model="), "{e}");
226        }
227    }
228
229    /// `IMPROV_TEST_PG` (a URL or a `key=value` DSN) as a `postgres://` URL.
230    fn pg_url() -> Option<String> {
231        let Ok(dsn) = std::env::var("IMPROV_TEST_PG") else {
232            println!("skipped: IMPROV_TEST_PG not set");
233            return None;
234        };
235        if is_postgres(&dsn) {
236            return Some(dsn);
237        }
238        let kv: std::collections::HashMap<_, _> = dsn
239            .split_whitespace()
240            .filter_map(|p| p.split_once('='))
241            .collect();
242        let get = |k| kv.get(k).copied().unwrap_or_default();
243        Some(format!(
244            "postgres://{}:{}@{}/{}",
245            get("user"),
246            get("password"),
247            get("host"),
248            get("dbname")
249        ))
250    }
251
252    #[test]
253    fn live_postgres_store() {
254        let Some(url) = pg_url() else { return };
255        let url = format!("{url}?model=store_{}", std::process::id());
256        assert!(matches!(Store::open(&url).unwrap(), Store::Pg(_)));
257        exercise(&url);
258    }
259}