es_entity/operation/
with_time.rs

1use super::{AtomicOperation, hooks};
2
3pub trait AtomicOperationWithTime: AtomicOperation {
4    fn now(&self) -> chrono::DateTime<chrono::Utc>;
5}
6
7/// Wrapper that guarantees time is available, borrowing the underlying operation.
8pub struct OpWithTime<'a, Op: AtomicOperation + ?Sized> {
9    inner: &'a mut Op,
10    now: chrono::DateTime<chrono::Utc>,
11}
12
13impl<'a, Op: AtomicOperation + ?Sized> AtomicOperationWithTime for OpWithTime<'a, Op> {
14    fn now(&self) -> chrono::DateTime<chrono::Utc> {
15        self.now
16    }
17}
18
19impl<'a, Op: AtomicOperation + ?Sized> OpWithTime<'a, Op> {
20    /// Wraps an operation, using existing time if present, otherwise fetching from DB.
21    pub async fn cached_or_db_time(op: &'a mut Op) -> Result<Self, sqlx::Error> {
22        let now = if let Some(time) = op.maybe_now() {
23            time
24        } else {
25            #[cfg(feature = "sim-time")]
26            {
27                crate::prelude::sim_time::now()
28            }
29            #[cfg(not(feature = "sim-time"))]
30            {
31                let res = sqlx::query!("SELECT NOW()")
32                    .fetch_one(op.as_executor())
33                    .await?;
34                res.now.expect("could not fetch now")
35            }
36        };
37        Ok(Self { inner: op, now })
38    }
39
40    /// Wraps with a specific time (uses existing if present).
41    pub fn cached_or_time(op: &'a mut Op, time: chrono::DateTime<chrono::Utc>) -> Self {
42        let now = op.maybe_now().unwrap_or(time);
43        Self { inner: op, now }
44    }
45
46    /// Wraps using system time (uses existing if present).
47    pub fn cached_or_system_time(op: &'a mut Op) -> Self {
48        let now = op.maybe_now().unwrap_or_else(|| {
49            #[cfg(feature = "sim-time")]
50            {
51                crate::prelude::sim_time::now()
52            }
53            #[cfg(not(feature = "sim-time"))]
54            {
55                chrono::Utc::now()
56            }
57        });
58        Self { inner: op, now }
59    }
60
61    pub fn now(&self) -> chrono::DateTime<chrono::Utc> {
62        self.now
63    }
64}
65
66impl<'a, Op: AtomicOperation + ?Sized> AtomicOperation for OpWithTime<'a, Op> {
67    fn maybe_now(&self) -> Option<chrono::DateTime<chrono::Utc>> {
68        Some(self.now)
69    }
70
71    fn as_executor(&mut self) -> &mut sqlx::PgConnection {
72        self.inner.as_executor()
73    }
74
75    fn add_commit_hook<H: hooks::CommitHook>(&mut self, hook: H) -> Result<(), H> {
76        self.inner.add_commit_hook(hook)
77    }
78}
es_entity/operation/with_time.rs

es_entity/operation/
with_time.rs
