scylla_cql/serialize/
batch.rs

1//! Contains the [`BatchValues`] and [`BatchValuesIterator`] trait and their
2//! implementations.
3
4// Note: When editing above doc-comment edit the corresponding comment on
5// re-export module in scylla crate too.
6
7use super::row::{RowSerializationContext, SerializeRow};
8use super::{RowWriter, SerializationError};
9
10/// Represents a list of sets of values for a batch statement.
11///
12/// The data in the object can be consumed with an iterator-like object returned
13/// by the [`BatchValues::batch_values_iter`] method.
14pub trait BatchValues {
15    /// An `Iterator`-like object over the values from the parent `BatchValues` object.
16    // For some unknown reason, this type, when not resolved to a concrete type for a given async function,
17    // cannot live across await boundaries while maintaining the corresponding future `Send`, unless `'r: 'static`
18    //
19    // See <https://github.com/scylladb/scylla-rust-driver/issues/599> for more details
20    type BatchValuesIter<'r>: BatchValuesIterator<'r>
21    where
22        Self: 'r;
23
24    /// Returns an iterator over the data contained in this object.
25    fn batch_values_iter(&self) -> Self::BatchValuesIter<'_>;
26}
27
28/// An `Iterator`-like object over the values from the parent [`BatchValues`] object.
29///
30/// It's not a true [`Iterator`] because it does not provide direct access to the
31/// items being iterated over, instead it allows calling methods of the underlying
32/// [`SerializeRow`] trait while advancing the iterator.
33pub trait BatchValuesIterator<'bv> {
34    /// Serializes the next set of values in the sequence and advances the iterator.
35    fn serialize_next(
36        &mut self,
37        ctx: &RowSerializationContext<'_>,
38        writer: &mut RowWriter,
39    ) -> Option<Result<(), SerializationError>>;
40
41    /// Returns whether the next set of values is empty or not and advances the iterator.
42    fn is_empty_next(&mut self) -> Option<bool>;
43
44    /// Skips the next set of values.
45    fn skip_next(&mut self) -> Option<()>;
46
47    /// Return the number of sets of values, consuming the iterator in the process.
48    #[inline]
49    fn count(mut self) -> usize
50    where
51        Self: Sized,
52    {
53        let mut count = 0;
54        while self.skip_next().is_some() {
55            count += 1;
56        }
57        count
58    }
59}
60
61/// Implements `BatchValuesIterator` from an `Iterator` over references to things that implement `SerializeRow`
62///
63/// Essentially used internally by this lib to provide implementers of `BatchValuesIterator` for cases
64/// that always serialize the same concrete `SerializeRow` type
65pub struct BatchValuesIteratorFromIterator<IT: Iterator> {
66    it: IT,
67}
68
69impl<'bv, 'sr: 'bv, IT, SR> BatchValuesIterator<'bv> for BatchValuesIteratorFromIterator<IT>
70where
71    IT: Iterator<Item = &'sr SR>,
72    SR: SerializeRow + 'sr,
73{
74    #[inline]
75    fn serialize_next(
76        &mut self,
77        ctx: &RowSerializationContext<'_>,
78        writer: &mut RowWriter,
79    ) -> Option<Result<(), SerializationError>> {
80        self.it.next().map(|sr| sr.serialize(ctx, writer))
81    }
82
83    #[inline]
84    fn is_empty_next(&mut self) -> Option<bool> {
85        self.it.next().map(|sr| sr.is_empty())
86    }
87
88    #[inline]
89    fn skip_next(&mut self) -> Option<()> {
90        self.it.next().map(|_| ())
91    }
92
93    #[inline]
94    fn count(self) -> usize
95    where
96        Self: Sized,
97    {
98        self.it.count()
99    }
100}
101
102impl<IT> From<IT> for BatchValuesIteratorFromIterator<IT>
103where
104    IT: Iterator,
105    IT::Item: SerializeRow,
106{
107    #[inline]
108    fn from(it: IT) -> Self {
109        BatchValuesIteratorFromIterator { it }
110    }
111}
112
113//
114// BatchValues impls
115//
116
117/// Implements `BatchValues` from an `Iterator` over references to things that implement `SerializeRow`
118///
119/// This is to avoid requiring allocating a new `Vec` containing all the `SerializeRow`s directly:
120/// with this, one can write:
121/// `session.batch(&batch, BatchValuesFromIter::from(lines_to_insert.iter().map(|l| &l.value_list)))`
122/// where `lines_to_insert` may also contain e.g. data to pick the statement...
123///
124/// The underlying iterator will always be cloned at least once, once to compute the length if it can't be known
125/// in advance, and be re-cloned at every retry.
126/// It is consequently expected that the provided iterator is cheap to clone (e.g. `slice.iter().map(...)`).
127pub struct BatchValuesFromIterator<'sr, IT> {
128    it: IT,
129
130    // Without artificially introducing a lifetime to the struct, I couldn't get
131    // impl BatchValues for BatchValuesFromIterator to work. I wish I understood
132    // why it's needed.
133    _phantom: std::marker::PhantomData<&'sr ()>,
134}
135
136impl<'sr, IT, SR> BatchValuesFromIterator<'sr, IT>
137where
138    IT: Iterator<Item = &'sr SR> + Clone,
139    SR: SerializeRow + 'sr,
140{
141    /// Creates a new `BatchValuesFromIter`` object.
142    #[inline]
143    pub fn new(into_iter: impl IntoIterator<IntoIter = IT>) -> Self {
144        Self {
145            it: into_iter.into_iter(),
146            _phantom: std::marker::PhantomData,
147        }
148    }
149}
150
151impl<'sr, IT, SR> From<IT> for BatchValuesFromIterator<'sr, IT>
152where
153    IT: Iterator<Item = &'sr SR> + Clone,
154    SR: SerializeRow + 'sr,
155{
156    #[inline]
157    fn from(it: IT) -> Self {
158        Self::new(it)
159    }
160}
161
162impl<'sr, IT, SR> BatchValues for BatchValuesFromIterator<'sr, IT>
163where
164    IT: Iterator<Item = &'sr SR> + Clone,
165    SR: SerializeRow + 'sr,
166{
167    type BatchValuesIter<'r>
168        = BatchValuesIteratorFromIterator<IT>
169    where
170        Self: 'r;
171
172    #[inline]
173    fn batch_values_iter(&self) -> Self::BatchValuesIter<'_> {
174        self.it.clone().into()
175    }
176}
177
178// Implement BatchValues for slices of SerializeRow types
179impl<T: SerializeRow> BatchValues for [T] {
180    type BatchValuesIter<'r>
181        = BatchValuesIteratorFromIterator<std::slice::Iter<'r, T>>
182    where
183        Self: 'r;
184
185    #[inline]
186    fn batch_values_iter(&self) -> Self::BatchValuesIter<'_> {
187        self.iter().into()
188    }
189}
190
191// Implement BatchValues for Vec<SerializeRow>
192impl<T: SerializeRow> BatchValues for Vec<T> {
193    type BatchValuesIter<'r>
194        = BatchValuesIteratorFromIterator<std::slice::Iter<'r, T>>
195    where
196        Self: 'r;
197
198    #[inline]
199    fn batch_values_iter(&self) -> Self::BatchValuesIter<'_> {
200        BatchValues::batch_values_iter(self.as_slice())
201    }
202}
203
204// Here is an example implementation for (T0, )
205// Further variants are done using a macro
206impl<T0: SerializeRow> BatchValues for (T0,) {
207    type BatchValuesIter<'r>
208        = BatchValuesIteratorFromIterator<std::iter::Once<&'r T0>>
209    where
210        Self: 'r;
211
212    #[inline]
213    fn batch_values_iter(&self) -> Self::BatchValuesIter<'_> {
214        std::iter::once(&self.0).into()
215    }
216}
217
218/// A [`BatchValuesIterator`] over a tuple.
219pub struct TupleValuesIter<'sr, T> {
220    tuple: &'sr T,
221    idx: usize,
222}
223
224macro_rules! impl_batch_values_for_tuple {
225    ( $($Ti:ident),* ; $($FieldI:tt),* ; $TupleSize:tt) => {
226        impl<$($Ti),+> BatchValues for ($($Ti,)+)
227        where
228            $($Ti: SerializeRow),+
229        {
230            type BatchValuesIter<'r> = TupleValuesIter<'r, ($($Ti,)+)> where Self: 'r;
231
232            #[inline]
233            fn batch_values_iter(&self) -> Self::BatchValuesIter<'_> {
234                TupleValuesIter {
235                    tuple: self,
236                    idx: 0,
237                }
238            }
239        }
240
241        impl<'bv, $($Ti),+> BatchValuesIterator<'bv> for TupleValuesIter<'bv, ($($Ti,)+)>
242        where
243            $($Ti: SerializeRow),+
244        {
245            #[inline]
246            fn serialize_next(
247                &mut self,
248                ctx: &RowSerializationContext<'_>,
249                writer: &mut RowWriter,
250            ) -> Option<Result<(), SerializationError>> {
251                let ret = match self.idx {
252                    $(
253                        $FieldI => self.tuple.$FieldI.serialize(ctx, writer),
254                    )*
255                    _ => return None,
256                };
257                self.idx += 1;
258                Some(ret)
259            }
260
261            #[inline]
262            fn is_empty_next(&mut self) -> Option<bool> {
263                let ret = match self.idx {
264                    $(
265                        $FieldI => self.tuple.$FieldI.is_empty(),
266                    )*
267                    _ => return None,
268                };
269                self.idx += 1;
270                Some(ret)
271            }
272
273            #[inline]
274            fn skip_next(&mut self) -> Option<()> {
275                if self.idx < $TupleSize {
276                    self.idx += 1;
277                    Some(())
278                } else {
279                    None
280                }
281            }
282
283            #[inline]
284            fn count(self) -> usize {
285                $TupleSize - self.idx
286            }
287        }
288    }
289}
290
291impl_batch_values_for_tuple!(T0, T1; 0, 1; 2);
292impl_batch_values_for_tuple!(T0, T1, T2; 0, 1, 2; 3);
293impl_batch_values_for_tuple!(T0, T1, T2, T3; 0, 1, 2, 3; 4);
294impl_batch_values_for_tuple!(T0, T1, T2, T3, T4; 0, 1, 2, 3, 4; 5);
295impl_batch_values_for_tuple!(T0, T1, T2, T3, T4, T5; 0, 1, 2, 3, 4, 5; 6);
296impl_batch_values_for_tuple!(T0, T1, T2, T3, T4, T5, T6; 0, 1, 2, 3, 4, 5, 6; 7);
297impl_batch_values_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7; 0, 1, 2, 3, 4, 5, 6, 7; 8);
298impl_batch_values_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8; 0, 1, 2, 3, 4, 5, 6, 7, 8; 9);
299impl_batch_values_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9;
300                             0, 1, 2, 3, 4, 5, 6, 7, 8, 9; 10);
301impl_batch_values_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10;
302                             0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10; 11);
303impl_batch_values_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11;
304                             0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11; 12);
305impl_batch_values_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12;
306                             0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12; 13);
307impl_batch_values_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13;
308                             0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13; 14);
309impl_batch_values_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14;
310                             0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14; 15);
311impl_batch_values_for_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15;
312                             0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15; 16);
313
314// Every &impl BatchValues should also implement BatchValues
315impl<T: BatchValues + ?Sized> BatchValues for &T {
316    type BatchValuesIter<'r>
317        = <T as BatchValues>::BatchValuesIter<'r>
318    where
319        Self: 'r;
320
321    #[inline]
322    fn batch_values_iter(&self) -> Self::BatchValuesIter<'_> {
323        <T as BatchValues>::batch_values_iter(*self)
324    }
325}
326
327#[cfg(test)]
328#[path = "batch_tests.rs"]
329pub(crate) mod tests;
scylla_cql/serialize/batch.rs

scylla_cql/serialize/
batch.rs
