oiseau/
postgres.rs

1#[cfg(not(feature = "redis"))]
2use crate::cache::no_cache::NoCache;
3#[cfg(feature = "redis")]
4use crate::cache::redis::RedisCache;
5
6use crate::{cache::Cache, config::Configuration};
7
8use bb8_postgres::{
9    PostgresConnectionManager,
10    bb8::{Pool, PooledConnection},
11};
12use std::str::FromStr;
13use tokio_postgres::{Config as PgConfig, NoTls, Row, types::ToSql};
14
15pub type Result<T> = std::result::Result<T, tokio_postgres::Error>;
16pub type Connection<'a> = PooledConnection<'a, PostgresConnectionManager<NoTls>>;
17
18#[derive(Clone)]
19pub struct DataManager<T: Clone + Configuration>(
20    pub T,
21    #[cfg(feature = "redis")] pub RedisCache,
22    #[cfg(not(feature = "redis"))] pub NoCache,
23    pub Pool<PostgresConnectionManager<NoTls>>,
24);
25
26impl<T: Clone + Configuration> DataManager<T> {
27    /// Obtain a connection to the staging database.
28    pub async fn connect(&self) -> Result<Connection> {
29        Ok(self.2.get().await.expect("ERROR_OISEAU_PSQL_CON"))
30    }
31
32    /// Create a new [`DataManager`] (and init database).
33    pub async fn new(config: T) -> Result<Self> {
34        let db_config = config.db_config();
35        let con_url = &format!(
36            "postgresql://{}:{}@{}/{}?target_session_attrs=read-write",
37            db_config.user, db_config.password, db_config.url, db_config.name
38        );
39
40        println!("attempting connection on: {con_url}");
41        let manager = PostgresConnectionManager::new(PgConfig::from_str(con_url).unwrap(), NoTls);
42
43        let pool = Pool::builder().max_size(15).build(manager).await.unwrap();
44        Ok(Self(
45            config.clone(),
46            #[cfg(feature = "redis")]
47            RedisCache::new().await,
48            #[cfg(not(feature = "redis"))]
49            NoCache::new().await,
50            pool,
51        ))
52    }
53}
54
55#[cfg(feature = "postgres")]
56#[macro_export]
57macro_rules! get {
58    ($row:ident->$idx:literal($t:ty)) => {
59        $row.get::<usize, Option<$t>>($idx).unwrap()
60    };
61}
62
63pub async fn query_row_helper<T, F>(
64    conn: &Connection<'_>,
65    sql: &str,
66    params: &[&(dyn ToSql + Sync)],
67    f: F,
68) -> Result<T>
69where
70    F: FnOnce(&Row) -> Result<T>,
71{
72    let query = conn.prepare(sql).await.unwrap();
73    let res = conn.query_one(&query, params).await;
74
75    if let Ok(row) = res {
76        Ok(f(&row).unwrap())
77    } else {
78        Err(res.unwrap_err())
79    }
80}
81
82#[macro_export]
83macro_rules! query_row {
84    ($conn:expr, $sql:expr, $params:expr, $f:expr) => {
85        $crate::postgres::query_row_helper($conn, $sql, $params, $f).await
86    };
87}
88
89pub async fn query_rows_helper<T, F>(
90    conn: &Connection<'_>,
91    sql: &str,
92    params: &[&(dyn ToSql + Sync)],
93    mut f: F,
94) -> Result<Vec<T>>
95where
96    F: FnMut(&Row) -> T,
97{
98    let query = conn.prepare(sql).await.unwrap();
99    let res = conn.query(&query, params).await;
100
101    if let Ok(rows) = res {
102        let mut out = Vec::new();
103
104        for row in rows {
105            out.push(f(&row));
106        }
107
108        return Ok(out);
109    } else {
110        Err(res.unwrap_err())
111    }
112}
113
114#[macro_export]
115macro_rules! query_rows {
116    ($conn:expr, $sql:expr, $params:expr, $f:expr) => {
117        $crate::postgres::query_rows_helper($conn, $sql, $params, $f).await
118    };
119}
120
121pub async fn execute_helper(
122    conn: &Connection<'_>,
123    sql: &str,
124    params: &[&(dyn ToSql + Sync)],
125) -> Result<()> {
126    let query = conn.prepare(sql).await.unwrap();
127    conn.execute(&query, params).await?;
128    Ok(())
129}
130
131#[macro_export]
132macro_rules! execute {
133    ($conn:expr, $sql:expr, $params:expr) => {
134        $crate::postgres::execute_helper($conn, $sql, $params).await
135    };
136
137    ($conn:expr, $sql:expr) => {
138        $crate::postgres::execute_helper($conn, $sql, &[]).await
139    };
140}
141
142#[macro_export]
143macro_rules! params {
144    () => {
145        &[]
146    };
147    ($($x:expr),+ $(,)?) => {
148        &[$(&$x),+]
149    };
150}