From fbd5b64ff3c0da95a2ae96b1f830d2ed7dc4c5a8 Mon Sep 17 00:00:00 2001 From: Alex Auvolat Date: Thu, 2 Jun 2022 17:38:30 +0200 Subject: [PATCH] Complete sled abstraction --- src/db/lib.rs | 55 ++++++++++++++++++++++-- src/db/sled_adapter.rs | 95 +++++++++++++++++++++++++++++++++++------- 2 files changed, 133 insertions(+), 17 deletions(-) diff --git a/src/db/lib.rs b/src/db/lib.rs index 75c6ffa2..c3e89ba4 100644 --- a/src/db/lib.rs +++ b/src/db/lib.rs @@ -24,6 +24,9 @@ pub type Value<'a> = Cow<'a, [u8]>; pub type ValueIter<'a> = Box, Value<'a>)>> + Send + Sync + 'a>; +pub type Exporter<'a> = + Box)>> + Send + Sync + 'a>; + // ---- #[derive(Debug, Error)] @@ -95,7 +98,6 @@ impl Tree { pub fn get<'a, T: AsRef<[u8]>>(&'a self, key: T) -> Result>> { self.0.get(self.1, key.as_ref()) } - pub fn len(&self) -> Result { self.0.len(self.1) } @@ -103,7 +105,6 @@ impl Tree { pub fn insert, U: AsRef<[u8]>>(&self, key: T, value: U) -> Result<()> { self.0.insert(self.1, key.as_ref(), value.as_ref()) } - pub fn remove<'a, T: AsRef<[u8]>>(&'a self, key: T) -> Result { self.0.remove(self.1, key.as_ref()) } @@ -139,6 +140,9 @@ impl<'a> Transaction<'a> { pub fn get>(&self, tree: &Tree, key: T) -> Result>> { self.0.get(tree.1, key.as_ref()) } + pub fn len(&self, tree: &Tree) -> Result { + self.0.len(tree.1) + } pub fn insert, U: AsRef<[u8]>>( &self, @@ -148,11 +152,38 @@ impl<'a> Transaction<'a> { ) -> Result<()> { self.0.insert(tree.1, key.as_ref(), value.as_ref()) } - pub fn remove>(&self, tree: &Tree, key: T) -> Result { self.0.remove(tree.1, key.as_ref()) } + pub fn iter(&self, tree: &Tree) -> Result> { + self.0.iter(tree.1) + } + pub fn iter_rev(&self, tree: &Tree) -> Result> { + self.0.iter_rev(tree.1) + } + + pub fn range(&self, tree: &Tree, range: R) -> Result> + where + K: AsRef<[u8]>, + R: RangeBounds, + { + let sb = range.start_bound(); + let eb = range.end_bound(); + self.0.range(tree.1, get_bound(sb), get_bound(eb)) + } + pub fn range_rev(&self, tree: &Tree, range: R) -> Result> + where + K: AsRef<[u8]>, + R: RangeBounds, + { + let sb = range.start_bound(); + let eb = range.end_bound(); + self.0.range_rev(tree.1, get_bound(sb), get_bound(eb)) + } + + // ---- + #[must_use] pub fn abort(self, e: E) -> TxResult where @@ -204,8 +235,26 @@ pub(crate) trait IDb: Send + Sync { pub(crate) trait ITx<'a> { fn get(&self, tree: usize, key: &[u8]) -> Result>>; + fn len(&self, tree: usize) -> Result; + fn insert(&self, tree: usize, key: &[u8], value: &[u8]) -> Result<()>; fn remove(&self, tree: usize, key: &[u8]) -> Result; + + fn iter(&self, tree: usize) -> Result>; + fn iter_rev(&self, tree: usize) -> Result>; + + fn range<'r>( + &self, + tree: usize, + low: Bound<&'r [u8]>, + high: Bound<&'r [u8]>, + ) -> Result>; + fn range_rev<'r>( + &self, + tree: usize, + low: Bound<&'r [u8]>, + high: Bound<&'r [u8]>, + ) -> Result>; } pub(crate) trait ITxFn { diff --git a/src/db/sled_adapter.rs b/src/db/sled_adapter.rs index cf69caba..9ee9ea58 100644 --- a/src/db/sled_adapter.rs +++ b/src/db/sled_adapter.rs @@ -10,7 +10,7 @@ use sled::transaction::{ UnabortableTransactionError, }; -use crate::{Db, Error, IDb, ITx, ITxFn, Result, TxError, TxFnResult, TxResult, Value, ValueIter}; +use crate::{Db, Error, IDb, ITx, ITxFn, Result, TxError, TxFnResult, TxResult, Value, ValueIter, Exporter}; pub use sled; @@ -43,6 +43,50 @@ impl SledDb { .cloned() .ok_or(Error("invalid tree id".into())) } + + pub fn export<'a>(&'a self) -> Result> { + let mut trees = vec![]; + for name in self.db.tree_names() { + let name = std::str::from_utf8(&name).map_err(|e| Error(format!("{}", e).into()))?.to_string(); + let tree = self.open_tree(&name)?; + let tree = self.trees.read().unwrap().0.get(tree).unwrap().clone(); + trees.push((name, tree)); + } + let trees_exporter: Exporter<'a> = Box::new(trees + .into_iter() + .map(|(name, tree)| { + let iter: ValueIter<'a> = Box::new(tree.iter().map(|v| { + v.map(|(x, y)| (x.to_vec().into(), y.to_vec().into())) + .map_err(Into::into) + })); + Ok((name.to_string(), iter)) + })); + Ok(trees_exporter) + } + + pub fn import<'a>(&self, ex: Exporter<'a>) -> Result<()> { + for ex_tree in ex { + let (name, data) = ex_tree?; + + let tree = self.open_tree(&name)?; + let tree = self.trees.read().unwrap().0.get(tree).unwrap().clone(); + if !tree.is_empty() { + return Err(Error(format!("tree {} already contains data", name).into())); + } + + let mut i = 0; + for item in data { + let (k, v) = item?; + tree.insert(k.as_ref(), v.as_ref())?; + i += 1; + if i % 1000 == 0 { + println!("{}: imported {}", name, i); + } + } + println!("{}: finished importing, {} items", name, i); + } + Ok(()) + } } impl IDb for SledDb { @@ -165,6 +209,12 @@ struct SledTx<'a> { } impl<'a> SledTx<'a> { + fn get_tree(&self, i: usize) -> Result<&TransactionalTree> { + self.trees + .get(i) + .ok_or(Error("invalid tree id (it might have been openned after the transaction started)".into())) + } + fn save_error(&self, v: std::result::Result) -> Result { match v { Ok(x) => Ok(x), @@ -179,28 +229,45 @@ impl<'a> SledTx<'a> { impl<'a> ITx<'a> for SledTx<'a> { fn get(&self, tree: usize, key: &[u8]) -> Result>> { - let tree = self - .trees - .get(tree) - .ok_or(Error("invalid tree id".into()))?; + let tree = self.get_tree(tree)?; let tmp = self.save_error(tree.get(key))?; Ok(tmp.map(|v| v.to_vec().into())) } + fn len(&self, _tree: usize) -> Result { + unimplemented!(".len() in transaction not supported with Sled backend") + } fn insert(&self, tree: usize, key: &[u8], value: &[u8]) -> Result<()> { - let tree = self - .trees - .get(tree) - .ok_or(Error("invalid tree id".into()))?; + let tree = self.get_tree(tree)?; self.save_error(tree.insert(key, value))?; Ok(()) } - fn remove(&self, tree: usize, key: &[u8]) -> Result { - let tree = self - .trees - .get(tree) - .ok_or(Error("invalid tree id".into()))?; + let tree = self.get_tree(tree)?; Ok(self.save_error(tree.remove(key))?.is_some()) } + + fn iter(&self, _tree: usize) -> Result> { + unimplemented!("Iterators in transactions not supported with Sled backend"); + } + fn iter_rev(&self, _tree: usize) -> Result> { + unimplemented!("Iterators in transactions not supported with Sled backend"); + } + + fn range<'r>( + &self, + _tree: usize, + _low: Bound<&'r [u8]>, + _high: Bound<&'r [u8]>, + ) -> Result> { + unimplemented!("Iterators in transactions not supported with Sled backend"); + } + fn range_rev<'r>( + &self, + _tree: usize, + _low: Bound<&'r [u8]>, + _high: Bound<&'r [u8]>, + ) -> Result> { + unimplemented!("Iterators in transactions not supported with Sled backend"); + } }