Skip to main content

uqa_storage/
transaction.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! In-memory transaction snapshots and shared transaction errors.
8
9use std::collections::BTreeMap;
10
11#[derive(Debug, thiserror::Error)]
12pub enum TransactionError {
13    #[error("transaction already finished")]
14    Finished,
15    #[error("savepoint `{0}` does not exist")]
16    UnknownSavepoint(String),
17    #[error(transparent)]
18    Storage(#[from] crate::StorageBackendError),
19}
20
21pub type TxResult<T> = std::result::Result<T, TransactionError>;
22
23/// Sources that can snapshot themselves and restore the snapshot.
24/// The engine implements this for its in-memory table state.
25pub trait Snapshotable {
26    type Snapshot;
27    fn snapshot(&self) -> Self::Snapshot;
28    fn restore(&self, snapshot: &Self::Snapshot);
29}
30
31/// Pure-in-memory transaction. Snapshots the source via
32/// [`Snapshotable::snapshot`] on `begin`, restores via
33/// [`Snapshotable::restore`] on rollback.
34pub struct InMemoryTransaction<S: Snapshotable> {
35    source: S,
36    snapshot: Option<S::Snapshot>,
37    savepoints: BTreeMap<String, S::Snapshot>,
38    finished: bool,
39}
40
41impl<S: Snapshotable> InMemoryTransaction<S> {
42    pub fn begin(source: S) -> Self {
43        let snapshot = source.snapshot();
44        Self {
45            source,
46            snapshot: Some(snapshot),
47            savepoints: BTreeMap::new(),
48            finished: false,
49        }
50    }
51
52    pub fn active(&self) -> bool {
53        !self.finished
54    }
55
56    pub fn commit(&mut self) -> TxResult<()> {
57        if self.finished {
58            return Err(TransactionError::Finished);
59        }
60        self.snapshot = None;
61        self.savepoints.clear();
62        self.finished = true;
63        Ok(())
64    }
65
66    pub fn rollback(&mut self) -> TxResult<()> {
67        if self.finished {
68            return Err(TransactionError::Finished);
69        }
70        if let Some(snap) = self.snapshot.take() {
71            self.source.restore(&snap);
72        }
73        self.savepoints.clear();
74        self.finished = true;
75        Ok(())
76    }
77
78    pub fn savepoint(&mut self, name: impl Into<String>) -> TxResult<()> {
79        if self.finished {
80            return Err(TransactionError::Finished);
81        }
82        self.savepoints.insert(name.into(), self.source.snapshot());
83        Ok(())
84    }
85
86    pub fn release_savepoint(&mut self, name: &str) -> TxResult<()> {
87        if self.finished {
88            return Err(TransactionError::Finished);
89        }
90        self.savepoints.remove(name);
91        Ok(())
92    }
93
94    pub fn rollback_to(&mut self, name: &str) -> TxResult<()> {
95        if self.finished {
96            return Err(TransactionError::Finished);
97        }
98        let snap = self
99            .savepoints
100            .get(name)
101            .ok_or_else(|| TransactionError::UnknownSavepoint(name.to_string()))?;
102        self.source.restore(snap);
103        Ok(())
104    }
105}
106
107impl<S: Snapshotable> Drop for InMemoryTransaction<S> {
108    fn drop(&mut self) {
109        if !self.finished {
110            if let Some(snap) = self.snapshot.take() {
111                self.source.restore(&snap);
112            }
113        }
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120    use std::cell::RefCell;
121    use std::rc::Rc;
122
123    /// Counter source: the snapshot is just the counter value at
124    /// snapshot time. `restore` overwrites the live counter.
125    #[derive(Clone)]
126    struct CounterSource {
127        v: Rc<RefCell<i64>>,
128    }
129    impl Snapshotable for CounterSource {
130        type Snapshot = i64;
131        fn snapshot(&self) -> i64 {
132            *self.v.borrow()
133        }
134        fn restore(&self, snap: &i64) {
135            *self.v.borrow_mut() = *snap;
136        }
137    }
138
139    #[test]
140    fn rollback_restores_state() {
141        let v = Rc::new(RefCell::new(10));
142        let src = CounterSource { v: v.clone() };
143        let mut tx = InMemoryTransaction::begin(src);
144        *v.borrow_mut() = 99;
145        tx.rollback().unwrap();
146        assert_eq!(*v.borrow(), 10);
147    }
148
149    #[test]
150    fn commit_keeps_changes() {
151        let v = Rc::new(RefCell::new(10));
152        let src = CounterSource { v: v.clone() };
153        let mut tx = InMemoryTransaction::begin(src);
154        *v.borrow_mut() = 99;
155        tx.commit().unwrap();
156        assert_eq!(*v.borrow(), 99);
157    }
158
159    #[test]
160    fn savepoint_rollback_undoes_partial_writes() {
161        let v = Rc::new(RefCell::new(0));
162        let src = CounterSource { v: v.clone() };
163        let mut tx = InMemoryTransaction::begin(src);
164        *v.borrow_mut() = 1;
165        tx.savepoint("sp").unwrap();
166        *v.borrow_mut() = 2;
167        tx.rollback_to("sp").unwrap();
168        assert_eq!(*v.borrow(), 1);
169        tx.commit().unwrap();
170    }
171
172    #[test]
173    fn drop_without_commit_rolls_back() {
174        let v = Rc::new(RefCell::new(10));
175        {
176            let src = CounterSource { v: v.clone() };
177            let _tx = InMemoryTransaction::begin(src);
178            *v.borrow_mut() = 50;
179        }
180        assert_eq!(*v.borrow(), 10);
181    }
182
183    #[test]
184    fn unknown_savepoint_errors() {
185        let v = Rc::new(RefCell::new(0));
186        let src = CounterSource { v };
187        let mut tx = InMemoryTransaction::begin(src);
188        let err = tx.rollback_to("missing").unwrap_err();
189        assert!(matches!(err, TransactionError::UnknownSavepoint(_)));
190        tx.commit().unwrap();
191    }
192}