diesel/sqlite/connection/
statement_iterator.rs

1use alloc::rc::Rc;
2use alloc::string::String;
3use core::cell::RefCell;
4use core::iter::FusedIterator;
5
6use super::row::{PrivateSqliteRow, SqliteRow};
7use super::stmt::StatementUse;
8use crate::result::QueryResult;
9
10#[allow(missing_debug_implementations)]
11pub struct StatementIterator<'stmt, 'query> {
12    inner: PrivateStatementIterator<'stmt, 'query>,
13    column_names: Option<Rc<[Option<String>]>>,
14    field_count: usize,
15}
16
17impl<'stmt, 'query> StatementIterator<'stmt, 'query> {
18    #[cold]
19    #[allow(unsafe_code)] // call to unsafe function
20    fn handle_duplicated_row_case(
21        outer_last_row: &mut Rc<RefCell<PrivateSqliteRow<'stmt, 'query>>>,
22        column_names: &mut Option<Rc<[Option<String>]>>,
23        field_count: usize,
24    ) -> Option<(QueryResult<SqliteRow<'stmt, 'query>>, bool)> {
25        // We don't own the statement. There is another existing reference, likely because
26        // a user stored the row in some long time container before calling next another time
27        // In this case we copy out the current values into a temporary store and advance
28        // the statement iterator internally afterwards
29        let last_row = {
30            let mut last_row = match outer_last_row.try_borrow_mut() {
31                Ok(o) => o,
32                Err(_e) => {
33                    return Some((Err(crate::result::Error::DeserializationError(
34                                    "Failed to reborrow row. Try to release any `SqliteField` or `SqliteValue` \
35                                     that exists at this point"
36                                        .into(),
37                                )), false));
38                }
39            };
40            let last_row = &mut *last_row;
41            let duplicated = match last_row.duplicate(column_names) {
42                Ok(duplicated) => duplicated,
43                // The statement is not advanced, so the row it holds stays current.
44                Err(e) => return Some((Err(e), false)),
45            };
46            core::mem::replace(last_row, duplicated)
47        };
48        if let PrivateSqliteRow::Direct(mut stmt) = last_row {
49            let res = unsafe {
50                // This is actually safe here as we've already
51                // performed one step. For the first step we would have
52                // used `PrivateStatementIterator::NotStarted` where we don't
53                // have access to `PrivateSqliteRow` at all
54                stmt.step(false)
55            };
56            *outer_last_row = Rc::new(RefCell::new(PrivateSqliteRow::Direct(stmt)));
57            match res {
58                Err(e) => Some((Err(e), true)),
59                Ok(false) => None,
60                Ok(true) => Some((
61                    Ok(SqliteRow {
62                        inner: Rc::clone(outer_last_row),
63                        field_count,
64                    }),
65                    false,
66                )),
67            }
68        } else {
69            // any other state than `PrivateSqliteRow::Direct` is invalid here
70            // and should not happen. If this ever happens this is a logic error
71            // in the code above
72            {
    ::core::panicking::panic_fmt(format_args!("internal error: entered unreachable code: {0}",
            format_args!("You\'ve reached an impossible internal state. If you ever see this error message please open an issue at https://github.com/diesel-rs/diesel providing example code how to trigger this error.")));
}unreachable!(
73                "You've reached an impossible internal state. \
74                             If you ever see this error message please open \
75                             an issue at https://github.com/diesel-rs/diesel \
76                             providing example code how to trigger this error."
77            )
78        }
79    }
80}
81
82enum PrivateStatementIterator<'stmt, 'query> {
83    NotStarted(Option<StatementUse<'stmt, 'query>>),
84    Started(Rc<RefCell<PrivateSqliteRow<'stmt, 'query>>>),
85    Finished,
86}
87
88impl<'stmt, 'query> StatementIterator<'stmt, 'query> {
89    pub fn new(stmt: StatementUse<'stmt, 'query>) -> StatementIterator<'stmt, 'query> {
90        Self {
91            inner: PrivateStatementIterator::NotStarted(Some(stmt)),
92            column_names: None,
93            field_count: 0,
94        }
95    }
96}
97
98impl<'stmt, 'query> Iterator for StatementIterator<'stmt, 'query> {
99    type Item = QueryResult<SqliteRow<'stmt, 'query>>;
100
101    #[allow(unsafe_code)] // call to unsafe function
102    fn next(&mut self) -> Option<Self::Item> {
103        use PrivateStatementIterator::{Finished, NotStarted, Started};
104        match &mut self.inner {
105            NotStarted(stmt @ Some(_)) => {
106                let mut stmt = stmt
107                    .take()
108                    .expect("It must be there because we checked that above");
109                let step = unsafe {
110                    // This is safe as we pass `first_step = true` to reset the cached column names
111                    stmt.step(true)
112                };
113                match step {
114                    Err(e) => {
115                        self.inner = Finished;
116                        Some(Err(e))
117                    }
118                    Ok(false) => {
119                        // we reached the end of the iterator
120                        // so we mark it done on rust side
121                        self.inner = Finished;
122                        None
123                    }
124                    Ok(true) => {
125                        let field_count = stmt
126                            .column_count()
127                            .try_into()
128                            .expect("Diesel expects to run at least on a 32 bit platform");
129                        self.field_count = field_count;
130                        let inner = Rc::new(RefCell::new(PrivateSqliteRow::Direct(stmt)));
131                        self.inner = Started(inner.clone());
132                        Some(Ok(SqliteRow { inner, field_count }))
133                    }
134                }
135            }
136            Started(last_row) => {
137                // There was already at least one iteration step
138                // We check here if the caller already released the row value or not
139                // by checking if our Rc owns the data or not
140                if let Some(last_row_ref) = Rc::get_mut(last_row) {
141                    // We own the statement, there is no other reference here.
142                    // This means we don't need to copy out values from the sqlite provided
143                    // datastructures for now
144                    // We don't need to use the runtime borrowing system of the RefCell here
145                    // as we have a mutable reference, so all of this below is checked at compile time
146                    if let PrivateSqliteRow::Direct(stmt) = last_row_ref.get_mut() {
147                        let step = unsafe {
148                            // This is actually safe here as we've already
149                            // performed one step. For the first step we would have
150                            // used `PrivateStatementIterator::NotStarted` where we don't
151                            // have access to `PrivateSqliteRow` at all
152
153                            stmt.step(false)
154                        };
155                        match step {
156                            Err(e) => {
157                                self.inner = Finished;
158                                Some(Err(e))
159                            }
160                            Ok(false) => {
161                                // we reached the end of the iterator
162                                // so we mark it done on rust side
163                                self.inner = Finished;
164                                None
165                            }
166                            Ok(true) => {
167                                let field_count = self.field_count;
168                                Some(Ok(SqliteRow {
169                                    inner: Rc::clone(last_row),
170                                    field_count,
171                                }))
172                            }
173                        }
174                    } else {
175                        // any other state than `PrivateSqliteRow::Direct` is invalid here
176                        // and should not happen. If this ever happens this is a logic error
177                        // in the code above
178                        {
    ::core::panicking::panic_fmt(format_args!("internal error: entered unreachable code: {0}",
            format_args!("You\'ve reached an impossible internal state. If you ever see this error message please open an issue at https://github.com/diesel-rs/diesel providing example code how to trigger this error.")));
}unreachable!(
179                            "You've reached an impossible internal state. \
180                             If you ever see this error message please open \
181                             an issue at https://github.com/diesel-rs/diesel \
182                             providing example code how to trigger this error."
183                        )
184                    }
185                } else {
186                    let result = Self::handle_duplicated_row_case(
187                        last_row,
188                        &mut self.column_names,
189                        self.field_count,
190                    );
191                    if result.as_ref().is_none_or(|(r, f)| *f && r.is_err()) {
192                        self.inner = Finished;
193                    }
194                    result.map(|(r, _)| r)
195                }
196            }
197            NotStarted(_s) => {
198                // we likely got an error while executing the other
199                // `NotStarted` branch above. In this case we just want to stop
200                // iterating here
201                None
202            }
203            Finished => None,
204        }
205    }
206}
207
208impl<'stmt, 'query> FusedIterator for StatementIterator<'stmt, 'query> {}
209
210#[cfg(test)]
211mod tests {
212    use crate::connection::{DefaultLoadingMode, SimpleConnection};
213    use crate::dsl::sql;
214    use crate::prelude::*;
215    use crate::sql_types::Integer;
216
217    #[diesel_test_helper::test]
218    fn statement_iterator_is_fused() {
219        let mut conn = SqliteConnection::establish(":memory:").unwrap();
220        conn.batch_execute("CREATE TABLE t(x INTEGER); INSERT INTO t VALUES (1), (2);")
221            .unwrap();
222        {
223            let mut it = sql::<Integer>("SELECT x FROM t")
224                .load_iter::<i32, DefaultLoadingMode>(&mut conn)
225                .unwrap();
226            assert_eq!(it.next().unwrap().unwrap(), 1);
227            assert_eq!(it.next().unwrap().unwrap(), 2);
228            assert!(it.next().is_none());
229            // Polling past the end must not restart (re-execute) the statement.
230            assert!(it.next().is_none(), "iterator re-executed the statement");
231        }
232        conn.batch_execute("DELETE FROM t").unwrap();
233        let mut it = sql::<Integer>("SELECT x FROM t")
234            .load_iter::<i32, DefaultLoadingMode>(&mut conn)
235            .unwrap();
236        assert!(it.next().is_none());
237        // Polling past the end must not restart (re-execute) the statement.
238        assert!(it.next().is_none(), "iterator re-executed the statement");
239    }
240
241    #[diesel_test_helper::test]
242    #[cfg(not(miri))] // returns an error which allocates and uses an ffi string
243    fn sqlite_iterator_terminates_after_step_error() {
244        let mut conn = SqliteConnection::establish(":memory:").unwrap();
245        let q = diesel::dsl::sql::<Integer>(
246            "WITH RECURSIVE c(x) AS (SELECT 1 UNION ALL SELECT x + 1 FROM c WHERE x < 5) \
247         SELECT CASE WHEN x = 3 THEN json('{') ELSE x END FROM c",
248        );
249        let items: Vec<Result<i32, String>> = q
250            .load_iter::<i32, DefaultLoadingMode>(&mut conn)
251            .unwrap()
252            .take(100)
253            .map(|r| r.map_err(|e| e.to_string()))
254            .collect();
255        assert_eq!(
256            items.len(),
257            3,
258            "iterator yielded {} items - restarts after error",
259            items.len()
260        );
261        assert_eq!(items[0], Ok(1));
262        assert_eq!(items[1], Ok(2));
263        assert!(items[2].is_err());
264        assert_eq!(items[2].as_ref().unwrap_err().to_string(), "malformed JSON");
265    }
266
267    #[diesel_test_helper::test]
268    #[cfg(not(miri))] // returns an error which allocates and uses an ffi string
269    fn held_rows_iterator_terminates_after_step_error() {
270        use crate::connection::LoadConnection;
271        use crate::row::Row;
272
273        let mut conn = SqliteConnection::establish(":memory:").unwrap();
274        let q = sql::<Integer>(
275            "WITH RECURSIVE c(x) AS (SELECT 1 UNION ALL SELECT x + 1 FROM c WHERE x < 5) \
276               SELECT CASE WHEN x = 3 THEN json('{') ELSE x END FROM c",
277        );
278        let mut iter = LoadConnection::<DefaultLoadingMode>::load(&mut conn, q).unwrap();
279
280        let first = iter.next().unwrap().unwrap();
281        let second = iter.next().unwrap().unwrap();
282        let err = iter.next().unwrap().err().expect("step error");
283        assert_eq!(err.to_string(), "malformed JSON");
284        assert!(iter.next().is_none(), "iterator restarted the statement");
285        assert!(iter.next().is_none());
286
287        assert_eq!(first.get_value::<Integer, i32, _>(0).unwrap(), 1);
288        assert_eq!(second.get_value::<Integer, i32, _>(0).unwrap(), 2);
289    }
290
291    #[diesel_test_helper::test]
292    #[cfg(not(miri))] // returns an error which allocates and uses an ffi string
293    fn held_rows_iterator_continues_after_reborrow_error() {
294        use crate::connection::LoadConnection;
295        use crate::row::Row;
296
297        let mut conn = SqliteConnection::establish(":memory:").unwrap();
298        conn.batch_execute("CREATE TABLE t(x INTEGER); INSERT INTO t VALUES (1), (2), (3);")
299            .unwrap();
300        let q = sql::<Integer>("SELECT x FROM t");
301        let mut iter = LoadConnection::<DefaultLoadingMode>::load(&mut conn, q).unwrap();
302
303        let first = iter.next().unwrap().unwrap();
304        let field = first.get(0).unwrap();
305        assert!(iter.next().unwrap().is_err());
306        drop(field);
307
308        let rest = iter
309            .map(|row| row.unwrap().get_value::<Integer, i32, _>(0).unwrap())
310            .collect::<Vec<_>>();
311        assert_eq!(rest, [2, 3]);
312    }
313}
diesel/sqlite/connection/statement_iterator.rs

diesel/sqlite/connection/
statement_iterator.rs
