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<bool>, usize>(ordinal) {
129 return v.map(Value::Bool).unwrap_or(Value::Null);
130 }
131 if let Ok(v) = row.try_get::<Option<i64>, usize>(ordinal) {
132 return v.map(Value::I64).unwrap_or(Value::Null);
133 }
134 if let Ok(v) = row.try_get::<Option<f64>, usize>(ordinal) {
135 return v.map(Value::F64).unwrap_or(Value::Null);
136 }
137 if let Ok(v) = row.try_get::<Option<String>, usize>(ordinal) {
138 return v.map(Value::String).unwrap_or(Value::Null);
139 }
140 Value::Null
141 }
142 _ => Value::Null,
144 }
145}
146
147pub struct SqlitePoolHandle {
148 pool: sqlx::SqlitePool,
149}
150
151impl SqlitePoolHandle {
152 pub async fn connect(url: &str) -> Result<Self, DbError> {
153 let opts = sqlx::sqlite::SqliteConnectOptions::from_str(url)
159 .map_err(map_sqlx_error)?
160 .journal_mode(sqlx::sqlite::SqliteJournalMode::Wal)
161 .synchronous(sqlx::sqlite::SqliteSynchronous::Normal)
162 .busy_timeout(std::time::Duration::from_secs(5))
163 .pragma("mmap_size", "268435456");
164 let pool = sqlx::sqlite::SqlitePoolOptions::new()
165 .max_connections(10)
166 .acquire_timeout(std::time::Duration::from_secs(30))
167 .idle_timeout(Some(std::time::Duration::from_secs(600)))
168 .max_lifetime(Some(std::time::Duration::from_secs(1800)))
169 .connect_with(opts)
170 .await
171 .map_err(map_sqlx_error)?;
172 Ok(Self { pool })
173 }
174
175 pub fn from_pool(pool: sqlx::SqlitePool) -> Self {
176 Self { pool }
177 }
178
179 pub fn pool(&self) -> &sqlx::SqlitePool {
180 &self.pool
181 }
182}
183
184pub struct SqlxSqliteConnectionFactory {
185 pool: Arc<SqlitePoolHandle>,
186}
187
188impl SqlxSqliteConnectionFactory {
189 pub fn new(pool: Arc<SqlitePoolHandle>) -> Self {
190 Self { pool }
191 }
192}
193
194#[async_trait]
195impl ConnectionFactory for SqlxSqliteConnectionFactory {
196 async fn create(&self) -> Result<Box<dyn Connection>, DbError> {
197 let conn = self.pool.pool.acquire().await.map_err(map_sqlx_error)?;
198 Ok(Box::new(SqlxSqliteConnection {
199 conn: Some(conn),
200 connected: true,
201 in_transaction: false,
202 }))
203 }
204}
205
206pub struct SqlxSqliteConnection {
207 conn: Option<sqlx::pool::PoolConnection<sqlx::Sqlite>>,
208 connected: bool,
209 in_transaction: bool,
210}
211
212impl Connection for SqlxSqliteConnection {
213 fn execute<'a>(
214 &'a mut self,
215 sql: &'a str,
216 ) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
217 Box::pin(async move {
218 let mut pool_conn = self
219 .conn
220 .take()
221 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
222 let result = if needs_raw_sql(sql) {
225 (&mut *pool_conn)
226 .execute(sqlx::raw_sql(sqlx::AssertSqlSafe(sql)))
227 .await
228 } else {
229 (&mut *pool_conn).execute(sqlx::AssertSqlSafe(sql)).await
230 };
231 self.conn = Some(pool_conn);
232
233 match result {
234 Ok(r) => Ok(r.rows_affected()),
235 Err(e) => {
236 let db_err = map_sqlx_error(e);
237 if matches!(db_err, DbError::ConnectionError(_) | DbError::IoError(_)) {
238 self.connected = false;
239 }
240 Err(db_err)
241 }
242 }
243 })
244 }
245
246 fn query<'a>(
247 &'a mut self,
248 sql: &'a str,
249 ) -> Pin<Box<dyn Future<Output = Result<Vec<HashMap<String, Value>>, DbError>> + Send + 'a>>
250 {
251 Box::pin(async move {
252 let mut pool_conn = self
253 .conn
254 .take()
255 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
256 let rows_result = (&mut *pool_conn).fetch_all(sqlx::AssertSqlSafe(sql)).await;
259 self.conn = Some(pool_conn);
260
261 let rows = rows_result.map_err(map_sqlx_error)?;
262 if rows.is_empty() {
263 return Ok(Vec::new());
264 }
265 let col_types: Vec<ColType> = rows[0]
268 .columns()
269 .iter()
270 .map(|col| ColType::parse_sqlite(col.type_info().name()))
271 .collect();
272 let col_names: Vec<String> = rows[0]
273 .columns()
274 .iter()
275 .map(|col| col.name().to_string())
276 .collect();
277 let mut result = Vec::with_capacity(rows.len());
278 for row in &rows {
279 let mut record = HashMap::with_capacity(col_types.len());
280 for (i, name) in col_names.iter().enumerate() {
281 let value = row_to_value_with_coltype_sqlite(row, i, col_types[i]);
282 record.insert(name.clone(), value);
283 }
284 result.push(record);
285 }
286 Ok(result)
287 })
288 }
289
290 fn begin_transaction<'a>(
291 &'a mut self,
292 ) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
293 Box::pin(async move {
294 if self.in_transaction {
295 return Err(DbError::Internal("transaction already started".to_string()));
296 }
297 self.execute("BEGIN").await?;
298 self.in_transaction = true;
299 Ok(())
300 })
301 }
302
303 fn commit<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
304 Box::pin(async move {
305 if self.in_transaction {
306 self.execute("COMMIT").await?;
307 self.in_transaction = false;
308 }
309 Ok(())
310 })
311 }
312
313 fn rollback<'a>(
314 &'a mut self,
315 ) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
316 Box::pin(async move {
317 if self.in_transaction {
318 let result = self.execute("ROLLBACK").await;
319 self.in_transaction = false;
320 result.map(|_| ())
321 } else {
322 Ok(())
323 }
324 })
325 }
326
327 fn is_connected(&self) -> bool {
328 self.connected
329 }
330
331 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
332 Box::pin(async move {
333 match self.execute("SELECT 1").await {
334 Ok(_) => true,
335 Err(_) => {
336 self.connected = false;
337 false
338 }
339 }
340 })
341 }
342
343 fn close<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
344 Box::pin(async move {
345 if let Some(conn) = self.conn.take() {
346 drop(conn);
347 }
348 self.connected = false;
349 self.in_transaction = false;
350 Ok(())
351 })
352 }
353
354 fn execute_with_params<'a>(
360 &'a mut self,
361 sql: &'a str,
362 params: &'a [Value],
363 ) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
364 Box::pin(async move {
365 if needs_raw_sql(sql) || params.is_empty() {
366 return self.execute(sql).await;
367 }
368 let mut pool_conn = self
369 .conn
370 .take()
371 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
372 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
373 for v in params {
374 q = match v {
375 Value::Null => q.bind(None::<i64>),
376 Value::Bool(b) => q.bind(*b),
377 Value::I8(n) => q.bind(*n),
378 Value::I16(n) => q.bind(*n),
379 Value::I32(n) => q.bind(*n),
380 Value::I64(n) => q.bind(*n),
381 Value::U8(n) => q.bind(*n),
382 Value::U16(n) => q.bind(*n),
383 Value::U32(n) => q.bind(*n),
384 Value::U64(n) => q.bind(*n as i64),
385 Value::F32(f) => q.bind(*f),
386 Value::F64(f) => q.bind(*f),
387 Value::String(s) => q.bind(s.as_str()),
388 Value::Decimal(s) => q.bind(s.as_str()),
389 Value::Bytes(b) => q.bind(b.as_slice()),
390 other => q.bind(other.to_string()),
391 };
392 }
393 let result = q.execute(&mut *pool_conn).await;
394 self.conn = Some(pool_conn);
395 match result {
396 Ok(r) => Ok(r.rows_affected()),
397 Err(e) => {
398 let db_err = map_sqlx_error(e);
399 if matches!(db_err, DbError::ConnectionError(_) | DbError::IoError(_)) {
400 self.connected = false;
401 }
402 Err(db_err)
403 }
404 }
405 })
406 }
407
408 fn query_with_params<'a>(
413 &'a mut self,
414 sql: &'a str,
415 params: &'a [Value],
416 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, DbError>> + Send + 'a>> {
417 Box::pin(async move {
418 if params.is_empty() {
419 return self.query(sql).await;
420 }
421 let mut pool_conn = self
422 .conn
423 .take()
424 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
425 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
426 for v in params {
427 q = match v {
428 Value::Null => q.bind(None::<i64>),
429 Value::Bool(b) => q.bind(*b),
430 Value::I8(n) => q.bind(*n),
431 Value::I16(n) => q.bind(*n),
432 Value::I32(n) => q.bind(*n),
433 Value::I64(n) => q.bind(*n),
434 Value::U8(n) => q.bind(*n),
435 Value::U16(n) => q.bind(*n),
436 Value::U32(n) => q.bind(*n),
437 Value::U64(n) => q.bind(*n as i64),
438 Value::F32(f) => q.bind(*f),
439 Value::F64(f) => q.bind(*f),
440 Value::String(s) => q.bind(s.as_str()),
441 Value::Decimal(s) => q.bind(s.as_str()),
442 Value::Bytes(b) => q.bind(b.as_slice()),
443 other => q.bind(other.to_string()),
444 };
445 }
446 let rows_result = q.fetch_all(&mut *pool_conn).await;
447 self.conn = Some(pool_conn);
448 let rows = rows_result.map_err(map_sqlx_error)?;
449 if rows.is_empty() {
450 return Ok(Vec::new());
451 }
452 let col_types: Vec<ColType> = rows[0]
455 .columns()
456 .iter()
457 .map(|col| ColType::parse_sqlite(col.type_info().name()))
458 .collect();
459 let col_names: Vec<String> = rows[0]
460 .columns()
461 .iter()
462 .map(|col| col.name().to_string())
463 .collect();
464 let mut result = Vec::with_capacity(rows.len());
465 for row in &rows {
466 let mut record = HashMap::with_capacity(col_types.len());
467 for (i, name) in col_names.iter().enumerate() {
468 let value = row_to_value_with_coltype_sqlite(row, i, col_types[i]);
469 record.insert(name.clone(), value);
470 }
471 result.push(record);
472 }
473 Ok(result)
474 })
475 }
476
477 fn query_values<'a>(
483 &'a mut self,
484 sql: &'a str,
485 ) -> Pin<Box<dyn Future<Output = Result<QueryValues, DbError>> + Send + 'a>> {
486 Box::pin(async move {
487 let mut pool_conn = self
488 .conn
489 .take()
490 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
491 let rows_result = (&mut *pool_conn).fetch_all(sqlx::AssertSqlSafe(sql)).await;
492 self.conn = Some(pool_conn);
493 let rows = rows_result.map_err(map_sqlx_error)?;
494 if rows.is_empty() {
495 return Ok((Vec::new(), Vec::new()));
496 }
497 let cols = rows[0].columns();
498 let mut col_names: Vec<String> = Vec::with_capacity(cols.len());
499 let mut col_types: Vec<ColType> = Vec::with_capacity(cols.len());
500 for col in cols {
501 col_names.push(col.name().to_string());
502 col_types.push(ColType::parse_sqlite(col.type_info().name()));
503 }
504 let mut result_rows: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
505 for row in &rows {
506 let mut row_values: Vec<Value> = Vec::with_capacity(col_names.len());
507 for (idx, _) in col_names.iter().enumerate() {
508 row_values.push(row_to_value_with_coltype_sqlite(row, idx, col_types[idx]));
509 }
510 result_rows.push(row_values);
511 }
512 Ok((col_names, result_rows))
513 })
514 }
515
516 fn query_values_with_params<'a>(
520 &'a mut self,
521 sql: &'a str,
522 params: &'a [Value],
523 ) -> Pin<Box<dyn Future<Output = Result<QueryValues, DbError>> + Send + 'a>> {
524 Box::pin(async move {
525 if params.is_empty() {
526 return self.query_values(sql).await;
527 }
528 let mut pool_conn = self
529 .conn
530 .take()
531 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
532 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
533 for v in params {
534 q = match v {
535 Value::Null => q.bind(None::<i64>),
536 Value::Bool(b) => q.bind(*b),
537 Value::I8(n) => q.bind(*n),
538 Value::I16(n) => q.bind(*n),
539 Value::I32(n) => q.bind(*n),
540 Value::I64(n) => q.bind(*n),
541 Value::U8(n) => q.bind(*n),
542 Value::U16(n) => q.bind(*n),
543 Value::U32(n) => q.bind(*n),
544 Value::U64(n) => q.bind(*n as i64),
545 Value::F32(f) => q.bind(*f),
546 Value::F64(f) => q.bind(*f),
547 Value::String(s) => q.bind(s.as_str()),
548 Value::Decimal(s) => q.bind(s.as_str()),
549 Value::Bytes(b) => q.bind(b.as_slice()),
550 other => q.bind(other.to_string()),
551 };
552 }
553 let rows_result = q.fetch_all(&mut *pool_conn).await;
554 self.conn = Some(pool_conn);
555 let rows = rows_result.map_err(map_sqlx_error)?;
556 if rows.is_empty() {
557 return Ok((Vec::new(), Vec::new()));
558 }
559 let cols = rows[0].columns();
560 let mut col_names: Vec<String> = Vec::with_capacity(cols.len());
561 let mut col_types: Vec<ColType> = Vec::with_capacity(cols.len());
562 for col in cols {
563 col_names.push(col.name().to_string());
564 col_types.push(ColType::parse_sqlite(col.type_info().name()));
565 }
566 let mut result_rows: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
567 for row in &rows {
568 let mut row_values: Vec<Value> = Vec::with_capacity(col_names.len());
569 for (idx, _) in col_names.iter().enumerate() {
570 row_values.push(row_to_value_with_coltype_sqlite(row, idx, col_types[idx]));
571 }
572 result_rows.push(row_values);
573 }
574 Ok((col_names, result_rows))
575 })
576 }
577
578 fn query_stream<'a>(
584 &'a mut self,
585 sql: &'a str,
586 ) -> Pin<Box<dyn futures::Stream<Item = Result<HashMap<String, Value>, DbError>> + Send + 'a>>
587 {
588 Box::pin(async_stream::try_stream! {
589 let mut pool_conn = self
590 .conn
591 .take()
592 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
593 let mut row_stream = sqlx::query(sqlx::AssertSqlSafe(sql)).fetch(&mut *pool_conn);
594 let mut col_types: Vec<ColType> = Vec::new();
595 let mut col_names: Vec<String> = Vec::new();
596 let mut first_row = true;
597 while let Some(row_result) = row_stream.next().await {
598 let row = row_result.map_err(map_sqlx_error)?;
599 if first_row {
600 for col in row.columns() {
601 col_names.push(col.name().to_string());
602 col_types.push(ColType::parse_sqlite(col.type_info().name()));
603 }
604 first_row = false;
605 }
606 let mut record = HashMap::with_capacity(col_names.len());
607 for (i, name) in col_names.iter().enumerate() {
608 let value = row_to_value_with_coltype_sqlite(&row, i, col_types[i]);
609 record.insert(name.clone(), value);
610 }
611 yield record;
612 }
613 drop(row_stream);
616 self.conn = Some(pool_conn);
617 })
618 }
619}
620
621impl Drop for SqlxSqliteConnection {
622 fn drop(&mut self) {
623 if let Some(conn) = self.conn.take() {
624 drop(conn);
625 }
626 }
627}
628
629pub async fn sqlite_backup(
634 conn: &mut SqlxSqliteConnection,
635 dest_path: &str,
636) -> Result<(), DbError> {
637 let escaped_path = dest_path.replace('\'', "''");
638 let sql = format!("VACUUM INTO '{}'", escaped_path);
639 conn.execute(&sql).await?;
640 Ok(())
641}
642
643fn row_to_value_mysql(row: &sqlx::mysql::MySqlRow, ordinal: usize) -> Value {
649 use sqlx::TypeInfo;
650 let type_name = row.columns()[ordinal].type_info().name();
651 match type_name {
652 "BOOLEAN" => match row.try_get::<Option<bool>, usize>(ordinal) {
653 Ok(v) => v.map(Value::Bool).unwrap_or(Value::Null),
654 Err(_) => Value::Null,
655 },
656 "TINYINT" | "TINYINT UNSIGNED" => match row.try_get::<Option<i8>, usize>(ordinal) {
657 Ok(v) => v.map(Value::I8).unwrap_or(Value::Null),
658 Err(_) => match row.try_get::<Option<u8>, usize>(ordinal) {
659 Ok(v) => v.map(Value::U8).unwrap_or(Value::Null),
660 Err(_) => Value::Null,
661 },
662 },
663 "SMALLINT" | "SMALLINT UNSIGNED" => match row.try_get::<Option<i16>, usize>(ordinal) {
664 Ok(v) => v.map(Value::I16).unwrap_or(Value::Null),
665 Err(_) => match row.try_get::<Option<u16>, usize>(ordinal) {
666 Ok(v) => v.map(Value::U16).unwrap_or(Value::Null),
667 Err(_) => Value::Null,
668 },
669 },
670 "INT" | "INT UNSIGNED" | "MEDIUMINT" | "MEDIUMINT UNSIGNED" => {
671 match row.try_get::<Option<i32>, usize>(ordinal) {
672 Ok(v) => v.map(Value::I32).unwrap_or(Value::Null),
673 Err(_) => match row.try_get::<Option<u32>, usize>(ordinal) {
674 Ok(v) => v.map(Value::U32).unwrap_or(Value::Null),
675 Err(_) => Value::Null,
676 },
677 }
678 }
679 "BIGINT" | "BIGINT UNSIGNED" => match row.try_get::<Option<i64>, usize>(ordinal) {
680 Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
681 Err(_) => match row.try_get::<Option<u64>, usize>(ordinal) {
682 Ok(v) => v.map(Value::U64).unwrap_or(Value::Null),
683 Err(_) => Value::Null,
684 },
685 },
686 "FLOAT" => match row.try_get::<Option<f32>, usize>(ordinal) {
687 Ok(v) => v.map(Value::F32).unwrap_or(Value::Null),
688 Err(_) => Value::Null,
689 },
690 "DOUBLE" => match row.try_get::<Option<f64>, usize>(ordinal) {
691 Ok(v) => v.map(Value::F64).unwrap_or(Value::Null),
692 Err(_) => Value::Null,
693 },
694 "VARCHAR" | "TEXT" | "CHAR" | "TINYTEXT" | "MEDIUMTEXT" | "LONGTEXT" | "ENUM" => {
695 match row.try_get::<Option<String>, usize>(ordinal) {
696 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
697 Err(_) => Value::Null,
698 }
699 }
700 "BLOB" | "TINYBLOB" | "MEDIUMBLOB" | "LONGBLOB" | "BINARY" | "VARBINARY" => {
701 match row.try_get::<Option<Vec<u8>>, usize>(ordinal) {
702 Ok(v) => v.map(Value::Bytes).unwrap_or(Value::Null),
703 Err(_) => Value::Null,
704 }
705 }
706 "DECIMAL" | "NUMERIC" | "NEWDECIMAL" => {
708 match row.try_get::<Option<rust_decimal::Decimal>, usize>(ordinal) {
709 Ok(Some(v)) => Value::Decimal(v.to_string()),
710 Ok(None) => Value::Null,
711 Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
712 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
713 Err(_) => Value::Null,
714 },
715 }
716 }
717 _ => {
718 if let Ok(v) = row.try_get::<Option<i64>, usize>(ordinal) {
720 return v.map(Value::I64).unwrap_or(Value::Null);
721 }
722 if let Ok(v) = row.try_get::<Option<f64>, usize>(ordinal) {
723 return v.map(Value::F64).unwrap_or(Value::Null);
724 }
725 if let Ok(v) = row.try_get::<Option<bool>, usize>(ordinal) {
726 return v.map(Value::Bool).unwrap_or(Value::Null);
727 }
728 if let Ok(v) = row.try_get::<Option<String>, usize>(ordinal) {
729 return v.map(Value::String).unwrap_or(Value::Null);
730 }
731 Value::Null
732 }
733 }
734}
735
736fn row_to_value_with_coltype_mysql(
744 row: &sqlx::mysql::MySqlRow,
745 ordinal: usize,
746 col_type: ColType,
747) -> Value {
748 match col_type {
749 ColType::Bool => match row.try_get::<Option<bool>, usize>(ordinal) {
750 Ok(v) => v.map(Value::Bool).unwrap_or(Value::Null),
751 Err(_) => Value::Null,
752 },
753 ColType::I8 => match row.try_get::<Option<i8>, usize>(ordinal) {
754 Ok(v) => v.map(Value::I8).unwrap_or(Value::Null),
755 Err(_) => match row.try_get::<Option<u8>, usize>(ordinal) {
756 Ok(v) => v.map(Value::U8).unwrap_or(Value::Null),
757 Err(_) => Value::Null,
758 },
759 },
760 ColType::I16 => match row.try_get::<Option<i16>, usize>(ordinal) {
761 Ok(v) => v.map(Value::I16).unwrap_or(Value::Null),
762 Err(_) => match row.try_get::<Option<u16>, usize>(ordinal) {
763 Ok(v) => v.map(Value::U16).unwrap_or(Value::Null),
764 Err(_) => Value::Null,
765 },
766 },
767 ColType::I32 => match row.try_get::<Option<i32>, usize>(ordinal) {
768 Ok(v) => v.map(Value::I32).unwrap_or(Value::Null),
769 Err(_) => match row.try_get::<Option<u32>, usize>(ordinal) {
770 Ok(v) => v.map(Value::U32).unwrap_or(Value::Null),
771 Err(_) => Value::Null,
772 },
773 },
774 ColType::I64 => match row.try_get::<Option<i64>, usize>(ordinal) {
775 Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
776 Err(_) => match row.try_get::<Option<u64>, usize>(ordinal) {
777 Ok(v) => v.map(Value::U64).unwrap_or(Value::Null),
778 Err(_) => Value::Null,
779 },
780 },
781 ColType::U8 => match row.try_get::<Option<u8>, usize>(ordinal) {
782 Ok(v) => v.map(Value::U8).unwrap_or(Value::Null),
783 Err(_) => Value::Null,
784 },
785 ColType::U16 => match row.try_get::<Option<u16>, usize>(ordinal) {
786 Ok(v) => v.map(Value::U16).unwrap_or(Value::Null),
787 Err(_) => Value::Null,
788 },
789 ColType::U32 => match row.try_get::<Option<u32>, usize>(ordinal) {
790 Ok(v) => v.map(Value::U32).unwrap_or(Value::Null),
791 Err(_) => Value::Null,
792 },
793 ColType::U64 => match row.try_get::<Option<u64>, usize>(ordinal) {
794 Ok(v) => v.map(Value::U64).unwrap_or(Value::Null),
795 Err(_) => Value::Null,
796 },
797 ColType::F32 => match row.try_get::<Option<f32>, usize>(ordinal) {
798 Ok(v) => v.map(Value::F32).unwrap_or(Value::Null),
799 Err(_) => Value::Null,
800 },
801 ColType::F64 => match row.try_get::<Option<f64>, usize>(ordinal) {
802 Ok(v) => v.map(Value::F64).unwrap_or(Value::Null),
803 Err(_) => Value::Null,
804 },
805 ColType::Decimal => match row.try_get::<Option<rust_decimal::Decimal>, usize>(ordinal) {
807 Ok(Some(v)) => Value::Decimal(v.to_string()),
808 Ok(None) => Value::Null,
809 Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
810 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
811 Err(_) => Value::Null,
812 },
813 },
814 ColType::String => match row.try_get::<Option<String>, usize>(ordinal) {
815 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
816 Err(_) => Value::Null,
817 },
818 ColType::Bytes => match row.try_get::<Option<Vec<u8>>, usize>(ordinal) {
819 Ok(v) => v.map(Value::Bytes).unwrap_or(Value::Null),
820 Err(_) => Value::Null,
821 },
822 ColType::Date | ColType::DateTime | ColType::Time | ColType::Json | ColType::Uuid => {
824 match row.try_get::<Option<String>, usize>(ordinal) {
825 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
826 Err(_) => Value::Null,
827 }
828 }
829 ColType::Unknown => {
830 if let Ok(v) = row.try_get::<Option<i64>, usize>(ordinal) {
832 return v.map(Value::I64).unwrap_or(Value::Null);
833 }
834 if let Ok(v) = row.try_get::<Option<f64>, usize>(ordinal) {
835 return v.map(Value::F64).unwrap_or(Value::Null);
836 }
837 if let Ok(v) = row.try_get::<Option<bool>, usize>(ordinal) {
838 return v.map(Value::Bool).unwrap_or(Value::Null);
839 }
840 if let Ok(v) = row.try_get::<Option<String>, usize>(ordinal) {
841 return v.map(Value::String).unwrap_or(Value::Null);
842 }
843 Value::Null
844 }
845 _ => Value::Null,
847 }
848}
849
850pub struct MySqlPoolHandle {
851 pool: sqlx::MySqlPool,
852}
853
854impl MySqlPoolHandle {
855 pub async fn connect(url: &str) -> Result<Self, DbError> {
856 let pool = sqlx::pool::PoolOptions::<sqlx::MySql>::new()
857 .max_connections(10)
858 .acquire_timeout(std::time::Duration::from_secs(30))
859 .idle_timeout(Some(std::time::Duration::from_secs(600)))
860 .max_lifetime(Some(std::time::Duration::from_secs(1800)))
861 .connect(url)
862 .await
863 .map_err(map_sqlx_error)?;
864 Ok(Self { pool })
865 }
866
867 pub fn from_pool(pool: sqlx::MySqlPool) -> Self {
868 Self { pool }
869 }
870
871 pub fn pool(&self) -> &sqlx::MySqlPool {
872 &self.pool
873 }
874}
875
876pub struct SqlxMySqlConnectionFactory {
877 pool: Arc<MySqlPoolHandle>,
878}
879
880impl SqlxMySqlConnectionFactory {
881 pub fn new(pool: Arc<MySqlPoolHandle>) -> Self {
882 Self { pool }
883 }
884}
885
886#[async_trait]
887impl ConnectionFactory for SqlxMySqlConnectionFactory {
888 async fn create(&self) -> Result<Box<dyn Connection>, DbError> {
889 let conn = self.pool.pool.acquire().await.map_err(map_sqlx_error)?;
890 Ok(Box::new(SqlxMySqlConnection {
891 conn: Some(conn),
892 connected: true,
893 in_transaction: false,
894 }))
895 }
896}
897
898pub struct SqlxMySqlConnection {
899 conn: Option<sqlx::pool::PoolConnection<sqlx::MySql>>,
900 connected: bool,
901 in_transaction: bool,
902}
903
904impl Connection for SqlxMySqlConnection {
905 fn execute<'a>(
906 &'a mut self,
907 sql: &'a str,
908 ) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
909 Box::pin(async move {
910 let mut pool_conn = self
911 .conn
912 .take()
913 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
914 let result = if needs_raw_sql(sql) {
917 (&mut *pool_conn)
918 .execute(sqlx::raw_sql(sqlx::AssertSqlSafe(sql)))
919 .await
920 } else {
921 (&mut *pool_conn).execute(sqlx::AssertSqlSafe(sql)).await
922 };
923 self.conn = Some(pool_conn);
924
925 match result {
926 Ok(r) => Ok(r.rows_affected()),
927 Err(e) => {
928 let db_err = map_sqlx_error(e);
929 if matches!(db_err, DbError::ConnectionError(_) | DbError::IoError(_)) {
930 self.connected = false;
931 }
932 Err(db_err)
933 }
934 }
935 })
936 }
937
938 fn query<'a>(
939 &'a mut self,
940 sql: &'a str,
941 ) -> Pin<Box<dyn Future<Output = Result<Vec<HashMap<String, Value>>, DbError>> + Send + 'a>>
942 {
943 Box::pin(async move {
944 let mut pool_conn = self
945 .conn
946 .take()
947 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
948 let rows_result = (&mut *pool_conn).fetch_all(sqlx::AssertSqlSafe(sql)).await;
951 self.conn = Some(pool_conn);
952
953 let rows = rows_result.map_err(map_sqlx_error)?;
954 if rows.is_empty() {
955 return Ok(Vec::new());
956 }
957 let col_types: Vec<ColType> = rows[0]
960 .columns()
961 .iter()
962 .map(|col| ColType::parse_mysql(col.type_info().name()))
963 .collect();
964 let col_names: Vec<String> = rows[0]
965 .columns()
966 .iter()
967 .map(|col| col.name().to_string())
968 .collect();
969 let mut result = Vec::with_capacity(rows.len());
970 for row in &rows {
971 let mut record = HashMap::with_capacity(col_types.len());
972 for (i, name) in col_names.iter().enumerate() {
973 let value = row_to_value_with_coltype_mysql(row, i, col_types[i]);
974 record.insert(name.clone(), value);
975 }
976 result.push(record);
977 }
978 Ok(result)
979 })
980 }
981
982 fn begin_transaction<'a>(
983 &'a mut self,
984 ) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
985 Box::pin(async move {
986 if self.in_transaction {
987 return Err(DbError::Internal("transaction already started".to_string()));
988 }
989 self.execute("BEGIN").await?;
990 self.in_transaction = true;
991 Ok(())
992 })
993 }
994
995 fn commit<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
996 Box::pin(async move {
997 if self.in_transaction {
998 self.execute("COMMIT").await?;
999 self.in_transaction = false;
1000 }
1001 Ok(())
1002 })
1003 }
1004
1005 fn rollback<'a>(
1006 &'a mut self,
1007 ) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
1008 Box::pin(async move {
1009 if self.in_transaction {
1010 let result = self.execute("ROLLBACK").await;
1011 self.in_transaction = false;
1012 result.map(|_| ())
1013 } else {
1014 Ok(())
1015 }
1016 })
1017 }
1018
1019 fn is_connected(&self) -> bool {
1020 self.connected
1021 }
1022
1023 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
1024 Box::pin(async move {
1025 match self.execute("SELECT 1").await {
1026 Ok(_) => true,
1027 Err(_) => {
1028 self.connected = false;
1029 false
1030 }
1031 }
1032 })
1033 }
1034
1035 fn close<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
1036 Box::pin(async move {
1037 if let Some(conn) = self.conn.take() {
1038 drop(conn);
1039 }
1040 self.connected = false;
1041 self.in_transaction = false;
1042 Ok(())
1043 })
1044 }
1045
1046 fn execute_with_params<'a>(
1052 &'a mut self,
1053 sql: &'a str,
1054 params: &'a [Value],
1055 ) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
1056 Box::pin(async move {
1057 if needs_raw_sql(sql) || params.is_empty() {
1058 return self.execute(sql).await;
1059 }
1060 let mut pool_conn = self
1061 .conn
1062 .take()
1063 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1064 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
1065 for v in params {
1066 q = match v {
1067 Value::Null => q.bind(None::<i64>),
1068 Value::Bool(b) => q.bind(*b),
1069 Value::I8(n) => q.bind(*n),
1070 Value::I16(n) => q.bind(*n),
1071 Value::I32(n) => q.bind(*n),
1072 Value::I64(n) => q.bind(*n),
1073 Value::U8(n) => q.bind(*n),
1074 Value::U16(n) => q.bind(*n),
1075 Value::U32(n) => q.bind(*n),
1076 Value::U64(n) => q.bind(*n as i64),
1077 Value::F32(f) => q.bind(*f),
1078 Value::F64(f) => q.bind(*f),
1079 Value::String(s) => q.bind(s.as_str()),
1080 Value::Decimal(s) => q.bind(s.as_str()),
1081 Value::Bytes(b) => q.bind(b.as_slice()),
1082 other => q.bind(other.to_string()),
1083 };
1084 }
1085 let result = q.execute(&mut *pool_conn).await;
1086 self.conn = Some(pool_conn);
1087 match result {
1088 Ok(r) => Ok(r.rows_affected()),
1089 Err(e) => {
1090 let db_err = map_sqlx_error(e);
1091 if matches!(db_err, DbError::ConnectionError(_) | DbError::IoError(_)) {
1092 self.connected = false;
1093 }
1094 Err(db_err)
1095 }
1096 }
1097 })
1098 }
1099
1100 fn query_with_params<'a>(
1102 &'a mut self,
1103 sql: &'a str,
1104 params: &'a [Value],
1105 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, DbError>> + Send + 'a>> {
1106 Box::pin(async move {
1107 if params.is_empty() {
1108 return self.query(sql).await;
1109 }
1110 let mut pool_conn = self
1111 .conn
1112 .take()
1113 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1114 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
1115 for v in params {
1116 q = match v {
1117 Value::Null => q.bind(None::<i64>),
1118 Value::Bool(b) => q.bind(*b),
1119 Value::I8(n) => q.bind(*n),
1120 Value::I16(n) => q.bind(*n),
1121 Value::I32(n) => q.bind(*n),
1122 Value::I64(n) => q.bind(*n),
1123 Value::U8(n) => q.bind(*n),
1124 Value::U16(n) => q.bind(*n),
1125 Value::U32(n) => q.bind(*n),
1126 Value::U64(n) => q.bind(*n as i64),
1127 Value::F32(f) => q.bind(*f),
1128 Value::F64(f) => q.bind(*f),
1129 Value::String(s) => q.bind(s.as_str()),
1130 Value::Decimal(s) => q.bind(s.as_str()),
1131 Value::Bytes(b) => q.bind(b.as_slice()),
1132 other => q.bind(other.to_string()),
1133 };
1134 }
1135 let rows_result = q.fetch_all(&mut *pool_conn).await;
1136 self.conn = Some(pool_conn);
1137 let rows = rows_result.map_err(map_sqlx_error)?;
1138 if rows.is_empty() {
1139 return Ok(Vec::new());
1140 }
1141 let col_names: Vec<String> = rows[0]
1142 .columns()
1143 .iter()
1144 .map(|col| col.name().to_string())
1145 .collect();
1146 let mut result = Vec::with_capacity(rows.len());
1147 for row in &rows {
1148 let mut record = HashMap::with_capacity(col_names.len());
1149 for (i, name) in col_names.iter().enumerate() {
1150 let value = row_to_value_mysql(row, i);
1151 record.insert(name.clone(), value);
1152 }
1153 result.push(record);
1154 }
1155 Ok(result)
1156 })
1157 }
1158
1159 fn query_values<'a>(
1163 &'a mut self,
1164 sql: &'a str,
1165 ) -> Pin<Box<dyn Future<Output = Result<QueryValues, DbError>> + Send + 'a>> {
1166 Box::pin(async move {
1167 let mut pool_conn = self
1168 .conn
1169 .take()
1170 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1171 let rows_result = (&mut *pool_conn).fetch_all(sqlx::AssertSqlSafe(sql)).await;
1172 self.conn = Some(pool_conn);
1173 let rows = rows_result.map_err(map_sqlx_error)?;
1174 if rows.is_empty() {
1175 return Ok((Vec::new(), Vec::new()));
1176 }
1177 let cols = rows[0].columns();
1178 let mut col_names: Vec<String> = Vec::with_capacity(cols.len());
1179 let mut col_types: Vec<ColType> = Vec::with_capacity(cols.len());
1180 for col in cols {
1181 col_names.push(col.name().to_string());
1182 col_types.push(ColType::parse_mysql(col.type_info().name()));
1183 }
1184 let mut result_rows: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
1185 for row in &rows {
1186 let mut row_values: Vec<Value> = Vec::with_capacity(col_names.len());
1187 for (idx, _) in col_names.iter().enumerate() {
1188 row_values.push(row_to_value_with_coltype_mysql(row, idx, col_types[idx]));
1189 }
1190 result_rows.push(row_values);
1191 }
1192 Ok((col_names, result_rows))
1193 })
1194 }
1195
1196 fn query_values_with_params<'a>(
1198 &'a mut self,
1199 sql: &'a str,
1200 params: &'a [Value],
1201 ) -> Pin<Box<dyn Future<Output = Result<QueryValues, DbError>> + Send + 'a>> {
1202 Box::pin(async move {
1203 if params.is_empty() {
1204 return self.query_values(sql).await;
1205 }
1206 let mut pool_conn = self
1207 .conn
1208 .take()
1209 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1210 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
1211 for v in params {
1212 q = match v {
1213 Value::Null => q.bind(None::<i64>),
1214 Value::Bool(b) => q.bind(*b),
1215 Value::I8(n) => q.bind(*n),
1216 Value::I16(n) => q.bind(*n),
1217 Value::I32(n) => q.bind(*n),
1218 Value::I64(n) => q.bind(*n),
1219 Value::U8(n) => q.bind(*n),
1220 Value::U16(n) => q.bind(*n),
1221 Value::U32(n) => q.bind(*n),
1222 Value::U64(n) => q.bind(*n as i64),
1223 Value::F32(f) => q.bind(*f),
1224 Value::F64(f) => q.bind(*f),
1225 Value::String(s) => q.bind(s.as_str()),
1226 Value::Decimal(s) => q.bind(s.as_str()),
1227 Value::Bytes(b) => q.bind(b.as_slice()),
1228 other => q.bind(other.to_string()),
1229 };
1230 }
1231 let rows_result = q.fetch_all(&mut *pool_conn).await;
1232 self.conn = Some(pool_conn);
1233 let rows = rows_result.map_err(map_sqlx_error)?;
1234 if rows.is_empty() {
1235 return Ok((Vec::new(), Vec::new()));
1236 }
1237 let cols = rows[0].columns();
1238 let mut col_names: Vec<String> = Vec::with_capacity(cols.len());
1239 let mut col_types: Vec<ColType> = Vec::with_capacity(cols.len());
1240 for col in cols {
1241 col_names.push(col.name().to_string());
1242 col_types.push(ColType::parse_mysql(col.type_info().name()));
1243 }
1244 let mut result_rows: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
1245 for row in &rows {
1246 let mut row_values: Vec<Value> = Vec::with_capacity(col_names.len());
1247 for (idx, _) in col_names.iter().enumerate() {
1248 row_values.push(row_to_value_with_coltype_mysql(row, idx, col_types[idx]));
1249 }
1250 result_rows.push(row_values);
1251 }
1252 Ok((col_names, result_rows))
1253 })
1254 }
1255
1256 fn query_stream<'a>(
1262 &'a mut self,
1263 sql: &'a str,
1264 ) -> Pin<Box<dyn futures::Stream<Item = Result<HashMap<String, Value>, DbError>> + Send + 'a>>
1265 {
1266 Box::pin(async_stream::try_stream! {
1267 let mut pool_conn = self
1268 .conn
1269 .take()
1270 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1271 let mut row_stream = sqlx::query(sqlx::AssertSqlSafe(sql)).fetch(&mut *pool_conn);
1272 let mut col_types: Vec<ColType> = Vec::new();
1273 let mut col_names: Vec<String> = Vec::new();
1274 let mut first_row = true;
1275 while let Some(row_result) = row_stream.next().await {
1276 let row = row_result.map_err(map_sqlx_error)?;
1277 if first_row {
1278 for col in row.columns() {
1279 col_names.push(col.name().to_string());
1280 col_types.push(ColType::parse_mysql(col.type_info().name()));
1281 }
1282 first_row = false;
1283 }
1284 let mut record = HashMap::with_capacity(col_names.len());
1285 for (i, name) in col_names.iter().enumerate() {
1286 let value = row_to_value_with_coltype_mysql(&row, i, col_types[i]);
1287 record.insert(name.clone(), value);
1288 }
1289 yield record;
1290 }
1291 drop(row_stream);
1294 self.conn = Some(pool_conn);
1295 })
1296 }
1297}
1298
1299impl Drop for SqlxMySqlConnection {
1300 fn drop(&mut self) {
1301 if let Some(conn) = self.conn.take() {
1302 drop(conn);
1303 }
1304 }
1305}
1306
1307pub async fn mysql_bulk_insert(
1314 conn: &mut SqlxMySqlConnection,
1315 table: &str,
1316 columns: &[&str],
1317 rows: &[Vec<Value>],
1318) -> Result<u64, DbError> {
1319 if rows.is_empty() {
1320 return Ok(0);
1321 }
1322 let col_list = columns.join(", ");
1323 let cols_per_row = columns.len();
1324 let row_placeholder = format!("({})", vec!["?"; cols_per_row].join(", "));
1326 let placeholders = vec![row_placeholder; rows.len()].join(", ");
1327 let sql = format!(
1328 "INSERT INTO {} ({}) VALUES {}",
1329 table, col_list, placeholders
1330 );
1331 let mut params: Vec<Value> = Vec::with_capacity(rows.len() * cols_per_row);
1333 for row in rows {
1334 for v in row {
1335 params.push(v.clone());
1336 }
1337 }
1338 conn.execute_with_params(&sql, ¶ms).await
1339}
1340
1341fn row_to_value_pg(row: &sqlx::postgres::PgRow, ordinal: usize) -> Value {
1347 use sqlx::TypeInfo;
1348 let type_name = row.columns()[ordinal].type_info().name();
1349 match type_name {
1350 "BOOL" => match row.try_get::<Option<bool>, usize>(ordinal) {
1351 Ok(v) => v.map(Value::Bool).unwrap_or(Value::Null),
1352 Err(_) => Value::Null,
1353 },
1354 "INT2" => match row.try_get::<Option<i16>, usize>(ordinal) {
1355 Ok(v) => v.map(Value::I16).unwrap_or(Value::Null),
1356 Err(_) => Value::Null,
1357 },
1358 "INT4" | "OID" => match row.try_get::<Option<i32>, usize>(ordinal) {
1359 Ok(v) => v.map(Value::I32).unwrap_or(Value::Null),
1360 Err(_) => Value::Null,
1361 },
1362 "INT8" => match row.try_get::<Option<i64>, usize>(ordinal) {
1363 Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
1364 Err(_) => Value::Null,
1365 },
1366 "FLOAT4" => match row.try_get::<Option<f32>, usize>(ordinal) {
1367 Ok(v) => v.map(Value::F32).unwrap_or(Value::Null),
1368 Err(_) => Value::Null,
1369 },
1370 "FLOAT8" => match row.try_get::<Option<f64>, usize>(ordinal) {
1371 Ok(v) => v.map(Value::F64).unwrap_or(Value::Null),
1372 Err(_) => Value::Null,
1373 },
1374 "TEXT" | "VARCHAR" | "CHAR" | "NAME" => match row.try_get::<Option<String>, usize>(ordinal)
1375 {
1376 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
1377 Err(_) => Value::Null,
1378 },
1379 "BYTEA" => match row.try_get::<Option<Vec<u8>>, usize>(ordinal) {
1380 Ok(v) => v.map(Value::Bytes).unwrap_or(Value::Null),
1381 Err(_) => Value::Null,
1382 },
1383 "NUMERIC" => match row.try_get::<Option<rust_decimal::Decimal>, usize>(ordinal) {
1384 Ok(Some(v)) => Value::Decimal(v.to_string()),
1385 Ok(None) => Value::Null,
1386 Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
1387 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
1388 Err(_) => Value::Null,
1389 },
1390 },
1391 "UUID" => match row.try_get::<Option<sqlx::types::Uuid>, usize>(ordinal) {
1393 Ok(v) => v
1394 .map(|uuid| Value::String(uuid.to_string()))
1395 .unwrap_or(Value::Null),
1396 Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
1397 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
1398 Err(_) => Value::Null,
1399 },
1400 },
1401 _ => {
1402 if let Ok(v) = row.try_get::<Option<i64>, usize>(ordinal) {
1404 return v.map(Value::I64).unwrap_or(Value::Null);
1405 }
1406 if let Ok(v) = row.try_get::<Option<f64>, usize>(ordinal) {
1407 return v.map(Value::F64).unwrap_or(Value::Null);
1408 }
1409 if let Ok(v) = row.try_get::<Option<bool>, usize>(ordinal) {
1410 return v.map(Value::Bool).unwrap_or(Value::Null);
1411 }
1412 if let Ok(v) = row.try_get::<Option<String>, usize>(ordinal) {
1413 return v.map(Value::String).unwrap_or(Value::Null);
1414 }
1415 Value::Null
1416 }
1417 }
1418}
1419
1420fn row_to_value_with_coltype_pg(
1428 row: &sqlx::postgres::PgRow,
1429 ordinal: usize,
1430 col_type: ColType,
1431) -> Value {
1432 match col_type {
1433 ColType::Bool => match row.try_get::<Option<bool>, usize>(ordinal) {
1434 Ok(v) => v.map(Value::Bool).unwrap_or(Value::Null),
1435 Err(_) => Value::Null,
1436 },
1437 ColType::I8 => match row.try_get::<Option<i8>, usize>(ordinal) {
1438 Ok(v) => v.map(Value::I8).unwrap_or(Value::Null),
1439 Err(_) => Value::Null,
1440 },
1441 ColType::I16 => match row.try_get::<Option<i16>, usize>(ordinal) {
1442 Ok(v) => v.map(Value::I16).unwrap_or(Value::Null),
1443 Err(_) => Value::Null,
1444 },
1445 ColType::I32 => match row.try_get::<Option<i32>, usize>(ordinal) {
1447 Ok(v) => v.map(Value::I32).unwrap_or(Value::Null),
1448 Err(_) => Value::Null,
1449 },
1450 ColType::I64 => match row.try_get::<Option<i64>, usize>(ordinal) {
1451 Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
1452 Err(_) => Value::Null,
1453 },
1454 ColType::U8 => match row.try_get::<Option<i16>, usize>(ordinal) {
1459 Ok(v) => v.map(Value::I16).unwrap_or(Value::Null),
1460 Err(_) => Value::Null,
1461 },
1462 ColType::U16 => match row.try_get::<Option<i32>, usize>(ordinal) {
1463 Ok(v) => v.map(Value::I32).unwrap_or(Value::Null),
1464 Err(_) => Value::Null,
1465 },
1466 ColType::U32 => match row.try_get::<Option<i64>, usize>(ordinal) {
1467 Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
1468 Err(_) => Value::Null,
1469 },
1470 ColType::U64 => match row.try_get::<Option<i64>, usize>(ordinal) {
1472 Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
1473 Err(_) => Value::Null,
1474 },
1475 ColType::F32 => match row.try_get::<Option<f32>, usize>(ordinal) {
1476 Ok(v) => v.map(Value::F32).unwrap_or(Value::Null),
1477 Err(_) => Value::Null,
1478 },
1479 ColType::F64 => match row.try_get::<Option<f64>, usize>(ordinal) {
1480 Ok(v) => v.map(Value::F64).unwrap_or(Value::Null),
1481 Err(_) => Value::Null,
1482 },
1483 ColType::Decimal => match row.try_get::<Option<rust_decimal::Decimal>, usize>(ordinal) {
1485 Ok(Some(v)) => Value::Decimal(v.to_string()),
1486 Ok(None) => Value::Null,
1487 Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
1488 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
1489 Err(_) => Value::Null,
1490 },
1491 },
1492 ColType::String => match row.try_get::<Option<String>, usize>(ordinal) {
1493 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
1494 Err(_) => Value::Null,
1495 },
1496 ColType::Bytes => match row.try_get::<Option<Vec<u8>>, usize>(ordinal) {
1497 Ok(v) => v.map(Value::Bytes).unwrap_or(Value::Null),
1498 Err(_) => Value::Null,
1499 },
1500 ColType::Date | ColType::DateTime | ColType::Time | ColType::Json => {
1502 match row.try_get::<Option<String>, usize>(ordinal) {
1503 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
1504 Err(_) => Value::Null,
1505 }
1506 }
1507 ColType::Uuid => match row.try_get::<Option<sqlx::types::Uuid>, usize>(ordinal) {
1509 Ok(v) => v
1510 .map(|uuid| Value::String(uuid.to_string()))
1511 .unwrap_or(Value::Null),
1512 Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
1513 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
1514 Err(_) => Value::Null,
1515 },
1516 },
1517 ColType::Unknown => {
1518 if let Ok(v) = row.try_get::<Option<i64>, usize>(ordinal) {
1520 return v.map(Value::I64).unwrap_or(Value::Null);
1521 }
1522 if let Ok(v) = row.try_get::<Option<f64>, usize>(ordinal) {
1523 return v.map(Value::F64).unwrap_or(Value::Null);
1524 }
1525 if let Ok(v) = row.try_get::<Option<bool>, usize>(ordinal) {
1526 return v.map(Value::Bool).unwrap_or(Value::Null);
1527 }
1528 if let Ok(v) = row.try_get::<Option<String>, usize>(ordinal) {
1529 return v.map(Value::String).unwrap_or(Value::Null);
1530 }
1531 Value::Null
1532 }
1533 _ => Value::Null,
1535 }
1536}
1537
1538pub struct PgPoolHandle {
1539 pool: sqlx::PgPool,
1540}
1541
1542impl PgPoolHandle {
1543 pub async fn connect(url: &str) -> Result<Self, DbError> {
1544 let pool = sqlx::pool::PoolOptions::<sqlx::Postgres>::new()
1545 .max_connections(10)
1546 .acquire_timeout(std::time::Duration::from_secs(30))
1547 .idle_timeout(Some(std::time::Duration::from_secs(600)))
1548 .max_lifetime(Some(std::time::Duration::from_secs(1800)))
1549 .connect(url)
1550 .await
1551 .map_err(map_sqlx_error)?;
1552 Ok(Self { pool })
1553 }
1554
1555 pub fn from_pool(pool: sqlx::PgPool) -> Self {
1556 Self { pool }
1557 }
1558
1559 pub fn pool(&self) -> &sqlx::PgPool {
1560 &self.pool
1561 }
1562}
1563
1564pub struct SqlxPgConnectionFactory {
1565 pool: Arc<PgPoolHandle>,
1566}
1567
1568impl SqlxPgConnectionFactory {
1569 pub fn new(pool: Arc<PgPoolHandle>) -> Self {
1570 Self { pool }
1571 }
1572}
1573
1574#[async_trait]
1575impl ConnectionFactory for SqlxPgConnectionFactory {
1576 async fn create(&self) -> Result<Box<dyn Connection>, DbError> {
1577 let conn = self.pool.pool.acquire().await.map_err(map_sqlx_error)?;
1578 Ok(Box::new(SqlxPgConnection {
1579 conn: Some(conn),
1580 connected: true,
1581 in_transaction: false,
1582 }))
1583 }
1584}
1585
1586pub struct SqlxPgConnection {
1587 conn: Option<sqlx::pool::PoolConnection<sqlx::Postgres>>,
1588 connected: bool,
1589 in_transaction: bool,
1590}
1591
1592impl Connection for SqlxPgConnection {
1593 fn execute<'a>(
1594 &'a mut self,
1595 sql: &'a str,
1596 ) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
1597 Box::pin(async move {
1598 let mut pool_conn = self
1599 .conn
1600 .take()
1601 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1602 let result = if needs_raw_sql(sql) {
1605 (&mut *pool_conn)
1606 .execute(sqlx::raw_sql(sqlx::AssertSqlSafe(sql)))
1607 .await
1608 } else {
1609 (&mut *pool_conn).execute(sqlx::AssertSqlSafe(sql)).await
1610 };
1611 self.conn = Some(pool_conn);
1612
1613 match result {
1614 Ok(r) => Ok(r.rows_affected()),
1615 Err(e) => {
1616 let db_err = map_sqlx_error(e);
1617 if matches!(db_err, DbError::ConnectionError(_) | DbError::IoError(_)) {
1618 self.connected = false;
1619 }
1620 Err(db_err)
1621 }
1622 }
1623 })
1624 }
1625
1626 fn query<'a>(
1627 &'a mut self,
1628 sql: &'a str,
1629 ) -> Pin<Box<dyn Future<Output = Result<Vec<HashMap<String, Value>>, DbError>> + Send + 'a>>
1630 {
1631 Box::pin(async move {
1632 let mut pool_conn = self
1633 .conn
1634 .take()
1635 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1636 let rows_result = (&mut *pool_conn).fetch_all(sqlx::AssertSqlSafe(sql)).await;
1639 self.conn = Some(pool_conn);
1640
1641 let rows = rows_result.map_err(map_sqlx_error)?;
1642 if rows.is_empty() {
1643 return Ok(Vec::new());
1644 }
1645 let col_types: Vec<ColType> = rows[0]
1648 .columns()
1649 .iter()
1650 .map(|col| ColType::parse_postgres(col.type_info().name()))
1651 .collect();
1652 let col_names: Vec<String> = rows[0]
1653 .columns()
1654 .iter()
1655 .map(|col| col.name().to_string())
1656 .collect();
1657 let mut result = Vec::with_capacity(rows.len());
1658 for row in &rows {
1659 let mut record = HashMap::with_capacity(col_types.len());
1660 for (i, name) in col_names.iter().enumerate() {
1661 let value = row_to_value_with_coltype_pg(row, i, col_types[i]);
1662 record.insert(name.clone(), value);
1663 }
1664 result.push(record);
1665 }
1666 Ok(result)
1667 })
1668 }
1669
1670 fn begin_transaction<'a>(
1671 &'a mut self,
1672 ) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
1673 Box::pin(async move {
1674 if self.in_transaction {
1675 return Err(DbError::Internal("transaction already started".to_string()));
1676 }
1677 self.execute("BEGIN").await?;
1678 self.in_transaction = true;
1679 Ok(())
1680 })
1681 }
1682
1683 fn commit<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
1684 Box::pin(async move {
1685 if self.in_transaction {
1686 self.execute("COMMIT").await?;
1687 self.in_transaction = false;
1688 }
1689 Ok(())
1690 })
1691 }
1692
1693 fn rollback<'a>(
1694 &'a mut self,
1695 ) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
1696 Box::pin(async move {
1697 if self.in_transaction {
1698 let result = self.execute("ROLLBACK").await;
1699 self.in_transaction = false;
1700 result.map(|_| ())
1701 } else {
1702 Ok(())
1703 }
1704 })
1705 }
1706
1707 fn is_connected(&self) -> bool {
1708 self.connected
1709 }
1710
1711 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
1712 Box::pin(async move {
1713 match self.execute("SELECT 1").await {
1714 Ok(_) => true,
1715 Err(_) => {
1716 self.connected = false;
1717 false
1718 }
1719 }
1720 })
1721 }
1722
1723 fn close<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
1724 Box::pin(async move {
1725 if let Some(conn) = self.conn.take() {
1726 drop(conn);
1727 }
1728 self.connected = false;
1729 self.in_transaction = false;
1730 Ok(())
1731 })
1732 }
1733
1734 fn execute_with_params<'a>(
1749 &'a mut self,
1750 sql: &'a str,
1751 params: &'a [Value],
1752 ) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
1753 Box::pin(async move {
1754 if needs_raw_sql(sql) || params.is_empty() {
1755 return self.execute(sql).await;
1756 }
1757 let mut pool_conn = self
1758 .conn
1759 .take()
1760 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1761 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
1762 for v in params {
1763 q = match v {
1764 Value::Null => q.bind(None::<i64>),
1765 Value::Bool(b) => q.bind(*b),
1766 Value::I8(n) => q.bind(*n),
1767 Value::I16(n) => q.bind(*n),
1768 Value::I32(n) => q.bind(*n),
1769 Value::I64(n) => q.bind(*n),
1770 Value::U8(n) => q.bind(*n as i16),
1772 Value::U16(n) => q.bind(*n as i32),
1773 Value::U32(n) => q.bind(*n as i64),
1774 Value::U64(n) => q.bind(*n as i64),
1775 Value::F32(f) => q.bind(*f),
1776 Value::F64(f) => q.bind(*f),
1777 Value::String(s) => q.bind(s.as_str()),
1778 Value::Decimal(s) => q.bind(s.as_str()),
1779 Value::Bytes(b) => q.bind(b.as_slice()),
1780 other => q.bind(other.to_string()),
1781 };
1782 }
1783 let result = q.execute(&mut *pool_conn).await;
1784 self.conn = Some(pool_conn);
1785 match result {
1786 Ok(r) => Ok(r.rows_affected()),
1787 Err(e) => {
1788 let db_err = map_sqlx_error(e);
1789 if matches!(db_err, DbError::ConnectionError(_) | DbError::IoError(_)) {
1790 self.connected = false;
1791 }
1792 Err(db_err)
1793 }
1794 }
1795 })
1796 }
1797
1798 fn query_with_params<'a>(
1802 &'a mut self,
1803 sql: &'a str,
1804 params: &'a [Value],
1805 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, DbError>> + Send + 'a>> {
1806 Box::pin(async move {
1807 if params.is_empty() {
1808 return self.query(sql).await;
1809 }
1810 let mut pool_conn = self
1811 .conn
1812 .take()
1813 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1814 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
1815 for v in params {
1816 q = match v {
1817 Value::Null => q.bind(None::<i64>),
1818 Value::Bool(b) => q.bind(*b),
1819 Value::I8(n) => q.bind(*n),
1820 Value::I16(n) => q.bind(*n),
1821 Value::I32(n) => q.bind(*n),
1822 Value::I64(n) => q.bind(*n),
1823 Value::U8(n) => q.bind(*n as i16),
1824 Value::U16(n) => q.bind(*n as i32),
1825 Value::U32(n) => q.bind(*n as i64),
1826 Value::U64(n) => q.bind(*n as i64),
1827 Value::F32(f) => q.bind(*f),
1828 Value::F64(f) => q.bind(*f),
1829 Value::String(s) => q.bind(s.as_str()),
1830 Value::Decimal(s) => q.bind(s.as_str()),
1831 Value::Bytes(b) => q.bind(b.as_slice()),
1832 other => q.bind(other.to_string()),
1833 };
1834 }
1835 let rows_result = q.fetch_all(&mut *pool_conn).await;
1836 self.conn = Some(pool_conn);
1837 let rows = rows_result.map_err(map_sqlx_error)?;
1838 if rows.is_empty() {
1839 return Ok(Vec::new());
1840 }
1841 let col_types: Vec<ColType> = rows[0]
1843 .columns()
1844 .iter()
1845 .map(|col| ColType::parse_postgres(col.type_info().name()))
1846 .collect();
1847 let col_names: Vec<String> = rows[0]
1848 .columns()
1849 .iter()
1850 .map(|col| col.name().to_string())
1851 .collect();
1852 let mut result = Vec::with_capacity(rows.len());
1853 for row in &rows {
1854 let mut record = HashMap::with_capacity(col_types.len());
1855 for (i, name) in col_names.iter().enumerate() {
1856 let value = row_to_value_with_coltype_pg(row, i, col_types[i]);
1857 record.insert(name.clone(), value);
1858 }
1859 result.push(record);
1860 }
1861 Ok(result)
1862 })
1863 }
1864
1865 fn query_values<'a>(
1869 &'a mut self,
1870 sql: &'a str,
1871 ) -> Pin<Box<dyn Future<Output = Result<QueryValues, DbError>> + Send + 'a>> {
1872 Box::pin(async move {
1873 let mut pool_conn = self
1874 .conn
1875 .take()
1876 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1877 let rows_result = (&mut *pool_conn).fetch_all(sqlx::AssertSqlSafe(sql)).await;
1878 self.conn = Some(pool_conn);
1879 let rows = rows_result.map_err(map_sqlx_error)?;
1880 if rows.is_empty() {
1881 return Ok((Vec::new(), Vec::new()));
1882 }
1883 let cols = rows[0].columns();
1884 let mut col_names: Vec<String> = Vec::with_capacity(cols.len());
1885 let mut col_types: Vec<ColType> = Vec::with_capacity(cols.len());
1886 for col in cols {
1887 col_names.push(col.name().to_string());
1888 col_types.push(ColType::parse_postgres(col.type_info().name()));
1889 }
1890 let mut result_rows: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
1891 for row in &rows {
1892 let mut row_values: Vec<Value> = Vec::with_capacity(col_names.len());
1893 for (idx, _) in col_names.iter().enumerate() {
1894 row_values.push(row_to_value_with_coltype_pg(row, idx, col_types[idx]));
1895 }
1896 result_rows.push(row_values);
1897 }
1898 Ok((col_names, result_rows))
1899 })
1900 }
1901
1902 fn query_values_with_params<'a>(
1906 &'a mut self,
1907 sql: &'a str,
1908 params: &'a [Value],
1909 ) -> Pin<Box<dyn Future<Output = Result<QueryValues, DbError>> + Send + 'a>> {
1910 Box::pin(async move {
1911 if params.is_empty() {
1912 return self.query_values(sql).await;
1913 }
1914 let mut pool_conn = self
1915 .conn
1916 .take()
1917 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1918 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
1919 for v in params {
1920 q = match v {
1921 Value::Null => q.bind(None::<i64>),
1922 Value::Bool(b) => q.bind(*b),
1923 Value::I8(n) => q.bind(*n),
1924 Value::I16(n) => q.bind(*n),
1925 Value::I32(n) => q.bind(*n),
1926 Value::I64(n) => q.bind(*n),
1927 Value::U8(n) => q.bind(*n as i16),
1928 Value::U16(n) => q.bind(*n as i32),
1929 Value::U32(n) => q.bind(*n as i64),
1930 Value::U64(n) => q.bind(*n as i64),
1931 Value::F32(f) => q.bind(*f),
1932 Value::F64(f) => q.bind(*f),
1933 Value::String(s) => q.bind(s.as_str()),
1934 Value::Decimal(s) => q.bind(s.as_str()),
1935 Value::Bytes(b) => q.bind(b.as_slice()),
1936 other => q.bind(other.to_string()),
1937 };
1938 }
1939 let rows_result = q.fetch_all(&mut *pool_conn).await;
1940 self.conn = Some(pool_conn);
1941 let rows = rows_result.map_err(map_sqlx_error)?;
1942 if rows.is_empty() {
1943 return Ok((Vec::new(), Vec::new()));
1944 }
1945 let cols = rows[0].columns();
1946 let mut col_names: Vec<String> = Vec::with_capacity(cols.len());
1947 for col in cols {
1948 col_names.push(col.name().to_string());
1949 }
1950 let mut result_rows: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
1951 for row in rows {
1952 let mut row_values: Vec<Value> = Vec::with_capacity(col_names.len());
1953 for (idx, _) in col_names.iter().enumerate() {
1954 let ordinal = row.columns()[idx].ordinal();
1955 row_values.push(row_to_value_pg(&row, ordinal));
1956 }
1957 result_rows.push(row_values);
1958 }
1959 Ok((col_names, result_rows))
1960 })
1961 }
1962
1963 fn query_stream<'a>(
1973 &'a mut self,
1974 sql: &'a str,
1975 ) -> Pin<Box<dyn futures::Stream<Item = Result<HashMap<String, Value>, DbError>> + Send + 'a>>
1976 {
1977 Box::pin(async_stream::try_stream! {
1978 let mut pool_conn = self
1979 .conn
1980 .take()
1981 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1982 let mut row_stream = sqlx::query(sqlx::AssertSqlSafe(sql)).fetch(&mut *pool_conn);
1983 while let Some(row_result) = row_stream.next().await {
1986 let row = row_result.map_err(map_sqlx_error)?;
1987 let cols = row.columns();
1988 let mut record = HashMap::with_capacity(cols.len());
1989 for (i, col) in cols.iter().enumerate() {
1990 let name = col.name().to_string();
1991 let value = row_to_value_pg(&row, i);
1993 record.insert(name, value);
1994 }
1995 yield record;
1996 }
1997 drop(row_stream);
2000 self.conn = Some(pool_conn);
2001 })
2002 }
2003}
2004
2005impl Drop for SqlxPgConnection {
2006 fn drop(&mut self) {
2007 if let Some(conn) = self.conn.take() {
2008 drop(conn);
2009 }
2010 }
2011}
2012
2013#[async_trait]
2034pub trait PgExtensions: Send + Sync {
2035 async fn listen(&mut self, channel: &str) -> Result<(), DbError>;
2046
2047 async fn notify(&mut self, channel: &str, payload: &str) -> Result<(), DbError>;
2058
2059 async fn copy_from_stdin(&mut self, sql: &str, data: &[u8]) -> Result<u64, DbError>;
2074}
2075
2076fn validate_pg_channel_name(channel: &str) -> Result<(), DbError> {
2081 if channel.is_empty() {
2082 return Err(DbError::Internal(
2083 "PG channel name must not be empty".to_string(),
2084 ));
2085 }
2086 if !channel
2087 .chars()
2088 .all(|c| c.is_ascii_alphanumeric() || c == '_')
2089 {
2090 return Err(DbError::Internal(format!(
2091 "invalid PG channel name: {} (only alphanumeric and underscore allowed)",
2092 channel
2093 )));
2094 }
2095 Ok(())
2096}
2097
2098#[async_trait]
2099impl PgExtensions for SqlxPgConnection {
2100 async fn listen(&mut self, channel: &str) -> Result<(), DbError> {
2101 validate_pg_channel_name(channel)?;
2102 self.execute(&format!("LISTEN {}", channel)).await?;
2104 Ok(())
2105 }
2106
2107 async fn notify(&mut self, channel: &str, payload: &str) -> Result<(), DbError> {
2108 validate_pg_channel_name(channel)?;
2109 let escaped_payload = payload.replace('\'', "''");
2111 self.execute(&format!("NOTIFY {}, '{}'", channel, escaped_payload))
2112 .await?;
2113 Ok(())
2114 }
2115
2116 async fn copy_from_stdin(&mut self, sql: &str, data: &[u8]) -> Result<u64, DbError> {
2117 let mut pool_conn = self
2118 .conn
2119 .take()
2120 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
2121 let mut copy = (*pool_conn)
2125 .copy_in_raw(sql)
2126 .await
2127 .map_err(map_sqlx_error)?;
2128 copy.send(data).await.map_err(map_sqlx_error)?;
2130 let result = copy.finish().await.map_err(map_sqlx_error)?;
2132 self.conn = Some(pool_conn);
2133 Ok(result)
2134 }
2135}
2136
2137pub async fn pg_bulk_insert(
2145 conn: &mut SqlxPgConnection,
2146 table: &str,
2147 columns: &[&str],
2148 rows: &[Vec<Value>],
2149) -> Result<u64, DbError> {
2150 if rows.is_empty() {
2151 return Ok(0);
2152 }
2153 let col_list = columns.join(", ");
2154 let cols_per_row = columns.len();
2155 let placeholders: Vec<String> = rows
2157 .iter()
2158 .enumerate()
2159 .map(|(row_idx, _)| {
2160 let base = row_idx * cols_per_row;
2161 let ph: Vec<String> = (0..cols_per_row)
2162 .map(|i| format!("${}", base + i + 1))
2163 .collect();
2164 format!("({})", ph.join(", "))
2165 })
2166 .collect();
2167 let sql = format!(
2168 "INSERT INTO {} ({}) VALUES {}",
2169 table,
2170 col_list,
2171 placeholders.join(", ")
2172 );
2173 let mut params: Vec<Value> = Vec::with_capacity(rows.len() * cols_per_row);
2175 for row in rows {
2176 for v in row {
2177 params.push(v.clone());
2178 }
2179 }
2180 conn.execute_with_params(&sql, ¶ms).await
2181}