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

ecat_data_sqlx/
lib.rs
