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::{ColType, Connection, ConnectionFactory, DbError, QueryRows, QueryValues, Value};
21
22use crate::error::map_sqlx_error;
23
24fn needs_raw_sql(sql: &str) -> bool {
27 let trimmed = sql.trim_start();
28 let upper = trimmed.to_uppercase();
29 upper.starts_with("BEGIN")
30 || upper.starts_with("COMMIT")
31 || upper.starts_with("ROLLBACK")
32 || upper.starts_with("SAVEPOINT")
33 || upper.starts_with("RELEASE")
34 || upper.starts_with("SET ")
35 || upper.starts_with("USE ")
36 || upper.starts_with("START TRANSACTION")
37}
38
39fn row_to_value_with_coltype_sqlite(
57 row: &sqlx::sqlite::SqliteRow,
58 ordinal: usize,
59 col_type: ColType,
60) -> Value {
61 match col_type {
62 ColType::Bool => match row.try_get::<Option<bool>, usize>(ordinal) {
63 Ok(v) => v.map(Value::Bool).unwrap_or(Value::Null),
64 Err(_) => Value::Null,
65 },
66 ColType::I8 => match row.try_get::<Option<i8>, usize>(ordinal) {
67 Ok(v) => v.map(Value::I8).unwrap_or(Value::Null),
68 Err(_) => Value::Null,
69 },
70 ColType::I16 => match row.try_get::<Option<i16>, usize>(ordinal) {
71 Ok(v) => v.map(Value::I16).unwrap_or(Value::Null),
72 Err(_) => Value::Null,
73 },
74 ColType::I32 => match row.try_get::<Option<i32>, usize>(ordinal) {
75 Ok(v) => v.map(Value::I32).unwrap_or(Value::Null),
76 Err(_) => Value::Null,
77 },
78 ColType::I64 => match row.try_get::<Option<i64>, usize>(ordinal) {
79 Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
80 Err(_) => Value::Null,
81 },
82 ColType::U8 => match row.try_get::<Option<u8>, usize>(ordinal) {
83 Ok(v) => v.map(Value::U8).unwrap_or(Value::Null),
84 Err(_) => Value::Null,
85 },
86 ColType::U16 => match row.try_get::<Option<u16>, usize>(ordinal) {
87 Ok(v) => v.map(Value::U16).unwrap_or(Value::Null),
88 Err(_) => Value::Null,
89 },
90 ColType::U32 => match row.try_get::<Option<u32>, usize>(ordinal) {
91 Ok(v) => v.map(Value::U32).unwrap_or(Value::Null),
92 Err(_) => Value::Null,
93 },
94 ColType::U64 => match row.try_get::<Option<i64>, usize>(ordinal) {
96 Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
97 Err(_) => Value::Null,
98 },
99 ColType::F32 => match row.try_get::<Option<f32>, usize>(ordinal) {
100 Ok(v) => v.map(Value::F32).unwrap_or(Value::Null),
101 Err(_) => Value::Null,
102 },
103 ColType::F64 => match row.try_get::<Option<f64>, usize>(ordinal) {
104 Ok(v) => v.map(Value::F64).unwrap_or(Value::Null),
105 Err(_) => Value::Null,
106 },
107 ColType::Decimal => match row.try_get::<Option<String>, usize>(ordinal) {
108 Ok(v) => v.map(Value::Decimal).unwrap_or(Value::Null),
109 Err(_) => Value::Null,
110 },
111 ColType::String => match row.try_get::<Option<String>, usize>(ordinal) {
112 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
113 Err(_) => Value::Null,
114 },
115 ColType::Bytes => match row.try_get::<Option<Vec<u8>>, usize>(ordinal) {
116 Ok(v) => v.map(Value::Bytes).unwrap_or(Value::Null),
117 Err(_) => Value::Null,
118 },
119 ColType::Date | ColType::DateTime | ColType::Time | ColType::Json | ColType::Uuid => {
121 match row.try_get::<Option<String>, usize>(ordinal) {
122 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
123 Err(_) => Value::Null,
124 }
125 }
126 ColType::Unknown => {
127 if let Ok(v) = row.try_get::<Option<i64>, usize>(ordinal) {
134 return v.map(Value::I64).unwrap_or(Value::Null);
135 }
136 if let Ok(v) = row.try_get::<Option<f64>, usize>(ordinal) {
137 return v.map(Value::F64).unwrap_or(Value::Null);
138 }
139 if let Ok(v) = row.try_get::<Option<bool>, usize>(ordinal) {
140 return v.map(Value::Bool).unwrap_or(Value::Null);
141 }
142 if let Ok(v) = row.try_get::<Option<String>, usize>(ordinal) {
143 return v.map(Value::String).unwrap_or(Value::Null);
144 }
145 Value::Null
146 }
147 _ => Value::Null,
149 }
150}
151
152pub struct SqlitePoolHandle {
153 pool: sqlx::SqlitePool,
154}
155
156impl SqlitePoolHandle {
157 pub async fn connect(url: &str) -> Result<Self, DbError> {
158 let opts = sqlx::sqlite::SqliteConnectOptions::from_str(url)
164 .map_err(map_sqlx_error)?
165 .journal_mode(sqlx::sqlite::SqliteJournalMode::Wal)
166 .synchronous(sqlx::sqlite::SqliteSynchronous::Normal)
167 .busy_timeout(std::time::Duration::from_secs(5))
168 .pragma("mmap_size", "268435456");
169 let pool = sqlx::sqlite::SqlitePoolOptions::new()
170 .max_connections(10)
171 .acquire_timeout(std::time::Duration::from_secs(30))
172 .idle_timeout(Some(std::time::Duration::from_secs(600)))
173 .max_lifetime(Some(std::time::Duration::from_secs(1800)))
174 .connect_with(opts)
175 .await
176 .map_err(map_sqlx_error)?;
177 Ok(Self { pool })
178 }
179
180 pub fn from_pool(pool: sqlx::SqlitePool) -> Self {
181 Self { pool }
182 }
183
184 pub fn pool(&self) -> &sqlx::SqlitePool {
185 &self.pool
186 }
187}
188
189pub struct SqlxSqliteConnectionFactory {
190 pool: Arc<SqlitePoolHandle>,
191}
192
193impl SqlxSqliteConnectionFactory {
194 pub fn new(pool: Arc<SqlitePoolHandle>) -> Self {
195 Self { pool }
196 }
197}
198
199#[async_trait]
200impl ConnectionFactory for SqlxSqliteConnectionFactory {
201 async fn create(&self) -> Result<Box<dyn Connection>, DbError> {
202 let conn = self.pool.pool.acquire().await.map_err(map_sqlx_error)?;
203 Ok(Box::new(SqlxSqliteConnection {
204 conn: Some(conn),
205 connected: true,
206 in_transaction: false,
207 }))
208 }
209}
210
211pub struct SqlxSqliteConnection {
212 conn: Option<sqlx::pool::PoolConnection<sqlx::Sqlite>>,
213 connected: bool,
214 in_transaction: bool,
215}
216
217impl Connection for SqlxSqliteConnection {
218 fn execute<'a>(
219 &'a mut self,
220 sql: &'a str,
221 ) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
222 Box::pin(async move {
223 let mut pool_conn = self
224 .conn
225 .take()
226 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
227 let result = if needs_raw_sql(sql) {
230 (&mut *pool_conn)
231 .execute(sqlx::raw_sql(sqlx::AssertSqlSafe(sql)))
232 .await
233 } else {
234 (&mut *pool_conn).execute(sqlx::AssertSqlSafe(sql)).await
235 };
236 self.conn = Some(pool_conn);
237
238 match result {
239 Ok(r) => Ok(r.rows_affected()),
240 Err(e) => {
241 let db_err = map_sqlx_error(e);
242 if matches!(db_err, DbError::ConnectionError(_) | DbError::IoError(_)) {
243 self.connected = false;
244 }
245 Err(db_err)
246 }
247 }
248 })
249 }
250
251 fn query<'a>(
252 &'a mut self,
253 sql: &'a str,
254 ) -> Pin<Box<dyn Future<Output = Result<Vec<HashMap<String, Value>>, DbError>> + Send + 'a>>
255 {
256 Box::pin(async move {
257 let mut pool_conn = self
258 .conn
259 .take()
260 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
261 let rows_result = (&mut *pool_conn).fetch_all(sqlx::AssertSqlSafe(sql)).await;
264 self.conn = Some(pool_conn);
265
266 let rows = rows_result.map_err(map_sqlx_error)?;
267 if rows.is_empty() {
268 return Ok(Vec::new());
269 }
270 let col_types: Vec<ColType> = rows[0]
273 .columns()
274 .iter()
275 .map(|col| ColType::parse_sqlite(col.type_info().name()))
276 .collect();
277 let col_names: Vec<String> = rows[0]
278 .columns()
279 .iter()
280 .map(|col| col.name().to_string())
281 .collect();
282 let mut result = Vec::with_capacity(rows.len());
283 for row in &rows {
284 let mut record = HashMap::with_capacity(col_types.len());
285 for (i, name) in col_names.iter().enumerate() {
286 let value = row_to_value_with_coltype_sqlite(row, i, col_types[i]);
287 record.insert(name.clone(), value);
288 }
289 result.push(record);
290 }
291 Ok(result)
292 })
293 }
294
295 fn begin_transaction<'a>(
296 &'a mut self,
297 ) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
298 Box::pin(async move {
299 if self.in_transaction {
300 return Err(DbError::Internal("transaction already started".to_string()));
301 }
302 self.execute("BEGIN").await?;
303 self.in_transaction = true;
304 Ok(())
305 })
306 }
307
308 fn commit<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
309 Box::pin(async move {
310 if self.in_transaction {
311 self.execute("COMMIT").await?;
312 self.in_transaction = false;
313 }
314 Ok(())
315 })
316 }
317
318 fn rollback<'a>(
319 &'a mut self,
320 ) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
321 Box::pin(async move {
322 if self.in_transaction {
323 let result = self.execute("ROLLBACK").await;
328 if result.is_ok() {
329 self.in_transaction = false;
330 }
331 result.map(|_| ())
332 } else {
333 Ok(())
334 }
335 })
336 }
337
338 fn is_connected(&self) -> bool {
339 self.connected && self.conn.is_some()
343 }
344
345 fn in_transaction(&self) -> bool {
346 self.in_transaction
347 }
348
349 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
350 Box::pin(async move {
351 match self.execute("SELECT 1").await {
352 Ok(_) => true,
353 Err(_) => {
354 self.connected = false;
355 false
356 }
357 }
358 })
359 }
360
361 fn close<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
362 Box::pin(async move {
363 if let Some(conn) = self.conn.take() {
364 drop(conn);
365 }
366 self.connected = false;
367 self.in_transaction = false;
368 Ok(())
369 })
370 }
371
372 fn execute_with_params<'a>(
378 &'a mut self,
379 sql: &'a str,
380 params: &'a [Value],
381 ) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
382 Box::pin(async move {
383 if needs_raw_sql(sql) || params.is_empty() {
384 return self.execute(sql).await;
385 }
386 let mut pool_conn = self
387 .conn
388 .take()
389 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
390 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
391 for v in params {
392 q = match v {
393 Value::Null => q.bind(None::<i64>),
394 Value::Bool(b) => q.bind(*b),
395 Value::I8(n) => q.bind(*n),
396 Value::I16(n) => q.bind(*n),
397 Value::I32(n) => q.bind(*n),
398 Value::I64(n) => q.bind(*n),
399 Value::U8(n) => q.bind(*n),
400 Value::U16(n) => q.bind(*n),
401 Value::U32(n) => q.bind(*n),
402 Value::U64(n) => q.bind(*n as i64),
403 Value::F32(f) => q.bind(*f),
404 Value::F64(f) => q.bind(*f),
405 Value::String(s) => q.bind(s.as_str()),
406 Value::Decimal(s) => q.bind(s.as_str()),
407 Value::Bytes(b) => q.bind(b.as_slice()),
408 other => q.bind(other.to_string()),
409 };
410 }
411 let result = q.execute(&mut *pool_conn).await;
412 self.conn = Some(pool_conn);
413 match result {
414 Ok(r) => Ok(r.rows_affected()),
415 Err(e) => {
416 let db_err = map_sqlx_error(e);
417 if matches!(db_err, DbError::ConnectionError(_) | DbError::IoError(_)) {
418 self.connected = false;
419 }
420 Err(db_err)
421 }
422 }
423 })
424 }
425
426 fn query_with_params<'a>(
431 &'a mut self,
432 sql: &'a str,
433 params: &'a [Value],
434 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, DbError>> + Send + 'a>> {
435 Box::pin(async move {
436 if params.is_empty() {
437 return self.query(sql).await;
438 }
439 let mut pool_conn = self
440 .conn
441 .take()
442 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
443 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
444 for v in params {
445 q = match v {
446 Value::Null => q.bind(None::<i64>),
447 Value::Bool(b) => q.bind(*b),
448 Value::I8(n) => q.bind(*n),
449 Value::I16(n) => q.bind(*n),
450 Value::I32(n) => q.bind(*n),
451 Value::I64(n) => q.bind(*n),
452 Value::U8(n) => q.bind(*n),
453 Value::U16(n) => q.bind(*n),
454 Value::U32(n) => q.bind(*n),
455 Value::U64(n) => q.bind(*n as i64),
456 Value::F32(f) => q.bind(*f),
457 Value::F64(f) => q.bind(*f),
458 Value::String(s) => q.bind(s.as_str()),
459 Value::Decimal(s) => q.bind(s.as_str()),
460 Value::Bytes(b) => q.bind(b.as_slice()),
461 other => q.bind(other.to_string()),
462 };
463 }
464 let rows_result = q.fetch_all(&mut *pool_conn).await;
465 self.conn = Some(pool_conn);
466 let rows = rows_result.map_err(map_sqlx_error)?;
467 if rows.is_empty() {
468 return Ok(Vec::new());
469 }
470 let col_types: Vec<ColType> = rows[0]
473 .columns()
474 .iter()
475 .map(|col| ColType::parse_sqlite(col.type_info().name()))
476 .collect();
477 let col_names: Vec<String> = rows[0]
478 .columns()
479 .iter()
480 .map(|col| col.name().to_string())
481 .collect();
482 let mut result = Vec::with_capacity(rows.len());
483 for row in &rows {
484 let mut record = HashMap::with_capacity(col_types.len());
485 for (i, name) in col_names.iter().enumerate() {
486 let value = row_to_value_with_coltype_sqlite(row, i, col_types[i]);
487 record.insert(name.clone(), value);
488 }
489 result.push(record);
490 }
491 Ok(result)
492 })
493 }
494
495 fn query_values<'a>(
501 &'a mut self,
502 sql: &'a str,
503 ) -> Pin<Box<dyn Future<Output = Result<QueryValues, DbError>> + Send + 'a>> {
504 Box::pin(async move {
505 let mut pool_conn = self
506 .conn
507 .take()
508 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
509 let rows_result = (&mut *pool_conn).fetch_all(sqlx::AssertSqlSafe(sql)).await;
510 self.conn = Some(pool_conn);
511 let rows = rows_result.map_err(map_sqlx_error)?;
512 if rows.is_empty() {
513 return Ok((Vec::new(), Vec::new()));
514 }
515 let cols = rows[0].columns();
516 let mut col_names: Vec<String> = Vec::with_capacity(cols.len());
517 let mut col_types: Vec<ColType> = Vec::with_capacity(cols.len());
518 for col in cols {
519 col_names.push(col.name().to_string());
520 col_types.push(ColType::parse_sqlite(col.type_info().name()));
521 }
522 let mut result_rows: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
523 for row in &rows {
524 let mut row_values: Vec<Value> = Vec::with_capacity(col_names.len());
525 for (idx, _) in col_names.iter().enumerate() {
526 row_values.push(row_to_value_with_coltype_sqlite(row, idx, col_types[idx]));
527 }
528 result_rows.push(row_values);
529 }
530 Ok((col_names, result_rows))
531 })
532 }
533
534 fn query_values_with_params<'a>(
538 &'a mut self,
539 sql: &'a str,
540 params: &'a [Value],
541 ) -> Pin<Box<dyn Future<Output = Result<QueryValues, DbError>> + Send + 'a>> {
542 Box::pin(async move {
543 if params.is_empty() {
544 return self.query_values(sql).await;
545 }
546 let mut pool_conn = self
547 .conn
548 .take()
549 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
550 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
551 for v in params {
552 q = match v {
553 Value::Null => q.bind(None::<i64>),
554 Value::Bool(b) => q.bind(*b),
555 Value::I8(n) => q.bind(*n),
556 Value::I16(n) => q.bind(*n),
557 Value::I32(n) => q.bind(*n),
558 Value::I64(n) => q.bind(*n),
559 Value::U8(n) => q.bind(*n),
560 Value::U16(n) => q.bind(*n),
561 Value::U32(n) => q.bind(*n),
562 Value::U64(n) => q.bind(*n as i64),
563 Value::F32(f) => q.bind(*f),
564 Value::F64(f) => q.bind(*f),
565 Value::String(s) => q.bind(s.as_str()),
566 Value::Decimal(s) => q.bind(s.as_str()),
567 Value::Bytes(b) => q.bind(b.as_slice()),
568 other => q.bind(other.to_string()),
569 };
570 }
571 let rows_result = q.fetch_all(&mut *pool_conn).await;
572 self.conn = Some(pool_conn);
573 let rows = rows_result.map_err(map_sqlx_error)?;
574 if rows.is_empty() {
575 return Ok((Vec::new(), Vec::new()));
576 }
577 let cols = rows[0].columns();
578 let mut col_names: Vec<String> = Vec::with_capacity(cols.len());
579 let mut col_types: Vec<ColType> = Vec::with_capacity(cols.len());
580 for col in cols {
581 col_names.push(col.name().to_string());
582 col_types.push(ColType::parse_sqlite(col.type_info().name()));
583 }
584 let mut result_rows: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
585 for row in &rows {
586 let mut row_values: Vec<Value> = Vec::with_capacity(col_names.len());
587 for (idx, _) in col_names.iter().enumerate() {
588 row_values.push(row_to_value_with_coltype_sqlite(row, idx, col_types[idx]));
589 }
590 result_rows.push(row_values);
591 }
592 Ok((col_names, result_rows))
593 })
594 }
595
596 fn query_stream<'a>(
602 &'a mut self,
603 sql: &'a str,
604 ) -> Pin<Box<dyn futures::Stream<Item = Result<HashMap<String, Value>, DbError>> + Send + 'a>>
605 {
606 Box::pin(async_stream::try_stream! {
607 let mut pool_conn = self
608 .conn
609 .take()
610 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
611 let mut row_stream = sqlx::query(sqlx::AssertSqlSafe(sql)).fetch(&mut *pool_conn);
612 let mut col_types: Vec<ColType> = Vec::new();
613 let mut col_names: Vec<String> = Vec::new();
614 let mut first_row = true;
615 while let Some(row_result) = row_stream.next().await {
616 let row = row_result.map_err(map_sqlx_error)?;
617 if first_row {
618 for col in row.columns() {
619 col_names.push(col.name().to_string());
620 col_types.push(ColType::parse_sqlite(col.type_info().name()));
621 }
622 first_row = false;
623 }
624 let mut record = HashMap::with_capacity(col_names.len());
625 for (i, name) in col_names.iter().enumerate() {
626 let value = row_to_value_with_coltype_sqlite(&row, i, col_types[i]);
627 record.insert(name.clone(), value);
628 }
629 yield record;
630 }
631 drop(row_stream);
634 self.conn = Some(pool_conn);
635 })
636 }
637}
638
639impl Drop for SqlxSqliteConnection {
640 fn drop(&mut self) {
641 if let Some(conn) = self.conn.take() {
642 drop(conn);
643 }
644 }
645}
646
647pub async fn sqlite_backup(
652 conn: &mut SqlxSqliteConnection,
653 dest_path: &str,
654) -> Result<(), DbError> {
655 let escaped_path = dest_path.replace('\'', "''");
656 let sql = format!("VACUUM INTO '{}'", escaped_path);
657 conn.execute(&sql).await?;
658 Ok(())
659}
660
661fn row_to_value_mysql(row: &sqlx::mysql::MySqlRow, ordinal: usize) -> Value {
667 use sqlx::TypeInfo;
668 let type_name = row.columns()[ordinal].type_info().name();
669 match type_name {
670 "BOOLEAN" => match row.try_get::<Option<bool>, usize>(ordinal) {
671 Ok(v) => v.map(Value::Bool).unwrap_or(Value::Null),
672 Err(_) => Value::Null,
673 },
674 "TINYINT" | "TINYINT UNSIGNED" => match row.try_get::<Option<i8>, usize>(ordinal) {
675 Ok(v) => v.map(Value::I8).unwrap_or(Value::Null),
676 Err(_) => match row.try_get::<Option<u8>, usize>(ordinal) {
677 Ok(v) => v.map(Value::U8).unwrap_or(Value::Null),
678 Err(_) => Value::Null,
679 },
680 },
681 "SMALLINT" | "SMALLINT UNSIGNED" => match row.try_get::<Option<i16>, usize>(ordinal) {
682 Ok(v) => v.map(Value::I16).unwrap_or(Value::Null),
683 Err(_) => match row.try_get::<Option<u16>, usize>(ordinal) {
684 Ok(v) => v.map(Value::U16).unwrap_or(Value::Null),
685 Err(_) => Value::Null,
686 },
687 },
688 "INT" | "INT UNSIGNED" | "MEDIUMINT" | "MEDIUMINT UNSIGNED" => {
689 match row.try_get::<Option<i32>, usize>(ordinal) {
690 Ok(v) => v.map(Value::I32).unwrap_or(Value::Null),
691 Err(_) => match row.try_get::<Option<u32>, usize>(ordinal) {
692 Ok(v) => v.map(Value::U32).unwrap_or(Value::Null),
693 Err(_) => Value::Null,
694 },
695 }
696 }
697 "BIGINT" | "BIGINT UNSIGNED" => match row.try_get::<Option<i64>, usize>(ordinal) {
698 Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
699 Err(_) => match row.try_get::<Option<u64>, usize>(ordinal) {
700 Ok(v) => v.map(Value::U64).unwrap_or(Value::Null),
701 Err(_) => Value::Null,
702 },
703 },
704 "FLOAT" => match row.try_get::<Option<f32>, usize>(ordinal) {
705 Ok(v) => v.map(Value::F32).unwrap_or(Value::Null),
706 Err(_) => Value::Null,
707 },
708 "DOUBLE" => match row.try_get::<Option<f64>, usize>(ordinal) {
709 Ok(v) => v.map(Value::F64).unwrap_or(Value::Null),
710 Err(_) => Value::Null,
711 },
712 "VARCHAR" | "TEXT" | "CHAR" | "TINYTEXT" | "MEDIUMTEXT" | "LONGTEXT" | "ENUM" => {
713 match row.try_get::<Option<String>, usize>(ordinal) {
714 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
715 Err(_) => Value::Null,
716 }
717 }
718 "BLOB" | "TINYBLOB" | "MEDIUMBLOB" | "LONGBLOB" | "BINARY" | "VARBINARY" => {
719 match row.try_get::<Option<Vec<u8>>, usize>(ordinal) {
720 Ok(v) => v.map(Value::Bytes).unwrap_or(Value::Null),
721 Err(_) => Value::Null,
722 }
723 }
724 "DECIMAL" | "NUMERIC" | "NEWDECIMAL" => {
726 match row.try_get::<Option<rust_decimal::Decimal>, usize>(ordinal) {
727 Ok(Some(v)) => Value::Decimal(v.to_string()),
728 Ok(None) => Value::Null,
729 Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
730 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
731 Err(_) => Value::Null,
732 },
733 }
734 }
735 _ => {
736 if let Ok(v) = row.try_get::<Option<i64>, usize>(ordinal) {
738 return v.map(Value::I64).unwrap_or(Value::Null);
739 }
740 if let Ok(v) = row.try_get::<Option<f64>, usize>(ordinal) {
741 return v.map(Value::F64).unwrap_or(Value::Null);
742 }
743 if let Ok(v) = row.try_get::<Option<bool>, usize>(ordinal) {
744 return v.map(Value::Bool).unwrap_or(Value::Null);
745 }
746 if let Ok(v) = row.try_get::<Option<String>, usize>(ordinal) {
747 return v.map(Value::String).unwrap_or(Value::Null);
748 }
749 Value::Null
750 }
751 }
752}
753
754fn row_to_value_with_coltype_mysql(
762 row: &sqlx::mysql::MySqlRow,
763 ordinal: usize,
764 col_type: ColType,
765) -> Value {
766 match col_type {
767 ColType::Bool => match row.try_get::<Option<bool>, usize>(ordinal) {
768 Ok(v) => v.map(Value::Bool).unwrap_or(Value::Null),
769 Err(_) => Value::Null,
770 },
771 ColType::I8 => match row.try_get::<Option<i8>, usize>(ordinal) {
772 Ok(v) => v.map(Value::I8).unwrap_or(Value::Null),
773 Err(_) => match row.try_get::<Option<u8>, usize>(ordinal) {
774 Ok(v) => v.map(Value::U8).unwrap_or(Value::Null),
775 Err(_) => Value::Null,
776 },
777 },
778 ColType::I16 => match row.try_get::<Option<i16>, usize>(ordinal) {
779 Ok(v) => v.map(Value::I16).unwrap_or(Value::Null),
780 Err(_) => match row.try_get::<Option<u16>, usize>(ordinal) {
781 Ok(v) => v.map(Value::U16).unwrap_or(Value::Null),
782 Err(_) => Value::Null,
783 },
784 },
785 ColType::I32 => match row.try_get::<Option<i32>, usize>(ordinal) {
786 Ok(v) => v.map(Value::I32).unwrap_or(Value::Null),
787 Err(_) => match row.try_get::<Option<u32>, usize>(ordinal) {
788 Ok(v) => v.map(Value::U32).unwrap_or(Value::Null),
789 Err(_) => Value::Null,
790 },
791 },
792 ColType::I64 => match row.try_get::<Option<i64>, usize>(ordinal) {
793 Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
794 Err(_) => match row.try_get::<Option<u64>, usize>(ordinal) {
795 Ok(v) => v.map(Value::U64).unwrap_or(Value::Null),
796 Err(_) => Value::Null,
797 },
798 },
799 ColType::U8 => match row.try_get::<Option<u8>, usize>(ordinal) {
800 Ok(v) => v.map(Value::U8).unwrap_or(Value::Null),
801 Err(_) => Value::Null,
802 },
803 ColType::U16 => match row.try_get::<Option<u16>, usize>(ordinal) {
804 Ok(v) => v.map(Value::U16).unwrap_or(Value::Null),
805 Err(_) => Value::Null,
806 },
807 ColType::U32 => match row.try_get::<Option<u32>, usize>(ordinal) {
808 Ok(v) => v.map(Value::U32).unwrap_or(Value::Null),
809 Err(_) => Value::Null,
810 },
811 ColType::U64 => match row.try_get::<Option<u64>, usize>(ordinal) {
812 Ok(v) => v.map(Value::U64).unwrap_or(Value::Null),
813 Err(_) => Value::Null,
814 },
815 ColType::F32 => match row.try_get::<Option<f32>, usize>(ordinal) {
816 Ok(v) => v.map(Value::F32).unwrap_or(Value::Null),
817 Err(_) => Value::Null,
818 },
819 ColType::F64 => match row.try_get::<Option<f64>, usize>(ordinal) {
820 Ok(v) => v.map(Value::F64).unwrap_or(Value::Null),
821 Err(_) => Value::Null,
822 },
823 ColType::Decimal => match row.try_get::<Option<rust_decimal::Decimal>, usize>(ordinal) {
825 Ok(Some(v)) => Value::Decimal(v.to_string()),
826 Ok(None) => Value::Null,
827 Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
828 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
829 Err(_) => Value::Null,
830 },
831 },
832 ColType::String => match row.try_get::<Option<String>, usize>(ordinal) {
833 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
834 Err(_) => Value::Null,
835 },
836 ColType::Bytes => match row.try_get::<Option<Vec<u8>>, usize>(ordinal) {
837 Ok(v) => v.map(Value::Bytes).unwrap_or(Value::Null),
838 Err(_) => Value::Null,
839 },
840 ColType::Date | ColType::DateTime | ColType::Time | ColType::Json | ColType::Uuid => {
842 match row.try_get::<Option<String>, usize>(ordinal) {
843 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
844 Err(_) => Value::Null,
845 }
846 }
847 ColType::Unknown => {
848 if let Ok(v) = row.try_get::<Option<i64>, usize>(ordinal) {
850 return v.map(Value::I64).unwrap_or(Value::Null);
851 }
852 if let Ok(v) = row.try_get::<Option<f64>, usize>(ordinal) {
853 return v.map(Value::F64).unwrap_or(Value::Null);
854 }
855 if let Ok(v) = row.try_get::<Option<bool>, usize>(ordinal) {
856 return v.map(Value::Bool).unwrap_or(Value::Null);
857 }
858 if let Ok(v) = row.try_get::<Option<String>, usize>(ordinal) {
859 return v.map(Value::String).unwrap_or(Value::Null);
860 }
861 Value::Null
862 }
863 _ => Value::Null,
865 }
866}
867
868pub struct MySqlPoolHandle {
869 pool: sqlx::MySqlPool,
870}
871
872impl MySqlPoolHandle {
873 pub async fn connect(url: &str) -> Result<Self, DbError> {
874 let pool = sqlx::pool::PoolOptions::<sqlx::MySql>::new()
875 .max_connections(10)
876 .acquire_timeout(std::time::Duration::from_secs(30))
877 .idle_timeout(Some(std::time::Duration::from_secs(600)))
878 .max_lifetime(Some(std::time::Duration::from_secs(1800)))
879 .connect(url)
880 .await
881 .map_err(map_sqlx_error)?;
882 Ok(Self { pool })
883 }
884
885 pub fn from_pool(pool: sqlx::MySqlPool) -> Self {
886 Self { pool }
887 }
888
889 pub fn pool(&self) -> &sqlx::MySqlPool {
890 &self.pool
891 }
892}
893
894pub struct SqlxMySqlConnectionFactory {
895 pool: Arc<MySqlPoolHandle>,
896}
897
898impl SqlxMySqlConnectionFactory {
899 pub fn new(pool: Arc<MySqlPoolHandle>) -> Self {
900 Self { pool }
901 }
902}
903
904#[async_trait]
905impl ConnectionFactory for SqlxMySqlConnectionFactory {
906 async fn create(&self) -> Result<Box<dyn Connection>, DbError> {
907 let conn = self.pool.pool.acquire().await.map_err(map_sqlx_error)?;
908 Ok(Box::new(SqlxMySqlConnection {
909 conn: Some(conn),
910 connected: true,
911 in_transaction: false,
912 }))
913 }
914}
915
916pub struct SqlxMySqlConnection {
917 conn: Option<sqlx::pool::PoolConnection<sqlx::MySql>>,
918 connected: bool,
919 in_transaction: bool,
920}
921
922impl Connection for SqlxMySqlConnection {
923 fn execute<'a>(
924 &'a mut self,
925 sql: &'a str,
926 ) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
927 Box::pin(async move {
928 let mut pool_conn = self
929 .conn
930 .take()
931 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
932 let result = if needs_raw_sql(sql) {
935 (&mut *pool_conn)
936 .execute(sqlx::raw_sql(sqlx::AssertSqlSafe(sql)))
937 .await
938 } else {
939 (&mut *pool_conn).execute(sqlx::AssertSqlSafe(sql)).await
940 };
941 self.conn = Some(pool_conn);
942
943 match result {
944 Ok(r) => Ok(r.rows_affected()),
945 Err(e) => {
946 let db_err = map_sqlx_error(e);
947 if matches!(db_err, DbError::ConnectionError(_) | DbError::IoError(_)) {
948 self.connected = false;
949 }
950 Err(db_err)
951 }
952 }
953 })
954 }
955
956 fn query<'a>(
957 &'a mut self,
958 sql: &'a str,
959 ) -> Pin<Box<dyn Future<Output = Result<Vec<HashMap<String, Value>>, DbError>> + Send + 'a>>
960 {
961 Box::pin(async move {
962 let mut pool_conn = self
963 .conn
964 .take()
965 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
966 let rows_result = (&mut *pool_conn).fetch_all(sqlx::AssertSqlSafe(sql)).await;
969 self.conn = Some(pool_conn);
970
971 let rows = rows_result.map_err(map_sqlx_error)?;
972 if rows.is_empty() {
973 return Ok(Vec::new());
974 }
975 let col_types: Vec<ColType> = rows[0]
978 .columns()
979 .iter()
980 .map(|col| ColType::parse_mysql(col.type_info().name()))
981 .collect();
982 let col_names: Vec<String> = rows[0]
983 .columns()
984 .iter()
985 .map(|col| col.name().to_string())
986 .collect();
987 let mut result = Vec::with_capacity(rows.len());
988 for row in &rows {
989 let mut record = HashMap::with_capacity(col_types.len());
990 for (i, name) in col_names.iter().enumerate() {
991 let value = row_to_value_with_coltype_mysql(row, i, col_types[i]);
992 record.insert(name.clone(), value);
993 }
994 result.push(record);
995 }
996 Ok(result)
997 })
998 }
999
1000 fn begin_transaction<'a>(
1001 &'a mut self,
1002 ) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
1003 Box::pin(async move {
1004 if self.in_transaction {
1005 return Err(DbError::Internal("transaction already started".to_string()));
1006 }
1007 self.execute("BEGIN").await?;
1008 self.in_transaction = true;
1009 Ok(())
1010 })
1011 }
1012
1013 fn commit<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
1014 Box::pin(async move {
1015 if self.in_transaction {
1016 self.execute("COMMIT").await?;
1017 self.in_transaction = false;
1018 }
1019 Ok(())
1020 })
1021 }
1022
1023 fn rollback<'a>(
1024 &'a mut self,
1025 ) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
1026 Box::pin(async move {
1027 if self.in_transaction {
1028 let result = self.execute("ROLLBACK").await;
1033 if result.is_ok() {
1034 self.in_transaction = false;
1035 }
1036 result.map(|_| ())
1037 } else {
1038 Ok(())
1039 }
1040 })
1041 }
1042
1043 fn is_connected(&self) -> bool {
1044 self.connected && self.conn.is_some()
1048 }
1049
1050 fn in_transaction(&self) -> bool {
1051 self.in_transaction
1052 }
1053
1054 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
1055 Box::pin(async move {
1056 match self.execute("SELECT 1").await {
1057 Ok(_) => true,
1058 Err(_) => {
1059 self.connected = false;
1060 false
1061 }
1062 }
1063 })
1064 }
1065
1066 fn close<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
1067 Box::pin(async move {
1068 if let Some(conn) = self.conn.take() {
1069 drop(conn);
1070 }
1071 self.connected = false;
1072 self.in_transaction = false;
1073 Ok(())
1074 })
1075 }
1076
1077 fn execute_with_params<'a>(
1083 &'a mut self,
1084 sql: &'a str,
1085 params: &'a [Value],
1086 ) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
1087 Box::pin(async move {
1088 if needs_raw_sql(sql) || params.is_empty() {
1089 return self.execute(sql).await;
1090 }
1091 let mut pool_conn = self
1092 .conn
1093 .take()
1094 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1095 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
1096 for v in params {
1097 q = match v {
1098 Value::Null => q.bind(None::<i64>),
1099 Value::Bool(b) => q.bind(*b),
1100 Value::I8(n) => q.bind(*n),
1101 Value::I16(n) => q.bind(*n),
1102 Value::I32(n) => q.bind(*n),
1103 Value::I64(n) => q.bind(*n),
1104 Value::U8(n) => q.bind(*n),
1105 Value::U16(n) => q.bind(*n),
1106 Value::U32(n) => q.bind(*n),
1107 Value::U64(n) => q.bind(*n as i64),
1108 Value::F32(f) => q.bind(*f),
1109 Value::F64(f) => q.bind(*f),
1110 Value::String(s) => q.bind(s.as_str()),
1111 Value::Decimal(s) => q.bind(s.as_str()),
1112 Value::Bytes(b) => q.bind(b.as_slice()),
1113 other => q.bind(other.to_string()),
1114 };
1115 }
1116 let result = q.execute(&mut *pool_conn).await;
1117 self.conn = Some(pool_conn);
1118 match result {
1119 Ok(r) => Ok(r.rows_affected()),
1120 Err(e) => {
1121 let db_err = map_sqlx_error(e);
1122 if matches!(db_err, DbError::ConnectionError(_) | DbError::IoError(_)) {
1123 self.connected = false;
1124 }
1125 Err(db_err)
1126 }
1127 }
1128 })
1129 }
1130
1131 fn query_with_params<'a>(
1133 &'a mut self,
1134 sql: &'a str,
1135 params: &'a [Value],
1136 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, DbError>> + Send + 'a>> {
1137 Box::pin(async move {
1138 if params.is_empty() {
1139 return self.query(sql).await;
1140 }
1141 let mut pool_conn = self
1142 .conn
1143 .take()
1144 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1145 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
1146 for v in params {
1147 q = match v {
1148 Value::Null => q.bind(None::<i64>),
1149 Value::Bool(b) => q.bind(*b),
1150 Value::I8(n) => q.bind(*n),
1151 Value::I16(n) => q.bind(*n),
1152 Value::I32(n) => q.bind(*n),
1153 Value::I64(n) => q.bind(*n),
1154 Value::U8(n) => q.bind(*n),
1155 Value::U16(n) => q.bind(*n),
1156 Value::U32(n) => q.bind(*n),
1157 Value::U64(n) => q.bind(*n as i64),
1158 Value::F32(f) => q.bind(*f),
1159 Value::F64(f) => q.bind(*f),
1160 Value::String(s) => q.bind(s.as_str()),
1161 Value::Decimal(s) => q.bind(s.as_str()),
1162 Value::Bytes(b) => q.bind(b.as_slice()),
1163 other => q.bind(other.to_string()),
1164 };
1165 }
1166 let rows_result = q.fetch_all(&mut *pool_conn).await;
1167 self.conn = Some(pool_conn);
1168 let rows = rows_result.map_err(map_sqlx_error)?;
1169 if rows.is_empty() {
1170 return Ok(Vec::new());
1171 }
1172 let col_names: Vec<String> = rows[0]
1173 .columns()
1174 .iter()
1175 .map(|col| col.name().to_string())
1176 .collect();
1177 let mut result = Vec::with_capacity(rows.len());
1178 for row in &rows {
1179 let mut record = HashMap::with_capacity(col_names.len());
1180 for (i, name) in col_names.iter().enumerate() {
1181 let value = row_to_value_mysql(row, i);
1182 record.insert(name.clone(), value);
1183 }
1184 result.push(record);
1185 }
1186 Ok(result)
1187 })
1188 }
1189
1190 fn query_values<'a>(
1194 &'a mut self,
1195 sql: &'a str,
1196 ) -> Pin<Box<dyn Future<Output = Result<QueryValues, DbError>> + Send + 'a>> {
1197 Box::pin(async move {
1198 let mut pool_conn = self
1199 .conn
1200 .take()
1201 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1202 let rows_result = (&mut *pool_conn).fetch_all(sqlx::AssertSqlSafe(sql)).await;
1203 self.conn = Some(pool_conn);
1204 let rows = rows_result.map_err(map_sqlx_error)?;
1205 if rows.is_empty() {
1206 return Ok((Vec::new(), Vec::new()));
1207 }
1208 let cols = rows[0].columns();
1209 let mut col_names: Vec<String> = Vec::with_capacity(cols.len());
1210 let mut col_types: Vec<ColType> = Vec::with_capacity(cols.len());
1211 for col in cols {
1212 col_names.push(col.name().to_string());
1213 col_types.push(ColType::parse_mysql(col.type_info().name()));
1214 }
1215 let mut result_rows: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
1216 for row in &rows {
1217 let mut row_values: Vec<Value> = Vec::with_capacity(col_names.len());
1218 for (idx, _) in col_names.iter().enumerate() {
1219 row_values.push(row_to_value_with_coltype_mysql(row, idx, col_types[idx]));
1220 }
1221 result_rows.push(row_values);
1222 }
1223 Ok((col_names, result_rows))
1224 })
1225 }
1226
1227 fn query_values_with_params<'a>(
1229 &'a mut self,
1230 sql: &'a str,
1231 params: &'a [Value],
1232 ) -> Pin<Box<dyn Future<Output = Result<QueryValues, DbError>> + Send + 'a>> {
1233 Box::pin(async move {
1234 if params.is_empty() {
1235 return self.query_values(sql).await;
1236 }
1237 let mut pool_conn = self
1238 .conn
1239 .take()
1240 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1241 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
1242 for v in params {
1243 q = match v {
1244 Value::Null => q.bind(None::<i64>),
1245 Value::Bool(b) => q.bind(*b),
1246 Value::I8(n) => q.bind(*n),
1247 Value::I16(n) => q.bind(*n),
1248 Value::I32(n) => q.bind(*n),
1249 Value::I64(n) => q.bind(*n),
1250 Value::U8(n) => q.bind(*n),
1251 Value::U16(n) => q.bind(*n),
1252 Value::U32(n) => q.bind(*n),
1253 Value::U64(n) => q.bind(*n as i64),
1254 Value::F32(f) => q.bind(*f),
1255 Value::F64(f) => q.bind(*f),
1256 Value::String(s) => q.bind(s.as_str()),
1257 Value::Decimal(s) => q.bind(s.as_str()),
1258 Value::Bytes(b) => q.bind(b.as_slice()),
1259 other => q.bind(other.to_string()),
1260 };
1261 }
1262 let rows_result = q.fetch_all(&mut *pool_conn).await;
1263 self.conn = Some(pool_conn);
1264 let rows = rows_result.map_err(map_sqlx_error)?;
1265 if rows.is_empty() {
1266 return Ok((Vec::new(), Vec::new()));
1267 }
1268 let cols = rows[0].columns();
1269 let mut col_names: Vec<String> = Vec::with_capacity(cols.len());
1270 let mut col_types: Vec<ColType> = Vec::with_capacity(cols.len());
1271 for col in cols {
1272 col_names.push(col.name().to_string());
1273 col_types.push(ColType::parse_mysql(col.type_info().name()));
1274 }
1275 let mut result_rows: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
1276 for row in &rows {
1277 let mut row_values: Vec<Value> = Vec::with_capacity(col_names.len());
1278 for (idx, _) in col_names.iter().enumerate() {
1279 row_values.push(row_to_value_with_coltype_mysql(row, idx, col_types[idx]));
1280 }
1281 result_rows.push(row_values);
1282 }
1283 Ok((col_names, result_rows))
1284 })
1285 }
1286
1287 fn query_stream<'a>(
1293 &'a mut self,
1294 sql: &'a str,
1295 ) -> Pin<Box<dyn futures::Stream<Item = Result<HashMap<String, Value>, DbError>> + Send + 'a>>
1296 {
1297 Box::pin(async_stream::try_stream! {
1298 let mut pool_conn = self
1299 .conn
1300 .take()
1301 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1302 let mut row_stream = sqlx::query(sqlx::AssertSqlSafe(sql)).fetch(&mut *pool_conn);
1303 let mut col_types: Vec<ColType> = Vec::new();
1304 let mut col_names: Vec<String> = Vec::new();
1305 let mut first_row = true;
1306 while let Some(row_result) = row_stream.next().await {
1307 let row = row_result.map_err(map_sqlx_error)?;
1308 if first_row {
1309 for col in row.columns() {
1310 col_names.push(col.name().to_string());
1311 col_types.push(ColType::parse_mysql(col.type_info().name()));
1312 }
1313 first_row = false;
1314 }
1315 let mut record = HashMap::with_capacity(col_names.len());
1316 for (i, name) in col_names.iter().enumerate() {
1317 let value = row_to_value_with_coltype_mysql(&row, i, col_types[i]);
1318 record.insert(name.clone(), value);
1319 }
1320 yield record;
1321 }
1322 drop(row_stream);
1325 self.conn = Some(pool_conn);
1326 })
1327 }
1328}
1329
1330impl Drop for SqlxMySqlConnection {
1331 fn drop(&mut self) {
1332 if let Some(conn) = self.conn.take() {
1333 drop(conn);
1334 }
1335 }
1336}
1337
1338pub async fn mysql_bulk_insert(
1345 conn: &mut SqlxMySqlConnection,
1346 table: &str,
1347 columns: &[&str],
1348 rows: &[Vec<Value>],
1349) -> Result<u64, DbError> {
1350 if rows.is_empty() {
1351 return Ok(0);
1352 }
1353 let col_list = columns.join(", ");
1354 let cols_per_row = columns.len();
1355 let row_placeholder = format!("({})", vec!["?"; cols_per_row].join(", "));
1357 let placeholders = vec![row_placeholder; rows.len()].join(", ");
1358 let sql = format!(
1359 "INSERT INTO {} ({}) VALUES {}",
1360 table, col_list, placeholders
1361 );
1362 let mut params: Vec<Value> = Vec::with_capacity(rows.len() * cols_per_row);
1364 for row in rows {
1365 for v in row {
1366 params.push(v.clone());
1367 }
1368 }
1369 conn.execute_with_params(&sql, ¶ms).await
1370}
1371
1372fn row_to_value_pg(row: &sqlx::postgres::PgRow, ordinal: usize) -> Value {
1378 use sqlx::TypeInfo;
1379 let type_name = row.columns()[ordinal].type_info().name();
1380 match type_name {
1381 "BOOL" => match row.try_get::<Option<bool>, usize>(ordinal) {
1382 Ok(v) => v.map(Value::Bool).unwrap_or(Value::Null),
1383 Err(_) => Value::Null,
1384 },
1385 "INT2" => match row.try_get::<Option<i16>, usize>(ordinal) {
1386 Ok(v) => v.map(Value::I16).unwrap_or(Value::Null),
1387 Err(_) => Value::Null,
1388 },
1389 "INT4" | "OID" => match row.try_get::<Option<i32>, usize>(ordinal) {
1390 Ok(v) => v.map(Value::I32).unwrap_or(Value::Null),
1391 Err(_) => Value::Null,
1392 },
1393 "INT8" => match row.try_get::<Option<i64>, usize>(ordinal) {
1394 Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
1395 Err(_) => Value::Null,
1396 },
1397 "FLOAT4" => match row.try_get::<Option<f32>, usize>(ordinal) {
1398 Ok(v) => v.map(Value::F32).unwrap_or(Value::Null),
1399 Err(_) => Value::Null,
1400 },
1401 "FLOAT8" => match row.try_get::<Option<f64>, usize>(ordinal) {
1402 Ok(v) => v.map(Value::F64).unwrap_or(Value::Null),
1403 Err(_) => Value::Null,
1404 },
1405 "TEXT" | "VARCHAR" | "CHAR" | "NAME" => match row.try_get::<Option<String>, usize>(ordinal)
1406 {
1407 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
1408 Err(_) => Value::Null,
1409 },
1410 "BYTEA" => match row.try_get::<Option<Vec<u8>>, usize>(ordinal) {
1411 Ok(v) => v.map(Value::Bytes).unwrap_or(Value::Null),
1412 Err(_) => Value::Null,
1413 },
1414 "NUMERIC" => match row.try_get::<Option<rust_decimal::Decimal>, usize>(ordinal) {
1415 Ok(Some(v)) => Value::Decimal(v.to_string()),
1416 Ok(None) => Value::Null,
1417 Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
1418 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
1419 Err(_) => Value::Null,
1420 },
1421 },
1422 "UUID" => match row.try_get::<Option<sqlx::types::Uuid>, usize>(ordinal) {
1424 Ok(v) => v
1425 .map(|uuid| Value::String(uuid.to_string()))
1426 .unwrap_or(Value::Null),
1427 Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
1428 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
1429 Err(_) => Value::Null,
1430 },
1431 },
1432 _ => {
1433 if let Ok(v) = row.try_get::<Option<i64>, usize>(ordinal) {
1435 return v.map(Value::I64).unwrap_or(Value::Null);
1436 }
1437 if let Ok(v) = row.try_get::<Option<f64>, usize>(ordinal) {
1438 return v.map(Value::F64).unwrap_or(Value::Null);
1439 }
1440 if let Ok(v) = row.try_get::<Option<bool>, usize>(ordinal) {
1441 return v.map(Value::Bool).unwrap_or(Value::Null);
1442 }
1443 if let Ok(v) = row.try_get::<Option<String>, usize>(ordinal) {
1444 return v.map(Value::String).unwrap_or(Value::Null);
1445 }
1446 Value::Null
1447 }
1448 }
1449}
1450
1451fn row_to_value_with_coltype_pg(
1459 row: &sqlx::postgres::PgRow,
1460 ordinal: usize,
1461 col_type: ColType,
1462) -> Value {
1463 match col_type {
1464 ColType::Bool => match row.try_get::<Option<bool>, usize>(ordinal) {
1465 Ok(v) => v.map(Value::Bool).unwrap_or(Value::Null),
1466 Err(_) => Value::Null,
1467 },
1468 ColType::I8 => match row.try_get::<Option<i8>, usize>(ordinal) {
1469 Ok(v) => v.map(Value::I8).unwrap_or(Value::Null),
1470 Err(_) => Value::Null,
1471 },
1472 ColType::I16 => match row.try_get::<Option<i16>, usize>(ordinal) {
1473 Ok(v) => v.map(Value::I16).unwrap_or(Value::Null),
1474 Err(_) => Value::Null,
1475 },
1476 ColType::I32 => match row.try_get::<Option<i32>, usize>(ordinal) {
1478 Ok(v) => v.map(Value::I32).unwrap_or(Value::Null),
1479 Err(_) => Value::Null,
1480 },
1481 ColType::I64 => match row.try_get::<Option<i64>, usize>(ordinal) {
1482 Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
1483 Err(_) => Value::Null,
1484 },
1485 ColType::U8 => match row.try_get::<Option<i16>, usize>(ordinal) {
1490 Ok(v) => v.map(Value::I16).unwrap_or(Value::Null),
1491 Err(_) => Value::Null,
1492 },
1493 ColType::U16 => match row.try_get::<Option<i32>, usize>(ordinal) {
1494 Ok(v) => v.map(Value::I32).unwrap_or(Value::Null),
1495 Err(_) => Value::Null,
1496 },
1497 ColType::U32 => match row.try_get::<Option<i64>, usize>(ordinal) {
1498 Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
1499 Err(_) => Value::Null,
1500 },
1501 ColType::U64 => match row.try_get::<Option<i64>, usize>(ordinal) {
1503 Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
1504 Err(_) => Value::Null,
1505 },
1506 ColType::F32 => match row.try_get::<Option<f32>, usize>(ordinal) {
1507 Ok(v) => v.map(Value::F32).unwrap_or(Value::Null),
1508 Err(_) => Value::Null,
1509 },
1510 ColType::F64 => match row.try_get::<Option<f64>, usize>(ordinal) {
1511 Ok(v) => v.map(Value::F64).unwrap_or(Value::Null),
1512 Err(_) => Value::Null,
1513 },
1514 ColType::Decimal => match row.try_get::<Option<rust_decimal::Decimal>, usize>(ordinal) {
1516 Ok(Some(v)) => Value::Decimal(v.to_string()),
1517 Ok(None) => Value::Null,
1518 Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
1519 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
1520 Err(_) => Value::Null,
1521 },
1522 },
1523 ColType::String => match row.try_get::<Option<String>, usize>(ordinal) {
1524 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
1525 Err(_) => Value::Null,
1526 },
1527 ColType::Bytes => match row.try_get::<Option<Vec<u8>>, usize>(ordinal) {
1528 Ok(v) => v.map(Value::Bytes).unwrap_or(Value::Null),
1529 Err(_) => Value::Null,
1530 },
1531 ColType::Date | ColType::DateTime | ColType::Time | ColType::Json => {
1533 match row.try_get::<Option<String>, usize>(ordinal) {
1534 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
1535 Err(_) => Value::Null,
1536 }
1537 }
1538 ColType::Uuid => match row.try_get::<Option<sqlx::types::Uuid>, usize>(ordinal) {
1540 Ok(v) => v
1541 .map(|uuid| Value::String(uuid.to_string()))
1542 .unwrap_or(Value::Null),
1543 Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
1544 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
1545 Err(_) => Value::Null,
1546 },
1547 },
1548 ColType::Unknown => {
1549 if let Ok(v) = row.try_get::<Option<i64>, usize>(ordinal) {
1551 return v.map(Value::I64).unwrap_or(Value::Null);
1552 }
1553 if let Ok(v) = row.try_get::<Option<f64>, usize>(ordinal) {
1554 return v.map(Value::F64).unwrap_or(Value::Null);
1555 }
1556 if let Ok(v) = row.try_get::<Option<bool>, usize>(ordinal) {
1557 return v.map(Value::Bool).unwrap_or(Value::Null);
1558 }
1559 if let Ok(v) = row.try_get::<Option<String>, usize>(ordinal) {
1560 return v.map(Value::String).unwrap_or(Value::Null);
1561 }
1562 Value::Null
1563 }
1564 _ => Value::Null,
1566 }
1567}
1568
1569pub struct PgPoolHandle {
1570 pool: sqlx::PgPool,
1571}
1572
1573impl PgPoolHandle {
1574 pub async fn connect(url: &str) -> Result<Self, DbError> {
1575 let pool = sqlx::pool::PoolOptions::<sqlx::Postgres>::new()
1576 .max_connections(10)
1577 .acquire_timeout(std::time::Duration::from_secs(30))
1578 .idle_timeout(Some(std::time::Duration::from_secs(600)))
1579 .max_lifetime(Some(std::time::Duration::from_secs(1800)))
1580 .connect(url)
1581 .await
1582 .map_err(map_sqlx_error)?;
1583 Ok(Self { pool })
1584 }
1585
1586 pub fn from_pool(pool: sqlx::PgPool) -> Self {
1587 Self { pool }
1588 }
1589
1590 pub fn pool(&self) -> &sqlx::PgPool {
1591 &self.pool
1592 }
1593}
1594
1595pub struct SqlxPgConnectionFactory {
1596 pool: Arc<PgPoolHandle>,
1597}
1598
1599impl SqlxPgConnectionFactory {
1600 pub fn new(pool: Arc<PgPoolHandle>) -> Self {
1601 Self { pool }
1602 }
1603}
1604
1605#[async_trait]
1606impl ConnectionFactory for SqlxPgConnectionFactory {
1607 async fn create(&self) -> Result<Box<dyn Connection>, DbError> {
1608 let conn = self.pool.pool.acquire().await.map_err(map_sqlx_error)?;
1609 Ok(Box::new(SqlxPgConnection {
1610 conn: Some(conn),
1611 connected: true,
1612 in_transaction: false,
1613 }))
1614 }
1615}
1616
1617pub struct SqlxPgConnection {
1618 conn: Option<sqlx::pool::PoolConnection<sqlx::Postgres>>,
1619 connected: bool,
1620 in_transaction: bool,
1621}
1622
1623impl Connection for SqlxPgConnection {
1624 fn execute<'a>(
1625 &'a mut self,
1626 sql: &'a str,
1627 ) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
1628 Box::pin(async move {
1629 let mut pool_conn = self
1630 .conn
1631 .take()
1632 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1633 let result = if needs_raw_sql(sql) {
1636 (&mut *pool_conn)
1637 .execute(sqlx::raw_sql(sqlx::AssertSqlSafe(sql)))
1638 .await
1639 } else {
1640 (&mut *pool_conn).execute(sqlx::AssertSqlSafe(sql)).await
1641 };
1642 self.conn = Some(pool_conn);
1643
1644 match result {
1645 Ok(r) => Ok(r.rows_affected()),
1646 Err(e) => {
1647 let db_err = map_sqlx_error(e);
1648 if matches!(db_err, DbError::ConnectionError(_) | DbError::IoError(_)) {
1649 self.connected = false;
1650 }
1651 Err(db_err)
1652 }
1653 }
1654 })
1655 }
1656
1657 fn query<'a>(
1658 &'a mut self,
1659 sql: &'a str,
1660 ) -> Pin<Box<dyn Future<Output = Result<Vec<HashMap<String, Value>>, DbError>> + Send + 'a>>
1661 {
1662 Box::pin(async move {
1663 let mut pool_conn = self
1664 .conn
1665 .take()
1666 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1667 let rows_result = (&mut *pool_conn).fetch_all(sqlx::AssertSqlSafe(sql)).await;
1670 self.conn = Some(pool_conn);
1671
1672 let rows = rows_result.map_err(map_sqlx_error)?;
1673 if rows.is_empty() {
1674 return Ok(Vec::new());
1675 }
1676 let col_types: Vec<ColType> = rows[0]
1679 .columns()
1680 .iter()
1681 .map(|col| ColType::parse_postgres(col.type_info().name()))
1682 .collect();
1683 let col_names: Vec<String> = rows[0]
1684 .columns()
1685 .iter()
1686 .map(|col| col.name().to_string())
1687 .collect();
1688 let mut result = Vec::with_capacity(rows.len());
1689 for row in &rows {
1690 let mut record = HashMap::with_capacity(col_types.len());
1691 for (i, name) in col_names.iter().enumerate() {
1692 let value = row_to_value_with_coltype_pg(row, i, col_types[i]);
1693 record.insert(name.clone(), value);
1694 }
1695 result.push(record);
1696 }
1697 Ok(result)
1698 })
1699 }
1700
1701 fn begin_transaction<'a>(
1702 &'a mut self,
1703 ) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
1704 Box::pin(async move {
1705 if self.in_transaction {
1706 return Err(DbError::Internal("transaction already started".to_string()));
1707 }
1708 self.execute("BEGIN").await?;
1709 self.in_transaction = true;
1710 Ok(())
1711 })
1712 }
1713
1714 fn commit<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
1715 Box::pin(async move {
1716 if self.in_transaction {
1717 self.execute("COMMIT").await?;
1718 self.in_transaction = false;
1719 }
1720 Ok(())
1721 })
1722 }
1723
1724 fn rollback<'a>(
1725 &'a mut self,
1726 ) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
1727 Box::pin(async move {
1728 if self.in_transaction {
1729 let result = self.execute("ROLLBACK").await;
1734 if result.is_ok() {
1735 self.in_transaction = false;
1736 }
1737 result.map(|_| ())
1738 } else {
1739 Ok(())
1740 }
1741 })
1742 }
1743
1744 fn is_connected(&self) -> bool {
1745 self.connected && self.conn.is_some()
1749 }
1750
1751 fn in_transaction(&self) -> bool {
1752 self.in_transaction
1753 }
1754
1755 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
1756 Box::pin(async move {
1757 match self.execute("SELECT 1").await {
1758 Ok(_) => true,
1759 Err(_) => {
1760 self.connected = false;
1761 false
1762 }
1763 }
1764 })
1765 }
1766
1767 fn close<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
1768 Box::pin(async move {
1769 if let Some(conn) = self.conn.take() {
1770 drop(conn);
1771 }
1772 self.connected = false;
1773 self.in_transaction = false;
1774 Ok(())
1775 })
1776 }
1777
1778 fn execute_with_params<'a>(
1793 &'a mut self,
1794 sql: &'a str,
1795 params: &'a [Value],
1796 ) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
1797 Box::pin(async move {
1798 if needs_raw_sql(sql) || params.is_empty() {
1799 return self.execute(sql).await;
1800 }
1801 let mut pool_conn = self
1802 .conn
1803 .take()
1804 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1805 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
1806 for v in params {
1807 q = match v {
1808 Value::Null => q.bind(None::<i64>),
1809 Value::Bool(b) => q.bind(*b),
1810 Value::I8(n) => q.bind(*n),
1811 Value::I16(n) => q.bind(*n),
1812 Value::I32(n) => q.bind(*n),
1813 Value::I64(n) => q.bind(*n),
1814 Value::U8(n) => q.bind(*n as i16),
1816 Value::U16(n) => q.bind(*n as i32),
1817 Value::U32(n) => q.bind(*n as i64),
1818 Value::U64(n) => q.bind(*n as i64),
1819 Value::F32(f) => q.bind(*f),
1820 Value::F64(f) => q.bind(*f),
1821 Value::String(s) => q.bind(s.as_str()),
1822 Value::Decimal(s) => q.bind(s.as_str()),
1823 Value::Bytes(b) => q.bind(b.as_slice()),
1824 other => q.bind(other.to_string()),
1825 };
1826 }
1827 let result = q.execute(&mut *pool_conn).await;
1828 self.conn = Some(pool_conn);
1829 match result {
1830 Ok(r) => Ok(r.rows_affected()),
1831 Err(e) => {
1832 let db_err = map_sqlx_error(e);
1833 if matches!(db_err, DbError::ConnectionError(_) | DbError::IoError(_)) {
1834 self.connected = false;
1835 }
1836 Err(db_err)
1837 }
1838 }
1839 })
1840 }
1841
1842 fn query_with_params<'a>(
1846 &'a mut self,
1847 sql: &'a str,
1848 params: &'a [Value],
1849 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, DbError>> + Send + 'a>> {
1850 Box::pin(async move {
1851 if params.is_empty() {
1852 return self.query(sql).await;
1853 }
1854 let mut pool_conn = self
1855 .conn
1856 .take()
1857 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1858 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
1859 for v in params {
1860 q = match v {
1861 Value::Null => q.bind(None::<i64>),
1862 Value::Bool(b) => q.bind(*b),
1863 Value::I8(n) => q.bind(*n),
1864 Value::I16(n) => q.bind(*n),
1865 Value::I32(n) => q.bind(*n),
1866 Value::I64(n) => q.bind(*n),
1867 Value::U8(n) => q.bind(*n as i16),
1868 Value::U16(n) => q.bind(*n as i32),
1869 Value::U32(n) => q.bind(*n as i64),
1870 Value::U64(n) => q.bind(*n as i64),
1871 Value::F32(f) => q.bind(*f),
1872 Value::F64(f) => q.bind(*f),
1873 Value::String(s) => q.bind(s.as_str()),
1874 Value::Decimal(s) => q.bind(s.as_str()),
1875 Value::Bytes(b) => q.bind(b.as_slice()),
1876 other => q.bind(other.to_string()),
1877 };
1878 }
1879 let rows_result = q.fetch_all(&mut *pool_conn).await;
1880 self.conn = Some(pool_conn);
1881 let rows = rows_result.map_err(map_sqlx_error)?;
1882 if rows.is_empty() {
1883 return Ok(Vec::new());
1884 }
1885 let col_types: Vec<ColType> = rows[0]
1887 .columns()
1888 .iter()
1889 .map(|col| ColType::parse_postgres(col.type_info().name()))
1890 .collect();
1891 let col_names: Vec<String> = rows[0]
1892 .columns()
1893 .iter()
1894 .map(|col| col.name().to_string())
1895 .collect();
1896 let mut result = Vec::with_capacity(rows.len());
1897 for row in &rows {
1898 let mut record = HashMap::with_capacity(col_types.len());
1899 for (i, name) in col_names.iter().enumerate() {
1900 let value = row_to_value_with_coltype_pg(row, i, col_types[i]);
1901 record.insert(name.clone(), value);
1902 }
1903 result.push(record);
1904 }
1905 Ok(result)
1906 })
1907 }
1908
1909 fn query_values<'a>(
1913 &'a mut self,
1914 sql: &'a str,
1915 ) -> Pin<Box<dyn Future<Output = Result<QueryValues, DbError>> + Send + 'a>> {
1916 Box::pin(async move {
1917 let mut pool_conn = self
1918 .conn
1919 .take()
1920 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1921 let rows_result = (&mut *pool_conn).fetch_all(sqlx::AssertSqlSafe(sql)).await;
1922 self.conn = Some(pool_conn);
1923 let rows = rows_result.map_err(map_sqlx_error)?;
1924 if rows.is_empty() {
1925 return Ok((Vec::new(), Vec::new()));
1926 }
1927 let cols = rows[0].columns();
1928 let mut col_names: Vec<String> = Vec::with_capacity(cols.len());
1929 let mut col_types: Vec<ColType> = Vec::with_capacity(cols.len());
1930 for col in cols {
1931 col_names.push(col.name().to_string());
1932 col_types.push(ColType::parse_postgres(col.type_info().name()));
1933 }
1934 let mut result_rows: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
1935 for row in &rows {
1936 let mut row_values: Vec<Value> = Vec::with_capacity(col_names.len());
1937 for (idx, _) in col_names.iter().enumerate() {
1938 row_values.push(row_to_value_with_coltype_pg(row, idx, col_types[idx]));
1939 }
1940 result_rows.push(row_values);
1941 }
1942 Ok((col_names, result_rows))
1943 })
1944 }
1945
1946 fn query_values_with_params<'a>(
1950 &'a mut self,
1951 sql: &'a str,
1952 params: &'a [Value],
1953 ) -> Pin<Box<dyn Future<Output = Result<QueryValues, DbError>> + Send + 'a>> {
1954 Box::pin(async move {
1955 if params.is_empty() {
1956 return self.query_values(sql).await;
1957 }
1958 let mut pool_conn = self
1959 .conn
1960 .take()
1961 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1962 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
1963 for v in params {
1964 q = match v {
1965 Value::Null => q.bind(None::<i64>),
1966 Value::Bool(b) => q.bind(*b),
1967 Value::I8(n) => q.bind(*n),
1968 Value::I16(n) => q.bind(*n),
1969 Value::I32(n) => q.bind(*n),
1970 Value::I64(n) => q.bind(*n),
1971 Value::U8(n) => q.bind(*n as i16),
1972 Value::U16(n) => q.bind(*n as i32),
1973 Value::U32(n) => q.bind(*n as i64),
1974 Value::U64(n) => q.bind(*n as i64),
1975 Value::F32(f) => q.bind(*f),
1976 Value::F64(f) => q.bind(*f),
1977 Value::String(s) => q.bind(s.as_str()),
1978 Value::Decimal(s) => q.bind(s.as_str()),
1979 Value::Bytes(b) => q.bind(b.as_slice()),
1980 other => q.bind(other.to_string()),
1981 };
1982 }
1983 let rows_result = q.fetch_all(&mut *pool_conn).await;
1984 self.conn = Some(pool_conn);
1985 let rows = rows_result.map_err(map_sqlx_error)?;
1986 if rows.is_empty() {
1987 return Ok((Vec::new(), Vec::new()));
1988 }
1989 let cols = rows[0].columns();
1990 let mut col_names: Vec<String> = Vec::with_capacity(cols.len());
1991 for col in cols {
1992 col_names.push(col.name().to_string());
1993 }
1994 let mut result_rows: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
1995 for row in rows {
1996 let mut row_values: Vec<Value> = Vec::with_capacity(col_names.len());
1997 for (idx, _) in col_names.iter().enumerate() {
1998 let ordinal = row.columns()[idx].ordinal();
1999 row_values.push(row_to_value_pg(&row, ordinal));
2000 }
2001 result_rows.push(row_values);
2002 }
2003 Ok((col_names, result_rows))
2004 })
2005 }
2006
2007 fn query_stream<'a>(
2017 &'a mut self,
2018 sql: &'a str,
2019 ) -> Pin<Box<dyn futures::Stream<Item = Result<HashMap<String, Value>, DbError>> + Send + 'a>>
2020 {
2021 Box::pin(async_stream::try_stream! {
2022 let mut pool_conn = self
2023 .conn
2024 .take()
2025 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
2026 let mut row_stream = sqlx::query(sqlx::AssertSqlSafe(sql)).fetch(&mut *pool_conn);
2027 while let Some(row_result) = row_stream.next().await {
2030 let row = row_result.map_err(map_sqlx_error)?;
2031 let cols = row.columns();
2032 let mut record = HashMap::with_capacity(cols.len());
2033 for (i, col) in cols.iter().enumerate() {
2034 let name = col.name().to_string();
2035 let value = row_to_value_pg(&row, i);
2037 record.insert(name, value);
2038 }
2039 yield record;
2040 }
2041 drop(row_stream);
2044 self.conn = Some(pool_conn);
2045 })
2046 }
2047}
2048
2049impl Drop for SqlxPgConnection {
2050 fn drop(&mut self) {
2051 if let Some(conn) = self.conn.take() {
2052 drop(conn);
2053 }
2054 }
2055}
2056
2057#[async_trait]
2078pub trait PgExtensions: Send + Sync {
2079 async fn listen(&mut self, channel: &str) -> Result<(), DbError>;
2090
2091 async fn notify(&mut self, channel: &str, payload: &str) -> Result<(), DbError>;
2102
2103 async fn copy_from_stdin(&mut self, sql: &str, data: &[u8]) -> Result<u64, DbError>;
2118}
2119
2120fn validate_pg_channel_name(channel: &str) -> Result<(), DbError> {
2125 if channel.is_empty() {
2126 return Err(DbError::Internal(
2127 "PG channel name must not be empty".to_string(),
2128 ));
2129 }
2130 if !channel
2131 .chars()
2132 .all(|c| c.is_ascii_alphanumeric() || c == '_')
2133 {
2134 return Err(DbError::Internal(format!(
2135 "invalid PG channel name: {} (only alphanumeric and underscore allowed)",
2136 channel
2137 )));
2138 }
2139 Ok(())
2140}
2141
2142#[async_trait]
2143impl PgExtensions for SqlxPgConnection {
2144 async fn listen(&mut self, channel: &str) -> Result<(), DbError> {
2145 validate_pg_channel_name(channel)?;
2146 self.execute(&format!("LISTEN {}", channel)).await?;
2148 Ok(())
2149 }
2150
2151 async fn notify(&mut self, channel: &str, payload: &str) -> Result<(), DbError> {
2152 validate_pg_channel_name(channel)?;
2153 let escaped_payload = payload.replace('\'', "''");
2155 self.execute(&format!("NOTIFY {}, '{}'", channel, escaped_payload))
2156 .await?;
2157 Ok(())
2158 }
2159
2160 async fn copy_from_stdin(&mut self, sql: &str, data: &[u8]) -> Result<u64, DbError> {
2161 let mut pool_conn = self
2162 .conn
2163 .take()
2164 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
2165 let mut copy = (*pool_conn)
2169 .copy_in_raw(sql)
2170 .await
2171 .map_err(map_sqlx_error)?;
2172 copy.send(data).await.map_err(map_sqlx_error)?;
2174 let result = copy.finish().await.map_err(map_sqlx_error)?;
2176 self.conn = Some(pool_conn);
2177 Ok(result)
2178 }
2179}
2180
2181pub async fn pg_bulk_insert(
2189 conn: &mut SqlxPgConnection,
2190 table: &str,
2191 columns: &[&str],
2192 rows: &[Vec<Value>],
2193) -> Result<u64, DbError> {
2194 if rows.is_empty() {
2195 return Ok(0);
2196 }
2197 let col_list = columns.join(", ");
2198 let cols_per_row = columns.len();
2199 let placeholders: Vec<String> = rows
2201 .iter()
2202 .enumerate()
2203 .map(|(row_idx, _)| {
2204 let base = row_idx * cols_per_row;
2205 let ph: Vec<String> = (0..cols_per_row)
2206 .map(|i| format!("${}", base + i + 1))
2207 .collect();
2208 format!("({})", ph.join(", "))
2209 })
2210 .collect();
2211 let sql = format!(
2212 "INSERT INTO {} ({}) VALUES {}",
2213 table,
2214 col_list,
2215 placeholders.join(", ")
2216 );
2217 let mut params: Vec<Value> = Vec::with_capacity(rows.len() * cols_per_row);
2219 for row in rows {
2220 for v in row {
2221 params.push(v.clone());
2222 }
2223 }
2224 conn.execute_with_params(&sql, ¶ms).await
2225}
2226#[cfg(test)]
2227mod tests {
2228 use super::*;
2229
2230 #[test]
2231 fn test_needs_raw_sql_begin() {
2232 assert!(needs_raw_sql("BEGIN"));
2233 assert!(needs_raw_sql("begin"));
2234 assert!(needs_raw_sql(" BEGIN"));
2235 assert!(needs_raw_sql("\tbegin"));
2236 }
2237
2238 #[test]
2239 fn test_needs_raw_sql_commit() {
2240 assert!(needs_raw_sql("COMMIT"));
2241 assert!(needs_raw_sql("commit"));
2242 assert!(needs_raw_sql(" COMMIT WORK"));
2243 }
2244
2245 #[test]
2246 fn test_needs_raw_sql_rollback() {
2247 assert!(needs_raw_sql("ROLLBACK"));
2248 assert!(needs_raw_sql("rollback"));
2249 assert!(needs_raw_sql(" ROLLBACK TO SAVEPOINT sp1"));
2250 }
2251
2252 #[test]
2253 fn test_needs_raw_sql_savepoint() {
2254 assert!(needs_raw_sql("SAVEPOINT sp1"));
2255 assert!(needs_raw_sql("savepoint s1"));
2256 }
2257
2258 #[test]
2259 fn test_needs_raw_sql_release() {
2260 assert!(needs_raw_sql("RELEASE sp1"));
2261 assert!(needs_raw_sql("release savepoint s1"));
2262 }
2263
2264 #[test]
2265 fn test_needs_raw_sql_set() {
2266 assert!(needs_raw_sql("SET autocommit = 0"));
2267 assert!(needs_raw_sql("set names utf8"));
2268 }
2269
2270 #[test]
2271 fn test_needs_raw_sql_use() {
2272 assert!(needs_raw_sql("USE mydb"));
2273 assert!(needs_raw_sql("use mydb"));
2274 }
2275
2276 #[test]
2277 fn test_needs_raw_sql_start_transaction() {
2278 assert!(needs_raw_sql("START TRANSACTION"));
2279 assert!(needs_raw_sql("start transaction"));
2280 assert!(needs_raw_sql(" START TRANSACTION READ ONLY"));
2281 }
2282
2283 #[test]
2284 fn test_needs_raw_sql_normal_queries() {
2285 assert!(!needs_raw_sql("SELECT * FROM users"));
2286 assert!(!needs_raw_sql("INSERT INTO t VALUES (1)"));
2287 assert!(!needs_raw_sql("UPDATE t SET x = 1"));
2288 assert!(!needs_raw_sql("DELETE FROM t"));
2289 assert!(!needs_raw_sql("CREATE TABLE t (id INT)"));
2290 assert!(!needs_raw_sql("DROP TABLE t"));
2291 assert!(!needs_raw_sql("ALTER TABLE t ADD COLUMN x INT"));
2292 }
2293
2294 #[test]
2295 fn test_needs_raw_sql_edge_cases() {
2296 assert!(!needs_raw_sql(""));
2297 assert!(!needs_raw_sql(" "));
2298 assert!(!needs_raw_sql("SELECT"));
2299 assert!(needs_raw_sql("BEGINNING"));
2300 assert!(needs_raw_sql("COMMITTEE"));
2301 assert!(!needs_raw_sql("INSERT"));
2302 assert!(!needs_raw_sql("DELETE"));
2303 }
2304
2305 #[test]
2308 fn test_validate_pg_channel_name_valid() {
2309 assert!(validate_pg_channel_name("my_channel").is_ok());
2310 assert!(validate_pg_channel_name("abc123").is_ok());
2311 assert!(validate_pg_channel_name("_test").is_ok());
2312 assert!(validate_pg_channel_name("a_b_c").is_ok());
2313 }
2314
2315 #[test]
2316 fn test_validate_pg_channel_name_empty() {
2317 let err = validate_pg_channel_name("").unwrap_err();
2318 assert!(matches!(err, DbError::Internal(ref s) if s.contains("empty")));
2319 }
2320
2321 #[test]
2322 fn test_validate_pg_channel_name_invalid_chars() {
2323 let err = validate_pg_channel_name("my-channel").unwrap_err();
2324 assert!(matches!(err, DbError::Internal(ref s) if s.contains("invalid")));
2325 assert!(validate_pg_channel_name("my.channel").is_err());
2326 assert!(validate_pg_channel_name("my channel").is_err());
2327 assert!(validate_pg_channel_name("my;channel").is_err());
2328 assert!(validate_pg_channel_name("my'channel").is_err());
2329 }
2330
2331 fn closed_sqlite() -> SqlxSqliteConnection {
2334 SqlxSqliteConnection {
2335 conn: None,
2336 connected: false,
2337 in_transaction: false,
2338 }
2339 }
2340
2341 #[tokio::test]
2342 async fn test_sqlite_execute_closed() {
2343 let mut c = closed_sqlite();
2344 let result = c.execute("SELECT 1").await;
2345 assert!(result.is_err());
2346 assert!(result.unwrap_err().to_string().contains("closed"));
2347 }
2348
2349 #[tokio::test]
2350 async fn test_sqlite_query_closed() {
2351 let mut c = closed_sqlite();
2352 let result = c.query("SELECT 1").await;
2353 assert!(result.is_err());
2354 }
2355
2356 #[tokio::test]
2357 async fn test_sqlite_begin_tx_already_started() {
2358 let mut c = SqlxSqliteConnection {
2359 conn: None,
2360 connected: false,
2361 in_transaction: true,
2362 };
2363 let result = c.begin_transaction().await;
2364 assert!(result.is_err());
2365 assert!(result.unwrap_err().to_string().contains("already started"));
2366 }
2367
2368 #[tokio::test]
2369 async fn test_sqlite_begin_tx_closed() {
2370 let mut c = closed_sqlite();
2371 let result = c.begin_transaction().await;
2372 assert!(result.is_err());
2373 }
2374
2375 #[tokio::test]
2376 async fn test_sqlite_commit_no_tx() {
2377 let mut c = closed_sqlite();
2378 let result = c.commit().await;
2379 assert!(result.is_ok());
2380 }
2381
2382 #[tokio::test]
2383 async fn test_sqlite_rollback_no_tx() {
2384 let mut c = closed_sqlite();
2385 let result = c.rollback().await;
2386 assert!(result.is_ok());
2387 }
2388
2389 #[test]
2390 fn test_sqlite_is_connected_false_when_conn_none() {
2391 let c = closed_sqlite();
2392 assert!(!c.is_connected());
2393 }
2394
2395 #[test]
2396 fn test_sqlite_is_connected_false_when_connected_false() {
2397 let c = SqlxSqliteConnection {
2398 conn: None,
2399 connected: true,
2400 in_transaction: false,
2401 };
2402 assert!(!c.is_connected());
2403 }
2404
2405 #[test]
2406 fn test_sqlite_in_transaction_flag() {
2407 let c = SqlxSqliteConnection {
2408 conn: None,
2409 connected: false,
2410 in_transaction: true,
2411 };
2412 assert!(c.in_transaction());
2413 let c2 = closed_sqlite();
2414 assert!(!c2.in_transaction());
2415 }
2416
2417 #[tokio::test]
2418 async fn test_sqlite_ping_closed() {
2419 let mut c = closed_sqlite();
2420 let result = c.ping().await;
2421 assert!(!result);
2422 assert!(!c.is_connected());
2423 }
2424
2425 #[tokio::test]
2426 async fn test_sqlite_close_closed() {
2427 let mut c = closed_sqlite();
2428 let result = c.close().await;
2429 assert!(result.is_ok());
2430 assert!(!c.is_connected());
2431 assert!(!c.in_transaction());
2432 }
2433
2434 #[tokio::test]
2435 async fn test_sqlite_execute_with_params_closed() {
2436 let mut c = closed_sqlite();
2437 let result = c
2438 .execute_with_params("INSERT INTO t VALUES (?)", &[Value::I64(1)])
2439 .await;
2440 assert!(result.is_err());
2441 }
2442
2443 #[tokio::test]
2444 async fn test_sqlite_execute_with_params_empty_delegates() {
2445 let mut c = closed_sqlite();
2446 let result = c.execute_with_params("INSERT INTO t VALUES (1)", &[]).await;
2447 assert!(result.is_err());
2448 }
2449
2450 #[tokio::test]
2451 async fn test_sqlite_query_with_params_empty_delegates() {
2452 let mut c = closed_sqlite();
2453 let result = c.query_with_params("SELECT * FROM t", &[]).await;
2454 assert!(result.is_err());
2455 }
2456
2457 #[tokio::test]
2458 async fn test_sqlite_query_with_params_closed() {
2459 let mut c = closed_sqlite();
2460 let result = c
2461 .query_with_params("SELECT * FROM t WHERE id = ?", &[Value::I64(1)])
2462 .await;
2463 assert!(result.is_err());
2464 }
2465
2466 fn closed_mysql() -> SqlxMySqlConnection {
2469 SqlxMySqlConnection {
2470 conn: None,
2471 connected: false,
2472 in_transaction: false,
2473 }
2474 }
2475
2476 #[tokio::test]
2477 async fn test_mysql_execute_closed() {
2478 let mut c = closed_mysql();
2479 assert!(c.execute("SELECT 1").await.is_err());
2480 }
2481
2482 #[tokio::test]
2483 async fn test_mysql_query_closed() {
2484 let mut c = closed_mysql();
2485 assert!(c.query("SELECT 1").await.is_err());
2486 }
2487
2488 #[test]
2489 fn test_mysql_is_connected_false() {
2490 let c = closed_mysql();
2491 assert!(!c.is_connected());
2492 }
2493
2494 #[tokio::test]
2495 async fn test_mysql_close_closed() {
2496 let mut c = closed_mysql();
2497 assert!(c.close().await.is_ok());
2498 assert!(!c.is_connected());
2499 }
2500
2501 #[tokio::test]
2502 async fn test_mysql_commit_no_tx() {
2503 let mut c = closed_mysql();
2504 assert!(c.commit().await.is_ok());
2505 }
2506
2507 #[tokio::test]
2508 async fn test_mysql_rollback_no_tx() {
2509 let mut c = closed_mysql();
2510 assert!(c.rollback().await.is_ok());
2511 }
2512
2513 #[tokio::test]
2514 async fn test_mysql_begin_tx_already_started() {
2515 let mut c = SqlxMySqlConnection {
2516 conn: None,
2517 connected: false,
2518 in_transaction: true,
2519 };
2520 let result = c.begin_transaction().await;
2521 assert!(result.is_err());
2522 }
2523
2524 #[tokio::test]
2525 async fn test_mysql_ping_closed() {
2526 let mut c = closed_mysql();
2527 assert!(!c.ping().await);
2528 }
2529
2530 fn closed_pg() -> SqlxPgConnection {
2533 SqlxPgConnection {
2534 conn: None,
2535 connected: false,
2536 in_transaction: false,
2537 }
2538 }
2539
2540 #[tokio::test]
2541 async fn test_pg_execute_closed() {
2542 let mut c = closed_pg();
2543 assert!(c.execute("SELECT 1").await.is_err());
2544 }
2545
2546 #[tokio::test]
2547 async fn test_pg_query_closed() {
2548 let mut c = closed_pg();
2549 assert!(c.query("SELECT 1").await.is_err());
2550 }
2551
2552 #[test]
2553 fn test_pg_is_connected_false() {
2554 let c = closed_pg();
2555 assert!(!c.is_connected());
2556 }
2557
2558 #[tokio::test]
2559 async fn test_pg_close_closed() {
2560 let mut c = closed_pg();
2561 assert!(c.close().await.is_ok());
2562 assert!(!c.is_connected());
2563 }
2564
2565 #[tokio::test]
2566 async fn test_pg_commit_no_tx() {
2567 let mut c = closed_pg();
2568 assert!(c.commit().await.is_ok());
2569 }
2570
2571 #[tokio::test]
2572 async fn test_pg_rollback_no_tx() {
2573 let mut c = closed_pg();
2574 assert!(c.rollback().await.is_ok());
2575 }
2576
2577 #[tokio::test]
2578 async fn test_pg_begin_tx_already_started() {
2579 let mut c = SqlxPgConnection {
2580 conn: None,
2581 connected: false,
2582 in_transaction: true,
2583 };
2584 assert!(c.begin_transaction().await.is_err());
2585 }
2586
2587 #[tokio::test]
2588 async fn test_pg_ping_closed() {
2589 let mut c = closed_pg();
2590 assert!(!c.ping().await);
2591 }
2592
2593 #[tokio::test]
2596 async fn test_mysql_execute_with_params_closed() {
2597 let mut c = closed_mysql();
2598 assert!(c
2599 .execute_with_params("INSERT INTO t VALUES (?)", &[Value::I64(1)])
2600 .await
2601 .is_err());
2602 }
2603
2604 #[tokio::test]
2605 async fn test_mysql_execute_with_params_empty_delegates() {
2606 let mut c = closed_mysql();
2607 assert!(c
2608 .execute_with_params("INSERT INTO t VALUES (1)", &[])
2609 .await
2610 .is_err());
2611 }
2612
2613 #[tokio::test]
2614 async fn test_mysql_query_with_params_closed() {
2615 let mut c = closed_mysql();
2616 assert!(c
2617 .query_with_params("SELECT * FROM t WHERE id = ?", &[Value::I64(1)])
2618 .await
2619 .is_err());
2620 }
2621
2622 #[tokio::test]
2623 async fn test_mysql_query_with_params_empty_delegates() {
2624 let mut c = closed_mysql();
2625 assert!(c.query_with_params("SELECT * FROM t", &[]).await.is_err());
2626 }
2627
2628 #[tokio::test]
2629 async fn test_mysql_query_values_closed() {
2630 let mut c = closed_mysql();
2631 assert!(c.query_values("SELECT * FROM t").await.is_err());
2632 }
2633
2634 #[tokio::test]
2635 async fn test_mysql_query_values_with_params_closed() {
2636 let mut c = closed_mysql();
2637 assert!(c
2638 .query_values_with_params("SELECT * FROM t WHERE id = ?", &[Value::I64(1)])
2639 .await
2640 .is_err());
2641 }
2642
2643 #[tokio::test]
2644 async fn test_mysql_query_values_with_params_empty_delegates() {
2645 let mut c = closed_mysql();
2646 assert!(c
2647 .query_values_with_params("SELECT * FROM t", &[])
2648 .await
2649 .is_err());
2650 }
2651
2652 #[tokio::test]
2655 async fn test_pg_execute_with_params_closed() {
2656 let mut c = closed_pg();
2657 assert!(c
2658 .execute_with_params("INSERT INTO t VALUES ($1)", &[Value::I64(1)])
2659 .await
2660 .is_err());
2661 }
2662
2663 #[tokio::test]
2664 async fn test_pg_execute_with_params_empty_delegates() {
2665 let mut c = closed_pg();
2666 assert!(c
2667 .execute_with_params("INSERT INTO t VALUES (1)", &[])
2668 .await
2669 .is_err());
2670 }
2671
2672 #[tokio::test]
2673 async fn test_pg_query_with_params_closed() {
2674 let mut c = closed_pg();
2675 assert!(c
2676 .query_with_params("SELECT * FROM t WHERE id = $1", &[Value::I64(1)])
2677 .await
2678 .is_err());
2679 }
2680
2681 #[tokio::test]
2682 async fn test_pg_query_with_params_empty_delegates() {
2683 let mut c = closed_pg();
2684 assert!(c.query_with_params("SELECT * FROM t", &[]).await.is_err());
2685 }
2686
2687 #[tokio::test]
2688 async fn test_pg_query_values_closed() {
2689 let mut c = closed_pg();
2690 assert!(c.query_values("SELECT * FROM t").await.is_err());
2691 }
2692
2693 #[tokio::test]
2694 async fn test_pg_query_values_with_params_closed() {
2695 let mut c = closed_pg();
2696 assert!(c
2697 .query_values_with_params("SELECT * FROM t WHERE id = $1", &[Value::I64(1)])
2698 .await
2699 .is_err());
2700 }
2701
2702 #[tokio::test]
2703 async fn test_pg_query_values_with_params_empty_delegates() {
2704 let mut c = closed_pg();
2705 assert!(c
2706 .query_values_with_params("SELECT * FROM t", &[])
2707 .await
2708 .is_err());
2709 }
2710
2711 #[tokio::test]
2714 async fn test_sqlite_backup_closed_conn() {
2715 let mut c = closed_sqlite();
2716 let result = sqlite_backup(&mut c, "/tmp/backup.db").await;
2717 assert!(result.is_err());
2718 }
2719
2720 #[tokio::test]
2723 async fn test_sqlite_query_values_closed() {
2724 let mut c = closed_sqlite();
2725 assert!(c.query_values("SELECT * FROM t").await.is_err());
2726 }
2727
2728 #[tokio::test]
2729 async fn test_sqlite_query_values_with_params_closed() {
2730 let mut c = closed_sqlite();
2731 assert!(c
2732 .query_values_with_params("SELECT * FROM t WHERE id = ?", &[Value::I64(1)])
2733 .await
2734 .is_err());
2735 }
2736
2737 #[tokio::test]
2738 async fn test_sqlite_query_values_with_params_empty_delegates() {
2739 let mut c = closed_sqlite();
2740 assert!(c
2741 .query_values_with_params("SELECT * FROM t", &[])
2742 .await
2743 .is_err());
2744 }
2745
2746 #[tokio::test]
2749 async fn test_mysql_bulk_insert_empty_rows() {
2750 let mut c = closed_mysql();
2751 let result = mysql_bulk_insert(&mut c, "t", &["a", "b"], &[]).await;
2752 assert_eq!(result.unwrap(), 0);
2753 }
2754
2755 #[tokio::test]
2756 async fn test_mysql_bulk_insert_closed_conn() {
2757 let mut c = closed_mysql();
2758 let rows = vec![vec![Value::I64(1), Value::String("x".to_string())]];
2759 let result = mysql_bulk_insert(&mut c, "t", &["a", "b"], &rows).await;
2760 assert!(result.is_err());
2761 }
2762
2763 #[tokio::test]
2764 async fn test_pg_bulk_insert_empty_rows() {
2765 let mut c = closed_pg();
2766 let result = pg_bulk_insert(&mut c, "t", &["a", "b"], &[]).await;
2767 assert_eq!(result.unwrap(), 0);
2768 }
2769
2770 #[tokio::test]
2771 async fn test_pg_bulk_insert_closed_conn() {
2772 let mut c = closed_pg();
2773 let rows = vec![vec![Value::I64(1), Value::String("x".to_string())]];
2774 let result = pg_bulk_insert(&mut c, "t", &["a", "b"], &rows).await;
2775 assert!(result.is_err());
2776 }
2777}