Skip to main content

improv_storage_sql/
duck.rs

1//! DuckDB backend (embedded, via the `duckdb` crate with the `bundled`
2//! feature so no system libduckdb is required).
3//!
4//! DuckDB is file/embedded like SQLite: no out-of-band secret. Its Rust API
5//! mirrors rusqlite (prepared statements, `?` placeholders, typed value refs),
6//! so this module reads the same way as the SQLite path in `backend.rs`.
7//! Values are always bound as parameters; only validated identifiers reach a
8//! query string (see `ident`/`sanitize_ident` in `lib.rs`).
9
10use crate::backend::{Backend, Cell, Table};
11use crate::{Result, SqlError};
12
13/// Open a DuckDB database at `path` (or `:memory:`). Embedded, no secret.
14///
15/// External access is disabled (`enable_external_access = false`), so a stored
16/// query cannot read host files (`read_csv('/etc/passwd')`), fetch URLs or
17/// `ATTACH` other databases. The database stays read-write because the same
18/// handle serves `export_measure`; refresh queries are confined by the
19/// read-only transaction in `duck_query` instead of `access_mode=READ_ONLY`.
20pub fn connect_duckdb(path: &str) -> Result<Backend> {
21    let err = |e: duckdb::Error| SqlError::Connect(format!("duckdb connect failed: {e}"));
22    let config = duckdb::Config::default()
23        .enable_external_access(false)
24        .and_then(|c| c.enable_autoload_extension(false))
25        .map_err(err)?;
26    let conn = duckdb::Connection::open_with_flags(path, config).map_err(err)?;
27    Ok(Backend::Duckdb(conn))
28}
29
30/// Run `sql` inside a DuckDB `READ ONLY` transaction, so it cannot write
31/// whatever statement type it turns out to be.
32pub(crate) fn duck_query(conn: &duckdb::Connection, sql: &str) -> Result<Table> {
33    let txn = |s: &str| {
34        conn.execute_batch(s)
35            .map_err(|e| SqlError::Import(format!("duckdb {s} failed: {e}")))
36    };
37    txn("BEGIN TRANSACTION READ ONLY")?;
38    let out = duck_query_inner(conn, sql);
39    match out {
40        Ok(_) => txn("COMMIT")?,
41        Err(_) => {
42            let _ = txn("ROLLBACK");
43        }
44    }
45    out
46}
47
48fn duck_query_inner(conn: &duckdb::Connection, sql: &str) -> Result<Table> {
49    let mut stmt = conn
50        .prepare(sql)
51        .map_err(|e| SqlError::Import(format!("duckdb prepare failed: {e}")))?;
52    let mut rows_iter = stmt
53        .query([])
54        .map_err(|e| SqlError::Import(format!("duckdb query failed: {e}")))?;
55
56    let mut rows = Vec::new();
57    let mut columns: Option<Vec<String>> = None;
58    while let Some(row) = rows_iter
59        .next()
60        .map_err(|e| SqlError::Import(format!("duckdb row failed: {e}")))?
61    {
62        // Column names are available once we have a statement/row; capture once.
63        if columns.is_none() {
64            let stmt_ref = row.as_ref();
65            columns = Some(
66                stmt_ref
67                    .column_names()
68                    .into_iter()
69                    .map(|s| s.to_string())
70                    .collect(),
71            );
72        }
73        let ncols = columns.as_ref().map(|c| c.len()).unwrap_or(0);
74        let mut cells = Vec::with_capacity(ncols);
75        for i in 0..ncols {
76            let vr = row
77                .get_ref(i)
78                .map_err(|e| SqlError::Import(format!("duckdb column {i} failed: {e}")))?;
79            cells.push(duck_cell(vr)?);
80        }
81        rows.push(cells);
82    }
83
84    // Empty result still needs column names: fall back to the statement metadata.
85    let columns = match columns {
86        Some(c) => c,
87        None => stmt
88            .column_names()
89            .into_iter()
90            .map(|s| s.to_string())
91            .collect(),
92    };
93    Ok(Table { columns, rows })
94}
95
96/// Map a DuckDB value ref to a [`Cell`]. Integer family → Int, float/decimal
97/// family → Float, text → Text, boolean → Bool, null → Null.
98fn duck_cell(vr: duckdb::types::ValueRef<'_>) -> Result<Cell> {
99    use duckdb::types::ValueRef as V;
100    Ok(match vr {
101        V::Null => Cell::Null,
102        V::Boolean(b) => Cell::Bool(b),
103        V::TinyInt(v) => Cell::Int(v as i64),
104        V::SmallInt(v) => Cell::Int(v as i64),
105        V::Int(v) => Cell::Int(v as i64),
106        V::BigInt(v) => Cell::Int(v),
107        V::HugeInt(v) => Cell::Float(v as f64),
108        V::UTinyInt(v) => Cell::Int(v as i64),
109        V::USmallInt(v) => Cell::Int(v as i64),
110        V::UInt(v) => Cell::Int(v as i64),
111        V::UBigInt(v) => Cell::Int(v as i64),
112        V::Float(v) => Cell::Float(v as f64),
113        V::Double(v) => Cell::Float(v),
114        V::Decimal(v) => Cell::Float(v.to_string().parse::<f64>().unwrap_or(0.0)),
115        V::Text(t) => Cell::Text(String::from_utf8_lossy(t).into_owned()),
116        other => {
117            // Dates/timestamps/blobs/etc.: stringify defensively rather than fail.
118            Cell::Text(format!("{other:?}"))
119        }
120    })
121}
122
123pub(crate) fn duck_execute(conn: &duckdb::Connection, sql: &str, params: &[Cell]) -> Result<()> {
124    let bound: Vec<duckdb::types::Value> = params.iter().map(cell_to_duck).collect();
125    conn.execute(sql, duckdb::params_from_iter(bound))
126        .map_err(|e| SqlError::Export(format!("duckdb execute failed: {e}")))?;
127    Ok(())
128}
129
130fn cell_to_duck(c: &Cell) -> duckdb::types::Value {
131    use duckdb::types::Value as V;
132    match c {
133        Cell::Null => V::Null,
134        Cell::Int(i) => V::BigInt(*i),
135        Cell::Float(f) => V::Double(*f),
136        Cell::Text(t) => V::Text(t.clone()),
137        Cell::Bool(b) => V::Boolean(*b),
138    }
139}
140
141#[cfg(test)]
142#[cfg_attr(coverage_nightly, coverage(off))]
143mod tests {
144    use super::*;
145
146    fn mem() -> duckdb::Connection {
147        duckdb::Connection::open_in_memory().unwrap()
148    }
149
150    #[test]
151    fn duck_maps_every_value_kind() {
152        let t = duck_query(
153            &mem(),
154            "SELECT NULL AS n, true AS b, 1::TINYINT AS ti, 2::SMALLINT AS si, \
155             3::INTEGER AS i, 4::BIGINT AS bi, 5::HUGEINT AS hi, 6::UTINYINT AS uti, \
156             7::USMALLINT AS usi, 8::UINTEGER AS ui, 9::UBIGINT AS ubi, \
157             1.5::FLOAT AS f, 2.5::DOUBLE AS d, 3.25::DECIMAL(5,2) AS dec, \
158             'x' AS t, DATE '2025-01-02' AS dt",
159        )
160        .unwrap();
161        assert_eq!(t.columns.len(), 16);
162        let r = &t.rows[0];
163        assert_eq!(
164            r[..15],
165            [
166                Cell::Null,
167                Cell::Bool(true),
168                Cell::Int(1),
169                Cell::Int(2),
170                Cell::Int(3),
171                Cell::Int(4),
172                Cell::Float(5.0),
173                Cell::Int(6),
174                Cell::Int(7),
175                Cell::Int(8),
176                Cell::Int(9),
177                Cell::Float(1.5),
178                Cell::Float(2.5),
179                Cell::Float(3.25),
180                Cell::Text("x".into()),
181            ]
182        );
183        // Dates are stringified defensively, not an error.
184        assert!(matches!(&r[15], Cell::Text(s) if !s.is_empty()));
185    }
186
187    #[test]
188    fn empty_result_keeps_column_names() {
189        let t = duck_query(&mem(), "SELECT 1 AS a, 'q' AS b WHERE false").unwrap();
190        assert_eq!(t.columns, ["a", "b"]);
191        assert!(t.rows.is_empty());
192    }
193
194    #[test]
195    fn execute_binds_every_cell_kind_and_reports_errors() {
196        let c = mem();
197        duck_execute(
198            &c,
199            "CREATE TABLE t (a INTEGER, b BIGINT, c DOUBLE, d VARCHAR, e BOOLEAN)",
200            &[],
201        )
202        .unwrap();
203        duck_execute(
204            &c,
205            "INSERT INTO t VALUES (?, ?, ?, ?, ?)",
206            &[
207                Cell::Null,
208                Cell::Int(2),
209                Cell::Float(0.5),
210                Cell::Text("s".into()),
211                Cell::Bool(false),
212            ],
213        )
214        .unwrap();
215        let t = duck_query(&c, "SELECT * FROM t").unwrap();
216        assert_eq!(
217            t.rows,
218            [[
219                Cell::Null,
220                Cell::Int(2),
221                Cell::Float(0.5),
222                Cell::Text("s".into()),
223                Cell::Bool(false)
224            ]]
225        );
226        let err = duck_execute(&c, "INSERT INTO nope VALUES (1)", &[]).unwrap_err();
227        assert!(err.to_string().contains("duckdb execute failed"), "{err}");
228        let err = duck_query(&c, "FROBNICATE everything").err().unwrap();
229        assert!(err.to_string().contains("duckdb prepare failed"), "{err}");
230        // The failed read rolled back; the connection still works.
231        assert_eq!(
232            duck_query(&c, "SELECT count(*) AS n FROM t")
233                .unwrap()
234                .rows
235                .len(),
236            1
237        );
238    }
239
240    #[test]
241    fn connect_failure_is_a_connect_error() {
242        let Err(err) = connect_duckdb("/nonexistent-dir-improv/x/y.duckdb") else {
243            panic!("opened a path under a missing directory");
244        };
245        assert!(err.to_string().contains("duckdb connect failed"), "{err}");
246    }
247}