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 mut result = Vec::with_capacity(rows.len());
273 for row in &rows {
274 let mut record = HashMap::with_capacity(col_types.len());
275 for (i, col) in row.columns().iter().enumerate() {
276 let name = col.name().to_string();
277 let value = row_to_value_with_coltype_sqlite(row, i, col_types[i]);
278 record.insert(name, value);
279 }
280 result.push(record);
281 }
282 Ok(result)
283 })
284 }
285
286 fn begin_transaction<'a>(
287 &'a mut self,
288 ) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
289 Box::pin(async move {
290 if self.in_transaction {
291 return Err(DbError::Internal("transaction already started".to_string()));
292 }
293 self.execute("BEGIN").await?;
294 self.in_transaction = true;
295 Ok(())
296 })
297 }
298
299 fn commit<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
300 Box::pin(async move {
301 if self.in_transaction {
302 self.execute("COMMIT").await?;
303 self.in_transaction = false;
304 }
305 Ok(())
306 })
307 }
308
309 fn rollback<'a>(
310 &'a mut self,
311 ) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
312 Box::pin(async move {
313 if self.in_transaction {
314 let result = self.execute("ROLLBACK").await;
315 self.in_transaction = false;
316 result.map(|_| ())
317 } else {
318 Ok(())
319 }
320 })
321 }
322
323 fn is_connected(&self) -> bool {
324 self.connected
325 }
326
327 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
328 Box::pin(async move {
329 match self.execute("SELECT 1").await {
330 Ok(_) => true,
331 Err(_) => {
332 self.connected = false;
333 false
334 }
335 }
336 })
337 }
338
339 fn close<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
340 Box::pin(async move {
341 if let Some(conn) = self.conn.take() {
342 drop(conn);
343 }
344 self.connected = false;
345 self.in_transaction = false;
346 Ok(())
347 })
348 }
349
350 fn execute_with_params<'a>(
356 &'a mut self,
357 sql: &'a str,
358 params: &'a [Value],
359 ) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
360 Box::pin(async move {
361 if needs_raw_sql(sql) || params.is_empty() {
362 return self.execute(sql).await;
363 }
364 let mut pool_conn = self
365 .conn
366 .take()
367 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
368 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
369 for v in params {
370 q = match v {
371 Value::Null => q.bind(None::<i64>),
372 Value::Bool(b) => q.bind(*b),
373 Value::I8(n) => q.bind(*n),
374 Value::I16(n) => q.bind(*n),
375 Value::I32(n) => q.bind(*n),
376 Value::I64(n) => q.bind(*n),
377 Value::U8(n) => q.bind(*n),
378 Value::U16(n) => q.bind(*n),
379 Value::U32(n) => q.bind(*n),
380 Value::U64(n) => q.bind(*n as i64),
381 Value::F32(f) => q.bind(*f),
382 Value::F64(f) => q.bind(*f),
383 Value::String(s) => q.bind(s.as_str()),
384 Value::Decimal(s) => q.bind(s.as_str()),
385 Value::Bytes(b) => q.bind(b.as_slice()),
386 other => q.bind(other.to_string()),
387 };
388 }
389 let result = q.execute(&mut *pool_conn).await;
390 self.conn = Some(pool_conn);
391 match result {
392 Ok(r) => Ok(r.rows_affected()),
393 Err(e) => {
394 let db_err = map_sqlx_error(e);
395 if matches!(db_err, DbError::ConnectionError(_) | DbError::IoError(_)) {
396 self.connected = false;
397 }
398 Err(db_err)
399 }
400 }
401 })
402 }
403
404 fn query_with_params<'a>(
409 &'a mut self,
410 sql: &'a str,
411 params: &'a [Value],
412 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, DbError>> + Send + 'a>> {
413 Box::pin(async move {
414 if params.is_empty() {
415 return self.query(sql).await;
416 }
417 let mut pool_conn = self
418 .conn
419 .take()
420 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
421 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
422 for v in params {
423 q = match v {
424 Value::Null => q.bind(None::<i64>),
425 Value::Bool(b) => q.bind(*b),
426 Value::I8(n) => q.bind(*n),
427 Value::I16(n) => q.bind(*n),
428 Value::I32(n) => q.bind(*n),
429 Value::I64(n) => q.bind(*n),
430 Value::U8(n) => q.bind(*n),
431 Value::U16(n) => q.bind(*n),
432 Value::U32(n) => q.bind(*n),
433 Value::U64(n) => q.bind(*n as i64),
434 Value::F32(f) => q.bind(*f),
435 Value::F64(f) => q.bind(*f),
436 Value::String(s) => q.bind(s.as_str()),
437 Value::Decimal(s) => q.bind(s.as_str()),
438 Value::Bytes(b) => q.bind(b.as_slice()),
439 other => q.bind(other.to_string()),
440 };
441 }
442 let rows_result = q.fetch_all(&mut *pool_conn).await;
443 self.conn = Some(pool_conn);
444 let rows = rows_result.map_err(map_sqlx_error)?;
445 if rows.is_empty() {
446 return Ok(Vec::new());
447 }
448 let col_types: Vec<ColType> = rows[0]
451 .columns()
452 .iter()
453 .map(|col| ColType::parse_sqlite(col.type_info().name()))
454 .collect();
455 let mut result = Vec::with_capacity(rows.len());
456 for row in &rows {
457 let mut record = HashMap::with_capacity(col_types.len());
458 for (i, col) in row.columns().iter().enumerate() {
459 let name = col.name().to_string();
460 let value = row_to_value_with_coltype_sqlite(row, i, col_types[i]);
461 record.insert(name, value);
462 }
463 result.push(record);
464 }
465 Ok(result)
466 })
467 }
468
469 fn query_values<'a>(
475 &'a mut self,
476 sql: &'a str,
477 ) -> Pin<Box<dyn Future<Output = Result<QueryValues, DbError>> + Send + 'a>> {
478 Box::pin(async move {
479 let mut pool_conn = self
480 .conn
481 .take()
482 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
483 let rows_result = (&mut *pool_conn).fetch_all(sqlx::AssertSqlSafe(sql)).await;
484 self.conn = Some(pool_conn);
485 let rows = rows_result.map_err(map_sqlx_error)?;
486 if rows.is_empty() {
487 return Ok((Vec::new(), Vec::new()));
488 }
489 let cols = rows[0].columns();
490 let mut col_names: Vec<String> = Vec::with_capacity(cols.len());
491 let mut col_types: Vec<ColType> = Vec::with_capacity(cols.len());
492 for col in cols {
493 col_names.push(col.name().to_string());
494 col_types.push(ColType::parse_sqlite(col.type_info().name()));
495 }
496 let mut result_rows: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
497 for row in &rows {
498 let mut row_values: Vec<Value> = Vec::with_capacity(col_names.len());
499 for (idx, _) in col_names.iter().enumerate() {
500 row_values.push(row_to_value_with_coltype_sqlite(row, idx, col_types[idx]));
501 }
502 result_rows.push(row_values);
503 }
504 Ok((col_names, result_rows))
505 })
506 }
507
508 fn query_values_with_params<'a>(
512 &'a mut self,
513 sql: &'a str,
514 params: &'a [Value],
515 ) -> Pin<Box<dyn Future<Output = Result<QueryValues, DbError>> + Send + 'a>> {
516 Box::pin(async move {
517 if params.is_empty() {
518 return self.query_values(sql).await;
519 }
520 let mut pool_conn = self
521 .conn
522 .take()
523 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
524 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
525 for v in params {
526 q = match v {
527 Value::Null => q.bind(None::<i64>),
528 Value::Bool(b) => q.bind(*b),
529 Value::I8(n) => q.bind(*n),
530 Value::I16(n) => q.bind(*n),
531 Value::I32(n) => q.bind(*n),
532 Value::I64(n) => q.bind(*n),
533 Value::U8(n) => q.bind(*n),
534 Value::U16(n) => q.bind(*n),
535 Value::U32(n) => q.bind(*n),
536 Value::U64(n) => q.bind(*n as i64),
537 Value::F32(f) => q.bind(*f),
538 Value::F64(f) => q.bind(*f),
539 Value::String(s) => q.bind(s.as_str()),
540 Value::Decimal(s) => q.bind(s.as_str()),
541 Value::Bytes(b) => q.bind(b.as_slice()),
542 other => q.bind(other.to_string()),
543 };
544 }
545 let rows_result = q.fetch_all(&mut *pool_conn).await;
546 self.conn = Some(pool_conn);
547 let rows = rows_result.map_err(map_sqlx_error)?;
548 if rows.is_empty() {
549 return Ok((Vec::new(), Vec::new()));
550 }
551 let cols = rows[0].columns();
552 let mut col_names: Vec<String> = Vec::with_capacity(cols.len());
553 let mut col_types: Vec<ColType> = Vec::with_capacity(cols.len());
554 for col in cols {
555 col_names.push(col.name().to_string());
556 col_types.push(ColType::parse_sqlite(col.type_info().name()));
557 }
558 let mut result_rows: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
559 for row in &rows {
560 let mut row_values: Vec<Value> = Vec::with_capacity(col_names.len());
561 for (idx, _) in col_names.iter().enumerate() {
562 row_values.push(row_to_value_with_coltype_sqlite(row, idx, col_types[idx]));
563 }
564 result_rows.push(row_values);
565 }
566 Ok((col_names, result_rows))
567 })
568 }
569
570 fn query_stream<'a>(
576 &'a mut self,
577 sql: &'a str,
578 ) -> Pin<Box<dyn futures::Stream<Item = Result<HashMap<String, Value>, DbError>> + Send + 'a>>
579 {
580 Box::pin(async_stream::try_stream! {
581 let mut pool_conn = self
582 .conn
583 .take()
584 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
585 let mut row_stream = sqlx::query(sqlx::AssertSqlSafe(sql)).fetch(&mut *pool_conn);
586 let mut col_types: Vec<ColType> = Vec::new();
587 let mut col_names: Vec<String> = Vec::new();
588 let mut first_row = true;
589 while let Some(row_result) = row_stream.next().await {
590 let row = row_result.map_err(map_sqlx_error)?;
591 if first_row {
592 for col in row.columns() {
593 col_names.push(col.name().to_string());
594 col_types.push(ColType::parse_sqlite(col.type_info().name()));
595 }
596 first_row = false;
597 }
598 let mut record = HashMap::with_capacity(col_names.len());
599 for (i, name) in col_names.iter().enumerate() {
600 let value = row_to_value_with_coltype_sqlite(&row, i, col_types[i]);
601 record.insert(name.clone(), value);
602 }
603 yield record;
604 }
605 drop(row_stream);
608 self.conn = Some(pool_conn);
609 })
610 }
611}
612
613impl Drop for SqlxSqliteConnection {
614 fn drop(&mut self) {
615 if let Some(conn) = self.conn.take() {
616 drop(conn);
617 }
618 }
619}
620
621pub async fn sqlite_backup(
626 conn: &mut SqlxSqliteConnection,
627 dest_path: &str,
628) -> Result<(), DbError> {
629 let escaped_path = dest_path.replace('\'', "''");
630 let sql = format!("VACUUM INTO '{}'", escaped_path);
631 conn.execute(&sql).await?;
632 Ok(())
633}
634
635fn row_to_value_mysql(row: &sqlx::mysql::MySqlRow, ordinal: usize) -> Value {
641 use sqlx::TypeInfo;
642 let type_name = row.columns()[ordinal].type_info().name();
643 match type_name {
644 "BOOLEAN" => match row.try_get::<Option<bool>, usize>(ordinal) {
645 Ok(v) => v.map(Value::Bool).unwrap_or(Value::Null),
646 Err(_) => Value::Null,
647 },
648 "TINYINT" | "TINYINT UNSIGNED" => match row.try_get::<Option<i8>, usize>(ordinal) {
649 Ok(v) => v.map(Value::I8).unwrap_or(Value::Null),
650 Err(_) => match row.try_get::<Option<u8>, usize>(ordinal) {
651 Ok(v) => v.map(Value::U8).unwrap_or(Value::Null),
652 Err(_) => Value::Null,
653 },
654 },
655 "SMALLINT" | "SMALLINT UNSIGNED" => match row.try_get::<Option<i16>, usize>(ordinal) {
656 Ok(v) => v.map(Value::I16).unwrap_or(Value::Null),
657 Err(_) => match row.try_get::<Option<u16>, usize>(ordinal) {
658 Ok(v) => v.map(Value::U16).unwrap_or(Value::Null),
659 Err(_) => Value::Null,
660 },
661 },
662 "INT" | "INT UNSIGNED" | "MEDIUMINT" | "MEDIUMINT UNSIGNED" => {
663 match row.try_get::<Option<i32>, usize>(ordinal) {
664 Ok(v) => v.map(Value::I32).unwrap_or(Value::Null),
665 Err(_) => match row.try_get::<Option<u32>, usize>(ordinal) {
666 Ok(v) => v.map(Value::U32).unwrap_or(Value::Null),
667 Err(_) => Value::Null,
668 },
669 }
670 }
671 "BIGINT" | "BIGINT UNSIGNED" => match row.try_get::<Option<i64>, usize>(ordinal) {
672 Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
673 Err(_) => match row.try_get::<Option<u64>, usize>(ordinal) {
674 Ok(v) => v.map(Value::U64).unwrap_or(Value::Null),
675 Err(_) => Value::Null,
676 },
677 },
678 "FLOAT" => match row.try_get::<Option<f32>, usize>(ordinal) {
679 Ok(v) => v.map(Value::F32).unwrap_or(Value::Null),
680 Err(_) => Value::Null,
681 },
682 "DOUBLE" => match row.try_get::<Option<f64>, usize>(ordinal) {
683 Ok(v) => v.map(Value::F64).unwrap_or(Value::Null),
684 Err(_) => Value::Null,
685 },
686 "VARCHAR" | "TEXT" | "CHAR" | "TINYTEXT" | "MEDIUMTEXT" | "LONGTEXT" | "ENUM" => {
687 match row.try_get::<Option<String>, usize>(ordinal) {
688 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
689 Err(_) => Value::Null,
690 }
691 }
692 "BLOB" | "TINYBLOB" | "MEDIUMBLOB" | "LONGBLOB" | "BINARY" | "VARBINARY" => {
693 match row.try_get::<Option<Vec<u8>>, usize>(ordinal) {
694 Ok(v) => v.map(Value::Bytes).unwrap_or(Value::Null),
695 Err(_) => Value::Null,
696 }
697 }
698 "DECIMAL" | "NUMERIC" | "NEWDECIMAL" => {
700 match row.try_get::<Option<rust_decimal::Decimal>, usize>(ordinal) {
701 Ok(Some(v)) => Value::Decimal(v.to_string()),
702 Ok(None) => Value::Null,
703 Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
704 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
705 Err(_) => Value::Null,
706 },
707 }
708 }
709 _ => {
710 if let Ok(v) = row.try_get::<Option<i64>, usize>(ordinal) {
712 return v.map(Value::I64).unwrap_or(Value::Null);
713 }
714 if let Ok(v) = row.try_get::<Option<f64>, usize>(ordinal) {
715 return v.map(Value::F64).unwrap_or(Value::Null);
716 }
717 if let Ok(v) = row.try_get::<Option<bool>, usize>(ordinal) {
718 return v.map(Value::Bool).unwrap_or(Value::Null);
719 }
720 if let Ok(v) = row.try_get::<Option<String>, usize>(ordinal) {
721 return v.map(Value::String).unwrap_or(Value::Null);
722 }
723 Value::Null
724 }
725 }
726}
727
728fn row_to_value_with_coltype_mysql(
736 row: &sqlx::mysql::MySqlRow,
737 ordinal: usize,
738 col_type: ColType,
739) -> Value {
740 match col_type {
741 ColType::Bool => match row.try_get::<Option<bool>, usize>(ordinal) {
742 Ok(v) => v.map(Value::Bool).unwrap_or(Value::Null),
743 Err(_) => Value::Null,
744 },
745 ColType::I8 => match row.try_get::<Option<i8>, usize>(ordinal) {
746 Ok(v) => v.map(Value::I8).unwrap_or(Value::Null),
747 Err(_) => match row.try_get::<Option<u8>, usize>(ordinal) {
748 Ok(v) => v.map(Value::U8).unwrap_or(Value::Null),
749 Err(_) => Value::Null,
750 },
751 },
752 ColType::I16 => match row.try_get::<Option<i16>, usize>(ordinal) {
753 Ok(v) => v.map(Value::I16).unwrap_or(Value::Null),
754 Err(_) => match row.try_get::<Option<u16>, usize>(ordinal) {
755 Ok(v) => v.map(Value::U16).unwrap_or(Value::Null),
756 Err(_) => Value::Null,
757 },
758 },
759 ColType::I32 => match row.try_get::<Option<i32>, usize>(ordinal) {
760 Ok(v) => v.map(Value::I32).unwrap_or(Value::Null),
761 Err(_) => match row.try_get::<Option<u32>, usize>(ordinal) {
762 Ok(v) => v.map(Value::U32).unwrap_or(Value::Null),
763 Err(_) => Value::Null,
764 },
765 },
766 ColType::I64 => match row.try_get::<Option<i64>, usize>(ordinal) {
767 Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
768 Err(_) => match row.try_get::<Option<u64>, usize>(ordinal) {
769 Ok(v) => v.map(Value::U64).unwrap_or(Value::Null),
770 Err(_) => Value::Null,
771 },
772 },
773 ColType::U8 => 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 ColType::U16 => match row.try_get::<Option<u16>, usize>(ordinal) {
778 Ok(v) => v.map(Value::U16).unwrap_or(Value::Null),
779 Err(_) => Value::Null,
780 },
781 ColType::U32 => match row.try_get::<Option<u32>, usize>(ordinal) {
782 Ok(v) => v.map(Value::U32).unwrap_or(Value::Null),
783 Err(_) => Value::Null,
784 },
785 ColType::U64 => match row.try_get::<Option<u64>, usize>(ordinal) {
786 Ok(v) => v.map(Value::U64).unwrap_or(Value::Null),
787 Err(_) => Value::Null,
788 },
789 ColType::F32 => match row.try_get::<Option<f32>, usize>(ordinal) {
790 Ok(v) => v.map(Value::F32).unwrap_or(Value::Null),
791 Err(_) => Value::Null,
792 },
793 ColType::F64 => match row.try_get::<Option<f64>, usize>(ordinal) {
794 Ok(v) => v.map(Value::F64).unwrap_or(Value::Null),
795 Err(_) => Value::Null,
796 },
797 ColType::Decimal => match row.try_get::<Option<rust_decimal::Decimal>, usize>(ordinal) {
799 Ok(Some(v)) => Value::Decimal(v.to_string()),
800 Ok(None) => Value::Null,
801 Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
802 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
803 Err(_) => Value::Null,
804 },
805 },
806 ColType::String => match row.try_get::<Option<String>, usize>(ordinal) {
807 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
808 Err(_) => Value::Null,
809 },
810 ColType::Bytes => match row.try_get::<Option<Vec<u8>>, usize>(ordinal) {
811 Ok(v) => v.map(Value::Bytes).unwrap_or(Value::Null),
812 Err(_) => Value::Null,
813 },
814 ColType::Date | ColType::DateTime | ColType::Time | ColType::Json | ColType::Uuid => {
816 match row.try_get::<Option<String>, usize>(ordinal) {
817 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
818 Err(_) => Value::Null,
819 }
820 }
821 ColType::Unknown => {
822 if let Ok(v) = row.try_get::<Option<i64>, usize>(ordinal) {
824 return v.map(Value::I64).unwrap_or(Value::Null);
825 }
826 if let Ok(v) = row.try_get::<Option<f64>, usize>(ordinal) {
827 return v.map(Value::F64).unwrap_or(Value::Null);
828 }
829 if let Ok(v) = row.try_get::<Option<bool>, usize>(ordinal) {
830 return v.map(Value::Bool).unwrap_or(Value::Null);
831 }
832 if let Ok(v) = row.try_get::<Option<String>, usize>(ordinal) {
833 return v.map(Value::String).unwrap_or(Value::Null);
834 }
835 Value::Null
836 }
837 _ => Value::Null,
839 }
840}
841
842pub struct MySqlPoolHandle {
843 pool: sqlx::MySqlPool,
844}
845
846impl MySqlPoolHandle {
847 pub async fn connect(url: &str) -> Result<Self, DbError> {
848 let pool = sqlx::pool::PoolOptions::<sqlx::MySql>::new()
849 .max_connections(10)
850 .acquire_timeout(std::time::Duration::from_secs(30))
851 .idle_timeout(Some(std::time::Duration::from_secs(600)))
852 .max_lifetime(Some(std::time::Duration::from_secs(1800)))
853 .connect(url)
854 .await
855 .map_err(map_sqlx_error)?;
856 Ok(Self { pool })
857 }
858
859 pub fn from_pool(pool: sqlx::MySqlPool) -> Self {
860 Self { pool }
861 }
862
863 pub fn pool(&self) -> &sqlx::MySqlPool {
864 &self.pool
865 }
866}
867
868pub struct SqlxMySqlConnectionFactory {
869 pool: Arc<MySqlPoolHandle>,
870}
871
872impl SqlxMySqlConnectionFactory {
873 pub fn new(pool: Arc<MySqlPoolHandle>) -> Self {
874 Self { pool }
875 }
876}
877
878#[async_trait]
879impl ConnectionFactory for SqlxMySqlConnectionFactory {
880 async fn create(&self) -> Result<Box<dyn Connection>, DbError> {
881 let conn = self.pool.pool.acquire().await.map_err(map_sqlx_error)?;
882 Ok(Box::new(SqlxMySqlConnection {
883 conn: Some(conn),
884 connected: true,
885 in_transaction: false,
886 }))
887 }
888}
889
890pub struct SqlxMySqlConnection {
891 conn: Option<sqlx::pool::PoolConnection<sqlx::MySql>>,
892 connected: bool,
893 in_transaction: bool,
894}
895
896impl Connection for SqlxMySqlConnection {
897 fn execute<'a>(
898 &'a mut self,
899 sql: &'a str,
900 ) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
901 Box::pin(async move {
902 let mut pool_conn = self
903 .conn
904 .take()
905 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
906 let result = if needs_raw_sql(sql) {
909 (&mut *pool_conn)
910 .execute(sqlx::raw_sql(sqlx::AssertSqlSafe(sql)))
911 .await
912 } else {
913 (&mut *pool_conn).execute(sqlx::AssertSqlSafe(sql)).await
914 };
915 self.conn = Some(pool_conn);
916
917 match result {
918 Ok(r) => Ok(r.rows_affected()),
919 Err(e) => {
920 let db_err = map_sqlx_error(e);
921 if matches!(db_err, DbError::ConnectionError(_) | DbError::IoError(_)) {
922 self.connected = false;
923 }
924 Err(db_err)
925 }
926 }
927 })
928 }
929
930 fn query<'a>(
931 &'a mut self,
932 sql: &'a str,
933 ) -> Pin<Box<dyn Future<Output = Result<Vec<HashMap<String, Value>>, DbError>> + Send + 'a>>
934 {
935 Box::pin(async move {
936 let mut pool_conn = self
937 .conn
938 .take()
939 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
940 let rows_result = (&mut *pool_conn).fetch_all(sqlx::AssertSqlSafe(sql)).await;
943 self.conn = Some(pool_conn);
944
945 let rows = rows_result.map_err(map_sqlx_error)?;
946 if rows.is_empty() {
947 return Ok(Vec::new());
948 }
949 let col_types: Vec<ColType> = rows[0]
952 .columns()
953 .iter()
954 .map(|col| ColType::parse_mysql(col.type_info().name()))
955 .collect();
956 let mut result = Vec::with_capacity(rows.len());
957 for row in &rows {
958 let mut record = HashMap::with_capacity(col_types.len());
959 for (i, col) in row.columns().iter().enumerate() {
960 let name = col.name().to_string();
961 let value = row_to_value_with_coltype_mysql(row, i, col_types[i]);
962 record.insert(name, value);
963 }
964 result.push(record);
965 }
966 Ok(result)
967 })
968 }
969
970 fn begin_transaction<'a>(
971 &'a mut self,
972 ) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
973 Box::pin(async move {
974 if self.in_transaction {
975 return Err(DbError::Internal("transaction already started".to_string()));
976 }
977 self.execute("BEGIN").await?;
978 self.in_transaction = true;
979 Ok(())
980 })
981 }
982
983 fn commit<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
984 Box::pin(async move {
985 if self.in_transaction {
986 self.execute("COMMIT").await?;
987 self.in_transaction = false;
988 }
989 Ok(())
990 })
991 }
992
993 fn rollback<'a>(
994 &'a mut self,
995 ) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
996 Box::pin(async move {
997 if self.in_transaction {
998 let result = self.execute("ROLLBACK").await;
999 self.in_transaction = false;
1000 result.map(|_| ())
1001 } else {
1002 Ok(())
1003 }
1004 })
1005 }
1006
1007 fn is_connected(&self) -> bool {
1008 self.connected
1009 }
1010
1011 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
1012 Box::pin(async move {
1013 match self.execute("SELECT 1").await {
1014 Ok(_) => true,
1015 Err(_) => {
1016 self.connected = false;
1017 false
1018 }
1019 }
1020 })
1021 }
1022
1023 fn close<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
1024 Box::pin(async move {
1025 if let Some(conn) = self.conn.take() {
1026 drop(conn);
1027 }
1028 self.connected = false;
1029 self.in_transaction = false;
1030 Ok(())
1031 })
1032 }
1033
1034 fn execute_with_params<'a>(
1040 &'a mut self,
1041 sql: &'a str,
1042 params: &'a [Value],
1043 ) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
1044 Box::pin(async move {
1045 if needs_raw_sql(sql) || params.is_empty() {
1046 return self.execute(sql).await;
1047 }
1048 let mut pool_conn = self
1049 .conn
1050 .take()
1051 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1052 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
1053 for v in params {
1054 q = match v {
1055 Value::Null => q.bind(None::<i64>),
1056 Value::Bool(b) => q.bind(*b),
1057 Value::I8(n) => q.bind(*n),
1058 Value::I16(n) => q.bind(*n),
1059 Value::I32(n) => q.bind(*n),
1060 Value::I64(n) => q.bind(*n),
1061 Value::U8(n) => q.bind(*n),
1062 Value::U16(n) => q.bind(*n),
1063 Value::U32(n) => q.bind(*n),
1064 Value::U64(n) => q.bind(*n as i64),
1065 Value::F32(f) => q.bind(*f),
1066 Value::F64(f) => q.bind(*f),
1067 Value::String(s) => q.bind(s.as_str()),
1068 Value::Decimal(s) => q.bind(s.as_str()),
1069 Value::Bytes(b) => q.bind(b.as_slice()),
1070 other => q.bind(other.to_string()),
1071 };
1072 }
1073 let result = q.execute(&mut *pool_conn).await;
1074 self.conn = Some(pool_conn);
1075 match result {
1076 Ok(r) => Ok(r.rows_affected()),
1077 Err(e) => {
1078 let db_err = map_sqlx_error(e);
1079 if matches!(db_err, DbError::ConnectionError(_) | DbError::IoError(_)) {
1080 self.connected = false;
1081 }
1082 Err(db_err)
1083 }
1084 }
1085 })
1086 }
1087
1088 fn query_with_params<'a>(
1090 &'a mut self,
1091 sql: &'a str,
1092 params: &'a [Value],
1093 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, DbError>> + Send + 'a>> {
1094 Box::pin(async move {
1095 if params.is_empty() {
1096 return self.query(sql).await;
1097 }
1098 let mut pool_conn = self
1099 .conn
1100 .take()
1101 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1102 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
1103 for v in params {
1104 q = match v {
1105 Value::Null => q.bind(None::<i64>),
1106 Value::Bool(b) => q.bind(*b),
1107 Value::I8(n) => q.bind(*n),
1108 Value::I16(n) => q.bind(*n),
1109 Value::I32(n) => q.bind(*n),
1110 Value::I64(n) => q.bind(*n),
1111 Value::U8(n) => q.bind(*n),
1112 Value::U16(n) => q.bind(*n),
1113 Value::U32(n) => q.bind(*n),
1114 Value::U64(n) => q.bind(*n as i64),
1115 Value::F32(f) => q.bind(*f),
1116 Value::F64(f) => q.bind(*f),
1117 Value::String(s) => q.bind(s.as_str()),
1118 Value::Decimal(s) => q.bind(s.as_str()),
1119 Value::Bytes(b) => q.bind(b.as_slice()),
1120 other => q.bind(other.to_string()),
1121 };
1122 }
1123 let rows_result = q.fetch_all(&mut *pool_conn).await;
1124 self.conn = Some(pool_conn);
1125 let rows = rows_result.map_err(map_sqlx_error)?;
1126 let mut result = Vec::with_capacity(rows.len());
1127 for row in rows {
1128 let columns = row.columns();
1130 let mut record = HashMap::with_capacity(columns.len());
1131 for col in columns {
1132 let name = col.name().to_string();
1133 let ordinal = col.ordinal();
1134 let value = row_to_value_mysql(&row, ordinal);
1135 record.insert(name, value);
1136 }
1137 result.push(record);
1138 }
1139 Ok(result)
1140 })
1141 }
1142
1143 fn query_values<'a>(
1147 &'a mut self,
1148 sql: &'a str,
1149 ) -> Pin<Box<dyn Future<Output = Result<QueryValues, DbError>> + Send + 'a>> {
1150 Box::pin(async move {
1151 let mut pool_conn = self
1152 .conn
1153 .take()
1154 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1155 let rows_result = (&mut *pool_conn).fetch_all(sqlx::AssertSqlSafe(sql)).await;
1156 self.conn = Some(pool_conn);
1157 let rows = rows_result.map_err(map_sqlx_error)?;
1158 if rows.is_empty() {
1159 return Ok((Vec::new(), Vec::new()));
1160 }
1161 let cols = rows[0].columns();
1162 let mut col_names: Vec<String> = Vec::with_capacity(cols.len());
1163 let mut col_types: Vec<ColType> = Vec::with_capacity(cols.len());
1164 for col in cols {
1165 col_names.push(col.name().to_string());
1166 col_types.push(ColType::parse_mysql(col.type_info().name()));
1167 }
1168 let mut result_rows: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
1169 for row in &rows {
1170 let mut row_values: Vec<Value> = Vec::with_capacity(col_names.len());
1171 for (idx, _) in col_names.iter().enumerate() {
1172 row_values.push(row_to_value_with_coltype_mysql(row, idx, col_types[idx]));
1173 }
1174 result_rows.push(row_values);
1175 }
1176 Ok((col_names, result_rows))
1177 })
1178 }
1179
1180 fn query_values_with_params<'a>(
1182 &'a mut self,
1183 sql: &'a str,
1184 params: &'a [Value],
1185 ) -> Pin<Box<dyn Future<Output = Result<QueryValues, DbError>> + Send + 'a>> {
1186 Box::pin(async move {
1187 if params.is_empty() {
1188 return self.query_values(sql).await;
1189 }
1190 let mut pool_conn = self
1191 .conn
1192 .take()
1193 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1194 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
1195 for v in params {
1196 q = match v {
1197 Value::Null => q.bind(None::<i64>),
1198 Value::Bool(b) => q.bind(*b),
1199 Value::I8(n) => q.bind(*n),
1200 Value::I16(n) => q.bind(*n),
1201 Value::I32(n) => q.bind(*n),
1202 Value::I64(n) => q.bind(*n),
1203 Value::U8(n) => q.bind(*n),
1204 Value::U16(n) => q.bind(*n),
1205 Value::U32(n) => q.bind(*n),
1206 Value::U64(n) => q.bind(*n as i64),
1207 Value::F32(f) => q.bind(*f),
1208 Value::F64(f) => q.bind(*f),
1209 Value::String(s) => q.bind(s.as_str()),
1210 Value::Decimal(s) => q.bind(s.as_str()),
1211 Value::Bytes(b) => q.bind(b.as_slice()),
1212 other => q.bind(other.to_string()),
1213 };
1214 }
1215 let rows_result = q.fetch_all(&mut *pool_conn).await;
1216 self.conn = Some(pool_conn);
1217 let rows = rows_result.map_err(map_sqlx_error)?;
1218 if rows.is_empty() {
1219 return Ok((Vec::new(), Vec::new()));
1220 }
1221 let cols = rows[0].columns();
1222 let mut col_names: Vec<String> = Vec::with_capacity(cols.len());
1223 let mut col_types: Vec<ColType> = Vec::with_capacity(cols.len());
1224 for col in cols {
1225 col_names.push(col.name().to_string());
1226 col_types.push(ColType::parse_mysql(col.type_info().name()));
1227 }
1228 let mut result_rows: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
1229 for row in &rows {
1230 let mut row_values: Vec<Value> = Vec::with_capacity(col_names.len());
1231 for (idx, _) in col_names.iter().enumerate() {
1232 row_values.push(row_to_value_with_coltype_mysql(row, idx, col_types[idx]));
1233 }
1234 result_rows.push(row_values);
1235 }
1236 Ok((col_names, result_rows))
1237 })
1238 }
1239
1240 fn query_stream<'a>(
1246 &'a mut self,
1247 sql: &'a str,
1248 ) -> Pin<Box<dyn futures::Stream<Item = Result<HashMap<String, Value>, DbError>> + Send + 'a>>
1249 {
1250 Box::pin(async_stream::try_stream! {
1251 let mut pool_conn = self
1252 .conn
1253 .take()
1254 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1255 let mut row_stream = sqlx::query(sqlx::AssertSqlSafe(sql)).fetch(&mut *pool_conn);
1256 let mut col_types: Vec<ColType> = Vec::new();
1257 let mut col_names: Vec<String> = Vec::new();
1258 let mut first_row = true;
1259 while let Some(row_result) = row_stream.next().await {
1260 let row = row_result.map_err(map_sqlx_error)?;
1261 if first_row {
1262 for col in row.columns() {
1263 col_names.push(col.name().to_string());
1264 col_types.push(ColType::parse_mysql(col.type_info().name()));
1265 }
1266 first_row = false;
1267 }
1268 let mut record = HashMap::with_capacity(col_names.len());
1269 for (i, name) in col_names.iter().enumerate() {
1270 let value = row_to_value_with_coltype_mysql(&row, i, col_types[i]);
1271 record.insert(name.clone(), value);
1272 }
1273 yield record;
1274 }
1275 drop(row_stream);
1278 self.conn = Some(pool_conn);
1279 })
1280 }
1281}
1282
1283impl Drop for SqlxMySqlConnection {
1284 fn drop(&mut self) {
1285 if let Some(conn) = self.conn.take() {
1286 drop(conn);
1287 }
1288 }
1289}
1290
1291pub async fn mysql_bulk_insert(
1298 conn: &mut SqlxMySqlConnection,
1299 table: &str,
1300 columns: &[&str],
1301 rows: &[Vec<Value>],
1302) -> Result<u64, DbError> {
1303 if rows.is_empty() {
1304 return Ok(0);
1305 }
1306 let col_list = columns.join(", ");
1307 let cols_per_row = columns.len();
1308 let row_placeholder = format!("({})", vec!["?"; cols_per_row].join(", "));
1310 let placeholders = vec![row_placeholder; rows.len()].join(", ");
1311 let sql = format!(
1312 "INSERT INTO {} ({}) VALUES {}",
1313 table, col_list, placeholders
1314 );
1315 let mut params: Vec<Value> = Vec::with_capacity(rows.len() * cols_per_row);
1317 for row in rows {
1318 for v in row {
1319 params.push(v.clone());
1320 }
1321 }
1322 conn.execute_with_params(&sql, ¶ms).await
1323}
1324
1325fn row_to_value_pg(row: &sqlx::postgres::PgRow, ordinal: usize) -> Value {
1331 use sqlx::TypeInfo;
1332 let type_name = row.columns()[ordinal].type_info().name();
1333 match type_name {
1334 "BOOL" => match row.try_get::<Option<bool>, usize>(ordinal) {
1335 Ok(v) => v.map(Value::Bool).unwrap_or(Value::Null),
1336 Err(_) => Value::Null,
1337 },
1338 "INT2" => match row.try_get::<Option<i16>, usize>(ordinal) {
1339 Ok(v) => v.map(Value::I16).unwrap_or(Value::Null),
1340 Err(_) => Value::Null,
1341 },
1342 "INT4" | "OID" => match row.try_get::<Option<i32>, usize>(ordinal) {
1343 Ok(v) => v.map(Value::I32).unwrap_or(Value::Null),
1344 Err(_) => Value::Null,
1345 },
1346 "INT8" => match row.try_get::<Option<i64>, usize>(ordinal) {
1347 Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
1348 Err(_) => Value::Null,
1349 },
1350 "FLOAT4" => match row.try_get::<Option<f32>, usize>(ordinal) {
1351 Ok(v) => v.map(Value::F32).unwrap_or(Value::Null),
1352 Err(_) => Value::Null,
1353 },
1354 "FLOAT8" => match row.try_get::<Option<f64>, usize>(ordinal) {
1355 Ok(v) => v.map(Value::F64).unwrap_or(Value::Null),
1356 Err(_) => Value::Null,
1357 },
1358 "TEXT" | "VARCHAR" | "CHAR" | "NAME" => match row.try_get::<Option<String>, usize>(ordinal)
1359 {
1360 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
1361 Err(_) => Value::Null,
1362 },
1363 "BYTEA" => match row.try_get::<Option<Vec<u8>>, usize>(ordinal) {
1364 Ok(v) => v.map(Value::Bytes).unwrap_or(Value::Null),
1365 Err(_) => Value::Null,
1366 },
1367 "NUMERIC" => match row.try_get::<Option<rust_decimal::Decimal>, usize>(ordinal) {
1368 Ok(Some(v)) => Value::Decimal(v.to_string()),
1369 Ok(None) => Value::Null,
1370 Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
1371 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
1372 Err(_) => Value::Null,
1373 },
1374 },
1375 "UUID" => match row.try_get::<Option<sqlx::types::Uuid>, usize>(ordinal) {
1377 Ok(v) => v
1378 .map(|uuid| Value::String(uuid.to_string()))
1379 .unwrap_or(Value::Null),
1380 Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
1381 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
1382 Err(_) => Value::Null,
1383 },
1384 },
1385 _ => {
1386 if let Ok(v) = row.try_get::<Option<i64>, usize>(ordinal) {
1388 return v.map(Value::I64).unwrap_or(Value::Null);
1389 }
1390 if let Ok(v) = row.try_get::<Option<f64>, usize>(ordinal) {
1391 return v.map(Value::F64).unwrap_or(Value::Null);
1392 }
1393 if let Ok(v) = row.try_get::<Option<bool>, usize>(ordinal) {
1394 return v.map(Value::Bool).unwrap_or(Value::Null);
1395 }
1396 if let Ok(v) = row.try_get::<Option<String>, usize>(ordinal) {
1397 return v.map(Value::String).unwrap_or(Value::Null);
1398 }
1399 Value::Null
1400 }
1401 }
1402}
1403
1404fn row_to_value_with_coltype_pg(
1412 row: &sqlx::postgres::PgRow,
1413 ordinal: usize,
1414 col_type: ColType,
1415) -> Value {
1416 match col_type {
1417 ColType::Bool => match row.try_get::<Option<bool>, usize>(ordinal) {
1418 Ok(v) => v.map(Value::Bool).unwrap_or(Value::Null),
1419 Err(_) => Value::Null,
1420 },
1421 ColType::I8 => match row.try_get::<Option<i8>, usize>(ordinal) {
1422 Ok(v) => v.map(Value::I8).unwrap_or(Value::Null),
1423 Err(_) => Value::Null,
1424 },
1425 ColType::I16 => match row.try_get::<Option<i16>, usize>(ordinal) {
1426 Ok(v) => v.map(Value::I16).unwrap_or(Value::Null),
1427 Err(_) => Value::Null,
1428 },
1429 ColType::I32 => match row.try_get::<Option<i32>, usize>(ordinal) {
1431 Ok(v) => v.map(Value::I32).unwrap_or(Value::Null),
1432 Err(_) => Value::Null,
1433 },
1434 ColType::I64 => match row.try_get::<Option<i64>, usize>(ordinal) {
1435 Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
1436 Err(_) => Value::Null,
1437 },
1438 ColType::U8 => match row.try_get::<Option<i16>, usize>(ordinal) {
1443 Ok(v) => v.map(Value::I16).unwrap_or(Value::Null),
1444 Err(_) => Value::Null,
1445 },
1446 ColType::U16 => 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::U32 => 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::U64 => match row.try_get::<Option<i64>, usize>(ordinal) {
1456 Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
1457 Err(_) => Value::Null,
1458 },
1459 ColType::F32 => match row.try_get::<Option<f32>, usize>(ordinal) {
1460 Ok(v) => v.map(Value::F32).unwrap_or(Value::Null),
1461 Err(_) => Value::Null,
1462 },
1463 ColType::F64 => match row.try_get::<Option<f64>, usize>(ordinal) {
1464 Ok(v) => v.map(Value::F64).unwrap_or(Value::Null),
1465 Err(_) => Value::Null,
1466 },
1467 ColType::Decimal => match row.try_get::<Option<rust_decimal::Decimal>, usize>(ordinal) {
1469 Ok(Some(v)) => Value::Decimal(v.to_string()),
1470 Ok(None) => Value::Null,
1471 Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
1472 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
1473 Err(_) => Value::Null,
1474 },
1475 },
1476 ColType::String => match row.try_get::<Option<String>, usize>(ordinal) {
1477 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
1478 Err(_) => Value::Null,
1479 },
1480 ColType::Bytes => match row.try_get::<Option<Vec<u8>>, usize>(ordinal) {
1481 Ok(v) => v.map(Value::Bytes).unwrap_or(Value::Null),
1482 Err(_) => Value::Null,
1483 },
1484 ColType::Date | ColType::DateTime | ColType::Time | ColType::Json => {
1486 match row.try_get::<Option<String>, usize>(ordinal) {
1487 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
1488 Err(_) => Value::Null,
1489 }
1490 }
1491 ColType::Uuid => match row.try_get::<Option<sqlx::types::Uuid>, usize>(ordinal) {
1493 Ok(v) => v
1494 .map(|uuid| Value::String(uuid.to_string()))
1495 .unwrap_or(Value::Null),
1496 Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
1497 Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
1498 Err(_) => Value::Null,
1499 },
1500 },
1501 ColType::Unknown => {
1502 if let Ok(v) = row.try_get::<Option<i64>, usize>(ordinal) {
1504 return v.map(Value::I64).unwrap_or(Value::Null);
1505 }
1506 if let Ok(v) = row.try_get::<Option<f64>, usize>(ordinal) {
1507 return v.map(Value::F64).unwrap_or(Value::Null);
1508 }
1509 if let Ok(v) = row.try_get::<Option<bool>, usize>(ordinal) {
1510 return v.map(Value::Bool).unwrap_or(Value::Null);
1511 }
1512 if let Ok(v) = row.try_get::<Option<String>, usize>(ordinal) {
1513 return v.map(Value::String).unwrap_or(Value::Null);
1514 }
1515 Value::Null
1516 }
1517 _ => Value::Null,
1519 }
1520}
1521
1522pub struct PgPoolHandle {
1523 pool: sqlx::PgPool,
1524}
1525
1526impl PgPoolHandle {
1527 pub async fn connect(url: &str) -> Result<Self, DbError> {
1528 let pool = sqlx::pool::PoolOptions::<sqlx::Postgres>::new()
1529 .max_connections(10)
1530 .acquire_timeout(std::time::Duration::from_secs(30))
1531 .idle_timeout(Some(std::time::Duration::from_secs(600)))
1532 .max_lifetime(Some(std::time::Duration::from_secs(1800)))
1533 .connect(url)
1534 .await
1535 .map_err(map_sqlx_error)?;
1536 Ok(Self { pool })
1537 }
1538
1539 pub fn from_pool(pool: sqlx::PgPool) -> Self {
1540 Self { pool }
1541 }
1542
1543 pub fn pool(&self) -> &sqlx::PgPool {
1544 &self.pool
1545 }
1546}
1547
1548pub struct SqlxPgConnectionFactory {
1549 pool: Arc<PgPoolHandle>,
1550}
1551
1552impl SqlxPgConnectionFactory {
1553 pub fn new(pool: Arc<PgPoolHandle>) -> Self {
1554 Self { pool }
1555 }
1556}
1557
1558#[async_trait]
1559impl ConnectionFactory for SqlxPgConnectionFactory {
1560 async fn create(&self) -> Result<Box<dyn Connection>, DbError> {
1561 let conn = self.pool.pool.acquire().await.map_err(map_sqlx_error)?;
1562 Ok(Box::new(SqlxPgConnection {
1563 conn: Some(conn),
1564 connected: true,
1565 in_transaction: false,
1566 }))
1567 }
1568}
1569
1570pub struct SqlxPgConnection {
1571 conn: Option<sqlx::pool::PoolConnection<sqlx::Postgres>>,
1572 connected: bool,
1573 in_transaction: bool,
1574}
1575
1576impl Connection for SqlxPgConnection {
1577 fn execute<'a>(
1578 &'a mut self,
1579 sql: &'a str,
1580 ) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
1581 Box::pin(async move {
1582 let mut pool_conn = self
1583 .conn
1584 .take()
1585 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1586 let result = if needs_raw_sql(sql) {
1589 (&mut *pool_conn)
1590 .execute(sqlx::raw_sql(sqlx::AssertSqlSafe(sql)))
1591 .await
1592 } else {
1593 (&mut *pool_conn).execute(sqlx::AssertSqlSafe(sql)).await
1594 };
1595 self.conn = Some(pool_conn);
1596
1597 match result {
1598 Ok(r) => Ok(r.rows_affected()),
1599 Err(e) => {
1600 let db_err = map_sqlx_error(e);
1601 if matches!(db_err, DbError::ConnectionError(_) | DbError::IoError(_)) {
1602 self.connected = false;
1603 }
1604 Err(db_err)
1605 }
1606 }
1607 })
1608 }
1609
1610 fn query<'a>(
1611 &'a mut self,
1612 sql: &'a str,
1613 ) -> Pin<Box<dyn Future<Output = Result<Vec<HashMap<String, Value>>, DbError>> + Send + 'a>>
1614 {
1615 Box::pin(async move {
1616 let mut pool_conn = self
1617 .conn
1618 .take()
1619 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1620 let rows_result = (&mut *pool_conn).fetch_all(sqlx::AssertSqlSafe(sql)).await;
1623 self.conn = Some(pool_conn);
1624
1625 let rows = rows_result.map_err(map_sqlx_error)?;
1626 if rows.is_empty() {
1627 return Ok(Vec::new());
1628 }
1629 let col_types: Vec<ColType> = rows[0]
1632 .columns()
1633 .iter()
1634 .map(|col| ColType::parse_postgres(col.type_info().name()))
1635 .collect();
1636 let mut result = Vec::with_capacity(rows.len());
1637 for row in &rows {
1638 let mut record = HashMap::with_capacity(col_types.len());
1639 for (i, col) in row.columns().iter().enumerate() {
1640 let name = col.name().to_string();
1641 let value = row_to_value_with_coltype_pg(row, i, col_types[i]);
1642 record.insert(name, value);
1643 }
1644 result.push(record);
1645 }
1646 Ok(result)
1647 })
1648 }
1649
1650 fn begin_transaction<'a>(
1651 &'a mut self,
1652 ) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
1653 Box::pin(async move {
1654 if self.in_transaction {
1655 return Err(DbError::Internal("transaction already started".to_string()));
1656 }
1657 self.execute("BEGIN").await?;
1658 self.in_transaction = true;
1659 Ok(())
1660 })
1661 }
1662
1663 fn commit<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
1664 Box::pin(async move {
1665 if self.in_transaction {
1666 self.execute("COMMIT").await?;
1667 self.in_transaction = false;
1668 }
1669 Ok(())
1670 })
1671 }
1672
1673 fn rollback<'a>(
1674 &'a mut self,
1675 ) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
1676 Box::pin(async move {
1677 if self.in_transaction {
1678 let result = self.execute("ROLLBACK").await;
1679 self.in_transaction = false;
1680 result.map(|_| ())
1681 } else {
1682 Ok(())
1683 }
1684 })
1685 }
1686
1687 fn is_connected(&self) -> bool {
1688 self.connected
1689 }
1690
1691 fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
1692 Box::pin(async move {
1693 match self.execute("SELECT 1").await {
1694 Ok(_) => true,
1695 Err(_) => {
1696 self.connected = false;
1697 false
1698 }
1699 }
1700 })
1701 }
1702
1703 fn close<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
1704 Box::pin(async move {
1705 if let Some(conn) = self.conn.take() {
1706 drop(conn);
1707 }
1708 self.connected = false;
1709 self.in_transaction = false;
1710 Ok(())
1711 })
1712 }
1713
1714 fn execute_with_params<'a>(
1729 &'a mut self,
1730 sql: &'a str,
1731 params: &'a [Value],
1732 ) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
1733 Box::pin(async move {
1734 if needs_raw_sql(sql) || params.is_empty() {
1735 return self.execute(sql).await;
1736 }
1737 let mut pool_conn = self
1738 .conn
1739 .take()
1740 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1741 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
1742 for v in params {
1743 q = match v {
1744 Value::Null => q.bind(None::<i64>),
1745 Value::Bool(b) => q.bind(*b),
1746 Value::I8(n) => q.bind(*n),
1747 Value::I16(n) => q.bind(*n),
1748 Value::I32(n) => q.bind(*n),
1749 Value::I64(n) => q.bind(*n),
1750 Value::U8(n) => q.bind(*n as i16),
1752 Value::U16(n) => q.bind(*n as i32),
1753 Value::U32(n) => q.bind(*n as i64),
1754 Value::U64(n) => q.bind(*n as i64),
1755 Value::F32(f) => q.bind(*f),
1756 Value::F64(f) => q.bind(*f),
1757 Value::String(s) => q.bind(s.as_str()),
1758 Value::Decimal(s) => q.bind(s.as_str()),
1759 Value::Bytes(b) => q.bind(b.as_slice()),
1760 other => q.bind(other.to_string()),
1761 };
1762 }
1763 let result = q.execute(&mut *pool_conn).await;
1764 self.conn = Some(pool_conn);
1765 match result {
1766 Ok(r) => Ok(r.rows_affected()),
1767 Err(e) => {
1768 let db_err = map_sqlx_error(e);
1769 if matches!(db_err, DbError::ConnectionError(_) | DbError::IoError(_)) {
1770 self.connected = false;
1771 }
1772 Err(db_err)
1773 }
1774 }
1775 })
1776 }
1777
1778 fn query_with_params<'a>(
1782 &'a mut self,
1783 sql: &'a str,
1784 params: &'a [Value],
1785 ) -> Pin<Box<dyn Future<Output = Result<QueryRows, DbError>> + Send + 'a>> {
1786 Box::pin(async move {
1787 if params.is_empty() {
1788 return self.query(sql).await;
1789 }
1790 let mut pool_conn = self
1791 .conn
1792 .take()
1793 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1794 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
1795 for v in params {
1796 q = match v {
1797 Value::Null => q.bind(None::<i64>),
1798 Value::Bool(b) => q.bind(*b),
1799 Value::I8(n) => q.bind(*n),
1800 Value::I16(n) => q.bind(*n),
1801 Value::I32(n) => q.bind(*n),
1802 Value::I64(n) => q.bind(*n),
1803 Value::U8(n) => q.bind(*n as i16),
1804 Value::U16(n) => q.bind(*n as i32),
1805 Value::U32(n) => q.bind(*n as i64),
1806 Value::U64(n) => q.bind(*n as i64),
1807 Value::F32(f) => q.bind(*f),
1808 Value::F64(f) => q.bind(*f),
1809 Value::String(s) => q.bind(s.as_str()),
1810 Value::Decimal(s) => q.bind(s.as_str()),
1811 Value::Bytes(b) => q.bind(b.as_slice()),
1812 other => q.bind(other.to_string()),
1813 };
1814 }
1815 let rows_result = q.fetch_all(&mut *pool_conn).await;
1816 self.conn = Some(pool_conn);
1817 let rows = rows_result.map_err(map_sqlx_error)?;
1818 if rows.is_empty() {
1819 return Ok(Vec::new());
1820 }
1821 let col_types: Vec<ColType> = rows[0]
1823 .columns()
1824 .iter()
1825 .map(|col| ColType::parse_postgres(col.type_info().name()))
1826 .collect();
1827 let mut result = Vec::with_capacity(rows.len());
1828 for row in &rows {
1829 let mut record = HashMap::with_capacity(col_types.len());
1830 for (i, col) in row.columns().iter().enumerate() {
1831 let name = col.name().to_string();
1832 let value = row_to_value_with_coltype_pg(row, i, col_types[i]);
1833 record.insert(name, value);
1834 }
1835 result.push(record);
1836 }
1837 Ok(result)
1838 })
1839 }
1840
1841 fn query_values<'a>(
1845 &'a mut self,
1846 sql: &'a str,
1847 ) -> Pin<Box<dyn Future<Output = Result<QueryValues, DbError>> + Send + 'a>> {
1848 Box::pin(async move {
1849 let mut pool_conn = self
1850 .conn
1851 .take()
1852 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1853 let rows_result = (&mut *pool_conn).fetch_all(sqlx::AssertSqlSafe(sql)).await;
1854 self.conn = Some(pool_conn);
1855 let rows = rows_result.map_err(map_sqlx_error)?;
1856 if rows.is_empty() {
1857 return Ok((Vec::new(), Vec::new()));
1858 }
1859 let cols = rows[0].columns();
1860 let mut col_names: Vec<String> = Vec::with_capacity(cols.len());
1861 let mut col_types: Vec<ColType> = Vec::with_capacity(cols.len());
1862 for col in cols {
1863 col_names.push(col.name().to_string());
1864 col_types.push(ColType::parse_postgres(col.type_info().name()));
1865 }
1866 let mut result_rows: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
1867 for row in &rows {
1868 let mut row_values: Vec<Value> = Vec::with_capacity(col_names.len());
1869 for (idx, _) in col_names.iter().enumerate() {
1870 row_values.push(row_to_value_with_coltype_pg(row, idx, col_types[idx]));
1871 }
1872 result_rows.push(row_values);
1873 }
1874 Ok((col_names, result_rows))
1875 })
1876 }
1877
1878 fn query_values_with_params<'a>(
1882 &'a mut self,
1883 sql: &'a str,
1884 params: &'a [Value],
1885 ) -> Pin<Box<dyn Future<Output = Result<QueryValues, DbError>> + Send + 'a>> {
1886 Box::pin(async move {
1887 if params.is_empty() {
1888 return self.query_values(sql).await;
1889 }
1890 let mut pool_conn = self
1891 .conn
1892 .take()
1893 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1894 let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
1895 for v in params {
1896 q = match v {
1897 Value::Null => q.bind(None::<i64>),
1898 Value::Bool(b) => q.bind(*b),
1899 Value::I8(n) => q.bind(*n),
1900 Value::I16(n) => q.bind(*n),
1901 Value::I32(n) => q.bind(*n),
1902 Value::I64(n) => q.bind(*n),
1903 Value::U8(n) => q.bind(*n as i16),
1904 Value::U16(n) => q.bind(*n as i32),
1905 Value::U32(n) => q.bind(*n as i64),
1906 Value::U64(n) => q.bind(*n as i64),
1907 Value::F32(f) => q.bind(*f),
1908 Value::F64(f) => q.bind(*f),
1909 Value::String(s) => q.bind(s.as_str()),
1910 Value::Decimal(s) => q.bind(s.as_str()),
1911 Value::Bytes(b) => q.bind(b.as_slice()),
1912 other => q.bind(other.to_string()),
1913 };
1914 }
1915 let rows_result = q.fetch_all(&mut *pool_conn).await;
1916 self.conn = Some(pool_conn);
1917 let rows = rows_result.map_err(map_sqlx_error)?;
1918 if rows.is_empty() {
1919 return Ok((Vec::new(), Vec::new()));
1920 }
1921 let cols = rows[0].columns();
1922 let mut col_names: Vec<String> = Vec::with_capacity(cols.len());
1923 for col in cols {
1924 col_names.push(col.name().to_string());
1925 }
1926 let mut result_rows: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
1927 for row in rows {
1928 let mut row_values: Vec<Value> = Vec::with_capacity(col_names.len());
1929 for (idx, _) in col_names.iter().enumerate() {
1930 let ordinal = row.columns()[idx].ordinal();
1931 row_values.push(row_to_value_pg(&row, ordinal));
1932 }
1933 result_rows.push(row_values);
1934 }
1935 Ok((col_names, result_rows))
1936 })
1937 }
1938
1939 fn query_stream<'a>(
1949 &'a mut self,
1950 sql: &'a str,
1951 ) -> Pin<Box<dyn futures::Stream<Item = Result<HashMap<String, Value>, DbError>> + Send + 'a>>
1952 {
1953 Box::pin(async_stream::try_stream! {
1954 let mut pool_conn = self
1955 .conn
1956 .take()
1957 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
1958 let mut row_stream = sqlx::query(sqlx::AssertSqlSafe(sql)).fetch(&mut *pool_conn);
1959 while let Some(row_result) = row_stream.next().await {
1962 let row = row_result.map_err(map_sqlx_error)?;
1963 let cols = row.columns();
1964 let mut record = HashMap::with_capacity(cols.len());
1965 for (i, col) in cols.iter().enumerate() {
1966 let name = col.name().to_string();
1967 let value = row_to_value_pg(&row, i);
1969 record.insert(name, value);
1970 }
1971 yield record;
1972 }
1973 drop(row_stream);
1976 self.conn = Some(pool_conn);
1977 })
1978 }
1979}
1980
1981impl Drop for SqlxPgConnection {
1982 fn drop(&mut self) {
1983 if let Some(conn) = self.conn.take() {
1984 drop(conn);
1985 }
1986 }
1987}
1988
1989#[async_trait]
2010pub trait PgExtensions: Send + Sync {
2011 async fn listen(&mut self, channel: &str) -> Result<(), DbError>;
2022
2023 async fn notify(&mut self, channel: &str, payload: &str) -> Result<(), DbError>;
2034
2035 async fn copy_from_stdin(&mut self, sql: &str, data: &[u8]) -> Result<u64, DbError>;
2050}
2051
2052fn validate_pg_channel_name(channel: &str) -> Result<(), DbError> {
2057 if channel.is_empty() {
2058 return Err(DbError::Internal(
2059 "PG channel name must not be empty".to_string(),
2060 ));
2061 }
2062 if !channel
2063 .chars()
2064 .all(|c| c.is_ascii_alphanumeric() || c == '_')
2065 {
2066 return Err(DbError::Internal(format!(
2067 "invalid PG channel name: {} (only alphanumeric and underscore allowed)",
2068 channel
2069 )));
2070 }
2071 Ok(())
2072}
2073
2074#[async_trait]
2075impl PgExtensions for SqlxPgConnection {
2076 async fn listen(&mut self, channel: &str) -> Result<(), DbError> {
2077 validate_pg_channel_name(channel)?;
2078 self.execute(&format!("LISTEN {}", channel)).await?;
2080 Ok(())
2081 }
2082
2083 async fn notify(&mut self, channel: &str, payload: &str) -> Result<(), DbError> {
2084 validate_pg_channel_name(channel)?;
2085 let escaped_payload = payload.replace('\'', "''");
2087 self.execute(&format!("NOTIFY {}, '{}'", channel, escaped_payload))
2088 .await?;
2089 Ok(())
2090 }
2091
2092 async fn copy_from_stdin(&mut self, sql: &str, data: &[u8]) -> Result<u64, DbError> {
2093 let mut pool_conn = self
2094 .conn
2095 .take()
2096 .ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
2097 let mut copy = (*pool_conn)
2101 .copy_in_raw(sql)
2102 .await
2103 .map_err(map_sqlx_error)?;
2104 copy.send(data).await.map_err(map_sqlx_error)?;
2106 let result = copy.finish().await.map_err(map_sqlx_error)?;
2108 self.conn = Some(pool_conn);
2109 Ok(result)
2110 }
2111}
2112
2113pub async fn pg_bulk_insert(
2121 conn: &mut SqlxPgConnection,
2122 table: &str,
2123 columns: &[&str],
2124 rows: &[Vec<Value>],
2125) -> Result<u64, DbError> {
2126 if rows.is_empty() {
2127 return Ok(0);
2128 }
2129 let col_list = columns.join(", ");
2130 let cols_per_row = columns.len();
2131 let placeholders: Vec<String> = rows
2133 .iter()
2134 .enumerate()
2135 .map(|(row_idx, _)| {
2136 let base = row_idx * cols_per_row;
2137 let ph: Vec<String> = (0..cols_per_row)
2138 .map(|i| format!("${}", base + i + 1))
2139 .collect();
2140 format!("({})", ph.join(", "))
2141 })
2142 .collect();
2143 let sql = format!(
2144 "INSERT INTO {} ({}) VALUES {}",
2145 table,
2146 col_list,
2147 placeholders.join(", ")
2148 );
2149 let mut params: Vec<Value> = Vec::with_capacity(rows.len() * cols_per_row);
2151 for row in rows {
2152 for v in row {
2153 params.push(v.clone());
2154 }
2155 }
2156 conn.execute_with_params(&sql, ¶ms).await
2157}