use crate::{
config::{DatabaseConfig, DatabaseType, PoolOptions},
error::Result,
};
use std::sync::Arc;
use tokio::sync::OnceCell;
#[cfg(feature = "mysql")]
use sqlx::{MySql, Pool as MySqlPool};
#[cfg(feature = "postgres")]
use sqlx::{Pool as PgPool, Postgres};
#[cfg(feature = "sqlite")]
use sqlx::{Pool as SqlitePool, Sqlite};
pub static DB_POOL: OnceCell<Arc<DbPool>> = OnceCell::const_new();
#[derive(Clone)]
pub struct DbPool {
#[cfg(feature = "mysql")]
mysql_pool: Option<MySqlPool<MySql>>,
#[cfg(feature = "postgres")]
pg_pool: Option<PgPool<Postgres>>,
#[cfg(feature = "sqlite")]
sqlite_pool: Option<SqlitePool<Sqlite>>,
db_type: DatabaseType,
}
impl DbPool {
pub async fn new(config: &DatabaseConfig) -> Result<Self> {
let db_type = config.get_database_type();
let pool_options = config.into_pool_options();
match pool_options {
#[cfg(feature = "mysql")]
PoolOptions::MySQL(options) => {
let pool = options
.connect(&config.url)
.await
.map_err(|e| crate::error::Error::Pool(e.to_string()))?;
Ok(Self {
mysql_pool: Some(pool),
#[cfg(feature = "postgres")]
pg_pool: None,
#[cfg(feature = "sqlite")]
sqlite_pool: None,
db_type,
})
}
#[cfg(feature = "postgres")]
PoolOptions::PostgreSQL(options) => {
let pool = options
.connect(&config.url)
.await
.map_err(|e| crate::error::Error::Pool(e.to_string()))?;
Ok(Self {
#[cfg(feature = "mysql")]
mysql_pool: None,
pg_pool: Some(pool),
#[cfg(feature = "sqlite")]
sqlite_pool: None,
db_type,
})
}
#[cfg(feature = "sqlite")]
PoolOptions::SQLite(options) => {
let pool = options
.connect(&config.url)
.await
.map_err(|e| crate::error::Error::Pool(e.to_string()))?;
Ok(Self {
#[cfg(feature = "mysql")]
mysql_pool: None,
#[cfg(feature = "postgres")]
pg_pool: None,
sqlite_pool: Some(pool),
db_type,
})
}
}
}
pub fn pool(&self) -> Result<&dyn sqlx::Pool> {
match self.db_type {
#[cfg(feature = "mysql")]
DatabaseType::MySQL => self
.mysql_pool
.as_ref()
.ok_or_else(|| crate::error::Error::Pool("MySQL pool not initialized".to_string())),
#[cfg(feature = "postgres")]
DatabaseType::PostgreSQL => self.pg_pool.as_ref().ok_or_else(|| {
crate::error::Error::Pool("PostgreSQL pool not initialized".to_string())
}),
#[cfg(feature = "sqlite")]
DatabaseType::SQLite => self.sqlite_pool.as_ref().ok_or_else(|| {
crate::error::Error::Pool("SQLite pool not initialized".to_string())
}),
}
}
pub async fn begin(&self) -> Result<sqlx::Transaction<'_, sqlx::Any>> {
self.pool()?
.begin()
.await
.map_err(|e| crate::error::Error::Transaction(e.to_string()))
}
pub async fn execute(&self, query: &str) -> Result<sqlx::any::AnyQueryResult> {
sqlx::query(query)
.execute(self.pool()?)
.await
.map_err(|e| crate::error::Error::Query(e.to_string()))
}
pub fn metrics(&self) -> DbMetrics {
DbMetrics {
connections: 0, idle_connections: 0,
active_connections: 0,
}
}
}
#[derive(Debug, Clone)]
pub struct DbMetrics {
pub connections: u32,
pub idle_connections: u32,
pub active_connections: u32,
}
pub async fn get_pool() -> Arc<DbPool> {
DB_POOL
.get()
.expect("Database pool not initialized")
.clone()
}
pub async fn init(config: DatabaseConfig) -> Result<Arc<DbPool>> {
let pool = DbPool::new(&config).await?;
let pool = Arc::new(pool);
if DB_POOL.set(pool.clone()).is_err() {
return Err(crate::error::Error::Pool(
"Database pool already initialized".to_string(),
));
}
Ok(pool)
}