Skip to main content

baichun_framework_db/
pool.rs

1//! 数据库连接池模块
2//!
3//! 本模块提供了数据库连接池的管理功能,包括:
4//! - 连接池的创建和初始化
5//! - 全局连接池实例的管理
6//! - 连接池状态监控
7//! - 事务管理
8
9use crate::{config::DatabaseConfig, error::Result};
10use sqlx::{MySql, Pool};
11use std::sync::Arc;
12use tokio::sync::OnceCell;
13
14/// 全局数据库连接池实例
15pub static DB_POOL: OnceCell<Arc<DbPool>> = OnceCell::const_new();
16
17/// 数据库连接池包装器
18///
19/// 提供了对 SQLx 连接池的高级封装,支持连接管理、查询执行和事务操作。
20///
21/// # 示例
22///
23/// ```rust
24/// use baichun_framework_db::{DatabaseConfig, DbPool};
25///
26/// async fn example() -> Result<()> {
27///     // 创建连接池
28///     let config = DatabaseConfig::new()
29///         .url("mysql://user:pass@localhost/db_name")
30///         .max_connections(10);
31///     let pool = DbPool::new(&config).await?;
32///
33///     // 执行查询
34///     pool.execute("CREATE TABLE users (id INT PRIMARY KEY)").await?;
35///
36///     // 开始事务
37///     let tx = pool.begin().await?;
38///     // ... 执行事务操作 ...
39///
40///     Ok(())
41/// }
42/// ```
43#[derive(Clone)]
44pub struct DbPool {
45    pool: Pool<MySql>,
46}
47
48impl DbPool {
49    /// 创建新的数据库连接池
50    ///
51    /// 根据提供的配置创建一个新的连接池实例。
52    pub async fn new(config: &DatabaseConfig) -> Result<Self> {
53        let pool = config
54            .into_pool_options()
55            .connect(&config.url)
56            .await
57            .map_err(|e| crate::error::Error::Pool(e.to_string()))?;
58
59        Ok(Self { pool })
60    }
61
62    /// 获取底层的 SQLx 连接池
63    ///
64    /// 用于需要直接访问 SQLx 功能的场景。
65    pub fn pool(&self) -> &Pool<MySql> {
66        &self.pool
67    }
68
69    /// 开始新的事务
70    ///
71    /// 返回一个事务对象,可以用于执行事务操作。
72    pub async fn begin(&self) -> Result<sqlx::Transaction<'_, MySql>> {
73        self.pool
74            .begin()
75            .await
76            .map_err(|e| crate::error::Error::Transaction(e.to_string()))
77    }
78
79    /// 执行不返回结果的查询
80    ///
81    /// 适用于 DDL 语句或不需要返回结果的 DML 语句。
82    pub async fn execute(&self, query: &str) -> Result<sqlx::mysql::MySqlQueryResult> {
83        sqlx::query(query)
84            .execute(&self.pool)
85            .await
86            .map_err(|e| crate::error::Error::Query(e.to_string()))
87    }
88
89    /// 获取数据库连接池指标
90    ///
91    /// 返回当前连接池的状态信息,包括总连接数、活动连接数等。
92    pub fn metrics(&self) -> DbMetrics {
93        DbMetrics {
94            connections: 0, // SQLx doesn't provide detailed metrics
95            idle_connections: 0,
96            active_connections: 0,
97        }
98    }
99}
100
101/// 数据库连接池指标
102///
103/// 提供了连接池的实时状态信息。
104#[derive(Debug, Clone)]
105pub struct DbMetrics {
106    /// 总连接数
107    pub connections: u32,
108    /// 空闲连接数
109    pub idle_connections: u32,
110    /// 活动连接数
111    pub active_connections: u32,
112}
113
114/// 获取全局数据库连接池实例
115///
116/// 如果连接池尚未初始化,将会 panic。
117///
118/// # 示例
119///
120/// ```rust
121/// use baichun_framework_db::get_pool;
122///
123/// async fn example() {
124///     let pool = get_pool().await;
125///     // 使用连接池...
126/// }
127/// ```
128pub async fn get_pool() -> Arc<DbPool> {
129    DB_POOL
130        .get()
131        .expect("Database pool not initialized")
132        .clone()
133}
134
135/// 初始化全局数据库连接池
136///
137/// 使用提供的配置初始化连接池。如果连接池已经初始化,将返回错误。
138///
139/// # 示例
140///
141/// ```rust
142/// use baichun_framework_db::{DatabaseConfig, init};
143///
144/// async fn example() -> Result<()> {
145///     let config = DatabaseConfig::new()
146///         .url("mysql://user:pass@localhost/db_name")
147///         .max_connections(10);
148///     
149///     init(config).await?;
150///     Ok(())
151/// }
152/// ```
153pub async fn init(config: DatabaseConfig) -> Result<Arc<DbPool>> {
154    let pool = DbPool::new(&config).await?;
155    let pool = Arc::new(pool);
156
157    if DB_POOL.set(pool.clone()).is_err() {
158        return Err(crate::error::Error::Pool(
159            "Database pool already initialized".to_string(),
160        ));
161    }
162
163    Ok(pool)
164}