kaizen/experiment/
store.rs

1// SPDX-License-Identifier: AGPL-3.0-or-later
2//! Persistence for experiments. IO at boundary; pure types in `types.rs`.
3
4use crate::experiment::binding::ManualTags;
5use crate::experiment::types::{Classification, Experiment, State};
6use crate::store::Store;
7use anyhow::{Context, Result};
8use rusqlite::{OptionalExtension, params};
9
10pub fn save_experiment(store: &Store, exp: &Experiment) -> Result<()> {
11    let json = serde_json::to_string(exp).context("serialize experiment")?;
12    store.conn().execute(
13        "INSERT INTO experiments (id, name, created_at_ms, metadata, state, concluded_at_ms)
14         VALUES (?1, ?2, ?3, ?4, ?5, ?6)
15         ON CONFLICT(id) DO UPDATE SET
16           name=excluded.name,
17           metadata=excluded.metadata,
18           state=excluded.state,
19           concluded_at_ms=excluded.concluded_at_ms",
20        params![
21            exp.id,
22            exp.name,
23            exp.created_at_ms as i64,
24            json,
25            format!("{:?}", exp.state),
26            exp.concluded_at_ms.map(|v| v as i64),
27        ],
28    )?;
29    Ok(())
30}
31
32pub fn load_experiment(store: &Store, id: &str) -> Result<Option<Experiment>> {
33    let row: Option<String> = store
34        .conn()
35        .query_row(
36            "SELECT metadata FROM experiments WHERE id = ?1",
37            params![id],
38            |r| r.get(0),
39        )
40        .optional()?;
41    match row {
42        Some(s) => Ok(Some(serde_json::from_str(&s)?)),
43        None => Ok(None),
44    }
45}
46
47pub fn list_experiments(store: &Store) -> Result<Vec<Experiment>> {
48    let mut stmt = store
49        .conn()
50        .prepare("SELECT metadata FROM experiments ORDER BY created_at_ms DESC")?;
51    let rows = stmt.query_map([], |r| r.get::<_, String>(0))?;
52    let mut out = Vec::new();
53    for row in rows {
54        let s = row?;
55        if let Ok(e) = serde_json::from_str::<Experiment>(&s) {
56            out.push(e);
57        }
58    }
59    Ok(out)
60}
61
62pub fn set_state(store: &Store, id: &str, state: State, now_ms: u64) -> Result<()> {
63    let Some(mut exp) = load_experiment(store, id)? else {
64        anyhow::bail!("experiment not found: {id}");
65    };
66    exp.state = state;
67    if matches!(state, State::Concluded) {
68        exp.concluded_at_ms = Some(now_ms);
69    }
70    save_experiment(store, &exp)
71}
72
73pub fn tag_session(
74    store: &Store,
75    exp_id: &str,
76    session_id: &str,
77    variant: Classification,
78) -> Result<()> {
79    store.conn().execute(
80        "INSERT INTO experiment_tags (experiment_id, session_id, variant)
81         VALUES (?1, ?2, ?3)
82         ON CONFLICT(experiment_id, session_id) DO UPDATE SET variant=excluded.variant",
83        params![exp_id, session_id, format!("{:?}", variant)],
84    )?;
85    Ok(())
86}
87
88pub fn manual_tags(store: &Store, exp_id: &str) -> Result<ManualTags> {
89    let mut stmt = store
90        .conn()
91        .prepare("SELECT session_id, variant FROM experiment_tags WHERE experiment_id = ?1")?;
92    let rows = stmt.query_map(params![exp_id], |r| {
93        Ok((r.get::<_, String>(0)?, r.get::<_, String>(1)?))
94    })?;
95    let mut out = ManualTags::new();
96    for row in rows {
97        let (sid, variant) = row?;
98        let v = match variant.as_str() {
99            "Control" => Classification::Control,
100            "Treatment" => Classification::Treatment,
101            _ => Classification::Excluded,
102        };
103        out.insert(sid, v);
104    }
105    Ok(out)
106}
107
108#[cfg(test)]
109mod tests {
110    use super::*;
111    use crate::experiment::types::{Binding, Criterion, Direction, Metric, State};
112    use tempfile::TempDir;
113
114    fn mk(id: &str) -> Experiment {
115        Experiment {
116            id: id.into(),
117            name: format!("exp-{id}"),
118            hypothesis: "h".into(),
119            change_description: "c".into(),
120            metric: Metric::TokensPerSession,
121            binding: Binding::GitCommit {
122                control_commit: "c1".into(),
123                treatment_commit: "c2".into(),
124            },
125            duration_days: 14,
126            success_criterion: Criterion::Delta {
127                direction: Direction::Decrease,
128                target_pct: 10.0,
129            },
130            state: State::Draft,
131            created_at_ms: 1000,
132            concluded_at_ms: None,
133        }
134    }
135
136    #[test]
137    fn round_trip_save_load() {
138        let dir = TempDir::new().unwrap();
139        let store = Store::open(&dir.path().join("k.db")).unwrap();
140        let e = mk("a");
141        save_experiment(&store, &e).unwrap();
142        let got = load_experiment(&store, "a").unwrap().unwrap();
143        assert_eq!(got.id, "a");
144        assert_eq!(got.state, State::Draft);
145    }
146
147    #[test]
148    fn set_state_transitions() {
149        let dir = TempDir::new().unwrap();
150        let store = Store::open(&dir.path().join("k.db")).unwrap();
151        save_experiment(&store, &mk("b")).unwrap();
152        set_state(&store, "b", State::Running, 5_000).unwrap();
153        let got = load_experiment(&store, "b").unwrap().unwrap();
154        assert_eq!(got.state, State::Running);
155        set_state(&store, "b", State::Concluded, 9_000).unwrap();
156        let got = load_experiment(&store, "b").unwrap().unwrap();
157        assert_eq!(got.state, State::Concluded);
158        assert_eq!(got.concluded_at_ms, Some(9_000));
159    }
160
161    #[test]
162    fn tags_round_trip() {
163        let dir = TempDir::new().unwrap();
164        let store = Store::open(&dir.path().join("k.db")).unwrap();
165        save_experiment(&store, &mk("e")).unwrap();
166        tag_session(&store, "e", "s1", Classification::Treatment).unwrap();
167        tag_session(&store, "e", "s2", Classification::Control).unwrap();
168        let tags = manual_tags(&store, "e").unwrap();
169        assert_eq!(tags.get("s1"), Some(&Classification::Treatment));
170        assert_eq!(tags.get("s2"), Some(&Classification::Control));
171    }
172}
kaizen/experiment/store.rs

kaizen/experiment/
store.rs
