chainseeker_server/
rocks_db_multi.rs

1/// C++'s `multimap`-like DB implementation backed by RocksDB.
2use crate::*;
3
4#[derive(Debug)]
5pub struct RocksDBMulti<K, V>
6    where K: Serialize + Deserialize + 'static,
7          V: Serialize + Deserialize + 'static + ConstantSize,
8{
9    db: RocksDB<K, Vec<V>>,
10}
11
12impl<K, V> RocksDBMulti<K, V>
13    where K: Serialize + Deserialize + 'static,
14          V: Serialize + Deserialize + 'static + ConstantSize + Clone + PartialEq,
15{
16    pub fn new(path: &str, temporary: bool) -> Self {
17        Self {
18            db: RocksDB::new(path, temporary),
19        }
20    }
21    pub fn get(&self, key: &K) -> Vec<V> {
22        self.db.get(key).unwrap_or_default()
23    }
24    pub fn put(&self, key: &K, values: &[V]) {
25        self.db.put(key, &values.to_vec());
26    }
27    pub fn push(&self, key: &K, value: V) {
28        let mut values = self.get(key);
29        values.push(value);
30        self.put(key, &values);
31    }
32    pub fn pop(&self, key: &K, value: &V) {
33        let values = self.get(key);
34        let values = values.iter().filter(|v| *v != value).cloned().collect::<Vec<V>>();
35        self.put(key, &values);
36    }
37}
38
39#[cfg(test)]
40mod tests {
41    use super::*;
42    #[test]
43    fn rocks_db_multi() {
44        let db = RocksDBMulti::<String, u32>::new("/tmp/chainseeker/test_rocks_db_multi", true);
45        let key1 = "bar".to_string();
46        let key2 = "foo".to_string();
47        db.push(&key1, 3939);
48        db.push(&key1, 4649);
49        db.push(&key2, 1234);
50        db.push(&key2, 5678);
51        assert_eq!(db.get(&key1), vec![3939, 4649]);
52        assert_eq!(db.get(&key2), vec![1234, 5678]);
53        db.pop(&key1, &3939);
54        assert_eq!(db.get(&key1), vec![4649]);
55    }
56}