rframe 0.1.1

A lightweight framework implemented in Rust
Documentation
//! 数据库连接池模块
//!
//! 本模块提供了数据库连接池的管理功能,包括:
//! - 连接池的创建和初始化
//! - 全局连接池实例的管理
//! - 连接池状态监控
//! - 事务管理

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();

/// 数据库连接池包装器
///
/// 提供了对 SQLx 连接池的高级封装,支持连接管理、查询执行和事务操作。
///
/// # 示例
///
/// ```rust
/// use baichun_framework_db::{DatabaseConfig, DbPool};
///
/// async fn example() -> Result<()> {
///     // 创建连接池
///     let config = DatabaseConfig::new()
///         .url("mysql://user:pass@localhost/db_name")
///         .max_connections(10);
///     let pool = DbPool::new(&config).await?;
///
///     // 执行查询
///     pool.execute("CREATE TABLE users (id INT PRIMARY KEY)").await?;
///
///     // 开始事务
///     let tx = pool.begin().await?;
///     // ... 执行事务操作 ...
///
///     Ok(())
/// }
/// ```
#[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,
                })
            }
        }
    }

    /// 获取底层的 SQLx 连接池
    ///
    /// 用于需要直接访问 SQLx 功能的场景。
    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()))
    }

    /// 执行不返回结果的查询
    ///
    /// 适用于 DDL 语句或不需要返回结果的 DML 语句。
    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, // SQLx doesn't provide detailed metrics
            idle_connections: 0,
            active_connections: 0,
        }
    }
}

/// 数据库连接池指标
///
/// 提供了连接池的实时状态信息。
#[derive(Debug, Clone)]
pub struct DbMetrics {
    /// 总连接数
    pub connections: u32,
    /// 空闲连接数
    pub idle_connections: u32,
    /// 活动连接数
    pub active_connections: u32,
}

/// 获取全局数据库连接池实例
///
/// 如果连接池尚未初始化,将会 panic。
///
/// # 示例
///
/// ```rust
/// use baichun_framework_db::get_pool;
///
/// async fn example() {
///     let pool = get_pool().await;
///     // 使用连接池...
/// }
/// ```
pub async fn get_pool() -> Arc<DbPool> {
    DB_POOL
        .get()
        .expect("Database pool not initialized")
        .clone()
}

/// 初始化全局数据库连接池
///
/// 使用提供的配置初始化连接池。如果连接池已经初始化,将返回错误。
///
/// # 示例
///
/// ```rust
/// use baichun_framework_db::{DatabaseConfig, init};
///
/// async fn example() -> Result<()> {
///     let config = DatabaseConfig::new()
///         .url("mysql://user:pass@localhost/db_name")
///         .max_connections(10);
///     
///     init(config).await?;
///     Ok(())
/// }
/// ```
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)
}