ecat_data_sqlx/
lib.rs

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

ecat_data_sqlx/
lib.rs
