Skip to main content

uqa_storage/
transaction.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Explicit user transactions with savepoint stacks.
8//!
9//! Transaction state and savepoint management.
10//!
11//! Two flavours:
12//! * [`SQLiteTransaction`] -- wraps a [`ManagedConnection`] and routes
13//!   begin / commit / rollback / savepoint calls to the underlying
14//!   connection. Auto-rollback on drop guards the `SQLite` write log
15//!   against in-flight panics.
16//! * [`InMemoryTransaction`] -- snapshots the per-table mutable
17//!   document store (or whatever the supplied [`Snapshotable`]
18//!   impl returns) and rolls it back on demand. Savepoints push
19//!   nested snapshots onto a stack.
20//!
21//! Both expose the same `commit` / `rollback` / `savepoint` /
22//! `release_savepoint` / `rollback_to` methods so the engine can
23//! treat them uniformly.
24
25use 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
42/// SQLite-backed transaction. Drops without commit roll back so
43/// panics never leak a half-applied write log.
44pub 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
123/// Sources that can snapshot themselves and restore the snapshot.
124/// The engine implements this for its in-memory table state.
125pub trait Snapshotable {
126    type Snapshot;
127    fn snapshot(&self) -> Self::Snapshot;
128    fn restore(&self, snapshot: &Self::Snapshot);
129}
130
131/// Pure-in-memory transaction. Snapshots the source via
132/// [`Snapshotable::snapshot`] on `begin`, restores via
133/// [`Snapshotable::restore`] on rollback.
134pub 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    /// Counter source: the snapshot is just the counter value at
224    /// snapshot time. `restore` overwrites the live counter.
225    #[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            // Tx drops without commit -> rollback fires automatically.
330        }
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}