sz_orm_sqlx/
any.rs

1//! sqlx 后端适配器实现
2//!
3//! 为 MySQL、PostgreSQL、SQLite 分别实现 Connection 和 ConnectionFactory。
4//! 不使用 sqlx::Any 以避免其类型限制和生命周期问题。
5//!
6//! 关键设计：
7//! Connection trait 已手动解糖（不使用 `#[async_trait]`），所有 async 方法
8//! 使用单一生命周期 `'a`（绑定 `&'a mut self` 和 `&'a str`），而非 HRTB。
9//! 这样 sqlx::Executor 对 `&'c mut XxxConnection` 的 impl（针对具体 `'c`）
10//! 即可满足约束，避免 "implementation of Executor is not general enough" 错误。
11
12use async_trait::async_trait;
13use sqlx::{Column, Executor, Row};
14use std::collections::HashMap;
15use std::future::Future;
16use std::pin::Pin;
17use std::sync::Arc;
18use sz_orm_core::{Connection, ConnectionFactory, DbError, Value};
19
20use crate::error::map_sqlx_error;
21
22/// 判断 SQL 是否需要走 raw_sql 路径
23/// MySQL prepared statement 协议不支持 BEGIN/COMMIT/ROLLBACK/SAVEPOINT 等命令
24fn needs_raw_sql(sql: &str) -> bool {
25    let trimmed = sql.trim_start();
26    let upper = trimmed.to_uppercase();
27    upper.starts_with("BEGIN")
28        || upper.starts_with("COMMIT")
29        || upper.starts_with("ROLLBACK")
30        || upper.starts_with("SAVEPOINT")
31        || upper.starts_with("RELEASE")
32        || upper.starts_with("SET ")
33        || upper.starts_with("USE ")
34        || upper.starts_with("START TRANSACTION")
35}
36
37// ===================== SQLite 适配器 =====================
38
39// 注：sqlx 0.9 起 Executor trait 要求 'static lifetime，
40// 原先的 execute_sqlite_boxed / query_sqlite_boxed 已内联到调用点。
41// 见 SqlxSqliteConnection::execute / query 实现。
42
43/// 将 SqliteRow 转换为 Value（按列序号）
44/// 使用列类型信息决定解码类型，避免 bool/int 混淆
45fn row_to_value_sqlite(row: &sqlx::sqlite::SqliteRow, ordinal: usize) -> Value {
46    use sqlx::TypeInfo;
47    let type_name = row.columns()[ordinal].type_info().name();
48    match type_name {
49        "BOOLEAN" => match row.try_get::<Option<bool>, usize>(ordinal) {
50            Ok(v) => v.map(Value::Bool).unwrap_or(Value::Null),
51            Err(_) => Value::Null,
52        },
53        "INTEGER" => match row.try_get::<Option<i64>, usize>(ordinal) {
54            Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
55            Err(_) => match row.try_get::<Option<i32>, usize>(ordinal) {
56                Ok(v) => v.map(Value::I32).unwrap_or(Value::Null),
57                Err(_) => Value::Null,
58            },
59        },
60        "REAL" => match row.try_get::<Option<f64>, usize>(ordinal) {
61            Ok(v) => v.map(Value::F64).unwrap_or(Value::Null),
62            Err(_) => match row.try_get::<Option<f32>, usize>(ordinal) {
63                Ok(v) => v.map(Value::F32).unwrap_or(Value::Null),
64                Err(_) => Value::Null,
65            },
66        },
67        "TEXT" => match row.try_get::<Option<String>, usize>(ordinal) {
68            Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
69            Err(_) => Value::Null,
70        },
71        "BLOB" => match row.try_get::<Option<Vec<u8>>, usize>(ordinal) {
72            Ok(v) => v.map(Value::Bytes).unwrap_or(Value::Null),
73            Err(_) => Value::Null,
74        },
75        _ => {
76            // 未知类型，按 bool → i64 → f64 → String 顺序回退
77            if let Ok(v) = row.try_get::<Option<bool>, usize>(ordinal) {
78                return v.map(Value::Bool).unwrap_or(Value::Null);
79            }
80            if let Ok(v) = row.try_get::<Option<i64>, usize>(ordinal) {
81                return v.map(Value::I64).unwrap_or(Value::Null);
82            }
83            if let Ok(v) = row.try_get::<Option<f64>, usize>(ordinal) {
84                return v.map(Value::F64).unwrap_or(Value::Null);
85            }
86            if let Ok(v) = row.try_get::<Option<String>, usize>(ordinal) {
87                return v.map(Value::String).unwrap_or(Value::Null);
88            }
89            Value::Null
90        }
91    }
92}
93
94pub struct SqlitePoolHandle {
95    pool: sqlx::SqlitePool,
96}
97
98impl SqlitePoolHandle {
99    pub async fn connect(url: &str) -> Result<Self, DbError> {
100        let pool = sqlx::pool::PoolOptions::<sqlx::Sqlite>::new()
101            .max_connections(10)
102            .acquire_timeout(std::time::Duration::from_secs(30))
103            .idle_timeout(Some(std::time::Duration::from_secs(600)))
104            .max_lifetime(Some(std::time::Duration::from_secs(1800)))
105            .connect(url)
106            .await
107            .map_err(map_sqlx_error)?;
108        Ok(Self { pool })
109    }
110
111    pub fn from_pool(pool: sqlx::SqlitePool) -> Self {
112        Self { pool }
113    }
114
115    pub fn pool(&self) -> &sqlx::SqlitePool {
116        &self.pool
117    }
118}
119
120pub struct SqlxSqliteConnectionFactory {
121    pool: Arc<SqlitePoolHandle>,
122}
123
124impl SqlxSqliteConnectionFactory {
125    pub fn new(pool: Arc<SqlitePoolHandle>) -> Self {
126        Self { pool }
127    }
128}
129
130#[async_trait]
131impl ConnectionFactory for SqlxSqliteConnectionFactory {
132    async fn create(&self) -> Result<Box<dyn Connection>, DbError> {
133        let conn = self.pool.pool.acquire().await.map_err(map_sqlx_error)?;
134        Ok(Box::new(SqlxSqliteConnection {
135            conn: Some(conn),
136            connected: true,
137            in_transaction: false,
138        }))
139    }
140}
141
142pub struct SqlxSqliteConnection {
143    conn: Option<sqlx::pool::PoolConnection<sqlx::Sqlite>>,
144    connected: bool,
145    in_transaction: bool,
146}
147
148impl Connection for SqlxSqliteConnection {
149    fn execute<'a>(
150        &'a mut self,
151        sql: &'a str,
152    ) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
153        Box::pin(async move {
154            let mut pool_conn = self
155                .conn
156                .take()
157                .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
158            // sqlx 0.9: PoolConnection 不再实现 Executor，需通过 DerefMut 解引用到内部连接
159            // sqlx 0.9: SqlSafeStr 只对 &'static str 直接实现，非 'static 的 &str 需用 AssertSqlSafe 包装
160            let result = if needs_raw_sql(sql) {
161                (&mut *pool_conn)
162                    .execute(sqlx::raw_sql(sqlx::AssertSqlSafe(sql)))
163                    .await
164            } else {
165                (&mut *pool_conn).execute(sqlx::AssertSqlSafe(sql)).await
166            };
167            self.conn = Some(pool_conn);
168
169            match result {
170                Ok(r) => Ok(r.rows_affected()),
171                Err(e) => {
172                    let db_err = map_sqlx_error(e);
173                    if matches!(db_err, DbError::ConnectionError(_) | DbError::IoError(_)) {
174                        self.connected = false;
175                    }
176                    Err(db_err)
177                }
178            }
179        })
180    }
181
182    fn query<'a>(
183        &'a mut self,
184        sql: &'a str,
185    ) -> Pin<Box<dyn Future<Output = Result<Vec<HashMap<String, Value>>, DbError>> + Send + 'a>>
186    {
187        Box::pin(async move {
188            let mut pool_conn = self
189                .conn
190                .take()
191                .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
192            // sqlx 0.9: PoolConnection 不再实现 Executor，需通过 DerefMut 解引用到内部连接
193            // sqlx 0.9: SqlSafeStr 只对 &'static str 直接实现，非 'static 的 &str 需用 AssertSqlSafe 包装
194            let rows_result = (&mut *pool_conn).fetch_all(sqlx::AssertSqlSafe(sql)).await;
195            self.conn = Some(pool_conn);
196
197            let rows = rows_result.map_err(map_sqlx_error)?;
198            let mut result = Vec::with_capacity(rows.len());
199            for row in rows {
200                let mut record = HashMap::new();
201                for col in row.columns() {
202                    let name = col.name().to_string();
203                    let ordinal = col.ordinal();
204                    let value = row_to_value_sqlite(&row, ordinal);
205                    record.insert(name, value);
206                }
207                result.push(record);
208            }
209            Ok(result)
210        })
211    }
212
213    fn begin_transaction<'a>(
214        &'a mut self,
215    ) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
216        Box::pin(async move {
217            if self.in_transaction {
218                return Err(DbError::Internal("transaction already started".to_string()));
219            }
220            self.execute("BEGIN").await?;
221            self.in_transaction = true;
222            Ok(())
223        })
224    }
225
226    fn commit<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
227        Box::pin(async move {
228            if self.in_transaction {
229                self.execute("COMMIT").await?;
230                self.in_transaction = false;
231            }
232            Ok(())
233        })
234    }
235
236    fn rollback<'a>(
237        &'a mut self,
238    ) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
239        Box::pin(async move {
240            if self.in_transaction {
241                let result = self.execute("ROLLBACK").await;
242                self.in_transaction = false;
243                result.map(|_| ())
244            } else {
245                Ok(())
246            }
247        })
248    }
249
250    fn is_connected(&self) -> bool {
251        self.connected
252    }
253
254    fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
255        Box::pin(async move {
256            match self.execute("SELECT 1").await {
257                Ok(_) => true,
258                Err(_) => {
259                    self.connected = false;
260                    false
261                }
262            }
263        })
264    }
265
266    fn close<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
267        Box::pin(async move {
268            if let Some(conn) = self.conn.take() {
269                drop(conn);
270            }
271            self.connected = false;
272            self.in_transaction = false;
273            Ok(())
274        })
275    }
276}
277
278impl Drop for SqlxSqliteConnection {
279    fn drop(&mut self) {
280        if let Some(conn) = self.conn.take() {
281            drop(conn);
282        }
283    }
284}
285
286// ===================== MySQL 适配器 =====================
287
288// 注：sqlx 0.9 起 Executor trait 要求 'static lifetime，
289// 原先的 execute_mysql_boxed / query_mysql_boxed 已内联到调用点。
290
291fn row_to_value_mysql(row: &sqlx::mysql::MySqlRow, ordinal: usize) -> Value {
292    use sqlx::TypeInfo;
293    let type_name = row.columns()[ordinal].type_info().name();
294    match type_name {
295        "BOOLEAN" => match row.try_get::<Option<bool>, usize>(ordinal) {
296            Ok(v) => v.map(Value::Bool).unwrap_or(Value::Null),
297            Err(_) => Value::Null,
298        },
299        "TINYINT" | "TINYINT UNSIGNED" => match row.try_get::<Option<i8>, usize>(ordinal) {
300            Ok(v) => v.map(Value::I8).unwrap_or(Value::Null),
301            Err(_) => match row.try_get::<Option<u8>, usize>(ordinal) {
302                Ok(v) => v.map(Value::U8).unwrap_or(Value::Null),
303                Err(_) => Value::Null,
304            },
305        },
306        "SMALLINT" | "SMALLINT UNSIGNED" => match row.try_get::<Option<i16>, usize>(ordinal) {
307            Ok(v) => v.map(Value::I16).unwrap_or(Value::Null),
308            Err(_) => match row.try_get::<Option<u16>, usize>(ordinal) {
309                Ok(v) => v.map(Value::U16).unwrap_or(Value::Null),
310                Err(_) => Value::Null,
311            },
312        },
313        "INT" | "INT UNSIGNED" | "MEDIUMINT" | "MEDIUMINT UNSIGNED" => {
314            match row.try_get::<Option<i32>, usize>(ordinal) {
315                Ok(v) => v.map(Value::I32).unwrap_or(Value::Null),
316                Err(_) => match row.try_get::<Option<u32>, usize>(ordinal) {
317                    Ok(v) => v.map(Value::U32).unwrap_or(Value::Null),
318                    Err(_) => Value::Null,
319                },
320            }
321        }
322        "BIGINT" | "BIGINT UNSIGNED" => match row.try_get::<Option<i64>, usize>(ordinal) {
323            Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
324            Err(_) => match row.try_get::<Option<u64>, usize>(ordinal) {
325                Ok(v) => v.map(Value::U64).unwrap_or(Value::Null),
326                Err(_) => Value::Null,
327            },
328        },
329        "FLOAT" => match row.try_get::<Option<f32>, usize>(ordinal) {
330            Ok(v) => v.map(Value::F32).unwrap_or(Value::Null),
331            Err(_) => Value::Null,
332        },
333        "DOUBLE" => match row.try_get::<Option<f64>, usize>(ordinal) {
334            Ok(v) => v.map(Value::F64).unwrap_or(Value::Null),
335            Err(_) => Value::Null,
336        },
337        "VARCHAR" | "TEXT" | "CHAR" | "TINYTEXT" | "MEDIUMTEXT" | "LONGTEXT" | "ENUM" => {
338            match row.try_get::<Option<String>, usize>(ordinal) {
339                Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
340                Err(_) => Value::Null,
341            }
342        }
343        "BLOB" | "TINYBLOB" | "MEDIUMBLOB" | "LONGBLOB" | "BINARY" | "VARBINARY" => {
344            match row.try_get::<Option<Vec<u8>>, usize>(ordinal) {
345                Ok(v) => v.map(Value::Bytes).unwrap_or(Value::Null),
346                Err(_) => Value::Null,
347            }
348        }
349        // DECIMAL/NUMERIC 使用 rust_decimal 解码
350        "DECIMAL" | "NUMERIC" | "NEWDECIMAL" => {
351            match row.try_get::<Option<rust_decimal::Decimal>, usize>(ordinal) {
352                Ok(Some(v)) => Value::F64(v.to_string().parse::<f64>().unwrap_or(0.0)),
353                Ok(None) => Value::Null,
354                Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
355                    Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
356                    Err(_) => Value::Null,
357                },
358            }
359        }
360        _ => {
361            // 未知类型回退：i64 → f64 → bool → String
362            if let Ok(v) = row.try_get::<Option<i64>, usize>(ordinal) {
363                return v.map(Value::I64).unwrap_or(Value::Null);
364            }
365            if let Ok(v) = row.try_get::<Option<f64>, usize>(ordinal) {
366                return v.map(Value::F64).unwrap_or(Value::Null);
367            }
368            if let Ok(v) = row.try_get::<Option<bool>, usize>(ordinal) {
369                return v.map(Value::Bool).unwrap_or(Value::Null);
370            }
371            if let Ok(v) = row.try_get::<Option<String>, usize>(ordinal) {
372                return v.map(Value::String).unwrap_or(Value::Null);
373            }
374            Value::Null
375        }
376    }
377}
378
379pub struct MySqlPoolHandle {
380    pool: sqlx::MySqlPool,
381}
382
383impl MySqlPoolHandle {
384    pub async fn connect(url: &str) -> Result<Self, DbError> {
385        let pool = sqlx::pool::PoolOptions::<sqlx::MySql>::new()
386            .max_connections(10)
387            .acquire_timeout(std::time::Duration::from_secs(30))
388            .idle_timeout(Some(std::time::Duration::from_secs(600)))
389            .max_lifetime(Some(std::time::Duration::from_secs(1800)))
390            .connect(url)
391            .await
392            .map_err(map_sqlx_error)?;
393        Ok(Self { pool })
394    }
395
396    pub fn from_pool(pool: sqlx::MySqlPool) -> Self {
397        Self { pool }
398    }
399
400    pub fn pool(&self) -> &sqlx::MySqlPool {
401        &self.pool
402    }
403}
404
405pub struct SqlxMySqlConnectionFactory {
406    pool: Arc<MySqlPoolHandle>,
407}
408
409impl SqlxMySqlConnectionFactory {
410    pub fn new(pool: Arc<MySqlPoolHandle>) -> Self {
411        Self { pool }
412    }
413}
414
415#[async_trait]
416impl ConnectionFactory for SqlxMySqlConnectionFactory {
417    async fn create(&self) -> Result<Box<dyn Connection>, DbError> {
418        let conn = self.pool.pool.acquire().await.map_err(map_sqlx_error)?;
419        Ok(Box::new(SqlxMySqlConnection {
420            conn: Some(conn),
421            connected: true,
422            in_transaction: false,
423        }))
424    }
425}
426
427pub struct SqlxMySqlConnection {
428    conn: Option<sqlx::pool::PoolConnection<sqlx::MySql>>,
429    connected: bool,
430    in_transaction: bool,
431}
432
433impl Connection for SqlxMySqlConnection {
434    fn execute<'a>(
435        &'a mut self,
436        sql: &'a str,
437    ) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
438        Box::pin(async move {
439            let mut pool_conn = self
440                .conn
441                .take()
442                .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
443            // sqlx 0.9: PoolConnection 不再实现 Executor，需通过 DerefMut 解引用到内部连接
444            // sqlx 0.9: SqlSafeStr 只对 &'static str 直接实现，非 'static 的 &str 需用 AssertSqlSafe 包装
445            let result = if needs_raw_sql(sql) {
446                (&mut *pool_conn)
447                    .execute(sqlx::raw_sql(sqlx::AssertSqlSafe(sql)))
448                    .await
449            } else {
450                (&mut *pool_conn).execute(sqlx::AssertSqlSafe(sql)).await
451            };
452            self.conn = Some(pool_conn);
453
454            match result {
455                Ok(r) => Ok(r.rows_affected()),
456                Err(e) => {
457                    let db_err = map_sqlx_error(e);
458                    if matches!(db_err, DbError::ConnectionError(_) | DbError::IoError(_)) {
459                        self.connected = false;
460                    }
461                    Err(db_err)
462                }
463            }
464        })
465    }
466
467    fn query<'a>(
468        &'a mut self,
469        sql: &'a str,
470    ) -> Pin<Box<dyn Future<Output = Result<Vec<HashMap<String, Value>>, DbError>> + Send + 'a>>
471    {
472        Box::pin(async move {
473            let mut pool_conn = self
474                .conn
475                .take()
476                .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
477            // sqlx 0.9: PoolConnection 不再实现 Executor，需通过 DerefMut 解引用到内部连接
478            // sqlx 0.9: SqlSafeStr 只对 &'static str 直接实现，非 'static 的 &str 需用 AssertSqlSafe 包装
479            let rows_result = (&mut *pool_conn).fetch_all(sqlx::AssertSqlSafe(sql)).await;
480            self.conn = Some(pool_conn);
481
482            let rows = rows_result.map_err(map_sqlx_error)?;
483            let mut result = Vec::with_capacity(rows.len());
484            for row in rows {
485                let mut record = HashMap::new();
486                for col in row.columns() {
487                    let name = col.name().to_string();
488                    let ordinal = col.ordinal();
489                    let value = row_to_value_mysql(&row, ordinal);
490                    record.insert(name, value);
491                }
492                result.push(record);
493            }
494            Ok(result)
495        })
496    }
497
498    fn begin_transaction<'a>(
499        &'a mut self,
500    ) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
501        Box::pin(async move {
502            if self.in_transaction {
503                return Err(DbError::Internal("transaction already started".to_string()));
504            }
505            self.execute("BEGIN").await?;
506            self.in_transaction = true;
507            Ok(())
508        })
509    }
510
511    fn commit<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
512        Box::pin(async move {
513            if self.in_transaction {
514                self.execute("COMMIT").await?;
515                self.in_transaction = false;
516            }
517            Ok(())
518        })
519    }
520
521    fn rollback<'a>(
522        &'a mut self,
523    ) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
524        Box::pin(async move {
525            if self.in_transaction {
526                let result = self.execute("ROLLBACK").await;
527                self.in_transaction = false;
528                result.map(|_| ())
529            } else {
530                Ok(())
531            }
532        })
533    }
534
535    fn is_connected(&self) -> bool {
536        self.connected
537    }
538
539    fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
540        Box::pin(async move {
541            match self.execute("SELECT 1").await {
542                Ok(_) => true,
543                Err(_) => {
544                    self.connected = false;
545                    false
546                }
547            }
548        })
549    }
550
551    fn close<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
552        Box::pin(async move {
553            if let Some(conn) = self.conn.take() {
554                drop(conn);
555            }
556            self.connected = false;
557            self.in_transaction = false;
558            Ok(())
559        })
560    }
561}
562
563impl Drop for SqlxMySqlConnection {
564    fn drop(&mut self) {
565        if let Some(conn) = self.conn.take() {
566            drop(conn);
567        }
568    }
569}
570
571// ===================== PostgreSQL 适配器 =====================
572
573// 注：sqlx 0.9 起 Executor trait 要求 'static lifetime，
574// 原先的 execute_pg_boxed / query_pg_boxed 已内联到调用点。
575
576fn row_to_value_pg(row: &sqlx::postgres::PgRow, ordinal: usize) -> Value {
577    use sqlx::TypeInfo;
578    let type_name = row.columns()[ordinal].type_info().name();
579    match type_name {
580        "BOOL" => match row.try_get::<Option<bool>, usize>(ordinal) {
581            Ok(v) => v.map(Value::Bool).unwrap_or(Value::Null),
582            Err(_) => Value::Null,
583        },
584        "INT2" => match row.try_get::<Option<i16>, usize>(ordinal) {
585            Ok(v) => v.map(Value::I16).unwrap_or(Value::Null),
586            Err(_) => Value::Null,
587        },
588        "INT4" | "OID" => match row.try_get::<Option<i32>, usize>(ordinal) {
589            Ok(v) => v.map(Value::I32).unwrap_or(Value::Null),
590            Err(_) => Value::Null,
591        },
592        "INT8" => match row.try_get::<Option<i64>, usize>(ordinal) {
593            Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
594            Err(_) => Value::Null,
595        },
596        "FLOAT4" => match row.try_get::<Option<f32>, usize>(ordinal) {
597            Ok(v) => v.map(Value::F32).unwrap_or(Value::Null),
598            Err(_) => Value::Null,
599        },
600        "FLOAT8" => match row.try_get::<Option<f64>, usize>(ordinal) {
601            Ok(v) => v.map(Value::F64).unwrap_or(Value::Null),
602            Err(_) => Value::Null,
603        },
604        "TEXT" | "VARCHAR" | "CHAR" | "NAME" => match row.try_get::<Option<String>, usize>(ordinal)
605        {
606            Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
607            Err(_) => Value::Null,
608        },
609        "BYTEA" => match row.try_get::<Option<Vec<u8>>, usize>(ordinal) {
610            Ok(v) => v.map(Value::Bytes).unwrap_or(Value::Null),
611            Err(_) => Value::Null,
612        },
613        "NUMERIC" => match row.try_get::<Option<rust_decimal::Decimal>, usize>(ordinal) {
614            Ok(Some(v)) => Value::F64(v.to_string().parse::<f64>().unwrap_or(0.0)),
615            Ok(None) => Value::Null,
616            Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
617                Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
618                Err(_) => Value::Null,
619            },
620        },
621        _ => {
622            // 未知类型回退
623            if let Ok(v) = row.try_get::<Option<i64>, usize>(ordinal) {
624                return v.map(Value::I64).unwrap_or(Value::Null);
625            }
626            if let Ok(v) = row.try_get::<Option<f64>, usize>(ordinal) {
627                return v.map(Value::F64).unwrap_or(Value::Null);
628            }
629            if let Ok(v) = row.try_get::<Option<bool>, usize>(ordinal) {
630                return v.map(Value::Bool).unwrap_or(Value::Null);
631            }
632            if let Ok(v) = row.try_get::<Option<String>, usize>(ordinal) {
633                return v.map(Value::String).unwrap_or(Value::Null);
634            }
635            Value::Null
636        }
637    }
638}
639
640pub struct PgPoolHandle {
641    pool: sqlx::PgPool,
642}
643
644impl PgPoolHandle {
645    pub async fn connect(url: &str) -> Result<Self, DbError> {
646        let pool = sqlx::pool::PoolOptions::<sqlx::Postgres>::new()
647            .max_connections(10)
648            .acquire_timeout(std::time::Duration::from_secs(30))
649            .idle_timeout(Some(std::time::Duration::from_secs(600)))
650            .max_lifetime(Some(std::time::Duration::from_secs(1800)))
651            .connect(url)
652            .await
653            .map_err(map_sqlx_error)?;
654        Ok(Self { pool })
655    }
656
657    pub fn from_pool(pool: sqlx::PgPool) -> Self {
658        Self { pool }
659    }
660
661    pub fn pool(&self) -> &sqlx::PgPool {
662        &self.pool
663    }
664}
665
666pub struct SqlxPgConnectionFactory {
667    pool: Arc<PgPoolHandle>,
668}
669
670impl SqlxPgConnectionFactory {
671    pub fn new(pool: Arc<PgPoolHandle>) -> Self {
672        Self { pool }
673    }
674}
675
676#[async_trait]
677impl ConnectionFactory for SqlxPgConnectionFactory {
678    async fn create(&self) -> Result<Box<dyn Connection>, DbError> {
679        let conn = self.pool.pool.acquire().await.map_err(map_sqlx_error)?;
680        Ok(Box::new(SqlxPgConnection {
681            conn: Some(conn),
682            connected: true,
683            in_transaction: false,
684        }))
685    }
686}
687
688pub struct SqlxPgConnection {
689    conn: Option<sqlx::pool::PoolConnection<sqlx::Postgres>>,
690    connected: bool,
691    in_transaction: bool,
692}
693
694impl Connection for SqlxPgConnection {
695    fn execute<'a>(
696        &'a mut self,
697        sql: &'a str,
698    ) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
699        Box::pin(async move {
700            let mut pool_conn = self
701                .conn
702                .take()
703                .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
704            // sqlx 0.9: PoolConnection 不再实现 Executor，需通过 DerefMut 解引用到内部连接
705            // sqlx 0.9: SqlSafeStr 只对 &'static str 直接实现，非 'static 的 &str 需用 AssertSqlSafe 包装
706            let result = if needs_raw_sql(sql) {
707                (&mut *pool_conn)
708                    .execute(sqlx::raw_sql(sqlx::AssertSqlSafe(sql)))
709                    .await
710            } else {
711                (&mut *pool_conn).execute(sqlx::AssertSqlSafe(sql)).await
712            };
713            self.conn = Some(pool_conn);
714
715            match result {
716                Ok(r) => Ok(r.rows_affected()),
717                Err(e) => {
718                    let db_err = map_sqlx_error(e);
719                    if matches!(db_err, DbError::ConnectionError(_) | DbError::IoError(_)) {
720                        self.connected = false;
721                    }
722                    Err(db_err)
723                }
724            }
725        })
726    }
727
728    fn query<'a>(
729        &'a mut self,
730        sql: &'a str,
731    ) -> Pin<Box<dyn Future<Output = Result<Vec<HashMap<String, Value>>, DbError>> + Send + 'a>>
732    {
733        Box::pin(async move {
734            let mut pool_conn = self
735                .conn
736                .take()
737                .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
738            // sqlx 0.9: PoolConnection 不再实现 Executor，需通过 DerefMut 解引用到内部连接
739            // sqlx 0.9: SqlSafeStr 只对 &'static str 直接实现，非 'static 的 &str 需用 AssertSqlSafe 包装
740            let rows_result = (&mut *pool_conn).fetch_all(sqlx::AssertSqlSafe(sql)).await;
741            self.conn = Some(pool_conn);
742
743            let rows = rows_result.map_err(map_sqlx_error)?;
744            let mut result = Vec::with_capacity(rows.len());
745            for row in rows {
746                let mut record = HashMap::new();
747                for col in row.columns() {
748                    let name = col.name().to_string();
749                    let ordinal = col.ordinal();
750                    let value = row_to_value_pg(&row, ordinal);
751                    record.insert(name, value);
752                }
753                result.push(record);
754            }
755            Ok(result)
756        })
757    }
758
759    fn begin_transaction<'a>(
760        &'a mut self,
761    ) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
762        Box::pin(async move {
763            if self.in_transaction {
764                return Err(DbError::Internal("transaction already started".to_string()));
765            }
766            self.execute("BEGIN").await?;
767            self.in_transaction = true;
768            Ok(())
769        })
770    }
771
772    fn commit<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
773        Box::pin(async move {
774            if self.in_transaction {
775                self.execute("COMMIT").await?;
776                self.in_transaction = false;
777            }
778            Ok(())
779        })
780    }
781
782    fn rollback<'a>(
783        &'a mut self,
784    ) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
785        Box::pin(async move {
786            if self.in_transaction {
787                let result = self.execute("ROLLBACK").await;
788                self.in_transaction = false;
789                result.map(|_| ())
790            } else {
791                Ok(())
792            }
793        })
794    }
795
796    fn is_connected(&self) -> bool {
797        self.connected
798    }
799
800    fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
801        Box::pin(async move {
802            match self.execute("SELECT 1").await {
803                Ok(_) => true,
804                Err(_) => {
805                    self.connected = false;
806                    false
807                }
808            }
809        })
810    }
811
812    fn close<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
813        Box::pin(async move {
814            if let Some(conn) = self.conn.take() {
815                drop(conn);
816            }
817            self.connected = false;
818            self.in_transaction = false;
819            Ok(())
820        })
821    }
822}
823
824impl Drop for SqlxPgConnection {
825    fn drop(&mut self) {
826        if let Some(conn) = self.conn.take() {
827            drop(conn);
828        }
829    }
830}
sz_orm_sqlx/any.rs

sz_orm_sqlx/
any.rs
