citadel_sql/
prepared.rs

1//! Prepared statements: parse + compile once, execute many times with parameters.
2
3use std::sync::Arc;
4
5use rustc_hash::FxHashMap;
6
7use crate::connection::Connection;
8use crate::error::{Result, SqlError};
9use crate::executor::compile::RowSourceIter;
10use crate::executor::helpers::expr_display_name;
11use crate::executor::{self, CompiledPlan};
12use crate::parser::{QueryBody, SelectColumn, SelectQuery, SelectStmt, Statement};
13use crate::schema::SchemaManager;
14use crate::types::{ExecutionResult, QueryResult, Value};
15
16/// A prepared SQL statement bound to a `Connection`.
17pub struct PreparedStatement<'c, 'db> {
18    conn: &'c Connection<'db>,
19    sql: String,
20    ast: Arc<Statement>,
21    compiled: Option<Arc<dyn CompiledPlan>>,
22    schema_gen: u64,
23    param_count: usize,
24    columns: Vec<String>,
25    column_index: FxHashMap<String, usize>,
26    readonly: bool,
27    is_explain: bool,
28}
29
30struct Compiled {
31    ast: Arc<Statement>,
32    plan: Option<Arc<dyn CompiledPlan>>,
33    schema_gen: u64,
34    param_count: usize,
35    columns: Vec<String>,
36}
37
38impl<'c, 'db> PreparedStatement<'c, 'db> {
39    pub(crate) fn new(conn: &'c Connection<'db>, sql: &str) -> Result<Self> {
40        let c = compile_for_sql(conn, sql)?;
41        let readonly = matches!(*c.ast, Statement::Select(_) | Statement::Explain(_));
42        let is_explain = matches!(*c.ast, Statement::Explain(_));
43        let mut column_index =
44            FxHashMap::with_capacity_and_hasher(c.columns.len(), Default::default());
45        for (i, name) in c.columns.iter().enumerate() {
46            column_index.entry(name.clone()).or_insert(i);
47        }
48        Ok(Self {
49            conn,
50            sql: sql.to_string(),
51            ast: c.ast,
52            compiled: c.plan,
53            schema_gen: c.schema_gen,
54            param_count: c.param_count,
55            columns: c.columns,
56            column_index,
57            readonly,
58            is_explain,
59        })
60    }
61
62    /// The original SQL text.
63    pub fn sql(&self) -> &str {
64        &self.sql
65    }
66
67    /// Number of positional parameters (`$1`, `$2`, ...) this statement expects.
68    pub fn param_count(&self) -> usize {
69        self.param_count
70    }
71
72    /// Alias of [`Self::param_count`] matching rusqlite's name.
73    pub fn parameter_count(&self) -> usize {
74        self.param_count
75    }
76
77    /// Number of output columns. Zero for non-SELECT statements.
78    pub fn column_count(&self) -> usize {
79        self.columns.len()
80    }
81
82    /// Output column names in declaration order.
83    pub fn column_names(&self) -> &[String] {
84        &self.columns
85    }
86
87    /// Output column name at index `i`, if any.
88    pub fn column_name(&self, i: usize) -> Option<&str> {
89        self.columns.get(i).map(|s| s.as_str())
90    }
91
92    /// Position of the column named `name`, if present.
93    pub fn column_index(&self, name: &str) -> Option<usize> {
94        self.column_index.get(name).copied()
95    }
96
97    /// True if the statement is read-only (SELECT or EXPLAIN).
98    pub fn readonly(&self) -> bool {
99        self.readonly
100    }
101
102    /// True if the statement is an EXPLAIN.
103    pub fn is_explain(&self) -> bool {
104        self.is_explain
105    }
106
107    /// Execute the statement; returns rows affected (0 for SELECT/DDL).
108    pub fn execute(&self, params: &[Value]) -> Result<u64> {
109        match self.run(params)? {
110            ExecutionResult::RowsAffected(n) => Ok(n),
111            ExecutionResult::Query(_) | ExecutionResult::Ok => Ok(0),
112        }
113    }
114
115    /// Execute and return a stepping `Rows<'_>` iterator.
116    pub fn query(&self, params: &[Value]) -> Result<Rows<'_>> {
117        if params.len() != self.param_count {
118            return Err(SqlError::ParameterCountMismatch {
119                expected: self.param_count,
120                got: params.len(),
121            });
122        }
123        if self.conn.inner.borrow().schema.generation() == self.schema_gen {
124            if let Some(plan) = &self.compiled {
125                if let Some(stream) = try_stream_via_plan(self, plan.as_ref(), params) {
126                    return Ok(Rows::streaming(stream));
127                }
128            }
129        }
130        let (columns, rows) = match self.run(params)? {
131            ExecutionResult::Query(qr) => (qr.columns, qr.rows),
132            ExecutionResult::RowsAffected(_) | ExecutionResult::Ok => {
133                (self.columns.clone(), Vec::new())
134            }
135        };
136        Ok(Rows::materialized(columns, rows))
137    }
138
139    /// Execute and return the fully-materialized `QueryResult`.
140    pub fn query_collect(&self, params: &[Value]) -> Result<QueryResult> {
141        match self.run(params)? {
142            ExecutionResult::Query(qr) => Ok(qr),
143            ExecutionResult::RowsAffected(n) => Ok(QueryResult {
144                columns: vec!["rows_affected".into()],
145                rows: vec![vec![Value::Integer(n as i64)]],
146            }),
147            ExecutionResult::Ok => Ok(QueryResult {
148                columns: vec![],
149                rows: vec![],
150            }),
151        }
152    }
153
154    /// Run the query and pass the first row to `f`.
155    pub fn query_row<T, F>(&self, params: &[Value], f: F) -> Result<T>
156    where
157        F: FnOnce(&Row<'_>) -> Result<T>,
158    {
159        let mut rows = self.query(params)?;
160        match rows.next()? {
161            Some(row) => f(&row),
162            None => Err(SqlError::QueryReturnedNoRows),
163        }
164    }
165
166    /// True if the query returns at least one row (DML returns `n > 0`).
167    pub fn exists(&self, params: &[Value]) -> Result<bool> {
168        if params.len() != self.param_count {
169            return Err(SqlError::ParameterCountMismatch {
170                expected: self.param_count,
171                got: params.len(),
172            });
173        }
174        if self.conn.inner.borrow().schema.generation() == self.schema_gen {
175            if let Some(plan) = &self.compiled {
176                if let Some(mut stream) = try_stream_via_plan(self, plan.as_ref(), params) {
177                    return Ok(stream.next_row()?.is_some());
178                }
179            }
180        }
181        match self.run(params)? {
182            ExecutionResult::Query(qr) => Ok(!qr.rows.is_empty()),
183            ExecutionResult::RowsAffected(n) => Ok(n > 0),
184            ExecutionResult::Ok => Ok(false),
185        }
186    }
187
188    fn run(&self, params: &[Value]) -> Result<ExecutionResult> {
189        if params.len() != self.param_count {
190            return Err(SqlError::ParameterCountMismatch {
191                expected: self.param_count,
192                got: params.len(),
193            });
194        }
195        let mut inner = self.conn.inner.borrow_mut();
196        if inner.schema.generation() == self.schema_gen {
197            return inner.execute_prepared(self.conn.db, &self.ast, self.compiled.as_ref(), params);
198        }
199        let c = compile_inside(&mut inner, &self.sql)?;
200        if c.param_count != self.param_count {
201            return Err(SqlError::ParameterCountMismatch {
202                expected: self.param_count,
203                got: c.param_count,
204            });
205        }
206        inner.execute_prepared(self.conn.db, &c.ast, c.plan.as_ref(), params)
207    }
208}
209
210/// Stepping iterator over query rows. Obtained from [`PreparedStatement::query`].
211pub struct Rows<'a> {
212    source: RowSource<'a>,
213    columns: Vec<String>,
214    buf: Vec<Value>,
215}
216
217enum RowSource<'a> {
218    Materialized(std::vec::IntoIter<Vec<Value>>),
219    Streaming(Box<dyn RowSourceIter + 'a>),
220}
221
222impl<'a> Rows<'a> {
223    fn materialized(columns: Vec<String>, rows: Vec<Vec<Value>>) -> Self {
224        Self {
225            source: RowSource::Materialized(rows.into_iter()),
226            columns,
227            buf: Vec::new(),
228        }
229    }
230
231    fn streaming(source: Box<dyn RowSourceIter + 'a>) -> Self {
232        let columns = source.columns().to_vec();
233        Self {
234            source: RowSource::Streaming(source),
235            columns,
236            buf: Vec::new(),
237        }
238    }
239
240    /// Step to the next row, if any.
241    #[allow(clippy::should_implement_trait)]
242    pub fn next(&mut self) -> Result<Option<Row<'_>>> {
243        let next: Option<Vec<Value>> = match &mut self.source {
244            RowSource::Materialized(iter) => iter.next(),
245            RowSource::Streaming(stream) => stream.next_row()?,
246        };
247        match next {
248            Some(values) => {
249                self.buf = values;
250                Ok(Some(Row {
251                    columns: &self.columns,
252                    values: &self.buf,
253                }))
254            }
255            None => Ok(None),
256        }
257    }
258
259    /// Number of output columns.
260    pub fn column_count(&self) -> usize {
261        self.columns.len()
262    }
263
264    /// Output column names.
265    pub fn column_names(&self) -> &[String] {
266        &self.columns
267    }
268
269    /// Drain all remaining rows into a [`QueryResult`].
270    pub fn collect(mut self) -> Result<QueryResult> {
271        let mut rows = Vec::new();
272        while let Some(row) = self.next()? {
273            rows.push(row.to_vec());
274        }
275        Ok(QueryResult {
276            columns: self.columns,
277            rows,
278        })
279    }
280}
281
282/// A single row produced by [`Rows::next`].
283pub struct Row<'a> {
284    columns: &'a [String],
285    values: &'a [Value],
286}
287
288impl<'a> Row<'a> {
289    /// Value at column index `i`, if present.
290    pub fn get(&self, i: usize) -> Option<&Value> {
291        self.values.get(i)
292    }
293
294    /// Value of the column named `name`, if present.
295    pub fn get_by_name(&self, name: &str) -> Option<&Value> {
296        self.columns
297            .iter()
298            .position(|c| c == name)
299            .and_then(|i| self.values.get(i))
300    }
301
302    /// Number of columns in this row.
303    pub fn column_count(&self) -> usize {
304        self.values.len()
305    }
306
307    /// Name of the column at index `i`, if present.
308    pub fn column_name(&self, i: usize) -> Option<&str> {
309        self.columns.get(i).map(|s| s.as_str())
310    }
311
312    /// Borrow the raw values slice.
313    pub fn as_slice(&self) -> &[Value] {
314        self.values
315    }
316
317    /// Clone into an owned `Vec<Value>`.
318    pub fn to_vec(&self) -> Vec<Value> {
319        self.values.to_vec()
320    }
321}
322
323fn compile_for_sql(conn: &Connection<'_>, sql: &str) -> Result<Compiled> {
324    let mut inner = conn.inner.borrow_mut();
325    compile_inside(&mut inner, sql)
326}
327
328fn compile_inside(
329    inner: &mut crate::connection::ConnectionInner<'_>,
330    sql: &str,
331) -> Result<Compiled> {
332    let (ast, param_count) = inner.get_or_parse(sql)?;
333    let schema_gen = inner.schema.generation();
334    let plan = executor::compile(&inner.schema, &ast);
335    if let Some(p) = &plan {
336        if let Some(entry) = inner.stmt_cache.get_mut(sql) {
337            entry.compiled = Some(Arc::clone(p));
338        }
339    }
340    let columns = derive_columns(&ast, &inner.schema);
341    Ok(Compiled {
342        ast,
343        plan,
344        schema_gen,
345        param_count,
346        columns,
347    })
348}
349
350fn derive_columns(stmt: &Statement, schema: &SchemaManager) -> Vec<String> {
351    match stmt {
352        Statement::Select(sq) => derive_select_columns(sq, schema),
353        Statement::Explain(_) => vec!["plan".into()],
354        _ => Vec::new(),
355    }
356}
357
358fn derive_select_columns(sq: &SelectQuery, schema: &SchemaManager) -> Vec<String> {
359    derive_body_columns(&sq.body, schema)
360}
361
362fn derive_body_columns(body: &QueryBody, schema: &SchemaManager) -> Vec<String> {
363    match body {
364        QueryBody::Select(sel) => derive_from_select_stmt(sel, schema),
365        QueryBody::Compound(cs) => derive_body_columns(&cs.left, schema),
366        QueryBody::Insert(_) | QueryBody::Update(_) | QueryBody::Delete(_) => Vec::new(),
367    }
368}
369
370fn try_stream_via_plan<'db>(
371    stmt: &PreparedStatement<'_, 'db>,
372    plan: &dyn CompiledPlan,
373    params: &[Value],
374) -> Option<Box<dyn RowSourceIter + 'db>> {
375    let inner = stmt.conn.inner.borrow();
376    if inner.active_txn_is_some() {
377        return None;
378    }
379    plan.try_stream(stmt.conn.db, &inner.schema, &stmt.ast, params)
380}
381
382fn derive_from_select_stmt(sel: &SelectStmt, schema: &SchemaManager) -> Vec<String> {
383    let lower = sel.from.to_ascii_lowercase();
384    let table_columns = schema.get(&lower).map(|ts| ts.columns.as_slice());
385    let mut out = Vec::new();
386    for col in &sel.columns {
387        match col {
388            SelectColumn::AllColumns | SelectColumn::AllFromOld | SelectColumn::AllFromNew => {
389                if let Some(cols) = table_columns {
390                    for c in cols {
391                        out.push(c.name.clone());
392                    }
393                }
394            }
395            SelectColumn::Expr { alias: Some(a), .. } => out.push(a.clone()),
396            SelectColumn::Expr { expr, alias: None } => out.push(expr_display_name(expr)),
397        }
398    }
399    out
400}
citadel_sql/prepared.rs

citadel_sql/
prepared.rs
