sled_tables/
tables.rs

1use crate::error::Error;
2use arrayref::array_ref;
3use serde::{de::DeserializeOwned, Serialize};
4use std::{collections::HashSet, hash::Hash, marker::PhantomData};
5
6/// Imitates collection table per key
7///
8pub struct SledEventTreeVec<T> {
9    tree: sled::Tree,
10    marker: PhantomData<T>,
11}
12
13impl<T> SledEventTreeVec<T> {
14    /// table constructor
15    ///
16    pub fn new(tree: sled::Tree) -> Self {
17        Self {
18            tree,
19            marker: PhantomData,
20        }
21    }
22}
23
24/// DB "Tables" functionality
25///
26impl<T> SledEventTreeVec<T>
27where
28    T: Serialize + DeserializeOwned,
29{
30    /// Gets all elements for given `key` as Vec<T>
31    ///
32    pub fn get(&self, key: u64) -> Result<Option<Vec<T>>, Error> {
33        if let Ok(Some(v)) = self.tree.get(key_bytes(key)) {
34            let set: Vec<T> = serde_cbor::from_slice(&v)?;
35            Ok(Some(set))
36        } else {
37            Ok(None)
38        }
39    }
40
41    /// Overwrites or adds new key<->value into the tree
42    ///
43    pub fn put(&self, key: u64, value: Vec<T>) -> Result<(), Error> {
44        self.tree
45            .insert(key_bytes(key), serde_cbor::to_vec(&value)?)?;
46        Ok(())
47    }
48
49    /// Pushes element to existing set of T
50    /// or creates new one with single element
51    ///
52    pub fn push(&self, key: u64, value: T) -> Result<(), Error> {
53        if let Ok(Some(mut set)) = self.get(key) {
54            set.push(value);
55            self.put(key, set)
56        } else {
57            self.put(key, vec![value])
58        }
59    }
60
61    /// Removes value `T` if present
62    ///
63    pub fn remove(&self, key: u64, value: T) -> Result<(), Error>
64    where
65        T: PartialEq,
66    {
67        if let Ok(Some(set)) = self.get(key) {
68            self.put(key, set.into_iter().filter(|e| e != &value).collect())
69        } else {
70            Ok(())
71        }
72    }
73
74    /// Appends one `Vec<T>` into DB present one
75    /// or `put()`s it if not present as is.
76    ///
77    pub fn append(&self, key: u64, value: Vec<T>) -> Result<(), Error>
78    where
79        T: ToOwned + Clone,
80    {
81        if let Ok(Some(mut set)) = self.get(key) {
82            Ok(set.append(&mut value.to_owned()))
83        } else {
84            self.put(key, value)
85        }
86    }
87
88    /// check if `T` is present in `Vec<T>` in the DB
89    ///
90    pub fn contains_value(&self, value: &T) -> bool
91    where
92        T: PartialEq,
93    {
94        self.tree.iter().flatten().any(|(_k, v)| {
95            serde_cbor::from_slice::<Vec<T>>(&v)
96                .unwrap()
97                .contains(value)
98        })
99    }
100
101    /// iterate inner collection under same key
102    ///
103    pub fn iter_values(&self, key: u64) -> Option<impl DoubleEndedIterator<Item = T>> {
104        if let Ok(Some(values)) = self.tree.get(key_bytes(key)) {
105            Some(
106                serde_cbor::from_slice::<Vec<T>>(&values)
107                    .unwrap()
108                    .into_iter(),
109            )
110        } else {
111            None
112        }
113    }
114}
115
116/// Direct singular key-value of T table
117///
118pub struct SledEventTree<T> {
119    tree: sled::Tree,
120    marker: PhantomData<T>,
121}
122
123impl<T> SledEventTree<T> {
124    /// table constructor
125    ///
126    pub fn new(tree: sled::Tree) -> Self {
127        Self {
128            tree,
129            marker: PhantomData,
130        }
131    }
132
133    /// Returns size of the underlying tree
134    pub fn len(&self) -> usize {
135        self.tree.len()
136    }
137}
138
139/// DB "Tables" functionality
140///
141impl<T> SledEventTree<T>
142where
143    T: Serialize + DeserializeOwned,
144{
145    /// get entire Vec<T> in one go
146    ///
147    pub fn get(&self, id: u64) -> Result<Option<T>, Error> {
148        match self.tree.get(key_bytes(id))? {
149            Some(value) => Ok(Some(serde_cbor::from_slice(&value)?)),
150            None => Ok(None),
151        }
152    }
153
154    /// check if sprovided `u64` key is present in the db
155    ///
156    pub fn contains_key(&self, id: u64) -> Result<bool, Error> {
157        Ok(self.tree.contains_key(key_bytes(id))?)
158    }
159
160    /// check if value `T` is present in the db
161    ///
162    pub fn contains_value(&self, value: &T) -> bool
163    where
164        T: PartialEq,
165    {
166        self.tree
167            .iter()
168            .flatten()
169            .find(|(_, v)| serde_cbor::from_slice::<T>(&v).unwrap().eq(value))
170            .is_some()
171    }
172
173    /// insert `T` with given `key`
174    /// Warning! This will rewrite existing value with the same `key`
175    ///
176    pub fn insert(&self, key: u64, value: &T) -> Result<(), Error> {
177        self.tree
178            .insert(key_bytes(key), serde_cbor::to_vec(value)?)?;
179        Ok(())
180    }
181
182    /// iterator over `T` deserialized from the db
183    ///
184    pub fn iter(&self) -> impl DoubleEndedIterator<Item = T> {
185        self.tree
186            .iter()
187            .flatten()
188            .flat_map(|(_, v)| serde_cbor::from_slice(&v))
189    }
190
191    /// provides which `u64` key to use to add NEW entry
192    ///
193    pub fn get_next_key(&self) -> u64 {
194        if let Ok(Some((k, _v))) = self.tree.last() {
195            u64::from_be_bytes(array_ref!(k, 0, 8).to_owned()) + 1
196        } else {
197            0
198        }
199    }
200
201    /// somewhat expensive! gets optional `u64` key for given `&T`
202    /// if present in the db
203    ///
204    pub fn get_key_by_value(&self, value: &T) -> Result<Option<u64>, Error>
205    where
206        T: Serialize,
207    {
208        let value = serde_cbor::to_vec(value)?;
209        if let Some((key, _)) = self.tree.iter().flatten().find(|(_k, v)| v.eq(&value)) {
210            Ok(Some(u64::from_be_bytes(array_ref!(key, 0, 8).to_owned())))
211        } else {
212            Ok(None)
213        }
214    }
215
216    /// combination of `get_key_by_value()` and `get_next_key()`
217    /// also expensive...
218    /// to be used when unsure if identifier is present in the db
219    ///
220    pub fn designated_key(&self, identifier: &T) -> u64
221    where
222        T: Serialize,
223    {
224        if let Ok(Some(key)) = self.get_key_by_value(identifier) {
225            key
226        } else {
227            let key = self.get_next_key();
228            self.tree
229                .insert(key_bytes(key), serde_cbor::to_vec(identifier).unwrap())
230                .unwrap();
231            key
232        }
233    }
234}
235
236pub struct SledEventTreeMap<T> {
237    tree: sled::Tree,
238    marker: PhantomData<T>,
239}
240
241impl<T> SledEventTreeMap<T> {
242    pub fn new(tree: sled::Tree) -> Self {
243        Self {
244            tree,
245            marker: PhantomData,
246        }
247    }
248}
249
250impl<T> SledEventTreeMap<T>
251where
252    T: Serialize + DeserializeOwned + Eq + Hash,
253{
254    pub fn get(&self, id: u64) -> Result<Option<HashSet<T>>, Error> {
255        match self.tree.get(key_bytes(id))? {
256            Some(val) => Ok(Some(serde_cbor::from_slice(&val)?)),
257            None => Ok(None),
258        }
259    }
260
261    pub fn replace_or_insert(&self, id: u64, new_set: HashSet<T>) -> Result<(), Error> {
262        self.tree
263            .insert(key_bytes(id), serde_cbor::to_vec(&new_set)?)?;
264        Ok(())
265    }
266
267    pub fn add_or_skip(&self, key: u64, value: T) -> Result<(), Error> {
268        if let Ok(Some(mut events)) = self.get(key) {
269            events.insert(value);
270            self.replace_or_insert(key, events)
271        } else {
272            let mut events = HashSet::new();
273            events.insert(value);
274            self.replace_or_insert(key, events)?;
275            Ok(())
276        }
277    }
278
279    pub fn iter_values(&self, id: u64) -> Option<impl IntoIterator<Item = T>> {
280        if let Ok(Some(data)) = self.tree.get(key_bytes(id)) {
281            Some(serde_cbor::from_slice::<HashSet<T>>(&data).unwrap())
282        } else {
283            None
284        }
285    }
286}
287
288fn key_bytes(key: u64) -> [u8; 8] {
289    key.to_be_bytes()
290}
sled_tables/tables.rs

sled_tables/
tables.rs
