es_entity/operation/
mod.rs

1//! Handle execution of database operations and transactions.
2
3pub mod hooks;
4mod with_time;
5
6use sqlx::{Acquire, Transaction};
7
8use crate::{clock::ClockHandle, db};
9
10pub use with_time::*;
11
12/// Default return type of the derived EsRepo::begin_op().
13///
14/// Used as a wrapper of a [`sqlx::Transaction`] but can also cache the time at which the
15/// transaction is taking place.
16///
17/// When a manual clock is provided, the transaction will automatically cache that
18/// clock's time, enabling deterministic testing. This cached time will be used in all
19/// time-dependent operations.
20pub struct DbOp<'c> {
21    tx: Transaction<'c, db::Db>,
22    clock: ClockHandle,
23    now: Option<chrono::DateTime<chrono::Utc>>,
24    commit_hooks: Option<hooks::CommitHooks>,
25}
26
27impl<'c> DbOp<'c> {
28    fn new(
29        tx: Transaction<'c, db::Db>,
30        clock: ClockHandle,
31        time: Option<chrono::DateTime<chrono::Utc>>,
32    ) -> Self {
33        Self {
34            tx,
35            clock,
36            now: time,
37            commit_hooks: Some(hooks::CommitHooks::new()),
38        }
39    }
40
41    /// Initializes a transaction using the global clock.
42    ///
43    /// Delegates to [`init_with_clock`](Self::init_with_clock) using the global clock handle.
44    pub async fn init(pool: &db::Pool) -> Result<DbOp<'static>, sqlx::Error> {
45        Self::init_with_clock(pool, crate::clock::Clock::handle()).await
46    }
47
48    /// Initializes a transaction with the specified clock.
49    ///
50    /// If the clock is manual, its current time will be cached in the transaction.
51    pub async fn init_with_clock(
52        pool: &db::Pool,
53        clock: &ClockHandle,
54    ) -> Result<DbOp<'static>, sqlx::Error> {
55        let tx = pool.begin().await?;
56
57        // If a manual clock is provided, cache its time for consistent
58        // timestamps within the transaction.
59        let time = clock.manual_now();
60
61        Ok(DbOp::new(tx, clock.clone(), time))
62    }
63
64    /// Transitions to a [`DbOpWithTime`] with the given time cached.
65    pub fn with_time(self, time: chrono::DateTime<chrono::Utc>) -> DbOpWithTime<'c> {
66        DbOpWithTime::new(self, time)
67    }
68
69    /// Transitions to a [`DbOpWithTime`] using the clock.
70    ///
71    /// Uses cached time if present, otherwise uses the clock's current time.
72    pub fn with_clock_time(self) -> DbOpWithTime<'c> {
73        let time = self.now.unwrap_or_else(|| self.clock.now());
74        DbOpWithTime::new(self, time)
75    }
76
77    /// Transitions to a [`DbOpWithTime`] using the database time.
78    ///
79    /// Priority order:
80    /// 1. Cached time if present
81    /// 2. Manual clock time if the clock is manual
82    /// 3. Database time via `SELECT NOW()`
83    pub async fn with_db_time(mut self) -> Result<DbOpWithTime<'c>, sqlx::Error> {
84        let time = if let Some(time) = self.now {
85            time
86        } else if let Some(manual_time) = self.clock.manual_now() {
87            manual_time
88        } else {
89            db::database_now(&mut *self.tx).await?
90        };
91
92        Ok(DbOpWithTime::new(self, time))
93    }
94
95    /// Returns the optionally cached [`chrono::DateTime`]
96    pub fn maybe_now(&self) -> Option<chrono::DateTime<chrono::Utc>> {
97        self.now
98    }
99
100    /// Begins a nested transaction.
101    pub async fn begin(&mut self) -> Result<DbOp<'_>, sqlx::Error> {
102        Ok(DbOp::new(
103            self.tx.begin().await?,
104            self.clock.clone(),
105            self.now,
106        ))
107    }
108
109    /// Commits the inner transaction.
110    pub async fn commit(mut self) -> Result<(), sqlx::Error> {
111        let commit_hooks = self.commit_hooks.take().expect("no hooks");
112        let post_hooks = commit_hooks.execute_pre(&mut self).await?;
113        self.tx.commit().await?;
114        post_hooks.execute();
115        Ok(())
116    }
117
118    /// Gets a mutable handle to the inner transaction
119    pub fn tx_mut(&mut self) -> &mut Transaction<'c, db::Db> {
120        &mut self.tx
121    }
122}
123
124impl<'o> AtomicOperation for DbOp<'o> {
125    fn maybe_now(&self) -> Option<chrono::DateTime<chrono::Utc>> {
126        self.maybe_now()
127    }
128
129    fn clock(&self) -> &ClockHandle {
130        &self.clock
131    }
132
133    fn as_executor(&mut self) -> &mut db::Connection {
134        self.tx.as_executor()
135    }
136
137    fn add_commit_hook<H: hooks::CommitHook>(&mut self, hook: H) -> Result<(), H> {
138        self.commit_hooks.as_mut().expect("no hooks").add(hook);
139        Ok(())
140    }
141
142    fn commit_hook<H: hooks::CommitHook>(&self) -> Option<&H> {
143        self.commit_hooks.as_ref()?.get_last::<H>()
144    }
145}
146
147/// Equivileant of [`DbOp`] just that the time is guaranteed to be cached.
148///
149/// Used as a wrapper of a [`sqlx::Transaction`] with cached time of the transaction.
150pub struct DbOpWithTime<'c> {
151    inner: DbOp<'c>,
152    now: chrono::DateTime<chrono::Utc>,
153}
154
155impl<'c> DbOpWithTime<'c> {
156    fn new(mut inner: DbOp<'c>, time: chrono::DateTime<chrono::Utc>) -> Self {
157        inner.now = Some(time);
158        Self { inner, now: time }
159    }
160
161    /// The cached [`chrono::DateTime`]
162    pub fn now(&self) -> chrono::DateTime<chrono::Utc> {
163        self.now
164    }
165
166    /// Begins a nested transaction.
167    pub async fn begin(&mut self) -> Result<DbOpWithTime<'_>, sqlx::Error> {
168        Ok(DbOpWithTime::new(self.inner.begin().await?, self.now))
169    }
170
171    /// Commits the inner transaction.
172    pub async fn commit(self) -> Result<(), sqlx::Error> {
173        self.inner.commit().await
174    }
175
176    /// Gets a mutable handle to the inner transaction
177    pub fn tx_mut(&mut self) -> &mut Transaction<'c, db::Db> {
178        self.inner.tx_mut()
179    }
180}
181
182impl<'o> AtomicOperation for DbOpWithTime<'o> {
183    fn maybe_now(&self) -> Option<chrono::DateTime<chrono::Utc>> {
184        Some(self.now())
185    }
186
187    fn clock(&self) -> &ClockHandle {
188        self.inner.clock()
189    }
190
191    fn as_executor(&mut self) -> &mut db::Connection {
192        self.inner.as_executor()
193    }
194
195    fn add_commit_hook<H: hooks::CommitHook>(&mut self, hook: H) -> Result<(), H> {
196        self.inner.add_commit_hook(hook)
197    }
198
199    fn commit_hook<H: hooks::CommitHook>(&self) -> Option<&H> {
200        self.inner.commit_hook::<H>()
201    }
202}
203
204impl<'o> AtomicOperationWithTime for DbOpWithTime<'o> {
205    fn now(&self) -> chrono::DateTime<chrono::Utc> {
206        self.now
207    }
208}
209
210/// Trait to signify we can make multiple consistent database roundtrips.
211///
212/// Its a stand in for [`&mut sqlx::Transaction<'_, DB>`](`sqlx::Transaction`).
213/// The reason for having a trait is to support custom types that wrap the inner
214/// transaction while providing additional functionality.
215///
216/// See [`DbOp`] or [`DbOpWithTime`].
217pub trait AtomicOperation: Send {
218    /// Function for querying when the operation is taking place - if it is cached.
219    fn maybe_now(&self) -> Option<chrono::DateTime<chrono::Utc>> {
220        None
221    }
222
223    /// Returns the clock handle for time operations.
224    ///
225    /// Default implementation returns the global clock handle.
226    fn clock(&self) -> &ClockHandle {
227        crate::clock::Clock::handle()
228    }
229
230    /// Returns the [`sqlx::Executor`] implementation.
231    /// The desired way to represent this would actually be as a GAT:
232    /// ```rust
233    /// trait AtomicOperation {
234    ///     type Executor<'c>: sqlx::PgExecutor<'c>
235    ///         where Self: 'c;
236    ///
237    ///     fn as_executor<'c>(&'c mut self) -> Self::Executor<'c>;
238    /// }
239    /// ```
240    ///
241    /// But GATs don't play well with `async_trait::async_trait` due to lifetime constraints
242    /// so we return the concrete [`&mut db::Connection`](`crate::db::Connection`) instead as a work around.
243    ///
244    /// Since this trait is generally applied to types that wrap a [`sqlx::Transaction`]
245    /// there is no variance in the return type - so its fine.
246    fn as_executor(&mut self) -> &mut sqlx::PgConnection;
247
248    /// Registers a commit hook that will run pre_commit before and post_commit after the transaction commits.
249    /// Returns Ok(()) if the hook was registered, Err(hook) if hooks are not supported.
250    fn add_commit_hook<H: hooks::CommitHook>(&mut self, hook: H) -> Result<(), H> {
251        Err(hook)
252    }
253
254    /// Typed shared access to the currently-accumulating commit hook of type `H`,
255    /// if this operation supports commit hooks and one is registered.
256    /// Returns the hook a subsequent `add_commit_hook::<H>` call would merge into.
257    fn commit_hook<H: hooks::CommitHook>(&self) -> Option<&H> {
258        None
259    }
260}
261
262impl<'c> AtomicOperation for sqlx::Transaction<'c, db::Db> {
263    fn as_executor(&mut self) -> &mut db::Connection {
264        &mut *self
265    }
266}
es_entity/operation/mod.rs

es_entity/operation/
mod.rs
