use std::sync::OnceLock;
use std::time::Duration;
use sqlx::Pool;
use sqlx::pool::PoolOptions;
use zirv_config::read_config;
#[cfg(feature = "mysql")]
pub type Db = sqlx::MySql;
#[cfg(all(feature = "postgres", not(feature = "mysql")))]
pub type Db = sqlx::Postgres;
#[cfg(all(feature = "sqlite", not(any(feature = "mysql", feature = "postgres"))))]
pub type Db = sqlx::Sqlite;
pub type DbPool = Pool<Db>;
static DB_POOL: OnceLock<DbPool> = OnceLock::new();
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PoolSettings {
pub url: String,
pub max_connections: u32,
pub min_connections: u32,
pub acquire_timeout: Duration,
pub idle_timeout: Option<Duration>,
pub max_lifetime: Option<Duration>,
pub test_before_acquire: bool,
}
impl PoolSettings {
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),
)
}
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 {
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),
}
}
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)
}
}
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");
}
pub fn get_db_pool() -> &'static DbPool {
DB_POOL
.get()
.expect("Database pool is not initialized; call init_db_pool first")
}
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());
}
}