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}