rust_query/
query.rs

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