Skip to main content

dbkit/
pg_handler.rs

1//! Native-Postgres handler — sqlx [`PgPool`] with full Postgres type support.
2//!
3//! Mirrors [`BaseHandler`](crate::BaseHandler)'s write surface, but binds the
4//! *rich* [`DbValue`] variants (date / timestamp / json / uuid) to their native
5//! Postgres types via sqlx, and returns native [`PgRow`](sqlx::postgres::PgRow)s.
6//! Use this when you need Postgres types the multi-backend `Any` pool can't
7//! represent.
8//!
9//! Reads use the ergonomic row-mapped [`ReadOp`] API over DuckDB (typically
10//! attached live to Postgres via [`with_duckdb_attached_postgres`]).
11//!
12//! [`with_duckdb_attached_postgres`]: PgHandler::with_duckdb_attached_postgres
13
14use crate::DbkitError;
15use crate::base_handler::{FetchMode, QueryResult, WriteOp};
16use crate::value::DbValue;
17use std::fmt::Write as _;
18use sqlx::postgres::{PgArguments, PgRow};
19use sqlx::query::Query;
20use sqlx::{AssertSqlSafe, PgPool, Postgres, Row as _};
21use tracing::warn;
22use unicode_normalization::UnicodeNormalization;
23
24#[cfg(feature = "duckdb")]
25use crate::analytical::RecordBatch;
26#[cfg(feature = "duckdb")]
27use crate::read::{ReadEngine, duckdb::DuckEngine};
28
29/// A typeless SQL `NULL`. Declares the Postgres parameter type as OID 0 so the
30/// server infers it from context — exactly like a bare `NULL` literal. This lets
31/// a `NULL` [`DbValue`] unify with any column type in `COALESCE` / `CASE` / etc.,
32/// instead of being pinned to one concrete type. (Binding `Option::<i64>::None`
33/// forced `int8`, which broke e.g. `COALESCE($1, external_id)` against a
34/// `varchar` column: "bigint and character varying cannot be matched".)
35struct PgNull;
36
37impl sqlx::Type<Postgres> for PgNull {
38    fn type_info() -> sqlx::postgres::PgTypeInfo {
39        // OID 0 → "unspecified", resolved from context by the server.
40        sqlx::postgres::PgTypeInfo::with_oid(sqlx::postgres::types::Oid(0))
41    }
42}
43
44impl<'q> sqlx::Encode<'q, Postgres> for PgNull {
45    fn encode_by_ref(
46        &self,
47        _buf: &mut sqlx::postgres::PgArgumentBuffer,
48    ) -> Result<sqlx::encode::IsNull, sqlx::error::BoxDynError> {
49        Ok(sqlx::encode::IsNull::Yes)
50    }
51}
52
53/// Bind a slice of [`DbValue`]s onto a sqlx Postgres query, in order, binding
54/// the rich variants to their native Postgres types (no text fallback).
55/// Text/bytes/json/array values are bound by reference (sqlx encodes them into
56/// the argument buffer immediately), so no per-value clone is paid; the
57/// returned query borrows `params` for `'q`.
58fn bind_pg<'q>(
59    mut q: Query<'q, Postgres, PgArguments>,
60    params: &'q [DbValue],
61) -> Query<'q, Postgres, PgArguments> {
62    for p in params {
63        q = match p {
64            DbValue::Null => q.bind(PgNull),
65            DbValue::Bool(b) => q.bind(*b),
66            DbValue::Int(i) => q.bind(*i),
67            DbValue::Float(f) => q.bind(*f),
68            DbValue::Text(s) => q.bind(s.as_str()),
69            DbValue::Bytes(b) => q.bind(b.as_slice()),
70            DbValue::Date(d) => q.bind(*d),
71            DbValue::DateTime(dt) => q.bind(*dt),
72            DbValue::TimestampTz(dt) => q.bind(*dt),
73            DbValue::Json(j) => q.bind(j),
74            DbValue::Uuid(u) => q.bind(*u),
75            DbValue::Time(t) => q.bind(*t),
76            // sqlx binds `Vec<T>` / `Vec<Option<T>>` as native Postgres arrays.
77            DbValue::TextArray(v) => q.bind(v),
78            DbValue::FloatArray(v) => q.bind(v),
79            DbValue::OptFloatArray(v) => q.bind(v),
80        };
81    }
82    q
83}
84
85/// Render one [`DbValue`] as a cell in Postgres `COPY` text format, appending to
86/// `out`. NULL is the `\N` sentinel; all other values are escaped.
87fn copy_render_cell(val: &DbValue, out: &mut String) {
88    match val {
89        DbValue::Null => out.push_str("\\N"),
90        DbValue::Bool(b) => out.push(if *b { 't' } else { 'f' }),
91        // Numbers contain only digits / sign / `.` / `e` — never a COPY escape
92        // char — so format straight into `out`, skipping a throwaway `String`.
93        DbValue::Int(i) => {
94            let _ = write!(out, "{i}");
95        }
96        DbValue::Float(f) => {
97            if f.is_nan() {
98                out.push_str("NaN");
99            } else if f.is_infinite() {
100                out.push_str(if *f > 0.0 { "Infinity" } else { "-Infinity" });
101            } else {
102                let _ = write!(out, "{f}");
103            }
104        }
105        DbValue::Text(s) => copy_escape_into(s, out),
106        DbValue::Bytes(b) => {
107            // bytea hex format `\x<hex>`. The backslash is COPY-escaped to `\\`,
108            // and hex digits never need escaping, so write the escaped form
109            // directly — no temporary `String` or per-byte allocation.
110            out.push_str("\\\\x");
111            for byte in b {
112                out.push(char::from_digit((byte >> 4) as u32, 16).unwrap());
113                out.push(char::from_digit((byte & 0x0f) as u32, 16).unwrap());
114            }
115        }
116        DbValue::Date(d) => copy_escape_into(&d.to_string(), out),
117        DbValue::DateTime(dt) => copy_escape_into(&dt.to_string(), out),
118        DbValue::TimestampTz(dt) => copy_escape_into(&dt.to_rfc3339(), out),
119        DbValue::Json(j) => copy_escape_into(&j.to_string(), out),
120        DbValue::Uuid(u) => copy_escape_into(&u.to_string(), out),
121        DbValue::Time(t) => copy_escape_into(&t.to_string(), out),
122        DbValue::TextArray(v) => copy_escape_into(&crate::value::pg_text_array_literal(v), out),
123        DbValue::FloatArray(v) => {
124            copy_escape_into(&crate::value::pg_float_array_literal(v.iter().map(|x| Some(*x))), out)
125        }
126        DbValue::OptFloatArray(v) => {
127            copy_escape_into(&crate::value::pg_float_array_literal(v.iter().copied()), out)
128        }
129    }
130}
131
132/// Flush threshold for streaming COPY payloads. Rendering is flushed to the
133/// sink whenever the buffer passes this size, so peak memory stays ~this bound
134/// regardless of row count, and rendering overlaps with network sends.
135const COPY_CHUNK_BYTES: usize = 4 * 1024 * 1024;
136
137/// Render `rows` as Postgres `COPY` text format (cells tab-separated, one row
138/// per line) and stream them into an open COPY sink in ~[`COPY_CHUNK_BYTES`]
139/// chunks. `ncols` is used only to pre-size the buffer.
140async fn send_copy_rows<C>(
141    sink: &mut sqlx::postgres::PgCopyIn<C>,
142    rows: &[Vec<DbValue>],
143    ncols: usize,
144) -> Result<(), DbkitError>
145where
146    C: std::ops::DerefMut<Target = sqlx::postgres::PgConnection>,
147{
148    // Pre-size to the estimated payload (~12 bytes/cell + tab/newline), capped
149    // at one chunk — beyond that the buffer is flushed and reused anyway.
150    let estimate = rows.len() * (ncols * 12 + 1);
151    let mut buf = String::with_capacity(estimate.min(COPY_CHUNK_BYTES + 1024));
152    for row in rows {
153        for (i, val) in row.iter().enumerate() {
154            if i > 0 {
155                buf.push('\t');
156            }
157            copy_render_cell(val, &mut buf);
158        }
159        buf.push('\n');
160        if buf.len() >= COPY_CHUNK_BYTES {
161            sink.send(buf.as_bytes()).await?;
162            buf.clear();
163        }
164    }
165    if !buf.is_empty() {
166        sink.send(buf.as_bytes()).await?;
167    }
168    Ok(())
169}
170
171/// Escape a value for Postgres `COPY` text format (backslash, tab, newline, CR).
172///
173/// Scans for the next byte needing an escape and copies the clean span before
174/// it in one `push_str`; a string with no escapable bytes (the common case) is
175/// appended in a single copy. The four escape chars are ASCII, so byte
176/// positions are always UTF-8 boundaries.
177fn copy_escape_into(s: &str, out: &mut String) {
178    let mut start = 0;
179    for (i, b) in s.bytes().enumerate() {
180        let esc = match b {
181            b'\\' => "\\\\",
182            b'\t' => "\\t",
183            b'\n' => "\\n",
184            b'\r' => "\\r",
185            _ => continue,
186        };
187        out.push_str(&s[start..i]);
188        out.push_str(esc);
189        start = i + 1;
190    }
191    out.push_str(&s[start..]);
192}
193
194/// Core query executor for native Postgres: rich-typed transactional writes via
195/// sqlx, and row-mapped analytical reads via DuckDB.
196pub struct PgHandler {
197    pool: PgPool,
198    #[cfg(feature = "duckdb")]
199    duck: Option<DuckEngine>,
200}
201
202impl PgHandler {
203    /// Create a handler for writes against the given native Postgres pool.
204    pub fn new(pool: PgPool) -> Self {
205        Self {
206            pool,
207            #[cfg(feature = "duckdb")]
208            duck: None,
209        }
210    }
211
212    /// Create a handler with an in-memory DuckDB analytical read engine.
213    #[cfg(feature = "duckdb")]
214    pub fn with_duckdb(pool: PgPool) -> Result<Self, DbkitError> {
215        Ok(Self {
216            pool,
217            duck: Some(DuckEngine::new_in_memory()?),
218        })
219    }
220
221    /// Create a handler with DuckDB and a live Postgres attachment, so DuckDB
222    /// queries the Postgres tables directly via the `pg` catalog
223    /// (`SELECT … FROM pg.<schema>.<table>`) without an explicit sync.
224    #[cfg(feature = "duckdb")]
225    pub fn with_duckdb_attached_postgres(
226        pool: PgPool,
227        pg_connection_string: &str,
228    ) -> Result<Self, DbkitError> {
229        let duck = DuckEngine::new_in_memory()?;
230        duck.attach_postgres(pg_connection_string)?;
231        Ok(Self {
232            pool,
233            duck: Some(duck),
234        })
235    }
236
237    /// Whether a DuckDB read engine is attached.
238    pub fn has_read_engine(&self) -> bool {
239        #[cfg(feature = "duckdb")]
240        {
241            self.duck.is_some()
242        }
243        #[cfg(not(feature = "duckdb"))]
244        {
245            false
246        }
247    }
248
249    /// Get a reference to the native Postgres write pool.
250    pub fn pool(&self) -> &PgPool {
251        &self.pool
252    }
253
254    /// Accent-insensitive name key: NFD-decompose, DROP combining marks, then
255    /// lowercase — "José Ramírez" and "Jose Ramirez" produce the same key.
256    /// See `BaseHandler::normalize_name` for the rationale (NFD alone leaves
257    /// combining marks, so accented names never matched stripped variants).
258    pub fn normalize_name(name: &str) -> String {
259        use unicode_normalization::char::is_combining_mark;
260        name.nfd()
261            .filter(|c| !is_combining_mark(*c))
262            .collect::<String>()
263            .to_lowercase()
264    }
265
266    // ==================== UNIFIED WRITE ====================
267
268    /// Execute a write operation against the Postgres pool. Placeholders are
269    /// Postgres-native (`$1, $2, …`).
270    pub async fn execute_write(
271        &self,
272        op: WriteOp<'_>,
273    ) -> Result<QueryResult<PgRow>, DbkitError> {
274        match op {
275            WriteOp::Single {
276                query,
277                params,
278                mode,
279            } => self.query(query, params, mode).await,
280
281            WriteOp::BatchDDL { queries } => {
282                let mut tx = self.pool.begin().await?;
283                for query in queries {
284                    sqlx::query(AssertSqlSafe(*query)).execute(&mut *tx).await?;
285                }
286                tx.commit().await?;
287                Ok(QueryResult::None)
288            }
289
290            WriteOp::BatchParams {
291                query,
292                params_list,
293                isolate_rows,
294            } => {
295                if params_list.is_empty() {
296                    return Ok(QueryResult::None);
297                }
298                let total = params_list.len();
299                let mut tx = self.pool.begin().await?;
300
301                if !isolate_rows {
302                    // Fast path: no per-row SAVEPOINT. The whole batch is one
303                    // transaction, so any error rolls back *everything*
304                    // (all-or-nothing) — the cost of dropping savepoints, but
305                    // ~2× faster than the isolated path below.
306                    //
307                    // Statement reuse: a typeless NULL (`PgNull`, OID 0) lets the
308                    // server pin the cached statement's parameter type from the
309                    // first *cached* row, so a later row binding a concrete type
310                    // for that same column fails with 22P03. Guard per ROW (not
311                    // per batch): rows with a NULL re-parse individually
312                    // (`persistent(false)`, letting the server infer that row's
313                    // NULL from context), while null-free rows — whose concrete
314                    // types are mutually consistent — keep reusing one cached
315                    // prepared statement. A batch that is 10% NULL rows keeps
316                    // statement reuse for the other 90%.
317                    for params in &params_list {
318                        let has_null = params.iter().any(|v| matches!(v, DbValue::Null));
319                        bind_pg(sqlx::query(AssertSqlSafe(query)), params)
320                            .persistent(!has_null)
321                            .execute(&mut *tx)
322                            .await?;
323                    }
324                    tx.commit().await?;
325                    return Ok(QueryResult::None);
326                }
327
328                let mut failed = 0usize;
329                for (idx, params) in params_list.iter().enumerate() {
330                    // Wrap each row in a SAVEPOINT so a bad row rolls back on its
331                    // own instead of aborting the whole transaction. Without this,
332                    // Postgres marks the transaction failed on the first error and
333                    // every following row dies with 25P02 ("current transaction is
334                    // aborted"), turning one bad row into a whole failed batch.
335                    sqlx::query(AssertSqlSafe("SAVEPOINT dbkit_row"))
336                        .execute(&mut *tx)
337                        .await?;
338                    // `.persistent(false)` re-parses per row instead of reusing one
339                    // cached prepared statement across the batch. Reuse pins each
340                    // parameter's type from the FIRST row: a row whose value is a
341                    // typeless NULL lets the server resolve that param to the column
342                    // type (e.g. int4), and a later row binding the same column's
343                    // value as int8 then fails with 22P03 ("incorrect binary data
344                    // format"). Per-row parse keeps each row's param types self-consistent.
345                    let q = bind_pg(sqlx::query(AssertSqlSafe(query)), params).persistent(false);
346                    match q.execute(&mut *tx).await {
347                        Ok(_) => {
348                            sqlx::query(AssertSqlSafe("RELEASE SAVEPOINT dbkit_row"))
349                                .execute(&mut *tx)
350                                .await?;
351                        }
352                        Err(e) => {
353                            warn!("BatchParams row {}/{} failed: {:?}", idx + 1, total, e);
354                            failed += 1;
355                            sqlx::query(AssertSqlSafe("ROLLBACK TO SAVEPOINT dbkit_row"))
356                                .execute(&mut *tx)
357                                .await?;
358                            sqlx::query(AssertSqlSafe("RELEASE SAVEPOINT dbkit_row"))
359                                .execute(&mut *tx)
360                                .await?;
361                        }
362                    }
363                }
364                tx.commit().await?;
365                if failed > 0 {
366                    warn!(
367                        "BatchParams: {}/{} succeeded, {} failed",
368                        total - failed,
369                        total,
370                        failed
371                    );
372                }
373                Ok(QueryResult::None)
374            }
375        }
376    }
377
378    /// Bulk-insert rows via Postgres `COPY ... FROM STDIN` (text format).
379    ///
380    /// **The fastest way to load many rows** — one streamed `COPY` instead of a
381    /// parse + execute (+ savepoint) per row like [`WriteOp::BatchParams`].
382    /// Benchmarks at roughly 30–50× the throughput of `BatchParams`. Each row in
383    /// `rows` must align positionally with `columns`. Returns the number of rows
384    /// copied.
385    ///
386    /// # `copy_in` vs [`WriteOp::BatchParams`] — which to use
387    ///
388    /// | Reach for `copy_in` when… | Reach for `BatchParams` when… |
389    /// |---|---|
390    /// | Plain bulk insert into one table | You need `INSERT … ON CONFLICT` (upsert) |
391    /// | Data is trusted; all-or-nothing is fine | You need per-row isolation (skip bad rows) |
392    /// | You want maximum throughput | The statement isn't a plain insert (`UPDATE`, `RETURNING`, computed `VALUES`) |
393    /// | Target is Postgres | Target is a non-Postgres backend (use the `Any` pool) |
394    ///
395    /// `COPY` is **not** an `INSERT` statement, so it does **not** support
396    /// `ON CONFLICT`, `RETURNING`, `DEFAULT` expressions, or `WHERE`, and it is
397    /// **all-or-nothing**: a constraint violation aborts the entire load (it does
398    /// not skip bad rows like `BatchParams`).
399    ///
400    /// To bulk-**upsert**, combine the two: `COPY` into a constraint-free staging
401    /// table, then run one set-based `INSERT … SELECT … ON CONFLICT` — far faster
402    /// than per-row `BatchParams` with `ON CONFLICT`:
403    ///
404    /// ```sql
405    /// CREATE TEMP TABLE stage (LIKE target INCLUDING DEFAULTS) ON COMMIT DROP;
406    /// COPY stage (id, name) FROM STDIN;            -- fast bulk load, no constraints
407    /// INSERT INTO target (id, name)
408    ///   SELECT id, name FROM stage                 -- one set-based upsert
409    ///   ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name;
410    /// ```
411    pub async fn copy_in(
412        &self,
413        table: &str,
414        columns: &[&str],
415        rows: &[Vec<DbValue>],
416    ) -> Result<u64, DbkitError> {
417        use sqlx::postgres::PgPoolCopyExt;
418
419        if rows.is_empty() {
420            return Ok(0);
421        }
422
423        let stmt = format!("COPY {table} ({}) FROM STDIN", columns.join(", "));
424
425        let mut sink = self.pool.copy_in_raw(&stmt).await?;
426        send_copy_rows(&mut sink, rows, columns.len()).await?;
427        Ok(sink.finish().await?)
428    }
429
430    /// Bulk-**upsert** rows: `COPY` into a staging table, then one set-based
431    /// `INSERT … SELECT … ON CONFLICT`, all in a single transaction.
432    ///
433    /// This is the fast path for `ON CONFLICT` at scale. Plain [`copy_in`] can't
434    /// do `ON CONFLICT` (it's not an `INSERT`), and per-row
435    /// [`WriteOp::BatchParams`] with `ON CONFLICT` pays per-row overhead. This
436    /// combines both strengths: COPY's bulk ingestion into a constraint-free
437    /// staging table, then a single set-based upsert into the target.
438    ///
439    /// - `columns` — columns present in `rows` (positional), copied into staging.
440    /// - `conflict_columns` — the conflict target (must back a unique/PK index).
441    /// - `update_columns` — columns to overwrite on conflict (set to the incoming
442    ///   `EXCLUDED` value). **Empty** ⇒ `DO NOTHING` (insert-or-ignore).
443    ///
444    /// Returns the number of rows inserted or updated.
445    ///
446    /// The staging table is `CREATE TEMP TABLE … AS SELECT {columns} FROM target
447    /// WITH NO DATA` (`ON COMMIT DROP`) — target column types, no constraints or
448    /// defaults — and vanishes at commit. The
449    /// final upsert is all-or-nothing: a non-conflict error (CHECK/FK/type) aborts
450    /// the batch. **Within a single call, `conflict_columns` must be unique across
451    /// `rows`** — duplicate keys make `ON CONFLICT DO UPDATE` error with "command
452    /// cannot affect row a second time"; de-duplicate before calling.
453    ///
454    /// [`copy_in`]: Self::copy_in
455    pub async fn copy_upsert(
456        &self,
457        table: &str,
458        columns: &[&str],
459        conflict_columns: &[&str],
460        update_columns: &[&str],
461        rows: &[Vec<DbValue>],
462    ) -> Result<u64, DbkitError> {
463        if rows.is_empty() {
464            return Ok(0);
465        }
466        if conflict_columns.is_empty() {
467            // Would render `ON CONFLICT () …` — a Postgres syntax error, but an
468            // opaque one; fail with a message that names the actual mistake.
469            return Err(DbkitError::InvalidArgument(
470                "copy_upsert: conflict_columns must not be empty".into(),
471            ));
472        }
473
474        let cols = columns.join(", ");
475        let stage = "dbkit_copy_stage";
476
477        let on_conflict = if update_columns.is_empty() {
478            format!("ON CONFLICT ({}) DO NOTHING", conflict_columns.join(", "))
479        } else {
480            let set = update_columns
481                .iter()
482                .map(|c| format!("{c} = EXCLUDED.{c}"))
483                .collect::<Vec<_>>()
484                .join(", ");
485            format!("ON CONFLICT ({}) DO UPDATE SET {set}", conflict_columns.join(", "))
486        };
487
488        let mut tx = self.pool.begin().await?;
489
490        // 1. Staging table with ONLY the copied columns (target types, no
491        //    constraints, no defaults), dropped at COMMIT. Temp tables are
492        //    connection-scoped, so the fixed name is safe even under concurrent
493        //    callers on separate connections.
494        //
495        //    `AS SELECT … WITH NO DATA` rather than `LIKE {table}`: LIKE always
496        //    copies NOT NULL constraints (so unlisted NOT-NULL columns would
497        //    reject COPY's NULL fill), and `INCLUDING DEFAULTS` — used before
498        //    0.5 to paper over that — made COPY fire a copied serial default,
499        //    burning one target-sequence value per staged row for nothing.
500        //    Narrowing staging to the listed columns sidesteps both; the
501        //    target's own defaults still apply at the INSERT below.
502        sqlx::query(AssertSqlSafe(format!(
503            "CREATE TEMP TABLE {stage} ON COMMIT DROP AS SELECT {cols} FROM {table} WITH NO DATA"
504        )))
505        .execute(&mut *tx)
506        .await?;
507
508        // 2. Bulk-load into staging via COPY on the SAME connection (so the temp
509        //    table is visible) — this is where the throughput comes from.
510        let mut copy = (*tx)
511            .copy_in_raw(&format!("COPY {stage} ({cols}) FROM STDIN"))
512            .await?;
513        send_copy_rows(&mut copy, rows, columns.len()).await?;
514        copy.finish().await?;
515
516        // 3. One set-based upsert from staging into the target.
517        let result = sqlx::query(AssertSqlSafe(format!(
518            "INSERT INTO {table} ({cols}) SELECT {cols} FROM {stage} {on_conflict}"
519        )))
520        .execute(&mut *tx)
521        .await?;
522
523        tx.commit().await?;
524        Ok(result.rows_affected())
525    }
526
527    // ==================== SCALAR READ HELPERS ====================
528
529    /// Run a single-value `SELECT` and return the first column of the single
530    /// matching row, or `None` if no row matched.
531    ///
532    /// Convenience over the `execute_write(WriteOp::Single { .., FetchMode::Optional })`
533    /// → `.optional()?` → `row.get(0)` dance that OLTP `get_*_id` / `get_or_create`
534    /// lookups hand-roll everywhere. Uses `Optional` (not `One`), so an empty
535    /// result is `Ok(None)` rather than an error; a genuine DB failure still
536    /// surfaces as `Err`.
537    pub async fn query_scalar_opt<T>(
538        &self,
539        query: &str,
540        params: Vec<DbValue>,
541    ) -> Result<Option<T>, DbkitError>
542    where
543        T: for<'r> sqlx::Decode<'r, Postgres> + sqlx::Type<Postgres>,
544    {
545        let result = self
546            .execute_write(WriteOp::Single {
547                query,
548                params,
549                mode: FetchMode::Optional,
550            })
551            .await?;
552        Ok(result.optional()?.map(|row| row.get(0)))
553    }
554
555    /// Run a single-value `SELECT` (or `INSERT … RETURNING`) that must yield
556    /// exactly one row, and return the first column. Errors if no row matched —
557    /// use [`query_scalar_opt`](Self::query_scalar_opt) when zero rows is valid.
558    pub async fn query_scalar<T>(
559        &self,
560        query: &str,
561        params: Vec<DbValue>,
562    ) -> Result<T, DbkitError>
563    where
564        T: for<'r> sqlx::Decode<'r, Postgres> + sqlx::Type<Postgres>,
565    {
566        let result = self
567            .execute_write(WriteOp::Single {
568                query,
569                params,
570                mode: FetchMode::One,
571            })
572            .await?;
573        Ok(result.one()?.get(0))
574    }
575
576    /// Run a `SELECT` and collect the first column of every row into a `Vec`.
577    pub async fn query_col<T>(
578        &self,
579        query: &str,
580        params: Vec<DbValue>,
581    ) -> Result<Vec<T>, DbkitError>
582    where
583        T: for<'r> sqlx::Decode<'r, Postgres> + sqlx::Type<Postgres>,
584    {
585        let result = self
586            .execute_write(WriteOp::Single {
587                query,
588                params,
589                mode: FetchMode::All,
590            })
591            .await?;
592        Ok(result.all()?.iter().map(|row| row.get(0)).collect())
593    }
594
595    // ==================== NATIVE POSTGRES READ ====================
596
597    /// Run a query against the native Postgres pool, returning rows per `mode`.
598    ///
599    /// This is the OLTP read path — single-row lookups and small result sets go
600    /// straight to Postgres (one round-trip → [`PgRow`]), no analytical engine.
601    /// Placeholders are Postgres-native (`$1, $2, …`); read columns off the
602    /// returned [`PgRow`]s with `row.get(i)` / `row.try_get(i)`.
603    pub async fn query(
604        &self,
605        query: &str,
606        params: Vec<DbValue>,
607        mode: FetchMode,
608    ) -> Result<QueryResult<PgRow>, DbkitError> {
609        // Statement reuse hazard: with the default `persistent(true)`, sqlx
610        // caches one prepared statement per (connection, SQL). A typeless NULL
611        // (`PgNull`, OID 0) lets the server pin that parameter's type from the
612        // FIRST execution, so a later call binding a concrete type for the same
613        // column fails with 22P03 ("incorrect binary data format"). Reuse the
614        // cached statement only when this call has no NULLs; otherwise re-parse
615        // so each call's param types stay self-consistent. (Same guard as the
616        // `BatchParams` write path.)
617        let has_null = params.iter().any(|v| matches!(v, DbValue::Null));
618        let q = bind_pg(sqlx::query(AssertSqlSafe(query)), &params).persistent(!has_null);
619        match mode {
620            FetchMode::None => {
621                q.execute(&self.pool).await?;
622                Ok(QueryResult::None)
623            }
624            FetchMode::One => Ok(QueryResult::One(q.fetch_one(&self.pool).await?)),
625            FetchMode::Optional => Ok(QueryResult::Optional(q.fetch_optional(&self.pool).await?)),
626            FetchMode::All => Ok(QueryResult::All(q.fetch_all(&self.pool).await?)),
627        }
628    }
629
630    // ==================== ANALYTICAL READ (DuckDB / Arrow) ====================
631
632    /// Run an analytical query against the attached DuckDB engine, returning
633    /// columnar Arrow [`RecordBatch`]es. For large joins/aggregations consumed as
634    /// DataFrames. Errors with [`DbkitError::NoReadEngine`] if no engine is
635    /// attached. For typed rows, deserialize the batches (see
636    /// [`BaseHandler::execute_read_as`](crate::BaseHandler::execute_read_as)).
637    #[cfg(feature = "duckdb")]
638    pub async fn execute_read(
639        &self,
640        sql: &str,
641        params: &[DbValue],
642    ) -> Result<Vec<RecordBatch>, DbkitError> {
643        self.duck
644            .as_ref()
645            .ok_or(DbkitError::NoReadEngine)?
646            .query_arrow(sql, params)
647            .await
648    }
649
650    /// Like [`execute_read`](Self::execute_read) but deserializes each row into
651    /// `T` via `serde_arrow` — the typed analytical read. `T`'s field names must
652    /// match the query's output column names. Use for DuckDB-side analytical
653    /// reads (large scans / aggregations) that map to typed rows. Errors with
654    /// [`DbkitError::NoReadEngine`] if no engine is attached.
655    #[cfg(feature = "duckdb")]
656    pub async fn execute_read_as<T>(
657        &self,
658        sql: &str,
659        params: &[DbValue],
660    ) -> Result<Vec<T>, DbkitError>
661    where
662        T: serde::de::DeserializeOwned,
663    {
664        let batches = self.execute_read(sql, params).await?;
665        crate::analytical::deserialize_batches(&batches)
666    }
667}
668
669#[cfg(test)]
670mod tests {
671    use super::*;
672
673    fn esc(s: &str) -> String {
674        let mut out = String::new();
675        copy_escape_into(s, &mut out);
676        out
677    }
678
679    #[test]
680    fn copy_escape_clean_passthrough_and_escapes() {
681        assert_eq!(esc(""), "");
682        assert_eq!(esc("plain text é ✓ 日本"), "plain text é ✓ 日本");
683        assert_eq!(esc("a\tb"), "a\\tb");
684        assert_eq!(esc("\\"), "\\\\");
685        assert_eq!(esc("\n\r\t"), "\\n\\r\\t");
686        assert_eq!(esc("trailing\\"), "trailing\\\\");
687        assert_eq!(esc("\\leading"), "\\\\leading");
688        assert_eq!(esc("mixé\tmulti✓\nbyte"), "mixé\\tmulti✓\\nbyte");
689    }
690}