ecat_data_sqlx/
pool.rs

1// Copyright (c) 2026 erik <erik@erik.xyz> — https://erik.xyz
2
3use crate::config::{MysqlCompat, PoolParams};
4use ecat_data::Dialect;
5use sqlx::mysql::{MySqlConnectOptions, MySqlPool, MySqlPoolOptions};
6use sqlx::postgres::{PgPool, PgPoolOptions};
7use sqlx::sqlite::{SqlitePool, SqlitePoolOptions};
8use std::str::FromStr;
9
10/// 三种原生池。方言由变体本身承载 —— 不再解析 URL 猜测。
11///
12/// `Clone`：sqlx 的池内部是 `Arc`，克隆只是多一个句柄（`metrics` feature 的
13/// collector 要长期持有池）。
14#[derive(Debug, Clone)]
15pub enum Pool {
16    Pg(PgPool),
17    My(MySqlPool),
18    Sq(SqlitePool),
19}
20
21/// 三种驱动的连接 guard 的聚合类型。Drop 即归还连接。
22/// `PoolConnection<DB>` 自身持有 `Arc`，无需生命周期参数。
23pub enum PoolGuard {
24    Pg(sqlx::pool::PoolConnection<sqlx::Postgres>),
25    My(sqlx::pool::PoolConnection<sqlx::MySql>),
26    Sq(sqlx::pool::PoolConnection<sqlx::Sqlite>),
27}
28
29impl Pool {
30    /// 按 URL scheme 建立对应的原生池。
31    pub async fn connect(url: &str, params: &PoolParams) -> Result<Self, sqlx::Error> {
32        Self::connect_with_mysql_compat(url, params, MysqlCompat::Mysql).await
33    }
34
35    /// 同 [`Pool::connect`]，但可指定 MySQL 协议兼容模式。
36    pub async fn connect_with_mysql_compat(
37        url: &str,
38        params: &PoolParams,
39        compat: MysqlCompat,
40    ) -> Result<Self, sqlx::Error> {
41        match Dialect::from_url(url) {
42            Dialect::Postgres => Ok(Self::Pg(
43                common_pg(PgPoolOptions::new(), params).connect(url).await?,
44            )),
45            Dialect::MySql => Ok(Self::My(
46                common_mysql(MySqlPoolOptions::new(), params)
47                    .connect_with(mysql_options(url, compat)?)
48                    .await?,
49            )),
50            Dialect::Sqlite => Ok(Self::Sq(
51                common_sqlite(SqlitePoolOptions::new(), params)
52                    .connect(url)
53                    .await?,
54            )),
55            // 报 scheme 原文而非 Dialect：`Dialect::Standard` 是「无法识别」的
56            // 回退值，把它打进错误信息等于丢掉唯一的诊断线索。
57            Dialect::Standard => Err(unsupported(url)),
58            // Mssql 走 tiberius 后端（ecat-data-mssql），不是 sqlx
59            Dialect::Mssql => Err(sqlx::Error::Configuration(Box::new(std::io::Error::new(
60                std::io::ErrorKind::InvalidInput,
61                "mssql:// is served by ecat-data-mssql, not the sqlx backend",
62            )))),
63        }
64    }
65
66    pub fn dialect(&self) -> Dialect {
67        match self {
68            Self::Pg(_) => Dialect::Postgres,
69            Self::My(_) => Dialect::MySql,
70            Self::Sq(_) => Dialect::Sqlite,
71        }
72    }
73
74    /// 池内已建立的连接总数。
75    pub fn size(&self) -> u32 {
76        match self {
77            Self::Pg(p) => p.size(),
78            Self::My(p) => p.size(),
79            Self::Sq(p) => p.size(),
80        }
81    }
82
83    /// 池配置的最大连接数（预算校验用）。注意 [`Pool::size`] 是当前已建立数，
84    /// 空闲池的 `size()` 不是上限。
85    pub fn max_connections(&self) -> u32 {
86        match self {
87            Self::Pg(p) => p.options().get_max_connections(),
88            Self::My(p) => p.options().get_max_connections(),
89            Self::Sq(p) => p.options().get_max_connections(),
90        }
91    }
92
93    /// 池内空闲连接数。
94    pub fn idle(&self) -> u32 {
95        match self {
96            Self::Pg(p) => p.num_idle() as u32,
97            Self::My(p) => p.num_idle() as u32,
98            Self::Sq(p) => p.num_idle() as u32,
99        }
100    }
101
102    /// 从池中取一条连接（预热用）。归还由 guard 的 Drop 完成。
103    pub async fn acquire(&self) -> Result<PoolGuard, sqlx::Error> {
104        Ok(match self {
105            Self::Pg(p) => PoolGuard::Pg(p.acquire().await?),
106            Self::My(p) => PoolGuard::My(p.acquire().await?),
107            Self::Sq(p) => PoolGuard::Sq(p.acquire().await?),
108        })
109    }
110}
111
112/// 无法识别的 scheme：错误信息要点名出问题的 scheme。
113fn unsupported(url: &str) -> sqlx::Error {
114    let scheme = url
115        .split("://")
116        .next()
117        .unwrap_or(url)
118        .split(':')
119        .next()
120        .unwrap_or(url);
121    sqlx::Error::Configuration(Box::new(std::io::Error::new(
122        std::io::ErrorKind::InvalidInput,
123        format!("unsupported database scheme for sqlx backend: {scheme}"),
124    )))
125}
126
127// 三个 common_* 把 PoolParams 落到各驱动的 PoolOptions 上。
128// 注意：三者的 Connection 类型不同，after_connect 闭包无法合并，各写一份。
129
130fn common_pg(opts: PgPoolOptions, p: &PoolParams) -> PgPoolOptions {
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
153/// 解析 MySQL URL，并按兼容模式关闭 sqlx 默认的 sql_mode 子查询。
154///
155/// StarRocks 不支持 `SET sql_mode=(SELECT CONCAT(...))`（sqlx issue #3205），
156/// 而 `pipes_as_concat` / `no_engine_substitution` 只能通过连接选项设置，URL
157/// 参数无法表达。Doris / StarRocks 同时显式使用 utf8mb4 通用排序规则，避免
158/// sqlx 默认的 `utf8mb4_0900_ai_ci` 在旧服务端不存在。
159fn mysql_options(url: &str, compat: MysqlCompat) -> Result<MySqlConnectOptions, sqlx::Error> {
160    let opts = MySqlConnectOptions::from_str(url)?;
161    if compat == MysqlCompat::Mysql {
162        return Ok(opts);
163    }
164    Ok(opts
165        .pipes_as_concat(false)
166        .no_engine_substitution(false)
167        .charset("utf8mb4")
168        .collation("utf8mb4_general_ci"))
169}
170
171fn common_mysql(opts: MySqlPoolOptions, p: &PoolParams) -> MySqlPoolOptions {
172    let opts = opts
173        .max_connections(p.max_connections)
174        .min_connections(p.min_connections)
175        .acquire_timeout(p.acquire_timeout)
176        .idle_timeout(Some(p.idle_timeout))
177        .max_lifetime(Some(p.max_lifetime))
178        .test_before_acquire(p.test_before_acquire);
179    if p.session_init.is_empty() {
180        return opts;
181    }
182    let stmts = p.session_init.clone();
183    opts.after_connect(move |conn, _meta| {
184        let stmts = stmts.clone();
185        Box::pin(async move {
186            for s in &stmts {
187                sqlx::query(s).execute(&mut *conn).await?;
188            }
189            Ok(())
190        })
191    })
192}
193
194fn common_sqlite(opts: SqlitePoolOptions, p: &PoolParams) -> SqlitePoolOptions {
195    let opts = opts
196        .max_connections(p.max_connections)
197        .min_connections(p.min_connections)
198        .acquire_timeout(p.acquire_timeout)
199        .idle_timeout(Some(p.idle_timeout))
200        .max_lifetime(Some(p.max_lifetime))
201        .test_before_acquire(p.test_before_acquire);
202    if p.session_init.is_empty() {
203        return opts;
204    }
205    let stmts = p.session_init.clone();
206    opts.after_connect(move |conn, _meta| {
207        let stmts = stmts.clone();
208        Box::pin(async move {
209            for s in &stmts {
210                sqlx::query(s).execute(&mut *conn).await?;
211            }
212            Ok(())
213        })
214    })
215}
216
217#[cfg(test)]
218mod tests {
219    use super::*;
220    use ecat_data::Dialect;
221
222    #[tokio::test]
223    async fn sqlite_url_builds_sqlite_pool() {
224        let pool = Pool::connect("sqlite::memory:", &PoolParams::for_url("sqlite::memory:"))
225            .await
226            .unwrap();
227        assert_eq!(pool.dialect(), Dialect::Sqlite);
228        assert!(matches!(pool, Pool::Sq(_)));
229    }
230
231    /// 无法识别的 scheme 必须被拒，且**错误信息要点名出问题的 scheme**。
232    #[test]
233    fn unsupported_scheme_is_rejected_and_names_the_scheme() {
234        let rt = tokio::runtime::Builder::new_current_thread()
235            .enable_all()
236            .build()
237            .unwrap();
238        let err = rt
239            .block_on(Pool::connect("oracle://h/db", &PoolParams::default()))
240            .unwrap_err();
241        let msg = err.to_string();
242        assert!(msg.contains("unsupported"), "got: {msg}");
243        assert!(
244            msg.contains("oracle"),
245            "错误信息必须点名 scheme，got: {msg}"
246        );
247    }
248
249    /// Doris / StarRocks 兼容模式必须关掉 sqlx 默认的 sql_mode 子查询；
250    /// `MySqlConnectOptions` 的字段是私有的，这里用 Debug 输出钉住三个开关。
251    #[test]
252    fn starrocks_compat_disables_sql_mode_subquery() {
253        let opts = mysql_options("mysql://h:9030/db", MysqlCompat::StarRocks).unwrap();
254        let debug = format!("{opts:?}");
255        assert!(debug.contains("pipes_as_concat: false"), "{debug}");
256        assert!(debug.contains("no_engine_substitution: false"), "{debug}");
257        assert!(
258            debug.contains("collation: Some(\"utf8mb4_general_ci\")"),
259            "{debug}"
260        );
261    }
262
263    #[test]
264    fn mysql_default_keeps_sqlx_session_settings() {
265        let opts = mysql_options("mysql://h/db", MysqlCompat::Mysql).unwrap();
266        let debug = format!("{opts:?}");
267        assert!(debug.contains("pipes_as_concat: true"), "{debug}");
268        assert!(debug.contains("no_engine_substitution: true"), "{debug}");
269    }
270
271    /// Mssql 归 ecat-data-mssql（tiberius），不应被 sqlx 后端静默接受。
272    #[test]
273    fn mssql_scheme_is_rejected_by_sqlx_backend() {
274        let rt = tokio::runtime::Builder::new_current_thread()
275            .enable_all()
276            .build()
277            .unwrap();
278        let err = rt
279            .block_on(Pool::connect("mssql://h:1433/db", &PoolParams::default()))
280            .unwrap_err();
281        assert!(err.to_string().contains("ecat-data-mssql"), "got: {err}");
282    }
283}
ecat_data_sqlx/pool.rs

ecat_data_sqlx/
pool.rs
