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

ecat_data_sqlx/
pool.rs
