treehouse/
lib.rs

1/// Error management
2pub mod err {
3
4    /// Error container
5    #[derive(thiserror::Error, Debug)]
6    pub enum Error {
7        #[error("I/O error: `{0}`")]
8        Io(#[from] std::io::Error),
9        #[error("Database error: `{0}`")]
10        Sled(#[from] sled::Error),
11        #[cfg(feature = "bincode")]
12        #[error("De/serialization error: `{0}`")]
13        Bincode(#[from] bincode::Error),
14        #[cfg(feature = "serde_cbor")]
15        #[error("De/serialization error: `{0}`")]
16        CBOR(#[from] serde_cbor::Error),
17        #[error("Error: `{0}`")]
18        Custom(Box<str>),
19    }
20
21    impl From<sled::transaction::TransactionError<Error>> for Error {
22        fn from(t: sled::transaction::TransactionError<Error>) -> Self {
23            match t {
24                sled::transaction::TransactionError::Abort(t) => t,
25                sled::transaction::TransactionError::Storage(t) => Error::Sled(t),
26            }
27        }
28    }
29
30    impl From<Error> for sled::transaction::ConflictableTransactionError<Error> {
31        fn from(t: Error) -> Self {
32            sled::transaction::ConflictableTransactionError::Abort(t)
33        }
34    }
35
36    pub fn custom<T: std::fmt::Display>(t: T) -> Error {
37        Error::Custom(t.to_string().into_boxed_str())
38    }
39
40    pub type Result<T> = std::result::Result<T, Error>;
41}
42
43mod utils {
44
45    use std::convert::TryInto;
46    use sled::transaction::TransactionalTree;
47
48    #[cfg(feature = "bincode")]
49    pub fn serialize<T: ?Sized + serde::Serialize>(value: &T) -> crate::err::Result<Vec<u8>> {
50        Ok(bincode::serialize(value)?)
51    }
52
53    #[cfg(feature = "bincode")]
54    pub fn deserialize<T: serde::de::DeserializeOwned>(bytes: &[u8]) -> crate::err::Result<T> {
55        Ok(bincode::deserialize(bytes)?)
56    }
57
58    pub fn u64_to_bytes(value: u64) -> [u8; 8] {
59        u64::to_be_bytes(value)
60    }
61
62    pub fn bytes_to_u64(value: &[u8]) -> crate::err::Result<u64> {
63        Ok(u64::from_be_bytes(value.try_into().map_err(crate::err::custom)?))
64    }
65}
66
67pub mod query;
68mod ext;
69
70use sled::transaction::Transactional;
71use sled::transaction::ConflictableTransactionError;
72use std::marker::PhantomData;
73use std::collections::BTreeSet;
74
75pub struct StoreBuilder<T> {
76    db: Option<sled::Db>,
77    tree_name: Option<Vec<u8>>,
78    meta_name: Option<Vec<u8>>,
79    marker: PhantomData<fn(T)>,
80}
81
82impl<T> Default for StoreBuilder<T> {
83    fn default() -> Self {
84        Self {
85            db: None,
86            tree_name: None,
87            meta_name: None,
88            marker: PhantomData,
89        }
90    }
91}
92
93impl<T> StoreBuilder<T> {
94    pub fn with_db(mut self, db: sled::Db) -> Self {
95        self.db = Some(db);
96        self
97    }
98
99    pub fn with_tree_name<S: Into<Vec<u8>>>(mut self, tree_name: S) -> Self {
100        self.tree_name = Some(tree_name.into());
101        self
102    }
103
104    pub fn with_meta_name<S: Into<Vec<u8>>>(mut self, meta_name: S) -> Self {
105        self.meta_name = Some(meta_name.into());
106        self
107    }
108
109    pub fn finish(self) -> err::Result<Store<T>> {
110        let Self { db, tree_name, meta_name, marker } = self;
111
112        let db = match db {
113            Some(db) => db,
114            None => { return Err(err::custom("No `db` provided")); }
115        };
116
117        let tree_name = match tree_name {
118            Some(tree_name) => tree_name,
119            None => { return Err(err::custom("No `tree_name` specified")); }
120        };
121
122        let meta_name = meta_name.unwrap_or_else(|| {
123            tree_name.iter().chain(b"\0meta").copied().collect()
124        });
125
126        Ok(Store {
127            tree: db.open_tree(&tree_name)?,
128            meta: db.open_tree(&meta_name)?,
129            db,
130            marker
131        })
132    }
133}
134
135pub struct Store<T> {
136    db: sled::Db,
137    tree: sled::Tree,
138    meta: sled::Tree,
139    marker: PhantomData<fn(T)>,
140}
141
142pub struct TransactionalStore<'a, T> {
143    tree: &'a sled::transaction::TransactionalTree,
144    meta: ext::TransactionalMeta<'a>,
145    marker: PhantomData<fn(T)>,
146}
147
148impl<'a, T: query::Queryable + serde::Serialize + serde::de::DeserializeOwned> TransactionalStore<'a, T> {
149    pub fn create(&self, id: u64, inner: &T) -> Result<u64, ConflictableTransactionError<err::Error>> {
150
151        let id_bytes = utils::u64_to_bytes(id);
152
153        let serialized_inner = utils::serialize(inner)?;
154
155        self.tree.insert(&id_bytes, serialized_inner)?;
156
157        let new_terms =
158            inner.query_terms().into_iter().map(|t| t.flatten()).collect::<BTreeSet<_>>();
159
160        self.meta.add(Some(new_terms), id, &id_bytes)?;
161
162        Ok(id)
163    }
164
165    pub fn create_multi(&self, ids: Vec<u64>, inners: &[T]) -> Result<Vec<u64>, ConflictableTransactionError<err::Error>> {
166
167        for (id, inner) in ids.clone().into_iter().zip(inners) {
168
169            let id_bytes = utils::u64_to_bytes(id);
170
171            let serialized_inner = utils::serialize(inner)?;
172
173            self.tree.insert(&id_bytes, serialized_inner)?;
174
175            let new_terms =
176                inner.query_terms().into_iter().map(|t| t.flatten()).collect::<BTreeSet<_>>();
177
178            self.meta.add(Some(new_terms), id, &id_bytes)?;
179        }
180
181        Ok(ids)
182    }
183
184    pub fn update(&self, object: &Object<T>) -> Result<(), ConflictableTransactionError<err::Error>> {
185        self.update_multi(std::slice::from_ref(object))
186    }
187
188    pub fn update_multi(&self, objects: &[Object<T>]) -> Result<(), ConflictableTransactionError<err::Error>> {
189
190        for Object { id, inner } in objects {
191
192            let id_bytes = utils::u64_to_bytes(*id);
193
194            let serialized_inner = utils::serialize(inner)?;
195
196            self.tree.insert(&id_bytes, serialized_inner)?;
197
198            let new_terms =
199                inner.query_terms().into_iter().map(|t| t.flatten()).collect::<BTreeSet<_>>();
200
201            self.meta.add(Some(new_terms), *id, &id_bytes)?;
202        }
203
204        Ok(())
205    }
206
207    pub fn delete(&self, id: u64) -> Result<(), ConflictableTransactionError<err::Error>> {
208        self.delete_multi(&[id])
209    }
210
211    pub fn delete_multi(&self, ids: &[u64]) -> Result<(), ConflictableTransactionError<err::Error>> {
212        for id in ids {
213
214            let id_bytes = utils::u64_to_bytes(*id);
215
216            self.meta.add(None, *id, &id_bytes)?;
217
218            self.tree.remove(&id_bytes)?;
219        }
220
221        Ok(())
222    }
223
224    pub fn find(&self, id: u64) -> Result<Option<Object<T>>, ConflictableTransactionError<err::Error>> {
225        Ok(self
226            .tree
227            .get(utils::u64_to_bytes(id))?
228            .map(|bytes| utils::deserialize(&bytes))
229            .transpose()?
230            .map(|inner| Object { id, inner }))
231    }
232
233    pub fn filter<Q: query::TransactionalQuery<sled::transaction::TransactionalTree>>(&self, query: Q) -> Result<query::TransactionalResults<T>, ConflictableTransactionError<err::Error>> {
234        let matching_ids = query.matching_ids(&self.meta.0)?;
235        Ok(query::TransactionalResults { matching_ids, store: self })
236    }
237
238}
239
240#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
241pub struct Object<T> {
242    pub id: u64,
243    pub inner: T,
244}
245
246impl<T> std::ops::Deref for Object<T> {
247    type Target = T;
248
249    fn deref(&self) -> &Self::Target {
250        &self.inner
251    }
252}
253
254impl<T> std::ops::DerefMut for Object<T> {
255    fn deref_mut(&mut self) -> &mut Self::Target {
256        &mut self.inner
257    }
258}
259
260impl<T: query::Queryable + serde::Serialize + serde::de::DeserializeOwned> Store<T> {
261
262    pub fn transaction<A, F>(&self, f: F) -> sled::transaction::TransactionResult<A, err::Error>
263        where F: Fn(&TransactionalStore<T>) -> sled::transaction::ConflictableTransactionResult<A, err::Error>,
264    {
265        [&self.tree, &self.meta].transaction(|trees| {
266
267            let tree = &trees[0];
268            let meta = &trees[1];
269
270            let transactional_store = TransactionalStore { tree, meta: ext::TransactionalMeta(meta), marker: PhantomData };
271
272            f(&transactional_store)
273        })
274    }
275
276    pub fn new(&self, inner: T) -> err::Result<Object<T>> {
277        let id = self.db.generate_id()?;
278        Ok(Object { id, inner })
279    }
280
281    pub fn create(&self, inner: &T) -> err::Result<u64> {
282        let id = self.db.generate_id()?;
283        Ok(self.transaction(|store| store.create(id, inner))?)
284    }
285
286    pub fn create_multi(&self, inners: &[T]) -> err::Result<Vec<u64>> {
287        let ids = (0..inners.len())
288            .map(|_| self.db.generate_id() )
289            .collect::<Result<Vec<_>, _>>()?;
290        Ok(self.transaction(|store| store.create_multi(ids.clone(), inners))?)
291    }
292
293    pub fn update(&self, object: &Object<T>) -> err::Result<()> {
294        self.update_multi(std::slice::from_ref(object))
295    }
296
297    pub fn update_multi(&self, objects: &[Object<T>]) -> err::Result<()> {
298        Ok(self.transaction(|store| store.update_multi(objects))?)
299    }
300
301    pub fn delete(&self, id: u64) -> err::Result<()> {
302        self.delete_multi(&[id])
303    }
304
305    pub fn delete_multi(&self, ids: &[u64]) -> err::Result<()> {
306        Ok(self.transaction(|store| store.delete_multi(ids))?)
307    }
308
309    pub fn all(&self) -> err::Result<Vec<Object<T>>> {
310        
311        Ok(self
312            .tree
313            .iter()
314            .flatten()
315            .map(|(k, v)| {
316                Ok(Object {
317                    id: utils::bytes_to_u64(k.as_ref())?,
318                    inner: utils::deserialize(&v)?,
319                })
320            })
321            .collect::<err::Result<Vec<_>>>()?)
322    }
323
324    pub fn find(&self, id: u64) -> err::Result<Option<Object<T>>> {
325        Ok(self
326            .tree
327            .get(utils::u64_to_bytes(id))?
328            .map(|bytes| utils::deserialize(&bytes))
329            .transpose()?
330            .map(|inner| Object { id, inner }))
331    }
332
333    pub fn filter<Q: query::Query>(&self, query: Q) -> err::Result<query::Results<T>> {
334        let matching_ids = query.matching_ids(&self.meta)?;
335        Ok(query::Results { matching_ids, store: self })
336    }
337
338    pub fn delete_all(&self) -> err::Result<()> {
339
340        let tree_keys = self.tree.iter().keys().collect::<Vec<_>>();
341        let meta_keys = self.meta.iter().keys().collect::<Vec<_>>();
342
343        [&self.tree, &self.meta].transaction(|trees| {
344            let tree = &trees[0];
345            let meta = &trees[1];
346
347            for key in tree_keys.clone() {
348                let key = key?;
349                tree.remove(key)?;
350            }
351
352            for key in meta_keys.clone() {
353                let key = key?;
354                meta.remove(key)?;
355            }
356
357
358            Ok(())
359        })?;
360
361        Ok(())
362    }
363
364}
treehouse/lib.rs

treehouse/
lib.rs
