resident_utils/sqlx/
holder.rs

1use std::{collections::HashMap, future::Future, hash::Hash, time::Duration};
2
3use chrono::prelude::*;
4
5use thiserror::Error;
6
7use super::SqlxPool;
8
9#[derive(Error, Debug)]
10pub enum Error {
11    #[error("Invalid {0}")]
12    Invalid(String),
13}
14
15pub struct HolderMap<K, V> {
16    map: HashMap<K, V>,
17    expire_interval: Duration,
18    expire_at: DateTime<Utc>,
19    pg_pool: SqlxPool,
20}
21
22impl<K, V> HolderMap<K, V>
23where
24    K: PartialEq + Eq + Hash + Clone,
25    V: Clone,
26{
27    pub fn new(pg_pool: SqlxPool, expire_interval: Duration, now: Option<DateTime<Utc>>) -> Self {
28        Self {
29            map: HashMap::new(),
30            expire_interval,
31            expire_at: now.unwrap_or(Utc::now()),
32            pg_pool,
33        }
34    }
35
36    pub async fn get<FutOne, FutAll>(
37        &mut self,
38        key: &K,
39        now: Option<DateTime<Utc>>,
40        f: impl FnOnce(SqlxPool, K) -> FutOne,
41        g: impl FnOnce(SqlxPool) -> FutAll,
42    ) -> Result<Option<V>, Error>
43    where
44        FutOne: Future<Output = Result<Option<V>, Error>>,
45        FutAll: Future<Output = Result<HashMap<K, V>, Error>>,
46    {
47        if get_now(now) >= self.expire_at {
48            let pg_client = self.pg_pool.clone();
49            self.map = g(pg_client).await?;
50            self.expire_at = expire_at(now, self.expire_interval);
51        }
52        if let Some(value) = self.map.get(key) {
53            return Ok(Some(value.clone()));
54        }
55        let pg_client = self.pg_pool.clone();
56        let Some(value) = f(pg_client, key.clone()).await? else {
57            return Ok(None);
58        };
59        self.map.insert(key.clone(), value.clone());
60        Ok(Some(value))
61    }
62}
63
64pub struct HolderMapEachExpire<K, V> {
65    map: HashMap<K, (V, DateTime<Utc>)>,
66    expire_interval: Duration,
67    pg_pool: SqlxPool,
68}
69
70impl<K, V> HolderMapEachExpire<K, V>
71where
72    K: PartialEq + Eq + Hash + Clone,
73    V: Clone,
74{
75    pub fn new(pg_pool: SqlxPool, expire_interval: Duration) -> Self {
76        Self {
77            map: HashMap::new(),
78            expire_interval,
79            pg_pool,
80        }
81    }
82
83    pub async fn get<FutOne>(
84        &mut self,
85        key: &K,
86        now: Option<DateTime<Utc>>,
87        f: impl FnOnce(SqlxPool, K) -> FutOne,
88    ) -> Result<Option<V>, Error>
89    where
90        FutOne: Future<Output = Result<Option<V>, Error>>,
91    {
92        match self.map.get(key) {
93            Some((value, expire_at)) if get_now(now) < *expire_at => {
94                return Ok(Some(value.clone()));
95            }
96            _ => {}
97        }
98        let pg_client = self.pg_pool.clone();
99        let Some(value) = f(pg_client, key.clone()).await? else {
100            return Ok(None);
101        };
102        self.map.insert(
103            key.clone(),
104            (value.clone(), expire_at(now, self.expire_interval)),
105        );
106        Ok(Some(value))
107    }
108}
109
110fn get_now(now: Option<DateTime<Utc>>) -> DateTime<Utc> {
111    now.unwrap_or(Utc::now())
112}
113
114fn expire_at(now: Option<DateTime<Utc>>, interval: Duration) -> DateTime<Utc> {
115    get_now(now) + interval
116}