agent_store/
generation.rs1use crate::backend::{as_u64, Backend, Value};
12use crate::error::{Result, StoreError};
13
14const META_TABLE: &str = "_agent_store_meta";
15
16#[derive(Clone, Debug)]
18pub struct Generation {
19 key: String,
20}
21
22impl Generation {
23 pub fn new(key: impl Into<String>) -> Self {
25 Self { key: key.into() }
26 }
27
28 pub fn ensure_schema(db: &dyn Backend) -> Result<()> {
30 db.exec(
31 &format!(
32 "CREATE TABLE IF NOT EXISTS {META_TABLE} (\
33 key TEXT PRIMARY KEY, ival INTEGER NOT NULL)"
34 ),
35 &[],
36 )?;
37 Ok(())
38 }
39
40 pub fn current(&self, db: &dyn Backend) -> Result<u64> {
42 let rows = db.query(
43 &format!("SELECT ival FROM {META_TABLE} WHERE key = ?"),
44 &[Value::Text(self.key.clone())],
45 )?;
46 match rows.first() {
47 Some(row) => as_u64(&row[0]),
48 None => Ok(0),
49 }
50 }
51
52 pub fn bump(&self, db: &dyn Backend) -> Result<u64> {
54 let rows = db.query(
57 &format!(
58 "INSERT INTO {META_TABLE} (key, ival) VALUES (?, 1) \
59 ON CONFLICT(key) DO UPDATE SET ival = ival + 1 RETURNING ival"
60 ),
61 &[Value::Text(self.key.clone())],
62 )?;
63 let row = rows
64 .first()
65 .ok_or_else(|| StoreError::Backend("bump: RETURNING produced no row".into()))?;
66 as_u64(&row[0])
67 }
68
69 pub fn set(&self, db: &dyn Backend, value: u64) -> Result<()> {
71 let current = self.current(db)?;
72 if value < current {
73 return Err(StoreError::NonMonotonicGeneration {
74 key: self.key.clone(),
75 current,
76 attempted: value,
77 });
78 }
79 db.exec(
80 &format!(
81 "INSERT INTO {META_TABLE} (key, ival) VALUES (?, ?) \
82 ON CONFLICT(key) DO UPDATE SET ival = excluded.ival"
83 ),
84 &[Value::Text(self.key.clone()), Value::Int(value as i64)],
85 )?;
86 Ok(())
87 }
88}
89
90#[cfg(test)]
91mod tests {
92 use super::*;
93 use crate::backend::SqliteBackend;
94
95 fn db() -> SqliteBackend {
96 let db = SqliteBackend::in_memory().unwrap();
97 Generation::ensure_schema(&db).unwrap();
98 db
99 }
100
101 #[test]
102 fn starts_at_zero_then_increments() {
103 let db = db();
104 let g = Generation::new("report");
105 assert_eq!(g.current(&db).unwrap(), 0);
106 assert_eq!(g.bump(&db).unwrap(), 1);
107 assert_eq!(g.bump(&db).unwrap(), 2);
108 assert_eq!(g.current(&db).unwrap(), 2);
109 }
110
111 #[test]
112 fn counters_are_independent() {
113 let db = db();
114 let a = Generation::new("a");
115 let b = Generation::new("b");
116 a.bump(&db).unwrap();
117 a.bump(&db).unwrap();
118 b.bump(&db).unwrap();
119 assert_eq!(a.current(&db).unwrap(), 2);
120 assert_eq!(b.current(&db).unwrap(), 1);
121 }
122
123 #[test]
124 fn set_advances_but_never_rewinds() {
125 let db = db();
126 let g = Generation::new("report");
127 g.set(&db, 10).unwrap();
128 assert_eq!(g.current(&db).unwrap(), 10);
129
130 let err = g.set(&db, 3).unwrap_err();
132 assert!(matches!(
133 err,
134 StoreError::NonMonotonicGeneration {
135 current: 10,
136 attempted: 3,
137 ..
138 }
139 ));
140 assert_eq!(g.current(&db).unwrap(), 10);
141 }
142}