resident_utils/postgres/
holder.rs

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