ecat_data_sqlx/
lib.rs

1// Copyright (c) 2026 erik <erik@erik.xyz> — https://erik.xyz
2use async_trait::async_trait;
3use base64::Engine as _;
4use ecat_data::{RdbmsClient, RdbmsError, Row, TransactionInner};
5use ecat_tls::TlsClientConfig;
6use serde::Deserialize;
7use sqlx::any::AnyRow;
8use sqlx::{AnyPool, Column as SqlxColumn, Row as SqlxRow};
9
10#[derive(Debug, Clone, Deserialize)]
11pub struct SqlxConfig {
12    pub url: String,
13    #[serde(default)]
14    pub username: Option<String>,
15    #[serde(default)]
16    pub password: Option<String>,
17    /// TLS — SQLx TLS is configured via URL params (e.g. ?sslmode=require).
18    /// This field is reserved for future programmatic TLS support.
19    #[serde(default)]
20    pub tls: Option<TlsClientConfig>,
21}
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/// AnyPool 首次连接前必须先安装驱动，否则 sqlx 直接 panic
35/// "No drivers installed"。install_default_drivers 内部对每个驱动
36/// 也是 Once 保护，这里再用 Once 显式保证只装一次。
37static DRIVERS_INSTALLED: std::sync::Once = std::sync::Once::new();
38
39fn ensure_drivers() {
40    DRIVERS_INSTALLED.call_once(sqlx::any::install_default_drivers);
41}
42
43pub struct SqlxClient {
44    pool: AnyPool,
45}
46
47impl SqlxClient {
48    pub async fn connect(url: &str) -> Result<Self, sqlx::Error> {
49        ensure_drivers();
50        let pool = AnyPool::connect(url).await?;
51        Ok(Self { pool })
52    }
53
54    pub async fn connect_with_auth(
55        url: &str,
56        username: &str,
57        password: &str,
58    ) -> Result<Self, sqlx::Error> {
59        let url = if url.contains('@') {
60            url.to_string()
61        } else {
62            let encoded_user = percent_encode(username);
63            let encoded_pass = percent_encode(password);
64            url.replacen("://", &format!("://{encoded_user}:{encoded_pass}@"), 1)
65        };
66        Self::connect(&url).await
67    }
68
69    pub async fn from_config(cfg: SqlxConfig) -> Result<Self, sqlx::Error> {
70        match (&cfg.username, &cfg.password) {
71            (Some(u), Some(p)) if !u.is_empty() || !p.is_empty() => {
72                Self::connect_with_auth(&cfg.url, u, p).await
73            }
74            _ => Self::connect(&cfg.url).await,
75        }
76    }
77
78    pub fn from_pool(pool: AnyPool) -> Self {
79        Self { pool }
80    }
81}
82
83/// 单格类型转换链：bool → i64 → i32 → f64（NaN/Inf 转字符串）→ String →
84/// Blob（base64）。此前 Blob/BYTEA 列会静默落到 Null。
85/// 注意：sqlx Any 驱动不支持时间类型，timestamp 列会在 fetch 时直接报错
86/// （AnyDriverError），不会静默——调用方可自行 CAST 成文本。
87fn cell_to_json(row: &AnyRow, col: &str) -> serde_json::Value {
88    row.try_get::<bool, _>(col)
89        .map(serde_json::Value::Bool)
90        .or_else(|_| {
91            row.try_get::<i64, _>(col)
92                .map(|n| serde_json::Value::Number(n.into()))
93        })
94        .or_else(|_| {
95            row.try_get::<i32, _>(col)
96                .map(|n| serde_json::Value::Number((n as i64).into()))
97        })
98        .or_else(|_| {
99            row.try_get::<f64, _>(col)
100                .ok()
101                .and_then(|n| {
102                    if n.is_finite() {
103                        serde_json::Number::from_f64(n).map(serde_json::Value::Number)
104                    } else if n.is_nan() {
105                        Some(serde_json::Value::String("NaN".into()))
106                    } else if n > 0.0 {
107                        Some(serde_json::Value::String("Infinity".into()))
108                    } else {
109                        Some(serde_json::Value::String("-Infinity".into()))
110                    }
111                })
112                .ok_or(())
113        })
114        .or_else(|_| row.try_get::<String, _>(col).map(serde_json::Value::String))
115        .or_else(|_| {
116            row.try_get::<Vec<u8>, _>(col).map(|b| {
117                serde_json::Value::String(base64::engine::general_purpose::STANDARD.encode(b))
118            })
119        })
120        .unwrap_or(serde_json::Value::Null)
121}
122
123fn rows_to_result(rows: Vec<AnyRow>) -> Vec<Row> {
124    if rows.is_empty() {
125        return Vec::new();
126    }
127    let columns: Vec<String> = rows[0]
128        .columns()
129        .iter()
130        .map(|c| c.name().to_string())
131        .collect();
132    rows.iter()
133        .map(|row| {
134            let values: Vec<serde_json::Value> =
135                columns.iter().map(|col| cell_to_json(row, col)).collect();
136            Row::new(columns.clone(), values)
137        })
138        .collect()
139}
140
141#[async_trait]
142impl RdbmsClient for SqlxClient {
143    async fn execute(&self, sql: &str) -> Result<u64, RdbmsError> {
144        sqlx::query(sql)
145            .execute(&self.pool)
146            .await
147            .map(|r| r.rows_affected())
148            .map_err(|e| RdbmsError::Database(e.to_string()))
149    }
150
151    async fn query(&self, sql: &str) -> Result<Vec<Row>, RdbmsError> {
152        let rows: Vec<AnyRow> = sqlx::query(sql)
153            .fetch_all(&self.pool)
154            .await
155            .map_err(|e| RdbmsError::Database(e.to_string()))?;
156        Ok(rows_to_result(rows))
157    }
158
159    async fn execute_with(
160        &self,
161        sql: &str,
162        params: &[serde_json::Value],
163    ) -> Result<u64, RdbmsError> {
164        let mut q = sqlx::query(sql);
165        for p in params {
166            q = match p {
167                serde_json::Value::String(s) => q.bind(s.as_str()),
168                serde_json::Value::Number(n) => {
169                    if let Some(i) = n.as_i64() {
170                        q.bind(i)
171                    } else if let Some(f) = n.as_f64() {
172                        q.bind(f)
173                    } else {
174                        q.bind(n.to_string())
175                    }
176                }
177                serde_json::Value::Bool(b) => q.bind(*b),
178                serde_json::Value::Null => q.bind(None::<String>),
179                _ => q.bind(p.to_string()),
180            };
181        }
182        q.execute(&self.pool)
183            .await
184            .map(|r| r.rows_affected())
185            .map_err(|e| RdbmsError::Database(e.to_string()))
186    }
187
188    async fn query_with(
189        &self,
190        sql: &str,
191        params: &[serde_json::Value],
192    ) -> Result<Vec<Row>, RdbmsError> {
193        let mut q = sqlx::query(sql);
194        for p in params {
195            q = match p {
196                serde_json::Value::String(s) => q.bind(s.as_str()),
197                serde_json::Value::Number(n) => {
198                    if let Some(i) = n.as_i64() {
199                        q.bind(i)
200                    } else if let Some(f) = n.as_f64() {
201                        q.bind(f)
202                    } else {
203                        q.bind(n.to_string())
204                    }
205                }
206                serde_json::Value::Bool(b) => q.bind(*b),
207                serde_json::Value::Null => q.bind(None::<String>),
208                _ => q.bind(p.to_string()),
209            };
210        }
211        let rows: Vec<sqlx::any::AnyRow> = q
212            .fetch_all(&self.pool)
213            .await
214            .map_err(|e| RdbmsError::Database(e.to_string()))?;
215        Ok(rows_to_result(rows))
216    }
217
218    async fn transaction(&self) -> Result<ecat_data::Transaction, RdbmsError> {
219        let tx = self
220            .pool
221            .begin()
222            .await
223            .map_err(|e| RdbmsError::Database(e.to_string()))?;
224        Ok(ecat_data::Transaction::with_inner(Box::new(
225            SqlxTransactionWrapper { inner: Some(tx) },
226        )))
227    }
228}
229
230struct SqlxTransactionWrapper {
231    inner: Option<sqlx::Transaction<'static, sqlx::Any>>,
232}
233
234#[async_trait]
235impl TransactionInner for SqlxTransactionWrapper {
236    async fn commit(&mut self) -> Result<(), RdbmsError> {
237        if let Some(tx) = self.inner.take() {
238            tx.commit()
239                .await
240                .map_err(|e| RdbmsError::Database(e.to_string()))?;
241        }
242        Ok(())
243    }
244
245    async fn rollback(&mut self) -> Result<(), RdbmsError> {
246        if let Some(tx) = self.inner.take() {
247            tx.rollback()
248                .await
249                .map_err(|e| RdbmsError::Database(e.to_string()))?;
250        }
251        Ok(())
252    }
253}
254
255#[cfg(test)]
256mod tests {
257    use super::*;
258
259    #[test]
260    fn percent_encode_special_chars() {
261        assert_eq!(percent_encode("user:pass"), "user%3Apass");
262        assert_eq!(percent_encode("a/b@c"), "a%2Fb%40c");
263        assert_eq!(percent_encode("a#b?c&d=e"), "a%23b%3Fc%26d%3De");
264        assert_eq!(percent_encode("100%"), "100%25");
265        assert_eq!(percent_encode("a+b"), "a%2Bb");
266        assert_eq!(percent_encode("hello world"), "hello%20world");
267    }
268
269    #[test]
270    fn percent_encode_no_special_chars() {
271        assert_eq!(percent_encode("simple"), "simple");
272        assert_eq!(percent_encode("user123"), "user123");
273        assert_eq!(percent_encode(""), "");
274    }
275
276    #[test]
277    fn config_deserialize_basic() {
278        let cfg: SqlxConfig =
279            serde_json::from_str(r#"{"url": "postgres://localhost/db"}"#).unwrap();
280        assert_eq!(cfg.url, "postgres://localhost/db");
281        assert!(cfg.username.is_none());
282        assert!(cfg.password.is_none());
283        assert!(cfg.tls.is_none());
284    }
285
286    #[test]
287    fn config_deserialize_with_auth() {
288        let cfg: SqlxConfig = serde_json::from_str(
289            r#"{"url": "mysql://localhost/db", "username": "root", "password": "secret"}"#,
290        )
291        .unwrap();
292        assert_eq!(cfg.url, "mysql://localhost/db");
293        assert_eq!(cfg.username.as_deref(), Some("root"));
294        assert_eq!(cfg.password.as_deref(), Some("secret"));
295    }
296
297    #[test]
298    fn config_deserialize_with_tls() {
299        let cfg: SqlxConfig = serde_json::from_str(
300            r#"{"url": "postgres://localhost/db", "tls": {"skip_verify": true}}"#,
301        )
302        .unwrap();
303        assert!(cfg.tls.is_some());
304        let tls = cfg.tls.unwrap();
305        assert_eq!(tls.skip_verify, Some(true));
306    }
307
308    #[test]
309    fn config_missing_url_is_error() {
310        let result: Result<SqlxConfig, _> = serde_json::from_str(r#"{}"#);
311        assert!(result.is_err());
312    }
313
314    #[test]
315    fn from_pool_is_constructible() {
316        // Compile-time check: SqlxClient::from_pool exists with correct signature.
317        fn _check_sig(pool: sqlx::AnyPool) -> SqlxClient {
318            SqlxClient::from_pool(pool)
319        }
320    }
321
322    /// 内存 SQLite：无外部服务即可做端到端往返。
323    /// 每次调用名唯一，避免并行测试互相干扰。
324    fn mem_sqlite(name: &str) -> String {
325        static N: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
326        let n = N.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
327        format!("sqlite:ecat-test-{name}{n}?mode=memory&cache=shared")
328    }
329
330    /// sqlx 0.8 要求先安装 Any driver（Once 保护，重复调用安全）。
331    fn init_drivers() {
332        sqlx::any::install_default_drivers();
333    }
334
335    /// 单连接池客户端：内存库随连接销毁，多连接池建的表会被后续
336    /// 连接丢失，单连接池保证同测试内所有语句命中同一库。
337    async fn single_conn_client(name: &str) -> SqlxClient {
338        init_drivers();
339        let pool = sqlx::any::AnyPoolOptions::new()
340            .max_connections(1)
341            .connect(&mem_sqlite(name))
342            .await
343            .unwrap();
344        SqlxClient::from_pool(pool)
345    }
346
347    #[tokio::test]
348    async fn connect_and_execute_query_round_trip() {
349        let client = single_conn_client("t1").await;
350        client
351            .execute("CREATE TABLE t (id INTEGER, name TEXT)")
352            .await
353            .unwrap();
354        client
355            .execute("INSERT INTO t VALUES (1, 'alice'), (2, 'bob')")
356            .await
357            .unwrap();
358        let rows = client
359            .query("SELECT id, name FROM t ORDER BY id")
360            .await
361            .unwrap();
362        assert_eq!(rows.len(), 2);
363        assert_eq!(rows[0].get("id"), Some(&serde_json::json!(1)));
364        assert_eq!(rows[0].get("name"), Some(&serde_json::json!("alice")));
365        assert_eq!(rows[1].get("id"), Some(&serde_json::json!(2)));
366        assert_eq!(rows[1].get("missing"), None);
367    }
368
369    #[tokio::test]
370    async fn execute_with_binds_all_json_value_types() {
371        let client = single_conn_client("t1").await;
372        client
373            .execute("CREATE TABLE t (s TEXT, i INTEGER, f REAL, b INTEGER, n TEXT)")
374            .await
375            .unwrap();
376        let affected = client
377            .execute_with(
378                "INSERT INTO t VALUES (?, ?, ?, ?, ?)",
379                &[
380                    serde_json::json!("str"),
381                    serde_json::json!(42),
382                    serde_json::json!(1.5),
383                    serde_json::json!(true),
384                    serde_json::Value::Null,
385                ],
386            )
387            .await
388            .unwrap();
389        assert_eq!(affected, 1);
390        let rows = client.query("SELECT * FROM t").await.unwrap();
391        assert_eq!(rows.len(), 1);
392        assert_eq!(rows[0].get("s"), Some(&serde_json::json!("str")));
393        assert_eq!(rows[0].get("i"), Some(&serde_json::json!(42)));
394        assert_eq!(rows[0].get("f"), Some(&serde_json::json!(1.5)));
395        // SQLite 无布尔类型：true 绑定后回读为整数 1
396        assert_eq!(rows[0].get("b"), Some(&serde_json::json!(1)));
397        assert_eq!(rows[0].get("n"), Some(&serde_json::Value::Null));
398    }
399
400    #[tokio::test]
401    async fn query_with_parameterized_sql() {
402        let client = single_conn_client("t1").await;
403        client.execute("CREATE TABLE t (name TEXT)").await.unwrap();
404        client
405            .execute_with("INSERT INTO t VALUES (?)", &[serde_json::json!("x")])
406            .await
407            .unwrap();
408        let rows = client
409            .query_with(
410                "SELECT name FROM t WHERE name = ?",
411                &[serde_json::json!("x")],
412            )
413            .await
414            .unwrap();
415        assert_eq!(rows.len(), 1);
416        assert_eq!(rows[0].get("name"), Some(&serde_json::json!("x")));
417    }
418
419    #[tokio::test]
420    async fn cell_to_json_encodes_blob_as_base64() {
421        let client = single_conn_client("t1").await;
422        client.execute("CREATE TABLE t (data BLOB)").await.unwrap();
423        // 真实 BLOB 字节（0x01 0x02 0x03），绕过 bind（bind 会把 JSON 值绑成文本）
424        client
425            .execute("INSERT INTO t VALUES (x'010203')")
426            .await
427            .unwrap();
428        let rows = client.query("SELECT data FROM t").await.unwrap();
429        assert_eq!(rows[0].get("data"), Some(&serde_json::json!("AQID")));
430    }
431
432    #[tokio::test]
433    async fn connect_with_auth_does_not_inject_into_non_url_scheme() {
434        // sqlite URL 不含 "://"，凭据注入分支不生效，直接连接成功
435        init_drivers();
436        let client = SqlxClient::connect_with_auth(&mem_sqlite("t2"), "user", "pass")
437            .await
438            .unwrap();
439        client.execute("SELECT 1").await.unwrap();
440    }
441
442    #[tokio::test]
443    async fn from_config_without_auth_connects_and_queries() {
444        // 不手动 init_drivers：验证 from_config 自身完成驱动安装后
445        // 能在内存 sqlite 上实际执行查询
446        // （注意：多连接池下 mode=memory 库随连接关闭而销毁，故只做单条查询）
447        let cfg = SqlxConfig {
448            url: mem_sqlite("t3"),
449            username: None,
450            password: None,
451            tls: None,
452        };
453        let client = SqlxClient::from_config(cfg).await.unwrap();
454        let rows = client.query("SELECT 1 AS one").await.unwrap();
455        assert_eq!(rows.len(), 1);
456        assert_eq!(rows[0].get("one"), Some(&serde_json::json!(1)));
457    }
458
459    #[tokio::test]
460    async fn from_config_with_empty_credentials_connects_plain() {
461        // (Some(""), Some("")) 或 (Some(""), None) 都走无认证分支；
462        // from_config 内部会装驱动，无需手动 init_drivers
463        for (u, p) in [
464            (Some("".to_string()), Some("".to_string())),
465            (Some("".to_string()), None),
466        ] {
467            let cfg = SqlxConfig {
468                url: mem_sqlite("t3"),
469                username: u,
470                password: p,
471                tls: None,
472            };
473            let client = SqlxClient::from_config(cfg).await.unwrap();
474            client.execute("SELECT 1").await.unwrap();
475        }
476    }
477}
ecat_data_sqlx/lib.rs

ecat_data_sqlx/
lib.rs
