secra-database 0.1.0

基于 SeaORM 的 Rust 数据库连接和管理库
Documentation
//! 连接池管理服务
//!
//! 提供数据库连接池的监控和统计功能

use sea_orm::{ConnectionTrait, DatabaseConnection, DatabaseBackend};
use serde::{Deserialize, Serialize};
use thiserror::Error;
use tracing::debug;

/// 连接池错误
#[derive(Error, Debug)]
pub enum PoolError {
    #[error("连接池信息获取失败: {0}")]
    InfoFailed(String),
}

/// 连接池统计信息
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PoolStats {
    /// 最大连接数
    pub max_connections: u32,
    /// 最小连接数
    pub min_connections: u32,
    /// 当前活跃连接数(估算)
    pub active_connections: Option<u32>,
    /// 当前空闲连接数(估算)
    pub idle_connections: Option<u32>,
    /// 等待连接的请求数(估算)
    pub waiting_requests: Option<u32>,
}

/// 连接池服务
pub struct ConnectionPoolService;

impl ConnectionPoolService {
    /// 获取连接池统计信息
    ///
    /// # 注意
    /// SeaORM 的连接池统计信息获取有限,部分信息可能无法准确获取。
    /// 对于 PostgreSQL,可以尝试通过系统表查询连接池状态。
    ///
    /// # 参数
    /// * `db` - 数据库连接
    /// * `max_connections` - 最大连接数(从配置中获取)
    /// * `min_connections` - 最小连接数(从配置中获取)
    pub async fn get_pool_stats(
        db: &DatabaseConnection,
        max_connections: Option<u32>,
        min_connections: Option<u32>,
    ) -> Result<PoolStats, PoolError> {
        let backend = db.get_database_backend();
        
        // 尝试从数据库获取实际的连接池统计信息
        let (active, idle, waiting) = match backend {
            DatabaseBackend::Postgres => {
                // PostgreSQL: 尝试从 pg_stat_activity 获取连接信息
                match Self::get_postgres_pool_stats(db).await {
                    Ok(stats) => stats,
                    Err(e) => {
                        debug!("无法获取 PostgreSQL 连接池统计信息: {}", e);
                        (None, None, None)
                    }
                }
            }
            _ => {
                // 其他数据库类型暂不支持详细统计
                (None, None, None)
            }
        };

        Ok(PoolStats {
            max_connections: max_connections.unwrap_or(100),
            min_connections: min_connections.unwrap_or(5),
            active_connections: active,
            idle_connections: idle,
            waiting_requests: waiting,
        })
    }

    /// 获取 PostgreSQL 连接池统计信息
    ///
    /// # 注意
    /// 此方法通过查询 PostgreSQL 系统表获取连接统计信息。
    /// 返回的是数据库级别的连接数,不是应用连接池的连接数。
    /// 应用连接池的连接数需要通过连接池配置获取。
    async fn get_postgres_pool_stats(
        _db: &DatabaseConnection,
    ) -> Result<(Option<u32>, Option<u32>, Option<u32>), PoolError> {
        // 注意:SeaORM 的 DatabaseConnection 不直接暴露查询结果解析 API
        // 要获取详细的连接池统计信息,需要使用 sqlx 连接池或通过其他方式
        // 这里返回 None,表示无法获取详细信息
        // 
        // 如果需要获取 PostgreSQL 连接统计,可以使用以下 SQL:
        // SELECT count(*) FROM pg_stat_activity WHERE state = 'active' AND datname = current_database();
        // SELECT count(*) FROM pg_stat_activity WHERE state = 'idle' AND datname = current_database();
        
        debug!("PostgreSQL 连接池统计信息需要通过 sqlx 连接池获取");
        Ok((None, None, None))
    }

    /// 测试连接池健康状态
    pub async fn health_check(db: &DatabaseConnection) -> Result<bool, PoolError> {
        // 尝试执行一个简单的查询来测试连接池
        db.execute_unprepared("SELECT 1")
            .await
            .map_err(|e| PoolError::InfoFailed(format!("连接池健康检查失败: {}", e)))?;

        Ok(true)
    }
}

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

    #[test]
    fn test_pool_stats_serialization() {
        let stats = PoolStats {
            max_connections: 100,
            min_connections: 5,
            active_connections: Some(10),
            idle_connections: Some(5),
            waiting_requests: Some(0),
        };

        let json = serde_json::to_string(&stats).unwrap();
        assert!(json.contains("max_connections"));
    }
}