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    alias::MyAlias,
16    dummy_impl::{Cacher, DynPrepared, IntoSelect, Prepared, Row, SelectImpl},
17    rows::Rows,
18};
19
20/// This is the type used by the [crate::Transaction::query] method.
21pub struct Query<'inner, S> {
22    pub(crate) phantom: PhantomData<&'inner ()>,
23    pub(crate) q: Rows<'inner, S>,
24    pub(crate) conn: &'inner rusqlite::Connection,
25}
26
27impl<'inner, S> Deref for Query<'inner, S> {
28    type Target = Rows<'inner, S>;
29
30    fn deref(&self) -> &Self::Target {
31        &self.q
32    }
33}
34
35impl<S> DerefMut for Query<'_, S> {
36    fn deref_mut(&mut self) -> &mut Self::Target {
37        &mut self.q
38    }
39}
40
41type Stmt<'x> = rusqlite::CachedStatement<'x>;
42type RRows<'a> = rusqlite::Rows<'a>;
43
44self_cell!(
45    struct OwnedRows<'x> {
46        owner: MutBorrow<Stmt<'x>>,
47
48        #[covariant]
49        dependent: RRows,
50    }
51);
52
53/// Lazy iterator over rows from a query.
54///
55/// This is currently invariant in `'inner` due to [MutBorrow].
56/// Would be nice to relax this variance in the future.
57pub struct Iter<'inner, O> {
58    inner: OwnedRows<'inner>,
59    prepared: DynPrepared<O>,
60    cached: Vec<MyAlias>,
61}
62
63impl<O> Iterator for Iter<'_, O> {
64    type Item = O;
65
66    fn next(&mut self) -> Option<Self::Item> {
67        self.inner.with_dependent_mut(|_, rows| {
68            let row = rows.next().unwrap()?;
69            Some(self.prepared.call(Row::new(row, &self.cached)))
70        })
71    }
72}
73
74impl<'inner, S> Query<'inner, S> {
75    /// Turn a database query into a [Vec] of results.
76    ///
77    /// The order of rows that is returned is unstable. This means that the order may change between any two
78    /// executions of the exact same query. If a specific order (or even a consistent order) is required,
79    /// then you have to use something like [slice::sort].
80    pub fn into_vec<O>(&self, select: impl IntoSelect<'inner, S, Out = O>) -> Vec<O> {
81        self.into_iter(select).collect()
82    }
83
84    /// Turn a database query into an iterator of results.
85    ///
86    /// The order of rows that is returned is unstable. This means that the order may change between any two
87    /// executions of the exact same query. If a specific order (or even a consistent order) is required,
88    /// then you have to use something like [slice::sort].
89    pub fn into_iter<O>(&self, select: impl IntoSelect<'inner, S, Out = O>) -> Iter<'inner, O> {
90        let mut cacher = Cacher::new();
91        let prepared = select.into_select().inner.prepare(&mut cacher);
92
93        let (select, cached) = self.ast.clone().full().simple(cacher.columns);
94        let (sql, values) = select.build_rusqlite(SqliteQueryBuilder);
95        track_stmt(self.conn, &sql, &values);
96
97        let statement = MutBorrow::new(self.conn.prepare_cached(&sql).unwrap());
98
99        Iter {
100            inner: OwnedRows::new(statement, |stmt| {
101                stmt.borrow_mut().query(&*values.as_params()).unwrap()
102            }),
103            prepared,
104            cached,
105        }
106    }
107}
108
109pub(crate) fn track_stmt(conn: &Connection, sql: &String, values: &RusqliteValues) {
110    if COLLECT.get() {
111        SQL_AND_PLAN.with_borrow_mut(|map| {
112            map.entry(sql.clone())
113                .or_insert_with(|| get_node(conn, values, sql));
114        });
115    }
116}
117
118thread_local! {
119    static COLLECT: Cell<bool> = const { Cell::new(false) };
120    static SQL_AND_PLAN: RefCell<BTreeMap<String, Node>> = const { RefCell::new(BTreeMap::new()) };
121}
122
123pub fn get_plan<R>(f: impl FnOnce() -> R) -> (R, BTreeMap<String, Node>) {
124    let old = COLLECT.get();
125    COLLECT.set(true);
126    let res = f();
127    COLLECT.set(old);
128    (res, SQL_AND_PLAN.take())
129}
130
131fn get_node(conn: &Connection, values: &RusqliteValues, sql: &str) -> Node {
132    let mut prepared = conn.prepare(&format!("EXPLAIN QUERY PLAN {sql}")).unwrap();
133    let rows = prepared
134        .query_map(&*values.as_params(), |row| {
135            Ok((
136                row.get_unwrap("parent"),
137                Node {
138                    id: row.get_unwrap("id"),
139                    detail: row.get_unwrap("detail"),
140                    children: vec![],
141                },
142            ))
143        })
144        .unwrap();
145    let mut out = Node {
146        id: 0,
147        detail: "QUERY PLAN".to_owned(),
148        children: vec![],
149    };
150    rows.for_each(|res| {
151        let (id, node) = res.unwrap();
152        out.get_mut(id).children.push(node);
153    });
154
155    out
156}
157
158pub struct Node {
159    id: i64,
160    detail: String,
161    children: Vec<Node>,
162}
163
164impl Node {
165    fn get_mut(&mut self, id: i64) -> &mut Node {
166        if self.id == id {
167            return self;
168        }
169        self.children.last_mut().unwrap().get_mut(id)
170    }
171}
172
173impl Debug for Node {
174    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
175        f.write_str(&self.detail)?;
176        if !self.children.is_empty() {
177            f.write_str(" ")?;
178            f.debug_list().entries(&self.children).finish()?;
179        }
180        Ok(())
181    }
182}
rust_query/query.rs

rust_query/
query.rs
