ecat_data_sqlx/
lib.rs

1// Copyright (c) 2026 erik <erik@erik.xyz> — https://erik.xyz
2mod cell;
3mod config;
4#[cfg(feature = "health")]
5mod health;
6#[cfg(test)]
7mod live_tests;
8#[cfg(feature = "metrics")]
9mod metrics;
10mod pool;
11#[cfg(test)]
12mod tests;
13// **不**随 feature 门控：feature 关闭时 `timed` 是直通函数（见模块文档），
14// 这样每个调用点不必各写一次 `#[cfg]`。
15mod tracing;
16mod transaction;
17
18use async_trait::async_trait;
19use cell::{mysql_rows_to_result, pg_rows_to_result, sqlite_rows_to_result};
20use ecat_data::{
21    BackendKind, Dialect, RdbmsClient, RdbmsError, Row, SqlExecutor, Transaction, TransactionInner,
22    run_with_timeout,
23};
24use std::time::Duration;
25use transaction::{MyTx, PgTx, SqTx};
26
27pub use config::{PoolParams, SqlxConfig};
28#[cfg(feature = "health")]
29pub use health::RdbmsHealthCheck;
30#[cfg(feature = "metrics")]
31pub use metrics::register_pool_metrics;
32pub use pool::{Pool, PoolGuard};
33
34fn percent_encode(s: &str) -> String {
35    s.chars()
36        .map(|c| match c {
37            ':' | '/' | '@' | '#' | '?' | '&' | '=' | '%' | '+' | ' ' => {
38                format!("%{:02X}", c as u8)
39            }
40            _ => c.to_string(),
41        })
42        .collect()
43}
44
45/// 把账号密码嵌入 URL（已有 `@` 时原样返回）。
46fn with_auth_in_url(url: &str, username: &str, password: &str) -> String {
47    if url.contains('@') {
48        return url.to_string();
49    }
50    let encoded_user = percent_encode(username);
51    let encoded_pass = percent_encode(password);
52    url.replacen("://", &format!("://{encoded_user}:{encoded_pass}@"), 1)
53}
54
55/// sqlx 错误 → [`RdbmsError::Database`]。三路分派后每个方法都有三个分支，
56/// 统一走这里避免 6 处重复。
57///
58/// 也是「等连接超时」的唯一漏斗：`PoolTimedOut` 由 sqlx 在内部取用连接超时时
59/// 抛出，流经此处时记一笔（`metrics` feature 的
60/// `ecat_rdbms_pool_timeouts_total`）。
61fn db_err(e: sqlx::Error) -> RdbmsError {
62    #[cfg(feature = "metrics")]
63    if matches!(e, sqlx::Error::PoolTimedOut) {
64        metrics::count_pool_timeout();
65    }
66    RdbmsError::Database(e.to_string())
67}
68
69pub struct SqlxClient {
70    pool: Pool,
71    query_timeout: Option<Duration>,
72    min_connections: u32,
73    /// 慢查询告警阈值；`None` = 不打。只有 `tracing` feature 会读它，
74    /// 但字段本身常驻 —— 否则配置项会随 feature 时有时无（见 [`SqlxConfig`]）。
75    slow_query: Option<Duration>,
76}
77
78impl SqlxClient {
79    pub async fn connect(url: &str) -> Result<Self, sqlx::Error> {
80        // 必须走 for_url 而非 default()：否则 connect() 拿不到方言默认的
81        // session_init，与 from_config() 行为静默不一致。
82        Self::connect_with_params(url, &PoolParams::for_url(url)).await
83    }
84
85    pub async fn connect_with_params(url: &str, params: &PoolParams) -> Result<Self, sqlx::Error> {
86        Ok(Self {
87            pool: Pool::connect(url, params).await?,
88            query_timeout: params.query_timeout,
89            min_connections: params.min_connections,
90            slow_query: params.slow_query,
91        })
92    }
93
94    pub async fn connect_with_auth(
95        url: &str,
96        username: &str,
97        password: &str,
98    ) -> Result<Self, sqlx::Error> {
99        Self::connect(&with_auth_in_url(url, username, password)).await
100    }
101
102    pub async fn from_config(cfg: SqlxConfig) -> Result<Self, sqlx::Error> {
103        // 响亮失败而非静默忽略：本字段曾被 serde 接受却从不读，
104        // 用户以为开了 TLS 实际什么都没发生。
105        if cfg.tls.is_some() {
106            return Err(sqlx::Error::Configuration(Box::new(std::io::Error::new(
107                std::io::ErrorKind::InvalidInput,
108                "SqlxConfig.tls is not supported: configure TLS through URL parameters \
109                 instead (e.g. postgres://...?sslmode=require)",
110            ))));
111        }
112        let params = cfg.pool();
113        let url = match (&cfg.username, &cfg.password) {
114            (Some(u), Some(p)) if !u.is_empty() || !p.is_empty() => {
115                with_auth_in_url(&cfg.url, u, p)
116            }
117            _ => cfg.url.clone(),
118        };
119        Self::connect_with_params(&url, &params).await
120    }
121
122    /// 用已有原生池构造客户端。
123    ///
124    /// 拿的是**裸池**，没有 `SqlxConfig` 可读，因此查询超时与 [`SqlxClient::warm_up`]
125    /// 在这里都是关的（`query_timeout: None`、`min_connections: 0`）——
126    /// 这是 API 形状决定的，不是遗漏。需要两者请用
127    /// [`SqlxClient::connect_with_params`] 或 [`SqlxClient::from_config`]。
128    pub fn from_pool(pool: Pool) -> Self {
129        Self {
130            pool,
131            query_timeout: None,
132            min_connections: 0,
133            slow_query: None,
134        }
135    }
136
137    /// 池本身（`metrics` feature 的 [`register_pool_metrics`] 要它）。
138    /// sqlx 的池内部是 `Arc`，克隆很轻。
139    pub fn pool(&self) -> &Pool {
140        &self.pool
141    }
142
143    /// 一次数据库调用的公共外壳：查询超时 + 慢查询告警。
144    /// 接在此处而不是各方法里，是为了两件事都只有一处定义。
145    async fn timed<F, T>(&self, sql: &str, fut: F) -> Result<T, RdbmsError>
146    where
147        F: std::future::Future<Output = Result<T, RdbmsError>>,
148    {
149        crate::tracing::timed(
150            self.slow_query,
151            sql,
152            run_with_timeout(BackendKind::Rdbms, self.query_timeout, fut),
153        )
154        .await
155    }
156
157    pub fn dialect(&self) -> Dialect {
158        self.pool.dialect()
159    }
160
161    /// 池内已建立的连接总数（供 metrics 与测试）。
162    pub fn pool_size(&self) -> u32 {
163        self.pool.size()
164    }
165
166    /// 同步建满 `min_connections` 条连接并归还，让服务启动后立刻处于就绪态。
167    /// 在服务启动时调用一次。
168    ///
169    /// 为什么 sqlx 已有 `min_connections` 还需要它：sqlx 的保底由**后台任务
170    /// 异步维护**（`sqlx-core-0.8.6/src/pool/inner.rs:514-517`），`connect()`
171    /// 返回时不保证已建满 —— 启动后第一波请求会与后台任务抢跑。
172    pub async fn warm_up(&self) -> Result<(), RdbmsError> {
173        let mut guards = Vec::with_capacity(self.min_connections as usize);
174        while (guards.len() as u32) < self.min_connections {
175            guards.push(
176                self.pool
177                    .acquire()
178                    .await
179                    .map_err(|e| RdbmsError::Connection(e.to_string()))?,
180            );
181        }
182        // guards 在此处 drop，连接归还池中
183        Ok(())
184    }
185}
186
187#[async_trait]
188impl SqlExecutor for SqlxClient {
189    async fn execute(&self, sql: &str) -> Result<u64, RdbmsError> {
190        self.timed(sql, async {
191            let affected = match &self.pool {
192                Pool::Pg(p) => sqlx::query(sql).execute(p).await.map(|r| r.rows_affected()),
193                Pool::My(p) => sqlx::query(sql).execute(p).await.map(|r| r.rows_affected()),
194                Pool::Sq(p) => sqlx::query(sql).execute(p).await.map(|r| r.rows_affected()),
195            }
196            .map_err(db_err)?;
197            Ok(affected)
198        })
199        .await
200    }
201
202    async fn query(&self, sql: &str) -> Result<Vec<Row>, RdbmsError> {
203        self.timed(sql, async {
204            let rows = match &self.pool {
205                Pool::Pg(p) => sqlx::query(sql)
206                    .fetch_all(p)
207                    .await
208                    .map_err(db_err)
209                    .and_then(pg_rows_to_result),
210                Pool::My(p) => sqlx::query(sql)
211                    .fetch_all(p)
212                    .await
213                    .map_err(db_err)
214                    .and_then(mysql_rows_to_result),
215                Pool::Sq(p) => sqlx::query(sql)
216                    .fetch_all(p)
217                    .await
218                    .map_err(db_err)
219                    .and_then(sqlite_rows_to_result),
220            }?;
221            Ok(rows)
222        })
223        .await
224    }
225
226    async fn execute_with(
227        &self,
228        sql: &str,
229        params: &[serde_json::Value],
230    ) -> Result<u64, RdbmsError> {
231        self.timed(sql, async {
232            let affected = match &self.pool {
233                Pool::Pg(p) => {
234                    let mut q = sqlx::query(sql);
235                    for param in params {
236                        q = match param {
237                            serde_json::Value::String(s) => q.bind(s.as_str()),
238                            serde_json::Value::Number(n) => {
239                                if let Some(i) = n.as_i64() {
240                                    q.bind(i)
241                                } else if let Some(f) = n.as_f64() {
242                                    q.bind(f)
243                                } else {
244                                    q.bind(n.to_string())
245                                }
246                            }
247                            serde_json::Value::Bool(b) => q.bind(*b),
248                            serde_json::Value::Null => q.bind(None::<String>),
249                            _ => q.bind(param.to_string()),
250                        };
251                    }
252                    q.execute(p).await.map(|r| r.rows_affected())
253                }
254                Pool::My(p) => {
255                    let mut q = sqlx::query(sql);
256                    for param in params {
257                        q = match param {
258                            serde_json::Value::String(s) => q.bind(s.as_str()),
259                            serde_json::Value::Number(n) => {
260                                if let Some(i) = n.as_i64() {
261                                    q.bind(i)
262                                } else if let Some(f) = n.as_f64() {
263                                    q.bind(f)
264                                } else {
265                                    q.bind(n.to_string())
266                                }
267                            }
268                            serde_json::Value::Bool(b) => q.bind(*b),
269                            serde_json::Value::Null => q.bind(None::<String>),
270                            _ => q.bind(param.to_string()),
271                        };
272                    }
273                    q.execute(p).await.map(|r| r.rows_affected())
274                }
275                Pool::Sq(p) => {
276                    let mut q = sqlx::query(sql);
277                    for param in params {
278                        q = match param {
279                            serde_json::Value::String(s) => q.bind(s.as_str()),
280                            serde_json::Value::Number(n) => {
281                                if let Some(i) = n.as_i64() {
282                                    q.bind(i)
283                                } else if let Some(f) = n.as_f64() {
284                                    q.bind(f)
285                                } else {
286                                    q.bind(n.to_string())
287                                }
288                            }
289                            serde_json::Value::Bool(b) => q.bind(*b),
290                            serde_json::Value::Null => q.bind(None::<String>),
291                            _ => q.bind(param.to_string()),
292                        };
293                    }
294                    q.execute(p).await.map(|r| r.rows_affected())
295                }
296            }
297            .map_err(db_err)?;
298            Ok(affected)
299        })
300        .await
301    }
302
303    async fn query_with(
304        &self,
305        sql: &str,
306        params: &[serde_json::Value],
307    ) -> Result<Vec<Row>, RdbmsError> {
308        self.timed(sql, async {
309            let rows = match &self.pool {
310                Pool::Pg(p) => {
311                    let mut q = sqlx::query(sql);
312                    for param in params {
313                        q = match param {
314                            serde_json::Value::String(s) => q.bind(s.as_str()),
315                            serde_json::Value::Number(n) => {
316                                if let Some(i) = n.as_i64() {
317                                    q.bind(i)
318                                } else if let Some(f) = n.as_f64() {
319                                    q.bind(f)
320                                } else {
321                                    q.bind(n.to_string())
322                                }
323                            }
324                            serde_json::Value::Bool(b) => q.bind(*b),
325                            serde_json::Value::Null => q.bind(None::<String>),
326                            _ => q.bind(param.to_string()),
327                        };
328                    }
329                    q.fetch_all(p)
330                        .await
331                        .map_err(db_err)
332                        .and_then(pg_rows_to_result)
333                }
334                Pool::My(p) => {
335                    let mut q = sqlx::query(sql);
336                    for param in params {
337                        q = match param {
338                            serde_json::Value::String(s) => q.bind(s.as_str()),
339                            serde_json::Value::Number(n) => {
340                                if let Some(i) = n.as_i64() {
341                                    q.bind(i)
342                                } else if let Some(f) = n.as_f64() {
343                                    q.bind(f)
344                                } else {
345                                    q.bind(n.to_string())
346                                }
347                            }
348                            serde_json::Value::Bool(b) => q.bind(*b),
349                            serde_json::Value::Null => q.bind(None::<String>),
350                            _ => q.bind(param.to_string()),
351                        };
352                    }
353                    q.fetch_all(p)
354                        .await
355                        .map_err(db_err)
356                        .and_then(mysql_rows_to_result)
357                }
358                Pool::Sq(p) => {
359                    let mut q = sqlx::query(sql);
360                    for param in params {
361                        q = match param {
362                            serde_json::Value::String(s) => q.bind(s.as_str()),
363                            serde_json::Value::Number(n) => {
364                                if let Some(i) = n.as_i64() {
365                                    q.bind(i)
366                                } else if let Some(f) = n.as_f64() {
367                                    q.bind(f)
368                                } else {
369                                    q.bind(n.to_string())
370                                }
371                            }
372                            serde_json::Value::Bool(b) => q.bind(*b),
373                            serde_json::Value::Null => q.bind(None::<String>),
374                            _ => q.bind(param.to_string()),
375                        };
376                    }
377                    q.fetch_all(p)
378                        .await
379                        .map_err(db_err)
380                        .and_then(sqlite_rows_to_result)
381                }
382            }?;
383            Ok(rows)
384        })
385        .await
386    }
387
388    /// 两步式原子执行：**自己开一个事务**包住两条语句并提交。
389    /// ORM 的 MySQL 主键回填走这里（`LAST_INSERT_ID()` 是连接作用域的，
390    /// 池下直发两条可能落到不同连接）；一步式后端用不到（默认实现报错）。
391    ///
392    /// 第一条失败时不提交，`tx` 在 drop 时由底层 sqlx 事务回滚。
393    async fn execute_then_query(
394        &self,
395        first: &str,
396        first_params: &[serde_json::Value],
397        second: &str,
398    ) -> Result<Vec<Row>, RdbmsError> {
399        let tx = self.transaction().await?;
400        tx.execute_with(first, first_params).await?;
401        let rows = tx.query(second).await?;
402        tx.commit().await?;
403        Ok(rows)
404    }
405
406    fn dialect(&self) -> Dialect {
407        // 注意是 `self.pool.dialect()`，写成 `self.dialect()` 会无限递归。
408        self.pool.dialect()
409    }
410}
411
412#[async_trait]
413impl RdbmsClient for SqlxClient {
414    async fn transaction(&self) -> Result<Transaction, RdbmsError> {
415        let dialect = self.pool.dialect();
416        // 超时值随事务一起带走：`transaction()` 之后调用方手里只有 `Transaction`，
417        // 事务 wrapper 是唯一知道该用什么超时的地方。
418        let query_timeout = self.query_timeout;
419        let inner: Box<dyn TransactionInner> = match &self.pool {
420            Pool::Pg(p) => Box::new(PgTx {
421                inner: Some(p.begin().await.map_err(db_err)?),
422                dialect,
423                query_timeout,
424            }),
425            Pool::My(p) => Box::new(MyTx {
426                inner: Some(p.begin().await.map_err(db_err)?),
427                dialect,
428                query_timeout,
429            }),
430            Pool::Sq(p) => Box::new(SqTx {
431                inner: Some(p.begin().await.map_err(db_err)?),
432                dialect,
433                query_timeout,
434            }),
435        };
436        Ok(Transaction::with_inner(inner))
437    }
438}
ecat_data_sqlx/lib.rs

ecat_data_sqlx/
lib.rs
