1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120

use crate::error::Error;
use crate::logger::QueryLogger;
use crate::sqlite::connection::{ConnectionHandle, ConnectionState};
use crate::sqlite::statement::{StatementHandle, VirtualStatement};
use crate::sqlite::{SqliteArguments, SqliteQueryResult, SqliteRow};
use either::Either;

pub struct ExecuteIter<'a> {
    handle: &'a mut ConnectionHandle,
    statement: &'a mut VirtualStatement,
    logger: QueryLogger<'a>,
    args: Option<SqliteArguments<'a>>,

    /// since a `VirtualStatement` can encompass multiple actual statements,
    /// this keeps track of the number of arguments so far
    args_used: usize,

    goto_next: bool,
}

pub(crate) fn iter<'a>(
    conn: &'a mut ConnectionState,
    query: &'a str,
    args: Option<SqliteArguments<'a>>,
    persistent: bool,
) -> Result<ExecuteIter<'a>, Error> {
    // fetch the cached statement or allocate a new one
    let statement = conn.statements.get(query, persistent)?;

    let logger = QueryLogger::new(query, conn.log_settings.clone());

    Ok(ExecuteIter {
        handle: &mut conn.handle,
        statement,
        logger,
        args,
        args_used: 0,
        goto_next: true,
    })
}

fn bind(
    statement: &mut StatementHandle,
    arguments: &Option<SqliteArguments<'_>>,
    offset: usize,
) -> Result<usize, Error> {
    let mut n = 0;

    if let Some(arguments) = arguments {
        n = arguments.bind(statement, offset)?;
    }

    Ok(n)
}

impl Iterator for ExecuteIter<'_> {
    type Item = Result<Either<SqliteQueryResult, SqliteRow>, Error>;

    fn next(&mut self) -> Option<Self::Item> {
        let statement = if self.goto_next {
            let mut statement = match self.statement.prepare_next(self.handle) {
                Ok(Some(statement)) => statement,
                Ok(None) => return None,
                Err(e) => return Some(Err(e.into())),
            };

            self.goto_next = false;

            // sanity check: ensure the VM is reset and the bindings are cleared
            if let Err(e) = statement.handle.reset() {
                return Some(Err(e.into()));
            }

            statement.handle.clear_bindings();

            match bind(&mut statement.handle, &self.args, self.args_used) {
                Ok(args_used) => self.args_used += args_used,
                Err(e) => return Some(Err(e)),
            }

            statement
        } else {
            self.statement.current()?
        };

        match statement.handle.step() {
            Ok(true) => {
                self.logger.increment_rows_returned();

                Some(Ok(Either::Right(SqliteRow::current(
                    &statement.handle,
                    &statement.columns,
                    &statement.column_names,
                ))))
            }
            Ok(false) => {
                let last_insert_rowid = self.handle.last_insert_rowid();

                let changes = statement.handle.changes();
                self.logger.increase_rows_affected(changes);

                let done = SqliteQueryResult {
                    changes,
                    last_insert_rowid,
                };

                self.goto_next = true;

                Some(Ok(Either::Left(done)))
            }
            Err(e) => Some(Err(e.into())),
        }
    }
}

impl Drop for ExecuteIter<'_> {
    fn drop(&mut self) {
        self.statement.reset().ok();
    }
}