resident_utils/postgres/
holder.rs1use 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}