qm_pg/
db.rs

1use crate::config::Config;
2use sqlx::postgres::PgPoolOptions;
3use sqlx::PgPool;
4use std::sync::Arc;
5use std::time::Duration;
6
7struct Inner {
8    pool: PgPool,
9}
10
11#[derive(Clone)]
12pub struct DB {
13    inner: Arc<Inner>,
14}
15
16impl DB {
17    pub async fn new(app_name: &str, cfg: &Config) -> anyhow::Result<Self> {
18        if let Some(database) = cfg.database() {
19            tracing::info!(
20                "'{app_name}' -> connects to postgresql '{database}' with {} max_connections",
21                cfg.max_connections(),
22            );
23        } else {
24            tracing::info!(
25                "'{app_name}' -> connects to postgresql with {} max_connections",
26                cfg.max_connections(),
27            );
28        }
29        let pool = PgPoolOptions::new()
30            .min_connections(cfg.min_connections())
31            .max_connections(cfg.max_connections())
32            .acquire_timeout(Duration::from_secs(cfg.acquire_timeout()))
33            .idle_timeout(Duration::from_secs(cfg.idle_timeout()))
34            .max_lifetime(Duration::from_secs(cfg.max_lifetime()))
35            .connect(cfg.address())
36            .await?;
37        Ok(Self {
38            inner: Arc::new(Inner { pool }),
39        })
40    }
41
42    pub async fn new_root(app_name: &str, cfg: &Config) -> anyhow::Result<Self> {
43        if let Some(database) = cfg.root_database() {
44            tracing::info!(
45                "'{app_name}' -> connects as root to postgresql '{database}' with {} max_connections",
46                cfg.max_connections(),
47            );
48        } else {
49            tracing::info!(
50                "'{app_name}' -> connects as root to postgresql with {} max_connections",
51                cfg.max_connections(),
52            );
53        }
54        let pool = PgPoolOptions::new()
55            .min_connections(1)
56            .max_connections(2)
57            .acquire_timeout(Duration::from_secs(cfg.acquire_timeout()))
58            .connect(cfg.root_address())
59            .await?;
60        Ok(Self {
61            inner: Arc::new(Inner { pool }),
62        })
63    }
64
65    pub fn database_connection(&self) -> sea_orm::DatabaseConnection {
66        sea_orm::SqlxPostgresConnector::from_sqlx_postgres_pool(self.inner.pool.clone())
67    }
68
69    pub fn pool(&self) -> &PgPool {
70        &self.inner.pool
71    }
72}