1#![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#[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 #[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
47pub enum Store {
49 Mentat(ModelStore),
50 Pg(PgModelStore),
51}
52
53pub 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 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 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 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#[derive(Debug, Clone, Default)]
146pub struct SaveToken {
147 tx: Option<i64>,
148 pg: Option<improv_storage_pg::PgSnapshot>,
149}
150
151impl SaveToken {
152 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 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 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 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}