Skip to main content

es_entity/operation/
with_time.rs

1use crate::clock::ClockHandle;
2
3use super::{AtomicOperation, hooks};
4
5pub trait AtomicOperationWithTime: AtomicOperation {
6    fn now(&self) -> chrono::DateTime<chrono::Utc>;
7}
8
9/// Wrapper that guarantees time is available, borrowing the underlying operation.
10pub struct OpWithTime<'a, Op: AtomicOperation + ?Sized> {
11    inner: &'a mut Op,
12    now: chrono::DateTime<chrono::Utc>,
13}
14
15impl<'a, Op: AtomicOperation + ?Sized> AtomicOperationWithTime for OpWithTime<'a, Op> {
16    fn now(&self) -> chrono::DateTime<chrono::Utc> {
17        self.now
18    }
19}
20
21impl<'a, Op: AtomicOperation + ?Sized> OpWithTime<'a, Op> {
22    /// Wraps an operation, using existing time if present, otherwise fetching from DB.
23    ///
24    /// Priority order:
25    /// 1. Cached time from operation
26    /// 2. Artificial clock time if the operation's clock is artificial (and hasn't transitioned)
27    /// 3. Database time via `SELECT NOW()`
28    pub async fn cached_or_db_time(op: &'a mut Op) -> Result<Self, sqlx::Error> {
29        let now = if let Some(time) = op.maybe_now() {
30            time
31        } else if let Some(artificial_time) = op.clock().artificial_now() {
32            artificial_time
33        } else {
34            sqlx::query_scalar::<_, chrono::DateTime<chrono::Utc>>("SELECT NOW()")
35                .fetch_one(op.as_executor())
36                .await?
37        };
38        Ok(Self { inner: op, now })
39    }
40
41    /// Wraps with a specific time (uses existing if present).
42    pub fn cached_or_time(op: &'a mut Op, time: chrono::DateTime<chrono::Utc>) -> Self {
43        let now = op.maybe_now().unwrap_or(time);
44        Self { inner: op, now }
45    }
46
47    /// Wraps using system time (uses existing if present).
48    ///
49    /// Uses cached time if present, otherwise uses the operation's clock.
50    pub fn cached_or_clock_time(op: &'a mut Op) -> Self {
51        let now = op.maybe_now().unwrap_or_else(|| op.clock().now());
52        Self { inner: op, now }
53    }
54
55    pub fn now(&self) -> chrono::DateTime<chrono::Utc> {
56        self.now
57    }
58}
59
60impl<'a, Op: AtomicOperation + ?Sized> AtomicOperation for OpWithTime<'a, Op> {
61    fn maybe_now(&self) -> Option<chrono::DateTime<chrono::Utc>> {
62        Some(self.now)
63    }
64
65    fn clock(&self) -> &ClockHandle {
66        self.inner.clock()
67    }
68
69    fn as_executor(&mut self) -> &mut sqlx::PgConnection {
70        self.inner.as_executor()
71    }
72
73    fn add_commit_hook<H: hooks::CommitHook>(&mut self, hook: H) -> Result<(), H> {
74        self.inner.add_commit_hook(hook)
75    }
76}