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::{MysqlCompat, 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        Self::connect_with_mysql_compat(url, params, MysqlCompat::Mysql).await
123    }
124
125    /// 同 [`SqlxClient::connect_with_params`]，但可指定 MySQL 协议兼容模式。
126    /// Doris / StarRocks 走 `MysqlCompat::Doris` / `MysqlCompat::StarRocks`。
127    pub async fn connect_with_mysql_compat(
128        url: &str,
129        params: &PoolParams,
130        compat: MysqlCompat,
131    ) -> Result<Self, sqlx::Error> {
132        Ok(Self {
133            pool: Pool::connect_with_mysql_compat(url, params, compat).await?,
134            query_timeout: params.query_timeout,
135            min_connections: params.min_connections,
136            slow_query: params.slow_query,
137            version: std::sync::Mutex::new(None),
138        })
139    }
140
141    pub async fn connect_with_auth(
142        url: &str,
143        username: &str,
144        password: &str,
145    ) -> Result<Self, sqlx::Error> {
146        Self::connect(&with_auth_in_url(url, username, password)).await
147    }
148
149    pub async fn from_config(cfg: SqlxConfig) -> Result<Self, sqlx::Error> {
150        // 响亮失败而非静默忽略：本字段曾被 serde 接受却从不读，
151        // 用户以为开了 TLS 实际什么都没发生。
152        if cfg.tls.is_some() {
153            return Err(sqlx::Error::Configuration(Box::new(std::io::Error::new(
154                std::io::ErrorKind::InvalidInput,
155                "SqlxConfig.tls is not supported: configure TLS through URL parameters \
156                 instead (e.g. postgres://...?sslmode=require)",
157            ))));
158        }
159        let params = cfg.pool();
160        let url = match (&cfg.username, &cfg.password) {
161            (Some(u), Some(p)) if !u.is_empty() || !p.is_empty() => {
162                with_auth_in_url(&cfg.url, u, p)
163            }
164            _ => cfg.url.clone(),
165        };
166        Self::connect_with_mysql_compat(&url, &params, cfg.mysql_compat.unwrap_or_default()).await
167    }
168
169    /// 用已有原生池构造客户端。
170    ///
171    /// 拿的是**裸池**，没有 `SqlxConfig` 可读，因此查询超时与 [`SqlxClient::warm_up`]
172    /// 在这里都是关的（`query_timeout: None`、`min_connections: 0`）——
173    /// 这是 API 形状决定的，不是遗漏。需要两者请用
174    /// [`SqlxClient::connect_with_params`] 或 [`SqlxClient::from_config`]。
175    pub fn from_pool(pool: Pool) -> Self {
176        Self {
177            pool,
178            query_timeout: None,
179            min_connections: 0,
180            slow_query: None,
181            version: std::sync::Mutex::new(None),
182        }
183    }
184
185    /// 池本身（`metrics` feature 的 [`register_pool_metrics`] 要它）。
186    /// sqlx 的池内部是 `Arc`，克隆很轻。
187    pub fn pool(&self) -> &Pool {
188        &self.pool
189    }
190
191    /// 一次数据库调用的公共外壳：查询超时 + 慢查询告警。
192    /// 接在此处而不是各方法里，是为了两件事都只有一处定义。
193    async fn timed<F, T>(&self, sql: &str, fut: F) -> Result<T, RdbmsError>
194    where
195        F: std::future::Future<Output = Result<T, RdbmsError>>,
196    {
197        crate::tracing::timed(
198            self.slow_query,
199            sql,
200            run_with_timeout(BackendKind::Rdbms, self.query_timeout, fut),
201        )
202        .await
203    }
204
205    /// `SELECT VERSION()` 的结果，每 client 只问一次（版本不会中途变）。
206    async fn mysql_version(&self) -> Option<String> {
207        if let Some(v) = self.version.lock().unwrap().clone() {
208            return Some(v);
209        }
210        let rows = self.query("SELECT VERSION()").await.ok()?;
211        let v = rows.first()?.values().first()?.as_str()?.to_string();
212        *self.version.lock().unwrap() = Some(v.clone());
213        Some(v)
214    }
215
216    pub fn dialect(&self) -> Dialect {
217        self.pool.dialect()
218    }
219
220    /// 池内已建立的连接总数（供 metrics 与测试）。
221    ///
222    /// 与 trait 方法 [`RdbmsClient::pool_size`]（配置的最大连接数）同名：
223    /// 具体类型上固有方法优先，预算校验请经由 `&dyn RdbmsClient` 调用。
224    pub fn pool_size(&self) -> u32 {
225        self.pool.size()
226    }
227
228    /// 同步建满 `min_connections` 条连接并归还，让服务启动后立刻处于就绪态。
229    /// 在服务启动时调用一次。
230    ///
231    /// 为什么 sqlx 已有 `min_connections` 还需要它：sqlx 的保底由**后台任务
232    /// 异步维护**（`sqlx-core-0.8.6/src/pool/inner.rs:514-517`），`connect()`
233    /// 返回时不保证已建满 —— 启动后第一波请求会与后台任务抢跑。
234    pub async fn warm_up(&self) -> Result<(), RdbmsError> {
235        let mut guards = Vec::with_capacity(self.min_connections as usize);
236        while (guards.len() as u32) < self.min_connections {
237            guards.push(
238                self.pool
239                    .acquire()
240                    .await
241                    .map_err(|e| RdbmsError::Connection(e.to_string()))?,
242            );
243        }
244        // guards 在此处 drop，连接归还池中
245        Ok(())
246    }
247}
248
249#[async_trait]
250impl SqlExecutor for SqlxClient {
251    async fn execute(&self, sql: &str) -> Result<u64, RdbmsError> {
252        self.timed(sql, async {
253            let affected = match &self.pool {
254                Pool::Pg(p) => sqlx::query(sql).execute(p).await.map(|r| r.rows_affected()),
255                Pool::My(p) => sqlx::query(sql).execute(p).await.map(|r| r.rows_affected()),
256                Pool::Sq(p) => sqlx::query(sql).execute(p).await.map(|r| r.rows_affected()),
257            }
258            .map_err(db_err)?;
259            Ok(affected)
260        })
261        .await
262    }
263
264    async fn query(&self, sql: &str) -> Result<Vec<Row>, RdbmsError> {
265        self.timed(sql, async {
266            let rows = match &self.pool {
267                Pool::Pg(p) => sqlx::query(sql)
268                    .fetch_all(p)
269                    .await
270                    .map_err(db_err)
271                    .and_then(pg_rows_to_result),
272                Pool::My(p) => sqlx::query(sql)
273                    .fetch_all(p)
274                    .await
275                    .map_err(db_err)
276                    .and_then(mysql_rows_to_result),
277                Pool::Sq(p) => sqlx::query(sql)
278                    .fetch_all(p)
279                    .await
280                    .map_err(db_err)
281                    .and_then(sqlite_rows_to_result),
282            }?;
283            Ok(rows)
284        })
285        .await
286    }
287
288    async fn execute_with(
289        &self,
290        sql: &str,
291        params: &[serde_json::Value],
292    ) -> Result<u64, RdbmsError> {
293        self.timed(sql, async {
294            let affected = match &self.pool {
295                Pool::Pg(p) => {
296                    let mut q = sqlx::query(sql);
297                    for param in params {
298                        q = match param {
299                            serde_json::Value::String(s) => q.bind(s.as_str()),
300                            serde_json::Value::Number(n) => {
301                                if let Some(i) = n.as_i64() {
302                                    q.bind(i)
303                                } else if let Some(f) = n.as_f64() {
304                                    q.bind(f)
305                                } else {
306                                    q.bind(n.to_string())
307                                }
308                            }
309                            serde_json::Value::Bool(b) => q.bind(*b),
310                            serde_json::Value::Null => q.bind(None::<String>),
311                            _ => q.bind(param.to_string()),
312                        };
313                    }
314                    q.execute(p).await.map(|r| r.rows_affected())
315                }
316                Pool::My(p) => {
317                    let mut q = sqlx::query(sql);
318                    for param in params {
319                        q = match param {
320                            serde_json::Value::String(s) => q.bind(s.as_str()),
321                            serde_json::Value::Number(n) => {
322                                if let Some(i) = n.as_i64() {
323                                    q.bind(i)
324                                } else if let Some(f) = n.as_f64() {
325                                    q.bind(f)
326                                } else {
327                                    q.bind(n.to_string())
328                                }
329                            }
330                            serde_json::Value::Bool(b) => q.bind(*b),
331                            serde_json::Value::Null => q.bind(None::<String>),
332                            _ => q.bind(param.to_string()),
333                        };
334                    }
335                    q.execute(p).await.map(|r| r.rows_affected())
336                }
337                Pool::Sq(p) => {
338                    let mut q = sqlx::query(sql);
339                    for param in params {
340                        q = match param {
341                            serde_json::Value::String(s) => q.bind(s.as_str()),
342                            serde_json::Value::Number(n) => {
343                                if let Some(i) = n.as_i64() {
344                                    q.bind(i)
345                                } else if let Some(f) = n.as_f64() {
346                                    q.bind(f)
347                                } else {
348                                    q.bind(n.to_string())
349                                }
350                            }
351                            serde_json::Value::Bool(b) => q.bind(*b),
352                            serde_json::Value::Null => q.bind(None::<String>),
353                            _ => q.bind(param.to_string()),
354                        };
355                    }
356                    q.execute(p).await.map(|r| r.rows_affected())
357                }
358            }
359            .map_err(db_err)?;
360            Ok(affected)
361        })
362        .await
363    }
364
365    async fn query_with(
366        &self,
367        sql: &str,
368        params: &[serde_json::Value],
369    ) -> Result<Vec<Row>, RdbmsError> {
370        self.timed(sql, async {
371            let rows = match &self.pool {
372                Pool::Pg(p) => {
373                    let mut q = sqlx::query(sql);
374                    for param in params {
375                        q = match param {
376                            serde_json::Value::String(s) => q.bind(s.as_str()),
377                            serde_json::Value::Number(n) => {
378                                if let Some(i) = n.as_i64() {
379                                    q.bind(i)
380                                } else if let Some(f) = n.as_f64() {
381                                    q.bind(f)
382                                } else {
383                                    q.bind(n.to_string())
384                                }
385                            }
386                            serde_json::Value::Bool(b) => q.bind(*b),
387                            serde_json::Value::Null => q.bind(None::<String>),
388                            _ => q.bind(param.to_string()),
389                        };
390                    }
391                    q.fetch_all(p)
392                        .await
393                        .map_err(db_err)
394                        .and_then(pg_rows_to_result)
395                }
396                Pool::My(p) => {
397                    let mut q = sqlx::query(sql);
398                    for param in params {
399                        q = match param {
400                            serde_json::Value::String(s) => q.bind(s.as_str()),
401                            serde_json::Value::Number(n) => {
402                                if let Some(i) = n.as_i64() {
403                                    q.bind(i)
404                                } else if let Some(f) = n.as_f64() {
405                                    q.bind(f)
406                                } else {
407                                    q.bind(n.to_string())
408                                }
409                            }
410                            serde_json::Value::Bool(b) => q.bind(*b),
411                            serde_json::Value::Null => q.bind(None::<String>),
412                            _ => q.bind(param.to_string()),
413                        };
414                    }
415                    q.fetch_all(p)
416                        .await
417                        .map_err(db_err)
418                        .and_then(mysql_rows_to_result)
419                }
420                Pool::Sq(p) => {
421                    let mut q = sqlx::query(sql);
422                    for param in params {
423                        q = match param {
424                            serde_json::Value::String(s) => q.bind(s.as_str()),
425                            serde_json::Value::Number(n) => {
426                                if let Some(i) = n.as_i64() {
427                                    q.bind(i)
428                                } else if let Some(f) = n.as_f64() {
429                                    q.bind(f)
430                                } else {
431                                    q.bind(n.to_string())
432                                }
433                            }
434                            serde_json::Value::Bool(b) => q.bind(*b),
435                            serde_json::Value::Null => q.bind(None::<String>),
436                            _ => q.bind(param.to_string()),
437                        };
438                    }
439                    q.fetch_all(p)
440                        .await
441                        .map_err(db_err)
442                        .and_then(sqlite_rows_to_result)
443                }
444            }?;
445            Ok(rows)
446        })
447        .await
448    }
449
450    /// 两步式原子执行：**自己开一个事务**包住两条语句并提交。
451    /// ORM 的 MySQL 主键回填走这里（`LAST_INSERT_ID()` 是连接作用域的，
452    /// 池下直发两条可能落到不同连接）；一步式后端用不到（默认实现报错）。
453    ///
454    /// 第一条失败时不提交，`tx` 在 drop 时由底层 sqlx 事务回滚。
455    async fn execute_then_query(
456        &self,
457        first: &str,
458        first_params: &[serde_json::Value],
459        second: &str,
460    ) -> Result<Vec<Row>, RdbmsError> {
461        let tx = self.transaction().await?;
462        tx.execute_with(first, first_params).await?;
463        let rows = tx.query(second).await?;
464        tx.commit().await?;
465        Ok(rows)
466    }
467
468    fn dialect(&self) -> Dialect {
469        // 注意是 `self.pool.dialect()`，写成 `self.dialect()` 会无限递归。
470        self.pool.dialect()
471    }
472}
ecat_data_sqlx/lib.rs

ecat_data_sqlx/
lib.rs
