diff options
Diffstat (limited to 'third_party/rust/sql-support/src/each_chunk.rs')
| -rw-r--r-- | third_party/rust/sql-support/src/each_chunk.rs | 311 |
1 files changed, 311 insertions, 0 deletions
diff --git a/third_party/rust/sql-support/src/each_chunk.rs b/third_party/rust/sql-support/src/each_chunk.rs new file mode 100644 index 0000000000..2d738bcb37 --- /dev/null +++ b/third_party/rust/sql-support/src/each_chunk.rs @@ -0,0 +1,311 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ + +use lazy_static::lazy_static; +use rusqlite::{self, limits::Limit, types::ToSql}; +use std::iter::Map; +use std::slice::Iter; + +/// Returns SQLITE_LIMIT_VARIABLE_NUMBER as read from an in-memory connection and cached. +/// connection and cached. That means this will return the wrong value if it's set to a lower +/// value for a connection using this will return the wrong thing, but doing so is rare enough +/// that we explicitly don't support it (why would you want to lower this at runtime?). +/// +/// If you call this and the actual value was set to a negative number or zero (nothing prevents +/// this beyond a warning in the SQLite documentation), we panic. However, it's unlikely you can +/// run useful queries if this happened anyway. +pub fn default_max_variable_number() -> usize { + lazy_static! { + static ref MAX_VARIABLE_NUMBER: usize = { + let conn = rusqlite::Connection::open_in_memory() + .expect("Failed to initialize in-memory connection (out of memory?)"); + + let limit = conn.limit(Limit::SQLITE_LIMIT_VARIABLE_NUMBER); + assert!( + limit > 0, + "Illegal value for SQLITE_LIMIT_VARIABLE_NUMBER (must be > 0) {}", + limit + ); + limit as usize + }; + } + *MAX_VARIABLE_NUMBER +} + +/// Helper for the case where you have a `&[impl ToSql]` of arbitrary length, but need one +/// of no more than the connection's `MAX_VARIABLE_NUMBER` (rather, +/// `default_max_variable_number()`). This is useful when performing batched updates. +/// +/// The `do_chunk` callback is called with a slice of no more than `default_max_variable_number()` +/// items as it's first argument, and the offset from the start as it's second. +/// +/// See `each_chunk_mapped` for the case where `T` doesn't implement `ToSql`, but can be +/// converted to something that does. +pub fn each_chunk<'a, T, E, F>(items: &'a [T], do_chunk: F) -> Result<(), E> +where + T: 'a, + F: FnMut(&'a [T], usize) -> Result<(), E>, +{ + each_sized_chunk(items, default_max_variable_number(), do_chunk) +} + +/// A version of `each_chunk` for the case when the conversion to `to_sql` requires an custom +/// intermediate step. For example, you might want to grab a property off of an arrray of records +pub fn each_chunk_mapped<'a, T, U, E, Mapper, DoChunk>( + items: &'a [T], + to_sql: Mapper, + do_chunk: DoChunk, +) -> Result<(), E> +where + T: 'a, + U: ToSql + 'a, + Mapper: Fn(&'a T) -> U, + DoChunk: FnMut(Map<Iter<'a, T>, &'_ Mapper>, usize) -> Result<(), E>, +{ + each_sized_chunk_mapped(items, default_max_variable_number(), to_sql, do_chunk) +} + +// Split out for testing. Separate so that we can pass an actual slice +// to the callback if they don't need mapping. We could probably unify +// this with each_sized_chunk_mapped with a lot of type system trickery, +// but one of the benefits to each_chunk over the mapped versions is +// that the declaration is simpler. +pub fn each_sized_chunk<'a, T, E, F>( + items: &'a [T], + chunk_size: usize, + mut do_chunk: F, +) -> Result<(), E> +where + T: 'a, + F: FnMut(&'a [T], usize) -> Result<(), E>, +{ + if items.is_empty() { + return Ok(()); + } + let mut offset = 0; + for chunk in items.chunks(chunk_size) { + do_chunk(chunk, offset)?; + offset += chunk.len(); + } + Ok(()) +} + +/// Utility to help perform batched updates, inserts, queries, etc. This is the low-level version +/// of this utility which is wrapped by `each_chunk` and `each_chunk_mapped`, and it allows you to +/// provide both the mapping function, and the chunk size. +/// +/// Note: `mapped` basically just refers to the translating of `T` to some `U` where `U: ToSql` +/// using the `to_sql` function. This is useful for e.g. inserting the IDs of a large list +/// of records. +pub fn each_sized_chunk_mapped<'a, T, U, E, Mapper, DoChunk>( + items: &'a [T], + chunk_size: usize, + to_sql: Mapper, + mut do_chunk: DoChunk, +) -> Result<(), E> +where + T: 'a, + U: ToSql + 'a, + Mapper: Fn(&'a T) -> U, + DoChunk: FnMut(Map<Iter<'a, T>, &'_ Mapper>, usize) -> Result<(), E>, +{ + if items.is_empty() { + return Ok(()); + } + let mut offset = 0; + for chunk in items.chunks(chunk_size) { + let mapped = chunk.iter().map(&to_sql); + do_chunk(mapped, offset)?; + offset += chunk.len(); + } + Ok(()) +} + +#[cfg(test)] +fn check_chunk<T, C>(items: C, expect: &[T], desc: &str) +where + C: IntoIterator, + <C as IntoIterator>::Item: ToSql, + T: ToSql, +{ + let items = items.into_iter().collect::<Vec<_>>(); + assert_eq!(items.len(), expect.len()); + // Can't quite make the borrowing work out here w/o a loop, oh well. + for (idx, (got, want)) in items.iter().zip(expect.iter()).enumerate() { + assert_eq!( + got.to_sql().unwrap(), + want.to_sql().unwrap(), + // ToSqlOutput::Owned(Value::Integer(*num)), + "{}: Bad value at index {}", + desc, + idx + ); + } +} + +#[cfg(test)] +mod test_mapped { + use super::*; + + #[test] + fn test_separate() { + let mut iteration = 0; + each_sized_chunk_mapped( + &[1, 2, 3, 4, 5], + 3, + |item| item as &dyn ToSql, + |chunk, offset| { + match offset { + 0 => { + assert_eq!(iteration, 0); + check_chunk(chunk, &[1, 2, 3], "first chunk"); + } + 3 => { + assert_eq!(iteration, 1); + check_chunk(chunk, &[4, 5], "second chunk"); + } + n => { + panic!("Unexpected offset {}", n); + } + } + iteration += 1; + Ok::<(), ()>(()) + }, + ) + .unwrap(); + } + + #[test] + fn test_leq_chunk_size() { + for &check_size in &[5, 6] { + let mut iteration = 0; + each_sized_chunk_mapped( + &[1, 2, 3, 4, 5], + check_size, + |item| item as &dyn ToSql, + |chunk, offset| { + assert_eq!(iteration, 0); + iteration += 1; + assert_eq!(offset, 0); + check_chunk(chunk, &[1, 2, 3, 4, 5], "only iteration"); + Ok::<(), ()>(()) + }, + ) + .unwrap(); + } + } + + #[test] + fn test_empty_chunk() { + let items: &[i64] = &[]; + each_sized_chunk_mapped::<_, _, (), _, _>( + items, + 100, + |item| item as &dyn ToSql, + |_, _| { + panic!("Should never be called"); + }, + ) + .unwrap(); + } + + #[test] + fn test_error() { + let mut iteration = 0; + let e = each_sized_chunk_mapped( + &[1, 2, 3, 4, 5, 6, 7], + 3, + |item| item as &dyn ToSql, + |_, offset| { + if offset == 0 { + assert_eq!(iteration, 0); + iteration += 1; + Ok(()) + } else if offset == 3 { + assert_eq!(iteration, 1); + iteration += 1; + Err("testing".to_string()) + } else { + // Make sure we stopped after the error. + panic!("Shouldn't get called with offset of {}", offset); + } + }, + ) + .expect_err("Should be an error"); + assert_eq!(e, "testing"); + } +} + +#[cfg(test)] +mod test_unmapped { + use super::*; + + #[test] + fn test_separate() { + let mut iteration = 0; + each_sized_chunk(&[1, 2, 3, 4, 5], 3, |chunk, offset| { + match offset { + 0 => { + assert_eq!(iteration, 0); + check_chunk(chunk, &[1, 2, 3], "first chunk"); + } + 3 => { + assert_eq!(iteration, 1); + check_chunk(chunk, &[4, 5], "second chunk"); + } + n => { + panic!("Unexpected offset {}", n); + } + } + iteration += 1; + Ok::<(), ()>(()) + }) + .unwrap(); + } + + #[test] + fn test_leq_chunk_size() { + for &check_size in &[5, 6] { + let mut iteration = 0; + each_sized_chunk(&[1, 2, 3, 4, 5], check_size, |chunk, offset| { + assert_eq!(iteration, 0); + iteration += 1; + assert_eq!(offset, 0); + check_chunk(chunk, &[1, 2, 3, 4, 5], "only iteration"); + Ok::<(), ()>(()) + }) + .unwrap(); + } + } + + #[test] + fn test_empty_chunk() { + let items: &[i64] = &[]; + each_sized_chunk::<_, (), _>(items, 100, |_, _| { + panic!("Should never be called"); + }) + .unwrap(); + } + + #[test] + fn test_error() { + let mut iteration = 0; + let e = each_sized_chunk(&[1, 2, 3, 4, 5, 6, 7], 3, |_, offset| { + if offset == 0 { + assert_eq!(iteration, 0); + iteration += 1; + Ok(()) + } else if offset == 3 { + assert_eq!(iteration, 1); + iteration += 1; + Err("testing".to_string()) + } else { + // Make sure we stopped after the error. + panic!("Shouldn't get called with offset of {}", offset); + } + }) + .expect_err("Should be an error"); + assert_eq!(e, "testing"); + } +} |