sz_orm_sqlx/
unified_pool.rs

1//! UnifiedPool — 统一连接池抽象（v2.2.0 A-3）
2//!
3//! 包装 `sz_orm_core::Pool`（完整连接池）+ `AnyBackend`，提供 5 后端透明的统一类型。
4//! 供 sz-rust AppState 持有单一类型 `Arc<UnifiedPool>`，业务代码无需感知后端类型。
5//!
6//! # 设计
7//!
8//! - `UnifiedPool` 是 `Pool` 的 newtype 包装，所有方法委托 `Pool`（零能力丢失）
9//! - `from_pool` 提供零成本迁移路径：sz-rust 从 `Arc<Pool>` 迁移到 `Arc<UnifiedPool>`
10//! - `connect`/`connect_with_config` 根据 DSN 自动识别后端并创建完整连接池
11//!
12//! # 用法
13//!
14//! ```ignore
15//! use sz_orm_sqlx::UnifiedPool;
16//!
17//! // 从 DSN 自动识别后端，创建完整连接池
18//! let pool = UnifiedPool::connect("mysql://root:pass@127.0.0.1/db").await?;
19//! let mut conn = pool.acquire().await?;
20//!
21//! // 运行时切换后端
22//! let pg_pool = UnifiedPool::connect("postgres://user:pass@127.0.0.1/db").await?;
23//! let d = pg_pool.dialect(); // 自动返回 PostgreSqlDialect
24//! ```
25
26use std::sync::Arc;
27use sz_orm_core::{
28    ConnectionFactory, DbError, Dialect, Pool, PoolConfig, PoolConfigBuilder, PoolError,
29    PoolStatus, PooledConnection,
30};
31
32use crate::any::{
33    MySqlPoolHandle, PgPoolHandle, SqlitePoolHandle, SqlxMySqlConnectionFactory,
34    SqlxPgConnectionFactory, SqlxSqliteConnectionFactory,
35};
36use crate::any_driver::AnyBackend;
37
38#[cfg(feature = "oracle")]
39use sz_orm_oracle::{OracleConnectionFactory, OraclePoolHandle};
40
41#[cfg(feature = "mssql")]
42use sz_orm_mssql::{MssqlConnectionFactory, MssqlPoolHandle};
43
44/// 统一连接池：包装 `Pool` + `AnyBackend`，5 后端透明切换（v2.2.0 新增）
45///
46/// 供 sz-rust AppState 持有 `Arc<UnifiedPool>`，业务代码无需感知后端类型。
47/// 所有方法委托内部 `Pool`，零能力丢失。
48pub struct UnifiedPool {
49    backend: AnyBackend,
50    pool: Pool,
51}
52
53impl UnifiedPool {
54    /// 连接数据库，根据 DSN 自动识别后端，使用默认 PoolConfig
55    ///
56    /// 默认配置：max_size=10, timeout=30s
57    pub async fn connect(dsn: &str) -> Result<Self, DbError> {
58        let config = PoolConfigBuilder::new()
59            .build()
60            .map_err(DbError::PoolError)?;
61        Self::connect_with_config(dsn, config).await
62    }
63
64    /// 连接数据库，根据 DSN 自动识别后端，使用自定义 PoolConfig
65    pub async fn connect_with_config(dsn: &str, config: PoolConfig) -> Result<Self, DbError> {
66        let backend = AnyBackend::from_dsn(dsn)?;
67        let factory: Arc<dyn ConnectionFactory> = match backend {
68            AnyBackend::MySql => {
69                let handle = Arc::new(MySqlPoolHandle::connect(dsn).await?);
70                Arc::new(SqlxMySqlConnectionFactory::new(handle))
71            }
72            AnyBackend::Postgres => {
73                let handle = Arc::new(PgPoolHandle::connect(dsn).await?);
74                Arc::new(SqlxPgConnectionFactory::new(handle))
75            }
76            AnyBackend::Sqlite => {
77                let handle = Arc::new(SqlitePoolHandle::connect(dsn).await?);
78                Arc::new(SqlxSqliteConnectionFactory::new(handle))
79            }
80            AnyBackend::Oracle => {
81                #[cfg(feature = "oracle")]
82                {
83                    let (username, password, connect_string) =
84                        crate::any_driver::parse_oracle_dsn(dsn)?;
85                    let handle = Arc::new(OraclePoolHandle::connect(
86                        &username,
87                        &password,
88                        &connect_string,
89                    )?);
90                    Arc::new(OracleConnectionFactory::new(handle))
91                }
92                #[cfg(not(feature = "oracle"))]
93                {
94                    return Err(DbError::ConnectionRefused(
95                        "Oracle 后端未启用，请在 Cargo.toml 中添加 features = [\"oracle\"]"
96                            .to_string(),
97                    ));
98                }
99            }
100            AnyBackend::Mssql => {
101                #[cfg(feature = "mssql")]
102                {
103                    let ado_string = crate::any_driver::parse_mssql_dsn(dsn)?;
104                    let handle = Arc::new(MssqlPoolHandle::connect(&ado_string).await?);
105                    Arc::new(MssqlConnectionFactory::new(handle))
106                }
107                #[cfg(not(feature = "mssql"))]
108                {
109                    return Err(DbError::ConnectionRefused(
110                        "MSSQL 后端未启用，请在 Cargo.toml 中添加 features = [\"mssql\"]"
111                            .to_string(),
112                    ));
113                }
114            }
115        };
116        let pool = Pool::new(config, factory).map_err(DbError::PoolError)?;
117        Ok(Self { backend, pool })
118    }
119
120    /// 从已有的 Pool 构造 UnifiedPool（零成本迁移）
121    ///
122    /// 供 sz-rust 从 `Arc<Pool>` 迁移到 `Arc<UnifiedPool>`：
123    /// ```ignore
124    /// let unified = UnifiedPool::from_pool(existing_pool, AnyBackend::MySql);
125    /// ```
126    pub fn from_pool(pool: Pool, backend: AnyBackend) -> Self {
127        Self { backend, pool }
128    }
129
130    /// 获取后端类型
131    #[inline]
132    pub fn backend(&self) -> AnyBackend {
133        self.backend
134    }
135
136    /// 返回对应后端的 Dialect 实例
137    #[inline]
138    pub fn dialect(&self) -> Box<dyn Dialect> {
139        self.backend.dialect()
140    }
141
142    /// 获取连接（委托 Pool::acquire）
143    #[inline]
144    pub async fn acquire(&self) -> Result<PooledConnection, PoolError> {
145        self.pool.acquire().await
146    }
147
148    /// 调整连接池大小（委托 Pool::resize）
149    #[inline]
150    pub fn resize(&self, new_max: usize) {
151        self.pool.resize(new_max);
152    }
153
154    /// 关闭所有连接（委托 Pool::close_all）
155    #[inline]
156    pub async fn close_all(&self) {
157        self.pool.close_all().await;
158    }
159
160    /// 获取连接池状态（委托 Pool::status）
161    #[inline]
162    pub async fn status(&self) -> PoolStatus {
163        self.pool.status().await
164    }
165
166    /// 预热连接池（委托 Pool::prewarm）
167    #[inline]
168    pub async fn prewarm(&self) {
169        self.pool.prewarm().await;
170    }
171
172    /// v3.2.0：渐进式分批预热（委托 Pool::progressive_prewarm）
173    #[cfg(feature = "auto-prewarm")]
174    pub async fn progressive_prewarm(
175        &self,
176        batch_size: u32,
177        interval: std::time::Duration,
178        total_timeout: std::time::Duration,
179        progress: &sz_orm_core::prewarm::PrewarmProgress,
180    ) {
181        self.pool
182            .progressive_prewarm(batch_size, interval, total_timeout, progress)
183            .await;
184    }
185}
186
187/// v3.2.0：多池注册表 — 统一管理多个后端的 UnifiedPool
188#[cfg(feature = "auto-prewarm")]
189pub struct MultiPoolRegistry {
190    pools: Vec<(String, UnifiedPool)>,
191}
192
193#[cfg(feature = "auto-prewarm")]
194impl MultiPoolRegistry {
195    pub fn new() -> Self {
196        Self { pools: Vec::new() }
197    }
198
199    pub fn register(&mut self, name: impl Into<String>, pool: UnifiedPool) {
200        self.pools.push((name.into(), pool));
201    }
202
203    /// 并行预热所有注册的池
204    pub async fn unified_prewarm_all(&self) -> sz_orm_core::prewarm::PrewarmSummary {
205        use std::time::Instant;
206        use sz_orm_core::prewarm::{BackendPrewarmResult, PrewarmSummary};
207
208        let mut summary = PrewarmSummary::new();
209        for (name, pool) in &self.pools {
210            let start = Instant::now();
211            let min_idle = pool.pool.config().min_idle;
212            let progress = sz_orm_core::prewarm::PrewarmProgress::new(min_idle);
213            pool.prewarm().await;
214            let snap = progress.snapshot();
215            summary.add(BackendPrewarmResult {
216                backend: name.clone(),
217                warmed: snap.warmed,
218                failed: snap.failed,
219                elapsed: start.elapsed(),
220                errors: vec![],
221            });
222        }
223        summary
224    }
225}
226
227#[cfg(feature = "auto-prewarm")]
228impl Default for MultiPoolRegistry {
229    fn default() -> Self {
230        Self::new()
231    }
232}
233
234impl std::fmt::Debug for UnifiedPool {
235    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
236        f.debug_struct("UnifiedPool")
237            .field("backend", &self.backend)
238            .finish_non_exhaustive()
239    }
240}
241
242#[cfg(test)]
243mod tests {
244    use super::*;
245
246    #[tokio::test]
247    async fn test_unified_pool_sqlite_connect() {
248        let pool = UnifiedPool::connect("sqlite::memory:").await.unwrap();
249        assert_eq!(pool.backend(), AnyBackend::Sqlite);
250
251        let mut conn = pool.acquire().await.unwrap();
252        conn.execute("CREATE TABLE t (id INTEGER PRIMARY KEY)")
253            .await
254            .unwrap();
255        conn.execute("INSERT INTO t (id) VALUES (1)").await.unwrap();
256        let rows = conn.query("SELECT * FROM t").await.unwrap();
257        assert_eq!(rows.len(), 1);
258    }
259
260    #[tokio::test]
261    async fn test_unified_pool_dialect() {
262        let pool = UnifiedPool::connect("sqlite::memory:").await.unwrap();
263        let d = pool.dialect();
264        assert_eq!(d.db_type(), sz_orm_core::DbType::Sqlite);
265    }
266
267    #[tokio::test]
268    async fn test_unified_pool_from_pool() {
269        let handle = Arc::new(SqlitePoolHandle::connect("sqlite::memory:").await.unwrap());
270        let factory = Arc::new(SqlxSqliteConnectionFactory::new(handle));
271        let config = PoolConfigBuilder::new().build().unwrap();
272        let pool = Pool::new(config, factory).unwrap();
273
274        let unified = UnifiedPool::from_pool(pool, AnyBackend::Sqlite);
275        assert_eq!(unified.backend(), AnyBackend::Sqlite);
276
277        let mut conn = unified.acquire().await.unwrap();
278        conn.execute("SELECT 1").await.unwrap();
279    }
280
281    #[tokio::test]
282    async fn test_unified_pool_resize_and_close() {
283        let pool = UnifiedPool::connect("sqlite::memory:").await.unwrap();
284        pool.resize(20);
285        let status = pool.status().await;
286        assert_eq!(status.max, 20);
287        pool.close_all().await;
288    }
289
290    #[tokio::test]
291    async fn test_unified_pool_invalid_dsn() {
292        let result = UnifiedPool::connect("invalid://dsn").await;
293        assert!(result.is_err());
294    }
295}
sz_orm_sqlx/unified_pool.rs

sz_orm_sqlx/
unified_pool.rs
