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

typed_store/
traits.rs
