uqa_storage/
transaction.rs1use std::collections::BTreeMap;
26
27use crate::sqlite::connection::ManagedConnection;
28use crate::SQLiteError;
29
30#[derive(Debug, thiserror::Error)]
31pub enum TransactionError {
32 #[error("transaction already finished")]
33 Finished,
34 #[error("savepoint `{0}` does not exist")]
35 UnknownSavepoint(String),
36 #[error(transparent)]
37 Storage(#[from] SQLiteError),
38}
39
40pub type TxResult<T> = std::result::Result<T, TransactionError>;
41
42pub struct SQLiteTransaction {
45 conn: ManagedConnection,
46 finished: bool,
47}
48
49impl SQLiteTransaction {
50 pub fn begin(conn: ManagedConnection) -> Result<Self, SQLiteError> {
51 conn.begin_transaction()?;
52 Ok(Self {
53 conn,
54 finished: false,
55 })
56 }
57
58 pub fn active(&self) -> bool {
59 !self.finished
60 }
61
62 pub fn commit(&mut self) -> TxResult<()> {
63 if self.finished {
64 return Err(TransactionError::Finished);
65 }
66 let result = self
67 .conn
68 .commit_transaction()
69 .map_err(TransactionError::from);
70 if result.is_ok() || !self.conn.in_transaction() {
71 self.finished = true;
72 }
73 result
74 }
75
76 pub fn rollback(&mut self) -> TxResult<()> {
77 if self.finished {
78 return Err(TransactionError::Finished);
79 }
80 let result = self
81 .conn
82 .rollback_transaction()
83 .map_err(TransactionError::from);
84 if result.is_ok() || !self.conn.in_transaction() {
85 self.finished = true;
86 }
87 result
88 }
89
90 pub fn savepoint(&self, name: &str) -> TxResult<()> {
91 if self.finished {
92 return Err(TransactionError::Finished);
93 }
94 self.conn.savepoint(name)?;
95 Ok(())
96 }
97
98 pub fn release_savepoint(&self, name: &str) -> TxResult<()> {
99 if self.finished {
100 return Err(TransactionError::Finished);
101 }
102 self.conn.release_savepoint(name)?;
103 Ok(())
104 }
105
106 pub fn rollback_to(&self, name: &str) -> TxResult<()> {
107 if self.finished {
108 return Err(TransactionError::Finished);
109 }
110 self.conn.rollback_to_savepoint(name)?;
111 Ok(())
112 }
113}
114
115impl Drop for SQLiteTransaction {
116 fn drop(&mut self) {
117 if !self.finished {
118 self.conn.rollback_transaction_on_drop();
119 }
120 }
121}
122
123pub trait Snapshotable {
126 type Snapshot;
127 fn snapshot(&self) -> Self::Snapshot;
128 fn restore(&self, snapshot: &Self::Snapshot);
129}
130
131pub struct InMemoryTransaction<S: Snapshotable> {
135 source: S,
136 snapshot: Option<S::Snapshot>,
137 savepoints: BTreeMap<String, S::Snapshot>,
138 finished: bool,
139}
140
141impl<S: Snapshotable> InMemoryTransaction<S> {
142 pub fn begin(source: S) -> Self {
143 let snapshot = source.snapshot();
144 Self {
145 source,
146 snapshot: Some(snapshot),
147 savepoints: BTreeMap::new(),
148 finished: false,
149 }
150 }
151
152 pub fn active(&self) -> bool {
153 !self.finished
154 }
155
156 pub fn commit(&mut self) -> TxResult<()> {
157 if self.finished {
158 return Err(TransactionError::Finished);
159 }
160 self.snapshot = None;
161 self.savepoints.clear();
162 self.finished = true;
163 Ok(())
164 }
165
166 pub fn rollback(&mut self) -> TxResult<()> {
167 if self.finished {
168 return Err(TransactionError::Finished);
169 }
170 if let Some(snap) = self.snapshot.take() {
171 self.source.restore(&snap);
172 }
173 self.savepoints.clear();
174 self.finished = true;
175 Ok(())
176 }
177
178 pub fn savepoint(&mut self, name: impl Into<String>) -> TxResult<()> {
179 if self.finished {
180 return Err(TransactionError::Finished);
181 }
182 self.savepoints.insert(name.into(), self.source.snapshot());
183 Ok(())
184 }
185
186 pub fn release_savepoint(&mut self, name: &str) -> TxResult<()> {
187 if self.finished {
188 return Err(TransactionError::Finished);
189 }
190 self.savepoints.remove(name);
191 Ok(())
192 }
193
194 pub fn rollback_to(&mut self, name: &str) -> TxResult<()> {
195 if self.finished {
196 return Err(TransactionError::Finished);
197 }
198 let snap = self
199 .savepoints
200 .get(name)
201 .ok_or_else(|| TransactionError::UnknownSavepoint(name.to_string()))?;
202 self.source.restore(snap);
203 Ok(())
204 }
205}
206
207impl<S: Snapshotable> Drop for InMemoryTransaction<S> {
208 fn drop(&mut self) {
209 if !self.finished {
210 if let Some(snap) = self.snapshot.take() {
211 self.source.restore(&snap);
212 }
213 }
214 }
215}
216
217#[cfg(test)]
218mod tests {
219 use super::*;
220 use std::cell::RefCell;
221 use std::rc::Rc;
222
223 #[derive(Clone)]
226 struct CounterSource {
227 v: Rc<RefCell<i64>>,
228 }
229 impl Snapshotable for CounterSource {
230 type Snapshot = i64;
231 fn snapshot(&self) -> i64 {
232 *self.v.borrow()
233 }
234 fn restore(&self, snap: &i64) {
235 *self.v.borrow_mut() = *snap;
236 }
237 }
238
239 #[test]
240 fn rollback_restores_state() {
241 let v = Rc::new(RefCell::new(10));
242 let src = CounterSource { v: v.clone() };
243 let mut tx = InMemoryTransaction::begin(src);
244 *v.borrow_mut() = 99;
245 tx.rollback().unwrap();
246 assert_eq!(*v.borrow(), 10);
247 }
248
249 #[test]
250 fn commit_keeps_changes() {
251 let v = Rc::new(RefCell::new(10));
252 let src = CounterSource { v: v.clone() };
253 let mut tx = InMemoryTransaction::begin(src);
254 *v.borrow_mut() = 99;
255 tx.commit().unwrap();
256 assert_eq!(*v.borrow(), 99);
257 }
258
259 #[test]
260 fn savepoint_rollback_undoes_partial_writes() {
261 let v = Rc::new(RefCell::new(0));
262 let src = CounterSource { v: v.clone() };
263 let mut tx = InMemoryTransaction::begin(src);
264 *v.borrow_mut() = 1;
265 tx.savepoint("sp").unwrap();
266 *v.borrow_mut() = 2;
267 tx.rollback_to("sp").unwrap();
268 assert_eq!(*v.borrow(), 1);
269 tx.commit().unwrap();
270 }
271
272 #[test]
273 fn drop_without_commit_rolls_back() {
274 let v = Rc::new(RefCell::new(10));
275 {
276 let src = CounterSource { v: v.clone() };
277 let _tx = InMemoryTransaction::begin(src);
278 *v.borrow_mut() = 50;
279 }
280 assert_eq!(*v.borrow(), 10);
281 }
282
283 #[test]
284 fn unknown_savepoint_errors() {
285 let v = Rc::new(RefCell::new(0));
286 let src = CounterSource { v };
287 let mut tx = InMemoryTransaction::begin(src);
288 let err = tx.rollback_to("missing").unwrap_err();
289 assert!(matches!(err, TransactionError::UnknownSavepoint(_)));
290 tx.commit().unwrap();
291 }
292
293 #[test]
294 fn sqlite_transaction_commits_writes() {
295 let conn = ManagedConnection::open_in_memory().unwrap();
296 conn.with(|c| {
297 c.execute("CREATE TABLE t (id INTEGER PRIMARY KEY, v TEXT)", [])?;
298 Ok(())
299 })
300 .unwrap();
301 let mut tx = SQLiteTransaction::begin(conn.clone()).unwrap();
302 conn.with(|c| {
303 c.execute("INSERT INTO t (id, v) VALUES (1, 'hi')", [])?;
304 Ok(())
305 })
306 .unwrap();
307 tx.commit().unwrap();
308 let got: i64 = conn
309 .with(|c| Ok(c.query_row("SELECT COUNT(*) FROM t", [], |r| r.get(0))?))
310 .unwrap();
311 assert_eq!(got, 1);
312 }
313
314 #[test]
315 fn sqlite_transaction_rolls_back_on_drop() {
316 let conn = ManagedConnection::open_in_memory().unwrap();
317 conn.with(|c| {
318 c.execute("CREATE TABLE t (id INTEGER PRIMARY KEY, v TEXT)", [])?;
319 Ok(())
320 })
321 .unwrap();
322 {
323 let _tx = SQLiteTransaction::begin(conn.clone()).unwrap();
324 conn.with(|c| {
325 c.execute("INSERT INTO t (id, v) VALUES (1, 'hi')", [])?;
326 Ok(())
327 })
328 .unwrap();
329 }
331 let got: i64 = conn
332 .with(|c| Ok(c.query_row("SELECT COUNT(*) FROM t", [], |r| r.get(0))?))
333 .unwrap();
334 assert_eq!(got, 0);
335 }
336}