use crate::DbkitError;
use crate::base_handler::{FetchMode, QueryResult, WriteOp};
use crate::value::DbValue;
use std::fmt::Write as _;
use sqlx::postgres::{PgArguments, PgRow};
use sqlx::query::Query;
use sqlx::{AssertSqlSafe, PgPool, Postgres, Row as _};
use tracing::warn;
use unicode_normalization::UnicodeNormalization;
#[cfg(feature = "duckdb")]
use crate::analytical::RecordBatch;
#[cfg(feature = "duckdb")]
use crate::read::{ReadEngine, duckdb::DuckEngine};
struct PgNull;
impl sqlx::Type<Postgres> for PgNull {
fn type_info() -> sqlx::postgres::PgTypeInfo {
sqlx::postgres::PgTypeInfo::with_oid(sqlx::postgres::types::Oid(0))
}
}
impl<'q> sqlx::Encode<'q, Postgres> for PgNull {
fn encode_by_ref(
&self,
_buf: &mut sqlx::postgres::PgArgumentBuffer,
) -> Result<sqlx::encode::IsNull, sqlx::error::BoxDynError> {
Ok(sqlx::encode::IsNull::Yes)
}
}
fn bind_pg<'q>(
mut q: Query<'q, Postgres, PgArguments>,
params: &'q [DbValue],
) -> Query<'q, Postgres, PgArguments> {
for p in params {
q = match p {
DbValue::Null => q.bind(PgNull),
DbValue::Bool(b) => q.bind(*b),
DbValue::Int(i) => q.bind(*i),
DbValue::Float(f) => q.bind(*f),
DbValue::Text(s) => q.bind(s.as_str()),
DbValue::Bytes(b) => q.bind(b.as_slice()),
DbValue::Date(d) => q.bind(*d),
DbValue::DateTime(dt) => q.bind(*dt),
DbValue::TimestampTz(dt) => q.bind(*dt),
DbValue::Json(j) => q.bind(j),
DbValue::Uuid(u) => q.bind(*u),
DbValue::Time(t) => q.bind(*t),
DbValue::TextArray(v) => q.bind(v),
DbValue::FloatArray(v) => q.bind(v),
DbValue::OptFloatArray(v) => q.bind(v),
};
}
q
}
fn copy_render_cell(val: &DbValue, out: &mut String) {
match val {
DbValue::Null => out.push_str("\\N"),
DbValue::Bool(b) => out.push(if *b { 't' } else { 'f' }),
DbValue::Int(i) => {
let _ = write!(out, "{i}");
}
DbValue::Float(f) => {
if f.is_nan() {
out.push_str("NaN");
} else if f.is_infinite() {
out.push_str(if *f > 0.0 { "Infinity" } else { "-Infinity" });
} else {
let _ = write!(out, "{f}");
}
}
DbValue::Text(s) => copy_escape_into(s, out),
DbValue::Bytes(b) => {
out.push_str("\\\\x");
for byte in b {
out.push(char::from_digit((byte >> 4) as u32, 16).unwrap());
out.push(char::from_digit((byte & 0x0f) as u32, 16).unwrap());
}
}
DbValue::Date(d) => copy_escape_into(&d.to_string(), out),
DbValue::DateTime(dt) => copy_escape_into(&dt.to_string(), out),
DbValue::TimestampTz(dt) => copy_escape_into(&dt.to_rfc3339(), out),
DbValue::Json(j) => copy_escape_into(&j.to_string(), out),
DbValue::Uuid(u) => copy_escape_into(&u.to_string(), out),
DbValue::Time(t) => copy_escape_into(&t.to_string(), out),
DbValue::TextArray(v) => copy_escape_into(&crate::value::pg_text_array_literal(v), out),
DbValue::FloatArray(v) => {
copy_escape_into(&crate::value::pg_float_array_literal(v.iter().map(|x| Some(*x))), out)
}
DbValue::OptFloatArray(v) => {
copy_escape_into(&crate::value::pg_float_array_literal(v.iter().copied()), out)
}
}
}
const COPY_CHUNK_BYTES: usize = 4 * 1024 * 1024;
async fn send_copy_rows<C>(
sink: &mut sqlx::postgres::PgCopyIn<C>,
rows: &[Vec<DbValue>],
ncols: usize,
) -> Result<(), DbkitError>
where
C: std::ops::DerefMut<Target = sqlx::postgres::PgConnection>,
{
let estimate = rows.len() * (ncols * 12 + 1);
let mut buf = String::with_capacity(estimate.min(COPY_CHUNK_BYTES + 1024));
for row in rows {
for (i, val) in row.iter().enumerate() {
if i > 0 {
buf.push('\t');
}
copy_render_cell(val, &mut buf);
}
buf.push('\n');
if buf.len() >= COPY_CHUNK_BYTES {
sink.send(buf.as_bytes()).await?;
buf.clear();
}
}
if !buf.is_empty() {
sink.send(buf.as_bytes()).await?;
}
Ok(())
}
fn copy_escape_into(s: &str, out: &mut String) {
let mut start = 0;
for (i, b) in s.bytes().enumerate() {
let esc = match b {
b'\\' => "\\\\",
b'\t' => "\\t",
b'\n' => "\\n",
b'\r' => "\\r",
_ => continue,
};
out.push_str(&s[start..i]);
out.push_str(esc);
start = i + 1;
}
out.push_str(&s[start..]);
}
pub struct PgHandler {
pool: PgPool,
#[cfg(feature = "duckdb")]
duck: Option<DuckEngine>,
}
impl PgHandler {
pub fn new(pool: PgPool) -> Self {
Self {
pool,
#[cfg(feature = "duckdb")]
duck: None,
}
}
#[cfg(feature = "duckdb")]
pub fn with_duckdb(pool: PgPool) -> Result<Self, DbkitError> {
Ok(Self {
pool,
duck: Some(DuckEngine::new_in_memory()?),
})
}
#[cfg(feature = "duckdb")]
pub fn with_duckdb_attached_postgres(
pool: PgPool,
pg_connection_string: &str,
) -> Result<Self, DbkitError> {
let duck = DuckEngine::new_in_memory()?;
duck.attach_postgres(pg_connection_string)?;
Ok(Self {
pool,
duck: Some(duck),
})
}
pub fn has_read_engine(&self) -> bool {
#[cfg(feature = "duckdb")]
{
self.duck.is_some()
}
#[cfg(not(feature = "duckdb"))]
{
false
}
}
pub fn pool(&self) -> &PgPool {
&self.pool
}
pub fn normalize_name(name: &str) -> String {
use unicode_normalization::char::is_combining_mark;
name.nfd()
.filter(|c| !is_combining_mark(*c))
.collect::<String>()
.to_lowercase()
}
pub async fn execute_write(
&self,
op: WriteOp<'_>,
) -> Result<QueryResult<PgRow>, DbkitError> {
match op {
WriteOp::Single {
query,
params,
mode,
} => self.query(query, params, mode).await,
WriteOp::BatchDDL { queries } => {
let mut tx = self.pool.begin().await?;
for query in queries {
sqlx::query(AssertSqlSafe(*query)).execute(&mut *tx).await?;
}
tx.commit().await?;
Ok(QueryResult::None)
}
WriteOp::BatchParams {
query,
params_list,
isolate_rows,
} => {
if params_list.is_empty() {
return Ok(QueryResult::None);
}
let total = params_list.len();
let mut tx = self.pool.begin().await?;
if !isolate_rows {
for params in ¶ms_list {
let has_null = params.iter().any(|v| matches!(v, DbValue::Null));
bind_pg(sqlx::query(AssertSqlSafe(query)), params)
.persistent(!has_null)
.execute(&mut *tx)
.await?;
}
tx.commit().await?;
return Ok(QueryResult::None);
}
let mut failed = 0usize;
for (idx, params) in params_list.iter().enumerate() {
sqlx::query(AssertSqlSafe("SAVEPOINT dbkit_row"))
.execute(&mut *tx)
.await?;
let q = bind_pg(sqlx::query(AssertSqlSafe(query)), params).persistent(false);
match q.execute(&mut *tx).await {
Ok(_) => {
sqlx::query(AssertSqlSafe("RELEASE SAVEPOINT dbkit_row"))
.execute(&mut *tx)
.await?;
}
Err(e) => {
warn!("BatchParams row {}/{} failed: {:?}", idx + 1, total, e);
failed += 1;
sqlx::query(AssertSqlSafe("ROLLBACK TO SAVEPOINT dbkit_row"))
.execute(&mut *tx)
.await?;
sqlx::query(AssertSqlSafe("RELEASE SAVEPOINT dbkit_row"))
.execute(&mut *tx)
.await?;
}
}
}
tx.commit().await?;
if failed > 0 {
warn!(
"BatchParams: {}/{} succeeded, {} failed",
total - failed,
total,
failed
);
}
Ok(QueryResult::None)
}
}
}
pub async fn copy_in(
&self,
table: &str,
columns: &[&str],
rows: &[Vec<DbValue>],
) -> Result<u64, DbkitError> {
use sqlx::postgres::PgPoolCopyExt;
if rows.is_empty() {
return Ok(0);
}
let stmt = format!("COPY {table} ({}) FROM STDIN", columns.join(", "));
let mut sink = self.pool.copy_in_raw(&stmt).await?;
send_copy_rows(&mut sink, rows, columns.len()).await?;
Ok(sink.finish().await?)
}
pub async fn copy_upsert(
&self,
table: &str,
columns: &[&str],
conflict_columns: &[&str],
update_columns: &[&str],
rows: &[Vec<DbValue>],
) -> Result<u64, DbkitError> {
if rows.is_empty() {
return Ok(0);
}
if conflict_columns.is_empty() {
return Err(DbkitError::InvalidArgument(
"copy_upsert: conflict_columns must not be empty".into(),
));
}
let cols = columns.join(", ");
let stage = "dbkit_copy_stage";
let on_conflict = if update_columns.is_empty() {
format!("ON CONFLICT ({}) DO NOTHING", conflict_columns.join(", "))
} else {
let set = update_columns
.iter()
.map(|c| format!("{c} = EXCLUDED.{c}"))
.collect::<Vec<_>>()
.join(", ");
format!("ON CONFLICT ({}) DO UPDATE SET {set}", conflict_columns.join(", "))
};
let mut tx = self.pool.begin().await?;
sqlx::query(AssertSqlSafe(format!(
"CREATE TEMP TABLE {stage} ON COMMIT DROP AS SELECT {cols} FROM {table} WITH NO DATA"
)))
.execute(&mut *tx)
.await?;
let mut copy = (*tx)
.copy_in_raw(&format!("COPY {stage} ({cols}) FROM STDIN"))
.await?;
send_copy_rows(&mut copy, rows, columns.len()).await?;
copy.finish().await?;
let result = sqlx::query(AssertSqlSafe(format!(
"INSERT INTO {table} ({cols}) SELECT {cols} FROM {stage} {on_conflict}"
)))
.execute(&mut *tx)
.await?;
tx.commit().await?;
Ok(result.rows_affected())
}
pub async fn query_scalar_opt<T>(
&self,
query: &str,
params: Vec<DbValue>,
) -> Result<Option<T>, DbkitError>
where
T: for<'r> sqlx::Decode<'r, Postgres> + sqlx::Type<Postgres>,
{
let result = self
.execute_write(WriteOp::Single {
query,
params,
mode: FetchMode::Optional,
})
.await?;
Ok(result.optional()?.map(|row| row.get(0)))
}
pub async fn query_scalar<T>(
&self,
query: &str,
params: Vec<DbValue>,
) -> Result<T, DbkitError>
where
T: for<'r> sqlx::Decode<'r, Postgres> + sqlx::Type<Postgres>,
{
let result = self
.execute_write(WriteOp::Single {
query,
params,
mode: FetchMode::One,
})
.await?;
Ok(result.one()?.get(0))
}
pub async fn query_col<T>(
&self,
query: &str,
params: Vec<DbValue>,
) -> Result<Vec<T>, DbkitError>
where
T: for<'r> sqlx::Decode<'r, Postgres> + sqlx::Type<Postgres>,
{
let result = self
.execute_write(WriteOp::Single {
query,
params,
mode: FetchMode::All,
})
.await?;
Ok(result.all()?.iter().map(|row| row.get(0)).collect())
}
pub async fn query(
&self,
query: &str,
params: Vec<DbValue>,
mode: FetchMode,
) -> Result<QueryResult<PgRow>, DbkitError> {
let has_null = params.iter().any(|v| matches!(v, DbValue::Null));
let q = bind_pg(sqlx::query(AssertSqlSafe(query)), ¶ms).persistent(!has_null);
match mode {
FetchMode::None => {
q.execute(&self.pool).await?;
Ok(QueryResult::None)
}
FetchMode::One => Ok(QueryResult::One(q.fetch_one(&self.pool).await?)),
FetchMode::Optional => Ok(QueryResult::Optional(q.fetch_optional(&self.pool).await?)),
FetchMode::All => Ok(QueryResult::All(q.fetch_all(&self.pool).await?)),
}
}
#[cfg(feature = "duckdb")]
pub async fn execute_read(
&self,
sql: &str,
params: &[DbValue],
) -> Result<Vec<RecordBatch>, DbkitError> {
self.duck
.as_ref()
.ok_or(DbkitError::NoReadEngine)?
.query_arrow(sql, params)
.await
}
#[cfg(feature = "duckdb")]
pub async fn execute_read_as<T>(
&self,
sql: &str,
params: &[DbValue],
) -> Result<Vec<T>, DbkitError>
where
T: serde::de::DeserializeOwned,
{
let batches = self.execute_read(sql, params).await?;
crate::analytical::deserialize_batches(&batches)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn esc(s: &str) -> String {
let mut out = String::new();
copy_escape_into(s, &mut out);
out
}
#[test]
fn copy_escape_clean_passthrough_and_escapes() {
assert_eq!(esc(""), "");
assert_eq!(esc("plain text é ✓ 日本"), "plain text é ✓ 日本");
assert_eq!(esc("a\tb"), "a\\tb");
assert_eq!(esc("\\"), "\\\\");
assert_eq!(esc("\n\r\t"), "\\n\\r\\t");
assert_eq!(esc("trailing\\"), "trailing\\\\");
assert_eq!(esc("\\leading"), "\\\\leading");
assert_eq!(esc("mixé\tmulti✓\nbyte"), "mixé\\tmulti✓\\nbyte");
}
}