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    /// 池内空闲连接数。
74    pub fn idle(&self) -> u32 {
75        match self {
76            Self::Pg(p) => p.num_idle() as u32,
77            Self::My(p) => p.num_idle() as u32,
78            Self::Sq(p) => p.num_idle() as u32,
79        }
80    }
81
82    /// 从池中取一条连接（预热用）。归还由 guard 的 Drop 完成。
83    pub async fn acquire(&self) -> Result<PoolGuard, sqlx::Error> {
84        Ok(match self {
85            Self::Pg(p) => PoolGuard::Pg(p.acquire().await?),
86            Self::My(p) => PoolGuard::My(p.acquire().await?),
87            Self::Sq(p) => PoolGuard::Sq(p.acquire().await?),
88        })
89    }
90}
91
92/// 无法识别的 scheme：错误信息要点名出问题的 scheme。
93fn unsupported(url: &str) -> sqlx::Error {
94    let scheme = url
95        .split("://")
96        .next()
97        .unwrap_or(url)
98        .split(':')
99        .next()
100        .unwrap_or(url);
101    sqlx::Error::Configuration(Box::new(std::io::Error::new(
102        std::io::ErrorKind::InvalidInput,
103        format!("unsupported database scheme for sqlx backend: {scheme}"),
104    )))
105}
106
107// 三个 common_* 把 PoolParams 落到各驱动的 PoolOptions 上。
108// 注意：三者的 Connection 类型不同，after_connect 闭包无法合并，各写一份。
109
110fn common_pg(opts: PgPoolOptions, p: &PoolParams) -> PgPoolOptions {
111    let opts = opts
112        .max_connections(p.max_connections)
113        .min_connections(p.min_connections)
114        .acquire_timeout(p.acquire_timeout)
115        .idle_timeout(Some(p.idle_timeout))
116        .max_lifetime(Some(p.max_lifetime))
117        .test_before_acquire(p.test_before_acquire);
118    if p.session_init.is_empty() {
119        return opts;
120    }
121    let stmts = p.session_init.clone();
122    opts.after_connect(move |conn, _meta| {
123        let stmts = stmts.clone();
124        Box::pin(async move {
125            for s in &stmts {
126                sqlx::query(s).execute(&mut *conn).await?;
127            }
128            Ok(())
129        })
130    })
131}
132
133fn common_mysql(opts: MySqlPoolOptions, p: &PoolParams) -> MySqlPoolOptions {
134    let opts = opts
135        .max_connections(p.max_connections)
136        .min_connections(p.min_connections)
137        .acquire_timeout(p.acquire_timeout)
138        .idle_timeout(Some(p.idle_timeout))
139        .max_lifetime(Some(p.max_lifetime))
140        .test_before_acquire(p.test_before_acquire);
141    if p.session_init.is_empty() {
142        return opts;
143    }
144    let stmts = p.session_init.clone();
145    opts.after_connect(move |conn, _meta| {
146        let stmts = stmts.clone();
147        Box::pin(async move {
148            for s in &stmts {
149                sqlx::query(s).execute(&mut *conn).await?;
150            }
151            Ok(())
152        })
153    })
154}
155
156fn common_sqlite(opts: SqlitePoolOptions, p: &PoolParams) -> SqlitePoolOptions {
157    let opts = opts
158        .max_connections(p.max_connections)
159        .min_connections(p.min_connections)
160        .acquire_timeout(p.acquire_timeout)
161        .idle_timeout(Some(p.idle_timeout))
162        .max_lifetime(Some(p.max_lifetime))
163        .test_before_acquire(p.test_before_acquire);
164    if p.session_init.is_empty() {
165        return opts;
166    }
167    let stmts = p.session_init.clone();
168    opts.after_connect(move |conn, _meta| {
169        let stmts = stmts.clone();
170        Box::pin(async move {
171            for s in &stmts {
172                sqlx::query(s).execute(&mut *conn).await?;
173            }
174            Ok(())
175        })
176    })
177}
178
179#[cfg(test)]
180mod tests {
181    use super::*;
182    use ecat_data::Dialect;
183
184    #[tokio::test]
185    async fn sqlite_url_builds_sqlite_pool() {
186        let pool = Pool::connect("sqlite::memory:", &PoolParams::for_url("sqlite::memory:"))
187            .await
188            .unwrap();
189        assert_eq!(pool.dialect(), Dialect::Sqlite);
190        assert!(matches!(pool, Pool::Sq(_)));
191    }
192
193    /// 无法识别的 scheme 必须被拒，且**错误信息要点名出问题的 scheme**。
194    #[test]
195    fn unsupported_scheme_is_rejected_and_names_the_scheme() {
196        let rt = tokio::runtime::Builder::new_current_thread()
197            .enable_all()
198            .build()
199            .unwrap();
200        let err = rt
201            .block_on(Pool::connect("oracle://h/db", &PoolParams::default()))
202            .unwrap_err();
203        let msg = err.to_string();
204        assert!(msg.contains("unsupported"), "got: {msg}");
205        assert!(
206            msg.contains("oracle"),
207            "错误信息必须点名 scheme，got: {msg}"
208        );
209    }
210
211    /// Mssql 归 ecat-data-mssql（tiberius），不应被 sqlx 后端静默接受。
212    #[test]
213    fn mssql_scheme_is_rejected_by_sqlx_backend() {
214        let rt = tokio::runtime::Builder::new_current_thread()
215            .enable_all()
216            .build()
217            .unwrap();
218        let err = rt
219            .block_on(Pool::connect("mssql://h:1433/db", &PoolParams::default()))
220            .unwrap_err();
221        assert!(err.to_string().contains("ecat-data-mssql"), "got: {err}");
222    }
223}
ecat_data_sqlx/pool.rs

ecat_data_sqlx/
pool.rs
