rust_query/
query.rs

1use std::{
2    cell::{Cell, RefCell},
3    collections::BTreeMap,
4    fmt::Debug,
5    marker::PhantomData,
6    ops::{Deref, DerefMut},
7};
8
9use rusqlite::Connection;
10use sea_query::SqliteQueryBuilder;
11use sea_query_rusqlite::{RusqliteBinder, RusqliteValues};
12use self_cell::{MutBorrow, self_cell};
13
14use crate::{
15    IntoExpr,
16    alias::MyAlias,
17    rows::Rows,
18    select::{Cacher, DynPrepared, IntoSelect, Prepared, Row, SelectImpl},
19    transaction::TXN,
20    value::{DynTypedExpr, OrdTyp},
21};
22
23/// This is the type used by the [crate::Transaction::query] method.
24pub struct Query<'t, 'inner, S> {
25    pub(crate) phantom: PhantomData<&'t ()>,
26    pub(crate) q: Rows<'inner, S>,
27}
28
29impl<'inner, S> Deref for Query<'_, 'inner, S> {
30    type Target = Rows<'inner, S>;
31
32    fn deref(&self) -> &Self::Target {
33        &self.q
34    }
35}
36
37impl<S> DerefMut for Query<'_, '_, S> {
38    fn deref_mut(&mut self) -> &mut Self::Target {
39        &mut self.q
40    }
41}
42
43type Stmt<'x> = rusqlite::CachedStatement<'x>;
44type RRows<'a> = rusqlite::Rows<'a>;
45
46self_cell!(
47    pub struct OwnedRows<'x> {
48        owner: MutBorrow<Stmt<'x>>,
49
50        #[covariant]
51        dependent: RRows,
52    }
53);
54
55/// Lazy iterator over rows from a query.
56///
57/// This is currently invariant in `'inner` due to [MutBorrow].
58/// Would be nice to relax this variance in the future.
59pub struct Iter<'inner, O> {
60    // The actual OwnedRows is stored in a thread local
61    inner_phantom: PhantomData<(OwnedRows<'inner>, *const ())>,
62    inner: usize,
63
64    prepared: DynPrepared<O>,
65    cached: Vec<MyAlias>,
66}
67
68impl<O> Iterator for Iter<'_, O> {
69    type Item = O;
70
71    fn next(&mut self) -> Option<Self::Item> {
72        TXN.with_borrow_mut(|combi| {
73            let combi = combi.as_mut().unwrap();
74            combi.with_dependent_mut(|_txn, row_store| {
75                // If rows is already dropped then we just return None.
76                // This can happen if this is called in a thread_local destructor or something.
77                let rows = row_store.get_mut(self.inner)?;
78                rows.with_dependent_mut(|_, rows| {
79                    let row = rows.next().unwrap()?;
80                    Some(self.prepared.call(Row::new(row, &self.cached)))
81                })
82            })
83        })
84    }
85}
86
87impl<O> Drop for Iter<'_, O> {
88    fn drop(&mut self) {
89        TXN.with_borrow_mut(|combi| {
90            let combi = combi.as_mut().unwrap();
91            combi.with_dependent_mut(|_txn, row_store| {
92                // If the rows is already dropped that is fine.
93                // This can happen if this is called in a thread_local destructor or something.
94                row_store.try_remove(self.inner);
95            })
96        })
97    }
98}
99
100impl<'t, 'inner, S> Query<'t, 'inner, S> {
101    /// Turn a database query into a [Vec] of results.
102    ///
103    /// The order of rows that is returned is unstable. This means that the order may change between any two
104    /// executions of the exact same query. If a specific order (or even a consistent order) is required,
105    /// then you have to use something like [slice::sort].
106    pub fn into_vec<O>(&self, select: impl IntoSelect<'inner, S, Out = O>) -> Vec<O> {
107        self.into_iter(select).collect()
108    }
109
110    /// Turn a database query into an iterator of results.
111    ///
112    /// The order of rows that is returned is unstable. This means that the order may change between any two
113    /// executions of the exact same query. If a specific order (or even a consistent order) is required,
114    /// then you have to use something like [slice::sort]. See also [Self::order_by].
115    pub fn into_iter<O>(&self, select: impl IntoSelect<'inner, S, Out = O>) -> Iter<'t, O> {
116        self.order_by().into_iter(select)
117    }
118
119    /// Use [Self::order_by] to refine the order or retrieved rows (partially).
120    /// This is useful if not all rows are retrieved, e.g. for top N style queries.
121    ///
122    /// Every additional call to [OrderBy::asc] and [OrderBy::desc] refines the order further.
123    /// This means that e.g. `order_by().asc(category).asc(priority)`, will have all items
124    /// with the same `category` grouped together and only within the group are items sorted
125    /// by priority.
126    pub fn order_by<'q>(&'q self) -> OrderBy<'q, 't, 'inner, S> {
127        OrderBy {
128            query: self,
129            order: Vec::new(),
130        }
131    }
132}
133
134/// This is an immutable borrow of [Query] that can be sorted.
135///
136/// Use [Self::asc] and [Self::desc] to refine the order of the returned rows.
137#[derive(Clone)]
138pub struct OrderBy<'q, 't, 'inner, S> {
139    query: &'q Query<'t, 'inner, S>,
140    order: Vec<(DynTypedExpr, sea_query::Order)>,
141}
142
143impl<'t, 'inner, S> OrderBy<'_, 't, 'inner, S> {
144    /// Add an additional value to sort on in ascending order.
145    pub fn asc<'q, T: OrdTyp>(mut self, key: impl IntoExpr<'inner, S, Typ = T>) -> Self {
146        self.order
147            .push((DynTypedExpr::erase(key), sea_query::Order::Asc));
148        self
149    }
150
151    /// Add an additional value to sort on in descending order.
152    pub fn desc<'q, T: OrdTyp>(mut self, key: impl IntoExpr<'inner, S, Typ = T>) -> Self {
153        self.order
154            .push((DynTypedExpr::erase(key), sea_query::Order::Desc));
155        self
156    }
157
158    /// Turn a database query into an iterator of results.
159    ///
160    /// Results are ordered as specified by [Self::asc] and [Self::desc].
161    ///
162    /// Rows of which the order is not determined by the calls to [Self::asc] and [Self::desc],
163    /// are returned in unspecified order. See also [Query::into_iter].
164    pub fn into_iter<O>(&self, select: impl IntoSelect<'inner, S, Out = O>) -> Iter<'t, O> {
165        let mut cacher = Cacher::new();
166        let prepared = select.into_select().inner.prepare(&mut cacher);
167        let (select, cached) = self
168            .query
169            .ast
170            .clone()
171            .full()
172            .simple_ordered(cacher.columns, self.order.clone());
173        let (sql, values) = select.build_rusqlite(SqliteQueryBuilder);
174
175        TXN.with_borrow_mut(|txn| {
176            let combi = txn.as_mut().unwrap();
177
178            combi.with_dependent_mut(|conn, rows_store| {
179                track_stmt(conn.get(), &sql, &values);
180                let statement = MutBorrow::new(conn.get().prepare_cached(&sql).unwrap());
181
182                let idx = rows_store.insert(OwnedRows::new(statement, |stmt| {
183                    stmt.borrow_mut().query(&*values.as_params()).unwrap()
184                }));
185
186                Iter {
187                    inner: idx,
188                    inner_phantom: PhantomData,
189                    prepared,
190                    cached,
191                }
192            })
193        })
194    }
195}
196
197pub(crate) fn track_stmt(conn: &Connection, sql: &String, values: &RusqliteValues) {
198    if COLLECT.get() {
199        SQL_AND_PLAN.with_borrow_mut(|map| {
200            map.entry(sql.clone())
201                .or_insert_with(|| get_node(conn, values, sql));
202        });
203    }
204}
205
206thread_local! {
207    static COLLECT: Cell<bool> = const { Cell::new(false) };
208    static SQL_AND_PLAN: RefCell<BTreeMap<String, Node>> = const { RefCell::new(BTreeMap::new()) };
209}
210
211pub fn get_plan<R>(f: impl FnOnce() -> R) -> (R, BTreeMap<String, Node>) {
212    let old = COLLECT.get();
213    COLLECT.set(true);
214    let res = f();
215    COLLECT.set(old);
216    (res, SQL_AND_PLAN.take())
217}
218
219fn get_node(conn: &Connection, values: &RusqliteValues, sql: &str) -> Node {
220    let mut prepared = conn.prepare(&format!("EXPLAIN QUERY PLAN {sql}")).unwrap();
221    let rows = prepared
222        .query_map(&*values.as_params(), |row| {
223            Ok((
224                row.get_unwrap("parent"),
225                Node {
226                    id: row.get_unwrap("id"),
227                    detail: row.get_unwrap("detail"),
228                    children: vec![],
229                },
230            ))
231        })
232        .unwrap();
233    let mut out = Node {
234        id: 0,
235        detail: "QUERY PLAN".to_owned(),
236        children: vec![],
237    };
238    rows.for_each(|res| {
239        let (id, node) = res.unwrap();
240        out.get_mut(id).children.push(node);
241    });
242
243    out
244}
245
246pub struct Node {
247    id: i64,
248    detail: String,
249    children: Vec<Node>,
250}
251
252impl Node {
253    fn get_mut(&mut self, id: i64) -> &mut Node {
254        if self.id == id {
255            return self;
256        }
257        self.children.last_mut().unwrap().get_mut(id)
258    }
259}
260
261impl Debug for Node {
262    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
263        f.write_str(&self.detail)?;
264        if !self.children.is_empty() {
265            f.write_str(" ")?;
266            f.debug_list().entries(&self.children).finish()?;
267        }
268        Ok(())
269    }
270}
rust_query/query.rs

rust_query/
query.rs
