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};
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 /// Unicode NFD normalization — decomposes characters then lowercases.
255 pub fn normalize_name(name: &str) -> String {
256 name.nfd().collect::<String>().to_lowercase()
257 }
258
259 // ==================== UNIFIED WRITE ====================
260
261 /// Execute a write operation against the Postgres pool. Placeholders are
262 /// Postgres-native (`$1, $2, …`).
263 pub async fn execute_write(
264 &self,
265 op: WriteOp<'_>,
266 ) -> Result<QueryResult<PgRow>, DbkitError> {
267 match op {
268 WriteOp::Single {
269 query,
270 params,
271 mode,
272 } => self.query(query, params, mode).await,
273
274 WriteOp::BatchDDL { queries } => {
275 let mut tx = self.pool.begin().await?;
276 for query in queries {
277 sqlx::query(AssertSqlSafe(*query)).execute(&mut *tx).await?;
278 }
279 tx.commit().await?;
280 Ok(QueryResult::None)
281 }
282
283 WriteOp::BatchParams {
284 query,
285 params_list,
286 isolate_rows,
287 } => {
288 if params_list.is_empty() {
289 return Ok(QueryResult::None);
290 }
291 let total = params_list.len();
292 let mut tx = self.pool.begin().await?;
293
294 if !isolate_rows {
295 // Fast path: no per-row SAVEPOINT. The whole batch is one
296 // transaction, so any error rolls back *everything*
297 // (all-or-nothing) — the cost of dropping savepoints, but
298 // ~2× faster than the isolated path below.
299 //
300 // Statement reuse: a typeless NULL (`PgNull`, OID 0) lets the
301 // server pin the cached statement's parameter type from the
302 // first *cached* row, so a later row binding a concrete type
303 // for that same column fails with 22P03. Guard per ROW (not
304 // per batch): rows with a NULL re-parse individually
305 // (`persistent(false)`, letting the server infer that row's
306 // NULL from context), while null-free rows — whose concrete
307 // types are mutually consistent — keep reusing one cached
308 // prepared statement. A batch that is 10% NULL rows keeps
309 // statement reuse for the other 90%.
310 for params in ¶ms_list {
311 let has_null = params.iter().any(|v| matches!(v, DbValue::Null));
312 bind_pg(sqlx::query(AssertSqlSafe(query)), params)
313 .persistent(!has_null)
314 .execute(&mut *tx)
315 .await?;
316 }
317 tx.commit().await?;
318 return Ok(QueryResult::None);
319 }
320
321 let mut failed = 0usize;
322 for (idx, params) in params_list.iter().enumerate() {
323 // Wrap each row in a SAVEPOINT so a bad row rolls back on its
324 // own instead of aborting the whole transaction. Without this,
325 // Postgres marks the transaction failed on the first error and
326 // every following row dies with 25P02 ("current transaction is
327 // aborted"), turning one bad row into a whole failed batch.
328 sqlx::query(AssertSqlSafe("SAVEPOINT dbkit_row"))
329 .execute(&mut *tx)
330 .await?;
331 // `.persistent(false)` re-parses per row instead of reusing one
332 // cached prepared statement across the batch. Reuse pins each
333 // parameter's type from the FIRST row: a row whose value is a
334 // typeless NULL lets the server resolve that param to the column
335 // type (e.g. int4), and a later row binding the same column's
336 // value as int8 then fails with 22P03 ("incorrect binary data
337 // format"). Per-row parse keeps each row's param types self-consistent.
338 let q = bind_pg(sqlx::query(AssertSqlSafe(query)), params).persistent(false);
339 match q.execute(&mut *tx).await {
340 Ok(_) => {
341 sqlx::query(AssertSqlSafe("RELEASE SAVEPOINT dbkit_row"))
342 .execute(&mut *tx)
343 .await?;
344 }
345 Err(e) => {
346 warn!("BatchParams row {}/{} failed: {:?}", idx + 1, total, e);
347 failed += 1;
348 sqlx::query(AssertSqlSafe("ROLLBACK TO SAVEPOINT dbkit_row"))
349 .execute(&mut *tx)
350 .await?;
351 sqlx::query(AssertSqlSafe("RELEASE SAVEPOINT dbkit_row"))
352 .execute(&mut *tx)
353 .await?;
354 }
355 }
356 }
357 tx.commit().await?;
358 if failed > 0 {
359 warn!(
360 "BatchParams: {}/{} succeeded, {} failed",
361 total - failed,
362 total,
363 failed
364 );
365 }
366 Ok(QueryResult::None)
367 }
368 }
369 }
370
371 /// Bulk-insert rows via Postgres `COPY ... FROM STDIN` (text format).
372 ///
373 /// **The fastest way to load many rows** — one streamed `COPY` instead of a
374 /// parse + execute (+ savepoint) per row like [`WriteOp::BatchParams`].
375 /// Benchmarks at roughly 30–50× the throughput of `BatchParams`. Each row in
376 /// `rows` must align positionally with `columns`. Returns the number of rows
377 /// copied.
378 ///
379 /// # `copy_in` vs [`WriteOp::BatchParams`] — which to use
380 ///
381 /// | Reach for `copy_in` when… | Reach for `BatchParams` when… |
382 /// |---|---|
383 /// | Plain bulk insert into one table | You need `INSERT … ON CONFLICT` (upsert) |
384 /// | Data is trusted; all-or-nothing is fine | You need per-row isolation (skip bad rows) |
385 /// | You want maximum throughput | The statement isn't a plain insert (`UPDATE`, `RETURNING`, computed `VALUES`) |
386 /// | Target is Postgres | Target is a non-Postgres backend (use the `Any` pool) |
387 ///
388 /// `COPY` is **not** an `INSERT` statement, so it does **not** support
389 /// `ON CONFLICT`, `RETURNING`, `DEFAULT` expressions, or `WHERE`, and it is
390 /// **all-or-nothing**: a constraint violation aborts the entire load (it does
391 /// not skip bad rows like `BatchParams`).
392 ///
393 /// To bulk-**upsert**, combine the two: `COPY` into a constraint-free staging
394 /// table, then run one set-based `INSERT … SELECT … ON CONFLICT` — far faster
395 /// than per-row `BatchParams` with `ON CONFLICT`:
396 ///
397 /// ```sql
398 /// CREATE TEMP TABLE stage (LIKE target INCLUDING DEFAULTS) ON COMMIT DROP;
399 /// COPY stage (id, name) FROM STDIN; -- fast bulk load, no constraints
400 /// INSERT INTO target (id, name)
401 /// SELECT id, name FROM stage -- one set-based upsert
402 /// ON CONFLICT (id) DO UPDATE SET name = EXCLUDED.name;
403 /// ```
404 pub async fn copy_in(
405 &self,
406 table: &str,
407 columns: &[&str],
408 rows: &[Vec<DbValue>],
409 ) -> Result<u64, DbkitError> {
410 use sqlx::postgres::PgPoolCopyExt;
411
412 if rows.is_empty() {
413 return Ok(0);
414 }
415
416 let stmt = format!("COPY {table} ({}) FROM STDIN", columns.join(", "));
417
418 let mut sink = self.pool.copy_in_raw(&stmt).await?;
419 send_copy_rows(&mut sink, rows, columns.len()).await?;
420 Ok(sink.finish().await?)
421 }
422
423 /// Bulk-**upsert** rows: `COPY` into a staging table, then one set-based
424 /// `INSERT … SELECT … ON CONFLICT`, all in a single transaction.
425 ///
426 /// This is the fast path for `ON CONFLICT` at scale. Plain [`copy_in`] can't
427 /// do `ON CONFLICT` (it's not an `INSERT`), and per-row
428 /// [`WriteOp::BatchParams`] with `ON CONFLICT` pays per-row overhead. This
429 /// combines both strengths: COPY's bulk ingestion into a constraint-free
430 /// staging table, then a single set-based upsert into the target.
431 ///
432 /// - `columns` — columns present in `rows` (positional), copied into staging.
433 /// - `conflict_columns` — the conflict target (must back a unique/PK index).
434 /// - `update_columns` — columns to overwrite on conflict (set to the incoming
435 /// `EXCLUDED` value). **Empty** ⇒ `DO NOTHING` (insert-or-ignore).
436 ///
437 /// Returns the number of rows inserted or updated.
438 ///
439 /// The staging table is `CREATE TEMP TABLE … AS SELECT {columns} FROM target
440 /// WITH NO DATA` (`ON COMMIT DROP`) — target column types, no constraints or
441 /// defaults — and vanishes at commit. The
442 /// final upsert is all-or-nothing: a non-conflict error (CHECK/FK/type) aborts
443 /// the batch. **Within a single call, `conflict_columns` must be unique across
444 /// `rows`** — duplicate keys make `ON CONFLICT DO UPDATE` error with "command
445 /// cannot affect row a second time"; de-duplicate before calling.
446 ///
447 /// [`copy_in`]: Self::copy_in
448 pub async fn copy_upsert(
449 &self,
450 table: &str,
451 columns: &[&str],
452 conflict_columns: &[&str],
453 update_columns: &[&str],
454 rows: &[Vec<DbValue>],
455 ) -> Result<u64, DbkitError> {
456 if rows.is_empty() {
457 return Ok(0);
458 }
459 if conflict_columns.is_empty() {
460 // Would render `ON CONFLICT () …` — a Postgres syntax error, but an
461 // opaque one; fail with a message that names the actual mistake.
462 return Err(DbkitError::InvalidArgument(
463 "copy_upsert: conflict_columns must not be empty".into(),
464 ));
465 }
466
467 let cols = columns.join(", ");
468 let stage = "dbkit_copy_stage";
469
470 let on_conflict = if update_columns.is_empty() {
471 format!("ON CONFLICT ({}) DO NOTHING", conflict_columns.join(", "))
472 } else {
473 let set = update_columns
474 .iter()
475 .map(|c| format!("{c} = EXCLUDED.{c}"))
476 .collect::<Vec<_>>()
477 .join(", ");
478 format!("ON CONFLICT ({}) DO UPDATE SET {set}", conflict_columns.join(", "))
479 };
480
481 let mut tx = self.pool.begin().await?;
482
483 // 1. Staging table with ONLY the copied columns (target types, no
484 // constraints, no defaults), dropped at COMMIT. Temp tables are
485 // connection-scoped, so the fixed name is safe even under concurrent
486 // callers on separate connections.
487 //
488 // `AS SELECT … WITH NO DATA` rather than `LIKE {table}`: LIKE always
489 // copies NOT NULL constraints (so unlisted NOT-NULL columns would
490 // reject COPY's NULL fill), and `INCLUDING DEFAULTS` — used before
491 // 0.5 to paper over that — made COPY fire a copied serial default,
492 // burning one target-sequence value per staged row for nothing.
493 // Narrowing staging to the listed columns sidesteps both; the
494 // target's own defaults still apply at the INSERT below.
495 sqlx::query(AssertSqlSafe(format!(
496 "CREATE TEMP TABLE {stage} ON COMMIT DROP AS SELECT {cols} FROM {table} WITH NO DATA"
497 )))
498 .execute(&mut *tx)
499 .await?;
500
501 // 2. Bulk-load into staging via COPY on the SAME connection (so the temp
502 // table is visible) — this is where the throughput comes from.
503 let mut copy = (*tx)
504 .copy_in_raw(&format!("COPY {stage} ({cols}) FROM STDIN"))
505 .await?;
506 send_copy_rows(&mut copy, rows, columns.len()).await?;
507 copy.finish().await?;
508
509 // 3. One set-based upsert from staging into the target.
510 let result = sqlx::query(AssertSqlSafe(format!(
511 "INSERT INTO {table} ({cols}) SELECT {cols} FROM {stage} {on_conflict}"
512 )))
513 .execute(&mut *tx)
514 .await?;
515
516 tx.commit().await?;
517 Ok(result.rows_affected())
518 }
519
520 // ==================== NATIVE POSTGRES READ ====================
521
522 /// Run a query against the native Postgres pool, returning rows per `mode`.
523 ///
524 /// This is the OLTP read path — single-row lookups and small result sets go
525 /// straight to Postgres (one round-trip → [`PgRow`]), no analytical engine.
526 /// Placeholders are Postgres-native (`$1, $2, …`); read columns off the
527 /// returned [`PgRow`]s with `row.get(i)` / `row.try_get(i)`.
528 pub async fn query(
529 &self,
530 query: &str,
531 params: Vec<DbValue>,
532 mode: FetchMode,
533 ) -> Result<QueryResult<PgRow>, DbkitError> {
534 // Statement reuse hazard: with the default `persistent(true)`, sqlx
535 // caches one prepared statement per (connection, SQL). A typeless NULL
536 // (`PgNull`, OID 0) lets the server pin that parameter's type from the
537 // FIRST execution, so a later call binding a concrete type for the same
538 // column fails with 22P03 ("incorrect binary data format"). Reuse the
539 // cached statement only when this call has no NULLs; otherwise re-parse
540 // so each call's param types stay self-consistent. (Same guard as the
541 // `BatchParams` write path.)
542 let has_null = params.iter().any(|v| matches!(v, DbValue::Null));
543 let q = bind_pg(sqlx::query(AssertSqlSafe(query)), ¶ms).persistent(!has_null);
544 match mode {
545 FetchMode::None => {
546 q.execute(&self.pool).await?;
547 Ok(QueryResult::None)
548 }
549 FetchMode::One => Ok(QueryResult::One(q.fetch_one(&self.pool).await?)),
550 FetchMode::Optional => Ok(QueryResult::Optional(q.fetch_optional(&self.pool).await?)),
551 FetchMode::All => Ok(QueryResult::All(q.fetch_all(&self.pool).await?)),
552 }
553 }
554
555 // ==================== ANALYTICAL READ (DuckDB / Arrow) ====================
556
557 /// Run an analytical query against the attached DuckDB engine, returning
558 /// columnar Arrow [`RecordBatch`]es. For large joins/aggregations consumed as
559 /// DataFrames. Errors with [`DbkitError::NoReadEngine`] if no engine is
560 /// attached. For typed rows, deserialize the batches (see
561 /// [`BaseHandler::execute_read_as`](crate::BaseHandler::execute_read_as)).
562 #[cfg(feature = "duckdb")]
563 pub async fn execute_read(
564 &self,
565 sql: &str,
566 params: &[DbValue],
567 ) -> Result<Vec<RecordBatch>, DbkitError> {
568 self.duck
569 .as_ref()
570 .ok_or(DbkitError::NoReadEngine)?
571 .query_arrow(sql, params)
572 .await
573 }
574
575 /// Like [`execute_read`](Self::execute_read) but deserializes each row into
576 /// `T` via `serde_arrow` — the typed analytical read. `T`'s field names must
577 /// match the query's output column names. Use for DuckDB-side analytical
578 /// reads (large scans / aggregations) that map to typed rows. Errors with
579 /// [`DbkitError::NoReadEngine`] if no engine is attached.
580 #[cfg(feature = "duckdb")]
581 pub async fn execute_read_as<T>(
582 &self,
583 sql: &str,
584 params: &[DbValue],
585 ) -> Result<Vec<T>, DbkitError>
586 where
587 T: serde::de::DeserializeOwned,
588 {
589 let batches = self.execute_read(sql, params).await?;
590 crate::analytical::deserialize_batches(&batches)
591 }
592}
593
594#[cfg(test)]
595mod tests {
596 use super::*;
597
598 fn esc(s: &str) -> String {
599 let mut out = String::new();
600 copy_escape_into(s, &mut out);
601 out
602 }
603
604 #[test]
605 fn copy_escape_clean_passthrough_and_escapes() {
606 assert_eq!(esc(""), "");
607 assert_eq!(esc("plain text é ✓ 日本"), "plain text é ✓ 日本");
608 assert_eq!(esc("a\tb"), "a\\tb");
609 assert_eq!(esc("\\"), "\\\\");
610 assert_eq!(esc("\n\r\t"), "\\n\\r\\t");
611 assert_eq!(esc("trailing\\"), "trailing\\\\");
612 assert_eq!(esc("\\leading"), "\\\\leading");
613 assert_eq!(esc("mixé\tmulti✓\nbyte"), "mixé\\tmulti✓\\nbyte");
614 }
615}