snops_common/db/
tree.rs

1use bytes::Buf;
2
3use super::error::DatabaseError;
4use crate::format::{read_dataformat, write_dataformat, DataFormat};
5
6pub struct DbTree<K, V> {
7    tree: sled::Tree,
8    _phantom: std::marker::PhantomData<(K, V)>,
9}
10
11impl<K: DataFormat, V: DataFormat> DbTree<K, V> {
12    pub fn new(tree: sled::Tree) -> Self {
13        Self {
14            tree,
15            _phantom: std::marker::PhantomData,
16        }
17    }
18
19    pub fn read_all(&self) -> impl Iterator<Item = (K, V)> {
20        self.tree.iter().filter_map(|row| {
21            let (key_bytes, value_bytes) = match row {
22                Ok((key, value)) => (key, value),
23                Err(e) => {
24                    tracing::error!("Error reading row from store: {e}");
25                    return None;
26                }
27            };
28
29            let key = match K::read_data(&mut key_bytes.reader(), &K::LATEST_HEADER) {
30                Ok(key) => key,
31                Err(e) => {
32                    tracing::error!("Error parsing key from store: {e}");
33                    return None;
34                }
35            };
36
37            let value = match read_dataformat(&mut value_bytes.reader()) {
38                Ok(value) => value,
39                Err(e) => {
40                    tracing::error!("Error parsing value from store: {e}");
41                    return None;
42                }
43            };
44
45            Some((key, value))
46        })
47    }
48
49    pub fn read_with_prefix<Prefix: DataFormat>(
50        &self,
51        prefix: &Prefix,
52    ) -> Result<impl Iterator<Item = (K, V)>, DatabaseError> {
53        Ok(self
54            .tree
55            .scan_prefix(prefix.to_byte_vec()?)
56            .filter_map(|row| {
57                let (key_bytes, value_bytes) = match row {
58                    Ok((key, value)) => (key, value),
59                    Err(e) => {
60                        tracing::error!("Error reading row from store: {e}");
61                        return None;
62                    }
63                };
64
65                let key = match K::read_data(&mut key_bytes.reader(), &K::LATEST_HEADER) {
66                    Ok(key) => key,
67                    Err(e) => {
68                        tracing::error!("Error parsing key from store: {e}");
69                        return None;
70                    }
71                };
72
73                let value = match read_dataformat(&mut value_bytes.reader()) {
74                    Ok(value) => value,
75                    Err(e) => {
76                        tracing::error!("Error parsing value from store: {e}");
77                        return None;
78                    }
79                };
80
81                Some((key, value))
82            }))
83    }
84
85    pub fn restore(&self, key: &K) -> Result<Option<V>, DatabaseError> {
86        Ok(self
87            .tree
88            .get(key.to_byte_vec()?)?
89            .map(|value_bytes| read_dataformat(&mut value_bytes.reader()))
90            .transpose()?)
91    }
92
93    pub fn save(&self, key: &K, value: &V) -> Result<(), DatabaseError> {
94        let key_bytes = key.to_byte_vec()?;
95        let mut value_bytes = Vec::new();
96        write_dataformat(&mut value_bytes, value)?;
97        self.tree.insert(key_bytes, value_bytes)?;
98        Ok(())
99    }
100
101    pub fn save_option(&self, key: &K, value: Option<&V>) -> Result<(), DatabaseError> {
102        match value {
103            Some(value) => self.save(key, value),
104            None => self.delete(key).map(|_| ()),
105        }
106    }
107
108    pub fn delete(&self, key: &K) -> Result<bool, DatabaseError> {
109        Ok(self.tree.remove(key.to_byte_vec()?)?.is_some())
110    }
111
112    pub fn delete_with_prefix<Prefix: DataFormat>(
113        &self,
114        prefix: &Prefix,
115    ) -> Result<usize, DatabaseError> {
116        Ok(self
117            .tree
118            .scan_prefix(prefix.to_byte_vec()?)
119            .map(|row| {
120                let key_bytes = match row {
121                    Ok((key, _)) => key,
122                    Err(e) => {
123                        tracing::error!("Error reading row from store: {e}");
124                        return 0;
125                    }
126                };
127
128                let key = match K::read_data(&mut key_bytes.reader(), &K::LATEST_HEADER) {
129                    Ok(key) => key,
130                    Err(e) => {
131                        tracing::error!("Error parsing key from store: {e}");
132                        return 0;
133                    }
134                };
135
136                if let Err(e) = self.delete(&key) {
137                    tracing::error!("Error deleting key from store: {e}");
138                    return 0;
139                }
140
141                1
142            })
143            .sum())
144    }
145}
snops_common/db/tree.rs

snops_common/db/
tree.rs
