dbkit/
base_handler.rs

1use crate::DbkitError;
2use crate::value::DbValue;
3use sqlx::any::{AnyArguments, AnyRow};
4use sqlx::query::Query;
5use sqlx::{Any, AnyPool, AssertSqlSafe};
6use tracing::warn;
7use unicode_normalization::UnicodeNormalization;
8
9#[cfg(any(feature = "duckdb", feature = "datafusion"))]
10use crate::analytical::RecordBatch;
11#[cfg(any(feature = "duckdb", feature = "datafusion"))]
12use crate::read::ReadEngine;
13
14// ---------------------------------------------------------------------------
15// Write operations
16// ---------------------------------------------------------------------------
17
18/// Unified write operation types.
19pub enum WriteOp<'a> {
20    /// Single query with optional return.
21    Single {
22        query: &'a str,
23        params: Vec<DbValue>,
24        mode: FetchMode,
25    },
26    /// Batch of DDL statements executed in a single transaction.
27    BatchDDL { queries: &'a [&'a str] },
28    /// Same query executed for each parameter set in a transaction.
29    BatchParams {
30        query: &'a str,
31        params_list: Vec<Vec<DbValue>>,
32    },
33}
34
35// ---------------------------------------------------------------------------
36// Query result types
37// ---------------------------------------------------------------------------
38
39/// How many rows to expect from a query.
40#[derive(Debug, Clone, Copy)]
41pub enum FetchMode {
42    None,
43    One,
44    Optional,
45    All,
46}
47
48/// Result wrapper for write queries.
49pub enum QueryResult<T> {
50    None,
51    One(T),
52    Optional(Option<T>),
53    All(Vec<T>),
54}
55
56impl<T> QueryResult<T> {
57    pub fn one(self) -> Result<T, DbkitError> {
58        match self {
59            Self::One(v) => Ok(v),
60            _ => Err(DbkitError::RowCount {
61                expected: "One".into(),
62                actual: 0,
63            }),
64        }
65    }
66
67    pub fn optional(self) -> Result<Option<T>, DbkitError> {
68        match self {
69            Self::Optional(v) => Ok(v),
70            Self::One(v) => Ok(Some(v)),
71            Self::None => Ok(None),
72            _ => Err(DbkitError::RowCount {
73                expected: "Optional".into(),
74                actual: 0,
75            }),
76        }
77    }
78
79    pub fn all(self) -> Result<Vec<T>, DbkitError> {
80        match self {
81            Self::All(v) => Ok(v),
82            _ => Err(DbkitError::RowCount {
83                expected: "All".into(),
84                actual: 0,
85            }),
86        }
87    }
88}
89
90// ---------------------------------------------------------------------------
91// Parameter binding
92// ---------------------------------------------------------------------------
93
94/// Bind a slice of [`DbValue`]s onto a sqlx query, in order.
95///
96/// Values are bound by owned copy, so the returned query does not borrow
97/// `params`.
98fn bind_params<'q>(
99    mut q: Query<'q, Any, AnyArguments>,
100    params: &[DbValue],
101) -> Query<'q, Any, AnyArguments> {
102    for p in params {
103        q = match p {
104            DbValue::Null => q.bind(Option::<i64>::None),
105            DbValue::Bool(b) => q.bind(*b),
106            DbValue::Int(i) => q.bind(*i),
107            DbValue::Float(f) => q.bind(*f),
108            DbValue::Text(s) => q.bind(s.clone()),
109            DbValue::Bytes(b) => q.bind(b.clone()),
110            // The Any driver can't carry native temporal/json/uuid types, so
111            // bind a text rendering — Postgres assignment casts handle the rest.
112            // For native rich-typed binds use `PgHandler` instead.
113            #[cfg(feature = "postgres-native")]
114            DbValue::Date(d) => q.bind(d.to_string()),
115            #[cfg(feature = "postgres-native")]
116            DbValue::DateTime(dt) => q.bind(dt.to_string()),
117            #[cfg(feature = "postgres-native")]
118            DbValue::TimestampTz(dt) => q.bind(dt.to_rfc3339()),
119            #[cfg(feature = "postgres-native")]
120            DbValue::Json(j) => q.bind(j.to_string()),
121            #[cfg(feature = "postgres-native")]
122            DbValue::Uuid(u) => q.bind(u.to_string()),
123            #[cfg(feature = "postgres-native")]
124            DbValue::Time(t) => q.bind(t.to_string()),
125            #[cfg(feature = "postgres-native")]
126            DbValue::TextArray(v) => q.bind(crate::value::pg_text_array_literal(v)),
127            #[cfg(feature = "postgres-native")]
128            DbValue::FloatArray(v) => {
129                q.bind(crate::value::pg_float_array_literal(v.iter().map(|x| Some(*x))))
130            }
131            #[cfg(feature = "postgres-native")]
132            DbValue::OptFloatArray(v) => {
133                q.bind(crate::value::pg_float_array_literal(v.iter().copied()))
134            }
135        };
136    }
137    q
138}
139
140// ---------------------------------------------------------------------------
141// BaseHandler
142// ---------------------------------------------------------------------------
143
144/// Core query executor: transactional writes via sqlx, and optionally
145/// analytical reads via a pluggable [`ReadEngine`] (DuckDB or DataFusion).
146pub struct BaseHandler {
147    pool: AnyPool,
148    #[cfg(any(feature = "duckdb", feature = "datafusion"))]
149    read_engine: Option<Box<dyn ReadEngine>>,
150}
151
152impl BaseHandler {
153    /// Create a handler for writes against the given sqlx pool.
154    pub fn new(pool: AnyPool) -> Self {
155        Self {
156            pool,
157            #[cfg(any(feature = "duckdb", feature = "datafusion"))]
158            read_engine: None,
159        }
160    }
161
162    /// Create a handler with an in-memory DuckDB analytical read engine.
163    #[cfg(feature = "duckdb")]
164    pub fn with_duckdb(pool: AnyPool) -> Result<Self, DbkitError> {
165        let engine = crate::read::duckdb::DuckEngine::new_in_memory()?;
166        Ok(Self {
167            pool,
168            read_engine: Some(Box::new(engine)),
169        })
170    }
171
172    /// Create a handler with DuckDB and a live Postgres attachment.
173    ///
174    /// DuckDB queries the Postgres tables directly via the `pg` catalog
175    /// (`SELECT … FROM pg.<schema>.<table>`) without an explicit sync — the
176    /// pre-rewrite zero-copy `ATTACH` pipeline. You can still also `sync_*`
177    /// tables into local memory for faster repeated analytics.
178    #[cfg(feature = "duckdb")]
179    pub fn with_duckdb_attached_postgres(
180        pool: AnyPool,
181        pg_connection_string: &str,
182    ) -> Result<Self, DbkitError> {
183        let engine = crate::read::duckdb::DuckEngine::new_in_memory()?;
184        engine.attach_postgres(pg_connection_string)?;
185        Ok(Self {
186            pool,
187            read_engine: Some(Box::new(engine)),
188        })
189    }
190
191    /// Create a handler with a DataFusion analytical read engine.
192    #[cfg(feature = "datafusion")]
193    pub fn with_datafusion(pool: AnyPool) -> Result<Self, DbkitError> {
194        let engine = crate::read::datafusion::DataFusionEngine::new();
195        Ok(Self {
196            pool,
197            read_engine: Some(Box::new(engine)),
198        })
199    }
200
201    /// Whether an analytical read engine is attached.
202    pub fn has_read_engine(&self) -> bool {
203        #[cfg(any(feature = "duckdb", feature = "datafusion"))]
204        {
205            self.read_engine.is_some()
206        }
207        #[cfg(not(any(feature = "duckdb", feature = "datafusion")))]
208        {
209            false
210        }
211    }
212
213    /// Get a reference to the write pool.
214    pub fn pool(&self) -> &AnyPool {
215        &self.pool
216    }
217
218    /// Unicode NFD normalization — decomposes characters then lowercases.
219    /// Useful for matching names with different Unicode representations.
220    pub fn normalize_name(name: &str) -> String {
221        name.nfd().collect::<String>().to_lowercase()
222    }
223
224    // ==================== UNIFIED WRITE ====================
225
226    /// Execute a write operation against the transactional pool.
227    ///
228    /// Placeholders are backend-native: `$1, $2, …` for Postgres, `?` for
229    /// MySQL/SQLite. sqlx's `Any` driver does not rewrite them, so write the
230    /// SQL for the backend you connected to.
231    pub async fn execute_write(
232        &self,
233        op: WriteOp<'_>,
234    ) -> Result<QueryResult<AnyRow>, DbkitError> {
235        match op {
236            WriteOp::Single {
237                query,
238                params,
239                mode,
240            } => {
241                let q = bind_params(sqlx::query(AssertSqlSafe(query)), &params);
242                match mode {
243                    FetchMode::None => {
244                        q.execute(&self.pool).await?;
245                        Ok(QueryResult::None)
246                    }
247                    FetchMode::One => {
248                        let row = q.fetch_one(&self.pool).await?;
249                        Ok(QueryResult::One(row))
250                    }
251                    FetchMode::Optional => {
252                        let row = q.fetch_optional(&self.pool).await?;
253                        Ok(QueryResult::Optional(row))
254                    }
255                    FetchMode::All => {
256                        let rows = q.fetch_all(&self.pool).await?;
257                        Ok(QueryResult::All(rows))
258                    }
259                }
260            }
261
262            WriteOp::BatchDDL { queries } => {
263                let mut tx = self.pool.begin().await?;
264                for query in queries {
265                    sqlx::query(AssertSqlSafe(*query)).execute(&mut *tx).await?;
266                }
267                tx.commit().await?;
268                Ok(QueryResult::None)
269            }
270
271            WriteOp::BatchParams {
272                query,
273                params_list,
274            } => {
275                if params_list.is_empty() {
276                    return Ok(QueryResult::None);
277                }
278
279                let total = params_list.len();
280                let mut tx = self.pool.begin().await?;
281                let mut failed = 0usize;
282
283                for (idx, params) in params_list.iter().enumerate() {
284                    // sqlx caches the prepared statement per query string on the
285                    // connection, so re-issuing the same query reuses it.
286                    let q = bind_params(sqlx::query(AssertSqlSafe(query)), params);
287                    if let Err(e) = q.execute(&mut *tx).await {
288                        warn!("BatchParams row {}/{} failed: {:?}", idx + 1, total, e);
289                        failed += 1;
290                    }
291                }
292
293                tx.commit().await?;
294
295                if failed > 0 {
296                    warn!(
297                        "BatchParams: {}/{} succeeded, {} failed",
298                        total - failed,
299                        total,
300                        failed
301                    );
302                }
303
304                Ok(QueryResult::None)
305            }
306        }
307    }
308
309    // ==================== UNIFIED READ ====================
310
311    /// Execute an analytical query against the attached read engine, returning
312    /// columnar Arrow [`RecordBatch`]es.
313    ///
314    /// Returns [`DbkitError::NoReadEngine`] if no engine is attached.
315    #[cfg(any(feature = "duckdb", feature = "datafusion"))]
316    pub async fn execute_read(
317        &self,
318        sql: &str,
319        params: &[DbValue],
320    ) -> Result<Vec<RecordBatch>, DbkitError> {
321        self.read_engine
322            .as_ref()
323            .ok_or(DbkitError::NoReadEngine)?
324            .query_arrow(sql, params)
325            .await
326    }
327
328    /// Execute an analytical query and deserialize each row into `T`.
329    ///
330    /// This is the typed-read replacement for the old closure-mapped
331    /// `ReadOp::Standard`: instead of a `|row| …` closure, derive
332    /// `serde::Deserialize` on your row struct. Works for any read engine,
333    /// since it deserializes from the Arrow batches via `serde_arrow`.
334    ///
335    /// ```ignore
336    /// #[derive(serde::Deserialize)]
337    /// struct Item { name: String, qty: i64 }
338    /// let items: Vec<Item> = handler.execute_read_as("SELECT name, qty FROM items", &[]).await?;
339    /// ```
340    #[cfg(any(feature = "duckdb", feature = "datafusion"))]
341    pub async fn execute_read_as<T>(
342        &self,
343        sql: &str,
344        params: &[DbValue],
345    ) -> Result<Vec<T>, DbkitError>
346    where
347        T: serde::de::DeserializeOwned,
348    {
349        let batches = self.execute_read(sql, params).await?;
350        crate::analytical::deserialize_batches(&batches)
351    }
352
353    // ==================== SYNC (transactional -> analytical) ====================
354
355    /// Run a query against the transactional pool and load its result into the
356    /// analytical engine as a named in-memory table.
357    ///
358    /// This is the engine-agnostic replacement for the old DuckDB `ATTACH`
359    /// sync: rows are fetched over sqlx, converted to Arrow, and handed to the
360    /// active read engine. Works for any backend × engine combination.
361    #[cfg(any(feature = "duckdb", feature = "datafusion"))]
362    pub async fn sync_query(
363        &self,
364        name: &str,
365        query: &str,
366        params: &[DbValue],
367    ) -> Result<(), DbkitError> {
368        let engine = self.read_engine.as_ref().ok_or(DbkitError::NoReadEngine)?;
369
370        let q = bind_params(sqlx::query(AssertSqlSafe(query.to_string())), params);
371        let rows = q.fetch_all(&self.pool).await?;
372
373        if let Some(batch) = crate::read::rows_to_record_batch(&rows)? {
374            engine.load_table(name, vec![batch]).await?;
375        }
376        Ok(())
377    }
378
379    /// Copy entire tables from the transactional store into the analytical
380    /// engine, one table per name (`SELECT * FROM {table}`).
381    #[cfg(any(feature = "duckdb", feature = "datafusion"))]
382    pub async fn sync_tables(&self, tables: &[&str]) -> Result<(), DbkitError> {
383        for table in tables {
384            self.sync_query(table, &format!("SELECT * FROM {table}"), &[])
385                .await?;
386        }
387        Ok(())
388    }
389}
dbkit/base_handler.rs

dbkit/
base_handler.rs
