1use async_trait::async_trait;
13use futures::StreamExt;
14use sqlx::{Column, Executor, Row, TypeInfo};
15use std::collections::HashMap;
16use std::future::Future;
17use std::pin::Pin;
18use std::str::FromStr;
19use std::sync::Arc;
20use sz_orm_core::{
21 ColType, Connection, ConnectionFactory, DbError, DbType, QueryRows, QueryValues, Value,
22};
23
24use crate::error::map_sqlx_error;
25
26fn needs_raw_sql(sql: &str) -> bool {
29 let trimmed = sql.trim_start();
30 let upper = trimmed.to_uppercase();
31 upper.starts_with("BEGIN")
32 || upper.starts_with("COMMIT")
33 || upper.starts_with("ROLLBACK")
34 || upper.starts_with("SAVEPOINT")
35 || upper.starts_with("RELEASE")
36 || upper.starts_with("SET ")
37 || upper.starts_with("USE ")
38 || upper.starts_with("START TRANSACTION")
39}
40
41fn row_to_value_with_coltype_sqlite(
59 row: &sqlx::sqlite::SqliteRow,
60 ordinal: usize,
61 col_type: ColType,
62) -> Value {
63 match col_type {
64 ColType::Bool => match row.try_get::<Option<bool>, usize>(ordinal) {
65 Ok(v) => v.map(Value::Bool).unwrap_or(Value::Null),
66 Err(_) => Value::Null,
67 },
68 ColType::I8 => match row.try_get::<Option<i8>, usize>(ordinal) {
69 Ok(v) => v.map(Value::I8).unwrap_or(Value::Null),
70 Err(_) => Value::Null,
71 },
72 ColType::I16 => match row.try_get::<Option<i16>, usize>(ordinal) {
73 Ok(v) => v.map(Value::I16).unwrap_or(Value::Null),
74 Err(_) => Value::Null,
75 },
76 ColType::I32 => match row.try_get::<Option<i32>, usize>(ordinal) {
77 Ok(v) => v.map(Value::I32).unwrap_or(Value::Null),
78 Err(_) => Value::Null,
79 },
80 ColType::I64 => match row.try_get::<Option<i64>, usize>(ordinal) {
81 Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
82 Err(_) => Value::Null,
83 },
84 ColType::U8 => match row.try_get::<Option<u8>, usize>(ordinal) {
85 Ok(v) => v.map(Value::U8).unwrap_or(Value::Null),
86 Err(_) => Value::Null,
87 },
88 ColType::U16 => match row.try_get::<Option<u16>, usize>(ordinal) {
89 Ok(v) => v.map(Value::U16).unwrap_or(Value::Null),
90 Err(_) => Value::Null,
91 },
92 ColType::U32 => match row.try_get::<Option<u32>, usize>(ordinal) {
93 Ok(v) => v.map(Value::U32).unwrap_or(Value::Null),
94 Err(_) => Value::Null,
95 },
96 ColType::U64 => match row.try_get::<Option<i64>, usize>(ordinal) {
98 Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
99 Err(_) => Value::Null,
100 },
101 ColType::F32 => match row.try_get::<Option<f32>, usize>(ordinal) {
102 Ok(v) => v.map(Value::F32).unwrap_or(Value::Null),
103 Err(_) => Value::Null,
104 },
105 ColType::F64 => match row.try_get::<Option<f64>, usize>(ordinal) {
106 Ok(v) => v.map(Value::F64).unwrap_or(Value::Null),
107 Err(_) => Value::Null,
108 },
109 ColType::Decimal => match row.try_get::<Option<String>, usize>(ordinal) {
110 Ok(v) => v.map(Value::Decimal).unwrap_or(Value::Null),
111 Err(_) => Value::Null,
112 },
113 ColType::String => match row.try_get::<Option<String>, usize>(ordinal) {
114 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
115 Err(_) => Value::Null,
116 },
117 ColType::Bytes => match row.try_get::<Option<Vec<u8>>, usize>(ordinal) {
118 Ok(v) => v.map(Value::Bytes).unwrap_or(Value::Null),
119 Err(_) => Value::Null,
120 },
121 ColType::Date | ColType::DateTime | ColType::Time | ColType::Json | ColType::Uuid => {
123 match row.try_get::<Option<String>, usize>(ordinal) {
124 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
125 Err(_) => Value::Null,
126 }
127 }
128 ColType::Unknown => {
129 if let Ok(v) = row.try_get::<Option<i64>, usize>(ordinal) {
136 return v.map(Value::I64).unwrap_or(Value::Null);
137 }
138 if let Ok(v) = row.try_get::<Option<f64>, usize>(ordinal) {
139 return v.map(Value::F64).unwrap_or(Value::Null);
140 }
141 if let Ok(v) = row.try_get::<Option<bool>, usize>(ordinal) {
142 return v.map(Value::Bool).unwrap_or(Value::Null);
143 }
144 if let Ok(v) = row.try_get::<Option<String>, usize>(ordinal) {
145 return v.map(Value::String).unwrap_or(Value::Null);
146 }
147 Value::Null
148 }
149 _ => Value::Null,
151 }
152}
153
154pub struct SqlitePoolHandle {
155 pool: sqlx::SqlitePool,
156}
157
158impl SqlitePoolHandle {
159 pub async fn connect(url: &str) -> Result<Self, DbError> {
160 Self::connect_with_max(url, 10).await
162 }
163
164 pub async fn connect_with_max(url: &str, max_connections: u32) -> Result<Self, DbError> {
171 let opts = sqlx::sqlite::SqliteConnectOptions::from_str(url)
177 .map_err(map_sqlx_error)?
178 .journal_mode(sqlx::sqlite::SqliteJournalMode::Wal)
179 .synchronous(sqlx::sqlite::SqliteSynchronous::Normal)
180 .busy_timeout(std::time::Duration::from_secs(5))
181 .pragma("mmap_size", "268435456");
182 let pool = sqlx::sqlite::SqlitePoolOptions::new()
183 .max_connections(max_connections)
184 .acquire_timeout(std::time::Duration::from_secs(30))
185 .idle_timeout(Some(std::time::Duration::from_secs(600)))
186 .max_lifetime(Some(std::time::Duration::from_secs(1800)))
187 .connect_with(opts)
188 .await
189 .map_err(map_sqlx_error)?;
190 Ok(Self { pool })
191 }
192
193 pub fn from_pool(pool: sqlx::SqlitePool) -> Self {
194 Self { pool }
195 }
196
197 pub fn pool(&self) -> &sqlx::SqlitePool {
198 &self.pool
199 }
200}
201
202pub struct SqlxSqliteConnectionFactory {
203 pool: Arc<SqlitePoolHandle>,
204}
205
206impl SqlxSqliteConnectionFactory {
207 pub fn new(pool: Arc<SqlitePoolHandle>) -> Self {
208 Self { pool }
209 }
210}
211
212#[async_trait]
213impl ConnectionFactory for SqlxSqliteConnectionFactory {
214 async fn create(&self) -> Result<Box<dyn Connection>, DbError> {
215 let conn = self.pool.pool.acquire().await.map_err(map_sqlx_error)?;
216 Ok(Box::new(crate::metrics_conn::MetricsConn::new(Box::new(
218 SqlxSqliteConnection {
219 conn: Some(conn),
220 connected: true,
221 in_transaction: false,
222 },
223 ))))
224 }
225}
226
227pub struct SqlxSqliteConnection {
228 conn: Option<sqlx::pool::PoolConnection<sqlx::Sqlite>>,
229 connected: bool,
230 in_transaction: bool,
231}
232
233impl Connection for SqlxSqliteConnection {
234 fn db_type(&self) -> Option<DbType> {
235 Some(DbType::Sqlite)
236 }
237
238 fn execute<'a>(
239 &'a mut self,
240 sql: &'a str,
241 ) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
242 Box::pin(async move {
243 let mut pool_conn = self
244 .conn
245 .take()
246 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
247 let result = if needs_raw_sql(sql) {
250 (&mut *pool_conn)
251 .execute(sqlx::raw_sql(sqlx::AssertSqlSafe(sql)))
252 .await
253 } else {
254 (&mut *pool_conn).execute(sqlx::AssertSqlSafe(sql)).await
255 };
256 self.conn = Some(pool_conn);
257
258 match result {
259 Ok(r) => Ok(r.rows_affected()),
260 Err(e) => {
261 let db_err = map_sqlx_error(e);
262 if matches!(db_err, DbError::ConnectionError(_) | DbError::IoError(_)) {
263 self.connected = false;
264 }
265 Err(db_err)
266 }
267 }
268 })
269 }
270
271 fn query<'a>(
272 &'a mut self,
273 sql: &'a str,
274 ) -> Pin<Box<dyn Future<Output = Result<Vec<HashMap<String, Value>>, DbError>> + Send + 'a>>
275 {
276 Box::pin(async move {
277 let mut pool_conn = self
278 .conn
279 .take()
280 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
281 let rows_result = (&mut *pool_conn).fetch_all(sqlx::AssertSqlSafe(sql)).await;
284 self.conn = Some(pool_conn);
285
286 let rows = rows_result.map_err(map_sqlx_error)?;
287 if rows.is_empty() {
288 return Ok(Vec::new());
289 }
290 let col_types: Vec<ColType> = rows[0]
293 .columns()
294 .iter()
295 .map(|col| ColType::parse_sqlite(col.type_info().name()))
296 .collect();
297 let col_names: Vec<String> = rows[0]
298 .columns()
299 .iter()
300 .map(|col| col.name().to_string())
301 .collect();
302 let mut result = Vec::with_capacity(rows.len());
303 for row in &rows {
304 let mut record = HashMap::with_capacity(col_types.len());
305 for (i, name) in col_names.iter().enumerate() {
306 let value = row_to_value_with_coltype_sqlite(row, i, col_types[i]);
307 record.insert(name.clone(), value);
308 }
309 result.push(record);
310 }
311 Ok(result)
312 })
313 }
314
315 fn begin_transaction<'a>(
316 &'a mut self,
317 ) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
318 Box::pin(async move {
319 if self.in_transaction {
320 return Err(DbError::Internal("transaction already started".to_string()));
321 }
322 self.execute("BEGIN").await?;
323 self.in_transaction = true;
324 Ok(())
325 })
326 }
327
328 fn commit<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
329 Box::pin(async move {
330 if self.in_transaction {
331 self.execute("COMMIT").await?;
332 self.in_transaction = false;
333 }
334 Ok(())
335 })
336 }
337
338 fn rollback<'a>(
339 &'a mut self,
340 ) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
341 Box::pin(async move {
342 if self.in_transaction {
343 let result = self.execute("ROLLBACK").await;
348 if result.is_ok() {
349 self.in_transaction = false;
350 }
351 result.map(|_| ())
352 } else {
353 Ok(())
354 }
355 })
356 }
357
358 fn is_connected(&self) -> bool {
359 self.connected && self.conn.is_some()
363 }
364
365 fn in_transaction(&self) -> bool {
366 self.in_transaction
367 }
368
369 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
370 Box::pin(async move {
371 match self.execute("SELECT 1").await {
372 Ok(_) => true,
373 Err(_) => {
374 self.connected = false;
375 false
376 }
377 }
378 })
379 }
380
381 fn close<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
382 Box::pin(async move {
383 if let Some(conn) = self.conn.take() {
384 drop(conn);
385 }
386 self.connected = false;
387 self.in_transaction = false;
388 Ok(())
389 })
390 }
391
392 fn execute_with_params<'a>(
398 &'a mut self,
399 sql: &'a str,
400 params: &'a [Value],
401 ) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
402 Box::pin(async move {
403 if needs_raw_sql(sql) || params.is_empty() {
404 return self.execute(sql).await;
405 }
406 let mut pool_conn = self
407 .conn
408 .take()
409 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
410 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
411 for v in params {
412 q = match v {
413 Value::Null => q.bind(None::<i64>),
414 Value::Bool(b) => q.bind(*b),
415 Value::I8(n) => q.bind(*n),
416 Value::I16(n) => q.bind(*n),
417 Value::I32(n) => q.bind(*n),
418 Value::I64(n) => q.bind(*n),
419 Value::U8(n) => q.bind(*n),
420 Value::U16(n) => q.bind(*n),
421 Value::U32(n) => q.bind(*n),
422 Value::U64(n) => q.bind(*n as i64),
423 Value::F32(f) => q.bind(*f),
424 Value::F64(f) => q.bind(*f),
425 Value::String(s) => q.bind(s.as_str()),
426 Value::Decimal(s) => q.bind(s.as_str()),
427 Value::Bytes(b) => q.bind(b.as_slice()),
428 Value::Uuid(s)
431 | Value::Date(s)
432 | Value::DateTime(s)
433 | Value::Time(s)
434 | Value::Json(s) => q.bind(s.as_str()),
435 other => q.bind(other.to_string()),
436 };
437 }
438 let result = q.execute(&mut *pool_conn).await;
439 self.conn = Some(pool_conn);
440 match result {
441 Ok(r) => Ok(r.rows_affected()),
442 Err(e) => {
443 let db_err = map_sqlx_error(e);
444 if matches!(db_err, DbError::ConnectionError(_) | DbError::IoError(_)) {
445 self.connected = false;
446 }
447 Err(db_err)
448 }
449 }
450 })
451 }
452
453 fn query_with_params<'a>(
458 &'a mut self,
459 sql: &'a str,
460 params: &'a [Value],
461 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, DbError>> + Send + 'a>> {
462 Box::pin(async move {
463 if params.is_empty() {
464 return self.query(sql).await;
465 }
466 let mut pool_conn = self
467 .conn
468 .take()
469 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
470 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
471 for v in params {
472 q = match v {
473 Value::Null => q.bind(None::<i64>),
474 Value::Bool(b) => q.bind(*b),
475 Value::I8(n) => q.bind(*n),
476 Value::I16(n) => q.bind(*n),
477 Value::I32(n) => q.bind(*n),
478 Value::I64(n) => q.bind(*n),
479 Value::U8(n) => q.bind(*n),
480 Value::U16(n) => q.bind(*n),
481 Value::U32(n) => q.bind(*n),
482 Value::U64(n) => q.bind(*n as i64),
483 Value::F32(f) => q.bind(*f),
484 Value::F64(f) => q.bind(*f),
485 Value::String(s) => q.bind(s.as_str()),
486 Value::Decimal(s) => q.bind(s.as_str()),
487 Value::Bytes(b) => q.bind(b.as_slice()),
488 Value::Uuid(s)
491 | Value::Date(s)
492 | Value::DateTime(s)
493 | Value::Time(s)
494 | Value::Json(s) => q.bind(s.as_str()),
495 other => q.bind(other.to_string()),
496 };
497 }
498 let rows_result = q.fetch_all(&mut *pool_conn).await;
499 self.conn = Some(pool_conn);
500 let rows = rows_result.map_err(map_sqlx_error)?;
501 if rows.is_empty() {
502 return Ok(Vec::new());
503 }
504 let col_types: Vec<ColType> = rows[0]
507 .columns()
508 .iter()
509 .map(|col| ColType::parse_sqlite(col.type_info().name()))
510 .collect();
511 let col_names: Vec<String> = rows[0]
512 .columns()
513 .iter()
514 .map(|col| col.name().to_string())
515 .collect();
516 let mut result = Vec::with_capacity(rows.len());
517 for row in &rows {
518 let mut record = HashMap::with_capacity(col_types.len());
519 for (i, name) in col_names.iter().enumerate() {
520 let value = row_to_value_with_coltype_sqlite(row, i, col_types[i]);
521 record.insert(name.clone(), value);
522 }
523 result.push(record);
524 }
525 Ok(result)
526 })
527 }
528
529 fn query_values<'a>(
535 &'a mut self,
536 sql: &'a str,
537 ) -> Pin<Box<dyn Future<Output = Result<QueryValues, DbError>> + Send + 'a>> {
538 Box::pin(async move {
539 let mut pool_conn = self
540 .conn
541 .take()
542 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
543 let rows_result = (&mut *pool_conn).fetch_all(sqlx::AssertSqlSafe(sql)).await;
544 self.conn = Some(pool_conn);
545 let rows = rows_result.map_err(map_sqlx_error)?;
546 if rows.is_empty() {
547 return Ok((Vec::new(), Vec::new()));
548 }
549 let cols = rows[0].columns();
550 let mut col_names: Vec<String> = Vec::with_capacity(cols.len());
551 let mut col_types: Vec<ColType> = Vec::with_capacity(cols.len());
552 for col in cols {
553 col_names.push(col.name().to_string());
554 col_types.push(ColType::parse_sqlite(col.type_info().name()));
555 }
556 let mut result_rows: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
557 for row in &rows {
558 let mut row_values: Vec<Value> = Vec::with_capacity(col_names.len());
559 for (idx, _) in col_names.iter().enumerate() {
560 row_values.push(row_to_value_with_coltype_sqlite(row, idx, col_types[idx]));
561 }
562 result_rows.push(row_values);
563 }
564 Ok((col_names, result_rows))
565 })
566 }
567
568 fn query_values_with_params<'a>(
572 &'a mut self,
573 sql: &'a str,
574 params: &'a [Value],
575 ) -> Pin<Box<dyn Future<Output = Result<QueryValues, DbError>> + Send + 'a>> {
576 Box::pin(async move {
577 if params.is_empty() {
578 return self.query_values(sql).await;
579 }
580 let mut pool_conn = self
581 .conn
582 .take()
583 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
584 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
585 for v in params {
586 q = match v {
587 Value::Null => q.bind(None::<i64>),
588 Value::Bool(b) => q.bind(*b),
589 Value::I8(n) => q.bind(*n),
590 Value::I16(n) => q.bind(*n),
591 Value::I32(n) => q.bind(*n),
592 Value::I64(n) => q.bind(*n),
593 Value::U8(n) => q.bind(*n),
594 Value::U16(n) => q.bind(*n),
595 Value::U32(n) => q.bind(*n),
596 Value::U64(n) => q.bind(*n as i64),
597 Value::F32(f) => q.bind(*f),
598 Value::F64(f) => q.bind(*f),
599 Value::String(s) => q.bind(s.as_str()),
600 Value::Decimal(s) => q.bind(s.as_str()),
601 Value::Bytes(b) => q.bind(b.as_slice()),
602 Value::Uuid(s)
605 | Value::Date(s)
606 | Value::DateTime(s)
607 | Value::Time(s)
608 | Value::Json(s) => q.bind(s.as_str()),
609 other => q.bind(other.to_string()),
610 };
611 }
612 let rows_result = q.fetch_all(&mut *pool_conn).await;
613 self.conn = Some(pool_conn);
614 let rows = rows_result.map_err(map_sqlx_error)?;
615 if rows.is_empty() {
616 return Ok((Vec::new(), Vec::new()));
617 }
618 let cols = rows[0].columns();
619 let mut col_names: Vec<String> = Vec::with_capacity(cols.len());
620 let mut col_types: Vec<ColType> = Vec::with_capacity(cols.len());
621 for col in cols {
622 col_names.push(col.name().to_string());
623 col_types.push(ColType::parse_sqlite(col.type_info().name()));
624 }
625 let mut result_rows: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
626 for row in &rows {
627 let mut row_values: Vec<Value> = Vec::with_capacity(col_names.len());
628 for (idx, _) in col_names.iter().enumerate() {
629 row_values.push(row_to_value_with_coltype_sqlite(row, idx, col_types[idx]));
630 }
631 result_rows.push(row_values);
632 }
633 Ok((col_names, result_rows))
634 })
635 }
636
637 fn query_stream<'a>(
643 &'a mut self,
644 sql: &'a str,
645 ) -> Pin<Box<dyn futures::Stream<Item = Result<HashMap<String, Value>, DbError>> + Send + 'a>>
646 {
647 Box::pin(async_stream::try_stream! {
648 let mut pool_conn = self
649 .conn
650 .take()
651 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
652 let mut row_stream = sqlx::query(sqlx::AssertSqlSafe(sql)).fetch(&mut *pool_conn);
653 let mut col_types: Vec<ColType> = Vec::new();
654 let mut col_names: Vec<String> = Vec::new();
655 let mut first_row = true;
656 while let Some(row_result) = row_stream.next().await {
657 let row = row_result.map_err(map_sqlx_error)?;
658 if first_row {
659 for col in row.columns() {
660 col_names.push(col.name().to_string());
661 col_types.push(ColType::parse_sqlite(col.type_info().name()));
662 }
663 first_row = false;
664 }
665 let mut record = HashMap::with_capacity(col_names.len());
666 for (i, name) in col_names.iter().enumerate() {
667 let value = row_to_value_with_coltype_sqlite(&row, i, col_types[i]);
668 record.insert(name.clone(), value);
669 }
670 yield record;
671 }
672 drop(row_stream);
675 self.conn = Some(pool_conn);
676 })
677 }
678}
679
680impl Drop for SqlxSqliteConnection {
681 fn drop(&mut self) {
682 if let Some(conn) = self.conn.take() {
683 drop(conn);
684 }
685 }
686}
687
688pub async fn sqlite_backup(
693 conn: &mut SqlxSqliteConnection,
694 dest_path: &str,
695) -> Result<(), DbError> {
696 let escaped_path = dest_path.replace('\'', "''");
697 let sql = format!("VACUUM INTO '{}'", escaped_path);
698 conn.execute(&sql).await?;
699 Ok(())
700}
701
702fn row_to_value_mysql(row: &sqlx::mysql::MySqlRow, ordinal: usize) -> Value {
708 use sqlx::TypeInfo;
709 let type_name = row.columns()[ordinal].type_info().name();
710 match type_name {
711 "BOOLEAN" => match row.try_get::<Option<bool>, usize>(ordinal) {
712 Ok(v) => v.map(Value::Bool).unwrap_or(Value::Null),
713 Err(_) => Value::Null,
714 },
715 "TINYINT" | "TINYINT UNSIGNED" => match row.try_get::<Option<i8>, usize>(ordinal) {
716 Ok(v) => v.map(Value::I8).unwrap_or(Value::Null),
717 Err(_) => match row.try_get::<Option<u8>, usize>(ordinal) {
718 Ok(v) => v.map(Value::U8).unwrap_or(Value::Null),
719 Err(_) => Value::Null,
720 },
721 },
722 "SMALLINT" | "SMALLINT UNSIGNED" => match row.try_get::<Option<i16>, usize>(ordinal) {
723 Ok(v) => v.map(Value::I16).unwrap_or(Value::Null),
724 Err(_) => match row.try_get::<Option<u16>, usize>(ordinal) {
725 Ok(v) => v.map(Value::U16).unwrap_or(Value::Null),
726 Err(_) => Value::Null,
727 },
728 },
729 "INT" | "INT UNSIGNED" | "MEDIUMINT" | "MEDIUMINT UNSIGNED" => {
730 match row.try_get::<Option<i32>, usize>(ordinal) {
731 Ok(v) => v.map(Value::I32).unwrap_or(Value::Null),
732 Err(_) => match row.try_get::<Option<u32>, usize>(ordinal) {
733 Ok(v) => v.map(Value::U32).unwrap_or(Value::Null),
734 Err(_) => Value::Null,
735 },
736 }
737 }
738 "BIGINT" | "BIGINT UNSIGNED" => match row.try_get::<Option<i64>, usize>(ordinal) {
739 Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
740 Err(_) => match row.try_get::<Option<u64>, usize>(ordinal) {
741 Ok(v) => v.map(Value::U64).unwrap_or(Value::Null),
742 Err(_) => Value::Null,
743 },
744 },
745 "FLOAT" => match row.try_get::<Option<f32>, usize>(ordinal) {
746 Ok(v) => v.map(Value::F32).unwrap_or(Value::Null),
747 Err(_) => Value::Null,
748 },
749 "DOUBLE" => match row.try_get::<Option<f64>, usize>(ordinal) {
750 Ok(v) => v.map(Value::F64).unwrap_or(Value::Null),
751 Err(_) => Value::Null,
752 },
753 "VARCHAR" | "TEXT" | "CHAR" | "TINYTEXT" | "MEDIUMTEXT" | "LONGTEXT" | "ENUM" => {
754 match row.try_get::<Option<String>, usize>(ordinal) {
755 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
756 Err(_) => Value::Null,
757 }
758 }
759 "BLOB" | "TINYBLOB" | "MEDIUMBLOB" | "LONGBLOB" | "BINARY" | "VARBINARY" => {
760 match row.try_get::<Option<Vec<u8>>, usize>(ordinal) {
761 Ok(v) => v.map(Value::Bytes).unwrap_or(Value::Null),
762 Err(_) => Value::Null,
763 }
764 }
765 "DECIMAL" | "NUMERIC" | "NEWDECIMAL" => {
767 match row.try_get::<Option<bigdecimal::BigDecimal>, usize>(ordinal) {
768 Ok(Some(v)) => Value::Decimal(v.to_string()),
769 Ok(None) => Value::Null,
770 Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
771 Ok(v) => v.map(Value::Decimal).unwrap_or(Value::Null),
772 Err(_) => Value::Null,
773 },
774 }
775 }
776 "DATE" => match row.try_get::<Option<sqlx::types::chrono::NaiveDate>, usize>(ordinal) {
778 Ok(v) => v.map(|d| Value::Date(d.to_string())).unwrap_or(Value::Null),
779 Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
780 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
781 Err(_) => Value::Null,
782 },
783 },
784 "DATETIME" | "TIMESTAMP" => {
785 match row.try_get::<Option<sqlx::types::chrono::NaiveDateTime>, usize>(ordinal) {
786 Ok(v) => v
787 .map(|d| Value::DateTime(d.to_string()))
788 .unwrap_or(Value::Null),
789 Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
790 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
791 Err(_) => Value::Null,
792 },
793 }
794 }
795 "TIME" => match row.try_get::<Option<sqlx::types::chrono::NaiveTime>, usize>(ordinal) {
796 Ok(v) => v.map(|t| Value::Time(t.to_string())).unwrap_or(Value::Null),
797 Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
798 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
799 Err(_) => Value::Null,
800 },
801 },
802 _ => {
803 if let Ok(v) = row.try_get::<Option<i64>, usize>(ordinal) {
805 return v.map(Value::I64).unwrap_or(Value::Null);
806 }
807 if let Ok(v) = row.try_get::<Option<f64>, usize>(ordinal) {
808 return v.map(Value::F64).unwrap_or(Value::Null);
809 }
810 if let Ok(v) = row.try_get::<Option<bool>, usize>(ordinal) {
811 return v.map(Value::Bool).unwrap_or(Value::Null);
812 }
813 if let Ok(v) = row.try_get::<Option<String>, usize>(ordinal) {
814 return v.map(Value::String).unwrap_or(Value::Null);
815 }
816 Value::Null
817 }
818 }
819}
820
821fn row_to_value_with_coltype_mysql(
829 row: &sqlx::mysql::MySqlRow,
830 ordinal: usize,
831 col_type: ColType,
832) -> Value {
833 match col_type {
834 ColType::Bool => match row.try_get::<Option<bool>, usize>(ordinal) {
835 Ok(v) => v.map(Value::Bool).unwrap_or(Value::Null),
836 Err(_) => Value::Null,
837 },
838 ColType::I8 => match row.try_get::<Option<i8>, usize>(ordinal) {
839 Ok(v) => v.map(Value::I8).unwrap_or(Value::Null),
840 Err(_) => match row.try_get::<Option<u8>, usize>(ordinal) {
841 Ok(v) => v.map(Value::U8).unwrap_or(Value::Null),
842 Err(_) => Value::Null,
843 },
844 },
845 ColType::I16 => match row.try_get::<Option<i16>, usize>(ordinal) {
846 Ok(v) => v.map(Value::I16).unwrap_or(Value::Null),
847 Err(_) => match row.try_get::<Option<u16>, usize>(ordinal) {
848 Ok(v) => v.map(Value::U16).unwrap_or(Value::Null),
849 Err(_) => Value::Null,
850 },
851 },
852 ColType::I32 => match row.try_get::<Option<i32>, usize>(ordinal) {
853 Ok(v) => v.map(Value::I32).unwrap_or(Value::Null),
854 Err(_) => match row.try_get::<Option<u32>, usize>(ordinal) {
855 Ok(v) => v.map(Value::U32).unwrap_or(Value::Null),
856 Err(_) => Value::Null,
857 },
858 },
859 ColType::I64 => match row.try_get::<Option<i64>, usize>(ordinal) {
860 Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
861 Err(_) => match row.try_get::<Option<u64>, usize>(ordinal) {
862 Ok(v) => v.map(Value::U64).unwrap_or(Value::Null),
863 Err(_) => Value::Null,
864 },
865 },
866 ColType::U8 => match row.try_get::<Option<u8>, usize>(ordinal) {
867 Ok(v) => v.map(Value::U8).unwrap_or(Value::Null),
868 Err(_) => Value::Null,
869 },
870 ColType::U16 => match row.try_get::<Option<u16>, usize>(ordinal) {
871 Ok(v) => v.map(Value::U16).unwrap_or(Value::Null),
872 Err(_) => Value::Null,
873 },
874 ColType::U32 => match row.try_get::<Option<u32>, usize>(ordinal) {
875 Ok(v) => v.map(Value::U32).unwrap_or(Value::Null),
876 Err(_) => Value::Null,
877 },
878 ColType::U64 => match row.try_get::<Option<u64>, usize>(ordinal) {
879 Ok(v) => v.map(Value::U64).unwrap_or(Value::Null),
880 Err(_) => Value::Null,
881 },
882 ColType::F32 => match row.try_get::<Option<f32>, usize>(ordinal) {
883 Ok(v) => v.map(Value::F32).unwrap_or(Value::Null),
884 Err(_) => Value::Null,
885 },
886 ColType::F64 => match row.try_get::<Option<f64>, usize>(ordinal) {
887 Ok(v) => v.map(Value::F64).unwrap_or(Value::Null),
888 Err(_) => Value::Null,
889 },
890 ColType::Decimal => match row.try_get::<Option<bigdecimal::BigDecimal>, usize>(ordinal) {
892 Ok(Some(v)) => Value::Decimal(v.to_string()),
893 Ok(None) => Value::Null,
894 Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
895 Ok(v) => v.map(Value::Decimal).unwrap_or(Value::Null),
896 Err(_) => Value::Null,
897 },
898 },
899 ColType::String => match row.try_get::<Option<String>, usize>(ordinal) {
900 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
901 Err(_) => Value::Null,
902 },
903 ColType::Bytes => match row.try_get::<Option<Vec<u8>>, usize>(ordinal) {
904 Ok(v) => v.map(Value::Bytes).unwrap_or(Value::Null),
905 Err(_) => Value::Null,
906 },
907 ColType::Date => {
909 match row.try_get::<Option<sqlx::types::chrono::NaiveDate>, usize>(ordinal) {
910 Ok(v) => v.map(|d| Value::Date(d.to_string())).unwrap_or(Value::Null),
911 Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
912 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
913 Err(_) => Value::Null,
914 },
915 }
916 }
917 ColType::DateTime => {
918 match row.try_get::<Option<sqlx::types::chrono::NaiveDateTime>, usize>(ordinal) {
919 Ok(v) => v
920 .map(|d| Value::DateTime(d.to_string()))
921 .unwrap_or(Value::Null),
922 Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
923 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
924 Err(_) => Value::Null,
925 },
926 }
927 }
928 ColType::Time => {
929 match row.try_get::<Option<sqlx::types::chrono::NaiveTime>, usize>(ordinal) {
930 Ok(v) => v.map(|t| Value::Time(t.to_string())).unwrap_or(Value::Null),
931 Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
932 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
933 Err(_) => Value::Null,
934 },
935 }
936 }
937 ColType::Json | ColType::Uuid => match row.try_get::<Option<String>, usize>(ordinal) {
939 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
940 Err(_) => Value::Null,
941 },
942 ColType::Unknown => {
943 if let Ok(v) = row.try_get::<Option<i64>, usize>(ordinal) {
945 return v.map(Value::I64).unwrap_or(Value::Null);
946 }
947 if let Ok(v) = row.try_get::<Option<f64>, usize>(ordinal) {
948 return v.map(Value::F64).unwrap_or(Value::Null);
949 }
950 if let Ok(v) = row.try_get::<Option<bool>, usize>(ordinal) {
951 return v.map(Value::Bool).unwrap_or(Value::Null);
952 }
953 if let Ok(v) = row.try_get::<Option<String>, usize>(ordinal) {
954 return v.map(Value::String).unwrap_or(Value::Null);
955 }
956 Value::Null
957 }
958 _ => Value::Null,
960 }
961}
962
963pub struct MySqlPoolHandle {
964 pool: sqlx::MySqlPool,
965}
966
967impl MySqlPoolHandle {
968 pub async fn connect(url: &str) -> Result<Self, DbError> {
969 Self::connect_with_max(url, 10).await
971 }
972
973 pub async fn connect_with_max(url: &str, max_connections: u32) -> Result<Self, DbError> {
976 let pool = sqlx::pool::PoolOptions::<sqlx::MySql>::new()
977 .max_connections(max_connections)
978 .acquire_timeout(std::time::Duration::from_secs(30))
979 .idle_timeout(Some(std::time::Duration::from_secs(600)))
980 .max_lifetime(Some(std::time::Duration::from_secs(1800)))
981 .connect(url)
982 .await
983 .map_err(map_sqlx_error)?;
984 Ok(Self { pool })
985 }
986
987 pub fn from_pool(pool: sqlx::MySqlPool) -> Self {
988 Self { pool }
989 }
990
991 pub fn pool(&self) -> &sqlx::MySqlPool {
992 &self.pool
993 }
994}
995
996pub struct SqlxMySqlConnectionFactory {
997 pool: Arc<MySqlPoolHandle>,
998}
999
1000impl SqlxMySqlConnectionFactory {
1001 pub fn new(pool: Arc<MySqlPoolHandle>) -> Self {
1002 Self { pool }
1003 }
1004}
1005
1006#[async_trait]
1007impl ConnectionFactory for SqlxMySqlConnectionFactory {
1008 async fn create(&self) -> Result<Box<dyn Connection>, DbError> {
1009 let conn = self.pool.pool.acquire().await.map_err(map_sqlx_error)?;
1010 Ok(Box::new(crate::metrics_conn::MetricsConn::new(Box::new(
1012 SqlxMySqlConnection {
1013 conn: Some(conn),
1014 connected: true,
1015 in_transaction: false,
1016 },
1017 ))))
1018 }
1019}
1020
1021pub struct SqlxMySqlConnection {
1022 conn: Option<sqlx::pool::PoolConnection<sqlx::MySql>>,
1023 connected: bool,
1024 in_transaction: bool,
1025}
1026
1027impl Connection for SqlxMySqlConnection {
1028 fn db_type(&self) -> Option<DbType> {
1029 Some(DbType::MySQL)
1030 }
1031
1032 fn session_reset_sql(&self) -> Option<&'static str> {
1038 Some(
1039 "SET SESSION sql_mode = DEFAULT, \
1040 SESSION character_set_client = @@GLOBAL.character_set_client, \
1041 SESSION character_set_results = @@GLOBAL.character_set_results, \
1042 SESSION character_set_connection = @@GLOBAL.character_set_connection, \
1043 SESSION time_zone = @@GLOBAL.time_zone, \
1044 SESSION autocommit = 1",
1045 )
1046 }
1047
1048 fn execute<'a>(
1049 &'a mut self,
1050 sql: &'a str,
1051 ) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
1052 Box::pin(async move {
1053 let mut pool_conn = self
1054 .conn
1055 .take()
1056 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1057 let result = if needs_raw_sql(sql) {
1060 (&mut *pool_conn)
1061 .execute(sqlx::raw_sql(sqlx::AssertSqlSafe(sql)))
1062 .await
1063 } else {
1064 (&mut *pool_conn).execute(sqlx::AssertSqlSafe(sql)).await
1065 };
1066 self.conn = Some(pool_conn);
1067
1068 match result {
1069 Ok(r) => Ok(r.rows_affected()),
1070 Err(e) => {
1071 let db_err = map_sqlx_error(e);
1072 if matches!(db_err, DbError::ConnectionError(_) | DbError::IoError(_)) {
1073 self.connected = false;
1074 }
1075 Err(db_err)
1076 }
1077 }
1078 })
1079 }
1080
1081 fn query<'a>(
1082 &'a mut self,
1083 sql: &'a str,
1084 ) -> Pin<Box<dyn Future<Output = Result<Vec<HashMap<String, Value>>, DbError>> + Send + 'a>>
1085 {
1086 Box::pin(async move {
1087 let mut pool_conn = self
1088 .conn
1089 .take()
1090 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1091 let rows_result = (&mut *pool_conn).fetch_all(sqlx::AssertSqlSafe(sql)).await;
1094 self.conn = Some(pool_conn);
1095
1096 let rows = rows_result.map_err(map_sqlx_error)?;
1097 if rows.is_empty() {
1098 return Ok(Vec::new());
1099 }
1100 let col_types: Vec<ColType> = rows[0]
1103 .columns()
1104 .iter()
1105 .map(|col| ColType::parse_mysql(col.type_info().name()))
1106 .collect();
1107 let col_names: Vec<String> = rows[0]
1108 .columns()
1109 .iter()
1110 .map(|col| col.name().to_string())
1111 .collect();
1112 let mut result = Vec::with_capacity(rows.len());
1113 for row in &rows {
1114 let mut record = HashMap::with_capacity(col_types.len());
1115 for (i, name) in col_names.iter().enumerate() {
1116 let value = row_to_value_with_coltype_mysql(row, i, col_types[i]);
1117 record.insert(name.clone(), value);
1118 }
1119 result.push(record);
1120 }
1121 Ok(result)
1122 })
1123 }
1124
1125 fn begin_transaction<'a>(
1126 &'a mut self,
1127 ) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
1128 Box::pin(async move {
1129 if self.in_transaction {
1130 return Err(DbError::Internal("transaction already started".to_string()));
1131 }
1132 self.execute("BEGIN").await?;
1133 self.in_transaction = true;
1134 Ok(())
1135 })
1136 }
1137
1138 fn commit<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
1139 Box::pin(async move {
1140 if self.in_transaction {
1141 self.execute("COMMIT").await?;
1142 self.in_transaction = false;
1143 }
1144 Ok(())
1145 })
1146 }
1147
1148 fn rollback<'a>(
1149 &'a mut self,
1150 ) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
1151 Box::pin(async move {
1152 if self.in_transaction {
1153 let result = self.execute("ROLLBACK").await;
1158 if result.is_ok() {
1159 self.in_transaction = false;
1160 }
1161 result.map(|_| ())
1162 } else {
1163 Ok(())
1164 }
1165 })
1166 }
1167
1168 fn is_connected(&self) -> bool {
1169 self.connected && self.conn.is_some()
1173 }
1174
1175 fn in_transaction(&self) -> bool {
1176 self.in_transaction
1177 }
1178
1179 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
1180 Box::pin(async move {
1181 match self.execute("SELECT 1").await {
1182 Ok(_) => true,
1183 Err(_) => {
1184 self.connected = false;
1185 false
1186 }
1187 }
1188 })
1189 }
1190
1191 fn close<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
1192 Box::pin(async move {
1193 if let Some(conn) = self.conn.take() {
1194 drop(conn);
1195 }
1196 self.connected = false;
1197 self.in_transaction = false;
1198 Ok(())
1199 })
1200 }
1201
1202 fn execute_with_params<'a>(
1208 &'a mut self,
1209 sql: &'a str,
1210 params: &'a [Value],
1211 ) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
1212 Box::pin(async move {
1213 if needs_raw_sql(sql) || params.is_empty() {
1214 return self.execute(sql).await;
1215 }
1216 let mut pool_conn = self
1217 .conn
1218 .take()
1219 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1220 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
1221 for v in params {
1222 q = match v {
1223 Value::Null => q.bind(None::<i64>),
1224 Value::Bool(b) => q.bind(*b),
1225 Value::I8(n) => q.bind(*n),
1226 Value::I16(n) => q.bind(*n),
1227 Value::I32(n) => q.bind(*n),
1228 Value::I64(n) => q.bind(*n),
1229 Value::U8(n) => q.bind(*n),
1230 Value::U16(n) => q.bind(*n),
1231 Value::U32(n) => q.bind(*n),
1232 Value::U64(n) => q.bind(*n as i64),
1233 Value::F32(f) => q.bind(*f),
1234 Value::F64(f) => q.bind(*f),
1235 Value::String(s) => q.bind(s.as_str()),
1236 Value::Decimal(s) => q.bind(s.as_str()),
1237 Value::Bytes(b) => q.bind(b.as_slice()),
1238 Value::Uuid(s)
1241 | Value::Date(s)
1242 | Value::DateTime(s)
1243 | Value::Time(s)
1244 | Value::Json(s) => q.bind(s.as_str()),
1245 other => q.bind(other.to_string()),
1246 };
1247 }
1248 let result = q.execute(&mut *pool_conn).await;
1249 self.conn = Some(pool_conn);
1250 match result {
1251 Ok(r) => Ok(r.rows_affected()),
1252 Err(e) => {
1253 let db_err = map_sqlx_error(e);
1254 if matches!(db_err, DbError::ConnectionError(_) | DbError::IoError(_)) {
1255 self.connected = false;
1256 }
1257 Err(db_err)
1258 }
1259 }
1260 })
1261 }
1262
1263 fn query_with_params<'a>(
1265 &'a mut self,
1266 sql: &'a str,
1267 params: &'a [Value],
1268 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, DbError>> + Send + 'a>> {
1269 Box::pin(async move {
1270 if params.is_empty() {
1271 return self.query(sql).await;
1272 }
1273 let mut pool_conn = self
1274 .conn
1275 .take()
1276 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1277 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
1278 for v in params {
1279 q = match v {
1280 Value::Null => q.bind(None::<i64>),
1281 Value::Bool(b) => q.bind(*b),
1282 Value::I8(n) => q.bind(*n),
1283 Value::I16(n) => q.bind(*n),
1284 Value::I32(n) => q.bind(*n),
1285 Value::I64(n) => q.bind(*n),
1286 Value::U8(n) => q.bind(*n),
1287 Value::U16(n) => q.bind(*n),
1288 Value::U32(n) => q.bind(*n),
1289 Value::U64(n) => q.bind(*n as i64),
1290 Value::F32(f) => q.bind(*f),
1291 Value::F64(f) => q.bind(*f),
1292 Value::String(s) => q.bind(s.as_str()),
1293 Value::Decimal(s) => q.bind(s.as_str()),
1294 Value::Bytes(b) => q.bind(b.as_slice()),
1295 Value::Uuid(s)
1298 | Value::Date(s)
1299 | Value::DateTime(s)
1300 | Value::Time(s)
1301 | Value::Json(s) => q.bind(s.as_str()),
1302 other => q.bind(other.to_string()),
1303 };
1304 }
1305 let rows_result = q.fetch_all(&mut *pool_conn).await;
1306 self.conn = Some(pool_conn);
1307 let rows = rows_result.map_err(map_sqlx_error)?;
1308 if rows.is_empty() {
1309 return Ok(Vec::new());
1310 }
1311 let col_names: Vec<String> = rows[0]
1312 .columns()
1313 .iter()
1314 .map(|col| col.name().to_string())
1315 .collect();
1316 let mut result = Vec::with_capacity(rows.len());
1317 for row in &rows {
1318 let mut record = HashMap::with_capacity(col_names.len());
1319 for (i, name) in col_names.iter().enumerate() {
1320 let value = row_to_value_mysql(row, i);
1321 record.insert(name.clone(), value);
1322 }
1323 result.push(record);
1324 }
1325 Ok(result)
1326 })
1327 }
1328
1329 fn query_values<'a>(
1333 &'a mut self,
1334 sql: &'a str,
1335 ) -> Pin<Box<dyn Future<Output = Result<QueryValues, DbError>> + Send + 'a>> {
1336 Box::pin(async move {
1337 let mut pool_conn = self
1338 .conn
1339 .take()
1340 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1341 let rows_result = (&mut *pool_conn).fetch_all(sqlx::AssertSqlSafe(sql)).await;
1342 self.conn = Some(pool_conn);
1343 let rows = rows_result.map_err(map_sqlx_error)?;
1344 if rows.is_empty() {
1345 return Ok((Vec::new(), Vec::new()));
1346 }
1347 let cols = rows[0].columns();
1348 let mut col_names: Vec<String> = Vec::with_capacity(cols.len());
1349 let mut col_types: Vec<ColType> = Vec::with_capacity(cols.len());
1350 for col in cols {
1351 col_names.push(col.name().to_string());
1352 col_types.push(ColType::parse_mysql(col.type_info().name()));
1353 }
1354 let mut result_rows: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
1355 for row in &rows {
1356 let mut row_values: Vec<Value> = Vec::with_capacity(col_names.len());
1357 for (idx, _) in col_names.iter().enumerate() {
1358 row_values.push(row_to_value_with_coltype_mysql(row, idx, col_types[idx]));
1359 }
1360 result_rows.push(row_values);
1361 }
1362 Ok((col_names, result_rows))
1363 })
1364 }
1365
1366 fn query_values_with_params<'a>(
1368 &'a mut self,
1369 sql: &'a str,
1370 params: &'a [Value],
1371 ) -> Pin<Box<dyn Future<Output = Result<QueryValues, DbError>> + Send + 'a>> {
1372 Box::pin(async move {
1373 if params.is_empty() {
1374 return self.query_values(sql).await;
1375 }
1376 let mut pool_conn = self
1377 .conn
1378 .take()
1379 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1380 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
1381 for v in params {
1382 q = match v {
1383 Value::Null => q.bind(None::<i64>),
1384 Value::Bool(b) => q.bind(*b),
1385 Value::I8(n) => q.bind(*n),
1386 Value::I16(n) => q.bind(*n),
1387 Value::I32(n) => q.bind(*n),
1388 Value::I64(n) => q.bind(*n),
1389 Value::U8(n) => q.bind(*n),
1390 Value::U16(n) => q.bind(*n),
1391 Value::U32(n) => q.bind(*n),
1392 Value::U64(n) => q.bind(*n as i64),
1393 Value::F32(f) => q.bind(*f),
1394 Value::F64(f) => q.bind(*f),
1395 Value::String(s) => q.bind(s.as_str()),
1396 Value::Decimal(s) => q.bind(s.as_str()),
1397 Value::Bytes(b) => q.bind(b.as_slice()),
1398 Value::Uuid(s)
1401 | Value::Date(s)
1402 | Value::DateTime(s)
1403 | Value::Time(s)
1404 | Value::Json(s) => q.bind(s.as_str()),
1405 other => q.bind(other.to_string()),
1406 };
1407 }
1408 let rows_result = q.fetch_all(&mut *pool_conn).await;
1409 self.conn = Some(pool_conn);
1410 let rows = rows_result.map_err(map_sqlx_error)?;
1411 if rows.is_empty() {
1412 return Ok((Vec::new(), Vec::new()));
1413 }
1414 let cols = rows[0].columns();
1415 let mut col_names: Vec<String> = Vec::with_capacity(cols.len());
1416 let mut col_types: Vec<ColType> = Vec::with_capacity(cols.len());
1417 for col in cols {
1418 col_names.push(col.name().to_string());
1419 col_types.push(ColType::parse_mysql(col.type_info().name()));
1420 }
1421 let mut result_rows: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
1422 for row in &rows {
1423 let mut row_values: Vec<Value> = Vec::with_capacity(col_names.len());
1424 for (idx, _) in col_names.iter().enumerate() {
1425 row_values.push(row_to_value_with_coltype_mysql(row, idx, col_types[idx]));
1426 }
1427 result_rows.push(row_values);
1428 }
1429 Ok((col_names, result_rows))
1430 })
1431 }
1432
1433 fn query_stream<'a>(
1439 &'a mut self,
1440 sql: &'a str,
1441 ) -> Pin<Box<dyn futures::Stream<Item = Result<HashMap<String, Value>, DbError>> + Send + 'a>>
1442 {
1443 Box::pin(async_stream::try_stream! {
1444 let mut pool_conn = self
1445 .conn
1446 .take()
1447 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1448 let mut row_stream = sqlx::query(sqlx::AssertSqlSafe(sql)).fetch(&mut *pool_conn);
1449 let mut col_types: Vec<ColType> = Vec::new();
1450 let mut col_names: Vec<String> = Vec::new();
1451 let mut first_row = true;
1452 while let Some(row_result) = row_stream.next().await {
1453 let row = row_result.map_err(map_sqlx_error)?;
1454 if first_row {
1455 for col in row.columns() {
1456 col_names.push(col.name().to_string());
1457 col_types.push(ColType::parse_mysql(col.type_info().name()));
1458 }
1459 first_row = false;
1460 }
1461 let mut record = HashMap::with_capacity(col_names.len());
1462 for (i, name) in col_names.iter().enumerate() {
1463 let value = row_to_value_with_coltype_mysql(&row, i, col_types[i]);
1464 record.insert(name.clone(), value);
1465 }
1466 yield record;
1467 }
1468 drop(row_stream);
1471 self.conn = Some(pool_conn);
1472 })
1473 }
1474}
1475
1476impl Drop for SqlxMySqlConnection {
1477 fn drop(&mut self) {
1478 if let Some(conn) = self.conn.take() {
1479 drop(conn);
1480 }
1481 }
1482}
1483
1484pub async fn mysql_bulk_insert(
1491 conn: &mut SqlxMySqlConnection,
1492 table: &str,
1493 columns: &[&str],
1494 rows: &[Vec<Value>],
1495) -> Result<u64, DbError> {
1496 if rows.is_empty() {
1497 return Ok(0);
1498 }
1499 let col_list = columns.join(", ");
1500 let cols_per_row = columns.len();
1501 let row_placeholder = format!("({})", vec!["?"; cols_per_row].join(", "));
1503 let placeholders = vec![row_placeholder; rows.len()].join(", ");
1504 let sql = format!(
1505 "INSERT INTO {} ({}) VALUES {}",
1506 table, col_list, placeholders
1507 );
1508 let mut params: Vec<Value> = Vec::with_capacity(rows.len() * cols_per_row);
1510 for row in rows {
1511 for v in row {
1512 params.push(v.clone());
1513 }
1514 }
1515 conn.execute_with_params(&sql, ¶ms).await
1516}
1517
1518fn pg_decode_err(
1525 row: &sqlx::postgres::PgRow,
1526 ordinal: usize,
1527 type_name: &str,
1528 err: sqlx::Error,
1529) -> DbError {
1530 let column = row
1531 .columns()
1532 .get(ordinal)
1533 .map(|c| c.name().to_string())
1534 .unwrap_or_else(|| format!("#{ordinal}"));
1535 DbError::QueryError(format!(
1536 "PostgreSQL 列 `{column}`(类型 {type_name},序号 {ordinal})解码失败:{err};\
1537 已拒绝以 NULL 代替真实值"
1538 ))
1539}
1540
1541fn row_to_value_pg(row: &sqlx::postgres::PgRow, ordinal: usize) -> Result<Value, DbError> {
1542 use sqlx::TypeInfo;
1543 let type_name = row.columns()[ordinal].type_info().name();
1544 macro_rules! dec {
1545 ($ty:ty) => {
1546 row.try_get::<Option<$ty>, usize>(ordinal)
1547 .map_err(|e| pg_decode_err(row, ordinal, type_name, e))
1548 };
1549 }
1550 match type_name {
1551 "BOOL" => Ok(dec!(bool)?.map(Value::Bool).unwrap_or(Value::Null)),
1552 "INT2" => Ok(dec!(i16)?.map(Value::I16).unwrap_or(Value::Null)),
1553 "INT4" | "OID" => Ok(dec!(i32)?.map(Value::I32).unwrap_or(Value::Null)),
1554 "INT8" => Ok(dec!(i64)?.map(Value::I64).unwrap_or(Value::Null)),
1555 "FLOAT4" => Ok(dec!(f32)?.map(Value::F32).unwrap_or(Value::Null)),
1556 "FLOAT8" => Ok(dec!(f64)?.map(Value::F64).unwrap_or(Value::Null)),
1557 "TEXT" | "VARCHAR" | "CHAR" | "NAME" | "BPCHAR" => {
1558 Ok(dec!(String)?.map(Value::String).unwrap_or(Value::Null))
1559 }
1560 "BYTEA" => Ok(dec!(Vec<u8>)?.map(Value::Bytes).unwrap_or(Value::Null)),
1561 "NUMERIC" => Ok(dec!(bigdecimal::BigDecimal)?
1563 .map(|v| Value::Decimal(v.to_string()))
1564 .unwrap_or(Value::Null)),
1565 "UUID" => Ok(dec!(sqlx::types::Uuid)?
1566 .map(|uuid| Value::String(uuid.to_string()))
1567 .unwrap_or(Value::Null)),
1568 "TIMESTAMP" => Ok(dec!(sqlx::types::chrono::NaiveDateTime)?
1570 .map(|v| Value::DateTime(v.to_string()))
1571 .unwrap_or(Value::Null)),
1572 "TIMESTAMPTZ" => Ok(
1573 dec!(sqlx::types::chrono::DateTime<sqlx::types::chrono::Utc>)?
1574 .map(|v| Value::DateTime(v.format("%Y-%m-%dT%H:%M:%S%.6fZ").to_string()))
1575 .unwrap_or(Value::Null),
1576 ),
1577 "DATE" => Ok(dec!(sqlx::types::chrono::NaiveDate)?
1578 .map(|v| Value::Date(v.to_string()))
1579 .unwrap_or(Value::Null)),
1580 "TIME" => Ok(dec!(sqlx::types::chrono::NaiveTime)?
1581 .map(|v| Value::Time(v.to_string()))
1582 .unwrap_or(Value::Null)),
1583 "JSON" | "JSONB" => Ok(dec!(serde_json::Value)?
1584 .map(|v| Value::Json(v.to_string()))
1585 .unwrap_or(Value::Null)),
1586 _ => {
1587 if let Ok(v) = row.try_get::<Option<i64>, usize>(ordinal) {
1589 return Ok(v.map(Value::I64).unwrap_or(Value::Null));
1590 }
1591 if let Ok(v) = row.try_get::<Option<f64>, usize>(ordinal) {
1592 return Ok(v.map(Value::F64).unwrap_or(Value::Null));
1593 }
1594 if let Ok(v) = row.try_get::<Option<bool>, usize>(ordinal) {
1595 return Ok(v.map(Value::Bool).unwrap_or(Value::Null));
1596 }
1597 match row.try_get::<Option<String>, usize>(ordinal) {
1598 Ok(v) => Ok(v.map(Value::String).unwrap_or(Value::Null)),
1599 Err(e) => Err(pg_decode_err(row, ordinal, type_name, e)),
1600 }
1601 }
1602 }
1603}
1604
1605fn row_to_value_with_coltype_pg(
1613 row: &sqlx::postgres::PgRow,
1614 ordinal: usize,
1615 col_type: ColType,
1616) -> Value {
1617 match col_type {
1618 ColType::Bool => match row.try_get::<Option<bool>, usize>(ordinal) {
1619 Ok(v) => v.map(Value::Bool).unwrap_or(Value::Null),
1620 Err(_) => Value::Null,
1621 },
1622 ColType::I8 => match row.try_get::<Option<i8>, usize>(ordinal) {
1623 Ok(v) => v.map(Value::I8).unwrap_or(Value::Null),
1624 Err(_) => Value::Null,
1625 },
1626 ColType::I16 => match row.try_get::<Option<i16>, usize>(ordinal) {
1627 Ok(v) => v.map(Value::I16).unwrap_or(Value::Null),
1628 Err(_) => Value::Null,
1629 },
1630 ColType::I32 => match row.try_get::<Option<i32>, usize>(ordinal) {
1632 Ok(v) => v.map(Value::I32).unwrap_or(Value::Null),
1633 Err(_) => Value::Null,
1634 },
1635 ColType::I64 => match row.try_get::<Option<i64>, usize>(ordinal) {
1636 Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
1637 Err(_) => Value::Null,
1638 },
1639 ColType::U8 => match row.try_get::<Option<i16>, usize>(ordinal) {
1644 Ok(v) => v.map(Value::I16).unwrap_or(Value::Null),
1645 Err(_) => Value::Null,
1646 },
1647 ColType::U16 => match row.try_get::<Option<i32>, usize>(ordinal) {
1648 Ok(v) => v.map(Value::I32).unwrap_or(Value::Null),
1649 Err(_) => Value::Null,
1650 },
1651 ColType::U32 => match row.try_get::<Option<i64>, usize>(ordinal) {
1652 Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
1653 Err(_) => Value::Null,
1654 },
1655 ColType::U64 => match row.try_get::<Option<i64>, usize>(ordinal) {
1657 Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
1658 Err(_) => Value::Null,
1659 },
1660 ColType::F32 => match row.try_get::<Option<f32>, usize>(ordinal) {
1661 Ok(v) => v.map(Value::F32).unwrap_or(Value::Null),
1662 Err(_) => Value::Null,
1663 },
1664 ColType::F64 => match row.try_get::<Option<f64>, usize>(ordinal) {
1665 Ok(v) => v.map(Value::F64).unwrap_or(Value::Null),
1666 Err(_) => Value::Null,
1667 },
1668 ColType::Decimal => match row.try_get::<Option<bigdecimal::BigDecimal>, usize>(ordinal) {
1670 Ok(Some(v)) => Value::Decimal(v.to_string()),
1671 Ok(None) => Value::Null,
1672 Err(_) => match row.try_get::<Option<rust_decimal::Decimal>, usize>(ordinal) {
1673 Ok(Some(v)) => Value::Decimal(v.to_string()),
1674 Ok(None) => Value::Null,
1675 Err(_) => Value::Null,
1676 },
1677 },
1678 ColType::String => match row.try_get::<Option<String>, usize>(ordinal) {
1679 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
1680 Err(_) => Value::Null,
1681 },
1682 ColType::Bytes => match row.try_get::<Option<Vec<u8>>, usize>(ordinal) {
1683 Ok(v) => v.map(Value::Bytes).unwrap_or(Value::Null),
1684 Err(_) => Value::Null,
1685 },
1686 ColType::Date => {
1688 match row.try_get::<Option<sqlx::types::chrono::NaiveDate>, usize>(ordinal) {
1689 Ok(v) => v.map(|d| Value::Date(d.to_string())).unwrap_or(Value::Null),
1690 Err(_) => Value::Null,
1691 }
1692 }
1693 ColType::DateTime => match row
1694 .try_get::<Option<sqlx::types::chrono::NaiveDateTime>, usize>(ordinal)
1695 {
1696 Ok(v) => v
1697 .map(|d| Value::DateTime(d.to_string()))
1698 .unwrap_or(Value::Null),
1699 Err(_) => match row
1700 .try_get::<Option<sqlx::types::chrono::DateTime<sqlx::types::chrono::Utc>>, usize>(
1701 ordinal,
1702 ) {
1703 Ok(v) => v
1704 .map(|d| Value::DateTime(d.format("%Y-%m-%dT%H:%M:%S%.6fZ").to_string()))
1705 .unwrap_or(Value::Null),
1706 Err(_) => Value::Null,
1707 },
1708 },
1709 ColType::Time => {
1710 match row.try_get::<Option<sqlx::types::chrono::NaiveTime>, usize>(ordinal) {
1711 Ok(v) => v.map(|t| Value::Time(t.to_string())).unwrap_or(Value::Null),
1712 Err(_) => Value::Null,
1713 }
1714 }
1715 ColType::Json => match row.try_get::<Option<serde_json::Value>, usize>(ordinal) {
1716 Ok(v) => v.map(|j| Value::Json(j.to_string())).unwrap_or(Value::Null),
1717 Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
1718 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
1719 Err(_) => Value::Null,
1720 },
1721 },
1722 ColType::Uuid => match row.try_get::<Option<sqlx::types::Uuid>, usize>(ordinal) {
1724 Ok(v) => v
1725 .map(|uuid| Value::String(uuid.to_string()))
1726 .unwrap_or(Value::Null),
1727 Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
1728 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
1729 Err(_) => Value::Null,
1730 },
1731 },
1732 ColType::Unknown => {
1733 if let Ok(v) = row.try_get::<Option<i64>, usize>(ordinal) {
1735 return v.map(Value::I64).unwrap_or(Value::Null);
1736 }
1737 if let Ok(v) = row.try_get::<Option<f64>, usize>(ordinal) {
1738 return v.map(Value::F64).unwrap_or(Value::Null);
1739 }
1740 if let Ok(v) = row.try_get::<Option<bool>, usize>(ordinal) {
1741 return v.map(Value::Bool).unwrap_or(Value::Null);
1742 }
1743 if let Ok(v) = row.try_get::<Option<String>, usize>(ordinal) {
1744 return v.map(Value::String).unwrap_or(Value::Null);
1745 }
1746 Value::Null
1747 }
1748 _ => Value::Null,
1750 }
1751}
1752
1753fn pg_rewrite_placeholders(sql: &str) -> std::borrow::Cow<'_, str> {
1762 if !sql.as_bytes().contains(&b'?') {
1763 return std::borrow::Cow::Borrowed(sql);
1764 }
1765 let mut out = String::with_capacity(sql.len() + 8);
1766 let mut chars = sql.chars().peekable();
1767 let mut in_single = false;
1768 let mut in_double = false;
1769 let mut in_line_comment = false;
1770 let mut in_block_comment = false;
1771 let mut dollar_tag: Option<String> = None;
1772 let mut prev = '\0';
1773 let mut n = 0usize;
1774 let mut saw_native = false;
1775 while let Some(c) = chars.next() {
1776 if in_line_comment {
1777 out.push(c);
1778 if c == '\n' {
1779 in_line_comment = false;
1780 }
1781 prev = c;
1782 continue;
1783 }
1784 if in_block_comment {
1785 out.push(c);
1786 if c == '*' && chars.peek() == Some(&'/') {
1787 out.push(chars.next().unwrap_or('/'));
1788 in_block_comment = false;
1789 }
1790 prev = c;
1791 continue;
1792 }
1793 if in_single {
1794 out.push(c);
1795 if c == '\'' {
1796 if chars.peek() == Some(&'\'') {
1797 out.push(chars.next().unwrap_or('\''));
1798 } else {
1799 in_single = false;
1800 }
1801 }
1802 prev = c;
1803 continue;
1804 }
1805 if in_double {
1806 out.push(c);
1807 if c == '"' {
1808 if chars.peek() == Some(&'"') {
1809 out.push(chars.next().unwrap_or('"'));
1810 } else {
1811 in_double = false;
1812 }
1813 }
1814 prev = c;
1815 continue;
1816 }
1817 if let Some(tag) = dollar_tag.as_deref() {
1818 if c == '$' {
1820 if let Some(extra) = pg_try_close_dollar_quote(&chars, tag) {
1821 out.push(c);
1822 for _ in 0..extra {
1823 out.push(chars.next().unwrap_or('$'));
1824 }
1825 dollar_tag = None;
1826 } else {
1827 out.push(c);
1828 }
1829 } else {
1830 out.push(c);
1831 }
1832 prev = c;
1833 continue;
1834 }
1835 match c {
1836 '\'' => {
1837 in_single = true;
1838 out.push(c);
1839 }
1840 '"' => {
1841 in_double = true;
1842 out.push(c);
1843 }
1844 '-' if chars.peek() == Some(&'-') => {
1845 out.push(c);
1846 out.push(chars.next().unwrap_or('-'));
1847 in_line_comment = true;
1848 }
1849 '/' if chars.peek() == Some(&'*') => {
1850 out.push(c);
1851 out.push(chars.next().unwrap_or('*'));
1852 in_block_comment = true;
1853 }
1854 '$' => {
1855 if let Some(rest) = pg_peek_dollar_quote_rest(&chars) {
1857 out.push('$');
1858 for ch in rest.chars() {
1859 out.push(ch);
1860 chars.next();
1861 }
1862 dollar_tag = Some(rest[..rest.len() - 1].to_string());
1863 } else if chars.peek().is_some_and(|next| next.is_ascii_digit()) {
1864 saw_native = true;
1866 out.push(c);
1867 } else {
1868 out.push(c);
1869 }
1870 }
1871 '?' if prev != '@'
1872 && !matches!(chars.peek(), Some('|') | Some('&'))
1873 && !pg_question_is_json_operator(&chars) =>
1874 {
1875 n += 1;
1876 out.push('$');
1877 out.push_str(&n.to_string());
1878 }
1879 _ => out.push(c),
1880 }
1881 prev = c;
1882 }
1883 if saw_native {
1884 return std::borrow::Cow::Borrowed(sql);
1885 }
1886 std::borrow::Cow::Owned(out)
1887}
1888
1889fn pg_question_is_json_operator(chars: &std::iter::Peekable<std::str::Chars<'_>>) -> bool {
1893 let mut lookahead = chars.clone();
1894 while let Some(&c) = lookahead.peek() {
1895 match c {
1896 ' ' | '\t' => {
1897 lookahead.next();
1898 }
1899 '\'' | '"' => return true,
1900 _ => return false,
1901 }
1902 }
1903 false
1904}
1905
1906fn pg_peek_dollar_quote_rest(chars: &std::iter::Peekable<std::str::Chars<'_>>) -> Option<String> {
1914 let mut lookahead = chars.clone();
1915 if lookahead.peek() == Some(&'$') {
1916 return Some("$".to_string());
1917 }
1918 let first = *lookahead.peek()?;
1919 if !(first.is_ascii_alphabetic() || first == '_') {
1920 return None;
1921 }
1922 let mut rest = String::new();
1923 while let Some(&c) = lookahead.peek() {
1924 match c {
1925 '$' => {
1926 rest.push('$');
1927 return Some(rest);
1928 }
1929 c if c.is_ascii_alphanumeric() || c == '_' => {
1930 rest.push(c);
1931 lookahead.next();
1932 }
1933 _ => return None,
1934 }
1935 }
1936 None
1937}
1938
1939fn pg_try_close_dollar_quote(
1944 chars: &std::iter::Peekable<std::str::Chars<'_>>,
1945 tag: &str,
1946) -> Option<usize> {
1947 let mut lookahead = chars.clone();
1948 for expected in tag.chars() {
1949 if lookahead.next() != Some(expected) {
1950 return None;
1951 }
1952 }
1953 if lookahead.next() == Some('$') {
1954 Some(tag.chars().count() + 1)
1955 } else {
1956 None
1957 }
1958}
1959
1960type PgQuery<'q> = sqlx::query::Query<'q, sqlx::Postgres, sqlx::postgres::PgArguments>;
1962
1963fn bind_pg_value<'q>(q: PgQuery<'q>, v: &'q Value) -> PgQuery<'q> {
1968 match v {
1969 Value::Null => q.bind(None::<i64>),
1970 Value::Bool(b) => q.bind(*b),
1971 Value::I8(n) => q.bind(*n),
1972 Value::I16(n) => q.bind(*n),
1973 Value::I32(n) => q.bind(*n),
1974 Value::I64(n) => q.bind(*n),
1975 Value::U8(n) => q.bind(*n as i16),
1977 Value::U16(n) => q.bind(*n as i32),
1978 Value::U32(n) => q.bind(*n as i64),
1979 Value::U64(n) => q.bind(*n as i64),
1980 Value::F32(f) => q.bind(*f),
1981 Value::F64(f) => q.bind(*f),
1982 Value::String(s) => q.bind(s.as_str()),
1983 Value::Decimal(s) => match s.parse::<bigdecimal::BigDecimal>() {
1985 Ok(d) => q.bind(d),
1986 Err(_) => q.bind(s.as_str()),
1987 },
1988 Value::Uuid(s) => match sqlx::types::Uuid::parse_str(s) {
1989 Ok(u) => q.bind(u),
1990 Err(_) => q.bind(s.as_str()),
1991 },
1992 Value::Date(s) => match parse_pg_date(s) {
1993 Some(d) => q.bind(d),
1994 None => q.bind(s.as_str()),
1995 },
1996 Value::DateTime(s) => match parse_pg_datetime(s) {
1997 Some(PgDateTime::Zoned(d)) => q.bind(d),
1999 Some(PgDateTime::Naive(d)) => q.bind(d),
2000 None => q.bind(s.as_str()),
2001 },
2002 Value::Time(s) => match parse_pg_time(s) {
2003 Some(t) => q.bind(t),
2004 None => q.bind(s.as_str()),
2005 },
2006 Value::Json(s) => match serde_json::from_str::<serde_json::Value>(s) {
2007 Ok(j) => q.bind(j),
2008 Err(_) => q.bind(s.as_str()),
2009 },
2010 Value::Bytes(b) => q.bind(b.as_slice()),
2011 other => q.bind(other.to_string()),
2012 }
2013}
2014
2015fn parse_pg_date(s: &str) -> Option<sqlx::types::chrono::NaiveDate> {
2017 use sqlx::types::chrono::NaiveDate;
2018 let head = s.get(..10)?;
2019 NaiveDate::parse_from_str(head, "%Y-%m-%d").ok()
2020}
2021
2022enum PgDateTime {
2024 Naive(sqlx::types::chrono::NaiveDateTime),
2026 Zoned(sqlx::types::chrono::DateTime<sqlx::types::chrono::Utc>),
2028}
2029
2030fn parse_pg_datetime(s: &str) -> Option<PgDateTime> {
2034 use sqlx::types::chrono::{DateTime, NaiveDateTime, Utc};
2035 let has_offset = s.ends_with('Z')
2036 || s.ends_with('z')
2037 || s.get(10..).is_some_and(|rest| {
2038 rest.char_indices().any(|(i, c)| {
2039 (c == '+' || c == '-')
2040 && rest[..i].chars().any(|p| p == ':')
2041 && i > 0
2043 })
2044 });
2045 if has_offset {
2046 if let Ok(dt) = DateTime::parse_from_rfc3339(s) {
2047 return Some(PgDateTime::Zoned(dt.with_timezone(&Utc)));
2048 }
2049 for fmt in [
2050 "%Y-%m-%d %H:%M:%S%.f%:z",
2051 "%Y-%m-%d %H:%M:%S%:z",
2052 "%Y-%m-%dT%H:%M:%S%.f%:z",
2053 "%Y-%m-%dT%H:%M:%S%:z",
2054 "%Y-%m-%d %H:%M:%S%.f%z",
2055 ] {
2056 if let Ok(dt) = DateTime::parse_from_str(s, fmt) {
2057 return Some(PgDateTime::Zoned(dt.with_timezone(&Utc)));
2058 }
2059 }
2060 }
2061 for fmt in ["%Y-%m-%d %H:%M:%S", "%Y-%m-%dT%H:%M:%S"] {
2062 if let Ok(dt) = NaiveDateTime::parse_from_str(s, fmt) {
2063 return Some(PgDateTime::Naive(dt));
2064 }
2065 let with_frac = format!("{fmt}%.f");
2066 if let Ok(dt) = NaiveDateTime::parse_from_str(s, &with_frac) {
2067 return Some(PgDateTime::Naive(dt));
2068 }
2069 }
2070 if let Ok(dt) = DateTime::parse_from_rfc3339(s) {
2071 return Some(PgDateTime::Zoned(dt.with_timezone(&Utc)));
2072 }
2073 None
2074}
2075
2076fn parse_pg_time(s: &str) -> Option<sqlx::types::chrono::NaiveTime> {
2078 use sqlx::types::chrono::NaiveTime;
2079 NaiveTime::parse_from_str(s, "%H:%M:%S")
2080 .or_else(|_| NaiveTime::parse_from_str(s, "%H:%M:%S%.f"))
2081 .ok()
2082}
2083
2084pub struct PgPoolHandle {
2085 pool: sqlx::PgPool,
2086}
2087
2088impl PgPoolHandle {
2089 pub async fn connect(url: &str) -> Result<Self, DbError> {
2090 Self::connect_with_max(url, 10).await
2092 }
2093
2094 pub async fn connect_with_max(url: &str, max_connections: u32) -> Result<Self, DbError> {
2097 let pool = sqlx::pool::PoolOptions::<sqlx::Postgres>::new()
2098 .max_connections(max_connections)
2099 .acquire_timeout(std::time::Duration::from_secs(30))
2100 .idle_timeout(Some(std::time::Duration::from_secs(600)))
2101 .max_lifetime(Some(std::time::Duration::from_secs(1800)))
2102 .connect(url)
2103 .await
2104 .map_err(map_sqlx_error)?;
2105 Ok(Self { pool })
2106 }
2107
2108 pub fn from_pool(pool: sqlx::PgPool) -> Self {
2109 Self { pool }
2110 }
2111
2112 pub fn pool(&self) -> &sqlx::PgPool {
2113 &self.pool
2114 }
2115}
2116
2117pub struct SqlxPgConnectionFactory {
2118 pool: Arc<PgPoolHandle>,
2119}
2120
2121impl SqlxPgConnectionFactory {
2122 pub fn new(pool: Arc<PgPoolHandle>) -> Self {
2123 Self { pool }
2124 }
2125}
2126
2127#[async_trait]
2128impl ConnectionFactory for SqlxPgConnectionFactory {
2129 async fn create(&self) -> Result<Box<dyn Connection>, DbError> {
2130 let conn = self.pool.pool.acquire().await.map_err(map_sqlx_error)?;
2131 Ok(Box::new(crate::metrics_conn::MetricsConn::new(Box::new(
2133 SqlxPgConnection {
2134 conn: Some(conn),
2135 connected: true,
2136 in_transaction: false,
2137 },
2138 ))))
2139 }
2140}
2141
2142pub struct SqlxPgConnection {
2143 conn: Option<sqlx::pool::PoolConnection<sqlx::Postgres>>,
2144 connected: bool,
2145 in_transaction: bool,
2146}
2147
2148impl Connection for SqlxPgConnection {
2149 fn db_type(&self) -> Option<DbType> {
2150 Some(DbType::PostgreSQL)
2151 }
2152
2153 fn session_reset_sql(&self) -> Option<&'static str> {
2157 Some("DISCARD ALL")
2158 }
2159
2160 fn execute<'a>(
2161 &'a mut self,
2162 sql: &'a str,
2163 ) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
2164 Box::pin(async move {
2165 let mut pool_conn = self
2166 .conn
2167 .take()
2168 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
2169 let result = if needs_raw_sql(sql) {
2172 (&mut *pool_conn)
2173 .execute(sqlx::raw_sql(sqlx::AssertSqlSafe(sql)))
2174 .await
2175 } else {
2176 (&mut *pool_conn).execute(sqlx::AssertSqlSafe(sql)).await
2177 };
2178 self.conn = Some(pool_conn);
2179
2180 match result {
2181 Ok(r) => Ok(r.rows_affected()),
2182 Err(e) => {
2183 let db_err = map_sqlx_error(e);
2184 if matches!(db_err, DbError::ConnectionError(_) | DbError::IoError(_)) {
2185 self.connected = false;
2186 }
2187 Err(db_err)
2188 }
2189 }
2190 })
2191 }
2192
2193 fn query<'a>(
2194 &'a mut self,
2195 sql: &'a str,
2196 ) -> Pin<Box<dyn Future<Output = Result<Vec<HashMap<String, Value>>, DbError>> + Send + 'a>>
2197 {
2198 Box::pin(async move {
2199 let mut pool_conn = self
2200 .conn
2201 .take()
2202 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
2203 let rows_result = (&mut *pool_conn).fetch_all(sqlx::AssertSqlSafe(sql)).await;
2206 self.conn = Some(pool_conn);
2207
2208 let rows = rows_result.map_err(map_sqlx_error)?;
2209 if rows.is_empty() {
2210 return Ok(Vec::new());
2211 }
2212 let col_types: Vec<ColType> = rows[0]
2215 .columns()
2216 .iter()
2217 .map(|col| ColType::parse_postgres(col.type_info().name()))
2218 .collect();
2219 let col_names: Vec<String> = rows[0]
2220 .columns()
2221 .iter()
2222 .map(|col| col.name().to_string())
2223 .collect();
2224 let mut result = Vec::with_capacity(rows.len());
2225 for row in &rows {
2226 let mut record = HashMap::with_capacity(col_types.len());
2227 for (i, name) in col_names.iter().enumerate() {
2228 let value = row_to_value_with_coltype_pg(row, i, col_types[i]);
2229 record.insert(name.clone(), value);
2230 }
2231 result.push(record);
2232 }
2233 Ok(result)
2234 })
2235 }
2236
2237 fn begin_transaction<'a>(
2238 &'a mut self,
2239 ) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
2240 Box::pin(async move {
2241 if self.in_transaction {
2242 return Err(DbError::Internal("transaction already started".to_string()));
2243 }
2244 self.execute("BEGIN").await?;
2245 self.in_transaction = true;
2246 Ok(())
2247 })
2248 }
2249
2250 fn commit<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
2251 Box::pin(async move {
2252 if self.in_transaction {
2253 self.execute("COMMIT").await?;
2254 self.in_transaction = false;
2255 }
2256 Ok(())
2257 })
2258 }
2259
2260 fn rollback<'a>(
2261 &'a mut self,
2262 ) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
2263 Box::pin(async move {
2264 if self.in_transaction {
2265 let result = self.execute("ROLLBACK").await;
2270 if result.is_ok() {
2271 self.in_transaction = false;
2272 }
2273 result.map(|_| ())
2274 } else {
2275 Ok(())
2276 }
2277 })
2278 }
2279
2280 fn is_connected(&self) -> bool {
2281 self.connected && self.conn.is_some()
2285 }
2286
2287 fn in_transaction(&self) -> bool {
2288 self.in_transaction
2289 }
2290
2291 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
2292 Box::pin(async move {
2293 match self.execute("SELECT 1").await {
2294 Ok(_) => true,
2295 Err(_) => {
2296 self.connected = false;
2297 false
2298 }
2299 }
2300 })
2301 }
2302
2303 fn close<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
2304 Box::pin(async move {
2305 if let Some(conn) = self.conn.take() {
2306 drop(conn);
2307 }
2308 self.connected = false;
2309 self.in_transaction = false;
2310 Ok(())
2311 })
2312 }
2313
2314 fn execute_with_params<'a>(
2332 &'a mut self,
2333 sql: &'a str,
2334 params: &'a [Value],
2335 ) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
2336 Box::pin(async move {
2337 if needs_raw_sql(sql) || params.is_empty() {
2338 return self.execute(sql).await;
2339 }
2340 let mut pool_conn = self
2341 .conn
2342 .take()
2343 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
2344 let rewritten = pg_rewrite_placeholders(sql);
2345 let mut q = sqlx::query(sqlx::AssertSqlSafe(rewritten.as_ref()));
2346 for v in params {
2347 q = bind_pg_value(q, v);
2348 }
2349 let result = q.execute(&mut *pool_conn).await;
2350 self.conn = Some(pool_conn);
2351 match result {
2352 Ok(r) => Ok(r.rows_affected()),
2353 Err(e) => {
2354 let db_err = map_sqlx_error(e);
2355 if matches!(db_err, DbError::ConnectionError(_) | DbError::IoError(_)) {
2356 self.connected = false;
2357 }
2358 Err(db_err)
2359 }
2360 }
2361 })
2362 }
2363
2364 fn query_with_params<'a>(
2368 &'a mut self,
2369 sql: &'a str,
2370 params: &'a [Value],
2371 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, DbError>> + Send + 'a>> {
2372 Box::pin(async move {
2373 if params.is_empty() {
2374 return self.query(sql).await;
2375 }
2376 let mut pool_conn = self
2377 .conn
2378 .take()
2379 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
2380 let rewritten = pg_rewrite_placeholders(sql);
2381 let mut q = sqlx::query(sqlx::AssertSqlSafe(rewritten.as_ref()));
2382 for v in params {
2383 q = bind_pg_value(q, v);
2384 }
2385 let rows_result = q.fetch_all(&mut *pool_conn).await;
2386 self.conn = Some(pool_conn);
2387 let rows = rows_result.map_err(map_sqlx_error)?;
2388 if rows.is_empty() {
2389 return Ok(Vec::new());
2390 }
2391 let col_types: Vec<ColType> = rows[0]
2393 .columns()
2394 .iter()
2395 .map(|col| ColType::parse_postgres(col.type_info().name()))
2396 .collect();
2397 let col_names: Vec<String> = rows[0]
2398 .columns()
2399 .iter()
2400 .map(|col| col.name().to_string())
2401 .collect();
2402 let mut result = Vec::with_capacity(rows.len());
2403 for row in &rows {
2404 let mut record = HashMap::with_capacity(col_types.len());
2405 for (i, name) in col_names.iter().enumerate() {
2406 let value = row_to_value_with_coltype_pg(row, i, col_types[i]);
2407 record.insert(name.clone(), value);
2408 }
2409 result.push(record);
2410 }
2411 Ok(result)
2412 })
2413 }
2414
2415 fn query_values<'a>(
2419 &'a mut self,
2420 sql: &'a str,
2421 ) -> Pin<Box<dyn Future<Output = Result<QueryValues, DbError>> + Send + 'a>> {
2422 Box::pin(async move {
2423 let mut pool_conn = self
2424 .conn
2425 .take()
2426 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
2427 let rows_result = (&mut *pool_conn).fetch_all(sqlx::AssertSqlSafe(sql)).await;
2428 self.conn = Some(pool_conn);
2429 let rows = rows_result.map_err(map_sqlx_error)?;
2430 if rows.is_empty() {
2431 return Ok((Vec::new(), Vec::new()));
2432 }
2433 let cols = rows[0].columns();
2434 let mut col_names: Vec<String> = Vec::with_capacity(cols.len());
2435 let mut col_types: Vec<ColType> = Vec::with_capacity(cols.len());
2436 for col in cols {
2437 col_names.push(col.name().to_string());
2438 col_types.push(ColType::parse_postgres(col.type_info().name()));
2439 }
2440 let mut result_rows: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
2441 for row in &rows {
2442 let mut row_values: Vec<Value> = Vec::with_capacity(col_names.len());
2443 for (idx, _) in col_names.iter().enumerate() {
2444 row_values.push(row_to_value_with_coltype_pg(row, idx, col_types[idx]));
2445 }
2446 result_rows.push(row_values);
2447 }
2448 Ok((col_names, result_rows))
2449 })
2450 }
2451
2452 fn query_values_with_params<'a>(
2456 &'a mut self,
2457 sql: &'a str,
2458 params: &'a [Value],
2459 ) -> Pin<Box<dyn Future<Output = Result<QueryValues, DbError>> + Send + 'a>> {
2460 Box::pin(async move {
2461 if params.is_empty() {
2462 return self.query_values(sql).await;
2463 }
2464 let mut pool_conn = self
2465 .conn
2466 .take()
2467 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
2468 let rewritten = pg_rewrite_placeholders(sql);
2469 let mut q = sqlx::query(sqlx::AssertSqlSafe(rewritten.as_ref()));
2470 for v in params {
2471 q = bind_pg_value(q, v);
2472 }
2473 let rows_result = q.fetch_all(&mut *pool_conn).await;
2474 self.conn = Some(pool_conn);
2475 let rows = rows_result.map_err(map_sqlx_error)?;
2476 if rows.is_empty() {
2477 return Ok((Vec::new(), Vec::new()));
2478 }
2479 let cols = rows[0].columns();
2480 let mut col_names: Vec<String> = Vec::with_capacity(cols.len());
2481 for col in cols {
2482 col_names.push(col.name().to_string());
2483 }
2484 let mut result_rows: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
2485 for row in rows {
2486 let mut row_values: Vec<Value> = Vec::with_capacity(col_names.len());
2487 for (idx, _) in col_names.iter().enumerate() {
2488 let ordinal = row.columns()[idx].ordinal();
2489 row_values.push(row_to_value_pg(&row, ordinal)?);
2490 }
2491 result_rows.push(row_values);
2492 }
2493 Ok((col_names, result_rows))
2494 })
2495 }
2496
2497 fn query_stream<'a>(
2507 &'a mut self,
2508 sql: &'a str,
2509 ) -> Pin<Box<dyn futures::Stream<Item = Result<HashMap<String, Value>, DbError>> + Send + 'a>>
2510 {
2511 Box::pin(async_stream::try_stream! {
2512 let mut pool_conn = self
2513 .conn
2514 .take()
2515 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
2516 let mut row_stream = sqlx::query(sqlx::AssertSqlSafe(sql)).fetch(&mut *pool_conn);
2517 while let Some(row_result) = row_stream.next().await {
2520 let row = row_result.map_err(map_sqlx_error)?;
2521 let cols = row.columns();
2522 let mut record = HashMap::with_capacity(cols.len());
2523 for (i, col) in cols.iter().enumerate() {
2524 let name = col.name().to_string();
2525 let value = row_to_value_pg(&row, i)?;
2527 record.insert(name, value);
2528 }
2529 yield record;
2530 }
2531 drop(row_stream);
2534 self.conn = Some(pool_conn);
2535 })
2536 }
2537}
2538
2539impl Drop for SqlxPgConnection {
2540 fn drop(&mut self) {
2541 if let Some(conn) = self.conn.take() {
2542 drop(conn);
2543 }
2544 }
2545}
2546
2547#[async_trait]
2568pub trait PgExtensions: Send + Sync {
2569 async fn listen(&mut self, channel: &str) -> Result<(), DbError>;
2580
2581 async fn notify(&mut self, channel: &str, payload: &str) -> Result<(), DbError>;
2592
2593 async fn copy_from_stdin(&mut self, sql: &str, data: &[u8]) -> Result<u64, DbError>;
2608}
2609
2610fn validate_pg_channel_name(channel: &str) -> Result<(), DbError> {
2615 if channel.is_empty() {
2616 return Err(DbError::Internal(
2617 "PG channel name must not be empty".to_string(),
2618 ));
2619 }
2620 if !channel
2621 .chars()
2622 .all(|c| c.is_ascii_alphanumeric() || c == '_')
2623 {
2624 return Err(DbError::Internal(format!(
2625 "invalid PG channel name: {} (only alphanumeric and underscore allowed)",
2626 channel
2627 )));
2628 }
2629 Ok(())
2630}
2631
2632#[async_trait]
2633impl PgExtensions for SqlxPgConnection {
2634 async fn listen(&mut self, channel: &str) -> Result<(), DbError> {
2635 validate_pg_channel_name(channel)?;
2636 self.execute(&format!("LISTEN {}", channel)).await?;
2638 Ok(())
2639 }
2640
2641 async fn notify(&mut self, channel: &str, payload: &str) -> Result<(), DbError> {
2642 validate_pg_channel_name(channel)?;
2643 let escaped_payload = payload.replace('\'', "''");
2645 self.execute(&format!("NOTIFY {}, '{}'", channel, escaped_payload))
2646 .await?;
2647 Ok(())
2648 }
2649
2650 async fn copy_from_stdin(&mut self, sql: &str, data: &[u8]) -> Result<u64, DbError> {
2651 let mut pool_conn = self
2652 .conn
2653 .take()
2654 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
2655 let mut copy = (*pool_conn)
2659 .copy_in_raw(sql)
2660 .await
2661 .map_err(map_sqlx_error)?;
2662 copy.send(data).await.map_err(map_sqlx_error)?;
2664 let result = copy.finish().await.map_err(map_sqlx_error)?;
2666 self.conn = Some(pool_conn);
2667 Ok(result)
2668 }
2669}
2670
2671pub async fn pg_bulk_insert(
2679 conn: &mut SqlxPgConnection,
2680 table: &str,
2681 columns: &[&str],
2682 rows: &[Vec<Value>],
2683) -> Result<u64, DbError> {
2684 if rows.is_empty() {
2685 return Ok(0);
2686 }
2687 let col_list = columns.join(", ");
2688 let cols_per_row = columns.len();
2689 let placeholders: Vec<String> = rows
2691 .iter()
2692 .enumerate()
2693 .map(|(row_idx, _)| {
2694 let base = row_idx * cols_per_row;
2695 let ph: Vec<String> = (0..cols_per_row)
2696 .map(|i| format!("${}", base + i + 1))
2697 .collect();
2698 format!("({})", ph.join(", "))
2699 })
2700 .collect();
2701 let sql = format!(
2702 "INSERT INTO {} ({}) VALUES {}",
2703 table,
2704 col_list,
2705 placeholders.join(", ")
2706 );
2707 let mut params: Vec<Value> = Vec::with_capacity(rows.len() * cols_per_row);
2709 for row in rows {
2710 for v in row {
2711 params.push(v.clone());
2712 }
2713 }
2714 conn.execute_with_params(&sql, ¶ms).await
2715}
2716#[cfg(test)]
2717mod tests {
2718 use super::*;
2719
2720 #[test]
2721 fn test_pg_rewrite_placeholders_basic() {
2722 assert_eq!(
2723 pg_rewrite_placeholders("SELECT * FROM t WHERE a = ? AND b = ?"),
2724 "SELECT * FROM t WHERE a = $1 AND b = $2"
2725 );
2726 }
2727
2728 #[test]
2729 fn test_pg_rewrite_placeholders_skips_native() {
2730 assert_eq!(
2731 pg_rewrite_placeholders("SELECT * FROM t WHERE a = $1"),
2732 "SELECT * FROM t WHERE a = $1"
2733 );
2734 }
2735
2736 #[test]
2737 fn test_pg_rewrite_placeholders_skips_literals_and_comments() {
2738 assert_eq!(
2739 pg_rewrite_placeholders("SELECT '?' AS q -- ?\nFROM t WHERE id = ?"),
2740 "SELECT '?' AS q -- ?\nFROM t WHERE id = $1"
2741 );
2742 assert_eq!(
2743 pg_rewrite_placeholders("SELECT /* ? */ \"a?\" FROM t WHERE id = ?"),
2744 "SELECT /* ? */ \"a?\" FROM t WHERE id = $1"
2745 );
2746 }
2747
2748 #[test]
2749 fn test_pg_rewrite_placeholders_keeps_json_operators() {
2750 assert_eq!(
2751 pg_rewrite_placeholders("SELECT * FROM t WHERE doc ? 'k' AND b = ?"),
2752 "SELECT * FROM t WHERE doc ? 'k' AND b = $1"
2753 );
2754 assert_eq!(
2755 pg_rewrite_placeholders("SELECT * FROM t WHERE doc ?| array['a'] AND b = ?"),
2756 "SELECT * FROM t WHERE doc ?| array['a'] AND b = $1"
2757 );
2758 assert_eq!(
2759 pg_rewrite_placeholders("SELECT * FROM t WHERE doc ?& array['a', 'b'] AND b = ?"),
2760 "SELECT * FROM t WHERE doc ?& array['a', 'b'] AND b = $1"
2761 );
2762 assert_eq!(
2763 pg_rewrite_placeholders("SELECT * FROM t WHERE doc @? '$.a' AND b = ?"),
2764 "SELECT * FROM t WHERE doc @? '$.a' AND b = $1"
2765 );
2766 }
2767
2768 #[test]
2769 fn test_pg_rewrite_placeholders_dollar_in_literal_still_rewrites() {
2770 assert_eq!(
2772 pg_rewrite_placeholders("SELECT '$1' AS lit, ? AS val"),
2773 "SELECT '$1' AS lit, $1 AS val"
2774 );
2775 assert_eq!(
2776 pg_rewrite_placeholders("SELECT 'price: $2' AS note FROM t WHERE id = ? AND code = ?"),
2777 "SELECT 'price: $2' AS note FROM t WHERE id = $1 AND code = $2"
2778 );
2779 }
2780
2781 #[test]
2782 fn test_pg_rewrite_placeholders_skips_dollar_quoted_strings() {
2783 assert_eq!(
2785 pg_rewrite_placeholders("SELECT $$?$$ AS lit, ? AS val"),
2786 "SELECT $$?$$ AS lit, $1 AS val"
2787 );
2788 assert_eq!(
2789 pg_rewrite_placeholders("SELECT $tag$a ? b$tag$ AS lit, ? AS val"),
2790 "SELECT $tag$a ? b$tag$ AS lit, $1 AS val"
2791 );
2792 let func = "CREATE FUNCTION f() RETURNS int AS $$ BEGIN RETURN (SELECT count(*) FROM t WHERE x = ?); END; $$ LANGUAGE plpgsql";
2794 assert_eq!(pg_rewrite_placeholders(func), func);
2795 }
2796
2797 #[test]
2798 fn test_pg_rewrite_placeholders_native_skips_even_with_question() {
2799 assert_eq!(
2801 pg_rewrite_placeholders("SELECT * FROM t WHERE a = $1 AND doc ? 'k'"),
2802 "SELECT * FROM t WHERE a = $1 AND doc ? 'k'"
2803 );
2804 }
2805
2806 #[test]
2807 fn test_parse_pg_date_and_datetime() {
2808 assert!(parse_pg_date("2026-10-28").is_some());
2809 assert!(parse_pg_date("2026-10-28 12:00:00").is_some());
2810 assert!(parse_pg_date("not-a-date").is_none());
2811 assert!(parse_pg_datetime("2026-10-28 12:00:00").is_some());
2812 assert!(parse_pg_datetime("2026-10-28T12:00:00").is_some());
2813 assert!(parse_pg_datetime("2026-10-28T12:00:00Z").is_some());
2814 assert!(parse_pg_datetime("bad").is_none());
2815 }
2816
2817 #[test]
2818 fn test_needs_raw_sql_begin() {
2819 assert!(needs_raw_sql("BEGIN"));
2820 assert!(needs_raw_sql("begin"));
2821 assert!(needs_raw_sql(" BEGIN"));
2822 assert!(needs_raw_sql("\tbegin"));
2823 }
2824
2825 #[test]
2826 fn test_needs_raw_sql_commit() {
2827 assert!(needs_raw_sql("COMMIT"));
2828 assert!(needs_raw_sql("commit"));
2829 assert!(needs_raw_sql(" COMMIT WORK"));
2830 }
2831
2832 #[test]
2833 fn test_needs_raw_sql_rollback() {
2834 assert!(needs_raw_sql("ROLLBACK"));
2835 assert!(needs_raw_sql("rollback"));
2836 assert!(needs_raw_sql(" ROLLBACK TO SAVEPOINT sp1"));
2837 }
2838
2839 #[test]
2840 fn test_needs_raw_sql_savepoint() {
2841 assert!(needs_raw_sql("SAVEPOINT sp1"));
2842 assert!(needs_raw_sql("savepoint s1"));
2843 }
2844
2845 #[test]
2846 fn test_needs_raw_sql_release() {
2847 assert!(needs_raw_sql("RELEASE sp1"));
2848 assert!(needs_raw_sql("release savepoint s1"));
2849 }
2850
2851 #[test]
2852 fn test_needs_raw_sql_set() {
2853 assert!(needs_raw_sql("SET autocommit = 0"));
2854 assert!(needs_raw_sql("set names utf8"));
2855 }
2856
2857 #[test]
2858 fn test_needs_raw_sql_use() {
2859 assert!(needs_raw_sql("USE mydb"));
2860 assert!(needs_raw_sql("use mydb"));
2861 }
2862
2863 #[test]
2864 fn test_needs_raw_sql_start_transaction() {
2865 assert!(needs_raw_sql("START TRANSACTION"));
2866 assert!(needs_raw_sql("start transaction"));
2867 assert!(needs_raw_sql(" START TRANSACTION READ ONLY"));
2868 }
2869
2870 #[test]
2871 fn test_needs_raw_sql_normal_queries() {
2872 assert!(!needs_raw_sql("SELECT * FROM users"));
2873 assert!(!needs_raw_sql("INSERT INTO t VALUES (1)"));
2874 assert!(!needs_raw_sql("UPDATE t SET x = 1"));
2875 assert!(!needs_raw_sql("DELETE FROM t"));
2876 assert!(!needs_raw_sql("CREATE TABLE t (id INT)"));
2877 assert!(!needs_raw_sql("DROP TABLE t"));
2878 assert!(!needs_raw_sql("ALTER TABLE t ADD COLUMN x INT"));
2879 }
2880
2881 #[test]
2882 fn test_needs_raw_sql_edge_cases() {
2883 assert!(!needs_raw_sql(""));
2884 assert!(!needs_raw_sql(" "));
2885 assert!(!needs_raw_sql("SELECT"));
2886 assert!(needs_raw_sql("BEGINNING"));
2887 assert!(needs_raw_sql("COMMITTEE"));
2888 assert!(!needs_raw_sql("INSERT"));
2889 assert!(!needs_raw_sql("DELETE"));
2890 }
2891
2892 #[test]
2895 fn test_validate_pg_channel_name_valid() {
2896 assert!(validate_pg_channel_name("my_channel").is_ok());
2897 assert!(validate_pg_channel_name("abc123").is_ok());
2898 assert!(validate_pg_channel_name("_test").is_ok());
2899 assert!(validate_pg_channel_name("a_b_c").is_ok());
2900 }
2901
2902 #[test]
2903 fn test_validate_pg_channel_name_empty() {
2904 let err = validate_pg_channel_name("").unwrap_err();
2905 assert!(matches!(err, DbError::Internal(ref s) if s.contains("empty")));
2906 }
2907
2908 #[test]
2909 fn test_validate_pg_channel_name_invalid_chars() {
2910 let err = validate_pg_channel_name("my-channel").unwrap_err();
2911 assert!(matches!(err, DbError::Internal(ref s) if s.contains("invalid")));
2912 assert!(validate_pg_channel_name("my.channel").is_err());
2913 assert!(validate_pg_channel_name("my channel").is_err());
2914 assert!(validate_pg_channel_name("my;channel").is_err());
2915 assert!(validate_pg_channel_name("my'channel").is_err());
2916 }
2917
2918 fn closed_sqlite() -> SqlxSqliteConnection {
2921 SqlxSqliteConnection {
2922 conn: None,
2923 connected: false,
2924 in_transaction: false,
2925 }
2926 }
2927
2928 #[tokio::test]
2929 async fn test_sqlite_execute_closed() {
2930 let mut c = closed_sqlite();
2931 let result = c.execute("SELECT 1").await;
2932 assert!(result.is_err());
2933 assert!(result.unwrap_err().to_string().contains("closed"));
2934 }
2935
2936 #[tokio::test]
2937 async fn test_sqlite_query_closed() {
2938 let mut c = closed_sqlite();
2939 let result = c.query("SELECT 1").await;
2940 assert!(result.is_err());
2941 }
2942
2943 #[tokio::test]
2944 async fn test_sqlite_begin_tx_already_started() {
2945 let mut c = SqlxSqliteConnection {
2946 conn: None,
2947 connected: false,
2948 in_transaction: true,
2949 };
2950 let result = c.begin_transaction().await;
2951 assert!(result.is_err());
2952 assert!(result.unwrap_err().to_string().contains("already started"));
2953 }
2954
2955 #[tokio::test]
2956 async fn test_sqlite_begin_tx_closed() {
2957 let mut c = closed_sqlite();
2958 let result = c.begin_transaction().await;
2959 assert!(result.is_err());
2960 }
2961
2962 #[tokio::test]
2963 async fn test_sqlite_commit_no_tx() {
2964 let mut c = closed_sqlite();
2965 let result = c.commit().await;
2966 assert!(result.is_ok());
2967 }
2968
2969 #[tokio::test]
2970 async fn test_sqlite_rollback_no_tx() {
2971 let mut c = closed_sqlite();
2972 let result = c.rollback().await;
2973 assert!(result.is_ok());
2974 }
2975
2976 #[test]
2977 fn test_sqlite_is_connected_false_when_conn_none() {
2978 let c = closed_sqlite();
2979 assert!(!c.is_connected());
2980 }
2981
2982 #[test]
2983 fn test_sqlite_is_connected_false_when_connected_false() {
2984 let c = SqlxSqliteConnection {
2985 conn: None,
2986 connected: true,
2987 in_transaction: false,
2988 };
2989 assert!(!c.is_connected());
2990 }
2991
2992 #[test]
2993 fn test_sqlite_in_transaction_flag() {
2994 let c = SqlxSqliteConnection {
2995 conn: None,
2996 connected: false,
2997 in_transaction: true,
2998 };
2999 assert!(c.in_transaction());
3000 let c2 = closed_sqlite();
3001 assert!(!c2.in_transaction());
3002 }
3003
3004 #[tokio::test]
3005 async fn test_sqlite_ping_closed() {
3006 let mut c = closed_sqlite();
3007 let result = c.ping().await;
3008 assert!(!result);
3009 assert!(!c.is_connected());
3010 }
3011
3012 #[tokio::test]
3013 async fn test_sqlite_close_closed() {
3014 let mut c = closed_sqlite();
3015 let result = c.close().await;
3016 assert!(result.is_ok());
3017 assert!(!c.is_connected());
3018 assert!(!c.in_transaction());
3019 }
3020
3021 #[tokio::test]
3022 async fn test_sqlite_execute_with_params_closed() {
3023 let mut c = closed_sqlite();
3024 let result = c
3025 .execute_with_params("INSERT INTO t VALUES (?)", &[Value::I64(1)])
3026 .await;
3027 assert!(result.is_err());
3028 }
3029
3030 #[tokio::test]
3031 async fn test_sqlite_execute_with_params_empty_delegates() {
3032 let mut c = closed_sqlite();
3033 let result = c.execute_with_params("INSERT INTO t VALUES (1)", &[]).await;
3034 assert!(result.is_err());
3035 }
3036
3037 #[tokio::test]
3038 async fn test_sqlite_query_with_params_empty_delegates() {
3039 let mut c = closed_sqlite();
3040 let result = c.query_with_params("SELECT * FROM t", &[]).await;
3041 assert!(result.is_err());
3042 }
3043
3044 #[tokio::test]
3045 async fn test_sqlite_query_with_params_closed() {
3046 let mut c = closed_sqlite();
3047 let result = c
3048 .query_with_params("SELECT * FROM t WHERE id = ?", &[Value::I64(1)])
3049 .await;
3050 assert!(result.is_err());
3051 }
3052
3053 fn closed_mysql() -> SqlxMySqlConnection {
3056 SqlxMySqlConnection {
3057 conn: None,
3058 connected: false,
3059 in_transaction: false,
3060 }
3061 }
3062
3063 #[tokio::test]
3064 async fn test_mysql_execute_closed() {
3065 let mut c = closed_mysql();
3066 assert!(c.execute("SELECT 1").await.is_err());
3067 }
3068
3069 #[tokio::test]
3070 async fn test_mysql_query_closed() {
3071 let mut c = closed_mysql();
3072 assert!(c.query("SELECT 1").await.is_err());
3073 }
3074
3075 #[test]
3076 fn test_mysql_is_connected_false() {
3077 let c = closed_mysql();
3078 assert!(!c.is_connected());
3079 }
3080
3081 #[tokio::test]
3082 async fn test_mysql_close_closed() {
3083 let mut c = closed_mysql();
3084 assert!(c.close().await.is_ok());
3085 assert!(!c.is_connected());
3086 }
3087
3088 #[tokio::test]
3089 async fn test_mysql_commit_no_tx() {
3090 let mut c = closed_mysql();
3091 assert!(c.commit().await.is_ok());
3092 }
3093
3094 #[tokio::test]
3095 async fn test_mysql_rollback_no_tx() {
3096 let mut c = closed_mysql();
3097 assert!(c.rollback().await.is_ok());
3098 }
3099
3100 #[tokio::test]
3101 async fn test_mysql_begin_tx_already_started() {
3102 let mut c = SqlxMySqlConnection {
3103 conn: None,
3104 connected: false,
3105 in_transaction: true,
3106 };
3107 let result = c.begin_transaction().await;
3108 assert!(result.is_err());
3109 }
3110
3111 #[tokio::test]
3112 async fn test_mysql_ping_closed() {
3113 let mut c = closed_mysql();
3114 assert!(!c.ping().await);
3115 }
3116
3117 fn closed_pg() -> SqlxPgConnection {
3120 SqlxPgConnection {
3121 conn: None,
3122 connected: false,
3123 in_transaction: false,
3124 }
3125 }
3126
3127 #[tokio::test]
3128 async fn test_pg_execute_closed() {
3129 let mut c = closed_pg();
3130 assert!(c.execute("SELECT 1").await.is_err());
3131 }
3132
3133 #[tokio::test]
3134 async fn test_pg_query_closed() {
3135 let mut c = closed_pg();
3136 assert!(c.query("SELECT 1").await.is_err());
3137 }
3138
3139 #[test]
3140 fn test_pg_is_connected_false() {
3141 let c = closed_pg();
3142 assert!(!c.is_connected());
3143 }
3144
3145 #[tokio::test]
3146 async fn test_pg_close_closed() {
3147 let mut c = closed_pg();
3148 assert!(c.close().await.is_ok());
3149 assert!(!c.is_connected());
3150 }
3151
3152 #[tokio::test]
3153 async fn test_pg_commit_no_tx() {
3154 let mut c = closed_pg();
3155 assert!(c.commit().await.is_ok());
3156 }
3157
3158 #[tokio::test]
3159 async fn test_pg_rollback_no_tx() {
3160 let mut c = closed_pg();
3161 assert!(c.rollback().await.is_ok());
3162 }
3163
3164 #[tokio::test]
3165 async fn test_pg_begin_tx_already_started() {
3166 let mut c = SqlxPgConnection {
3167 conn: None,
3168 connected: false,
3169 in_transaction: true,
3170 };
3171 assert!(c.begin_transaction().await.is_err());
3172 }
3173
3174 #[tokio::test]
3175 async fn test_pg_ping_closed() {
3176 let mut c = closed_pg();
3177 assert!(!c.ping().await);
3178 }
3179
3180 #[tokio::test]
3183 async fn test_mysql_execute_with_params_closed() {
3184 let mut c = closed_mysql();
3185 assert!(c
3186 .execute_with_params("INSERT INTO t VALUES (?)", &[Value::I64(1)])
3187 .await
3188 .is_err());
3189 }
3190
3191 #[tokio::test]
3192 async fn test_mysql_execute_with_params_empty_delegates() {
3193 let mut c = closed_mysql();
3194 assert!(c
3195 .execute_with_params("INSERT INTO t VALUES (1)", &[])
3196 .await
3197 .is_err());
3198 }
3199
3200 #[tokio::test]
3201 async fn test_mysql_query_with_params_closed() {
3202 let mut c = closed_mysql();
3203 assert!(c
3204 .query_with_params("SELECT * FROM t WHERE id = ?", &[Value::I64(1)])
3205 .await
3206 .is_err());
3207 }
3208
3209 #[tokio::test]
3210 async fn test_mysql_query_with_params_empty_delegates() {
3211 let mut c = closed_mysql();
3212 assert!(c.query_with_params("SELECT * FROM t", &[]).await.is_err());
3213 }
3214
3215 #[tokio::test]
3216 async fn test_mysql_query_values_closed() {
3217 let mut c = closed_mysql();
3218 assert!(c.query_values("SELECT * FROM t").await.is_err());
3219 }
3220
3221 #[tokio::test]
3222 async fn test_mysql_query_values_with_params_closed() {
3223 let mut c = closed_mysql();
3224 assert!(c
3225 .query_values_with_params("SELECT * FROM t WHERE id = ?", &[Value::I64(1)])
3226 .await
3227 .is_err());
3228 }
3229
3230 #[tokio::test]
3231 async fn test_mysql_query_values_with_params_empty_delegates() {
3232 let mut c = closed_mysql();
3233 assert!(c
3234 .query_values_with_params("SELECT * FROM t", &[])
3235 .await
3236 .is_err());
3237 }
3238
3239 #[tokio::test]
3242 async fn test_pg_execute_with_params_closed() {
3243 let mut c = closed_pg();
3244 assert!(c
3245 .execute_with_params("INSERT INTO t VALUES ($1)", &[Value::I64(1)])
3246 .await
3247 .is_err());
3248 }
3249
3250 #[tokio::test]
3251 async fn test_pg_execute_with_params_empty_delegates() {
3252 let mut c = closed_pg();
3253 assert!(c
3254 .execute_with_params("INSERT INTO t VALUES (1)", &[])
3255 .await
3256 .is_err());
3257 }
3258
3259 #[tokio::test]
3260 async fn test_pg_query_with_params_closed() {
3261 let mut c = closed_pg();
3262 assert!(c
3263 .query_with_params("SELECT * FROM t WHERE id = $1", &[Value::I64(1)])
3264 .await
3265 .is_err());
3266 }
3267
3268 #[tokio::test]
3269 async fn test_pg_query_with_params_empty_delegates() {
3270 let mut c = closed_pg();
3271 assert!(c.query_with_params("SELECT * FROM t", &[]).await.is_err());
3272 }
3273
3274 #[tokio::test]
3275 async fn test_pg_query_values_closed() {
3276 let mut c = closed_pg();
3277 assert!(c.query_values("SELECT * FROM t").await.is_err());
3278 }
3279
3280 #[tokio::test]
3281 async fn test_pg_query_values_with_params_closed() {
3282 let mut c = closed_pg();
3283 assert!(c
3284 .query_values_with_params("SELECT * FROM t WHERE id = $1", &[Value::I64(1)])
3285 .await
3286 .is_err());
3287 }
3288
3289 #[tokio::test]
3290 async fn test_pg_query_values_with_params_empty_delegates() {
3291 let mut c = closed_pg();
3292 assert!(c
3293 .query_values_with_params("SELECT * FROM t", &[])
3294 .await
3295 .is_err());
3296 }
3297
3298 #[tokio::test]
3301 async fn test_sqlite_backup_closed_conn() {
3302 let mut c = closed_sqlite();
3303 let result = sqlite_backup(&mut c, "/tmp/backup.db").await;
3304 assert!(result.is_err());
3305 }
3306
3307 #[tokio::test]
3310 async fn test_sqlite_query_values_closed() {
3311 let mut c = closed_sqlite();
3312 assert!(c.query_values("SELECT * FROM t").await.is_err());
3313 }
3314
3315 #[tokio::test]
3316 async fn test_sqlite_query_values_with_params_closed() {
3317 let mut c = closed_sqlite();
3318 assert!(c
3319 .query_values_with_params("SELECT * FROM t WHERE id = ?", &[Value::I64(1)])
3320 .await
3321 .is_err());
3322 }
3323
3324 #[tokio::test]
3325 async fn test_sqlite_query_values_with_params_empty_delegates() {
3326 let mut c = closed_sqlite();
3327 assert!(c
3328 .query_values_with_params("SELECT * FROM t", &[])
3329 .await
3330 .is_err());
3331 }
3332
3333 #[tokio::test]
3336 async fn test_mysql_bulk_insert_empty_rows() {
3337 let mut c = closed_mysql();
3338 let result = mysql_bulk_insert(&mut c, "t", &["a", "b"], &[]).await;
3339 assert_eq!(result.unwrap(), 0);
3340 }
3341
3342 #[tokio::test]
3343 async fn test_mysql_bulk_insert_closed_conn() {
3344 let mut c = closed_mysql();
3345 let rows = vec![vec![Value::I64(1), Value::String("x".to_string())]];
3346 let result = mysql_bulk_insert(&mut c, "t", &["a", "b"], &rows).await;
3347 assert!(result.is_err());
3348 }
3349
3350 #[tokio::test]
3351 async fn test_pg_bulk_insert_empty_rows() {
3352 let mut c = closed_pg();
3353 let result = pg_bulk_insert(&mut c, "t", &["a", "b"], &[]).await;
3354 assert_eq!(result.unwrap(), 0);
3355 }
3356
3357 #[tokio::test]
3358 async fn test_pg_bulk_insert_closed_conn() {
3359 let mut c = closed_pg();
3360 let rows = vec![vec![Value::I64(1), Value::String("x".to_string())]];
3361 let result = pg_bulk_insert(&mut c, "t", &["a", "b"], &rows).await;
3362 assert!(result.is_err());
3363 }
3364
3365 #[test]
3368 fn test_parse_pg_datetime_naive() {
3369 match parse_pg_datetime("2026-10-28 13:14:15.123456") {
3370 Some(PgDateTime::Naive(dt)) => {
3371 assert_eq!(
3372 dt.format("%Y-%m-%d %H:%M:%S%.6f").to_string(),
3373 "2026-10-28 13:14:15.123456"
3374 );
3375 }
3376 _ => panic!("无偏移输入必须解析为 Naive"),
3377 }
3378 }
3379
3380 #[test]
3381 fn test_parse_pg_datetime_zoned_forms() {
3382 for s in [
3384 "2026-10-28T13:14:15+08:00",
3385 "2026-10-28 13:14:15+08:00",
3386 "2026-10-28 13:14:15 +0800",
3387 "2026-10-28T05:14:15Z",
3388 ] {
3389 match parse_pg_datetime(s) {
3390 Some(PgDateTime::Zoned(dt)) => {
3391 assert_eq!(
3393 dt.format("%Y-%m-%d %H:%M:%S").to_string(),
3394 "2026-10-28 05:14:15",
3395 "输入 {s}"
3396 );
3397 }
3398 other => panic!("{s} 必须解析为 Zoned,实际 {:?}", other.is_some()),
3399 }
3400 }
3401 }
3402
3403 #[test]
3404 fn test_parse_pg_time_fractional() {
3405 assert!(parse_pg_time("13:14:15").is_some());
3406 assert!(parse_pg_time("13:14:15.123456").is_some());
3407 assert!(parse_pg_time("13:14").is_none());
3408 }
3409
3410 #[test]
3413 fn test_bigdecimal_38_digits_no_truncation() {
3414 let s = "1234567890123456789012345678.9012345678"; let d: bigdecimal::BigDecimal = s.parse().unwrap();
3416 assert_eq!(d.to_string(), s, "BigDecimal 往返必须逐位一致");
3417 assert!(
3419 s.parse::<rust_decimal::Decimal>().is_err() || {
3420 let rd: rust_decimal::Decimal = s.parse().unwrap();
3421 rd.to_string() != s
3422 }
3423 );
3424 }
3425}