Skip to main content

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/// PostgreSQL database connection wrapper.
12///
13/// Manages connection pool and provides access to sqlx and SeaORM connections.
14#[derive(Clone)]
15pub struct DB {
16    inner: Arc<Inner>,
17}
18
19impl DB {
20    /// Creates a new PostgreSQL database connection with the given config.
21    pub async fn new(app_name: &str, cfg: &Config) -> anyhow::Result<Self> {
22        if let Some(database) = cfg.database() {
23            tracing::info!(
24                "'{app_name}' -> connects to postgresql '{database}' with {} max_connections",
25                cfg.max_connections(),
26            );
27        } else {
28            tracing::info!(
29                "'{app_name}' -> connects to postgresql with {} max_connections",
30                cfg.max_connections(),
31            );
32        }
33        let pool = PgPoolOptions::new()
34            .min_connections(cfg.min_connections())
35            .max_connections(cfg.max_connections())
36            .acquire_timeout(Duration::from_secs(cfg.acquire_timeout()))
37            .idle_timeout(Duration::from_secs(cfg.idle_timeout()))
38            .max_lifetime(Duration::from_secs(cfg.max_lifetime()))
39            .connect(cfg.address())
40            .await?;
41        Ok(Self {
42            inner: Arc::new(Inner { pool }),
43        })
44    }
45
46    /// Creates a new PostgreSQL database connection with root credentials.
47    pub async fn new_root(app_name: &str, cfg: &Config) -> anyhow::Result<Self> {
48        if let Some(database) = cfg.root_database() {
49            tracing::info!(
50                "'{app_name}' -> connects as root to postgresql '{database}' with {} max_connections",
51                cfg.max_connections(),
52            );
53        } else {
54            tracing::info!(
55                "'{app_name}' -> connects as root to postgresql with {} max_connections",
56                cfg.max_connections(),
57            );
58        }
59        let pool = PgPoolOptions::new()
60            .min_connections(1)
61            .max_connections(2)
62            .acquire_timeout(Duration::from_secs(cfg.acquire_timeout()))
63            .connect(cfg.root_address())
64            .await?;
65        Ok(Self {
66            inner: Arc::new(Inner { pool }),
67        })
68    }
69
70    /// Returns a SeaORM database connection.
71    pub fn database_connection(&self) -> sea_orm::DatabaseConnection {
72        sea_orm::SqlxPostgresConnector::from_sqlx_postgres_pool(self.inner.pool.clone())
73    }
74
75    /// Returns a reference to the sqlx connection pool.
76    pub fn pool(&self) -> &PgPool {
77        &self.inner.pool
78    }
79}