slatedb/
batch.rs

1//! # Batch
2//!
3//! This module contains types for batch operations. A batch operation is a
4//! collection of write operations (puts and/or deletes) that are applied
5//! atomically to the database.
6
7use crate::config::{MergeOptions, PutOptions};
8use crate::db_common::extract_segment_prefix;
9use crate::error::SlateDBError;
10use crate::iter::{IterationOrder, RowEntryIterator};
11use crate::merge_operator::{MergeOperatorIterator, MergeOperatorType};
12use crate::prefix_extractor::PrefixExtractor;
13use crate::types::{RowEntry, ValueDeletable};
14use async_trait::async_trait;
15use bytes::Bytes;
16use smallvec::{smallvec, SmallVec};
17use std::collections::{BTreeMap, BTreeSet, HashSet};
18use std::iter::Peekable;
19use std::ops::RangeBounds;
20
21/// A batch of write operations (puts, deletes, and merges). All operations in
22/// the batch are applied atomically to the database. Put and delete operations
23/// replace earlier operations for the same key, while merge operations are
24/// preserved until a later put or delete replaces them.
25///
26/// # Examples
27/// ```rust
28/// # async fn run() -> Result<(), slatedb::Error> {
29/// #     use std::sync::Arc;
30/// #     use slatedb::object_store::memory::InMemory;
31///     use slatedb::Db;
32///     use slatedb::WriteBatch;
33///
34///     let object_store = Arc::new(InMemory::new());
35///     let db = Db::open("path/to/db", object_store).await?;
36///
37///     let mut batch = WriteBatch::new();
38///     batch.put(b"key1", b"value1");
39///     batch.put(b"key2", b"value2");
40///     batch.delete(b"key3");
41///
42///     db.write(batch).await;
43///     Ok(())
44/// # };
45/// # tokio_test::block_on(run());
46/// ```
47///
48/// Note that the `WriteBatch` has an unlimited size. This means that batch
49/// writes can exceed `l0_sst_size_bytes` (when `WAL` is disabled). It also
50/// means that WAL SSTs could get large if there's a large batch write.
51#[derive(Clone, Debug)]
52pub struct WriteBatch {
53    pub(crate) ops: BTreeMap<Bytes, SmallVec<[WriteOp; 1]>>,
54    pub(crate) op_count: usize,
55    pub(crate) has_merge_ops: bool,
56}
57
58impl Default for WriteBatch {
59    fn default() -> Self {
60        Self::new()
61    }
62}
63
64/// A write operation in a batch.
65#[derive(PartialEq, Clone)]
66pub(crate) enum WriteOp {
67    Put(Bytes, PutOptions),
68    Delete,
69    Merge(Bytes, MergeOptions),
70}
71
72impl std::fmt::Debug for WriteOp {
73    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
74        fn trunc(bytes: &Bytes) -> String {
75            if bytes.len() > 10 {
76                format!("{:?}...", &bytes[..10])
77            } else {
78                format!("{:?}", bytes)
79            }
80        }
81
82        match self {
83            WriteOp::Put(value, options) => {
84                let value = trunc(value);
85                write!(f, "Put({value}, {:?})", options)
86            }
87            WriteOp::Delete => write!(f, "Delete"),
88            WriteOp::Merge(value, options) => {
89                let value = trunc(value);
90                write!(f, "Merge({value}, {:?})", options)
91            }
92        }
93    }
94}
95
96impl WriteOp {
97    /// Convert WriteOp to RowEntry for queries
98    pub(crate) fn to_row_entry(
99        &self,
100        key: &Bytes,
101        seq: u64,
102        create_ts: Option<i64>,
103        expire_ts: Option<i64>,
104    ) -> RowEntry {
105        match self {
106            WriteOp::Put(value, _options) => {
107                // For queries, we don't need to compute expiration time here
108                // since these are uncommitted writes. The expiration will be
109                // computed when the batch is actually written.
110                RowEntry::new(
111                    key.clone(),
112                    ValueDeletable::Value(value.clone()),
113                    seq,
114                    create_ts,
115                    expire_ts,
116                )
117            }
118            WriteOp::Delete => RowEntry::new(
119                key.clone(),
120                ValueDeletable::Tombstone,
121                seq,
122                create_ts,
123                expire_ts,
124            ),
125            WriteOp::Merge(value, _options) => RowEntry::new(
126                key.clone(),
127                ValueDeletable::Merge(value.clone()),
128                seq,
129                create_ts,
130                expire_ts,
131            ),
132        }
133    }
134}
135
136impl WriteBatch {
137    pub fn new() -> Self {
138        WriteBatch {
139            ops: BTreeMap::new(),
140            op_count: 0,
141            has_merge_ops: false,
142        }
143    }
144
145    fn assert_kv<K, V>(&self, key: &K, value: &V)
146    where
147        K: AsRef<[u8]>,
148        V: AsRef<[u8]>,
149    {
150        let key = key.as_ref();
151        let value = value.as_ref();
152
153        assert!(!key.is_empty(), "key cannot be empty");
154        assert!(
155            key.len() <= u16::MAX as usize,
156            "key size must be <= u16::MAX"
157        );
158        assert!(
159            value.len() <= u32::MAX as usize,
160            "value size must be <= u32::MAX"
161        );
162    }
163
164    /// Put a key-value pair into the batch. Keys must not be empty.
165    ///
166    /// # Panics
167    /// - if the key is empty
168    /// - if the key size is larger than u16::MAX
169    /// - if the value size is larger than u32::MAX
170    pub fn put<K, V>(&mut self, key: K, value: V)
171    where
172        K: AsRef<[u8]>,
173        V: AsRef<[u8]>,
174    {
175        self.put_with_options(key, value, &PutOptions::default())
176    }
177
178    /// Put a key-value pair into the batch. Keys must not be empty.
179    ///
180    /// # Panics
181    /// - if the key is empty
182    /// - if the key size is larger than u16::MAX
183    /// - if the value size is larger than u32::MAX
184    pub fn put_with_options<K, V>(&mut self, key: K, value: V, options: &PutOptions)
185    where
186        K: AsRef<[u8]>,
187        V: AsRef<[u8]>,
188    {
189        self.put_bytes_with_options(
190            Bytes::copy_from_slice(key.as_ref()),
191            Bytes::copy_from_slice(value.as_ref()),
192            options,
193        )
194    }
195
196    /// Put a key-value pair into the batch using owned [`Bytes`], avoiding
197    /// the copies that [`WriteBatch::put`] performs via
198    /// `Bytes::copy_from_slice`. Prefer this form when the caller already
199    /// holds the data as [`Bytes`] (e.g. from a prior read, a zero-copy
200    /// buffer pool, or a client that produces [`Bytes`] directly). Keys must
201    /// not be empty.
202    ///
203    /// # Panics
204    /// - if the key is empty
205    /// - if the key size is larger than u16::MAX
206    /// - if the value size is larger than u32::MAX
207    pub fn put_bytes(&mut self, key: Bytes, value: Bytes) {
208        self.put_bytes_with_options(key, value, &PutOptions::default())
209    }
210
211    /// Put a key-value pair into the batch using owned [`Bytes`] with custom
212    /// options. See [`WriteBatch::put_bytes`] for why this form exists.
213    ///
214    /// # Panics
215    /// - if the key is empty
216    /// - if the key size is larger than u16::MAX
217    /// - if the value size is larger than u32::MAX
218    pub fn put_bytes_with_options(&mut self, key: Bytes, value: Bytes, options: &PutOptions) {
219        self.assert_kv(&key, &value);
220
221        // put will overwrite the existing key so we can safely
222        // remove all previous entries.
223        let previous = self
224            .ops
225            .insert(key, smallvec![WriteOp::Put(value, options.clone())]);
226        self.op_count += 1;
227        if let Some(previous) = previous {
228            self.op_count -= previous.len();
229        }
230    }
231
232    /// Merge a key-value pair into the batch. Keys must not be empty.
233    pub fn merge<K, V>(&mut self, key: K, value: V)
234    where
235        K: AsRef<[u8]>,
236        V: AsRef<[u8]>,
237    {
238        self.merge_with_options(key, value, &MergeOptions::default());
239    }
240
241    /// Merge a key-value pair into the batch with custom options.
242    pub fn merge_with_options<K, V>(&mut self, key: K, value: V, options: &MergeOptions)
243    where
244        K: AsRef<[u8]>,
245        V: AsRef<[u8]>,
246    {
247        self.assert_kv(&key, &value);
248
249        let key = key.as_ref();
250        let value = value.as_ref();
251        let op = WriteOp::Merge(Bytes::copy_from_slice(value), options.clone());
252        if let Some(ops) = self.ops.get_mut(key) {
253            ops.push(op);
254        } else {
255            self.ops.insert(Bytes::copy_from_slice(key), smallvec![op]);
256        }
257
258        self.has_merge_ops = true;
259        self.op_count += 1;
260    }
261
262    /// Delete a key-value pair into the batch. Keys must not be empty.
263    pub fn delete<K: AsRef<[u8]>>(&mut self, key: K) {
264        self.assert_kv(&key, &[]);
265
266        let key = Bytes::copy_from_slice(key.as_ref());
267
268        // delete will overwrite the existing key so we can safely
269        // remove all previous entries.
270        let previous = self.ops.insert(key, smallvec![WriteOp::Delete]);
271        self.op_count += 1;
272        if let Some(previous) = previous {
273            self.op_count -= previous.len();
274        }
275    }
276
277    pub fn is_empty(&self) -> bool {
278        self.ops.is_empty()
279    }
280
281    pub(crate) fn has_merge_ops(&self) -> bool {
282        self.has_merge_ops
283    }
284
285    pub(crate) fn op_count(&self) -> usize {
286        self.op_count
287    }
288
289    pub(crate) fn keys(&self) -> HashSet<Bytes> {
290        self.ops.keys().cloned().collect()
291    }
292
293    /// Converts a WriteBatch into a vector of RowEntry objects with
294    /// seq and timestamp set, applying the merge operator to any
295    /// mergeable entries.
296    ///
297    /// Also returning the size of keys and values in the batch
298    /// after merge iterator and overwrites are collapsed.
299    ///
300    /// When `extractor` is `Some`, the same iteration also derives
301    /// each entry's segment prefix (RFC-0024) and accumulates the
302    /// distinct prefixes into a `BTreeSet`. Returns
303    /// `EmptySegmentPrefix` if the extractor produces a zero-length
304    /// or absent prefix for any key.
305    ///
306    /// When `extractor` is `None`, the returned prefix set is empty.
307    #[allow(clippy::panic)]
308    pub(crate) async fn extract_entries(
309        &self,
310        seq: u64,
311        now: i64,
312        default_ttl: Option<u64>,
313        merger: Option<MergeOperatorType>,
314        extractor: Option<&dyn PrefixExtractor>,
315    ) -> Result<(Vec<RowEntry>, BTreeSet<Bytes>, u64), SlateDBError> {
316        let mut entries_bytes: u64 = 0;
317        let mut it: Box<dyn RowEntryIterator> = Box::new(WriteBatchIterator::new(
318            self,
319            ..,
320            IterationOrder::Ascending,
321            seq,
322            Some(now),
323            default_ttl,
324        ));
325        if self.has_merge_ops() {
326            if let Some(ref merge_operator) = merger {
327                it = Box::new(MergeOperatorIterator::new(
328                    merge_operator.clone(),
329                    it,
330                    false,
331                    None,
332                ));
333            } else {
334                return Err(SlateDBError::MergeOperatorMissing);
335            }
336        }
337
338        let mut entries: Vec<RowEntry> = Vec::new();
339        let mut touched_segments: BTreeSet<Bytes> = BTreeSet::new();
340        while let Some(entry) = it.next().await? {
341            // Verify that user has not supplied merge entries with different TTLs.
342            // This is not allowed. See #1581 for details.
343            match entries.last() {
344                Some(previous) if previous.key == entry.key => {
345                    if previous.expire_ts == entry.expire_ts {
346                        panic!("iterator emitted duplicate keys [key={:?}]", entry.key);
347                    }
348                    return Err(SlateDBError::IncompatibleMergeTtls {
349                        key: entry.key.clone(),
350                        previous_expire_ts: previous.expire_ts,
351                        current_expire_ts: entry.expire_ts,
352                    });
353                }
354                _ => {}
355            }
356            if let Some(ext) = extractor {
357                touched_segments.insert(extract_segment_prefix(ext, &entry.key)?);
358            }
359            entries_bytes += (entry.key.len() + entry.value.len()) as u64;
360            entries.push(entry);
361        }
362        Ok((entries, touched_segments, entries_bytes))
363    }
364}
365
366#[cfg(feature = "bench-internal")]
367pub mod benches {
368    use super::WriteBatch;
369    use crate::{Error, MergeOperator, PrefixExtractor, RowEntry};
370    use bytes::Bytes;
371    use std::collections::{BTreeSet, HashSet};
372    use std::sync::Arc;
373
374    pub fn keys(batch: &WriteBatch) -> HashSet<Bytes> {
375        batch.keys()
376    }
377
378    pub async fn extract_entries(
379        batch: &WriteBatch,
380        seq: u64,
381        now: i64,
382        default_ttl: Option<u64>,
383        merger: Option<Arc<dyn MergeOperator + Send + Sync>>,
384        extractor: Option<&dyn PrefixExtractor>,
385    ) -> Result<(Vec<RowEntry>, BTreeSet<Bytes>, u64), Error> {
386        batch
387            .extract_entries(seq, now, default_ttl, merger, extractor)
388            .await
389            .map_err(Into::into)
390    }
391}
392
393/// Iterator over `WriteBatch` entries.
394pub(crate) struct WriteBatchIterator {
395    iter: Peekable<Box<dyn Iterator<Item = RowEntry> + Send + Sync>>,
396    ordering: IterationOrder,
397}
398
399impl WriteBatchIterator {
400    /// Converts a batch write operation into a [`RowEntry`] with the supplied
401    /// sequence number and optional batch-local timestamp metadata.
402    ///
403    /// This wrapper exists because a [`WriteOp`] stores only the value-level
404    /// operation payload while the key lives in the batch map. The iterator is
405    /// where that key is available alongside the write options, current time,
406    /// and default TTL, so it computes the effective expiration timestamp and
407    /// passes both the key and timestamp metadata to [`WriteOp::to_row_entry`].
408    fn write_op_to_row_entry(
409        key: &Bytes,
410        op: &WriteOp,
411        seq: u64,
412        now: Option<i64>,
413        default_ttl: Option<u64>,
414    ) -> RowEntry {
415        let expire_ts = match (op, now) {
416            (WriteOp::Put(_, opts), Some(now)) => opts.expire_ts_from(default_ttl, now),
417            (WriteOp::Merge(_, opts), Some(now)) => opts.expire_ts_from(default_ttl, now),
418            _ => None,
419        };
420        op.to_row_entry(key, seq, now, expire_ts)
421    }
422
423    pub(crate) fn new(
424        batch: &WriteBatch,
425        range: impl RangeBounds<Bytes>,
426        ordering: IterationOrder,
427        seq: u64,
428        now: Option<i64>,
429        default_ttl: Option<u64>,
430    ) -> Self {
431        let entries: Vec<RowEntry> = match ordering {
432            IterationOrder::Ascending => batch
433                .ops
434                .range(range)
435                .flat_map(|(key, ops)| ops.iter().rev().map(move |op| (key, op)))
436                .map(|(key, op)| Self::write_op_to_row_entry(key, op, seq, now, default_ttl))
437                .collect(),
438            IterationOrder::Descending => batch
439                .ops
440                .range(range)
441                .rev()
442                .flat_map(|(key, ops)| ops.iter().rev().map(move |op| (key, op)))
443                .map(|(key, op)| Self::write_op_to_row_entry(key, op, seq, now, default_ttl))
444                .collect(),
445        };
446
447        let iter: Box<dyn Iterator<Item = RowEntry> + Send + Sync> = Box::new(entries.into_iter());
448
449        Self {
450            iter: iter.peekable(),
451            ordering,
452        }
453    }
454}
455
456#[async_trait]
457impl RowEntryIterator for WriteBatchIterator {
458    async fn init(&mut self) -> Result<(), crate::error::SlateDBError> {
459        Ok(())
460    }
461
462    async fn next(&mut self) -> Result<Option<RowEntry>, crate::error::SlateDBError> {
463        Ok(self.iter.next())
464    }
465
466    async fn seek(&mut self, next_key: &[u8]) -> Result<(), crate::error::SlateDBError> {
467        while let Some(entry) = self.iter.peek() {
468            if match self.ordering {
469                IterationOrder::Ascending => entry.key.as_ref() < next_key,
470                IterationOrder::Descending => entry.key.as_ref() > next_key,
471            } {
472                self.iter.next();
473            } else {
474                break;
475            }
476        }
477        Ok(())
478    }
479}
480
481#[cfg(test)]
482mod tests {
483    use rstest::rstest;
484
485    use super::*;
486    use crate::bytes_range::BytesRange;
487    use crate::config::Ttl;
488    use crate::test_utils::assert_iterator;
489    use crate::types::RowEntry;
490
491    struct WriteOpTestCase {
492        key: Vec<u8>,
493        // None is a delete and options will be ignored
494        value: Option<Vec<u8>>,
495        options: PutOptions,
496    }
497
498    fn ops_for_key<'a>(batch: &'a WriteBatch, key: &[u8]) -> &'a [WriteOp] {
499        let key = Bytes::copy_from_slice(key);
500        batch.ops.get(&key).expect("missing key").as_slice()
501    }
502
503    fn only_op<'a>(batch: &'a WriteBatch, key: &[u8]) -> &'a WriteOp {
504        let ops = ops_for_key(batch, key);
505        assert_eq!(ops.len(), 1);
506        &ops[0]
507    }
508
509    #[rstest]
510    #[case(vec![WriteOpTestCase {
511        key: b"key".to_vec(),
512        value: Some(b"value".to_vec()),
513        options: PutOptions::default(),
514    }])]
515    #[case(vec![WriteOpTestCase {
516        key: b"key".to_vec(),
517        value: None,
518        options: PutOptions::default(),
519    }])]
520    #[should_panic(expected = "key size must be <= u16::MAX")]
521    #[case(vec![WriteOpTestCase {
522        key: vec![b'k'; 65_536], // 2^16
523        value: None,
524        options: PutOptions::default(),
525    }])]
526    #[should_panic(expected = "value size must be <= u32::MAX")]
527    #[case(vec![WriteOpTestCase {
528        key: b"key".to_vec(),
529        value: Some(vec![b'x'; u32::MAX as usize + 1]), // 2^32
530        options: PutOptions::default(),
531    }])]
532    #[should_panic(expected = "key cannot be empty")]
533    #[case(vec![WriteOpTestCase {
534        key: b"".to_vec(),
535        value: Some(b"value".to_vec()),
536        options: PutOptions::default(),
537    }])]
538    #[should_panic(expected = "key cannot be empty")]
539    #[case(vec![WriteOpTestCase {
540        key: b"".to_vec(),
541        value: None,
542        options: PutOptions::default(),
543    }])]
544    fn test_put_delete_batch(#[case] test_case: Vec<WriteOpTestCase>) {
545        let mut batch = WriteBatch::new();
546        let mut expected_ops: BTreeMap<Bytes, SmallVec<[WriteOp; 1]>> = BTreeMap::new();
547        for test_case in test_case {
548            if let Some(value) = test_case.value {
549                batch.put_with_options(
550                    test_case.key.as_slice(),
551                    value.as_slice(),
552                    &test_case.options,
553                );
554                expected_ops.insert(
555                    Bytes::from(test_case.key.clone()),
556                    smallvec![WriteOp::Put(Bytes::from(value), test_case.options)],
557                );
558            } else {
559                batch.delete(test_case.key.as_slice());
560                expected_ops.insert(
561                    Bytes::from(test_case.key.clone()),
562                    smallvec![WriteOp::Delete],
563                );
564            }
565        }
566        assert_eq!(batch.ops, expected_ops);
567    }
568
569    #[test]
570    fn test_writebatch_deduplication() {
571        let mut batch = WriteBatch::new();
572        batch.put(b"key1", b"value1");
573        batch.put(b"key1", b"value2"); // Should overwrite previous
574
575        assert_eq!(batch.ops.len(), 1); // Only one entry due to deduplication
576        assert_eq!(batch.op_count(), 1);
577        let op = only_op(&batch, b"key1");
578        match op {
579            WriteOp::Put(value, _) => assert_eq!(value.as_ref(), b"value2"),
580            _ => panic!("Expected Put operation"),
581        }
582    }
583
584    #[tokio::test]
585    async fn test_writebatch_iterator_basic() {
586        let mut batch = WriteBatch::new();
587        batch.put(b"key1", b"value1");
588        batch.put(b"key3", b"value3");
589        batch.put(b"key2", b"value2");
590        batch.delete(b"key4");
591
592        let mut iter =
593            WriteBatchIterator::new(&batch, .., IterationOrder::Ascending, u64::MAX, None, None);
594
595        let expected = vec![
596            RowEntry::new_value(b"key1", b"value1", u64::MAX),
597            RowEntry::new_value(b"key2", b"value2", u64::MAX),
598            RowEntry::new_value(b"key3", b"value3", u64::MAX),
599            RowEntry::new(
600                Bytes::from("key4"),
601                ValueDeletable::Tombstone,
602                u64::MAX,
603                None,
604                None,
605            ),
606        ];
607
608        assert_iterator(&mut iter, expected).await;
609    }
610
611    #[tokio::test]
612    async fn test_writebatch_iterator_range() {
613        let mut batch = WriteBatch::new();
614        batch.put(b"key1", b"value1");
615        batch.put(b"key3", b"value3");
616        batch.put(b"key5", b"value5");
617
618        // Test range [key2, key4)
619        let mut iter = WriteBatchIterator::new(
620            &batch,
621            BytesRange::from(Bytes::from_static(b"key2")..Bytes::from_static(b"key4")),
622            IterationOrder::Ascending,
623            u64::MAX,
624            None,
625            None,
626        );
627
628        let expected = vec![RowEntry::new_value(b"key3", b"value3", u64::MAX)];
629
630        assert_iterator(&mut iter, expected).await;
631    }
632
633    #[tokio::test]
634    async fn test_writebatch_iterator_descending() {
635        let mut batch = WriteBatch::new();
636        batch.put(b"key1", b"value1");
637        batch.put(b"key3", b"value3");
638        batch.put(b"key2", b"value2");
639
640        let mut iter =
641            WriteBatchIterator::new(&batch, .., IterationOrder::Descending, u64::MAX, None, None);
642
643        let expected = vec![
644            RowEntry::new_value(b"key3", b"value3", u64::MAX),
645            RowEntry::new_value(b"key2", b"value2", u64::MAX),
646            RowEntry::new_value(b"key1", b"value1", u64::MAX),
647        ];
648
649        assert_iterator(&mut iter, expected).await;
650    }
651
652    #[tokio::test]
653    async fn test_writebatch_iterator_seek_ascending() {
654        let mut batch = WriteBatch::new();
655        batch.put(b"key1", b"value1");
656        batch.put(b"key3", b"value3");
657        batch.put(b"key5", b"value5");
658
659        let mut iter =
660            WriteBatchIterator::new(&batch, .., IterationOrder::Ascending, u64::MAX, None, None);
661
662        // Seek to key3
663        iter.seek(b"key3").await.unwrap();
664
665        // Should get key3 and key5
666        let expected = vec![
667            RowEntry::new_value(b"key3", b"value3", u64::MAX),
668            RowEntry::new_value(b"key5", b"value5", u64::MAX),
669        ];
670
671        assert_iterator(&mut iter, expected).await;
672    }
673
674    #[tokio::test]
675    async fn test_writebatch_iterator_seek_descending() {
676        let mut batch = WriteBatch::new();
677        batch.put(b"key1", b"value1");
678        batch.put(b"key3", b"value3");
679        batch.put(b"key5", b"value5");
680
681        let mut iter =
682            WriteBatchIterator::new(&batch, .., IterationOrder::Descending, u64::MAX, None, None);
683
684        // Seek to key3 (in descending, we want keys <= key3)
685        iter.seek(b"key3").await.unwrap();
686
687        // Should get key3 and key1 (in descending order)
688        let expected = vec![
689            RowEntry::new_value(b"key3", b"value3", u64::MAX),
690            RowEntry::new_value(b"key1", b"value1", u64::MAX),
691        ];
692
693        assert_iterator(&mut iter, expected).await;
694    }
695
696    #[tokio::test]
697    async fn test_writebatch_iterator_empty_batch() {
698        let batch = WriteBatch::new();
699        let mut iter =
700            WriteBatchIterator::new(&batch, .., IterationOrder::Ascending, u64::MAX, None, None);
701
702        let result = iter.next().await.unwrap();
703        assert!(result.is_none());
704    }
705
706    #[tokio::test]
707    async fn test_writebatch_iterator_seek_to_nonexistent() {
708        let mut batch = WriteBatch::new();
709        batch.put(b"key1", b"value1");
710        batch.put(b"key3", b"value3");
711
712        let mut iter =
713            WriteBatchIterator::new(&batch, .., IterationOrder::Ascending, u64::MAX, None, None);
714
715        // Seek to key2 (doesn't exist)
716        iter.seek(b"key2").await.unwrap();
717
718        // Should get key3 (next available)
719        let result = iter.next().await.unwrap();
720        assert!(result.is_some());
721        let entry = result.unwrap();
722        assert_eq!(entry.key, Bytes::from_static(b"key3"));
723        assert_eq!(
724            entry.value,
725            ValueDeletable::Value(Bytes::from_static(b"value3"))
726        );
727    }
728
729    #[tokio::test]
730    async fn test_writebatch_iterator_seek_beyond_end() {
731        let mut batch = WriteBatch::new();
732        batch.put(b"key1", b"value1");
733        batch.put(b"key3", b"value3");
734
735        let mut iter =
736            WriteBatchIterator::new(&batch, .., IterationOrder::Ascending, u64::MAX, None, None);
737
738        // Seek beyond maximum key
739        iter.seek(b"key9").await.unwrap();
740
741        // Should be exhausted
742        let result = iter.next().await.unwrap();
743        assert!(result.is_none());
744    }
745
746    #[tokio::test]
747    async fn test_writebatch_iterator_multiple_tombstones() {
748        let mut batch = WriteBatch::new();
749        batch.put(b"key1", b"value1");
750        batch.delete(b"key2");
751        batch.put(b"key3", b"value3");
752        batch.delete(b"key4");
753
754        let mut iter =
755            WriteBatchIterator::new(&batch, .., IterationOrder::Ascending, u64::MAX, None, None);
756
757        let expected = vec![
758            RowEntry::new_value(b"key1", b"value1", u64::MAX),
759            RowEntry::new(
760                Bytes::from_static(b"key2"),
761                ValueDeletable::Tombstone,
762                u64::MAX,
763                None,
764                None,
765            ),
766            RowEntry::new_value(b"key3", b"value3", u64::MAX),
767            RowEntry::new(
768                Bytes::from_static(b"key4"),
769                ValueDeletable::Tombstone,
770                u64::MAX,
771                None,
772                None,
773            ),
774        ];
775
776        assert_iterator(&mut iter, expected).await;
777    }
778
779    #[tokio::test]
780    async fn test_writebatch_iterator_seek_before_first() {
781        let mut batch = WriteBatch::new();
782        batch.put(b"key2", b"value2");
783        batch.put(b"key3", b"value3");
784
785        let mut iter =
786            WriteBatchIterator::new(&batch, .., IterationOrder::Ascending, u64::MAX, None, None);
787
788        // Seek before first key
789        iter.seek(b"key1").await.unwrap();
790
791        // Should get key2 (first available)
792        let result = iter.next().await.unwrap();
793        assert!(result.is_some());
794        let entry = result.unwrap();
795        assert_eq!(entry.key, Bytes::from_static(b"key2"));
796    }
797
798    #[tokio::test]
799    async fn test_writebatch_iterator_range_with_tombstones() {
800        let mut batch = WriteBatch::new();
801        batch.put(b"key1", b"value1");
802        batch.delete(b"key2");
803        batch.put(b"key3", b"value3");
804        batch.delete(b"key4");
805        batch.put(b"key5", b"value5");
806
807        // Range [key2, key4) should include tombstones
808        let mut iter = WriteBatchIterator::new(
809            &batch,
810            BytesRange::from(Bytes::from_static(b"key2")..Bytes::from_static(b"key4")),
811            IterationOrder::Ascending,
812            u64::MAX,
813            None,
814            None,
815        );
816
817        let expected = vec![
818            RowEntry::new(
819                Bytes::from_static(b"key2"),
820                ValueDeletable::Tombstone,
821                u64::MAX,
822                None,
823                None,
824            ),
825            RowEntry::new_value(b"key3", b"value3", u64::MAX),
826        ];
827
828        assert_iterator(&mut iter, expected).await;
829    }
830
831    #[test]
832    fn should_create_merge_operation_with_default_options() {
833        // Given: an empty WriteBatch
834        let mut batch = WriteBatch::new();
835
836        // When: adding a merge operation
837        batch.merge(b"key1", b"value1");
838
839        // Then: the batch should contain one merge operation with default options
840        assert_eq!(batch.ops.len(), 1);
841        let op = only_op(&batch, b"key1");
842        match op {
843            WriteOp::Merge(value, options) => {
844                assert_eq!(value.as_ref(), b"value1");
845                assert_eq!(options, &MergeOptions::default());
846            }
847            _ => panic!("Expected Merge operation"),
848        }
849    }
850
851    #[test]
852    fn should_create_merge_operation_with_custom_options() {
853        // Given: an empty WriteBatch and custom merge options
854        let mut batch = WriteBatch::new();
855        let merge_options = MergeOptions {
856            ttl: Ttl::ExpireAfter(3600), // 1 hour
857        };
858
859        // When: adding a merge operation with custom options
860        batch.merge_with_options(b"key1", b"value1", &merge_options);
861
862        // Then: the batch should contain one merge operation with the custom options
863        assert_eq!(batch.ops.len(), 1);
864        let op = only_op(&batch, b"key1");
865        match op {
866            WriteOp::Merge(value, options) => {
867                assert_eq!(value.as_ref(), b"value1");
868                assert_eq!(options.ttl, Ttl::ExpireAfter(3600));
869            }
870            _ => panic!("Expected Merge operation"),
871        }
872    }
873
874    #[test]
875    fn should_allow_multiple_merges_for_same_key() {
876        // Given: an empty WriteBatch
877        let mut batch = WriteBatch::new();
878
879        // When: adding multiple merge operations for the same key
880        batch.merge(b"key1", b"value1");
881        batch.merge(b"key1", b"value2");
882        batch.merge(b"key1", b"value3");
883
884        // Then: all merge operations should be preserved (not deduplicated)
885        assert_eq!(batch.ops.len(), 1);
886        assert_eq!(batch.op_count(), 3);
887
888        // And: all three operations should exist for the key in insertion order
889        let ops = ops_for_key(&batch, b"key1");
890        let op1 = &ops[0];
891        let op2 = &ops[1];
892        let op3 = &ops[2];
893
894        match (op1, op2, op3) {
895            (WriteOp::Merge(v1, _), WriteOp::Merge(v2, _), WriteOp::Merge(v3, _)) => {
896                assert_eq!(v1.as_ref(), b"value1");
897                assert_eq!(v2.as_ref(), b"value2");
898                assert_eq!(v3.as_ref(), b"value3");
899            }
900            _ => panic!("Expected all Merge operations"),
901        }
902    }
903
904    #[test]
905    fn should_allow_merges_for_different_keys() {
906        // Given: an empty WriteBatch
907        let mut batch = WriteBatch::new();
908
909        // When: adding merge operations for different keys
910        batch.merge(b"key1", b"value1");
911        batch.merge(b"key2", b"value2");
912        batch.merge(b"key3", b"value3");
913
914        // Then: all merge operations should be preserved
915        assert_eq!(batch.ops.len(), 3);
916        assert_eq!(batch.op_count(), 3);
917
918        // And: all keys should exist
919        assert!(batch.ops.contains_key(&Bytes::from_static(b"key1")));
920        assert!(batch.ops.contains_key(&Bytes::from_static(b"key2")));
921        assert!(batch.ops.contains_key(&Bytes::from_static(b"key3")));
922    }
923
924    #[test]
925    fn should_track_merge_presence() {
926        let mut batch = WriteBatch::new();
927        assert!(!batch.has_merge_ops());
928
929        batch.merge(b"key1", b"value1");
930
931        assert!(batch.has_merge_ops());
932    }
933
934    #[test]
935    fn should_remember_merge_presence_when_put_overwrites_merge_ops() {
936        let mut batch = WriteBatch::new();
937        batch.merge(b"key1", b"value1");
938        batch.merge(b"key1", b"value2");
939        assert!(batch.has_merge_ops());
940
941        batch.put(b"key1", b"final");
942
943        assert!(batch.has_merge_ops());
944        assert_eq!(batch.op_count(), 1);
945        match only_op(&batch, b"key1") {
946            WriteOp::Put(value, _) => {
947                assert_eq!(value.as_ref(), b"final");
948            }
949            _ => panic!("Expected Put operation"),
950        }
951    }
952
953    #[test]
954    fn should_remember_merge_presence_when_delete_overwrites_merge_ops() {
955        let mut batch = WriteBatch::new();
956        batch.merge(b"key1", b"value1");
957        batch.merge(b"key2", b"value2");
958        assert!(batch.has_merge_ops());
959
960        batch.delete(b"key1");
961
962        assert!(batch.has_merge_ops());
963        assert_eq!(batch.ops.len(), 2);
964        assert_eq!(batch.op_count(), 2);
965
966        match only_op(&batch, b"key1") {
967            WriteOp::Delete => {}
968            _ => panic!("Expected Delete operation"),
969        }
970
971        match only_op(&batch, b"key2") {
972            WriteOp::Merge(value, _) => {
973                assert_eq!(value.as_ref(), b"value2");
974            }
975            _ => panic!("Expected Merge operation"),
976        }
977    }
978
979    #[test]
980    fn should_preserve_both_put_and_merge_when_merge_comes_after_put() {
981        // Given: a WriteBatch with a put operation
982        let mut batch = WriteBatch::new();
983        batch.put(b"key1", b"put_value");
984
985        // When: adding a merge operation for the same key
986        batch.merge(b"key1", b"merge_value");
987
988        // Then: both operations should remain (merge accumulates on top of put)
989        assert_eq!(batch.ops.len(), 1);
990        assert_eq!(batch.op_count(), 2);
991        let ops = ops_for_key(&batch, b"key1");
992
993        // Verify the put operation exists
994        let put_op = &ops[0];
995        match put_op {
996            WriteOp::Put(value, _) => {
997                assert_eq!(value.as_ref(), b"put_value");
998            }
999            _ => panic!("Expected Put operation"),
1000        }
1001
1002        // Verify the merge operation exists
1003        let merge_op = &ops[1];
1004        match merge_op {
1005            WriteOp::Merge(value, _) => {
1006                assert_eq!(value.as_ref(), b"merge_value");
1007            }
1008            _ => panic!("Expected Merge operation"),
1009        }
1010    }
1011
1012    #[test]
1013    fn should_preserve_both_delete_and_merge_when_merge_comes_after_delete() {
1014        // Given: a WriteBatch with a delete operation
1015        let mut batch = WriteBatch::new();
1016        batch.delete(b"key1");
1017
1018        // When: adding a merge operation for the same key
1019        batch.merge(b"key1", b"merge_value");
1020
1021        // Then: both operations should remain (merge accumulates on top of delete)
1022        assert_eq!(batch.ops.len(), 1);
1023        assert_eq!(batch.op_count(), 2);
1024        let ops = ops_for_key(&batch, b"key1");
1025
1026        // Verify the delete operation exists
1027        let delete_op = &ops[0];
1028        match delete_op {
1029            WriteOp::Delete => {}
1030            _ => panic!("Expected Delete operation"),
1031        }
1032
1033        // Verify the merge operation exists
1034        let merge_op = &ops[1];
1035        match merge_op {
1036            WriteOp::Merge(value, _) => {
1037                assert_eq!(value.as_ref(), b"merge_value");
1038            }
1039            _ => panic!("Expected Merge operation"),
1040        }
1041    }
1042
1043    #[test]
1044    fn should_deduplicate_merge_when_put_is_added_for_same_key() {
1045        // Given: a WriteBatch with a merge operation
1046        let mut batch = WriteBatch::new();
1047        batch.merge(b"key1", b"merge_value");
1048
1049        // When: adding a put operation for the same key
1050        batch.put(b"key1", b"put_value");
1051
1052        // Then: only the put operation should remain (merge is deduplicated by put)
1053        assert_eq!(batch.ops.len(), 1);
1054        assert_eq!(batch.op_count(), 1);
1055
1056        let op = only_op(&batch, b"key1");
1057        match op {
1058            WriteOp::Put(value, _) => {
1059                assert_eq!(value.as_ref(), b"put_value");
1060            }
1061            _ => panic!("Expected Put operation"),
1062        }
1063    }
1064
1065    #[test]
1066    fn should_deduplicate_merge_when_delete_is_added_for_same_key() {
1067        // Given: a WriteBatch with a merge operation
1068        let mut batch = WriteBatch::new();
1069        batch.merge(b"key1", b"merge_value");
1070
1071        // When: adding a delete operation for the same key
1072        batch.delete(b"key1");
1073
1074        // Then: only the delete operation should remain (merge is deduplicated by delete)
1075        assert_eq!(batch.ops.len(), 1);
1076        assert_eq!(batch.op_count(), 1);
1077
1078        let op = only_op(&batch, b"key1");
1079        match op {
1080            WriteOp::Delete => {}
1081            _ => panic!("Expected Delete operation"),
1082        }
1083    }
1084
1085    #[tokio::test]
1086    async fn should_iterate_over_mixed_operations_including_merges() {
1087        // Given: a WriteBatch with mixed operations including multiple merges
1088        let mut batch = WriteBatch::new();
1089        batch.put(b"key1", b"value1");
1090        batch.merge(b"key2", b"merge1");
1091        batch.merge(b"key2", b"merge2");
1092        batch.delete(b"key3");
1093
1094        // When: creating an iterator
1095        let mut iter =
1096            WriteBatchIterator::new(&batch, .., IterationOrder::Ascending, u64::MAX, None, None);
1097
1098        // Then: the iterator should return all operations in order
1099        let expected = vec![
1100            RowEntry::new_value(b"key1", b"value1", u64::MAX),
1101            RowEntry::new(
1102                Bytes::from_static(b"key2"),
1103                ValueDeletable::Merge(Bytes::from_static(b"merge2")),
1104                u64::MAX,
1105                None,
1106                None,
1107            ),
1108            RowEntry::new(
1109                Bytes::from_static(b"key2"),
1110                ValueDeletable::Merge(Bytes::from_static(b"merge1")),
1111                u64::MAX,
1112                None,
1113                None,
1114            ),
1115            RowEntry::new(
1116                Bytes::from_static(b"key3"),
1117                ValueDeletable::Tombstone,
1118                u64::MAX,
1119                None,
1120                None,
1121            ),
1122        ];
1123
1124        assert_iterator(&mut iter, expected).await;
1125    }
1126
1127    #[tokio::test]
1128    async fn should_iterate_over_multiple_merges_for_same_key() {
1129        // Given: a WriteBatch with multiple merges for the same key
1130        let mut batch = WriteBatch::new();
1131        batch.merge(b"key1", b"merge1");
1132        batch.merge(b"key1", b"merge2");
1133        batch.merge(b"key1", b"merge3");
1134
1135        // When: creating an iterator
1136        let mut iter =
1137            WriteBatchIterator::new(&batch, .., IterationOrder::Ascending, u64::MAX, None, None);
1138
1139        // Then: the iterator should return all merge operations
1140        let expected = vec![
1141            RowEntry::new(
1142                Bytes::from_static(b"key1"),
1143                ValueDeletable::Merge(Bytes::from_static(b"merge3")),
1144                u64::MAX,
1145                None,
1146                None,
1147            ),
1148            RowEntry::new(
1149                Bytes::from_static(b"key1"),
1150                ValueDeletable::Merge(Bytes::from_static(b"merge2")),
1151                u64::MAX,
1152                None,
1153                None,
1154            ),
1155            RowEntry::new(
1156                Bytes::from_static(b"key1"),
1157                ValueDeletable::Merge(Bytes::from_static(b"merge1")),
1158                u64::MAX,
1159                None,
1160                None,
1161            ),
1162        ];
1163
1164        assert_iterator(&mut iter, expected).await;
1165    }
1166
1167    #[tokio::test]
1168    async fn should_iterate_over_merges_in_descending_order() {
1169        // Given: a WriteBatch with merges for different keys
1170        let mut batch = WriteBatch::new();
1171        batch.merge(b"key1", b"merge1");
1172        batch.merge(b"key2", b"merge2");
1173        batch.merge(b"key1", b"merge3");
1174
1175        // When: creating a descending iterator
1176        let mut iter =
1177            WriteBatchIterator::new(&batch, .., IterationOrder::Descending, u64::MAX, None, None);
1178
1179        // Then: the iterator should return operations in descending order
1180        let expected = vec![
1181            RowEntry::new(
1182                Bytes::from_static(b"key2"),
1183                ValueDeletable::Merge(Bytes::from_static(b"merge2")),
1184                u64::MAX,
1185                None,
1186                None,
1187            ),
1188            RowEntry::new(
1189                Bytes::from_static(b"key1"),
1190                ValueDeletable::Merge(Bytes::from_static(b"merge3")),
1191                u64::MAX,
1192                None,
1193                None,
1194            ),
1195            RowEntry::new(
1196                Bytes::from_static(b"key1"),
1197                ValueDeletable::Merge(Bytes::from_static(b"merge1")),
1198                u64::MAX,
1199                None,
1200                None,
1201            ),
1202        ];
1203
1204        assert_iterator(&mut iter, expected).await;
1205    }
1206
1207    #[tokio::test]
1208    async fn should_iterate_over_merges_in_range() {
1209        // Given: a WriteBatch with merges for different keys
1210        let mut batch = WriteBatch::new();
1211        batch.merge(b"key1", b"merge1");
1212        batch.merge(b"key3", b"merge3");
1213        batch.merge(b"key5", b"merge5");
1214
1215        // When: creating an iterator with a range filter
1216        let mut iter = WriteBatchIterator::new(
1217            &batch,
1218            BytesRange::from(Bytes::from_static(b"key2")..Bytes::from_static(b"key4")),
1219            IterationOrder::Ascending,
1220            u64::MAX,
1221            None,
1222            None,
1223        );
1224
1225        // Then: the iterator should only return merges within the range
1226        let expected = vec![RowEntry::new(
1227            Bytes::from_static(b"key3"),
1228            ValueDeletable::Merge(Bytes::from_static(b"merge3")),
1229            u64::MAX,
1230            None,
1231            None,
1232        )];
1233
1234        assert_iterator(&mut iter, expected).await;
1235    }
1236
1237    #[tokio::test]
1238    async fn should_seek_to_merge_operations() {
1239        // Given: a WriteBatch with merges for different keys
1240        let mut batch = WriteBatch::new();
1241        batch.merge(b"key1", b"merge1");
1242        batch.merge(b"key3", b"merge3");
1243        batch.merge(b"key5", b"merge5");
1244
1245        // When: creating an iterator and seeking to key3
1246        let mut iter =
1247            WriteBatchIterator::new(&batch, .., IterationOrder::Ascending, u64::MAX, None, None);
1248        iter.seek(b"key3").await.unwrap();
1249
1250        // Then: the iterator should return key3 and key5
1251        let expected = vec![
1252            RowEntry::new(
1253                Bytes::from_static(b"key3"),
1254                ValueDeletable::Merge(Bytes::from_static(b"merge3")),
1255                u64::MAX,
1256                None,
1257                None,
1258            ),
1259            RowEntry::new(
1260                Bytes::from_static(b"key5"),
1261                ValueDeletable::Merge(Bytes::from_static(b"merge5")),
1262                u64::MAX,
1263                None,
1264                None,
1265            ),
1266        ];
1267
1268        assert_iterator(&mut iter, expected).await;
1269    }
1270
1271    // Tests for extract_entries
1272    struct StringConcatMergeOperator;
1273
1274    impl crate::merge_operator::MergeOperator for StringConcatMergeOperator {
1275        fn merge(
1276            &self,
1277            _key: &Bytes,
1278            existing_value: Option<Bytes>,
1279            operand: Bytes,
1280        ) -> Result<Bytes, crate::merge_operator::MergeOperatorError> {
1281            match existing_value {
1282                Some(base) => {
1283                    let mut merged = base.to_vec();
1284                    merged.extend_from_slice(&operand);
1285                    Ok(Bytes::from(merged))
1286                }
1287                None => Ok(operand),
1288            }
1289        }
1290    }
1291
1292    #[tokio::test]
1293    async fn should_preserve_interleaved_merge_order_by_key() {
1294        let mut batch = WriteBatch::new();
1295        batch.merge(b"key1", b"a");
1296        batch.merge(b"key2", b"x");
1297        batch.merge(b"key1", b"b");
1298
1299        let mut iter =
1300            WriteBatchIterator::new(&batch, .., IterationOrder::Ascending, u64::MAX, None, None);
1301
1302        let expected = vec![
1303            RowEntry::new(
1304                Bytes::from_static(b"key1"),
1305                ValueDeletable::Merge(Bytes::from_static(b"b")),
1306                u64::MAX,
1307                None,
1308                None,
1309            ),
1310            RowEntry::new(
1311                Bytes::from_static(b"key1"),
1312                ValueDeletable::Merge(Bytes::from_static(b"a")),
1313                u64::MAX,
1314                None,
1315                None,
1316            ),
1317            RowEntry::new(
1318                Bytes::from_static(b"key2"),
1319                ValueDeletable::Merge(Bytes::from_static(b"x")),
1320                u64::MAX,
1321                None,
1322                None,
1323            ),
1324        ];
1325        assert_iterator(&mut iter, expected).await;
1326
1327        let merge_operator = Some(std::sync::Arc::new(StringConcatMergeOperator)
1328            as crate::merge_operator::MergeOperatorType);
1329        let (result, _, _) = batch
1330            .extract_entries(100, 1000, None, merge_operator, None)
1331            .await
1332            .unwrap();
1333
1334        assert_eq!(result.len(), 2);
1335        assert_eq!(result[0].key, Bytes::from_static(b"key1"));
1336        assert_eq!(
1337            result[0].value,
1338            ValueDeletable::Merge(Bytes::from_static(b"ab"))
1339        );
1340        assert_eq!(result[1].key, Bytes::from_static(b"key2"));
1341        assert_eq!(
1342            result[1].value,
1343            ValueDeletable::Merge(Bytes::from_static(b"x"))
1344        );
1345    }
1346
1347    #[tokio::test]
1348    async fn should_extract_entries_no_merges() {
1349        // Given: a WriteBatch with no merge operations
1350        let mut batch = WriteBatch::new();
1351        batch.put(b"key1", b"value1");
1352        batch.put(b"key2", b"value2");
1353        batch.delete(b"key3");
1354
1355        // When: extracting entries
1356        let (result, _, _) = batch
1357            .extract_entries(100, 1000, None, None, None)
1358            .await
1359            .unwrap();
1360
1361        // Then: should return entries for all operations without merging
1362        let mut entries = result.into_iter().collect::<Vec<_>>();
1363        entries.sort_by_key(|e| e.key.clone());
1364
1365        assert_eq!(entries.len(), 3);
1366        assert_eq!(entries[0].key, Bytes::from_static(b"key1"));
1367        assert_eq!(
1368            entries[0].value,
1369            ValueDeletable::Value(Bytes::from_static(b"value1"))
1370        );
1371        assert_eq!(entries[1].key, Bytes::from_static(b"key2"));
1372        assert_eq!(
1373            entries[1].value,
1374            ValueDeletable::Value(Bytes::from_static(b"value2"))
1375        );
1376        assert_eq!(entries[2].key, Bytes::from_static(b"key3"));
1377        assert_eq!(entries[2].value, ValueDeletable::Tombstone);
1378    }
1379
1380    #[tokio::test]
1381    async fn should_error_extracting_entries_with_merges_without_merge_operator() {
1382        let mut batch = WriteBatch::new();
1383        batch.put(b"key1", b"value1");
1384        batch.merge(b"key2", b"merge1");
1385
1386        let err = batch
1387            .extract_entries(100, 1000, None, None, None)
1388            .await
1389            .unwrap_err();
1390        assert!(matches!(err, SlateDBError::MergeOperatorMissing));
1391    }
1392
1393    #[tokio::test]
1394    async fn should_extract_entries_multiple_merges() {
1395        // Given: a WriteBatch with multiple merge operations for the same key
1396        let mut batch = WriteBatch::new();
1397        batch.merge(b"key1", b"merge1");
1398        batch.merge(b"key1", b"merge2");
1399        batch.merge(b"key1", b"merge3");
1400
1401        // When: extracting entries with a merge operator
1402        let merge_operator = Some(std::sync::Arc::new(StringConcatMergeOperator)
1403            as crate::merge_operator::MergeOperatorType);
1404        let (result, _, _) = batch
1405            .extract_entries(100, 1000, None, merge_operator, None)
1406            .await
1407            .unwrap();
1408
1409        // Then: should return merged entry
1410        assert_eq!(result.len(), 1);
1411        assert_eq!(result[0].key, Bytes::from_static(b"key1"));
1412        assert_eq!(
1413            result[0].value,
1414            ValueDeletable::Merge(Bytes::from_static(b"merge1merge2merge3"))
1415        );
1416    }
1417
1418    #[tokio::test]
1419    async fn should_extract_same_key_merges_with_same_effective_ttl() {
1420        let mut batch = WriteBatch::new();
1421        batch.merge_with_options(
1422            b"key1",
1423            b"a",
1424            &MergeOptions {
1425                ttl: Ttl::ExpireAfter(3600),
1426            },
1427        );
1428        batch.merge_with_options(
1429            b"key1",
1430            b"b",
1431            &MergeOptions {
1432                ttl: Ttl::ExpireAt(4600),
1433            },
1434        );
1435
1436        let merge_operator = Some(std::sync::Arc::new(StringConcatMergeOperator)
1437            as crate::merge_operator::MergeOperatorType);
1438        let (result, _, _) = batch
1439            .extract_entries(100, 1000, None, merge_operator, None)
1440            .await
1441            .unwrap();
1442
1443        assert_eq!(result.len(), 1);
1444        assert_eq!(result[0].key, Bytes::from_static(b"key1"));
1445        assert_eq!(
1446            result[0].value,
1447            ValueDeletable::Merge(Bytes::from_static(b"ab"))
1448        );
1449        assert_eq!(result[0].expire_ts, Some(4600));
1450    }
1451
1452    #[tokio::test]
1453    async fn should_error_extracting_same_key_merges_with_different_ttls() {
1454        let mut batch = WriteBatch::new();
1455        batch.merge_with_options(
1456            b"key1",
1457            b"a",
1458            &MergeOptions {
1459                ttl: Ttl::ExpireAfter(3600),
1460            },
1461        );
1462        batch.merge_with_options(
1463            b"key1",
1464            b"b",
1465            &MergeOptions {
1466                ttl: Ttl::ExpireAfter(7200),
1467            },
1468        );
1469
1470        let merge_operator = Some(std::sync::Arc::new(StringConcatMergeOperator)
1471            as crate::merge_operator::MergeOperatorType);
1472        let err = batch
1473            .extract_entries(100, 1000, None, merge_operator, None)
1474            .await
1475            .unwrap_err();
1476
1477        match err {
1478            SlateDBError::IncompatibleMergeTtls {
1479                key,
1480                previous_expire_ts,
1481                current_expire_ts,
1482            } => {
1483                assert_eq!(key, Bytes::from_static(b"key1"));
1484                assert_eq!(previous_expire_ts, Some(8200));
1485                assert_eq!(current_expire_ts, Some(4600));
1486            }
1487            other => panic!("expected IncompatibleMergeTtls, got {other:?}"),
1488        }
1489    }
1490
1491    #[tokio::test]
1492    async fn should_allow_different_keys_with_different_merge_ttls() {
1493        let mut batch = WriteBatch::new();
1494        batch.merge_with_options(
1495            b"key1",
1496            b"a",
1497            &MergeOptions {
1498                ttl: Ttl::ExpireAfter(3600),
1499            },
1500        );
1501        batch.merge_with_options(
1502            b"key2",
1503            b"b",
1504            &MergeOptions {
1505                ttl: Ttl::ExpireAfter(7200),
1506            },
1507        );
1508
1509        let merge_operator = Some(std::sync::Arc::new(StringConcatMergeOperator)
1510            as crate::merge_operator::MergeOperatorType);
1511        let (result, _, _) = batch
1512            .extract_entries(100, 1000, None, merge_operator, None)
1513            .await
1514            .unwrap();
1515
1516        assert_eq!(result.len(), 2);
1517        assert_eq!(result[0].key, Bytes::from_static(b"key1"));
1518        assert_eq!(result[0].expire_ts, Some(4600));
1519        assert_eq!(result[1].key, Bytes::from_static(b"key2"));
1520        assert_eq!(result[1].expire_ts, Some(8200));
1521    }
1522
1523    #[tokio::test]
1524    async fn should_error_extracting_put_then_merge_with_different_ttls() {
1525        let mut batch = WriteBatch::new();
1526        batch.put(b"key1", b"base");
1527        batch.merge_with_options(
1528            b"key1",
1529            b"a",
1530            &MergeOptions {
1531                ttl: Ttl::ExpireAfter(7200),
1532            },
1533        );
1534
1535        let merge_operator = Some(std::sync::Arc::new(StringConcatMergeOperator)
1536            as crate::merge_operator::MergeOperatorType);
1537        let err = batch
1538            .extract_entries(100, 1000, None, merge_operator, None)
1539            .await
1540            .unwrap_err();
1541
1542        assert!(matches!(err, SlateDBError::IncompatibleMergeTtls { .. }));
1543    }
1544
1545    #[tokio::test]
1546    async fn should_error_extracting_delete_then_merge_with_different_ttls() {
1547        let mut batch = WriteBatch::new();
1548        batch.delete(b"key1");
1549        batch.merge_with_options(
1550            b"key1",
1551            b"a",
1552            &MergeOptions {
1553                ttl: Ttl::ExpireAfter(7200),
1554            },
1555        );
1556
1557        let merge_operator = Some(std::sync::Arc::new(StringConcatMergeOperator)
1558            as crate::merge_operator::MergeOperatorType);
1559        let err = batch
1560            .extract_entries(100, 1000, None, merge_operator, None)
1561            .await
1562            .unwrap_err();
1563
1564        assert!(matches!(err, SlateDBError::IncompatibleMergeTtls { .. }));
1565    }
1566
1567    #[tokio::test]
1568    async fn should_extract_entries_delete_then_merge() {
1569        // Given: a WriteBatch with a delete followed by merges
1570        let mut batch = WriteBatch::new();
1571        batch.delete(b"key1");
1572        batch.merge(b"key1", b"merge1");
1573        batch.merge(b"key1", b"merge2");
1574
1575        // When: extracting entries with a merge operator
1576        let merge_operator = Some(std::sync::Arc::new(StringConcatMergeOperator)
1577            as crate::merge_operator::MergeOperatorType);
1578        let (result, _, _) = batch
1579            .extract_entries(100, 1000, None, merge_operator, None)
1580            .await
1581            .unwrap();
1582
1583        // Then: should return merged entry (delete gets merged with merges)
1584        assert_eq!(result.len(), 1);
1585        assert_eq!(result[0].key, Bytes::from_static(b"key1"));
1586        // The delete (tombstone) gets converted to empty bytes, then merged
1587        assert_eq!(
1588            result[0].value,
1589            ValueDeletable::Value(Bytes::from_static(b"merge1merge2"))
1590        );
1591    }
1592
1593    #[tokio::test]
1594    async fn should_extract_entries_value_then_merge() {
1595        // Given: a WriteBatch with a put followed by merges
1596        let mut batch = WriteBatch::new();
1597        batch.put(b"key1", b"value");
1598        batch.merge(b"key1", b"merge1");
1599        batch.merge(b"key1", b"merge2");
1600
1601        // When: extracting entries with a merge operator
1602        let merge_operator = Some(std::sync::Arc::new(StringConcatMergeOperator)
1603            as crate::merge_operator::MergeOperatorType);
1604        let (result, _, _) = batch
1605            .extract_entries(100, 1000, None, merge_operator, None)
1606            .await
1607            .unwrap();
1608
1609        // Then: should return merged entry
1610        assert_eq!(result.len(), 1);
1611        assert_eq!(result[0].key, Bytes::from_static(b"key1"));
1612        // Value gets merged with merges
1613        assert_eq!(
1614            result[0].value,
1615            ValueDeletable::Value(Bytes::from_static(b"valuemerge1merge2"))
1616        );
1617    }
1618}
slatedb/batch.rs

slatedb/
batch.rs
