db.rs (2072B)
1 use anyhow::Result; 2 use serde::{de::DeserializeOwned, Serialize}; 3 use sled::{Config as SledConfig, Db, IVec}; 4 use std::collections::HashMap; 5 use std::sync::Mutex; 6 7 use crate::utils::serde_utils::{deserialize, serialize}; 8 9 pub struct IndexerDb { 10 pub db: Db, 11 trees: Mutex<HashMap<String, sled::Tree>>, 12 } 13 14 impl IndexerDb { 15 pub fn open(path: &str) -> Result<Self> { 16 let db = SledConfig::new().path(path).open()?; 17 Ok(Self { 18 db, 19 trees: Mutex::new(HashMap::new()), 20 }) 21 } 22 23 fn tree(&self, name: &str) -> Result<sled::Tree> { 24 let mut trees = self.trees.lock().unwrap_or_else(|err| err.into_inner()); 25 if let Some(tree) = trees.get(name) { 26 return Ok(tree.clone()); 27 } 28 let tree = self.db.open_tree(name)?; 29 trees.insert(name.to_string(), tree.clone()); 30 Ok(tree) 31 } 32 33 pub fn insert<T: Serialize>(&self, tree: &str, key: &str, value: &T) -> Result<()> { 34 let t = self.tree(tree)?; 35 let blob: Vec<u8> = serialize(value)?; 36 t.insert(key, blob)?; 37 Ok(()) 38 } 39 40 pub fn get<T: DeserializeOwned>(&self, tree: &str, key: &str) -> Result<Option<T>> { 41 let t = self.tree(tree)?; 42 if let Some(bytes) = t.get(key)? { 43 let v: T = deserialize(&bytes)?; 44 Ok(Some(v)) 45 } else { 46 Ok(None) 47 } 48 } 49 50 pub fn get_all<T: DeserializeOwned>(&self, tree: &str) -> Result<Vec<T>> { 51 let t = self.tree(tree)?; 52 let mut out = Vec::new(); 53 for res in t.iter().values() { 54 let bytes = res?; 55 let v: T = deserialize(&bytes)?; 56 out.push(v); 57 } 58 Ok(out) 59 } 60 61 pub fn insert_raw(&self, tree: &str, key: &str, bytes: &[u8]) -> Result<()> { 62 self.tree(tree)?.insert(key, bytes)?; 63 Ok(()) 64 } 65 66 pub fn get_raw(&self, tree: &str, key: &str) -> Result<Option<IVec>> { 67 Ok(self.tree(tree)?.get(key)?) 68 } 69 70 pub fn flush(&self) -> Result<()> { 71 self.db.flush()?; 72 Ok(()) 73 } 74 }