sql_support/
each_chunk.rs

Help
1/* This Source Code Form is subject to the terms of the Mozilla Public
2 * License, v. 2.0. If a copy of the MPL was not distributed with this
3 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
4
5use lazy_static::lazy_static;
6use rusqlite::{self, limits::Limit, types::ToSql};
7use std::iter::Map;
8use std::slice::Iter;
9
10/// Returns SQLITE_LIMIT_VARIABLE_NUMBER as read from an in-memory connection and cached.
11/// connection and cached. That means this will return the wrong value if it's set to a lower
12/// value for a connection using this will return the wrong thing, but doing so is rare enough
13/// that we explicitly don't support it (why would you want to lower this at runtime?).
14///
15/// If you call this and the actual value was set to a negative number or zero (nothing prevents
16/// this beyond a warning in the SQLite documentation), we panic. However, it's unlikely you can
17/// run useful queries if this happened anyway.
18pub fn default_max_variable_number() -> usize {
19    lazy_static! {
20        static ref MAX_VARIABLE_NUMBER: usize = {
21            let conn = rusqlite::Connection::open_in_memory()
22                .expect("Failed to initialize in-memory connection (out of memory?)");
23
24            let limit = conn.limit(Limit::SQLITE_LIMIT_VARIABLE_NUMBER).unwrap();
25            assert!(
26                limit > 0,
27                "Illegal value for SQLITE_LIMIT_VARIABLE_NUMBER (must be > 0) {}",
28                limit
29            );
30            limit as usize
31        };
32    }
33    *MAX_VARIABLE_NUMBER
34}
35
36/// Helper for the case where you have a `&[impl ToSql]` of arbitrary length, but need one
37/// of no more than the connection's `MAX_VARIABLE_NUMBER` (rather,
38/// `default_max_variable_number()`). This is useful when performing batched updates.
39///
40/// The `do_chunk` callback is called with a slice of no more than `default_max_variable_number()`
41/// items as it's first argument, and the offset from the start as it's second.
42///
43/// See `each_chunk_mapped` for the case where `T` doesn't implement `ToSql`, but can be
44/// converted to something that does.
45pub fn each_chunk<'a, T, E, F>(items: &'a [T], do_chunk: F) -> Result<(), E>
46where
47    T: 'a,
48    F: FnMut(&'a [T], usize) -> Result<(), E>,
49{
50    each_sized_chunk(items, default_max_variable_number(), do_chunk)
51}
52
53/// A version of `each_chunk` for the case when the conversion to `to_sql` requires an custom
54/// intermediate step. For example, you might want to grab a property off of an array of records
55pub fn each_chunk_mapped<'a, T, U, E, Mapper, DoChunk>(
56    items: &'a [T],
57    to_sql: Mapper,
58    do_chunk: DoChunk,
59) -> Result<(), E>
60where
61    T: 'a,
62    U: ToSql + 'a,
63    Mapper: Fn(&'a T) -> U,
64    DoChunk: FnMut(Map<Iter<'a, T>, &'_ Mapper>, usize) -> Result<(), E>,
65{
66    each_sized_chunk_mapped(items, default_max_variable_number(), to_sql, do_chunk)
67}
68
69// Split out for testing. Separate so that we can pass an actual slice
70// to the callback if they don't need mapping. We could probably unify
71// this with each_sized_chunk_mapped with a lot of type system trickery,
72// but one of the benefits to each_chunk over the mapped versions is
73// that the declaration is simpler.
74pub fn each_sized_chunk<'a, T, E, F>(
75    items: &'a [T],
76    chunk_size: usize,
77    mut do_chunk: F,
78) -> Result<(), E>
79where
80    T: 'a,
81    F: FnMut(&'a [T], usize) -> Result<(), E>,
82{
83    if items.is_empty() {
84        return Ok(());
85    }
86    let mut offset = 0;
87    for chunk in items.chunks(chunk_size) {
88        do_chunk(chunk, offset)?;
89        offset += chunk.len();
90    }
91    Ok(())
92}
93
94/// Utility to help perform batched updates, inserts, queries, etc. This is the low-level version
95/// of this utility which is wrapped by `each_chunk` and `each_chunk_mapped`, and it allows you to
96/// provide both the mapping function, and the chunk size.
97///
98/// Note: `mapped` basically just refers to the translating of `T` to some `U` where `U: ToSql`
99/// using the `to_sql` function. This is useful for e.g. inserting the IDs of a large list
100/// of records.
101pub fn each_sized_chunk_mapped<'a, T, U, E, Mapper, DoChunk>(
102    items: &'a [T],
103    chunk_size: usize,
104    to_sql: Mapper,
105    mut do_chunk: DoChunk,
106) -> Result<(), E>
107where
108    T: 'a,
109    U: ToSql + 'a,
110    Mapper: Fn(&'a T) -> U,
111    DoChunk: FnMut(Map<Iter<'a, T>, &'_ Mapper>, usize) -> Result<(), E>,
112{
113    if items.is_empty() {
114        return Ok(());
115    }
116    let mut offset = 0;
117    for chunk in items.chunks(chunk_size) {
118        let mapped = chunk.iter().map(&to_sql);
119        do_chunk(mapped, offset)?;
120        offset += chunk.len();
121    }
122    Ok(())
123}
124
125#[cfg(test)]
126fn check_chunk<T, C>(items: C, expect: &[T], desc: &str)
127where
128    C: IntoIterator,
129    <C as IntoIterator>::Item: ToSql,
130    T: ToSql,
131{
132    let items = items.into_iter().collect::<Vec<_>>();
133    assert_eq!(items.len(), expect.len());
134    // Can't quite make the borrowing work out here w/o a loop, oh well.
135    for (idx, (got, want)) in items.iter().zip(expect.iter()).enumerate() {
136        assert_eq!(
137            got.to_sql().unwrap(),
138            want.to_sql().unwrap(),
139            // ToSqlOutput::Owned(Value::Integer(*num)),
140            "{}: Bad value at index {}",
141            desc,
142            idx
143        );
144    }
145}
146
147#[cfg(test)]
148mod test_mapped {
149    use super::*;
150
151    #[test]
152    fn test_separate() {
153        let mut iteration = 0;
154        each_sized_chunk_mapped(
155            &[1, 2, 3, 4, 5],
156            3,
157            |item| item as &dyn ToSql,
158            |chunk, offset| {
159                match offset {
160                    0 => {
161                        assert_eq!(iteration, 0);
162                        check_chunk(chunk, &[1, 2, 3], "first chunk");
163                    }
164                    3 => {
165                        assert_eq!(iteration, 1);
166                        check_chunk(chunk, &[4, 5], "second chunk");
167                    }
168                    n => {
169                        panic!("Unexpected offset {}", n);
170                    }
171                }
172                iteration += 1;
173                Ok::<(), ()>(())
174            },
175        )
176        .unwrap();
177    }
178
179    #[test]
180    fn test_leq_chunk_size() {
181        for &check_size in &[5, 6] {
182            let mut iteration = 0;
183            each_sized_chunk_mapped(
184                &[1, 2, 3, 4, 5],
185                check_size,
186                |item| item as &dyn ToSql,
187                |chunk, offset| {
188                    assert_eq!(iteration, 0);
189                    iteration += 1;
190                    assert_eq!(offset, 0);
191                    check_chunk(chunk, &[1, 2, 3, 4, 5], "only iteration");
192                    Ok::<(), ()>(())
193                },
194            )
195            .unwrap();
196        }
197    }
198
199    #[test]
200    fn test_empty_chunk() {
201        let items: &[i64] = &[];
202        each_sized_chunk_mapped::<_, _, (), _, _>(
203            items,
204            100,
205            |item| item as &dyn ToSql,
206            |_, _| {
207                panic!("Should never be called");
208            },
209        )
210        .unwrap();
211    }
212
213    #[test]
214    fn test_error() {
215        let mut iteration = 0;
216        let e = each_sized_chunk_mapped(
217            &[1, 2, 3, 4, 5, 6, 7],
218            3,
219            |item| item as &dyn ToSql,
220            |_, offset| {
221                if offset == 0 {
222                    assert_eq!(iteration, 0);
223                    iteration += 1;
224                    Ok(())
225                } else if offset == 3 {
226                    assert_eq!(iteration, 1);
227                    iteration += 1;
228                    Err("testing".to_string())
229                } else {
230                    // Make sure we stopped after the error.
231                    panic!("Shouldn't get called with offset of {}", offset);
232                }
233            },
234        )
235        .expect_err("Should be an error");
236        assert_eq!(e, "testing");
237    }
238}
239
240#[cfg(test)]
241mod test_unmapped {
242    use super::*;
243
244    #[test]
245    fn test_separate() {
246        let mut iteration = 0;
247        each_sized_chunk(&[1, 2, 3, 4, 5], 3, |chunk, offset| {
248            match offset {
249                0 => {
250                    assert_eq!(iteration, 0);
251                    check_chunk(chunk, &[1, 2, 3], "first chunk");
252                }
253                3 => {
254                    assert_eq!(iteration, 1);
255                    check_chunk(chunk, &[4, 5], "second chunk");
256                }
257                n => {
258                    panic!("Unexpected offset {}", n);
259                }
260            }
261            iteration += 1;
262            Ok::<(), ()>(())
263        })
264        .unwrap();
265    }
266
267    #[test]
268    fn test_leq_chunk_size() {
269        for &check_size in &[5, 6] {
270            let mut iteration = 0;
271            each_sized_chunk(&[1, 2, 3, 4, 5], check_size, |chunk, offset| {
272                assert_eq!(iteration, 0);
273                iteration += 1;
274                assert_eq!(offset, 0);
275                check_chunk(chunk, &[1, 2, 3, 4, 5], "only iteration");
276                Ok::<(), ()>(())
277            })
278            .unwrap();
279        }
280    }
281
282    #[test]
283    fn test_empty_chunk() {
284        let items: &[i64] = &[];
285        each_sized_chunk::<_, (), _>(items, 100, |_, _| {
286            panic!("Should never be called");
287        })
288        .unwrap();
289    }
290
291    #[test]
292    fn test_error() {
293        let mut iteration = 0;
294        let e = each_sized_chunk(&[1, 2, 3, 4, 5, 6, 7], 3, |_, offset| {
295            if offset == 0 {
296                assert_eq!(iteration, 0);
297                iteration += 1;
298                Ok(())
299            } else if offset == 3 {
300                assert_eq!(iteration, 1);
301                iteration += 1;
302                Err("testing".to_string())
303            } else {
304                // Make sure we stopped after the error.
305                panic!("Shouldn't get called with offset of {}", offset);
306            }
307        })
308        .expect_err("Should be an error");
309        assert_eq!(e, "testing");
310    }
311}
sql_support/each_chunk.rs

sql_support/
each_chunk.rs
