zirv_db_sqlx/

db.rs

//! Global connection-pool management for the compile-time selected database backend.
//!
//! Exactly one database backend feature (`mysql`, `postgres`, or `sqlite`) selects
//! the concrete sqlx [`Database`](sqlx::Database) used throughout this module. The
//! pool is stored in a process-wide [`OnceLock`] and shared by reference, so
//! [`get_db_pool`] is effectively free to call from anywhere.

use std::sync::OnceLock;
use std::time::Duration;

use sqlx::Pool;
use sqlx::pool::PoolOptions;
use zirv_config::read_config;

// The backend is selected at compile time. The `not(...)` guards make the
// selection unambiguous even if several backend features are enabled at once
// (the `compile_error!` in `lib.rs` still reports the misconfiguration, but we
// avoid a confusing duplicate-type-definition error on top of it).
#[cfg(feature = "mysql")]
/// The sqlx database backend selected at compile time.
pub type Db = sqlx::MySql;
#[cfg(all(feature = "postgres", not(feature = "mysql")))]
/// The sqlx database backend selected at compile time.
pub type Db = sqlx::Postgres;
#[cfg(all(feature = "sqlite", not(any(feature = "mysql", feature = "postgres"))))]
/// The sqlx database backend selected at compile time.
pub type Db = sqlx::Sqlite;

/// Connection pool for the selected [`Db`] backend.
pub type DbPool = Pool<Db>;

/// The global, one-time-initialized connection pool.
static DB_POOL: OnceLock<DbPool> = OnceLock::new();

/// Connection-pool tuning resolved from the `database.*` configuration namespace.
///
/// Every field except [`url`](Self::url) is optional in configuration and falls
/// back to the same default sqlx itself uses, so an application that only sets
/// `database.url` keeps sqlx's well-chosen defaults. The recognised keys are:
///
/// | Config key | Type | Default | Meaning |
/// |------------|------|---------|---------|
/// | `database.url` | string | (required) | Connection string passed to sqlx. |
/// | `database.max_connections` | u32 | `10` | Hard cap on open connections. |
/// | `database.min_connections` | u32 | `0` | Connections kept warm when idle. |
/// | `database.acquire_timeout_seconds` | u64 | `30` | How long `acquire` waits before erroring. |
/// | `database.idle_timeout_seconds` | u64 | `600` | Reap idle connections after this long. `0` disables. |
/// | `database.max_lifetime_seconds` | u64 | `1800` | Recycle connections older than this. `0` disables. |
/// | `database.test_before_acquire` | bool | `true` | Ping a connection before handing it out. |
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PoolSettings {
    /// Database connection string (e.g. `mysql://user:pass@host/db`).
    pub url: String,
    /// Maximum number of connections the pool will open.
    pub max_connections: u32,
    /// Minimum number of connections kept open while the pool is idle.
    pub min_connections: u32,
    /// Maximum time to wait when acquiring a connection from the pool.
    pub acquire_timeout: Duration,
    /// Close a connection after it has been idle for this long. `None` disables.
    pub idle_timeout: Option<Duration>,
    /// Close a connection once it reaches this age. `None` disables.
    pub max_lifetime: Option<Duration>,
    /// Whether to validate a connection (with a ping) before handing it out.
    pub test_before_acquire: bool,
}

impl PoolSettings {
    /// Reads pool settings from the registered `database` configuration block.
    ///
    /// # Panics
    /// Panics if `database.url` is not configured.
    pub fn from_config() -> Self {
        Self::resolve(
            read_config!("database.url", String)
                .expect("`database.url` must be set in the configuration"),
            read_config!("database.max_connections", u32),
            read_config!("database.min_connections", u32),
            read_config!("database.acquire_timeout_seconds", u64),
            read_config!("database.idle_timeout_seconds", u64),
            read_config!("database.max_lifetime_seconds", u64),
            read_config!("database.test_before_acquire", bool),
        )
    }

    /// Applies defaults to raw, optional configuration values. Kept separate from
    /// [`from_config`](Self::from_config) so the defaulting logic is unit-testable
    /// without a global configuration store.
    fn resolve(
        url: String,
        max_connections: Option<u32>,
        min_connections: Option<u32>,
        acquire_timeout_secs: Option<u64>,
        idle_timeout_secs: Option<u64>,
        max_lifetime_secs: Option<u64>,
        test_before_acquire: Option<bool>,
    ) -> Self {
        // A zero-second duration means "disabled" (no timeout / no max lifetime).
        let opt_duration = |secs: u64| (secs > 0).then(|| Duration::from_secs(secs));
        Self {
            url,
            max_connections: max_connections.unwrap_or(10),
            min_connections: min_connections.unwrap_or(0),
            acquire_timeout: Duration::from_secs(acquire_timeout_secs.unwrap_or(30)),
            idle_timeout: opt_duration(idle_timeout_secs.unwrap_or(600)),
            max_lifetime: opt_duration(max_lifetime_secs.unwrap_or(1800)),
            test_before_acquire: test_before_acquire.unwrap_or(true),
        }
    }

    /// Builds [`PoolOptions`] for the selected backend from these settings, without
    /// opening any connections.
    pub fn to_pool_options(&self) -> PoolOptions<Db> {
        PoolOptions::<Db>::new()
            .max_connections(self.max_connections)
            .min_connections(self.min_connections)
            .acquire_timeout(self.acquire_timeout)
            .idle_timeout(self.idle_timeout)
            .max_lifetime(self.max_lifetime)
            .test_before_acquire(self.test_before_acquire)
    }
}

/// Initializes the global database pool exactly once.
///
/// Reads [`PoolSettings`] from configuration and eagerly opens the pool so that a
/// misconfigured or unreachable database fails fast at startup. Call this early in
/// your application's lifecycle (for example, in `main`).
///
/// # Panics
/// - If `database.url` is not configured.
/// - If the pool fails to connect.
/// - If the pool has already been initialized.
pub async fn init_db_pool() {
    let settings = PoolSettings::from_config();
    let pool = settings
        .to_pool_options()
        .connect(&settings.url)
        .await
        .expect("Failed to create the database connection pool");

    DB_POOL
        .set(pool)
        .expect("Database pool is already initialized; call init_db_pool only once");
}

/// Returns a reference to the global database pool.
///
/// # Panics
/// Panics if [`init_db_pool`] has not been called yet.
pub fn get_db_pool() -> &'static DbPool {
    DB_POOL
        .get()
        .expect("Database pool is not initialized; call init_db_pool first")
}

/// Gracefully closes the global pool, waiting for in-flight connections to be
/// returned and closed. Intended to be called during shutdown. This is a no-op if
/// the pool was never initialized, so it is always safe to call.
pub async fn close_db_pool() {
    if let Some(pool) = DB_POOL.get() {
        pool.close().await;
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn resolve_applies_sqlx_matching_defaults_when_unset() {
        let s = PoolSettings::resolve(
            "mysql://localhost/db".to_string(),
            None,
            None,
            None,
            None,
            None,
            None,
        );
        assert_eq!(s.max_connections, 10);
        assert_eq!(s.min_connections, 0);
        assert_eq!(s.acquire_timeout, Duration::from_secs(30));
        assert_eq!(s.idle_timeout, Some(Duration::from_secs(600)));
        assert_eq!(s.max_lifetime, Some(Duration::from_secs(1800)));
        assert!(s.test_before_acquire);
    }

    #[test]
    fn resolve_honours_explicit_overrides() {
        let s = PoolSettings::resolve(
            "postgres://localhost/db".to_string(),
            Some(50),
            Some(5),
            Some(10),
            Some(120),
            Some(3600),
            Some(false),
        );
        assert_eq!(s.max_connections, 50);
        assert_eq!(s.min_connections, 5);
        assert_eq!(s.acquire_timeout, Duration::from_secs(10));
        assert_eq!(s.idle_timeout, Some(Duration::from_secs(120)));
        assert_eq!(s.max_lifetime, Some(Duration::from_secs(3600)));
        assert!(!s.test_before_acquire);
    }

    #[test]
    fn resolve_treats_zero_durations_as_disabled() {
        let s = PoolSettings::resolve(
            "sqlite::memory:".to_string(),
            None,
            None,
            None,
            Some(0),
            Some(0),
            None,
        );
        assert_eq!(s.idle_timeout, None);
        assert_eq!(s.max_lifetime, None);
    }

    #[test]
    fn to_pool_options_maps_every_field() {
        let s = PoolSettings::resolve(
            "mysql://localhost/db".to_string(),
            Some(42),
            Some(7),
            Some(15),
            Some(0),
            Some(900),
            Some(false),
        );
        let opts = s.to_pool_options();
        assert_eq!(opts.get_max_connections(), 42);
        assert_eq!(opts.get_min_connections(), 7);
        assert_eq!(opts.get_acquire_timeout(), Duration::from_secs(15));
        assert_eq!(opts.get_idle_timeout(), None);
        assert_eq!(opts.get_max_lifetime(), Some(Duration::from_secs(900)));
        assert!(!opts.get_test_before_acquire());
    }
}