1use serde::{Deserialize, Serialize};
7use thiserror::Error;
8
9#[cfg(with_testing)]
10use crate::store::TestKeyValueDatabase;
11use crate::{
12 batch::Batch,
13 store::{
14 KeyValueDatabase, KeyValueStoreError, ReadableKeyValueStore, WithError,
15 WritableKeyValueStore,
16 },
17};
18
19#[derive(Debug, Clone, Serialize, Deserialize)]
21pub struct DualDatabase<D1, D2, A> {
22 pub first_database: D1,
24 pub second_database: D2,
26 _marker: std::marker::PhantomData<A>,
28}
29
30#[derive(Debug, Clone, Serialize, Deserialize)]
32pub struct DualStoreConfig<C1, C2> {
33 pub first_config: C1,
35 pub second_config: C2,
37}
38
39#[derive(Clone, Copy, Debug)]
41pub enum StoreInUse {
42 First,
44 Second,
46}
47
48pub trait DualStoreRootKeyAssignment {
50 fn assigned_store(root_key: &[u8]) -> Result<StoreInUse, bcs::Error>;
52}
53
54#[derive(Clone)]
56pub enum DualStore<S1, S2> {
57 First(S1),
59 Second(S2),
61}
62
63impl<D1, D2, A> WithError for DualDatabase<D1, D2, A>
64where
65 D1: WithError,
66 D2: WithError,
67{
68 type Error = DualStoreError<D1::Error, D2::Error>;
69}
70
71impl<S1, S2> WithError for DualStore<S1, S2>
72where
73 S1: WithError,
74 S2: WithError,
75{
76 type Error = DualStoreError<S1::Error, S2::Error>;
77}
78
79impl<S1, S2> ReadableKeyValueStore for DualStore<S1, S2>
80where
81 S1: ReadableKeyValueStore,
82 S2: ReadableKeyValueStore,
83{
84 const MAX_KEY_SIZE: usize = if S1::MAX_KEY_SIZE < S2::MAX_KEY_SIZE {
86 S1::MAX_KEY_SIZE
87 } else {
88 S2::MAX_KEY_SIZE
89 };
90
91 fn root_key(&self) -> Result<Vec<u8>, Self::Error> {
92 Ok(match self {
93 Self::First(store) => store.root_key().map_err(DualStoreError::First)?,
94 Self::Second(store) => store.root_key().map_err(DualStoreError::Second)?,
95 })
96 }
97
98 async fn read_value_bytes(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
99 let result = match self {
100 Self::First(store) => store
101 .read_value_bytes(key)
102 .await
103 .map_err(DualStoreError::First)?,
104 Self::Second(store) => store
105 .read_value_bytes(key)
106 .await
107 .map_err(DualStoreError::Second)?,
108 };
109 Ok(result)
110 }
111
112 async fn contains_key(&self, key: &[u8]) -> Result<bool, Self::Error> {
113 let result = match self {
114 Self::First(store) => store
115 .contains_key(key)
116 .await
117 .map_err(DualStoreError::First)?,
118 Self::Second(store) => store
119 .contains_key(key)
120 .await
121 .map_err(DualStoreError::Second)?,
122 };
123 Ok(result)
124 }
125
126 async fn contains_keys(&self, keys: &[Vec<u8>]) -> Result<Vec<bool>, Self::Error> {
127 let result = match self {
128 Self::First(store) => store
129 .contains_keys(keys)
130 .await
131 .map_err(DualStoreError::First)?,
132 Self::Second(store) => store
133 .contains_keys(keys)
134 .await
135 .map_err(DualStoreError::Second)?,
136 };
137 Ok(result)
138 }
139
140 async fn read_multi_values_bytes(
141 &self,
142 keys: &[Vec<u8>],
143 ) -> Result<Vec<Option<Vec<u8>>>, Self::Error> {
144 let result = match self {
145 Self::First(store) => store
146 .read_multi_values_bytes(keys)
147 .await
148 .map_err(DualStoreError::First)?,
149 Self::Second(store) => store
150 .read_multi_values_bytes(keys)
151 .await
152 .map_err(DualStoreError::Second)?,
153 };
154 Ok(result)
155 }
156
157 async fn find_keys_by_prefix(&self, key_prefix: &[u8]) -> Result<Vec<Vec<u8>>, Self::Error> {
158 let result = match self {
159 Self::First(store) => store
160 .find_keys_by_prefix(key_prefix)
161 .await
162 .map_err(DualStoreError::First)?,
163 Self::Second(store) => store
164 .find_keys_by_prefix(key_prefix)
165 .await
166 .map_err(DualStoreError::Second)?,
167 };
168 Ok(result)
169 }
170
171 async fn find_key_values_by_prefix(
172 &self,
173 key_prefix: &[u8],
174 ) -> Result<Vec<(Vec<u8>, Vec<u8>)>, Self::Error> {
175 let result = match self {
176 Self::First(store) => store
177 .find_key_values_by_prefix(key_prefix)
178 .await
179 .map_err(DualStoreError::First)?,
180 Self::Second(store) => store
181 .find_key_values_by_prefix(key_prefix)
182 .await
183 .map_err(DualStoreError::Second)?,
184 };
185 Ok(result)
186 }
187}
188
189impl<S1, S2> WritableKeyValueStore for DualStore<S1, S2>
190where
191 S1: WritableKeyValueStore,
192 S2: WritableKeyValueStore,
193{
194 const MAX_VALUE_SIZE: usize = usize::MAX;
195
196 async fn write_batch(&self, batch: Batch) -> Result<(), Self::Error> {
197 match self {
198 Self::First(store) => store
199 .write_batch(batch)
200 .await
201 .map_err(DualStoreError::First)?,
202 Self::Second(store) => store
203 .write_batch(batch)
204 .await
205 .map_err(DualStoreError::Second)?,
206 }
207 Ok(())
208 }
209
210 async fn clear_journal(&self) -> Result<(), Self::Error> {
211 match self {
212 Self::First(store) => store.clear_journal().await.map_err(DualStoreError::First)?,
213 Self::Second(store) => store
214 .clear_journal()
215 .await
216 .map_err(DualStoreError::Second)?,
217 }
218 Ok(())
219 }
220}
221
222impl<D1, D2, A> KeyValueDatabase for DualDatabase<D1, D2, A>
223where
224 D1: KeyValueDatabase,
225 D2: KeyValueDatabase,
226 A: DualStoreRootKeyAssignment + linera_base::util::traits::AutoTraits,
227{
228 type Config = DualStoreConfig<D1::Config, D2::Config>;
229 type Store = DualStore<D1::Store, D2::Store>;
230
231 fn get_name() -> String {
232 format!("dual {} and {}", D1::get_name(), D2::get_name())
233 }
234
235 async fn connect(config: &Self::Config, namespace: &str) -> Result<Self, Self::Error> {
236 let first_database = D1::connect(&config.first_config, namespace)
237 .await
238 .map_err(DualStoreError::First)?;
239 let second_database = D2::connect(&config.second_config, namespace)
240 .await
241 .map_err(DualStoreError::Second)?;
242 let database = Self {
243 first_database,
244 second_database,
245 _marker: std::marker::PhantomData,
246 };
247 Ok(database)
248 }
249
250 fn open_shared(&self, root_key: &[u8]) -> Result<Self::Store, Self::Error> {
251 match A::assigned_store(root_key)? {
252 StoreInUse::First => {
253 let store = self
254 .first_database
255 .open_shared(root_key)
256 .map_err(DualStoreError::First)?;
257 Ok(DualStore::First(store))
258 }
259 StoreInUse::Second => {
260 let store = self
261 .second_database
262 .open_shared(root_key)
263 .map_err(DualStoreError::Second)?;
264 Ok(DualStore::Second(store))
265 }
266 }
267 }
268
269 fn open_exclusive(&self, root_key: &[u8]) -> Result<Self::Store, Self::Error> {
270 match A::assigned_store(root_key)? {
271 StoreInUse::First => {
272 let store = self
273 .first_database
274 .open_exclusive(root_key)
275 .map_err(DualStoreError::First)?;
276 Ok(DualStore::First(store))
277 }
278 StoreInUse::Second => {
279 let store = self
280 .second_database
281 .open_exclusive(root_key)
282 .map_err(DualStoreError::Second)?;
283 Ok(DualStore::Second(store))
284 }
285 }
286 }
287
288 async fn list_all(config: &Self::Config) -> Result<Vec<String>, Self::Error> {
289 let namespaces1 = D1::list_all(&config.first_config)
290 .await
291 .map_err(DualStoreError::First)?;
292 let mut namespaces = Vec::new();
293 for namespace in namespaces1 {
294 if D2::exists(&config.second_config, &namespace)
295 .await
296 .map_err(DualStoreError::Second)?
297 {
298 namespaces.push(namespace);
299 } else {
300 tracing::warn!("Namespace {} only exists in the first store", namespace);
301 }
302 }
303 Ok(namespaces)
304 }
305
306 async fn list_root_keys(&self) -> Result<Vec<Vec<u8>>, Self::Error> {
307 let mut root_keys = self
308 .first_database
309 .list_root_keys()
310 .await
311 .map_err(DualStoreError::First)?;
312 root_keys.extend(
313 self.second_database
314 .list_root_keys()
315 .await
316 .map_err(DualStoreError::Second)?,
317 );
318 Ok(root_keys)
319 }
320
321 async fn exists(config: &Self::Config, namespace: &str) -> Result<bool, Self::Error> {
322 Ok(D1::exists(&config.first_config, namespace)
323 .await
324 .map_err(DualStoreError::First)?
325 && D2::exists(&config.second_config, namespace)
326 .await
327 .map_err(DualStoreError::Second)?)
328 }
329
330 async fn create(config: &Self::Config, namespace: &str) -> Result<(), Self::Error> {
331 let exists1 = D1::exists(&config.first_config, namespace)
332 .await
333 .map_err(DualStoreError::First)?;
334 let exists2 = D2::exists(&config.second_config, namespace)
335 .await
336 .map_err(DualStoreError::Second)?;
337 if exists1 && exists2 {
338 return Err(DualStoreError::StoreAlreadyExists);
339 }
340 if !exists1 {
341 D1::create(&config.first_config, namespace)
342 .await
343 .map_err(DualStoreError::First)?;
344 }
345 if !exists2 {
346 D2::create(&config.second_config, namespace)
347 .await
348 .map_err(DualStoreError::Second)?;
349 }
350 Ok(())
351 }
352
353 async fn delete(config: &Self::Config, namespace: &str) -> Result<(), Self::Error> {
354 D1::delete(&config.first_config, namespace)
355 .await
356 .map_err(DualStoreError::First)?;
357 D2::delete(&config.second_config, namespace)
358 .await
359 .map_err(DualStoreError::Second)?;
360 Ok(())
361 }
362}
363
364#[cfg(with_testing)]
365impl<D1, D2, A> TestKeyValueDatabase for DualDatabase<D1, D2, A>
366where
367 D1: TestKeyValueDatabase,
368 D2: TestKeyValueDatabase,
369 A: DualStoreRootKeyAssignment + linera_base::util::traits::AutoTraits,
370{
371 async fn new_test_config() -> Result<Self::Config, Self::Error> {
372 let first_config = D1::new_test_config().await.map_err(DualStoreError::First)?;
373 let second_config = D2::new_test_config()
374 .await
375 .map_err(DualStoreError::Second)?;
376 Ok(DualStoreConfig {
377 first_config,
378 second_config,
379 })
380 }
381}
382
383#[derive(Error, Debug)]
385pub enum DualStoreError<E1, E2> {
386 #[error("Store already exists during a create operation")]
388 StoreAlreadyExists,
389
390 #[error(transparent)]
392 BcsError(#[from] bcs::Error),
393
394 #[error("Error in first store: {0}")]
396 First(E1),
397
398 #[error("Error in second store: {0}")]
400 Second(E2),
401}
402
403impl<E1, E2> KeyValueStoreError for DualStoreError<E1, E2>
404where
405 E1: KeyValueStoreError,
406 E2: KeyValueStoreError,
407{
408 const BACKEND: &'static str = "dual_store";
409
410 fn must_reload_view(&self) -> bool {
411 match self {
412 DualStoreError::First(e) => e.must_reload_view(),
413 DualStoreError::Second(e) => e.must_reload_view(),
414 _ => false,
415 }
416 }
417}