Support `multi_get_*` methods (#572)

4 files changed, 138 insertions, 6 deletions

diff --git a/src/db.rs b/src/db.rs
index 0738b7d..9b673ec 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -145,6 +145,25 @@ pub trait DBAccess {
         key: K,
         readopts: &ReadOptions,
     ) -> Result<Option<Vec<u8>>, Error>;
+
+    fn multi_get_opt<K, I>(
+        &self,
+        keys: I,
+        readopts: &ReadOptions,
+    ) -> Vec<Result<Option<Vec<u8>>, Error>>
+    where
+        K: AsRef<[u8]>,
+        I: IntoIterator<Item = K>;
+
+    fn multi_get_cf_opt<'b, K, I, W>(
+        &self,
+        keys_cf: I,
+        readopts: &ReadOptions,
+    ) -> Vec<Result<Option<Vec<u8>>, Error>>
+    where
+        K: AsRef<[u8]>,
+        I: IntoIterator<Item = (&'b W, K)>,
+        W: AsColumnFamilyRef + 'b;
 }
 
 impl<T: ThreadMode> DBAccess for DBWithThreadMode<T> {
@@ -168,6 +187,31 @@ impl<T: ThreadMode> DBAccess for DBWithThreadMode<T> {
     ) -> Result<Option<Vec<u8>>, Error> {
         self.get_cf_opt(cf, key, readopts)
     }
+
+    fn multi_get_opt<K, I>(
+        &self,
+        keys: I,
+        readopts: &ReadOptions,
+    ) -> Vec<Result<Option<Vec<u8>>, Error>>
+    where
+        K: AsRef<[u8]>,
+        I: IntoIterator<Item = K>,
+    {
+        self.multi_get_opt(keys, readopts)
+    }
+
+    fn multi_get_cf_opt<'b, K, I, W>(
+        &self,
+        keys_cf: I,
+        readopts: &ReadOptions,
+    ) -> Vec<Result<Option<Vec<u8>>, Error>>
+    where
+        K: AsRef<[u8]>,
+        I: IntoIterator<Item = (&'b W, K)>,
+        W: AsColumnFamilyRef + 'b,
+    {
+        self.multi_get_cf_opt(keys_cf, readopts)
+    }
 }
 
 /// A type alias to DB instance type with the single-threaded column family
diff --git a/src/lib.rs b/src/lib.rs
index 1b2fba6..c84f834 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -98,7 +98,7 @@ pub use crate::{
         ColumnFamilyRef, DEFAULT_COLUMN_FAMILY_NAME,
     },
     compaction_filter::Decision as CompactionDecision,
-    db::{DBWithThreadMode, LiveFile, MultiThreaded, SingleThreaded, ThreadMode, DB},
+    db::{DBAccess, DBWithThreadMode, LiveFile, MultiThreaded, SingleThreaded, ThreadMode, DB},
     db_iterator::{
         DBIterator, DBIteratorWithThreadMode, DBRawIterator, DBRawIteratorWithThreadMode,
         DBWALIterator, Direction, IteratorMode,
diff --git a/src/snapshot.rs b/src/snapshot.rs
index a4356ca..0a0ea06 100644
--- a/src/snapshot.rs
+++ b/src/snapshot.rs
@@ -160,6 +160,55 @@ impl<'a, D: DBAccess> SnapshotWithThreadMode<'a, D> {
         readopts.set_snapshot(self);
         self.db.get_cf_opt(cf, key.as_ref(), &readopts)
     }
+
+    /// Returns the bytes associated with the given key values and default read options.
+    pub fn multi_get<K: AsRef<[u8]>, I>(&self, keys: I) -> Vec<Result<Option<Vec<u8>>, Error>>
+    where
+        I: IntoIterator<Item = K>,
+    {
+        let readopts = ReadOptions::default();
+        self.multi_get_opt(keys, readopts)
+    }
+
+    /// Returns the bytes associated with the given key values and default read options.
+    pub fn multi_get_cf<'b, K, I, W>(&self, keys_cf: I) -> Vec<Result<Option<Vec<u8>>, Error>>
+    where
+        K: AsRef<[u8]>,
+        I: IntoIterator<Item = (&'b W, K)>,
+        W: AsColumnFamilyRef + 'b,
+    {
+        let readopts = ReadOptions::default();
+        self.multi_get_cf_opt(keys_cf, readopts)
+    }
+
+    /// Returns the bytes associated with the given key values and given read options.
+    pub fn multi_get_opt<K, I>(
+        &self,
+        keys: I,
+        mut readopts: ReadOptions,
+    ) -> Vec<Result<Option<Vec<u8>>, Error>>
+    where
+        K: AsRef<[u8]>,
+        I: IntoIterator<Item = K>,
+    {
+        readopts.set_snapshot(self);
+        self.db.multi_get_opt(keys, &readopts)
+    }
+
+    /// Returns the bytes associated with the given key values, given column family and read options.
+    pub fn multi_get_cf_opt<'b, K, I, W>(
+        &self,
+        keys_cf: I,
+        mut readopts: ReadOptions,
+    ) -> Vec<Result<Option<Vec<u8>>, Error>>
+    where
+        K: AsRef<[u8]>,
+        I: IntoIterator<Item = (&'b W, K)>,
+        W: AsColumnFamilyRef + 'b,
+    {
+        readopts.set_snapshot(self);
+        self.db.multi_get_cf_opt(keys_cf, &readopts)
+    }
 }
 
 impl<'a, D: DBAccess> Drop for SnapshotWithThreadMode<'a, D> {
diff --git a/tests/test_db.rs b/tests/test_db.rs
index 394d05e..a362d46 100644
--- a/tests/test_db.rs
+++ b/tests/test_db.rs
@@ -20,7 +20,7 @@ use pretty_assertions::assert_eq;
 
 use rocksdb::{
     perf::get_memory_usage_stats, BlockBasedOptions, BottommostLevelCompaction, Cache,
-    CompactOptions, CuckooTableOptions, DBCompactionStyle, DBWithThreadMode, Env, Error,
+    CompactOptions, CuckooTableOptions, DBAccess, DBCompactionStyle, DBWithThreadMode, Env, Error,
     FifoCompactOptions, IteratorMode, MultiThreaded, Options, PerfContext, PerfMetric, ReadOptions,
     SingleThreaded, SliceTransform, Snapshot, UniversalCompactOptions,
     UniversalCompactionStopStyle, WriteBatch, DB,
@@ -926,20 +926,59 @@ fn multi_get() {
 
     {
         let db = DB::open_default(&path).unwrap();
+        let initial_snap = db.snapshot();
         db.put(b"k1", b"v1").unwrap();
+        let k1_snap = db.snapshot();
         db.put(b"k2", b"v2").unwrap();
 
         let _ = db.multi_get(&[b"k0"; 40]);
 
+        let assert_values = |values: Vec<_>| {
+            assert_eq!(3, values.len());
+            assert_eq!(values[0], None);
+            assert_eq!(values[1], Some(b"v1".to_vec()));
+            assert_eq!(values[2], Some(b"v2".to_vec()));
+        };
+
         let values = db
             .multi_get(&[b"k0", b"k1", b"k2"])
             .into_iter()
             .map(Result::unwrap)
             .collect::<Vec<_>>();
-        assert_eq!(3, values.len());
-        assert_eq!(values[0], None);
-        assert_eq!(values[1], Some(b"v1".to_vec()));
-        assert_eq!(values[2], Some(b"v2".to_vec()));
+
+        assert_values(values);
+
+        let values = DBAccess::multi_get_opt(&db, &[b"k0", b"k1", b"k2"], &Default::default())
+            .into_iter()
+            .map(Result::unwrap)
+            .collect::<Vec<_>>();
+
+        assert_values(values);
+
+        let values = db
+            .snapshot()
+            .multi_get(&[b"k0", b"k1", b"k2"])
+            .into_iter()
+            .map(Result::unwrap)
+            .collect::<Vec<_>>();
+
+        assert_values(values);
+
+        let none_values = initial_snap
+            .multi_get(&[b"k0", b"k1", b"k2"])
+            .into_iter()
+            .map(Result::unwrap)
+            .collect::<Vec<_>>();
+
+        assert_eq!(none_values, vec![None; 3]);
+
+        let k1_only = k1_snap
+            .multi_get(&[b"k0", b"k1", b"k2"])
+            .into_iter()
+            .map(Result::unwrap)
+            .collect::<Vec<_>>();
+
+        assert_eq!(k1_only, vec![None, Some(b"v1".to_vec()), None]);
     }
 }