tulip_sled_util/
lib.rs

1//! an embedded database using the sled framework
2//!
3use borsh::{BorshDeserialize, BorshSerialize};
4pub mod config;
5pub mod types;
6use anyhow::{anyhow, Result};
7use config::DbOpts;
8use sled::{IVec, Tree};
9use std::sync::Arc;
10
11use self::types::{DbKey, DbTrees};
12
13/// Database is the main embedded database object using the
14/// sled db
15#[derive(Clone)]
16pub struct Database {
17    db: sled::Db,
18}
19
20/// DbTree is a wrapper around the sled::Tree type providing
21/// convenience functions
22#[derive(Clone)]
23pub struct DbTree {
24    pub tree: Tree,
25}
26
27/// DbBatch is a wrapper around the sled::Batch type providing
28/// convenience functions
29#[derive(Default, Clone)]
30pub struct DbBatch {
31    batch: sled::Batch,
32    count: u64,
33}
34
35impl Database {
36    /// returns a new sled database
37    pub fn new(cfg: &DbOpts) -> Result<Arc<Self>> {
38        let sled_config: sled::Config = cfg.into();
39        let db = sled_config.open()?;
40        drop(sled_config);
41        Ok(Arc::new(Database { db }))
42    }
43    /// opens the given database tree
44    pub fn open_tree(self: &Arc<Self>, tree: DbTrees) -> Result<Arc<DbTree>> {
45        DbTree::open(&self.db, tree)
46    }
47    /// opens the given db tree, return a vector of (key, value)
48    pub fn list_values(self: &Arc<Self>, tree: DbTrees) -> Result<Vec<(IVec, IVec)>> {
49        let tree = self.open_tree(tree)?;
50        Ok(tree
51            .iter()
52            .filter_map(|entry| {
53                if let Ok((key, value)) = entry {
54                    Some((key, value))
55                } else {
56                    None
57                }
58            })
59            .collect())
60    }
61    /// flushes teh database
62    pub fn flush(self: &Arc<Self>) -> Result<usize> {
63        Ok(self.db.flush()?)
64    }
65    /// returns a clone of the inner database
66    pub fn inner(self: &Arc<Self>) -> sled::Db {
67        self.db.clone()
68    }
69    /// destroys all trees except the default tree
70    pub fn destroy(self: &Arc<Self>) {
71        const SLED_DEFAULT_TREE: &[u8] = b"__sled__default";
72        self.db
73            .tree_names()
74            .iter()
75            .filter(|tree_name| tree_name.as_ref().ne(SLED_DEFAULT_TREE))
76            .for_each(|tree_name| {
77                if let Err(err) = self.db.drop_tree(tree_name) {
78                    log::error!("failed to drop tree {:?}: {:#?}", tree_name.as_ref(), err);
79                }
80            });
81    }
82    pub fn get<K: AsRef<[u8]>>(&self, key: K) -> sled::Result<Option<sled::IVec>> {
83        self.db.get(key)
84    }
85    pub fn deserialize<K: AsRef<[u8]>, T>(&self, key: K) -> Result<T>
86    where
87        T: BorshDeserialize,
88    {
89        let value = self.get(key)?;
90        if let Some(value) = value {
91            Ok(borsh::de::BorshDeserialize::try_from_slice(&value)?)
92        } else {
93            Err(anyhow!("value for key is None"))
94        }
95    }
96    pub fn apply_batch(&self, batch: &mut DbBatch) -> sled::Result<()> {
97        self.db.apply_batch(batch.take_inner())
98    }
99    /// inserts a value into the default tree
100    pub fn insert<T>(&mut self, value: &T) -> Result<()>
101    where
102        T: BorshSerialize + DbKey,
103    {
104        self.db.insert(
105            value.key()?,
106            match borsh::to_vec(value) {
107                Ok(data) => data,
108                Err(err) => return Err(anyhow!("failed to insert entry into batch {:#?}", err)),
109            },
110        )?;
111        Ok(())
112    }
113}
114
115impl DbTree {
116    pub fn open(db: &sled::Db, tree: DbTrees) -> Result<Arc<Self>> {
117        let tree = db.open_tree(tree.str())?;
118        Ok(Arc::new(Self { tree }))
119    }
120    pub fn len(&self) -> usize {
121        self.tree.len()
122    }
123    pub fn is_empty(&self) -> bool {
124        self.tree.is_empty()
125    }
126    pub fn iter(&self) -> sled::Iter {
127        self.tree.iter()
128    }
129    pub fn contains_key<K: AsRef<[u8]>>(&self, key: K) -> sled::Result<bool> {
130        self.tree.contains_key(key)
131    }
132    pub fn flush(&self) -> sled::Result<usize> {
133        self.tree.flush()
134    }
135    pub fn apply_batch(&self, batch: &mut DbBatch) -> sled::Result<()> {
136        self.tree.apply_batch(batch.take_inner())
137    }
138    pub fn insert<T>(&self, value: &T) -> Result<Option<sled::IVec>>
139    where
140        T: BorshSerialize + DbKey,
141    {
142        Ok(self.tree.insert(
143            value.key()?,
144            match borsh::to_vec(value) {
145                Ok(data) => data,
146                Err(err) => return Err(anyhow!("failed to insert entry {:#?}", err)),
147            },
148        )?)
149    }
150    pub fn get<K: AsRef<[u8]>>(&self, key: K) -> sled::Result<Option<sled::IVec>> {
151        self.tree.get(key)
152    }
153    pub fn deserialize<K: AsRef<[u8]>, T>(&self, key: K) -> Result<T>
154    where
155        T: BorshDeserialize,
156    {
157        let value = self.get(key)?;
158        if let Some(value) = value {
159            Ok(borsh::de::BorshDeserialize::try_from_slice(&value)?)
160        } else {
161            Err(anyhow!("value for key is None"))
162        }
163    }
164}
165
166impl DbBatch {
167    pub fn new() -> DbBatch {
168        DbBatch {
169            batch: Default::default(),
170            count: 0,
171        }
172    }
173    pub fn insert<T>(&mut self, value: &T) -> Result<()>
174    where
175        T: BorshSerialize + DbKey,
176    {
177        self.batch.insert(
178            value.key()?,
179            match borsh::to_vec(value) {
180                Ok(data) => data,
181                Err(err) => return Err(anyhow!("failed to insert entry into batch {:#?}", err)),
182            },
183        );
184        self.count += 1;
185        Ok(())
186    }
187    /// returns the inner batch, and should only be used when the batch object
188    /// is finished with and the batch needs to be applied, as it replaces the inner
189    /// batch with its default version
190    pub fn take_inner(&mut self) -> sled::Batch {
191        std::mem::take(&mut self.batch)
192    }
193    pub fn inner(&self) -> &sled::Batch {
194        &self.batch
195    }
196    pub fn count(&self) -> u64 {
197        self.count
198    }
199}
200
201#[cfg(test)]
202mod test {
203    use super::*;
204    use borsh::{BorshDeserialize, BorshSchema, BorshSerialize};
205    use config::DbOpts;
206    use std::fs::remove_dir_all;
207
208    #[derive(BorshSerialize, BorshDeserialize, BorshSchema)]
209    pub struct TestData {
210        pub key: String,
211        pub foo: String,
212    }
213
214    impl DbKey for TestData {
215        fn key(&self) -> anyhow::Result<Vec<u8>> {
216            Ok(self.key.as_bytes().to_vec())
217        }
218    }
219
220    // performs very basic database testing
221    #[test]
222    fn test_db_basic() {
223        let db_opts = DbOpts::default();
224
225        let db = Database::new(&db_opts).unwrap();
226        let insert = || {
227            let mut db_batch = DbBatch::new();
228            db_batch
229                .insert(&TestData {
230                    key: "key1".to_string(),
231                    foo: "foo1".to_string(),
232                })
233                .unwrap();
234            {
235                let tree = db.open_tree(DbTrees::Custom("foobar")).unwrap();
236                tree.apply_batch(&mut db_batch).unwrap();
237                assert_eq!(tree.len(), 1);
238            }
239
240            db_batch
241                .insert(&TestData {
242                    key: "key2".to_string(),
243                    foo: "foo2".to_string(),
244                })
245                .unwrap();
246            {
247                let tree = db.open_tree(DbTrees::Custom("foobar")).unwrap();
248                tree.apply_batch(&mut db_batch).unwrap();
249                assert_eq!(tree.len(), 2);
250            }
251
252            db_batch
253                .insert(&TestData {
254                    key: "key3".to_string(),
255                    foo: "foo3".to_string(),
256                })
257                .unwrap();
258            {
259                let tree = db.open_tree(DbTrees::Custom("foobarbaz")).unwrap();
260                tree.apply_batch(&mut db_batch).unwrap();
261                assert_eq!(tree.len(), 1);
262            }
263            db_batch
264                .insert(&TestData {
265                    key: "key4".to_string(),
266                    foo: "foo4".to_string(),
267                })
268                .unwrap();
269            {
270                db.apply_batch(&mut db_batch).unwrap();
271            }
272        };
273        let query = || {
274            let foobar_values = db.list_values(DbTrees::Custom("foobar")).unwrap();
275            assert_eq!(foobar_values.len(), 2);
276            let test_data_one: TestData = db
277                .open_tree(DbTrees::Custom("foobar"))
278                .unwrap()
279                .deserialize(foobar_values[0].0.clone())
280                .unwrap();
281            assert_eq!(test_data_one.key, "key1".to_string());
282            assert_eq!(test_data_one.foo, "foo1".to_string());
283            let test_data_two: TestData = db
284                .open_tree(DbTrees::Custom("foobar"))
285                .unwrap()
286                .deserialize(foobar_values[1].0.clone())
287                .unwrap();
288            assert_eq!(test_data_two.key, "key2".to_string());
289            assert_eq!(test_data_two.foo, "foo2".to_string());
290            let foobarbaz_values = db.list_values(DbTrees::Custom("foobarbaz")).unwrap();
291            assert_eq!(foobarbaz_values.len(), 1);
292            let test_data_three: TestData = db
293                .open_tree(DbTrees::Custom("foobarbaz"))
294                .unwrap()
295                .deserialize(foobarbaz_values[0].0.clone())
296                .unwrap();
297            assert_eq!(test_data_three.key, "key3".to_string());
298            assert_eq!(test_data_three.foo, "foo3".to_string());
299            let default_tree_values = db.list_values(DbTrees::Default).unwrap();
300            assert_eq!(default_tree_values.len(), 1);
301        };
302        insert();
303        query();
304        db.destroy();
305        remove_dir_all("test_infos.db").unwrap();
306    }
307}
tulip_sled_util/lib.rs

tulip_sled_util/
lib.rs
