typed_store/
traits.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{borrow::Borrow, collections::BTreeMap, error::Error, ops::RangeBounds};
6
7use async_trait::async_trait;
8use serde::{Serialize, de::DeserializeOwned};
9
10use crate::TypedStoreError;
11
12pub trait Map<'a, K, V>
13where
14    K: Serialize + DeserializeOwned,
15    V: Serialize + DeserializeOwned,
16{
17    type Error: Error;
18    type Iterator: Iterator<Item = (K, V)>;
19    type SafeIterator: Iterator<Item = Result<(K, V), TypedStoreError>>;
20    type Keys: Iterator<Item = Result<K, TypedStoreError>>;
21    type Values: Iterator<Item = Result<V, TypedStoreError>>;
22
23    /// Returns true if the map contains a value for the specified key.
24    fn contains_key(&self, key: &K) -> Result<bool, Self::Error>;
25
26    /// Returns true if the map contains a value for the specified key.
27    fn multi_contains_keys<J>(
28        &self,
29        keys: impl IntoIterator<Item = J>,
30    ) -> Result<Vec<bool>, Self::Error>
31    where
32        J: Borrow<K>,
33    {
34        keys.into_iter()
35            .map(|key| self.contains_key(key.borrow()))
36            .collect()
37    }
38
39    /// Returns the value for the given key from the map, if it exists.
40    fn get(&self, key: &K) -> Result<Option<V>, Self::Error>;
41
42    /// Returns the raw value (serialized bytes) for the given key from the map,
43    /// if it exists.
44    fn get_raw_bytes(&self, key: &K) -> Result<Option<Vec<u8>>, Self::Error>;
45
46    /// Inserts the given key-value pair into the map.
47    fn insert(&self, key: &K, value: &V) -> Result<(), Self::Error>;
48
49    /// Removes the entry for the given key from the map.
50    fn remove(&self, key: &K) -> Result<(), Self::Error>;
51
52    /// Removes every key-value pair from the map.
53    fn unsafe_clear(&self) -> Result<(), Self::Error>;
54
55    /// Removes every key-value pair from the map by deleting the underlying
56    /// file.
57    fn delete_file_in_range(&self, from: &K, to: &K) -> Result<(), TypedStoreError>;
58
59    /// Uses delete range on the entire key range
60    fn schedule_delete_all(&self) -> Result<(), TypedStoreError>;
61
62    /// Returns true if the map is empty, otherwise false.
63    fn is_empty(&self) -> bool;
64
65    /// Returns an unbounded iterator visiting each key-value pair in the map.
66    /// This is potentially unsafe as it can perform a full table scan
67    fn unbounded_iter(&'a self) -> Self::Iterator;
68
69    /// Returns an iterator visiting each key-value pair within the specified
70    /// bounds in the map.
71    fn iter_with_bounds(&'a self, lower_bound: Option<K>, upper_bound: Option<K>)
72    -> Self::Iterator;
73
74    /// Similar to `iter_with_bounds` but allows specifying
75    /// inclusivity/exclusivity of ranges explicitly. TODO: find better name
76    fn range_iter(&'a self, range: impl RangeBounds<K>) -> Self::Iterator;
77
78    /// Same as `iter` but performs status check.
79    fn safe_iter(&'a self) -> Self::SafeIterator;
80
81    // Same as `iter_with_bounds` but performs status check.
82    fn safe_iter_with_bounds(
83        &'a self,
84        lower_bound: Option<K>,
85        upper_bound: Option<K>,
86    ) -> Self::SafeIterator;
87
88    // Same as `range_iter` but performs status check.
89    fn safe_range_iter(&'a self, range: impl RangeBounds<K>) -> Self::SafeIterator;
90
91    /// Returns an iterator over each key in the map.
92    fn keys(&'a self) -> Self::Keys;
93
94    /// Returns an iterator over each value in the map.
95    fn values(&'a self) -> Self::Values;
96
97    /// Returns a vector of values corresponding to the keys provided,
98    /// non-atomically.
99    fn multi_get<J>(&self, keys: impl IntoIterator<Item = J>) -> Result<Vec<Option<V>>, Self::Error>
100    where
101        J: Borrow<K>,
102    {
103        keys.into_iter().map(|key| self.get(key.borrow())).collect()
104    }
105
106    /// Returns a vector of raw values corresponding to the keys provided,
107    /// non-atomically.
108    fn multi_get_raw_bytes<J>(
109        &self,
110        keys: impl IntoIterator<Item = J>,
111    ) -> Result<Vec<Option<Vec<u8>>>, Self::Error>
112    where
113        J: Borrow<K>,
114    {
115        keys.into_iter()
116            .map(|key| self.get_raw_bytes(key.borrow()))
117            .collect()
118    }
119
120    /// Returns a vector of values corresponding to the keys provided,
121    /// non-atomically.
122    fn chunked_multi_get<J>(
123        &self,
124        keys: impl IntoIterator<Item = J>,
125        _chunk_size: usize,
126    ) -> Result<Vec<Option<V>>, Self::Error>
127    where
128        J: Borrow<K>,
129    {
130        keys.into_iter().map(|key| self.get(key.borrow())).collect()
131    }
132
133    /// Inserts key-value pairs, non-atomically.
134    fn multi_insert<J, U>(
135        &self,
136        key_val_pairs: impl IntoIterator<Item = (J, U)>,
137    ) -> Result<(), Self::Error>
138    where
139        J: Borrow<K>,
140        U: Borrow<V>,
141    {
142        key_val_pairs
143            .into_iter()
144            .try_for_each(|(key, value)| self.insert(key.borrow(), value.borrow()))
145    }
146
147    /// Removes keys, non-atomically.
148    fn multi_remove<J>(&self, keys: impl IntoIterator<Item = J>) -> Result<(), Self::Error>
149    where
150        J: Borrow<K>,
151    {
152        keys.into_iter()
153            .try_for_each(|key| self.remove(key.borrow()))
154    }
155
156    /// Try to catch up with primary when running as secondary
157    fn try_catch_up_with_primary(&self) -> Result<(), Self::Error>;
158}
159
160#[async_trait]
161pub trait AsyncMap<'a, K, V>
162where
163    K: Serialize + DeserializeOwned + std::marker::Sync,
164    V: Serialize + DeserializeOwned + std::marker::Sync + std::marker::Send,
165{
166    type Error: Error;
167    type Iterator: Iterator<Item = Result<(K, V), TypedStoreError>>;
168    type Keys: Iterator<Item = Result<K, TypedStoreError>>;
169    type Values: Iterator<Item = Result<V, TypedStoreError>>;
170
171    /// Returns true if the map contains a value for the specified key.
172    async fn contains_key(&self, key: &K) -> Result<bool, Self::Error>;
173
174    /// Returns the value for the given key from the map, if it exists.
175    async fn get(&self, key: &K) -> Result<Option<V>, Self::Error>;
176
177    /// Returns the raw value (serialized bytes) for the given key from the map,
178    /// if it exists.
179    async fn get_raw_bytes(&self, key: &K) -> Result<Option<Vec<u8>>, Self::Error>;
180
181    /// Returns true if the map is empty, otherwise false.
182    async fn is_empty(&self) -> bool;
183
184    /// Returns an iterator visiting each key-value pair in the map.
185    async fn iter(&'a self) -> Self::Iterator;
186
187    /// Returns an iterator over each key in the map.
188    async fn keys(&'a self) -> Self::Keys;
189
190    /// Returns an iterator over each value in the map.
191    async fn values(&'a self) -> Self::Values;
192
193    /// Returns a vector of values corresponding to the keys provided,
194    /// non-atomically.
195    async fn multi_get<J>(
196        &self,
197        keys: impl IntoIterator<Item = J> + std::marker::Send,
198    ) -> Result<Vec<Option<V>>, Self::Error>
199    where
200        J: Borrow<K>;
201
202    /// Try to catch up with primary when running as secondary
203    async fn try_catch_up_with_primary(&self) -> Result<(), Self::Error>;
204}
205
206pub struct TableSummary {
207    pub num_keys: u64,
208    pub key_bytes_total: usize,
209    pub value_bytes_total: usize,
210    pub key_hist: hdrhistogram::Histogram<u64>,
211    pub value_hist: hdrhistogram::Histogram<u64>,
212}
213
214pub trait TypedStoreDebug {
215    /// Dump a DB table with pagination
216    fn dump_table(
217        &self,
218        table_name: String,
219        page_size: u16,
220        page_number: usize,
221    ) -> eyre::Result<BTreeMap<String, String>>;
222
223    /// Get the name of the DB. This is simply the name of the struct
224    fn primary_db_name(&self) -> String;
225
226    /// Get a map of table names to key-value types
227    fn describe_all_tables(&self) -> BTreeMap<String, (String, String)>;
228
229    /// Count the entries in the table
230    fn count_table_keys(&self, table_name: String) -> eyre::Result<usize>;
231
232    /// Return table summary of the input table
233    fn table_summary(&self, table_name: String) -> eyre::Result<TableSummary>;
234}
typed_store/traits.rs

typed_store/
traits.rs
