ecat_data_sqlx/
pool.rs

1// Copyright (c) 2026 erik <erik@erik.xyz> — https://erik.xyz
2
3use crate::config::PoolParams;
4use ecat_data::Dialect;
5use sqlx::mysql::{MySqlPool, MySqlPoolOptions};
6use sqlx::postgres::{PgPool, PgPoolOptions};
7use sqlx::sqlite::{SqlitePool, SqlitePoolOptions};
8
9/// 三种原生池。方言由变体本身承载 —— 不再解析 URL 猜测。
10#[derive(Debug)]
11pub enum Pool {
12    Pg(PgPool),
13    My(MySqlPool),
14    Sq(SqlitePool),
15}
16
17/// 三种驱动的连接 guard 的聚合类型。Drop 即归还连接。
18/// `PoolConnection<DB>` 自身持有 `Arc`，无需生命周期参数。
19pub enum PoolGuard {
20    Pg(sqlx::pool::PoolConnection<sqlx::Postgres>),
21    My(sqlx::pool::PoolConnection<sqlx::MySql>),
22    Sq(sqlx::pool::PoolConnection<sqlx::Sqlite>),
23}
24
25impl Pool {
26    /// 按 URL scheme 建立对应的原生池。
27    pub async fn connect(url: &str, params: &PoolParams) -> Result<Self, sqlx::Error> {
28        match Dialect::from_url(url) {
29            Dialect::Postgres => Ok(Self::Pg(
30                common_pg(PgPoolOptions::new(), params).connect(url).await?,
31            )),
32            Dialect::MySql => Ok(Self::My(
33                common_mysql(MySqlPoolOptions::new(), params)
34                    .connect(url)
35                    .await?,
36            )),
37            Dialect::Sqlite => Ok(Self::Sq(
38                common_sqlite(SqlitePoolOptions::new(), params)
39                    .connect(url)
40                    .await?,
41            )),
42            // 报 scheme 原文而非 Dialect：`Dialect::Standard` 是「无法识别」的
43            // 回退值，把它打进错误信息等于丢掉唯一的诊断线索。
44            Dialect::Standard => Err(unsupported(url)),
45            // Mssql 走 tiberius 后端（ecat-data-mssql），不是 sqlx
46            Dialect::Mssql => Err(sqlx::Error::Configuration(Box::new(std::io::Error::new(
47                std::io::ErrorKind::InvalidInput,
48                "mssql:// is served by ecat-data-mssql, not the sqlx backend",
49            )))),
50        }
51    }
52
53    pub fn dialect(&self) -> Dialect {
54        match self {
55            Self::Pg(_) => Dialect::Postgres,
56            Self::My(_) => Dialect::MySql,
57            Self::Sq(_) => Dialect::Sqlite,
58        }
59    }
60
61    /// 池内已建立的连接总数。
62    pub fn size(&self) -> u32 {
63        match self {
64            Self::Pg(p) => p.size(),
65            Self::My(p) => p.size(),
66            Self::Sq(p) => p.size(),
67        }
68    }
69
70    /// 池内空闲连接数。
71    pub fn idle(&self) -> u32 {
72        match self {
73            Self::Pg(p) => p.num_idle() as u32,
74            Self::My(p) => p.num_idle() as u32,
75            Self::Sq(p) => p.num_idle() as u32,
76        }
77    }
78
79    /// 从池中取一条连接（预热用）。归还由 guard 的 Drop 完成。
80    pub async fn acquire(&self) -> Result<PoolGuard, sqlx::Error> {
81        Ok(match self {
82            Self::Pg(p) => PoolGuard::Pg(p.acquire().await?),
83            Self::My(p) => PoolGuard::My(p.acquire().await?),
84            Self::Sq(p) => PoolGuard::Sq(p.acquire().await?),
85        })
86    }
87}
88
89/// 无法识别的 scheme：错误信息要点名出问题的 scheme。
90fn unsupported(url: &str) -> sqlx::Error {
91    let scheme = url
92        .split("://")
93        .next()
94        .unwrap_or(url)
95        .split(':')
96        .next()
97        .unwrap_or(url);
98    sqlx::Error::Configuration(Box::new(std::io::Error::new(
99        std::io::ErrorKind::InvalidInput,
100        format!("unsupported database scheme for sqlx backend: {scheme}"),
101    )))
102}
103
104// 三个 common_* 把 PoolParams 落到各驱动的 PoolOptions 上。
105// 注意：三者的 Connection 类型不同，after_connect 闭包无法合并，各写一份。
106
107fn common_pg(opts: PgPoolOptions, p: &PoolParams) -> PgPoolOptions {
108    let opts = opts
109        .max_connections(p.max_connections)
110        .min_connections(p.min_connections)
111        .acquire_timeout(p.acquire_timeout)
112        .idle_timeout(Some(p.idle_timeout))
113        .max_lifetime(Some(p.max_lifetime))
114        .test_before_acquire(p.test_before_acquire);
115    if p.session_init.is_empty() {
116        return opts;
117    }
118    let stmts = p.session_init.clone();
119    opts.after_connect(move |conn, _meta| {
120        let stmts = stmts.clone();
121        Box::pin(async move {
122            for s in &stmts {
123                sqlx::query(s).execute(&mut *conn).await?;
124            }
125            Ok(())
126        })
127    })
128}
129
130fn common_mysql(opts: MySqlPoolOptions, p: &PoolParams) -> MySqlPoolOptions {
131    let opts = opts
132        .max_connections(p.max_connections)
133        .min_connections(p.min_connections)
134        .acquire_timeout(p.acquire_timeout)
135        .idle_timeout(Some(p.idle_timeout))
136        .max_lifetime(Some(p.max_lifetime))
137        .test_before_acquire(p.test_before_acquire);
138    if p.session_init.is_empty() {
139        return opts;
140    }
141    let stmts = p.session_init.clone();
142    opts.after_connect(move |conn, _meta| {
143        let stmts = stmts.clone();
144        Box::pin(async move {
145            for s in &stmts {
146                sqlx::query(s).execute(&mut *conn).await?;
147            }
148            Ok(())
149        })
150    })
151}
152
153fn common_sqlite(opts: SqlitePoolOptions, p: &PoolParams) -> SqlitePoolOptions {
154    let opts = opts
155        .max_connections(p.max_connections)
156        .min_connections(p.min_connections)
157        .acquire_timeout(p.acquire_timeout)
158        .idle_timeout(Some(p.idle_timeout))
159        .max_lifetime(Some(p.max_lifetime))
160        .test_before_acquire(p.test_before_acquire);
161    if p.session_init.is_empty() {
162        return opts;
163    }
164    let stmts = p.session_init.clone();
165    opts.after_connect(move |conn, _meta| {
166        let stmts = stmts.clone();
167        Box::pin(async move {
168            for s in &stmts {
169                sqlx::query(s).execute(&mut *conn).await?;
170            }
171            Ok(())
172        })
173    })
174}
175
176#[cfg(test)]
177mod tests {
178    use super::*;
179    use ecat_data::Dialect;
180
181    #[tokio::test]
182    async fn sqlite_url_builds_sqlite_pool() {
183        let pool = Pool::connect("sqlite::memory:", &PoolParams::for_url("sqlite::memory:"))
184            .await
185            .unwrap();
186        assert_eq!(pool.dialect(), Dialect::Sqlite);
187        assert!(matches!(pool, Pool::Sq(_)));
188    }
189
190    /// 无法识别的 scheme 必须被拒，且**错误信息要点名出问题的 scheme**。
191    #[test]
192    fn unsupported_scheme_is_rejected_and_names_the_scheme() {
193        let rt = tokio::runtime::Builder::new_current_thread()
194            .enable_all()
195            .build()
196            .unwrap();
197        let err = rt
198            .block_on(Pool::connect("oracle://h/db", &PoolParams::default()))
199            .unwrap_err();
200        let msg = err.to_string();
201        assert!(msg.contains("unsupported"), "got: {msg}");
202        assert!(
203            msg.contains("oracle"),
204            "错误信息必须点名 scheme，got: {msg}"
205        );
206    }
207
208    /// Mssql 归 ecat-data-mssql（tiberius），不应被 sqlx 后端静默接受。
209    #[test]
210    fn mssql_scheme_is_rejected_by_sqlx_backend() {
211        let rt = tokio::runtime::Builder::new_current_thread()
212            .enable_all()
213            .build()
214            .unwrap();
215        let err = rt
216            .block_on(Pool::connect("mssql://h:1433/db", &PoolParams::default()))
217            .unwrap_err();
218        assert!(err.to_string().contains("ecat-data-mssql"), "got: {err}");
219    }
220}
ecat_data_sqlx/pool.rs

ecat_data_sqlx/
pool.rs
