Complete sled abstraction

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<dyn std::iter::Iterator<Item = Result<(Value<'a>, Value<'a>)>> + Send + Sync + 'a>;
 
+pub type Exporter<'a> =
+	Box<dyn std::iter::Iterator<Item=Result<(String, ValueIter<'a>)>> + Send + Sync + 'a>;
+
 // ----
 
 #[derive(Debug, Error)]
@@ -95,7 +98,6 @@ impl Tree {
 	pub fn get<'a, T: AsRef<[u8]>>(&'a self, key: T) -> Result<Option<Value<'a>>> {
 		self.0.get(self.1, key.as_ref())
 	}
-
 	pub fn len(&self) -> Result<usize> {
 		self.0.len(self.1)
 	}
@@ -103,7 +105,6 @@ impl Tree {
 	pub fn insert<T: AsRef<[u8]>, 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<bool> {
 		self.0.remove(self.1, key.as_ref())
 	}
@@ -139,6 +140,9 @@ impl<'a> Transaction<'a> {
 	pub fn get<T: AsRef<[u8]>>(&self, tree: &Tree, key: T) -> Result<Option<Value<'a>>> {
 		self.0.get(tree.1, key.as_ref())
 	}
+	pub fn len(&self, tree: &Tree) -> Result<usize> {
+		self.0.len(tree.1)
+	}
 
 	pub fn insert<T: AsRef<[u8]>, 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<T: AsRef<[u8]>>(&self, tree: &Tree, key: T) -> Result<bool> {
 		self.0.remove(tree.1, key.as_ref())
 	}
 
+	pub fn iter(&self, tree: &Tree) -> Result<ValueIter<'a>> {
+		self.0.iter(tree.1)
+	}
+	pub fn iter_rev(&self, tree: &Tree) -> Result<ValueIter<'a>> {
+		self.0.iter_rev(tree.1)
+	}
+
+	pub fn range<K, R>(&self, tree: &Tree, range: R) -> Result<ValueIter<'a>>
+	where
+		K: AsRef<[u8]>,
+		R: RangeBounds<K>,
+	{
+		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<K, R>(&self, tree: &Tree, range: R) -> Result<ValueIter<'a>>
+	where
+		K: AsRef<[u8]>,
+		R: RangeBounds<K>,
+	{
+		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<R, E>(self, e: E) -> TxResult<R, E>
 	where
@@ -204,8 +235,26 @@ pub(crate) trait IDb: Send + Sync {
 
 pub(crate) trait ITx<'a> {
 	fn get(&self, tree: usize, key: &[u8]) -> Result<Option<Value<'a>>>;
+	fn len(&self, tree: usize) -> Result<usize>;
+
 	fn insert(&self, tree: usize, key: &[u8], value: &[u8]) -> Result<()>;
 	fn remove(&self, tree: usize, key: &[u8]) -> Result<bool>;
+
+	fn iter(&self, tree: usize) -> Result<ValueIter<'a>>;
+	fn iter_rev(&self, tree: usize) -> Result<ValueIter<'a>>;
+
+	fn range<'r>(
+		&self,
+		tree: usize,
+		low: Bound<&'r [u8]>,
+		high: Bound<&'r [u8]>,
+	) -> Result<ValueIter<'a>>;
+	fn range_rev<'r>(
+		&self,
+		tree: usize,
+		low: Bound<&'r [u8]>,
+		high: Bound<&'r [u8]>,
+	) -> Result<ValueIter<'a>>;
 }
 
 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<Exporter<'a>> {
+		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<R>(&self, v: std::result::Result<R, UnabortableTransactionError>) -> Result<R> {
 		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<Option<Value<'a>>> {
-		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<usize> {
+		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<bool> {
-		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<ValueIter<'a>> {
+		unimplemented!("Iterators in transactions not supported with Sled backend");
+	}
+	fn iter_rev(&self, _tree: usize) -> Result<ValueIter<'a>> {
+		unimplemented!("Iterators in transactions not supported with Sled backend");
+	}
+
+	fn range<'r>(
+		&self,
+		_tree: usize,
+		_low: Bound<&'r [u8]>,
+		_high: Bound<&'r [u8]>,
+	) -> Result<ValueIter<'a>> {
+		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<ValueIter<'a>> {
+		unimplemented!("Iterators in transactions not supported with Sled backend");
+	}
 }