sql_traits/
conn.rs

1// Copyright 2026 Greg Burd
2//
3// Licensed under the Apache License, Version 2.0 (the "License"); you may not use
4// this file except in compliance with the License. You may obtain a copy of the
5// License at http://www.apache.org/licenses/LICENSE-2.0
6// Unless required by applicable law or agreed to in writing, software distributed
7// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
8// CONDITIONS OF ANY KIND, either express or implied. See the License for the
9// specific language governing permissions and limitations under the License.
10
11//! The storage seam: the few SQL operations the mentat engine (transactor,
12//! query runner, projector, pull) needs from a database, so the same engine can
13//! run on SQLite (the CLI and library) and on DuckDB (the DuckDB extension).
14//!
15//! `mentat_sql::SqliteConn` implements [`SqlConn`] over a `rusqlite::Connection`,
16//! keeping its statement cache, so the SQLite path runs the same SQL it always
17//! has. Where another engine
18//! needs different SQL, callers branch on [`SqlConn::dialect`].
19
20use std::fmt;
21
22/// A SQL value as the engine stores it: SQLite's five storage classes. DuckDB
23/// maps these onto the members of its `v` column's `UNION` type.
24#[derive(Clone, Debug, PartialEq)]
25pub enum SqlValue {
26    Null,
27    Integer(i64),
28    Real(f64),
29    Text(String),
30    Blob(Vec<u8>),
31}
32
33impl From<i64> for SqlValue {
34    fn from(v: i64) -> Self {
35        SqlValue::Integer(v)
36    }
37}
38
39impl From<i32> for SqlValue {
40    fn from(v: i32) -> Self {
41        SqlValue::Integer(v.into())
42    }
43}
44
45impl From<bool> for SqlValue {
46    fn from(v: bool) -> Self {
47        SqlValue::Integer(v.into())
48    }
49}
50
51impl From<&str> for SqlValue {
52    fn from(v: &str) -> Self {
53        SqlValue::Text(v.to_string())
54    }
55}
56
57impl From<String> for SqlValue {
58    fn from(v: String) -> Self {
59        SqlValue::Text(v)
60    }
61}
62
63impl SqlValue {
64    pub fn as_i64(&self) -> Option<i64> {
65        match self {
66            SqlValue::Integer(i) => Some(*i),
67            _ => None,
68        }
69    }
70}
71
72/// Which SQL dialect a connection speaks. The engine's SQL is SQLite's; a
73/// handful of statements (DDL, temp tables, fulltext) differ on DuckDB.
74#[derive(Clone, Copy, Debug, PartialEq, Eq)]
75pub enum Dialect {
76    Sqlite,
77    DuckDb,
78}
79
80/// An error from the underlying database, carried as text so the trait stays
81/// engine-neutral. `sqlite_code` is set for SQLite errors that callers inspect
82/// (e.g. a constraint violation).
83#[derive(Clone, Debug, PartialEq)]
84pub struct SqlError {
85    pub message: String,
86    pub sqlite_code: Option<i32>,
87}
88
89impl fmt::Display for SqlError {
90    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
91        f.write_str(&self.message)
92    }
93}
94
95impl std::error::Error for SqlError {}
96
97impl SqlError {
98    pub fn new(message: impl Into<String>) -> Self {
99        SqlError {
100            message: message.into(),
101            sqlite_code: None,
102        }
103    }
104}
105
106/// One result row. Columns are 0-based.
107pub trait SqlRow {
108    fn get_value(&self, idx: usize) -> Result<SqlValue, SqlError>;
109
110    fn column_count(&self) -> usize;
111
112    fn get_i64(&self, idx: usize) -> Result<i64, SqlError> {
113        match self.get_value(idx)? {
114            SqlValue::Integer(i) => Ok(i),
115            other => Err(SqlError::new(format!(
116                "column {idx}: expected an integer, got {other:?}"
117            ))),
118        }
119    }
120
121    fn get_bool(&self, idx: usize) -> Result<bool, SqlError> {
122        Ok(self.get_i64(idx)? != 0)
123    }
124
125    fn get_text(&self, idx: usize) -> Result<String, SqlError> {
126        match self.get_value(idx)? {
127            SqlValue::Text(s) => Ok(s),
128            other => Err(SqlError::new(format!(
129                "column {idx}: expected text, got {other:?}"
130            ))),
131        }
132    }
133
134    fn get_opt_i64(&self, idx: usize) -> Result<Option<i64>, SqlError> {
135        match self.get_value(idx)? {
136            SqlValue::Null => Ok(None),
137            SqlValue::Integer(i) => Ok(Some(i)),
138            other => Err(SqlError::new(format!(
139                "column {idx}: expected an integer or NULL, got {other:?}"
140            ))),
141        }
142    }
143}
144
145impl SqlRow for Vec<SqlValue> {
146    fn get_value(&self, idx: usize) -> Result<SqlValue, SqlError> {
147        self.get(idx)
148            .cloned()
149            .ok_or_else(|| SqlError::new(format!("column {idx} out of range")))
150    }
151
152    fn column_count(&self) -> usize {
153        self.len()
154    }
155}
156
157/// The SQL operations the engine needs. Positional parameters are `?`; named
158/// parameters (`$name`, used by the query builder) go through
159/// [`SqlConn::query_named`].
160pub trait SqlConn {
161    fn dialect(&self) -> Dialect;
162
163    /// Run statements with no parameters and no results.
164    fn execute_batch(&self, sql: &str) -> Result<(), SqlError>;
165
166    /// Run one statement; returns the number of rows changed.
167    fn execute(&self, sql: &str, params: &[SqlValue]) -> Result<usize, SqlError>;
168
169    /// Run a query, calling `f` for every row in order. `f` returning an error
170    /// stops the scan and propagates it.
171    fn query(
172        &self,
173        sql: &str,
174        params: &[SqlValue],
175        f: &mut dyn FnMut(&dyn SqlRow) -> Result<(), SqlError>,
176    ) -> Result<(), SqlError>;
177
178    /// As [`SqlConn::query`], with named parameters (`$name` -> value).
179    fn query_named(
180        &self,
181        sql: &str,
182        params: &[(&str, SqlValue)],
183        f: &mut dyn FnMut(&dyn SqlRow) -> Result<(), SqlError>,
184    ) -> Result<(), SqlError>;
185
186    /// The most bound parameters one statement may take.
187    fn max_variables(&self) -> usize;
188
189    /// Number of result columns of `sql` (prepared, not run).
190    fn column_count(&self, sql: &str) -> Result<usize, SqlError>;
191
192    /// Resolve unqualified table names in `schema` (DuckDB's search path;
193    /// SQLite has no schemas and never calls this).
194    fn use_schema(&self, schema: &str) -> Result<(), SqlError> {
195        self.execute_batch(&format!("SET search_path = '{schema}'"))
196    }
197
198    /// As [`SqlConn::query`], but the connection may keep the statement
199    /// prepared for reuse. Only for a single-table statement on the store's
200    /// own bookkeeping tables in which every value that varies is a bound
201    /// parameter: on DuckDB 1.5.x a reused plan keeps the statistics it was
202    /// planned with (duckdb/duckdb#24696; joins: #24182), so a literal, or a
203    /// join, can keep matching nothing after rows arrive.
204    fn query_cached(
205        &self,
206        sql: &str,
207        params: &[SqlValue],
208        f: &mut dyn FnMut(&dyn SqlRow) -> Result<(), SqlError>,
209    ) -> Result<(), SqlError> {
210        self.query(sql, params, f)
211    }
212
213    /// As [`SqlConn::execute`], cached under the rules of
214    /// [`SqlConn::query_cached`].
215    fn execute_cached(&self, sql: &str, params: &[SqlValue]) -> Result<usize, SqlError> {
216        self.execute(sql, params)
217    }
218}
219
220/// Convenience helpers over any [`SqlConn`].
221pub trait SqlConnExt: SqlConn {
222    /// Collect every row (`f` maps one row).
223    fn query_rows<T>(
224        &self,
225        sql: &str,
226        params: &[SqlValue],
227        mut f: impl FnMut(&dyn SqlRow) -> Result<T, SqlError>,
228    ) -> Result<Vec<T>, SqlError> {
229        let mut out = Vec::new();
230        self.query(sql, params, &mut |row| {
231            out.push(f(row)?);
232            Ok(())
233        })?;
234        Ok(out)
235    }
236
237    /// The first row, if any.
238    fn query_opt<T>(
239        &self,
240        sql: &str,
241        params: &[SqlValue],
242        mut f: impl FnMut(&dyn SqlRow) -> Result<T, SqlError>,
243    ) -> Result<Option<T>, SqlError> {
244        let mut out = None;
245        self.query(sql, params, &mut |row| {
246            if out.is_none() {
247                out = Some(f(row)?);
248            }
249            Ok(())
250        })?;
251        Ok(out)
252    }
253}
254
255impl<C: SqlConn + ?Sized> SqlConnExt for C {}
sql_traits/conn.rs

sql_traits/
conn.rs
