use async_trait::async_trait;
use futures::StreamExt;
use sqlx::{Column, Executor, Row, TypeInfo};
use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::str::FromStr;
use std::sync::Arc;
use sz_orm_core::{
ColType, Connection, ConnectionFactory, DbError, DbType, QueryRows, QueryValues, Value,
};
use crate::error::map_sqlx_error;
fn needs_raw_sql(sql: &str) -> bool {
let trimmed = sql.trim_start();
let upper = trimmed.to_uppercase();
upper.starts_with("BEGIN")
|| upper.starts_with("COMMIT")
|| upper.starts_with("ROLLBACK")
|| upper.starts_with("SAVEPOINT")
|| upper.starts_with("RELEASE")
|| upper.starts_with("SET ")
|| upper.starts_with("USE ")
|| upper.starts_with("START TRANSACTION")
}
fn row_to_value_with_coltype_sqlite(
row: &sqlx::sqlite::SqliteRow,
ordinal: usize,
col_type: ColType,
) -> Value {
match col_type {
ColType::Bool => match row.try_get::<Option<bool>, usize>(ordinal) {
Ok(v) => v.map(Value::Bool).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::I8 => match row.try_get::<Option<i8>, usize>(ordinal) {
Ok(v) => v.map(Value::I8).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::I16 => match row.try_get::<Option<i16>, usize>(ordinal) {
Ok(v) => v.map(Value::I16).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::I32 => match row.try_get::<Option<i32>, usize>(ordinal) {
Ok(v) => v.map(Value::I32).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::I64 => match row.try_get::<Option<i64>, usize>(ordinal) {
Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::U8 => match row.try_get::<Option<u8>, usize>(ordinal) {
Ok(v) => v.map(Value::U8).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::U16 => match row.try_get::<Option<u16>, usize>(ordinal) {
Ok(v) => v.map(Value::U16).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::U32 => match row.try_get::<Option<u32>, usize>(ordinal) {
Ok(v) => v.map(Value::U32).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::U64 => match row.try_get::<Option<i64>, usize>(ordinal) {
Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::F32 => match row.try_get::<Option<f32>, usize>(ordinal) {
Ok(v) => v.map(Value::F32).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::F64 => match row.try_get::<Option<f64>, usize>(ordinal) {
Ok(v) => v.map(Value::F64).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::Decimal => match row.try_get::<Option<String>, usize>(ordinal) {
Ok(v) => v.map(Value::Decimal).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::String => match row.try_get::<Option<String>, usize>(ordinal) {
Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::Bytes => match row.try_get::<Option<Vec<u8>>, usize>(ordinal) {
Ok(v) => v.map(Value::Bytes).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::Date | ColType::DateTime | ColType::Time | ColType::Json | ColType::Uuid => {
match row.try_get::<Option<String>, usize>(ordinal) {
Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
Err(_) => Value::Null,
}
}
ColType::Unknown => {
if let Ok(v) = row.try_get::<Option<i64>, usize>(ordinal) {
return v.map(Value::I64).unwrap_or(Value::Null);
}
if let Ok(v) = row.try_get::<Option<f64>, usize>(ordinal) {
return v.map(Value::F64).unwrap_or(Value::Null);
}
if let Ok(v) = row.try_get::<Option<bool>, usize>(ordinal) {
return v.map(Value::Bool).unwrap_or(Value::Null);
}
if let Ok(v) = row.try_get::<Option<String>, usize>(ordinal) {
return v.map(Value::String).unwrap_or(Value::Null);
}
Value::Null
}
_ => Value::Null,
}
}
pub struct SqlitePoolHandle {
pool: sqlx::SqlitePool,
}
impl SqlitePoolHandle {
pub async fn connect(url: &str) -> Result<Self, DbError> {
Self::connect_with_max(url, 10).await
}
pub async fn connect_with_max(url: &str, max_connections: u32) -> Result<Self, DbError> {
let opts = sqlx::sqlite::SqliteConnectOptions::from_str(url)
.map_err(map_sqlx_error)?
.journal_mode(sqlx::sqlite::SqliteJournalMode::Wal)
.synchronous(sqlx::sqlite::SqliteSynchronous::Normal)
.busy_timeout(std::time::Duration::from_secs(5))
.pragma("mmap_size", "268435456");
let pool = sqlx::sqlite::SqlitePoolOptions::new()
.max_connections(max_connections)
.acquire_timeout(std::time::Duration::from_secs(30))
.idle_timeout(Some(std::time::Duration::from_secs(600)))
.max_lifetime(Some(std::time::Duration::from_secs(1800)))
.connect_with(opts)
.await
.map_err(map_sqlx_error)?;
Ok(Self { pool })
}
pub fn from_pool(pool: sqlx::SqlitePool) -> Self {
Self { pool }
}
pub fn pool(&self) -> &sqlx::SqlitePool {
&self.pool
}
}
pub struct SqlxSqliteConnectionFactory {
pool: Arc<SqlitePoolHandle>,
}
impl SqlxSqliteConnectionFactory {
pub fn new(pool: Arc<SqlitePoolHandle>) -> Self {
Self { pool }
}
}
#[async_trait]
impl ConnectionFactory for SqlxSqliteConnectionFactory {
async fn create(&self) -> Result<Box<dyn Connection>, DbError> {
let conn = self.pool.pool.acquire().await.map_err(map_sqlx_error)?;
Ok(Box::new(crate::metrics_conn::MetricsConn::new(Box::new(
SqlxSqliteConnection {
conn: Some(conn),
connected: true,
in_transaction: false,
},
))))
}
}
pub struct SqlxSqliteConnection {
conn: Option<sqlx::pool::PoolConnection<sqlx::Sqlite>>,
connected: bool,
in_transaction: bool,
}
impl Connection for SqlxSqliteConnection {
fn db_type(&self) -> Option<DbType> {
Some(DbType::Sqlite)
}
fn execute<'a>(
&'a mut self,
sql: &'a str,
) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
Box::pin(async move {
let mut pool_conn = self
.conn
.take()
.ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
let result = if needs_raw_sql(sql) {
(&mut *pool_conn)
.execute(sqlx::raw_sql(sqlx::AssertSqlSafe(sql)))
.await
} else {
(&mut *pool_conn).execute(sqlx::AssertSqlSafe(sql)).await
};
self.conn = Some(pool_conn);
match result {
Ok(r) => Ok(r.rows_affected()),
Err(e) => {
let db_err = map_sqlx_error(e);
if matches!(db_err, DbError::ConnectionError(_) | DbError::IoError(_)) {
self.connected = false;
}
Err(db_err)
}
}
})
}
fn query<'a>(
&'a mut self,
sql: &'a str,
) -> Pin<Box<dyn Future<Output = Result<Vec<HashMap<String, Value>>, DbError>> + Send + 'a>>
{
Box::pin(async move {
let mut pool_conn = self
.conn
.take()
.ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
let rows_result = (&mut *pool_conn).fetch_all(sqlx::AssertSqlSafe(sql)).await;
self.conn = Some(pool_conn);
let rows = rows_result.map_err(map_sqlx_error)?;
if rows.is_empty() {
return Ok(Vec::new());
}
let col_types: Vec<ColType> = rows[0]
.columns()
.iter()
.map(|col| ColType::parse_sqlite(col.type_info().name()))
.collect();
let col_names: Vec<String> = rows[0]
.columns()
.iter()
.map(|col| col.name().to_string())
.collect();
let mut result = Vec::with_capacity(rows.len());
for row in &rows {
let mut record = HashMap::with_capacity(col_types.len());
for (i, name) in col_names.iter().enumerate() {
let value = row_to_value_with_coltype_sqlite(row, i, col_types[i]);
record.insert(name.clone(), value);
}
result.push(record);
}
Ok(result)
})
}
fn begin_transaction<'a>(
&'a mut self,
) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
Box::pin(async move {
if self.in_transaction {
return Err(DbError::Internal("transaction already started".to_string()));
}
self.execute("BEGIN").await?;
self.in_transaction = true;
Ok(())
})
}
fn commit<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
Box::pin(async move {
if self.in_transaction {
self.execute("COMMIT").await?;
self.in_transaction = false;
}
Ok(())
})
}
fn rollback<'a>(
&'a mut self,
) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
Box::pin(async move {
if self.in_transaction {
let result = self.execute("ROLLBACK").await;
if result.is_ok() {
self.in_transaction = false;
}
result.map(|_| ())
} else {
Ok(())
}
})
}
fn is_connected(&self) -> bool {
self.connected && self.conn.is_some()
}
fn in_transaction(&self) -> bool {
self.in_transaction
}
fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
Box::pin(async move {
match self.execute("SELECT 1").await {
Ok(_) => true,
Err(_) => {
self.connected = false;
false
}
}
})
}
fn close<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
Box::pin(async move {
if let Some(conn) = self.conn.take() {
drop(conn);
}
self.connected = false;
self.in_transaction = false;
Ok(())
})
}
fn execute_with_params<'a>(
&'a mut self,
sql: &'a str,
params: &'a [Value],
) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
Box::pin(async move {
if needs_raw_sql(sql) || params.is_empty() {
return self.execute(sql).await;
}
let mut pool_conn = self
.conn
.take()
.ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
for v in params {
q = match v {
Value::Null => q.bind(None::<i64>),
Value::Bool(b) => q.bind(*b),
Value::I8(n) => q.bind(*n),
Value::I16(n) => q.bind(*n),
Value::I32(n) => q.bind(*n),
Value::I64(n) => q.bind(*n),
Value::U8(n) => q.bind(*n),
Value::U16(n) => q.bind(*n),
Value::U32(n) => q.bind(*n),
Value::U64(n) => q.bind(*n as i64),
Value::F32(f) => q.bind(*f),
Value::F64(f) => q.bind(*f),
Value::String(s) => q.bind(s.as_str()),
Value::Decimal(s) => q.bind(s.as_str()),
Value::Bytes(b) => q.bind(b.as_slice()),
Value::Uuid(s)
| Value::Date(s)
| Value::DateTime(s)
| Value::Time(s)
| Value::Json(s) => q.bind(s.as_str()),
other => q.bind(other.to_string()),
};
}
let result = q.execute(&mut *pool_conn).await;
self.conn = Some(pool_conn);
match result {
Ok(r) => Ok(r.rows_affected()),
Err(e) => {
let db_err = map_sqlx_error(e);
if matches!(db_err, DbError::ConnectionError(_) | DbError::IoError(_)) {
self.connected = false;
}
Err(db_err)
}
}
})
}
fn query_with_params<'a>(
&'a mut self,
sql: &'a str,
params: &'a [Value],
) -> Pin<Box<dyn Future<Output = Result<QueryRows, DbError>> + Send + 'a>> {
Box::pin(async move {
if params.is_empty() {
return self.query(sql).await;
}
let mut pool_conn = self
.conn
.take()
.ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
for v in params {
q = match v {
Value::Null => q.bind(None::<i64>),
Value::Bool(b) => q.bind(*b),
Value::I8(n) => q.bind(*n),
Value::I16(n) => q.bind(*n),
Value::I32(n) => q.bind(*n),
Value::I64(n) => q.bind(*n),
Value::U8(n) => q.bind(*n),
Value::U16(n) => q.bind(*n),
Value::U32(n) => q.bind(*n),
Value::U64(n) => q.bind(*n as i64),
Value::F32(f) => q.bind(*f),
Value::F64(f) => q.bind(*f),
Value::String(s) => q.bind(s.as_str()),
Value::Decimal(s) => q.bind(s.as_str()),
Value::Bytes(b) => q.bind(b.as_slice()),
Value::Uuid(s)
| Value::Date(s)
| Value::DateTime(s)
| Value::Time(s)
| Value::Json(s) => q.bind(s.as_str()),
other => q.bind(other.to_string()),
};
}
let rows_result = q.fetch_all(&mut *pool_conn).await;
self.conn = Some(pool_conn);
let rows = rows_result.map_err(map_sqlx_error)?;
if rows.is_empty() {
return Ok(Vec::new());
}
let col_types: Vec<ColType> = rows[0]
.columns()
.iter()
.map(|col| ColType::parse_sqlite(col.type_info().name()))
.collect();
let col_names: Vec<String> = rows[0]
.columns()
.iter()
.map(|col| col.name().to_string())
.collect();
let mut result = Vec::with_capacity(rows.len());
for row in &rows {
let mut record = HashMap::with_capacity(col_types.len());
for (i, name) in col_names.iter().enumerate() {
let value = row_to_value_with_coltype_sqlite(row, i, col_types[i]);
record.insert(name.clone(), value);
}
result.push(record);
}
Ok(result)
})
}
fn query_values<'a>(
&'a mut self,
sql: &'a str,
) -> Pin<Box<dyn Future<Output = Result<QueryValues, DbError>> + Send + 'a>> {
Box::pin(async move {
let mut pool_conn = self
.conn
.take()
.ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
let rows_result = (&mut *pool_conn).fetch_all(sqlx::AssertSqlSafe(sql)).await;
self.conn = Some(pool_conn);
let rows = rows_result.map_err(map_sqlx_error)?;
if rows.is_empty() {
return Ok((Vec::new(), Vec::new()));
}
let cols = rows[0].columns();
let mut col_names: Vec<String> = Vec::with_capacity(cols.len());
let mut col_types: Vec<ColType> = Vec::with_capacity(cols.len());
for col in cols {
col_names.push(col.name().to_string());
col_types.push(ColType::parse_sqlite(col.type_info().name()));
}
let mut result_rows: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
for row in &rows {
let mut row_values: Vec<Value> = Vec::with_capacity(col_names.len());
for (idx, _) in col_names.iter().enumerate() {
row_values.push(row_to_value_with_coltype_sqlite(row, idx, col_types[idx]));
}
result_rows.push(row_values);
}
Ok((col_names, result_rows))
})
}
fn query_values_with_params<'a>(
&'a mut self,
sql: &'a str,
params: &'a [Value],
) -> Pin<Box<dyn Future<Output = Result<QueryValues, DbError>> + Send + 'a>> {
Box::pin(async move {
if params.is_empty() {
return self.query_values(sql).await;
}
let mut pool_conn = self
.conn
.take()
.ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
for v in params {
q = match v {
Value::Null => q.bind(None::<i64>),
Value::Bool(b) => q.bind(*b),
Value::I8(n) => q.bind(*n),
Value::I16(n) => q.bind(*n),
Value::I32(n) => q.bind(*n),
Value::I64(n) => q.bind(*n),
Value::U8(n) => q.bind(*n),
Value::U16(n) => q.bind(*n),
Value::U32(n) => q.bind(*n),
Value::U64(n) => q.bind(*n as i64),
Value::F32(f) => q.bind(*f),
Value::F64(f) => q.bind(*f),
Value::String(s) => q.bind(s.as_str()),
Value::Decimal(s) => q.bind(s.as_str()),
Value::Bytes(b) => q.bind(b.as_slice()),
Value::Uuid(s)
| Value::Date(s)
| Value::DateTime(s)
| Value::Time(s)
| Value::Json(s) => q.bind(s.as_str()),
other => q.bind(other.to_string()),
};
}
let rows_result = q.fetch_all(&mut *pool_conn).await;
self.conn = Some(pool_conn);
let rows = rows_result.map_err(map_sqlx_error)?;
if rows.is_empty() {
return Ok((Vec::new(), Vec::new()));
}
let cols = rows[0].columns();
let mut col_names: Vec<String> = Vec::with_capacity(cols.len());
let mut col_types: Vec<ColType> = Vec::with_capacity(cols.len());
for col in cols {
col_names.push(col.name().to_string());
col_types.push(ColType::parse_sqlite(col.type_info().name()));
}
let mut result_rows: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
for row in &rows {
let mut row_values: Vec<Value> = Vec::with_capacity(col_names.len());
for (idx, _) in col_names.iter().enumerate() {
row_values.push(row_to_value_with_coltype_sqlite(row, idx, col_types[idx]));
}
result_rows.push(row_values);
}
Ok((col_names, result_rows))
})
}
fn query_stream<'a>(
&'a mut self,
sql: &'a str,
) -> Pin<Box<dyn futures::Stream<Item = Result<HashMap<String, Value>, DbError>> + Send + 'a>>
{
Box::pin(async_stream::try_stream! {
let mut pool_conn = self
.conn
.take()
.ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
let mut row_stream = sqlx::query(sqlx::AssertSqlSafe(sql)).fetch(&mut *pool_conn);
let mut col_types: Vec<ColType> = Vec::new();
let mut col_names: Vec<String> = Vec::new();
let mut first_row = true;
while let Some(row_result) = row_stream.next().await {
let row = row_result.map_err(map_sqlx_error)?;
if first_row {
for col in row.columns() {
col_names.push(col.name().to_string());
col_types.push(ColType::parse_sqlite(col.type_info().name()));
}
first_row = false;
}
let mut record = HashMap::with_capacity(col_names.len());
for (i, name) in col_names.iter().enumerate() {
let value = row_to_value_with_coltype_sqlite(&row, i, col_types[i]);
record.insert(name.clone(), value);
}
yield record;
}
drop(row_stream);
self.conn = Some(pool_conn);
})
}
}
impl Drop for SqlxSqliteConnection {
fn drop(&mut self) {
if let Some(conn) = self.conn.take() {
drop(conn);
}
}
}
pub async fn sqlite_backup(
conn: &mut SqlxSqliteConnection,
dest_path: &str,
) -> Result<(), DbError> {
let escaped_path = dest_path.replace('\'', "''");
let sql = format!("VACUUM INTO '{}'", escaped_path);
conn.execute(&sql).await?;
Ok(())
}
fn row_to_value_mysql(row: &sqlx::mysql::MySqlRow, ordinal: usize) -> Value {
use sqlx::TypeInfo;
let type_name = row.columns()[ordinal].type_info().name();
match type_name {
"BOOLEAN" => match row.try_get::<Option<bool>, usize>(ordinal) {
Ok(v) => v.map(Value::Bool).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
"TINYINT" | "TINYINT UNSIGNED" => match row.try_get::<Option<i8>, usize>(ordinal) {
Ok(v) => v.map(Value::I8).unwrap_or(Value::Null),
Err(_) => match row.try_get::<Option<u8>, usize>(ordinal) {
Ok(v) => v.map(Value::U8).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
},
"SMALLINT" | "SMALLINT UNSIGNED" => match row.try_get::<Option<i16>, usize>(ordinal) {
Ok(v) => v.map(Value::I16).unwrap_or(Value::Null),
Err(_) => match row.try_get::<Option<u16>, usize>(ordinal) {
Ok(v) => v.map(Value::U16).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
},
"INT" | "INT UNSIGNED" | "MEDIUMINT" | "MEDIUMINT UNSIGNED" => {
match row.try_get::<Option<i32>, usize>(ordinal) {
Ok(v) => v.map(Value::I32).unwrap_or(Value::Null),
Err(_) => match row.try_get::<Option<u32>, usize>(ordinal) {
Ok(v) => v.map(Value::U32).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
}
}
"BIGINT" | "BIGINT UNSIGNED" => match row.try_get::<Option<i64>, usize>(ordinal) {
Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
Err(_) => match row.try_get::<Option<u64>, usize>(ordinal) {
Ok(v) => v.map(Value::U64).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
},
"FLOAT" => match row.try_get::<Option<f32>, usize>(ordinal) {
Ok(v) => v.map(Value::F32).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
"DOUBLE" => match row.try_get::<Option<f64>, usize>(ordinal) {
Ok(v) => v.map(Value::F64).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
"VARCHAR" | "TEXT" | "CHAR" | "TINYTEXT" | "MEDIUMTEXT" | "LONGTEXT" | "ENUM" => {
match row.try_get::<Option<String>, usize>(ordinal) {
Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
Err(_) => Value::Null,
}
}
"BLOB" | "TINYBLOB" | "MEDIUMBLOB" | "LONGBLOB" | "BINARY" | "VARBINARY" => {
match row.try_get::<Option<Vec<u8>>, usize>(ordinal) {
Ok(v) => v.map(Value::Bytes).unwrap_or(Value::Null),
Err(_) => Value::Null,
}
}
"DECIMAL" | "NUMERIC" | "NEWDECIMAL" => {
match row.try_get::<Option<bigdecimal::BigDecimal>, usize>(ordinal) {
Ok(Some(v)) => Value::Decimal(v.to_string()),
Ok(None) => Value::Null,
Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
Ok(v) => v.map(Value::Decimal).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
}
}
"DATE" => match row.try_get::<Option<sqlx::types::chrono::NaiveDate>, usize>(ordinal) {
Ok(v) => v.map(|d| Value::Date(d.to_string())).unwrap_or(Value::Null),
Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
},
"DATETIME" | "TIMESTAMP" => {
match row.try_get::<Option<sqlx::types::chrono::NaiveDateTime>, usize>(ordinal) {
Ok(v) => v
.map(|d| Value::DateTime(d.to_string()))
.unwrap_or(Value::Null),
Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
}
}
"TIME" => match row.try_get::<Option<sqlx::types::chrono::NaiveTime>, usize>(ordinal) {
Ok(v) => v.map(|t| Value::Time(t.to_string())).unwrap_or(Value::Null),
Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
},
_ => {
if let Ok(v) = row.try_get::<Option<i64>, usize>(ordinal) {
return v.map(Value::I64).unwrap_or(Value::Null);
}
if let Ok(v) = row.try_get::<Option<f64>, usize>(ordinal) {
return v.map(Value::F64).unwrap_or(Value::Null);
}
if let Ok(v) = row.try_get::<Option<bool>, usize>(ordinal) {
return v.map(Value::Bool).unwrap_or(Value::Null);
}
if let Ok(v) = row.try_get::<Option<String>, usize>(ordinal) {
return v.map(Value::String).unwrap_or(Value::Null);
}
Value::Null
}
}
}
fn row_to_value_with_coltype_mysql(
row: &sqlx::mysql::MySqlRow,
ordinal: usize,
col_type: ColType,
) -> Value {
match col_type {
ColType::Bool => match row.try_get::<Option<bool>, usize>(ordinal) {
Ok(v) => v.map(Value::Bool).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::I8 => match row.try_get::<Option<i8>, usize>(ordinal) {
Ok(v) => v.map(Value::I8).unwrap_or(Value::Null),
Err(_) => match row.try_get::<Option<u8>, usize>(ordinal) {
Ok(v) => v.map(Value::U8).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
},
ColType::I16 => match row.try_get::<Option<i16>, usize>(ordinal) {
Ok(v) => v.map(Value::I16).unwrap_or(Value::Null),
Err(_) => match row.try_get::<Option<u16>, usize>(ordinal) {
Ok(v) => v.map(Value::U16).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
},
ColType::I32 => match row.try_get::<Option<i32>, usize>(ordinal) {
Ok(v) => v.map(Value::I32).unwrap_or(Value::Null),
Err(_) => match row.try_get::<Option<u32>, usize>(ordinal) {
Ok(v) => v.map(Value::U32).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
},
ColType::I64 => match row.try_get::<Option<i64>, usize>(ordinal) {
Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
Err(_) => match row.try_get::<Option<u64>, usize>(ordinal) {
Ok(v) => v.map(Value::U64).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
},
ColType::U8 => match row.try_get::<Option<u8>, usize>(ordinal) {
Ok(v) => v.map(Value::U8).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::U16 => match row.try_get::<Option<u16>, usize>(ordinal) {
Ok(v) => v.map(Value::U16).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::U32 => match row.try_get::<Option<u32>, usize>(ordinal) {
Ok(v) => v.map(Value::U32).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::U64 => match row.try_get::<Option<u64>, usize>(ordinal) {
Ok(v) => v.map(Value::U64).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::F32 => match row.try_get::<Option<f32>, usize>(ordinal) {
Ok(v) => v.map(Value::F32).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::F64 => match row.try_get::<Option<f64>, usize>(ordinal) {
Ok(v) => v.map(Value::F64).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::Decimal => match row.try_get::<Option<bigdecimal::BigDecimal>, usize>(ordinal) {
Ok(Some(v)) => Value::Decimal(v.to_string()),
Ok(None) => Value::Null,
Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
Ok(v) => v.map(Value::Decimal).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
},
ColType::String => match row.try_get::<Option<String>, usize>(ordinal) {
Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::Bytes => match row.try_get::<Option<Vec<u8>>, usize>(ordinal) {
Ok(v) => v.map(Value::Bytes).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::Date => {
match row.try_get::<Option<sqlx::types::chrono::NaiveDate>, usize>(ordinal) {
Ok(v) => v.map(|d| Value::Date(d.to_string())).unwrap_or(Value::Null),
Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
}
}
ColType::DateTime => {
match row.try_get::<Option<sqlx::types::chrono::NaiveDateTime>, usize>(ordinal) {
Ok(v) => v
.map(|d| Value::DateTime(d.to_string()))
.unwrap_or(Value::Null),
Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
}
}
ColType::Time => {
match row.try_get::<Option<sqlx::types::chrono::NaiveTime>, usize>(ordinal) {
Ok(v) => v.map(|t| Value::Time(t.to_string())).unwrap_or(Value::Null),
Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
}
}
ColType::Json | ColType::Uuid => match row.try_get::<Option<String>, usize>(ordinal) {
Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::Unknown => {
if let Ok(v) = row.try_get::<Option<i64>, usize>(ordinal) {
return v.map(Value::I64).unwrap_or(Value::Null);
}
if let Ok(v) = row.try_get::<Option<f64>, usize>(ordinal) {
return v.map(Value::F64).unwrap_or(Value::Null);
}
if let Ok(v) = row.try_get::<Option<bool>, usize>(ordinal) {
return v.map(Value::Bool).unwrap_or(Value::Null);
}
if let Ok(v) = row.try_get::<Option<String>, usize>(ordinal) {
return v.map(Value::String).unwrap_or(Value::Null);
}
Value::Null
}
_ => Value::Null,
}
}
pub struct MySqlPoolHandle {
pool: sqlx::MySqlPool,
}
impl MySqlPoolHandle {
pub async fn connect(url: &str) -> Result<Self, DbError> {
Self::connect_with_max(url, 10).await
}
pub async fn connect_with_max(url: &str, max_connections: u32) -> Result<Self, DbError> {
let pool = sqlx::pool::PoolOptions::<sqlx::MySql>::new()
.max_connections(max_connections)
.acquire_timeout(std::time::Duration::from_secs(30))
.idle_timeout(Some(std::time::Duration::from_secs(600)))
.max_lifetime(Some(std::time::Duration::from_secs(1800)))
.connect(url)
.await
.map_err(map_sqlx_error)?;
Ok(Self { pool })
}
pub fn from_pool(pool: sqlx::MySqlPool) -> Self {
Self { pool }
}
pub fn pool(&self) -> &sqlx::MySqlPool {
&self.pool
}
}
pub struct SqlxMySqlConnectionFactory {
pool: Arc<MySqlPoolHandle>,
}
impl SqlxMySqlConnectionFactory {
pub fn new(pool: Arc<MySqlPoolHandle>) -> Self {
Self { pool }
}
}
#[async_trait]
impl ConnectionFactory for SqlxMySqlConnectionFactory {
async fn create(&self) -> Result<Box<dyn Connection>, DbError> {
let conn = self.pool.pool.acquire().await.map_err(map_sqlx_error)?;
Ok(Box::new(crate::metrics_conn::MetricsConn::new(Box::new(
SqlxMySqlConnection {
conn: Some(conn),
connected: true,
in_transaction: false,
},
))))
}
}
pub struct SqlxMySqlConnection {
conn: Option<sqlx::pool::PoolConnection<sqlx::MySql>>,
connected: bool,
in_transaction: bool,
}
impl Connection for SqlxMySqlConnection {
fn db_type(&self) -> Option<DbType> {
Some(DbType::MySQL)
}
fn session_reset_sql(&self) -> Option<&'static str> {
Some(
"SET SESSION sql_mode = DEFAULT, \
SESSION character_set_client = @@GLOBAL.character_set_client, \
SESSION character_set_results = @@GLOBAL.character_set_results, \
SESSION character_set_connection = @@GLOBAL.character_set_connection, \
SESSION time_zone = @@GLOBAL.time_zone, \
SESSION autocommit = 1",
)
}
fn execute<'a>(
&'a mut self,
sql: &'a str,
) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
Box::pin(async move {
let mut pool_conn = self
.conn
.take()
.ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
let result = if needs_raw_sql(sql) {
(&mut *pool_conn)
.execute(sqlx::raw_sql(sqlx::AssertSqlSafe(sql)))
.await
} else {
(&mut *pool_conn).execute(sqlx::AssertSqlSafe(sql)).await
};
self.conn = Some(pool_conn);
match result {
Ok(r) => Ok(r.rows_affected()),
Err(e) => {
let db_err = map_sqlx_error(e);
if matches!(db_err, DbError::ConnectionError(_) | DbError::IoError(_)) {
self.connected = false;
}
Err(db_err)
}
}
})
}
fn query<'a>(
&'a mut self,
sql: &'a str,
) -> Pin<Box<dyn Future<Output = Result<Vec<HashMap<String, Value>>, DbError>> + Send + 'a>>
{
Box::pin(async move {
let mut pool_conn = self
.conn
.take()
.ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
let rows_result = (&mut *pool_conn).fetch_all(sqlx::AssertSqlSafe(sql)).await;
self.conn = Some(pool_conn);
let rows = rows_result.map_err(map_sqlx_error)?;
if rows.is_empty() {
return Ok(Vec::new());
}
let col_types: Vec<ColType> = rows[0]
.columns()
.iter()
.map(|col| ColType::parse_mysql(col.type_info().name()))
.collect();
let col_names: Vec<String> = rows[0]
.columns()
.iter()
.map(|col| col.name().to_string())
.collect();
let mut result = Vec::with_capacity(rows.len());
for row in &rows {
let mut record = HashMap::with_capacity(col_types.len());
for (i, name) in col_names.iter().enumerate() {
let value = row_to_value_with_coltype_mysql(row, i, col_types[i]);
record.insert(name.clone(), value);
}
result.push(record);
}
Ok(result)
})
}
fn begin_transaction<'a>(
&'a mut self,
) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
Box::pin(async move {
if self.in_transaction {
return Err(DbError::Internal("transaction already started".to_string()));
}
self.execute("BEGIN").await?;
self.in_transaction = true;
Ok(())
})
}
fn commit<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
Box::pin(async move {
if self.in_transaction {
self.execute("COMMIT").await?;
self.in_transaction = false;
}
Ok(())
})
}
fn rollback<'a>(
&'a mut self,
) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
Box::pin(async move {
if self.in_transaction {
let result = self.execute("ROLLBACK").await;
if result.is_ok() {
self.in_transaction = false;
}
result.map(|_| ())
} else {
Ok(())
}
})
}
fn is_connected(&self) -> bool {
self.connected && self.conn.is_some()
}
fn in_transaction(&self) -> bool {
self.in_transaction
}
fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
Box::pin(async move {
match self.execute("SELECT 1").await {
Ok(_) => true,
Err(_) => {
self.connected = false;
false
}
}
})
}
fn close<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
Box::pin(async move {
if let Some(conn) = self.conn.take() {
drop(conn);
}
self.connected = false;
self.in_transaction = false;
Ok(())
})
}
fn execute_with_params<'a>(
&'a mut self,
sql: &'a str,
params: &'a [Value],
) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
Box::pin(async move {
if needs_raw_sql(sql) || params.is_empty() {
return self.execute(sql).await;
}
let mut pool_conn = self
.conn
.take()
.ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
for v in params {
q = match v {
Value::Null => q.bind(None::<i64>),
Value::Bool(b) => q.bind(*b),
Value::I8(n) => q.bind(*n),
Value::I16(n) => q.bind(*n),
Value::I32(n) => q.bind(*n),
Value::I64(n) => q.bind(*n),
Value::U8(n) => q.bind(*n),
Value::U16(n) => q.bind(*n),
Value::U32(n) => q.bind(*n),
Value::U64(n) => q.bind(*n as i64),
Value::F32(f) => q.bind(*f),
Value::F64(f) => q.bind(*f),
Value::String(s) => q.bind(s.as_str()),
Value::Decimal(s) => q.bind(s.as_str()),
Value::Bytes(b) => q.bind(b.as_slice()),
Value::Uuid(s)
| Value::Date(s)
| Value::DateTime(s)
| Value::Time(s)
| Value::Json(s) => q.bind(s.as_str()),
other => q.bind(other.to_string()),
};
}
let result = q.execute(&mut *pool_conn).await;
self.conn = Some(pool_conn);
match result {
Ok(r) => Ok(r.rows_affected()),
Err(e) => {
let db_err = map_sqlx_error(e);
if matches!(db_err, DbError::ConnectionError(_) | DbError::IoError(_)) {
self.connected = false;
}
Err(db_err)
}
}
})
}
fn query_with_params<'a>(
&'a mut self,
sql: &'a str,
params: &'a [Value],
) -> Pin<Box<dyn Future<Output = Result<QueryRows, DbError>> + Send + 'a>> {
Box::pin(async move {
if params.is_empty() {
return self.query(sql).await;
}
let mut pool_conn = self
.conn
.take()
.ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
for v in params {
q = match v {
Value::Null => q.bind(None::<i64>),
Value::Bool(b) => q.bind(*b),
Value::I8(n) => q.bind(*n),
Value::I16(n) => q.bind(*n),
Value::I32(n) => q.bind(*n),
Value::I64(n) => q.bind(*n),
Value::U8(n) => q.bind(*n),
Value::U16(n) => q.bind(*n),
Value::U32(n) => q.bind(*n),
Value::U64(n) => q.bind(*n as i64),
Value::F32(f) => q.bind(*f),
Value::F64(f) => q.bind(*f),
Value::String(s) => q.bind(s.as_str()),
Value::Decimal(s) => q.bind(s.as_str()),
Value::Bytes(b) => q.bind(b.as_slice()),
Value::Uuid(s)
| Value::Date(s)
| Value::DateTime(s)
| Value::Time(s)
| Value::Json(s) => q.bind(s.as_str()),
other => q.bind(other.to_string()),
};
}
let rows_result = q.fetch_all(&mut *pool_conn).await;
self.conn = Some(pool_conn);
let rows = rows_result.map_err(map_sqlx_error)?;
if rows.is_empty() {
return Ok(Vec::new());
}
let col_names: Vec<String> = rows[0]
.columns()
.iter()
.map(|col| col.name().to_string())
.collect();
let mut result = Vec::with_capacity(rows.len());
for row in &rows {
let mut record = HashMap::with_capacity(col_names.len());
for (i, name) in col_names.iter().enumerate() {
let value = row_to_value_mysql(row, i);
record.insert(name.clone(), value);
}
result.push(record);
}
Ok(result)
})
}
fn query_values<'a>(
&'a mut self,
sql: &'a str,
) -> Pin<Box<dyn Future<Output = Result<QueryValues, DbError>> + Send + 'a>> {
Box::pin(async move {
let mut pool_conn = self
.conn
.take()
.ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
let rows_result = (&mut *pool_conn).fetch_all(sqlx::AssertSqlSafe(sql)).await;
self.conn = Some(pool_conn);
let rows = rows_result.map_err(map_sqlx_error)?;
if rows.is_empty() {
return Ok((Vec::new(), Vec::new()));
}
let cols = rows[0].columns();
let mut col_names: Vec<String> = Vec::with_capacity(cols.len());
let mut col_types: Vec<ColType> = Vec::with_capacity(cols.len());
for col in cols {
col_names.push(col.name().to_string());
col_types.push(ColType::parse_mysql(col.type_info().name()));
}
let mut result_rows: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
for row in &rows {
let mut row_values: Vec<Value> = Vec::with_capacity(col_names.len());
for (idx, _) in col_names.iter().enumerate() {
row_values.push(row_to_value_with_coltype_mysql(row, idx, col_types[idx]));
}
result_rows.push(row_values);
}
Ok((col_names, result_rows))
})
}
fn query_values_with_params<'a>(
&'a mut self,
sql: &'a str,
params: &'a [Value],
) -> Pin<Box<dyn Future<Output = Result<QueryValues, DbError>> + Send + 'a>> {
Box::pin(async move {
if params.is_empty() {
return self.query_values(sql).await;
}
let mut pool_conn = self
.conn
.take()
.ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
let mut q = sqlx::query(sqlx::AssertSqlSafe(sql));
for v in params {
q = match v {
Value::Null => q.bind(None::<i64>),
Value::Bool(b) => q.bind(*b),
Value::I8(n) => q.bind(*n),
Value::I16(n) => q.bind(*n),
Value::I32(n) => q.bind(*n),
Value::I64(n) => q.bind(*n),
Value::U8(n) => q.bind(*n),
Value::U16(n) => q.bind(*n),
Value::U32(n) => q.bind(*n),
Value::U64(n) => q.bind(*n as i64),
Value::F32(f) => q.bind(*f),
Value::F64(f) => q.bind(*f),
Value::String(s) => q.bind(s.as_str()),
Value::Decimal(s) => q.bind(s.as_str()),
Value::Bytes(b) => q.bind(b.as_slice()),
Value::Uuid(s)
| Value::Date(s)
| Value::DateTime(s)
| Value::Time(s)
| Value::Json(s) => q.bind(s.as_str()),
other => q.bind(other.to_string()),
};
}
let rows_result = q.fetch_all(&mut *pool_conn).await;
self.conn = Some(pool_conn);
let rows = rows_result.map_err(map_sqlx_error)?;
if rows.is_empty() {
return Ok((Vec::new(), Vec::new()));
}
let cols = rows[0].columns();
let mut col_names: Vec<String> = Vec::with_capacity(cols.len());
let mut col_types: Vec<ColType> = Vec::with_capacity(cols.len());
for col in cols {
col_names.push(col.name().to_string());
col_types.push(ColType::parse_mysql(col.type_info().name()));
}
let mut result_rows: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
for row in &rows {
let mut row_values: Vec<Value> = Vec::with_capacity(col_names.len());
for (idx, _) in col_names.iter().enumerate() {
row_values.push(row_to_value_with_coltype_mysql(row, idx, col_types[idx]));
}
result_rows.push(row_values);
}
Ok((col_names, result_rows))
})
}
fn query_stream<'a>(
&'a mut self,
sql: &'a str,
) -> Pin<Box<dyn futures::Stream<Item = Result<HashMap<String, Value>, DbError>> + Send + 'a>>
{
Box::pin(async_stream::try_stream! {
let mut pool_conn = self
.conn
.take()
.ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
let mut row_stream = sqlx::query(sqlx::AssertSqlSafe(sql)).fetch(&mut *pool_conn);
let mut col_types: Vec<ColType> = Vec::new();
let mut col_names: Vec<String> = Vec::new();
let mut first_row = true;
while let Some(row_result) = row_stream.next().await {
let row = row_result.map_err(map_sqlx_error)?;
if first_row {
for col in row.columns() {
col_names.push(col.name().to_string());
col_types.push(ColType::parse_mysql(col.type_info().name()));
}
first_row = false;
}
let mut record = HashMap::with_capacity(col_names.len());
for (i, name) in col_names.iter().enumerate() {
let value = row_to_value_with_coltype_mysql(&row, i, col_types[i]);
record.insert(name.clone(), value);
}
yield record;
}
drop(row_stream);
self.conn = Some(pool_conn);
})
}
}
impl Drop for SqlxMySqlConnection {
fn drop(&mut self) {
if let Some(conn) = self.conn.take() {
drop(conn);
}
}
}
pub async fn mysql_bulk_insert(
conn: &mut SqlxMySqlConnection,
table: &str,
columns: &[&str],
rows: &[Vec<Value>],
) -> Result<u64, DbError> {
if rows.is_empty() {
return Ok(0);
}
let col_list = columns.join(", ");
let cols_per_row = columns.len();
let row_placeholder = format!("({})", vec!["?"; cols_per_row].join(", "));
let placeholders = vec![row_placeholder; rows.len()].join(", ");
let sql = format!(
"INSERT INTO {} ({}) VALUES {}",
table, col_list, placeholders
);
let mut params: Vec<Value> = Vec::with_capacity(rows.len() * cols_per_row);
for row in rows {
for v in row {
params.push(v.clone());
}
}
conn.execute_with_params(&sql, ¶ms).await
}
fn pg_decode_err(
row: &sqlx::postgres::PgRow,
ordinal: usize,
type_name: &str,
err: sqlx::Error,
) -> DbError {
let column = row
.columns()
.get(ordinal)
.map(|c| c.name().to_string())
.unwrap_or_else(|| format!("#{ordinal}"));
DbError::QueryError(format!(
"PostgreSQL 列 `{column}`(类型 {type_name},序号 {ordinal})解码失败:{err};\
已拒绝以 NULL 代替真实值"
))
}
fn row_to_value_pg(row: &sqlx::postgres::PgRow, ordinal: usize) -> Result<Value, DbError> {
use sqlx::TypeInfo;
let type_name = row.columns()[ordinal].type_info().name();
macro_rules! dec {
($ty:ty) => {
row.try_get::<Option<$ty>, usize>(ordinal)
.map_err(|e| pg_decode_err(row, ordinal, type_name, e))
};
}
match type_name {
"BOOL" => Ok(dec!(bool)?.map(Value::Bool).unwrap_or(Value::Null)),
"INT2" => Ok(dec!(i16)?.map(Value::I16).unwrap_or(Value::Null)),
"INT4" | "OID" => Ok(dec!(i32)?.map(Value::I32).unwrap_or(Value::Null)),
"INT8" => Ok(dec!(i64)?.map(Value::I64).unwrap_or(Value::Null)),
"FLOAT4" => Ok(dec!(f32)?.map(Value::F32).unwrap_or(Value::Null)),
"FLOAT8" => Ok(dec!(f64)?.map(Value::F64).unwrap_or(Value::Null)),
"TEXT" | "VARCHAR" | "CHAR" | "NAME" | "BPCHAR" => {
Ok(dec!(String)?.map(Value::String).unwrap_or(Value::Null))
}
"BYTEA" => Ok(dec!(Vec<u8>)?.map(Value::Bytes).unwrap_or(Value::Null)),
"NUMERIC" => Ok(dec!(bigdecimal::BigDecimal)?
.map(|v| Value::Decimal(v.to_string()))
.unwrap_or(Value::Null)),
"UUID" => Ok(dec!(sqlx::types::Uuid)?
.map(|uuid| Value::String(uuid.to_string()))
.unwrap_or(Value::Null)),
"TIMESTAMP" => Ok(dec!(sqlx::types::chrono::NaiveDateTime)?
.map(|v| Value::DateTime(v.to_string()))
.unwrap_or(Value::Null)),
"TIMESTAMPTZ" => Ok(
dec!(sqlx::types::chrono::DateTime<sqlx::types::chrono::Utc>)?
.map(|v| Value::DateTime(v.format("%Y-%m-%dT%H:%M:%S%.6fZ").to_string()))
.unwrap_or(Value::Null),
),
"DATE" => Ok(dec!(sqlx::types::chrono::NaiveDate)?
.map(|v| Value::Date(v.to_string()))
.unwrap_or(Value::Null)),
"TIME" => Ok(dec!(sqlx::types::chrono::NaiveTime)?
.map(|v| Value::Time(v.to_string()))
.unwrap_or(Value::Null)),
"JSON" | "JSONB" => Ok(dec!(serde_json::Value)?
.map(|v| Value::Json(v.to_string()))
.unwrap_or(Value::Null)),
_ => {
if let Ok(v) = row.try_get::<Option<i64>, usize>(ordinal) {
return Ok(v.map(Value::I64).unwrap_or(Value::Null));
}
if let Ok(v) = row.try_get::<Option<f64>, usize>(ordinal) {
return Ok(v.map(Value::F64).unwrap_or(Value::Null));
}
if let Ok(v) = row.try_get::<Option<bool>, usize>(ordinal) {
return Ok(v.map(Value::Bool).unwrap_or(Value::Null));
}
match row.try_get::<Option<String>, usize>(ordinal) {
Ok(v) => Ok(v.map(Value::String).unwrap_or(Value::Null)),
Err(e) => Err(pg_decode_err(row, ordinal, type_name, e)),
}
}
}
}
fn row_to_value_with_coltype_pg(
row: &sqlx::postgres::PgRow,
ordinal: usize,
col_type: ColType,
) -> Value {
match col_type {
ColType::Bool => match row.try_get::<Option<bool>, usize>(ordinal) {
Ok(v) => v.map(Value::Bool).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::I8 => match row.try_get::<Option<i8>, usize>(ordinal) {
Ok(v) => v.map(Value::I8).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::I16 => match row.try_get::<Option<i16>, usize>(ordinal) {
Ok(v) => v.map(Value::I16).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::I32 => match row.try_get::<Option<i32>, usize>(ordinal) {
Ok(v) => v.map(Value::I32).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::I64 => match row.try_get::<Option<i64>, usize>(ordinal) {
Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::U8 => match row.try_get::<Option<i16>, usize>(ordinal) {
Ok(v) => v.map(Value::I16).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::U16 => match row.try_get::<Option<i32>, usize>(ordinal) {
Ok(v) => v.map(Value::I32).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::U32 => match row.try_get::<Option<i64>, usize>(ordinal) {
Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::U64 => match row.try_get::<Option<i64>, usize>(ordinal) {
Ok(v) => v.map(Value::I64).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::F32 => match row.try_get::<Option<f32>, usize>(ordinal) {
Ok(v) => v.map(Value::F32).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::F64 => match row.try_get::<Option<f64>, usize>(ordinal) {
Ok(v) => v.map(Value::F64).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::Decimal => match row.try_get::<Option<bigdecimal::BigDecimal>, usize>(ordinal) {
Ok(Some(v)) => Value::Decimal(v.to_string()),
Ok(None) => Value::Null,
Err(_) => match row.try_get::<Option<rust_decimal::Decimal>, usize>(ordinal) {
Ok(Some(v)) => Value::Decimal(v.to_string()),
Ok(None) => Value::Null,
Err(_) => Value::Null,
},
},
ColType::String => match row.try_get::<Option<String>, usize>(ordinal) {
Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::Bytes => match row.try_get::<Option<Vec<u8>>, usize>(ordinal) {
Ok(v) => v.map(Value::Bytes).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
ColType::Date => {
match row.try_get::<Option<sqlx::types::chrono::NaiveDate>, usize>(ordinal) {
Ok(v) => v.map(|d| Value::Date(d.to_string())).unwrap_or(Value::Null),
Err(_) => Value::Null,
}
}
ColType::DateTime => match row
.try_get::<Option<sqlx::types::chrono::NaiveDateTime>, usize>(ordinal)
{
Ok(v) => v
.map(|d| Value::DateTime(d.to_string()))
.unwrap_or(Value::Null),
Err(_) => match row
.try_get::<Option<sqlx::types::chrono::DateTime<sqlx::types::chrono::Utc>>, usize>(
ordinal,
) {
Ok(v) => v
.map(|d| Value::DateTime(d.format("%Y-%m-%dT%H:%M:%S%.6fZ").to_string()))
.unwrap_or(Value::Null),
Err(_) => Value::Null,
},
},
ColType::Time => {
match row.try_get::<Option<sqlx::types::chrono::NaiveTime>, usize>(ordinal) {
Ok(v) => v.map(|t| Value::Time(t.to_string())).unwrap_or(Value::Null),
Err(_) => Value::Null,
}
}
ColType::Json => match row.try_get::<Option<serde_json::Value>, usize>(ordinal) {
Ok(v) => v.map(|j| Value::Json(j.to_string())).unwrap_or(Value::Null),
Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
},
ColType::Uuid => match row.try_get::<Option<sqlx::types::Uuid>, usize>(ordinal) {
Ok(v) => v
.map(|uuid| Value::String(uuid.to_string()))
.unwrap_or(Value::Null),
Err(_) => match row.try_get::<Option<String>, usize>(ordinal) {
Ok(v) => v.map(Value::String).unwrap_or(Value::Null),
Err(_) => Value::Null,
},
},
ColType::Unknown => {
if let Ok(v) = row.try_get::<Option<i64>, usize>(ordinal) {
return v.map(Value::I64).unwrap_or(Value::Null);
}
if let Ok(v) = row.try_get::<Option<f64>, usize>(ordinal) {
return v.map(Value::F64).unwrap_or(Value::Null);
}
if let Ok(v) = row.try_get::<Option<bool>, usize>(ordinal) {
return v.map(Value::Bool).unwrap_or(Value::Null);
}
if let Ok(v) = row.try_get::<Option<String>, usize>(ordinal) {
return v.map(Value::String).unwrap_or(Value::Null);
}
Value::Null
}
_ => Value::Null,
}
}
fn pg_rewrite_placeholders(sql: &str) -> std::borrow::Cow<'_, str> {
if !sql.as_bytes().contains(&b'?') {
return std::borrow::Cow::Borrowed(sql);
}
let mut out = String::with_capacity(sql.len() + 8);
let mut chars = sql.chars().peekable();
let mut in_single = false;
let mut in_double = false;
let mut in_line_comment = false;
let mut in_block_comment = false;
let mut dollar_tag: Option<String> = None;
let mut prev = '\0';
let mut n = 0usize;
let mut saw_native = false;
while let Some(c) = chars.next() {
if in_line_comment {
out.push(c);
if c == '\n' {
in_line_comment = false;
}
prev = c;
continue;
}
if in_block_comment {
out.push(c);
if c == '*' && chars.peek() == Some(&'/') {
out.push(chars.next().unwrap_or('/'));
in_block_comment = false;
}
prev = c;
continue;
}
if in_single {
out.push(c);
if c == '\'' {
if chars.peek() == Some(&'\'') {
out.push(chars.next().unwrap_or('\''));
} else {
in_single = false;
}
}
prev = c;
continue;
}
if in_double {
out.push(c);
if c == '"' {
if chars.peek() == Some(&'"') {
out.push(chars.next().unwrap_or('"'));
} else {
in_double = false;
}
}
prev = c;
continue;
}
if let Some(tag) = dollar_tag.as_deref() {
if c == '$' {
if let Some(extra) = pg_try_close_dollar_quote(&chars, tag) {
out.push(c);
for _ in 0..extra {
out.push(chars.next().unwrap_or('$'));
}
dollar_tag = None;
} else {
out.push(c);
}
} else {
out.push(c);
}
prev = c;
continue;
}
match c {
'\'' => {
in_single = true;
out.push(c);
}
'"' => {
in_double = true;
out.push(c);
}
'-' if chars.peek() == Some(&'-') => {
out.push(c);
out.push(chars.next().unwrap_or('-'));
in_line_comment = true;
}
'/' if chars.peek() == Some(&'*') => {
out.push(c);
out.push(chars.next().unwrap_or('*'));
in_block_comment = true;
}
'$' => {
if let Some(rest) = pg_peek_dollar_quote_rest(&chars) {
out.push('$');
for ch in rest.chars() {
out.push(ch);
chars.next();
}
dollar_tag = Some(rest[..rest.len() - 1].to_string());
} else if chars.peek().is_some_and(|next| next.is_ascii_digit()) {
saw_native = true;
out.push(c);
} else {
out.push(c);
}
}
'?' if prev != '@'
&& !matches!(chars.peek(), Some('|') | Some('&'))
&& !pg_question_is_json_operator(&chars) =>
{
n += 1;
out.push('$');
out.push_str(&n.to_string());
}
_ => out.push(c),
}
prev = c;
}
if saw_native {
return std::borrow::Cow::Borrowed(sql);
}
std::borrow::Cow::Owned(out)
}
fn pg_question_is_json_operator(chars: &std::iter::Peekable<std::str::Chars<'_>>) -> bool {
let mut lookahead = chars.clone();
while let Some(&c) = lookahead.peek() {
match c {
' ' | '\t' => {
lookahead.next();
}
'\'' | '"' => return true,
_ => return false,
}
}
false
}
fn pg_peek_dollar_quote_rest(chars: &std::iter::Peekable<std::str::Chars<'_>>) -> Option<String> {
let mut lookahead = chars.clone();
if lookahead.peek() == Some(&'$') {
return Some("$".to_string());
}
let first = *lookahead.peek()?;
if !(first.is_ascii_alphabetic() || first == '_') {
return None;
}
let mut rest = String::new();
while let Some(&c) = lookahead.peek() {
match c {
'$' => {
rest.push('$');
return Some(rest);
}
c if c.is_ascii_alphanumeric() || c == '_' => {
rest.push(c);
lookahead.next();
}
_ => return None,
}
}
None
}
fn pg_try_close_dollar_quote(
chars: &std::iter::Peekable<std::str::Chars<'_>>,
tag: &str,
) -> Option<usize> {
let mut lookahead = chars.clone();
for expected in tag.chars() {
if lookahead.next() != Some(expected) {
return None;
}
}
if lookahead.next() == Some('$') {
Some(tag.chars().count() + 1)
} else {
None
}
}
type PgQuery<'q> = sqlx::query::Query<'q, sqlx::Postgres, sqlx::postgres::PgArguments>;
fn bind_pg_value<'q>(q: PgQuery<'q>, v: &'q Value) -> PgQuery<'q> {
match v {
Value::Null => q.bind(None::<i64>),
Value::Bool(b) => q.bind(*b),
Value::I8(n) => q.bind(*n),
Value::I16(n) => q.bind(*n),
Value::I32(n) => q.bind(*n),
Value::I64(n) => q.bind(*n),
Value::U8(n) => q.bind(*n as i16),
Value::U16(n) => q.bind(*n as i32),
Value::U32(n) => q.bind(*n as i64),
Value::U64(n) => q.bind(*n as i64),
Value::F32(f) => q.bind(*f),
Value::F64(f) => q.bind(*f),
Value::String(s) => q.bind(s.as_str()),
Value::Decimal(s) => match s.parse::<bigdecimal::BigDecimal>() {
Ok(d) => q.bind(d),
Err(_) => q.bind(s.as_str()),
},
Value::Uuid(s) => match sqlx::types::Uuid::parse_str(s) {
Ok(u) => q.bind(u),
Err(_) => q.bind(s.as_str()),
},
Value::Date(s) => match parse_pg_date(s) {
Some(d) => q.bind(d),
None => q.bind(s.as_str()),
},
Value::DateTime(s) => match parse_pg_datetime(s) {
Some(PgDateTime::Zoned(d)) => q.bind(d),
Some(PgDateTime::Naive(d)) => q.bind(d),
None => q.bind(s.as_str()),
},
Value::Time(s) => match parse_pg_time(s) {
Some(t) => q.bind(t),
None => q.bind(s.as_str()),
},
Value::Json(s) => match serde_json::from_str::<serde_json::Value>(s) {
Ok(j) => q.bind(j),
Err(_) => q.bind(s.as_str()),
},
Value::Bytes(b) => q.bind(b.as_slice()),
other => q.bind(other.to_string()),
}
}
fn parse_pg_date(s: &str) -> Option<sqlx::types::chrono::NaiveDate> {
use sqlx::types::chrono::NaiveDate;
let head = s.get(..10)?;
NaiveDate::parse_from_str(head, "%Y-%m-%d").ok()
}
enum PgDateTime {
Naive(sqlx::types::chrono::NaiveDateTime),
Zoned(sqlx::types::chrono::DateTime<sqlx::types::chrono::Utc>),
}
fn parse_pg_datetime(s: &str) -> Option<PgDateTime> {
use sqlx::types::chrono::{DateTime, NaiveDateTime, Utc};
let has_offset = s.ends_with('Z')
|| s.ends_with('z')
|| s.get(10..).is_some_and(|rest| {
rest.char_indices().any(|(i, c)| {
(c == '+' || c == '-')
&& rest[..i].chars().any(|p| p == ':')
&& i > 0
})
});
if has_offset {
if let Ok(dt) = DateTime::parse_from_rfc3339(s) {
return Some(PgDateTime::Zoned(dt.with_timezone(&Utc)));
}
for fmt in [
"%Y-%m-%d %H:%M:%S%.f%:z",
"%Y-%m-%d %H:%M:%S%:z",
"%Y-%m-%dT%H:%M:%S%.f%:z",
"%Y-%m-%dT%H:%M:%S%:z",
"%Y-%m-%d %H:%M:%S%.f%z",
] {
if let Ok(dt) = DateTime::parse_from_str(s, fmt) {
return Some(PgDateTime::Zoned(dt.with_timezone(&Utc)));
}
}
}
for fmt in ["%Y-%m-%d %H:%M:%S", "%Y-%m-%dT%H:%M:%S"] {
if let Ok(dt) = NaiveDateTime::parse_from_str(s, fmt) {
return Some(PgDateTime::Naive(dt));
}
let with_frac = format!("{fmt}%.f");
if let Ok(dt) = NaiveDateTime::parse_from_str(s, &with_frac) {
return Some(PgDateTime::Naive(dt));
}
}
if let Ok(dt) = DateTime::parse_from_rfc3339(s) {
return Some(PgDateTime::Zoned(dt.with_timezone(&Utc)));
}
None
}
fn parse_pg_time(s: &str) -> Option<sqlx::types::chrono::NaiveTime> {
use sqlx::types::chrono::NaiveTime;
NaiveTime::parse_from_str(s, "%H:%M:%S")
.or_else(|_| NaiveTime::parse_from_str(s, "%H:%M:%S%.f"))
.ok()
}
pub struct PgPoolHandle {
pool: sqlx::PgPool,
}
impl PgPoolHandle {
pub async fn connect(url: &str) -> Result<Self, DbError> {
Self::connect_with_max(url, 10).await
}
pub async fn connect_with_max(url: &str, max_connections: u32) -> Result<Self, DbError> {
let pool = sqlx::pool::PoolOptions::<sqlx::Postgres>::new()
.max_connections(max_connections)
.acquire_timeout(std::time::Duration::from_secs(30))
.idle_timeout(Some(std::time::Duration::from_secs(600)))
.max_lifetime(Some(std::time::Duration::from_secs(1800)))
.connect(url)
.await
.map_err(map_sqlx_error)?;
Ok(Self { pool })
}
pub fn from_pool(pool: sqlx::PgPool) -> Self {
Self { pool }
}
pub fn pool(&self) -> &sqlx::PgPool {
&self.pool
}
}
pub struct SqlxPgConnectionFactory {
pool: Arc<PgPoolHandle>,
}
impl SqlxPgConnectionFactory {
pub fn new(pool: Arc<PgPoolHandle>) -> Self {
Self { pool }
}
}
#[async_trait]
impl ConnectionFactory for SqlxPgConnectionFactory {
async fn create(&self) -> Result<Box<dyn Connection>, DbError> {
let conn = self.pool.pool.acquire().await.map_err(map_sqlx_error)?;
Ok(Box::new(crate::metrics_conn::MetricsConn::new(Box::new(
SqlxPgConnection {
conn: Some(conn),
connected: true,
in_transaction: false,
},
))))
}
}
pub struct SqlxPgConnection {
conn: Option<sqlx::pool::PoolConnection<sqlx::Postgres>>,
connected: bool,
in_transaction: bool,
}
impl Connection for SqlxPgConnection {
fn db_type(&self) -> Option<DbType> {
Some(DbType::PostgreSQL)
}
fn session_reset_sql(&self) -> Option<&'static str> {
Some("DISCARD ALL")
}
fn execute<'a>(
&'a mut self,
sql: &'a str,
) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
Box::pin(async move {
let mut pool_conn = self
.conn
.take()
.ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
let result = if needs_raw_sql(sql) {
(&mut *pool_conn)
.execute(sqlx::raw_sql(sqlx::AssertSqlSafe(sql)))
.await
} else {
(&mut *pool_conn).execute(sqlx::AssertSqlSafe(sql)).await
};
self.conn = Some(pool_conn);
match result {
Ok(r) => Ok(r.rows_affected()),
Err(e) => {
let db_err = map_sqlx_error(e);
if matches!(db_err, DbError::ConnectionError(_) | DbError::IoError(_)) {
self.connected = false;
}
Err(db_err)
}
}
})
}
fn query<'a>(
&'a mut self,
sql: &'a str,
) -> Pin<Box<dyn Future<Output = Result<Vec<HashMap<String, Value>>, DbError>> + Send + 'a>>
{
Box::pin(async move {
let mut pool_conn = self
.conn
.take()
.ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
let rows_result = (&mut *pool_conn).fetch_all(sqlx::AssertSqlSafe(sql)).await;
self.conn = Some(pool_conn);
let rows = rows_result.map_err(map_sqlx_error)?;
if rows.is_empty() {
return Ok(Vec::new());
}
let col_types: Vec<ColType> = rows[0]
.columns()
.iter()
.map(|col| ColType::parse_postgres(col.type_info().name()))
.collect();
let col_names: Vec<String> = rows[0]
.columns()
.iter()
.map(|col| col.name().to_string())
.collect();
let mut result = Vec::with_capacity(rows.len());
for row in &rows {
let mut record = HashMap::with_capacity(col_types.len());
for (i, name) in col_names.iter().enumerate() {
let value = row_to_value_with_coltype_pg(row, i, col_types[i]);
record.insert(name.clone(), value);
}
result.push(record);
}
Ok(result)
})
}
fn begin_transaction<'a>(
&'a mut self,
) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
Box::pin(async move {
if self.in_transaction {
return Err(DbError::Internal("transaction already started".to_string()));
}
self.execute("BEGIN").await?;
self.in_transaction = true;
Ok(())
})
}
fn commit<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
Box::pin(async move {
if self.in_transaction {
self.execute("COMMIT").await?;
self.in_transaction = false;
}
Ok(())
})
}
fn rollback<'a>(
&'a mut self,
) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
Box::pin(async move {
if self.in_transaction {
let result = self.execute("ROLLBACK").await;
if result.is_ok() {
self.in_transaction = false;
}
result.map(|_| ())
} else {
Ok(())
}
})
}
fn is_connected(&self) -> bool {
self.connected && self.conn.is_some()
}
fn in_transaction(&self) -> bool {
self.in_transaction
}
fn ping<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = bool> + Send + 'a>> {
Box::pin(async move {
match self.execute("SELECT 1").await {
Ok(_) => true,
Err(_) => {
self.connected = false;
false
}
}
})
}
fn close<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), DbError>> + Send + 'a>> {
Box::pin(async move {
if let Some(conn) = self.conn.take() {
drop(conn);
}
self.connected = false;
self.in_transaction = false;
Ok(())
})
}
fn execute_with_params<'a>(
&'a mut self,
sql: &'a str,
params: &'a [Value],
) -> Pin<Box<dyn Future<Output = Result<u64, DbError>> + Send + 'a>> {
Box::pin(async move {
if needs_raw_sql(sql) || params.is_empty() {
return self.execute(sql).await;
}
let mut pool_conn = self
.conn
.take()
.ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
let rewritten = pg_rewrite_placeholders(sql);
let mut q = sqlx::query(sqlx::AssertSqlSafe(rewritten.as_ref()));
for v in params {
q = bind_pg_value(q, v);
}
let result = q.execute(&mut *pool_conn).await;
self.conn = Some(pool_conn);
match result {
Ok(r) => Ok(r.rows_affected()),
Err(e) => {
let db_err = map_sqlx_error(e);
if matches!(db_err, DbError::ConnectionError(_) | DbError::IoError(_)) {
self.connected = false;
}
Err(db_err)
}
}
})
}
fn query_with_params<'a>(
&'a mut self,
sql: &'a str,
params: &'a [Value],
) -> Pin<Box<dyn Future<Output = Result<QueryRows, DbError>> + Send + 'a>> {
Box::pin(async move {
if params.is_empty() {
return self.query(sql).await;
}
let mut pool_conn = self
.conn
.take()
.ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
let rewritten = pg_rewrite_placeholders(sql);
let mut q = sqlx::query(sqlx::AssertSqlSafe(rewritten.as_ref()));
for v in params {
q = bind_pg_value(q, v);
}
let rows_result = q.fetch_all(&mut *pool_conn).await;
self.conn = Some(pool_conn);
let rows = rows_result.map_err(map_sqlx_error)?;
if rows.is_empty() {
return Ok(Vec::new());
}
let col_types: Vec<ColType> = rows[0]
.columns()
.iter()
.map(|col| ColType::parse_postgres(col.type_info().name()))
.collect();
let col_names: Vec<String> = rows[0]
.columns()
.iter()
.map(|col| col.name().to_string())
.collect();
let mut result = Vec::with_capacity(rows.len());
for row in &rows {
let mut record = HashMap::with_capacity(col_types.len());
for (i, name) in col_names.iter().enumerate() {
let value = row_to_value_with_coltype_pg(row, i, col_types[i]);
record.insert(name.clone(), value);
}
result.push(record);
}
Ok(result)
})
}
fn query_values<'a>(
&'a mut self,
sql: &'a str,
) -> Pin<Box<dyn Future<Output = Result<QueryValues, DbError>> + Send + 'a>> {
Box::pin(async move {
let mut pool_conn = self
.conn
.take()
.ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
let rows_result = (&mut *pool_conn).fetch_all(sqlx::AssertSqlSafe(sql)).await;
self.conn = Some(pool_conn);
let rows = rows_result.map_err(map_sqlx_error)?;
if rows.is_empty() {
return Ok((Vec::new(), Vec::new()));
}
let cols = rows[0].columns();
let mut col_names: Vec<String> = Vec::with_capacity(cols.len());
let mut col_types: Vec<ColType> = Vec::with_capacity(cols.len());
for col in cols {
col_names.push(col.name().to_string());
col_types.push(ColType::parse_postgres(col.type_info().name()));
}
let mut result_rows: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
for row in &rows {
let mut row_values: Vec<Value> = Vec::with_capacity(col_names.len());
for (idx, _) in col_names.iter().enumerate() {
row_values.push(row_to_value_with_coltype_pg(row, idx, col_types[idx]));
}
result_rows.push(row_values);
}
Ok((col_names, result_rows))
})
}
fn query_values_with_params<'a>(
&'a mut self,
sql: &'a str,
params: &'a [Value],
) -> Pin<Box<dyn Future<Output = Result<QueryValues, DbError>> + Send + 'a>> {
Box::pin(async move {
if params.is_empty() {
return self.query_values(sql).await;
}
let mut pool_conn = self
.conn
.take()
.ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
let rewritten = pg_rewrite_placeholders(sql);
let mut q = sqlx::query(sqlx::AssertSqlSafe(rewritten.as_ref()));
for v in params {
q = bind_pg_value(q, v);
}
let rows_result = q.fetch_all(&mut *pool_conn).await;
self.conn = Some(pool_conn);
let rows = rows_result.map_err(map_sqlx_error)?;
if rows.is_empty() {
return Ok((Vec::new(), Vec::new()));
}
let cols = rows[0].columns();
let mut col_names: Vec<String> = Vec::with_capacity(cols.len());
for col in cols {
col_names.push(col.name().to_string());
}
let mut result_rows: Vec<Vec<Value>> = Vec::with_capacity(rows.len());
for row in rows {
let mut row_values: Vec<Value> = Vec::with_capacity(col_names.len());
for (idx, _) in col_names.iter().enumerate() {
let ordinal = row.columns()[idx].ordinal();
row_values.push(row_to_value_pg(&row, ordinal)?);
}
result_rows.push(row_values);
}
Ok((col_names, result_rows))
})
}
fn query_stream<'a>(
&'a mut self,
sql: &'a str,
) -> Pin<Box<dyn futures::Stream<Item = Result<HashMap<String, Value>, DbError>> + Send + 'a>>
{
Box::pin(async_stream::try_stream! {
let mut pool_conn = self
.conn
.take()
.ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
let mut row_stream = sqlx::query(sqlx::AssertSqlSafe(sql)).fetch(&mut *pool_conn);
while let Some(row_result) = row_stream.next().await {
let row = row_result.map_err(map_sqlx_error)?;
let cols = row.columns();
let mut record = HashMap::with_capacity(cols.len());
for (i, col) in cols.iter().enumerate() {
let name = col.name().to_string();
let value = row_to_value_pg(&row, i)?;
record.insert(name, value);
}
yield record;
}
drop(row_stream);
self.conn = Some(pool_conn);
})
}
}
impl Drop for SqlxPgConnection {
fn drop(&mut self) {
if let Some(conn) = self.conn.take() {
drop(conn);
}
}
}
#[async_trait]
pub trait PgExtensions: Send + Sync {
async fn listen(&mut self, channel: &str) -> Result<(), DbError>;
async fn notify(&mut self, channel: &str, payload: &str) -> Result<(), DbError>;
async fn copy_from_stdin(&mut self, sql: &str, data: &[u8]) -> Result<u64, DbError>;
}
fn validate_pg_channel_name(channel: &str) -> Result<(), DbError> {
if channel.is_empty() {
return Err(DbError::Internal(
"PG channel name must not be empty".to_string(),
));
}
if !channel
.chars()
.all(|c| c.is_ascii_alphanumeric() || c == '_')
{
return Err(DbError::Internal(format!(
"invalid PG channel name: {} (only alphanumeric and underscore allowed)",
channel
)));
}
Ok(())
}
#[async_trait]
impl PgExtensions for SqlxPgConnection {
async fn listen(&mut self, channel: &str) -> Result<(), DbError> {
validate_pg_channel_name(channel)?;
self.execute(&format!("LISTEN {}", channel)).await?;
Ok(())
}
async fn notify(&mut self, channel: &str, payload: &str) -> Result<(), DbError> {
validate_pg_channel_name(channel)?;
let escaped_payload = payload.replace('\'', "''");
self.execute(&format!("NOTIFY {}, '{}'", channel, escaped_payload))
.await?;
Ok(())
}
async fn copy_from_stdin(&mut self, sql: &str, data: &[u8]) -> Result<u64, DbError> {
let mut pool_conn = self
.conn
.take()
.ok_or_else(|| DbError::Internal("connection already closed".to_string()))?;
let mut copy = (*pool_conn)
.copy_in_raw(sql)
.await
.map_err(map_sqlx_error)?;
copy.send(data).await.map_err(map_sqlx_error)?;
let result = copy.finish().await.map_err(map_sqlx_error)?;
self.conn = Some(pool_conn);
Ok(result)
}
}
pub async fn pg_bulk_insert(
conn: &mut SqlxPgConnection,
table: &str,
columns: &[&str],
rows: &[Vec<Value>],
) -> Result<u64, DbError> {
if rows.is_empty() {
return Ok(0);
}
let col_list = columns.join(", ");
let cols_per_row = columns.len();
let placeholders: Vec<String> = rows
.iter()
.enumerate()
.map(|(row_idx, _)| {
let base = row_idx * cols_per_row;
let ph: Vec<String> = (0..cols_per_row)
.map(|i| format!("${}", base + i + 1))
.collect();
format!("({})", ph.join(", "))
})
.collect();
let sql = format!(
"INSERT INTO {} ({}) VALUES {}",
table,
col_list,
placeholders.join(", ")
);
let mut params: Vec<Value> = Vec::with_capacity(rows.len() * cols_per_row);
for row in rows {
for v in row {
params.push(v.clone());
}
}
conn.execute_with_params(&sql, ¶ms).await
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_pg_rewrite_placeholders_basic() {
assert_eq!(
pg_rewrite_placeholders("SELECT * FROM t WHERE a = ? AND b = ?"),
"SELECT * FROM t WHERE a = $1 AND b = $2"
);
}
#[test]
fn test_pg_rewrite_placeholders_skips_native() {
assert_eq!(
pg_rewrite_placeholders("SELECT * FROM t WHERE a = $1"),
"SELECT * FROM t WHERE a = $1"
);
}
#[test]
fn test_pg_rewrite_placeholders_skips_literals_and_comments() {
assert_eq!(
pg_rewrite_placeholders("SELECT '?' AS q -- ?\nFROM t WHERE id = ?"),
"SELECT '?' AS q -- ?\nFROM t WHERE id = $1"
);
assert_eq!(
pg_rewrite_placeholders("SELECT /* ? */ \"a?\" FROM t WHERE id = ?"),
"SELECT /* ? */ \"a?\" FROM t WHERE id = $1"
);
}
#[test]
fn test_pg_rewrite_placeholders_keeps_json_operators() {
assert_eq!(
pg_rewrite_placeholders("SELECT * FROM t WHERE doc ? 'k' AND b = ?"),
"SELECT * FROM t WHERE doc ? 'k' AND b = $1"
);
assert_eq!(
pg_rewrite_placeholders("SELECT * FROM t WHERE doc ?| array['a'] AND b = ?"),
"SELECT * FROM t WHERE doc ?| array['a'] AND b = $1"
);
assert_eq!(
pg_rewrite_placeholders("SELECT * FROM t WHERE doc ?& array['a', 'b'] AND b = ?"),
"SELECT * FROM t WHERE doc ?& array['a', 'b'] AND b = $1"
);
assert_eq!(
pg_rewrite_placeholders("SELECT * FROM t WHERE doc @? '$.a' AND b = ?"),
"SELECT * FROM t WHERE doc @? '$.a' AND b = $1"
);
}
#[test]
fn test_pg_rewrite_placeholders_dollar_in_literal_still_rewrites() {
assert_eq!(
pg_rewrite_placeholders("SELECT '$1' AS lit, ? AS val"),
"SELECT '$1' AS lit, $1 AS val"
);
assert_eq!(
pg_rewrite_placeholders("SELECT 'price: $2' AS note FROM t WHERE id = ? AND code = ?"),
"SELECT 'price: $2' AS note FROM t WHERE id = $1 AND code = $2"
);
}
#[test]
fn test_pg_rewrite_placeholders_skips_dollar_quoted_strings() {
assert_eq!(
pg_rewrite_placeholders("SELECT $$?$$ AS lit, ? AS val"),
"SELECT $$?$$ AS lit, $1 AS val"
);
assert_eq!(
pg_rewrite_placeholders("SELECT $tag$a ? b$tag$ AS lit, ? AS val"),
"SELECT $tag$a ? b$tag$ AS lit, $1 AS val"
);
let func = "CREATE FUNCTION f() RETURNS int AS $$ BEGIN RETURN (SELECT count(*) FROM t WHERE x = ?); END; $$ LANGUAGE plpgsql";
assert_eq!(pg_rewrite_placeholders(func), func);
}
#[test]
fn test_pg_rewrite_placeholders_native_skips_even_with_question() {
assert_eq!(
pg_rewrite_placeholders("SELECT * FROM t WHERE a = $1 AND doc ? 'k'"),
"SELECT * FROM t WHERE a = $1 AND doc ? 'k'"
);
}
#[test]
fn test_parse_pg_date_and_datetime() {
assert!(parse_pg_date("2026-10-28").is_some());
assert!(parse_pg_date("2026-10-28 12:00:00").is_some());
assert!(parse_pg_date("not-a-date").is_none());
assert!(parse_pg_datetime("2026-10-28 12:00:00").is_some());
assert!(parse_pg_datetime("2026-10-28T12:00:00").is_some());
assert!(parse_pg_datetime("2026-10-28T12:00:00Z").is_some());
assert!(parse_pg_datetime("bad").is_none());
}
#[test]
fn test_needs_raw_sql_begin() {
assert!(needs_raw_sql("BEGIN"));
assert!(needs_raw_sql("begin"));
assert!(needs_raw_sql(" BEGIN"));
assert!(needs_raw_sql("\tbegin"));
}
#[test]
fn test_needs_raw_sql_commit() {
assert!(needs_raw_sql("COMMIT"));
assert!(needs_raw_sql("commit"));
assert!(needs_raw_sql(" COMMIT WORK"));
}
#[test]
fn test_needs_raw_sql_rollback() {
assert!(needs_raw_sql("ROLLBACK"));
assert!(needs_raw_sql("rollback"));
assert!(needs_raw_sql(" ROLLBACK TO SAVEPOINT sp1"));
}
#[test]
fn test_needs_raw_sql_savepoint() {
assert!(needs_raw_sql("SAVEPOINT sp1"));
assert!(needs_raw_sql("savepoint s1"));
}
#[test]
fn test_needs_raw_sql_release() {
assert!(needs_raw_sql("RELEASE sp1"));
assert!(needs_raw_sql("release savepoint s1"));
}
#[test]
fn test_needs_raw_sql_set() {
assert!(needs_raw_sql("SET autocommit = 0"));
assert!(needs_raw_sql("set names utf8"));
}
#[test]
fn test_needs_raw_sql_use() {
assert!(needs_raw_sql("USE mydb"));
assert!(needs_raw_sql("use mydb"));
}
#[test]
fn test_needs_raw_sql_start_transaction() {
assert!(needs_raw_sql("START TRANSACTION"));
assert!(needs_raw_sql("start transaction"));
assert!(needs_raw_sql(" START TRANSACTION READ ONLY"));
}
#[test]
fn test_needs_raw_sql_normal_queries() {
assert!(!needs_raw_sql("SELECT * FROM users"));
assert!(!needs_raw_sql("INSERT INTO t VALUES (1)"));
assert!(!needs_raw_sql("UPDATE t SET x = 1"));
assert!(!needs_raw_sql("DELETE FROM t"));
assert!(!needs_raw_sql("CREATE TABLE t (id INT)"));
assert!(!needs_raw_sql("DROP TABLE t"));
assert!(!needs_raw_sql("ALTER TABLE t ADD COLUMN x INT"));
}
#[test]
fn test_needs_raw_sql_edge_cases() {
assert!(!needs_raw_sql(""));
assert!(!needs_raw_sql(" "));
assert!(!needs_raw_sql("SELECT"));
assert!(needs_raw_sql("BEGINNING"));
assert!(needs_raw_sql("COMMITTEE"));
assert!(!needs_raw_sql("INSERT"));
assert!(!needs_raw_sql("DELETE"));
}
#[test]
fn test_validate_pg_channel_name_valid() {
assert!(validate_pg_channel_name("my_channel").is_ok());
assert!(validate_pg_channel_name("abc123").is_ok());
assert!(validate_pg_channel_name("_test").is_ok());
assert!(validate_pg_channel_name("a_b_c").is_ok());
}
#[test]
fn test_validate_pg_channel_name_empty() {
let err = validate_pg_channel_name("").unwrap_err();
assert!(matches!(err, DbError::Internal(ref s) if s.contains("empty")));
}
#[test]
fn test_validate_pg_channel_name_invalid_chars() {
let err = validate_pg_channel_name("my-channel").unwrap_err();
assert!(matches!(err, DbError::Internal(ref s) if s.contains("invalid")));
assert!(validate_pg_channel_name("my.channel").is_err());
assert!(validate_pg_channel_name("my channel").is_err());
assert!(validate_pg_channel_name("my;channel").is_err());
assert!(validate_pg_channel_name("my'channel").is_err());
}
fn closed_sqlite() -> SqlxSqliteConnection {
SqlxSqliteConnection {
conn: None,
connected: false,
in_transaction: false,
}
}
#[tokio::test]
async fn test_sqlite_execute_closed() {
let mut c = closed_sqlite();
let result = c.execute("SELECT 1").await;
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("closed"));
}
#[tokio::test]
async fn test_sqlite_query_closed() {
let mut c = closed_sqlite();
let result = c.query("SELECT 1").await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_sqlite_begin_tx_already_started() {
let mut c = SqlxSqliteConnection {
conn: None,
connected: false,
in_transaction: true,
};
let result = c.begin_transaction().await;
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("already started"));
}
#[tokio::test]
async fn test_sqlite_begin_tx_closed() {
let mut c = closed_sqlite();
let result = c.begin_transaction().await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_sqlite_commit_no_tx() {
let mut c = closed_sqlite();
let result = c.commit().await;
assert!(result.is_ok());
}
#[tokio::test]
async fn test_sqlite_rollback_no_tx() {
let mut c = closed_sqlite();
let result = c.rollback().await;
assert!(result.is_ok());
}
#[test]
fn test_sqlite_is_connected_false_when_conn_none() {
let c = closed_sqlite();
assert!(!c.is_connected());
}
#[test]
fn test_sqlite_is_connected_false_when_connected_false() {
let c = SqlxSqliteConnection {
conn: None,
connected: true,
in_transaction: false,
};
assert!(!c.is_connected());
}
#[test]
fn test_sqlite_in_transaction_flag() {
let c = SqlxSqliteConnection {
conn: None,
connected: false,
in_transaction: true,
};
assert!(c.in_transaction());
let c2 = closed_sqlite();
assert!(!c2.in_transaction());
}
#[tokio::test]
async fn test_sqlite_ping_closed() {
let mut c = closed_sqlite();
let result = c.ping().await;
assert!(!result);
assert!(!c.is_connected());
}
#[tokio::test]
async fn test_sqlite_close_closed() {
let mut c = closed_sqlite();
let result = c.close().await;
assert!(result.is_ok());
assert!(!c.is_connected());
assert!(!c.in_transaction());
}
#[tokio::test]
async fn test_sqlite_execute_with_params_closed() {
let mut c = closed_sqlite();
let result = c
.execute_with_params("INSERT INTO t VALUES (?)", &[Value::I64(1)])
.await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_sqlite_execute_with_params_empty_delegates() {
let mut c = closed_sqlite();
let result = c.execute_with_params("INSERT INTO t VALUES (1)", &[]).await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_sqlite_query_with_params_empty_delegates() {
let mut c = closed_sqlite();
let result = c.query_with_params("SELECT * FROM t", &[]).await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_sqlite_query_with_params_closed() {
let mut c = closed_sqlite();
let result = c
.query_with_params("SELECT * FROM t WHERE id = ?", &[Value::I64(1)])
.await;
assert!(result.is_err());
}
fn closed_mysql() -> SqlxMySqlConnection {
SqlxMySqlConnection {
conn: None,
connected: false,
in_transaction: false,
}
}
#[tokio::test]
async fn test_mysql_execute_closed() {
let mut c = closed_mysql();
assert!(c.execute("SELECT 1").await.is_err());
}
#[tokio::test]
async fn test_mysql_query_closed() {
let mut c = closed_mysql();
assert!(c.query("SELECT 1").await.is_err());
}
#[test]
fn test_mysql_is_connected_false() {
let c = closed_mysql();
assert!(!c.is_connected());
}
#[tokio::test]
async fn test_mysql_close_closed() {
let mut c = closed_mysql();
assert!(c.close().await.is_ok());
assert!(!c.is_connected());
}
#[tokio::test]
async fn test_mysql_commit_no_tx() {
let mut c = closed_mysql();
assert!(c.commit().await.is_ok());
}
#[tokio::test]
async fn test_mysql_rollback_no_tx() {
let mut c = closed_mysql();
assert!(c.rollback().await.is_ok());
}
#[tokio::test]
async fn test_mysql_begin_tx_already_started() {
let mut c = SqlxMySqlConnection {
conn: None,
connected: false,
in_transaction: true,
};
let result = c.begin_transaction().await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_mysql_ping_closed() {
let mut c = closed_mysql();
assert!(!c.ping().await);
}
fn closed_pg() -> SqlxPgConnection {
SqlxPgConnection {
conn: None,
connected: false,
in_transaction: false,
}
}
#[tokio::test]
async fn test_pg_execute_closed() {
let mut c = closed_pg();
assert!(c.execute("SELECT 1").await.is_err());
}
#[tokio::test]
async fn test_pg_query_closed() {
let mut c = closed_pg();
assert!(c.query("SELECT 1").await.is_err());
}
#[test]
fn test_pg_is_connected_false() {
let c = closed_pg();
assert!(!c.is_connected());
}
#[tokio::test]
async fn test_pg_close_closed() {
let mut c = closed_pg();
assert!(c.close().await.is_ok());
assert!(!c.is_connected());
}
#[tokio::test]
async fn test_pg_commit_no_tx() {
let mut c = closed_pg();
assert!(c.commit().await.is_ok());
}
#[tokio::test]
async fn test_pg_rollback_no_tx() {
let mut c = closed_pg();
assert!(c.rollback().await.is_ok());
}
#[tokio::test]
async fn test_pg_begin_tx_already_started() {
let mut c = SqlxPgConnection {
conn: None,
connected: false,
in_transaction: true,
};
assert!(c.begin_transaction().await.is_err());
}
#[tokio::test]
async fn test_pg_ping_closed() {
let mut c = closed_pg();
assert!(!c.ping().await);
}
#[tokio::test]
async fn test_mysql_execute_with_params_closed() {
let mut c = closed_mysql();
assert!(c
.execute_with_params("INSERT INTO t VALUES (?)", &[Value::I64(1)])
.await
.is_err());
}
#[tokio::test]
async fn test_mysql_execute_with_params_empty_delegates() {
let mut c = closed_mysql();
assert!(c
.execute_with_params("INSERT INTO t VALUES (1)", &[])
.await
.is_err());
}
#[tokio::test]
async fn test_mysql_query_with_params_closed() {
let mut c = closed_mysql();
assert!(c
.query_with_params("SELECT * FROM t WHERE id = ?", &[Value::I64(1)])
.await
.is_err());
}
#[tokio::test]
async fn test_mysql_query_with_params_empty_delegates() {
let mut c = closed_mysql();
assert!(c.query_with_params("SELECT * FROM t", &[]).await.is_err());
}
#[tokio::test]
async fn test_mysql_query_values_closed() {
let mut c = closed_mysql();
assert!(c.query_values("SELECT * FROM t").await.is_err());
}
#[tokio::test]
async fn test_mysql_query_values_with_params_closed() {
let mut c = closed_mysql();
assert!(c
.query_values_with_params("SELECT * FROM t WHERE id = ?", &[Value::I64(1)])
.await
.is_err());
}
#[tokio::test]
async fn test_mysql_query_values_with_params_empty_delegates() {
let mut c = closed_mysql();
assert!(c
.query_values_with_params("SELECT * FROM t", &[])
.await
.is_err());
}
#[tokio::test]
async fn test_pg_execute_with_params_closed() {
let mut c = closed_pg();
assert!(c
.execute_with_params("INSERT INTO t VALUES ($1)", &[Value::I64(1)])
.await
.is_err());
}
#[tokio::test]
async fn test_pg_execute_with_params_empty_delegates() {
let mut c = closed_pg();
assert!(c
.execute_with_params("INSERT INTO t VALUES (1)", &[])
.await
.is_err());
}
#[tokio::test]
async fn test_pg_query_with_params_closed() {
let mut c = closed_pg();
assert!(c
.query_with_params("SELECT * FROM t WHERE id = $1", &[Value::I64(1)])
.await
.is_err());
}
#[tokio::test]
async fn test_pg_query_with_params_empty_delegates() {
let mut c = closed_pg();
assert!(c.query_with_params("SELECT * FROM t", &[]).await.is_err());
}
#[tokio::test]
async fn test_pg_query_values_closed() {
let mut c = closed_pg();
assert!(c.query_values("SELECT * FROM t").await.is_err());
}
#[tokio::test]
async fn test_pg_query_values_with_params_closed() {
let mut c = closed_pg();
assert!(c
.query_values_with_params("SELECT * FROM t WHERE id = $1", &[Value::I64(1)])
.await
.is_err());
}
#[tokio::test]
async fn test_pg_query_values_with_params_empty_delegates() {
let mut c = closed_pg();
assert!(c
.query_values_with_params("SELECT * FROM t", &[])
.await
.is_err());
}
#[tokio::test]
async fn test_sqlite_backup_closed_conn() {
let mut c = closed_sqlite();
let result = sqlite_backup(&mut c, "/tmp/backup.db").await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_sqlite_query_values_closed() {
let mut c = closed_sqlite();
assert!(c.query_values("SELECT * FROM t").await.is_err());
}
#[tokio::test]
async fn test_sqlite_query_values_with_params_closed() {
let mut c = closed_sqlite();
assert!(c
.query_values_with_params("SELECT * FROM t WHERE id = ?", &[Value::I64(1)])
.await
.is_err());
}
#[tokio::test]
async fn test_sqlite_query_values_with_params_empty_delegates() {
let mut c = closed_sqlite();
assert!(c
.query_values_with_params("SELECT * FROM t", &[])
.await
.is_err());
}
#[tokio::test]
async fn test_mysql_bulk_insert_empty_rows() {
let mut c = closed_mysql();
let result = mysql_bulk_insert(&mut c, "t", &["a", "b"], &[]).await;
assert_eq!(result.unwrap(), 0);
}
#[tokio::test]
async fn test_mysql_bulk_insert_closed_conn() {
let mut c = closed_mysql();
let rows = vec![vec![Value::I64(1), Value::String("x".to_string())]];
let result = mysql_bulk_insert(&mut c, "t", &["a", "b"], &rows).await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_pg_bulk_insert_empty_rows() {
let mut c = closed_pg();
let result = pg_bulk_insert(&mut c, "t", &["a", "b"], &[]).await;
assert_eq!(result.unwrap(), 0);
}
#[tokio::test]
async fn test_pg_bulk_insert_closed_conn() {
let mut c = closed_pg();
let rows = vec![vec![Value::I64(1), Value::String("x".to_string())]];
let result = pg_bulk_insert(&mut c, "t", &["a", "b"], &rows).await;
assert!(result.is_err());
}
#[test]
fn test_parse_pg_datetime_naive() {
match parse_pg_datetime("2026-10-28 13:14:15.123456") {
Some(PgDateTime::Naive(dt)) => {
assert_eq!(
dt.format("%Y-%m-%d %H:%M:%S%.6f").to_string(),
"2026-10-28 13:14:15.123456"
);
}
_ => panic!("无偏移输入必须解析为 Naive"),
}
}
#[test]
fn test_parse_pg_datetime_zoned_forms() {
for s in [
"2026-10-28T13:14:15+08:00",
"2026-10-28 13:14:15+08:00",
"2026-10-28 13:14:15 +0800",
"2026-10-28T05:14:15Z",
] {
match parse_pg_datetime(s) {
Some(PgDateTime::Zoned(dt)) => {
assert_eq!(
dt.format("%Y-%m-%d %H:%M:%S").to_string(),
"2026-10-28 05:14:15",
"输入 {s}"
);
}
other => panic!("{s} 必须解析为 Zoned,实际 {:?}", other.is_some()),
}
}
}
#[test]
fn test_parse_pg_time_fractional() {
assert!(parse_pg_time("13:14:15").is_some());
assert!(parse_pg_time("13:14:15.123456").is_some());
assert!(parse_pg_time("13:14").is_none());
}
#[test]
fn test_bigdecimal_38_digits_no_truncation() {
let s = "1234567890123456789012345678.9012345678"; let d: bigdecimal::BigDecimal = s.parse().unwrap();
assert_eq!(d.to_string(), s, "BigDecimal 往返必须逐位一致");
assert!(
s.parse::<rust_decimal::Decimal>().is_err() || {
let rd: rust_decimal::Decimal = s.parse().unwrap();
rd.to_string() != s
}
);
}
}