gravitydb/
mem_kv_store.rs

1use crate::KVStore;
2use std::{collections::BTreeMap, str::Utf8Error};
3use thiserror::Error;
4use std::ops::Bound::Included;
5
6type HashId = String;
7
8/// A Backend for the graph database running in memory only.
9///
10/// Use this for testing purposes or if your data does not need to be
11/// persisted permanently.
12#[derive(Debug, Default)]
13pub struct MemoryKvStore {
14  data: BTreeMap<HashId, Vec<u8>>,
15}
16
17impl MemoryKvStore {
18  pub fn get_inner(self) -> BTreeMap<HashId, Vec<u8>> {
19    self.data
20  }
21}
22
23impl KVStore<Error> for MemoryKvStore
24{
25  fn create_bucket(&mut self, _key: &[u8]) -> Result<(), Error> {
26    Ok(())
27  }
28
29  fn delete_record(&mut self, key: &[u8]) -> Result<(), Error> {
30    self.data.remove(std::str::from_utf8(key)?);
31    Ok(())
32  }
33
34  fn store_record(&mut self, key: &[u8], value: &[u8]) -> Result<(), Error> {
35    self.data.insert(key_to_string(key)?, value.to_vec());
36    Ok(())
37  }
38
39  fn fetch_record(&self, key: &[u8]) -> Result<Vec<u8>, Error> {
40    let key = key_to_string(key)?;
41    let out = self.data.get(&key).ok_or(Error::Missing(key))?;
42    Ok(out.to_vec())
43  }
44
45  fn list_records(&self, from: &[u8], to: &[u8]) -> Result<Vec<Vec<u8>>, Error> {
46    let to = if to.len() != 0 {
47      key_to_string(to)?
48    } else {
49      let mut to: Vec<u8> = from.to_vec();
50      *to.last_mut().unwrap() += 1;
51      key_to_string(&to)?
52    };
53    let from = key_to_string(from)?;
54    let iter: Vec<Vec<u8>>  = self.data
55      .range((Included(from), Included(to)))
56      .map(|(k, _v)| k.as_bytes().to_vec())
57      .collect();
58    Ok(iter)
59  }
60
61  fn exists(&self, key: &[u8]) -> Result<bool, Error> {
62    Ok(self.data.contains_key(&key_to_string(key)?))
63  }
64}
65
66#[derive(Error, Debug)]
67pub enum Error {
68  #[error("the record {0} could not be found")]
69  Missing(String),
70  #[error(transparent)]
71  Utf8(#[from] Utf8Error),
72}
73
74fn key_to_string(key: &[u8]) -> Result<String, Error> {
75  Ok(std::str::from_utf8(key)?.to_string())
76}
77
gravitydb/mem_kv_store.rs

gravitydb/
mem_kv_store.rs
