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<'inner, S> {
21    pub(crate) phantom: PhantomData<&'inner ()>,
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<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<'inner, S> Query<'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, S, Out = O>) -> Vec<O> {
80        self.into_iter(select).collect()
81    }
82
83    /// Turn a database query into an iterator 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>(&self, select: impl IntoSelect<'inner, S, Out = O>) -> Iter<'inner, O> {
89        let mut cacher = Cacher::new();
90        let prepared = select.into_select().inner.prepare(&mut cacher);
91
92        let (select, cached) = self.ast.clone().full().simple(cacher.columns);
93        let (sql, values) = select.build_rusqlite(SqliteQueryBuilder);
94        if SHOW_SQL.get() {
95            println!("{sql}");
96            println!("{values:?}");
97        }
98        if GET_PLAN.get() {
99            let node = get_node(self.conn, &values, &sql);
100            PLAN.set(Some(node));
101        }
102
103        let statement = MutBorrow::new(self.conn.prepare_cached(&sql).unwrap());
104
105        Iter {
106            inner: OwnedRows::new(statement, |stmt| {
107                stmt.borrow_mut().query(&*values.as_params()).unwrap()
108            }),
109            prepared,
110            cached,
111        }
112    }
113}
114
115thread_local! {
116    static SHOW_SQL: Cell<bool> = const { Cell::new(false) };
117    static GET_PLAN: Cell<bool> = const { Cell::new(false) };
118    static PLAN: Cell<Option<Node>> = const { Cell::new(None) };
119}
120
121pub fn show_sql<R>(f: impl FnOnce() -> R) -> R {
122    let old = SHOW_SQL.get();
123    SHOW_SQL.set(true);
124    let res = f();
125    SHOW_SQL.set(old);
126    res
127}
128
129pub fn get_plan<R>(f: impl FnOnce() -> R) -> (R, Node) {
130    let old = GET_PLAN.get();
131    GET_PLAN.set(true);
132    let res = f();
133    GET_PLAN.set(old);
134    (res, PLAN.take().unwrap())
135}
136
137fn get_node(conn: &Connection, values: &RusqliteValues, sql: &str) -> Node {
138    let mut prepared = conn.prepare(&format!("EXPLAIN QUERY PLAN {sql}")).unwrap();
139    let rows = prepared
140        .query_map(&*values.as_params(), |row| {
141            Ok((
142                row.get_unwrap("parent"),
143                Node {
144                    id: row.get_unwrap("id"),
145                    detail: row.get_unwrap("detail"),
146                    children: vec![],
147                },
148            ))
149        })
150        .unwrap();
151    let mut out = Node {
152        id: 0,
153        detail: "QUERY PLAN".to_owned(),
154        children: vec![],
155    };
156    rows.for_each(|res| {
157        let (id, node) = res.unwrap();
158        out.get_mut(id).children.push(node);
159    });
160
161    out
162}
163
164pub struct Node {
165    id: i64,
166    detail: String,
167    children: Vec<Node>,
168}
169
170impl Node {
171    fn get_mut(&mut self, id: i64) -> &mut Node {
172        if self.id == id {
173            return self;
174        }
175        self.children.last_mut().unwrap().get_mut(id)
176    }
177}
178
179impl Debug for Node {
180    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
181        f.write_str(&self.detail)?;
182        if !self.children.is_empty() {
183            f.write_str(" ")?;
184            f.debug_list().entries(&self.children).finish()?;
185        }
186        Ok(())
187    }
188}