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
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 op.clock().is_artificial() {
32            op.clock().now()
33        } else {
34            let res = sqlx::query!("SELECT NOW()")
35                .fetch_one(op.as_executor())
36                .await?;
37            res.now.expect("could not fetch now")
38        };
39        Ok(Self { inner: op, now })
40    }
41
42    /// Wraps with a specific time (uses existing if present).
43    pub fn cached_or_time(op: &'a mut Op, time: chrono::DateTime<chrono::Utc>) -> Self {
44        let now = op.maybe_now().unwrap_or(time);
45        Self { inner: op, now }
46    }
47
48    /// Wraps using system time (uses existing if present).
49    ///
50    /// Uses cached time if present, otherwise uses the operation's clock.
51    pub fn cached_or_clock_time(op: &'a mut Op) -> Self {
52        let now = op.maybe_now().unwrap_or_else(|| op.clock().now());
53        Self { inner: op, now }
54    }
55
56    pub fn now(&self) -> chrono::DateTime<chrono::Utc> {
57        self.now
58    }
59}
60
61impl<'a, Op: AtomicOperation + ?Sized> AtomicOperation for OpWithTime<'a, Op> {
62    fn maybe_now(&self) -> Option<chrono::DateTime<chrono::Utc>> {
63        Some(self.now)
64    }
65
66    fn clock(&self) -> &ClockHandle {
67        self.inner.clock()
68    }
69
70    fn as_executor(&mut self) -> &mut sqlx::PgConnection {
71        self.inner.as_executor()
72    }
73
74    fn add_commit_hook<H: hooks::CommitHook>(&mut self, hook: H) -> Result<(), H> {
75        self.inner.add_commit_hook(hook)
76    }
77}