Skip to main content

prolly_store_slatedb/
lib.rs

1#![doc = include_str!("../README.md")]
2
3use std::collections::{hash_map::Entry, HashMap};
4use std::future::Future;
5use std::sync::{Arc, Mutex};
6
7use futures_util::stream::{self, StreamExt};
8use slatedb::bytes::Bytes;
9use slatedb::config::Settings;
10use slatedb::object_store::ObjectStore;
11use slatedb::{Db, WriteBatch};
12use tokio::runtime::{Builder, Runtime};
13
14use prolly::{
15    BatchOp, Cid, Error, ManifestStore, ManifestStoreScan, ManifestUpdate, NamedRootManifest,
16    NodeStoreScan, RootCondition, RootManifest, RootWrite, Store, TransactionConflict,
17    TransactionNodeWrite, TransactionUpdate, TransactionalStore,
18};
19
20struct OrderedBatchReadPlan<'a> {
21    unique_keys: Vec<&'a [u8]>,
22    positions: Option<Vec<usize>>,
23}
24
25impl<'a> OrderedBatchReadPlan<'a> {
26    fn new(keys: &[&'a [u8]]) -> Self {
27        let mut unique_indexes = HashMap::with_capacity(keys.len());
28        let mut unique_keys = Vec::with_capacity(keys.len());
29        let mut positions = None;
30        for key in keys {
31            match unique_indexes.entry(*key) {
32                Entry::Occupied(entry) => positions
33                    .get_or_insert_with(|| (0..unique_keys.len()).collect::<Vec<_>>())
34                    .push(*entry.get()),
35                Entry::Vacant(entry) => {
36                    let index = unique_keys.len();
37                    unique_keys.push(*key);
38                    if let Some(positions) = positions.as_mut() {
39                        positions.push(index);
40                    }
41                    entry.insert(index);
42                }
43            }
44        }
45        Self {
46            unique_keys,
47            positions,
48        }
49    }
50
51    fn unique_keys(&self) -> &[&'a [u8]] {
52        &self.unique_keys
53    }
54
55    fn expand_owned<T: Clone>(&self, values: Vec<Option<T>>) -> Vec<Option<T>> {
56        match &self.positions {
57            Some(positions) => positions
58                .iter()
59                .map(|&index| values[index].clone())
60                .collect(),
61            None => values,
62        }
63    }
64}
65
66fn cid_from_store_key(key: &[u8], context: &str) -> Result<Cid, String> {
67    let bytes: [u8; 32] = key.try_into().map_err(|_| {
68        format!(
69            "{context} key has invalid CID length {}, expected 32",
70            key.len()
71        )
72    })?;
73    Ok(Cid(bytes))
74}
75
76fn sort_cids(cids: &mut [Cid]) {
77    cids.sort_by(|left, right| left.as_bytes().cmp(right.as_bytes()));
78}
79
80fn sort_named_root_manifests(roots: &mut [NamedRootManifest]) {
81    roots.sort_by(|left, right| left.name.cmp(&right.name));
82}
83
84const NODE_PREFIX: &[u8] = b"node:";
85const HINT_PREFIX: &[u8] = b"hint:";
86const ROOT_PREFIX: &[u8] = b"root:";
87
88/// Configuration options for [`SlateDbStore`].
89#[derive(Debug, Clone)]
90pub struct SlateDbStoreConfig {
91    /// SlateDB engine settings.
92    pub settings: Settings,
93    /// Flush writes to object storage before returning from write operations.
94    pub flush_after_write: bool,
95    /// Close the SlateDB instance when this store is dropped.
96    pub close_on_drop: bool,
97    /// Maximum number of concurrent reads used by `batch_get` operations.
98    pub read_parallelism: usize,
99}
100
101impl Default for SlateDbStoreConfig {
102    fn default() -> Self {
103        Self {
104            settings: Settings::default(),
105            flush_after_write: true,
106            close_on_drop: true,
107            read_parallelism: 64,
108        }
109    }
110}
111
112/// Error type for SlateDB store operations.
113#[derive(Debug)]
114pub struct SlateDbStoreError {
115    message: String,
116    source: Option<Box<dyn std::error::Error + Send + Sync>>,
117}
118
119impl SlateDbStoreError {
120    /// Create a new error with a message.
121    pub fn new(message: impl Into<String>) -> Self {
122        Self {
123            message: message.into(),
124            source: None,
125        }
126    }
127
128    /// Create a new error with a source error.
129    pub fn with_source(
130        message: impl Into<String>,
131        source: impl std::error::Error + Send + Sync + 'static,
132    ) -> Self {
133        Self {
134            message: message.into(),
135            source: Some(Box::new(source)),
136        }
137    }
138
139    fn from_slatedb(err: slatedb::Error, context: impl Into<String>) -> Self {
140        Self::with_source(format!("{}: {}", context.into(), err), err)
141    }
142}
143
144impl std::fmt::Display for SlateDbStoreError {
145    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
146        write!(f, "SlateDB error: {}", self.message)
147    }
148}
149
150impl std::error::Error for SlateDbStoreError {
151    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
152        self.source
153            .as_ref()
154            .map(|e| e.as_ref() as &(dyn std::error::Error + 'static))
155    }
156}
157
158impl From<slatedb::Error> for SlateDbStoreError {
159    fn from(err: slatedb::Error) -> Self {
160        Self::from_slatedb(err, "SlateDB operation failed")
161    }
162}
163
164/// SlateDB-backed storage backend for Prolly Trees.
165///
166/// SlateDB is async-first; this adapter owns a private Tokio runtime so it can
167/// implement the synchronous [`Store`] trait used by the rest of the prolly
168/// tree engine.
169pub struct SlateDbStore {
170    db: Db,
171    runtime: Runtime,
172    manifest_lock: Mutex<()>,
173    flush_after_write: bool,
174    close_on_drop: bool,
175    read_parallelism: usize,
176}
177
178impl SlateDbStore {
179    /// Open or create a SlateDB database at `path` in the provided object store.
180    pub fn open(
181        path: impl Into<String>,
182        object_store: Arc<dyn ObjectStore>,
183    ) -> Result<Self, SlateDbStoreError> {
184        Self::open_with_config(path, object_store, SlateDbStoreConfig::default())
185    }
186
187    /// Open or create a SlateDB database with custom configuration.
188    pub fn open_with_config(
189        path: impl Into<String>,
190        object_store: Arc<dyn ObjectStore>,
191        config: SlateDbStoreConfig,
192    ) -> Result<Self, SlateDbStoreError> {
193        let runtime = Builder::new_multi_thread()
194            .thread_name("prolly-slatedb")
195            .enable_all()
196            .build()
197            .map_err(|e| SlateDbStoreError::with_source("failed to create Tokio runtime", e))?;
198
199        let path = path.into();
200        let settings = config.settings;
201        let db = runtime
202            .block_on(async move {
203                Db::builder(path, object_store)
204                    .with_settings(settings)
205                    .build()
206                    .await
207            })
208            .map_err(|e| SlateDbStoreError::from_slatedb(e, "failed to open SlateDB"))?;
209
210        Ok(Self {
211            db,
212            runtime,
213            manifest_lock: Mutex::new(()),
214            flush_after_write: config.flush_after_write,
215            close_on_drop: config.close_on_drop,
216            read_parallelism: config.read_parallelism.max(1),
217        })
218    }
219
220    /// Flush outstanding writes to object storage.
221    pub fn flush(&self) -> Result<(), SlateDbStoreError> {
222        self.block_on(self.db.flush(), "failed to flush SlateDB")
223    }
224
225    fn block_on<T, F>(&self, future: F, context: &'static str) -> Result<T, SlateDbStoreError>
226    where
227        F: Future<Output = Result<T, slatedb::Error>>,
228    {
229        self.runtime
230            .block_on(future)
231            .map_err(|e| SlateDbStoreError::from_slatedb(e, context))
232    }
233
234    fn flush_after_write_if_configured(&self) -> Result<(), SlateDbStoreError> {
235        if self.flush_after_write {
236            self.flush()?;
237        }
238        Ok(())
239    }
240
241    fn write_batch(
242        &self,
243        batch: WriteBatch,
244        context: &'static str,
245    ) -> Result<(), SlateDbStoreError> {
246        self.block_on(async { self.db.write(batch).await.map(|_| ()) }, context)?;
247        self.flush_after_write_if_configured()
248    }
249
250    fn batch_read_ordered(
251        &self,
252        storage_keys: Vec<Vec<u8>>,
253        context: &'static str,
254    ) -> Result<Vec<Option<Vec<u8>>>, SlateDbStoreError> {
255        let len = storage_keys.len();
256        if len == 0 {
257            return Ok(Vec::new());
258        }
259
260        if len == 1 {
261            let key = storage_keys.into_iter().next().expect("one key");
262            return self
263                .block_on(async { self.db.get(key).await }, context)
264                .map(|value| vec![value.map(|bytes| bytes.to_vec())]);
265        }
266
267        let db = self.db.clone();
268        let parallelism = self.read_parallelism;
269
270        self.block_on(
271            async move {
272                let indexed_values = stream::iter(storage_keys.into_iter().enumerate())
273                    .map(|(idx, key)| {
274                        let db = db.clone();
275                        async move {
276                            db.get(key)
277                                .await
278                                .map(|value| (idx, value.map(|bytes| bytes.to_vec())))
279                        }
280                    })
281                    .buffer_unordered(parallelism)
282                    .collect::<Vec<_>>()
283                    .await;
284
285                let mut ordered = vec![None; len];
286                for result in indexed_values {
287                    let (idx, value) = result?;
288                    ordered[idx] = Some(value);
289                }
290
291                Ok(ordered
292                    .into_iter()
293                    .map(|value| value.expect("all batch reads must fill their result slot"))
294                    .collect())
295            },
296            context,
297        )
298    }
299}
300
301impl Drop for SlateDbStore {
302    fn drop(&mut self) {
303        if self.close_on_drop {
304            let _ = self.runtime.block_on(self.db.close());
305        }
306    }
307}
308
309impl Store for SlateDbStore {
310    type Error = SlateDbStoreError;
311
312    fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
313        let key = node_key(key);
314        self.block_on(async { self.db.get(key).await }, "failed to read key")
315            .map(|value| value.map(|bytes| bytes.to_vec()))
316    }
317
318    fn put(&self, key: &[u8], value: &[u8]) -> Result<(), Self::Error> {
319        let key = Bytes::from(node_key(key));
320        let value = Bytes::copy_from_slice(value);
321        self.block_on(
322            async { self.db.put_bytes(key, value).await.map(|_| ()) },
323            "failed to write key",
324        )?;
325        self.flush_after_write_if_configured()
326    }
327
328    fn delete(&self, key: &[u8]) -> Result<(), Self::Error> {
329        let key = node_key(key);
330        self.block_on(
331            async { self.db.delete(key).await.map(|_| ()) },
332            "failed to delete key",
333        )?;
334        self.flush_after_write_if_configured()
335    }
336
337    fn batch(&self, ops: &[BatchOp]) -> Result<(), Self::Error> {
338        if ops.is_empty() {
339            return Ok(());
340        }
341
342        let mut batch = WriteBatch::new();
343        for op in ops {
344            match op {
345                BatchOp::Upsert { key, value } => {
346                    batch.put(node_key(key), value);
347                }
348                BatchOp::Delete { key } => {
349                    batch.delete(node_key(key));
350                }
351            }
352        }
353
354        self.write_batch(batch, "batch operation failed")
355    }
356
357    fn batch_get(&self, keys: &[&[u8]]) -> Result<HashMap<Vec<u8>, Vec<u8>>, Self::Error> {
358        if keys.is_empty() {
359            return Ok(HashMap::new());
360        }
361
362        if keys.len() == 1 {
363            let mut results = HashMap::with_capacity(1);
364            if let Some(value) = self.get(keys[0])? {
365                results.insert(keys[0].to_vec(), value);
366            }
367            return Ok(results);
368        }
369
370        let plan = OrderedBatchReadPlan::new(keys);
371        let storage_keys = plan
372            .unique_keys()
373            .iter()
374            .map(|key| node_key(key))
375            .collect::<Vec<_>>();
376        let values = self.batch_read_ordered(storage_keys, "failed to read keys in batch")?;
377
378        let mut results = HashMap::with_capacity(plan.unique_keys().len());
379        for (key, value) in plan.unique_keys().iter().zip(values) {
380            if let Some(value) = value {
381                results.insert(key.to_vec(), value);
382            }
383        }
384        Ok(results)
385    }
386
387    fn batch_get_ordered(&self, keys: &[&[u8]]) -> Result<Vec<Option<Vec<u8>>>, Self::Error> {
388        if keys.is_empty() {
389            return Ok(Vec::new());
390        }
391
392        if keys.len() == 1 {
393            return Ok(vec![self.get(keys[0])?]);
394        }
395
396        let plan = OrderedBatchReadPlan::new(keys);
397        let storage_keys = plan
398            .unique_keys()
399            .iter()
400            .map(|key| node_key(key))
401            .collect::<Vec<_>>();
402        let values =
403            self.batch_read_ordered(storage_keys, "failed to read keys in ordered batch")?;
404        Ok(plan.expand_owned(values))
405    }
406
407    fn batch_get_ordered_unique(
408        &self,
409        keys: &[&[u8]],
410    ) -> Result<Vec<Option<Vec<u8>>>, Self::Error> {
411        let storage_keys = keys.iter().map(|key| node_key(key)).collect::<Vec<_>>();
412        self.batch_read_ordered(storage_keys, "failed to read keys in unique ordered batch")
413    }
414
415    fn prefers_batch_reads(&self) -> bool {
416        true
417    }
418
419    fn batch_put(&self, entries: &[(&[u8], &[u8])]) -> Result<(), Self::Error> {
420        if entries.is_empty() {
421            return Ok(());
422        }
423
424        let mut batch = WriteBatch::new();
425        for (key, value) in entries {
426            batch.put(node_key(key), value);
427        }
428
429        self.write_batch(batch, "batch put operation failed")
430    }
431
432    fn supports_hints(&self) -> bool {
433        true
434    }
435
436    fn get_hint(&self, namespace: &[u8], key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
437        let key = hint_key(namespace, key);
438        self.block_on(async { self.db.get(key).await }, "failed to read hint")
439            .map(|value| value.map(|bytes| bytes.to_vec()))
440    }
441
442    fn put_hint(&self, namespace: &[u8], key: &[u8], value: &[u8]) -> Result<(), Self::Error> {
443        let key = Bytes::from(hint_key(namespace, key));
444        let value = Bytes::copy_from_slice(value);
445        self.block_on(
446            async { self.db.put_bytes(key, value).await.map(|_| ()) },
447            "failed to write hint",
448        )?;
449        self.flush_after_write_if_configured()
450    }
451
452    fn batch_put_with_hint(
453        &self,
454        entries: &[(&[u8], &[u8])],
455        namespace: &[u8],
456        key: &[u8],
457        value: &[u8],
458    ) -> Result<(), Self::Error> {
459        let mut batch = WriteBatch::new();
460        for (key, value) in entries {
461            batch.put(node_key(key), value);
462        }
463        batch.put(hint_key(namespace, key), value);
464
465        self.write_batch(batch, "batch put with hint operation failed")
466    }
467}
468
469impl NodeStoreScan for SlateDbStore {
470    type Error = SlateDbStoreError;
471
472    fn list_node_cids(&self) -> Result<Vec<Cid>, Self::Error> {
473        self.block_on(
474            async {
475                let mut iter = self.db.scan_prefix(NODE_PREFIX, ..).await?;
476                let mut cids = Vec::new();
477                while let Some(kv) = iter.next().await? {
478                    let key = kv.key.as_ref();
479                    let cid = key
480                        .strip_prefix(NODE_PREFIX)
481                        .ok_or_else(|| {
482                            slatedb::Error::invalid(
483                                "SlateDB node scan returned key without node prefix".to_string(),
484                            )
485                        })
486                        .and_then(|key| {
487                            cid_from_store_key(key, "SlateDB node").map_err(slatedb::Error::invalid)
488                        })?;
489                    cids.push(cid);
490                }
491                sort_cids(&mut cids);
492                Ok(cids)
493            },
494            "failed to list node CIDs",
495        )
496    }
497}
498
499impl ManifestStore for SlateDbStore {
500    type Error = SlateDbStoreError;
501
502    fn get_root(&self, name: &[u8]) -> Result<Option<RootManifest>, Self::Error> {
503        let key = root_key(name);
504        let bytes = self
505            .block_on(
506                async { self.db.get(key).await },
507                "failed to read root manifest",
508            )?
509            .map(|bytes| bytes.to_vec());
510        decode_root_manifest(bytes)
511    }
512
513    fn put_root(&self, name: &[u8], manifest: &RootManifest) -> Result<(), Self::Error> {
514        let _guard = self
515            .manifest_lock
516            .lock()
517            .map_err(|e| SlateDbStoreError::new(format!("manifest lock poisoned: {e}")))?;
518
519        let key = Bytes::from(root_key(name));
520        let bytes = encode_root_manifest(manifest)?;
521        let value = Bytes::from(bytes);
522        self.block_on(
523            async { self.db.put_bytes(key, value).await.map(|_| ()) },
524            "failed to write root manifest",
525        )?;
526        self.flush_after_write_if_configured()
527    }
528
529    fn delete_root(&self, name: &[u8]) -> Result<(), Self::Error> {
530        let _guard = self
531            .manifest_lock
532            .lock()
533            .map_err(|e| SlateDbStoreError::new(format!("manifest lock poisoned: {e}")))?;
534
535        let key = root_key(name);
536        self.block_on(
537            async { self.db.delete(key).await.map(|_| ()) },
538            "failed to delete root manifest",
539        )?;
540        self.flush_after_write_if_configured()
541    }
542
543    fn compare_and_swap_root(
544        &self,
545        name: &[u8],
546        expected: Option<&RootManifest>,
547        new: Option<&RootManifest>,
548    ) -> Result<ManifestUpdate, Self::Error> {
549        let _guard = self
550            .manifest_lock
551            .lock()
552            .map_err(|e| SlateDbStoreError::new(format!("manifest lock poisoned: {e}")))?;
553
554        let key = root_key(name);
555        let current_bytes = self
556            .block_on(
557                async { self.db.get(key.clone()).await },
558                "failed to read root manifest",
559            )?
560            .map(|bytes| bytes.to_vec());
561        let current = decode_root_manifest(current_bytes)?;
562        if current.as_ref() != expected {
563            return Ok(ManifestUpdate::Conflict { current });
564        }
565
566        match new {
567            Some(manifest) => {
568                let key = Bytes::from(key);
569                let value = Bytes::from(encode_root_manifest(manifest)?);
570                self.block_on(
571                    async { self.db.put_bytes(key, value).await.map(|_| ()) },
572                    "failed to write root manifest",
573                )?;
574            }
575            None => {
576                self.block_on(
577                    async { self.db.delete(key).await.map(|_| ()) },
578                    "failed to delete root manifest",
579                )?;
580            }
581        }
582
583        self.flush_after_write_if_configured()?;
584        Ok(ManifestUpdate::Applied)
585    }
586}
587
588impl ManifestStoreScan for SlateDbStore {
589    fn list_roots(&self) -> Result<Vec<NamedRootManifest>, Self::Error> {
590        let raw_roots = self.block_on(
591            async {
592                let mut iter = self.db.scan_prefix(ROOT_PREFIX, ..).await?;
593                let mut roots = Vec::new();
594                while let Some(kv) = iter.next().await? {
595                    let key = kv.key.as_ref();
596                    let name = key
597                        .strip_prefix(ROOT_PREFIX)
598                        .ok_or_else(|| {
599                            slatedb::Error::invalid(
600                                "SlateDB root scan returned key without root prefix".to_string(),
601                            )
602                        })?
603                        .to_vec();
604                    roots.push((name, kv.value.to_vec()));
605                }
606                Ok(roots)
607            },
608            "failed to list root manifests",
609        )?;
610
611        let mut roots = raw_roots
612            .into_iter()
613            .map(|(name, bytes)| {
614                let manifest = RootManifest::from_bytes(&bytes)
615                    .map_err(|err| SlateDbStoreError::new(err.to_string()))?;
616                Ok(NamedRootManifest::new(name, manifest))
617            })
618            .collect::<Result<Vec<_>, SlateDbStoreError>>()?;
619        sort_named_root_manifests(&mut roots);
620        Ok(roots)
621    }
622}
623
624impl TransactionalStore for SlateDbStore {
625    fn supports_transactions(&self) -> bool {
626        true
627    }
628
629    fn commit_transaction(
630        &self,
631        node_writes: &[TransactionNodeWrite],
632        root_conditions: &[RootCondition],
633        root_writes: &[RootWrite],
634    ) -> Result<TransactionUpdate, Error> {
635        let _guard = self.manifest_lock.lock().map_err(|err| {
636            Error::Store(Box::new(SlateDbStoreError::new(format!(
637                "manifest lock poisoned: {err}"
638            ))))
639        })?;
640
641        for condition in root_conditions {
642            let key = root_key(&condition.name);
643            let current_bytes = self
644                .block_on(
645                    async { self.db.get(key).await },
646                    "failed to read root manifest during transaction commit",
647                )
648                .map_err(|err| Error::Store(Box::new(err)))?
649                .map(|bytes| bytes.to_vec());
650            let current =
651                decode_root_manifest(current_bytes).map_err(|err| Error::Store(Box::new(err)))?;
652            if current != condition.expected {
653                return Ok(TransactionUpdate::Conflict(Box::new(
654                    TransactionConflict::new(
655                        condition.name.clone(),
656                        condition.expected.clone(),
657                        current,
658                    ),
659                )));
660            }
661        }
662
663        let mut batch = WriteBatch::new();
664        for write in node_writes {
665            match write {
666                TransactionNodeWrite::Upsert { key, value } => batch.put(node_key(key), value),
667                TransactionNodeWrite::Delete { key } => batch.delete(node_key(key)),
668            }
669        }
670
671        for write in root_writes {
672            match write {
673                RootWrite::Put { name, manifest } => {
674                    let bytes = encode_root_manifest(manifest)
675                        .map_err(|err| Error::Store(Box::new(err)))?;
676                    batch.put(root_key(name), bytes);
677                }
678                RootWrite::Delete { name } => batch.delete(root_key(name)),
679            }
680        }
681
682        self.write_batch(batch, "failed to commit transaction")
683            .map_err(|err| Error::Store(Box::new(err)))?;
684        Ok(TransactionUpdate::Applied {
685            nodes_written: node_writes.len(),
686            roots_written: root_writes.len(),
687        })
688    }
689}
690
691fn node_key(key: &[u8]) -> Vec<u8> {
692    let mut storage_key = Vec::with_capacity(NODE_PREFIX.len() + key.len());
693    storage_key.extend_from_slice(NODE_PREFIX);
694    storage_key.extend_from_slice(key);
695    storage_key
696}
697
698fn hint_key(namespace: &[u8], key: &[u8]) -> Vec<u8> {
699    let mut storage_key = Vec::with_capacity(HINT_PREFIX.len() + 4 + namespace.len() + key.len());
700    storage_key.extend_from_slice(HINT_PREFIX);
701    storage_key.extend_from_slice(&(namespace.len() as u32).to_be_bytes());
702    storage_key.extend_from_slice(namespace);
703    storage_key.extend_from_slice(key);
704    storage_key
705}
706
707fn root_key(name: &[u8]) -> Vec<u8> {
708    let mut storage_key = Vec::with_capacity(ROOT_PREFIX.len() + name.len());
709    storage_key.extend_from_slice(ROOT_PREFIX);
710    storage_key.extend_from_slice(name);
711    storage_key
712}
713
714fn encode_root_manifest(manifest: &RootManifest) -> Result<Vec<u8>, SlateDbStoreError> {
715    manifest
716        .to_bytes()
717        .map_err(|e| SlateDbStoreError::new(format!("failed to encode root manifest: {e}")))
718}
719
720fn decode_root_manifest(bytes: Option<Vec<u8>>) -> Result<Option<RootManifest>, SlateDbStoreError> {
721    bytes
722        .as_deref()
723        .map(RootManifest::from_bytes)
724        .transpose()
725        .map_err(|e| SlateDbStoreError::new(format!("failed to decode root manifest: {e}")))
726}
727
728#[cfg(test)]
729mod tests {
730    use super::*;
731    use prolly::{Config, Prolly};
732
733    fn in_memory_store(path: &str) -> SlateDbStore {
734        let object_store: Arc<dyn ObjectStore> =
735            Arc::new(slatedb::object_store::memory::InMemory::new());
736        SlateDbStore::open(path, object_store).unwrap()
737    }
738
739    #[test]
740    fn slatedb_store_put_get_delete() {
741        let store = in_memory_store("test_put_get_delete");
742
743        store.put(b"key", b"value").unwrap();
744        assert_eq!(store.get(b"key").unwrap(), Some(b"value".to_vec()));
745
746        store.delete(b"key").unwrap();
747        assert_eq!(store.get(b"key").unwrap(), None);
748    }
749
750    #[test]
751    fn slatedb_store_batch_is_order_preserving_for_reads() {
752        let store = in_memory_store("test_batch_reads");
753        let ops = vec![
754            BatchOp::Upsert {
755                key: b"a",
756                value: b"1",
757            },
758            BatchOp::Upsert {
759                key: b"b",
760                value: b"2",
761            },
762            BatchOp::Upsert {
763                key: b"c",
764                value: b"3",
765            },
766        ];
767
768        store.batch(&ops).unwrap();
769
770        let keys: Vec<&[u8]> = vec![b"c", b"missing", b"a", b"c", b"missing", b"b"];
771        assert_eq!(
772            store.batch_get_ordered(&keys).unwrap(),
773            vec![
774                Some(b"3".to_vec()),
775                None,
776                Some(b"1".to_vec()),
777                Some(b"3".to_vec()),
778                None,
779                Some(b"2".to_vec())
780            ]
781        );
782    }
783
784    #[test]
785    fn slatedb_store_fast_paths_empty_and_single_batch_reads() {
786        let store = in_memory_store("test_empty_single_batch_reads");
787        let empty: Vec<&[u8]> = Vec::new();
788
789        assert_eq!(store.batch_get_ordered(&empty).unwrap(), Vec::new());
790        assert!(store.batch_get(&empty).unwrap().is_empty());
791
792        store.put(b"a", b"1").unwrap();
793        let existing: Vec<&[u8]> = vec![b"a"];
794        let missing: Vec<&[u8]> = vec![b"missing"];
795
796        assert_eq!(
797            store.batch_get_ordered(&existing).unwrap(),
798            vec![Some(b"1".to_vec())]
799        );
800        assert_eq!(store.batch_get_ordered(&missing).unwrap(), vec![None]);
801
802        let values = store.batch_get(&existing).unwrap();
803        assert_eq!(values.get(b"a".as_slice()), Some(&b"1".to_vec()));
804        assert!(store.batch_get(&missing).unwrap().is_empty());
805    }
806
807    #[test]
808    fn slatedb_store_persists_hints_separately_from_nodes() {
809        let store = in_memory_store("test_hints");
810
811        store.put_hint(b"rightmost", b"root", b"hint-v1").unwrap();
812        assert_eq!(
813            store.get_hint(b"rightmost", b"root").unwrap(),
814            Some(b"hint-v1".to_vec())
815        );
816        assert_eq!(store.get_hint(b"rightmost", b"missing").unwrap(), None);
817        assert_eq!(store.get(b"root").unwrap(), None);
818
819        store.put_hint(b"rightmost", b"root", b"hint-v2").unwrap();
820        assert_eq!(
821            store.get_hint(b"rightmost", b"root").unwrap(),
822            Some(b"hint-v2".to_vec())
823        );
824    }
825
826    #[test]
827    fn slatedb_store_reopens_from_same_object_store() {
828        let object_store: Arc<dyn ObjectStore> =
829            Arc::new(slatedb::object_store::memory::InMemory::new());
830        let path = "test_reopen";
831        {
832            let store = SlateDbStore::open(path, object_store.clone()).unwrap();
833            store.put(b"key", b"value").unwrap();
834        }
835
836        let store = SlateDbStore::open(path, object_store).unwrap();
837        assert_eq!(store.get(b"key").unwrap(), Some(b"value".to_vec()));
838    }
839
840    #[test]
841    fn slatedb_store_supports_prolly_tree_round_trip() {
842        let store = in_memory_store("test_prolly_round_trip");
843        let config = Config::default();
844        let prolly = Prolly::new(store, config);
845        let tree = prolly.create();
846        let tree = prolly
847            .put(&tree, b"name".to_vec(), b"Alice".to_vec())
848            .unwrap();
849        let tree = prolly.put(&tree, b"age".to_vec(), b"30".to_vec()).unwrap();
850
851        assert_eq!(prolly.get(&tree, b"name").unwrap(), Some(b"Alice".to_vec()));
852        assert_eq!(prolly.get(&tree, b"age").unwrap(), Some(b"30".to_vec()));
853    }
854}