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#[derive(Debug, Clone)]
90pub struct SlateDbStoreConfig {
91 pub settings: Settings,
93 pub flush_after_write: bool,
95 pub close_on_drop: bool,
97 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#[derive(Debug)]
114pub struct SlateDbStoreError {
115 message: String,
116 source: Option<Box<dyn std::error::Error + Send + Sync>>,
117}
118
119impl SlateDbStoreError {
120 pub fn new(message: impl Into<String>) -> Self {
122 Self {
123 message: message.into(),
124 source: None,
125 }
126 }
127
128 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
164pub 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 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 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 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}