1#[cfg(with_testing)]
7use std::sync::Arc;
8
9use linera_base::ensure;
10use linera_views::{
11 batch::Batch,
12 store::{ReadableKeyValueStore, WithError, WritableKeyValueStore},
13};
14use thiserror::Error;
15
16#[cfg(with_testing)]
17use super::mock_key_value_store::MockKeyValueStore;
18use crate::{
19 contract::wit::{
20 base_runtime_api::{self as contract_wit},
21 contract_runtime_api::{self, WriteOperation},
22 },
23 service::wit::base_runtime_api as service_wit,
24 util::yield_once,
25};
26
27const MAX_KEY_SIZE: usize = 900;
34
35#[derive(Clone)]
37pub struct KeyValueStore {
38 wit_api: WitInterface,
39}
40
41#[cfg_attr(with_testing, allow(dead_code))]
42impl KeyValueStore {
43 pub(crate) fn for_contracts() -> Self {
45 KeyValueStore {
46 wit_api: WitInterface::Contract,
47 }
48 }
49
50 pub(crate) fn for_services() -> Self {
52 KeyValueStore {
53 wit_api: WitInterface::Service,
54 }
55 }
56
57 #[cfg(with_testing)]
59 pub fn mock() -> Self {
60 KeyValueStore {
61 wit_api: WitInterface::Mock {
62 store: Arc::new(MockKeyValueStore::default()),
63 read_only: true,
64 },
65 }
66 }
67
68 #[cfg(with_testing)]
71 pub fn to_mut(&self) -> Self {
72 let WitInterface::Mock { store, .. } = &self.wit_api else {
73 panic!("Real `KeyValueStore` should not be used in unit tests");
74 };
75
76 KeyValueStore {
77 wit_api: WitInterface::Mock {
78 store: store.clone(),
79 read_only: false,
80 },
81 }
82 }
83}
84
85impl WithError for KeyValueStore {
86 type Error = KeyValueStoreError;
87}
88
89#[derive(Error, Debug)]
91pub enum KeyValueStoreError {
92 #[error("Key too long")]
94 KeyTooLong,
95
96 #[error(transparent)]
98 BcsError(#[from] bcs::Error),
99}
100
101impl linera_views::store::KeyValueStoreError for KeyValueStoreError {
102 const BACKEND: &'static str = "key_value_store";
103}
104
105impl ReadableKeyValueStore for KeyValueStore {
106 const MAX_KEY_SIZE: usize = MAX_KEY_SIZE;
109
110 fn max_stream_queries(&self) -> usize {
111 1
112 }
113
114 fn root_key(&self) -> Result<Vec<u8>, KeyValueStoreError> {
115 Ok(Vec::new())
116 }
117
118 async fn contains_key(&self, key: &[u8]) -> Result<bool, KeyValueStoreError> {
119 ensure!(
120 key.len() <= Self::MAX_KEY_SIZE,
121 KeyValueStoreError::KeyTooLong
122 );
123 let promise = self.wit_api.contains_key_new(key);
124 yield_once().await;
125 Ok(self.wit_api.contains_key_wait(promise))
126 }
127
128 async fn contains_keys(&self, keys: &[Vec<u8>]) -> Result<Vec<bool>, KeyValueStoreError> {
129 for key in keys {
130 ensure!(
131 key.len() <= Self::MAX_KEY_SIZE,
132 KeyValueStoreError::KeyTooLong
133 );
134 }
135 let promise = self.wit_api.contains_keys_new(keys);
136 yield_once().await;
137 Ok(self.wit_api.contains_keys_wait(promise))
138 }
139
140 async fn read_multi_values_bytes(
141 &self,
142 keys: &[Vec<u8>],
143 ) -> Result<Vec<Option<Vec<u8>>>, KeyValueStoreError> {
144 for key in keys {
145 ensure!(
146 key.len() <= Self::MAX_KEY_SIZE,
147 KeyValueStoreError::KeyTooLong
148 );
149 }
150 let promise = self.wit_api.read_multi_values_bytes_new(keys);
151 yield_once().await;
152 Ok(self.wit_api.read_multi_values_bytes_wait(promise))
153 }
154
155 async fn read_value_bytes(&self, key: &[u8]) -> Result<Option<Vec<u8>>, KeyValueStoreError> {
156 ensure!(
157 key.len() <= Self::MAX_KEY_SIZE,
158 KeyValueStoreError::KeyTooLong
159 );
160 let promise = self.wit_api.read_value_bytes_new(key);
161 yield_once().await;
162 Ok(self.wit_api.read_value_bytes_wait(promise))
163 }
164
165 async fn find_keys_by_prefix(
166 &self,
167 key_prefix: &[u8],
168 ) -> Result<Vec<Vec<u8>>, KeyValueStoreError> {
169 ensure!(
170 key_prefix.len() <= Self::MAX_KEY_SIZE,
171 KeyValueStoreError::KeyTooLong
172 );
173 let promise = self.wit_api.find_keys_new(key_prefix);
174 yield_once().await;
175 Ok(self.wit_api.find_keys_wait(promise))
176 }
177
178 async fn find_key_values_by_prefix(
179 &self,
180 key_prefix: &[u8],
181 ) -> Result<Vec<(Vec<u8>, Vec<u8>)>, KeyValueStoreError> {
182 ensure!(
183 key_prefix.len() <= Self::MAX_KEY_SIZE,
184 KeyValueStoreError::KeyTooLong
185 );
186 let promise = self.wit_api.find_key_values_new(key_prefix);
187 yield_once().await;
188 Ok(self.wit_api.find_key_values_wait(promise))
189 }
190}
191
192impl WritableKeyValueStore for KeyValueStore {
193 const MAX_VALUE_SIZE: usize = usize::MAX;
194
195 async fn write_batch(&self, batch: Batch) -> Result<(), KeyValueStoreError> {
196 self.wit_api.write_batch(batch);
197 Ok(())
198 }
199
200 async fn clear_journal(&self) -> Result<(), KeyValueStoreError> {
201 Ok(())
202 }
203}
204
205#[derive(Clone)]
207#[cfg_attr(with_testing, allow(dead_code))]
208enum WitInterface {
209 Contract,
211 Service,
213 #[cfg(with_testing)]
214 Mock {
216 store: Arc<MockKeyValueStore>,
217 read_only: bool,
218 },
219}
220
221impl WitInterface {
222 fn contains_key_new(&self, key: &[u8]) -> u32 {
224 match self {
225 WitInterface::Contract => contract_wit::contains_key_new(key),
226 WitInterface::Service => service_wit::contains_key_new(key),
227 #[cfg(with_testing)]
228 WitInterface::Mock { store, .. } => store.contains_key_new(key),
229 }
230 }
231
232 fn contains_key_wait(&self, promise: u32) -> bool {
234 match self {
235 WitInterface::Contract => contract_wit::contains_key_wait(promise),
236 WitInterface::Service => service_wit::contains_key_wait(promise),
237 #[cfg(with_testing)]
238 WitInterface::Mock { store, .. } => store.contains_key_wait(promise),
239 }
240 }
241
242 fn contains_keys_new(&self, keys: &[Vec<u8>]) -> u32 {
244 match self {
245 WitInterface::Contract => contract_wit::contains_keys_new(keys),
246 WitInterface::Service => service_wit::contains_keys_new(keys),
247 #[cfg(with_testing)]
248 WitInterface::Mock { store, .. } => store.contains_keys_new(keys),
249 }
250 }
251
252 fn contains_keys_wait(&self, promise: u32) -> Vec<bool> {
254 match self {
255 WitInterface::Contract => contract_wit::contains_keys_wait(promise),
256 WitInterface::Service => service_wit::contains_keys_wait(promise),
257 #[cfg(with_testing)]
258 WitInterface::Mock { store, .. } => store.contains_keys_wait(promise),
259 }
260 }
261
262 fn read_multi_values_bytes_new(&self, keys: &[Vec<u8>]) -> u32 {
264 match self {
265 WitInterface::Contract => contract_wit::read_multi_values_bytes_new(keys),
266 WitInterface::Service => service_wit::read_multi_values_bytes_new(keys),
267 #[cfg(with_testing)]
268 WitInterface::Mock { store, .. } => store.read_multi_values_bytes_new(keys),
269 }
270 }
271
272 fn read_multi_values_bytes_wait(&self, promise: u32) -> Vec<Option<Vec<u8>>> {
274 match self {
275 WitInterface::Contract => contract_wit::read_multi_values_bytes_wait(promise),
276 WitInterface::Service => service_wit::read_multi_values_bytes_wait(promise),
277 #[cfg(with_testing)]
278 WitInterface::Mock { store, .. } => store.read_multi_values_bytes_wait(promise),
279 }
280 }
281
282 fn read_value_bytes_new(&self, key: &[u8]) -> u32 {
284 match self {
285 WitInterface::Contract => contract_wit::read_value_bytes_new(key),
286 WitInterface::Service => service_wit::read_value_bytes_new(key),
287 #[cfg(with_testing)]
288 WitInterface::Mock { store, .. } => store.read_value_bytes_new(key),
289 }
290 }
291
292 fn read_value_bytes_wait(&self, promise: u32) -> Option<Vec<u8>> {
294 match self {
295 WitInterface::Contract => contract_wit::read_value_bytes_wait(promise),
296 WitInterface::Service => service_wit::read_value_bytes_wait(promise),
297 #[cfg(with_testing)]
298 WitInterface::Mock { store, .. } => store.read_value_bytes_wait(promise),
299 }
300 }
301
302 fn find_keys_new(&self, key_prefix: &[u8]) -> u32 {
304 match self {
305 WitInterface::Contract => contract_wit::find_keys_new(key_prefix),
306 WitInterface::Service => service_wit::find_keys_new(key_prefix),
307 #[cfg(with_testing)]
308 WitInterface::Mock { store, .. } => store.find_keys_new(key_prefix),
309 }
310 }
311
312 fn find_keys_wait(&self, promise: u32) -> Vec<Vec<u8>> {
314 match self {
315 WitInterface::Contract => contract_wit::find_keys_wait(promise),
316 WitInterface::Service => service_wit::find_keys_wait(promise),
317 #[cfg(with_testing)]
318 WitInterface::Mock { store, .. } => store.find_keys_wait(promise),
319 }
320 }
321
322 fn find_key_values_new(&self, key_prefix: &[u8]) -> u32 {
324 match self {
325 WitInterface::Contract => contract_wit::find_key_values_new(key_prefix),
326 WitInterface::Service => service_wit::find_key_values_new(key_prefix),
327 #[cfg(with_testing)]
328 WitInterface::Mock { store, .. } => store.find_key_values_new(key_prefix),
329 }
330 }
331
332 fn find_key_values_wait(&self, promise: u32) -> Vec<(Vec<u8>, Vec<u8>)> {
334 match self {
335 WitInterface::Contract => contract_wit::find_key_values_wait(promise),
336 WitInterface::Service => service_wit::find_key_values_wait(promise),
337 #[cfg(with_testing)]
338 WitInterface::Mock { store, .. } => store.find_key_values_wait(promise),
339 }
340 }
341
342 fn write_batch(&self, batch: Batch) {
344 match self {
345 WitInterface::Contract => {
346 let batch_operations = batch
347 .operations
348 .into_iter()
349 .map(WriteOperation::from)
350 .collect::<Vec<_>>();
351
352 contract_runtime_api::write_batch(&batch_operations);
353 }
354 WitInterface::Service => panic!("Attempt to modify storage from a service"),
355 #[cfg(with_testing)]
356 WitInterface::Mock {
357 store,
358 read_only: false,
359 } => {
360 store.write_batch(batch);
361 }
362 #[cfg(with_testing)]
363 WitInterface::Mock {
364 read_only: true, ..
365 } => {
366 panic!("Attempt to modify storage from a service")
367 }
368 }
369 }
370}
371
372pub type ViewStorageContext = linera_views::context::ViewContext<(), KeyValueStore>;
375
376#[cfg(all(test, not(target_arch = "wasm32")))]
377mod tests {
378 use super::*;
379
380 #[tokio::test]
381 async fn test_key_value_store_mock() -> anyhow::Result<()> {
382 let store = KeyValueStore::mock();
384 let mock_store = store.to_mut();
385
386 let is_key_existing = mock_store.contains_key(b"foo").await?;
388 assert!(!is_key_existing);
389
390 let is_keys_existing = mock_store
392 .contains_keys(&[b"foo".to_vec(), b"bar".to_vec()])
393 .await?;
394 assert!(!is_keys_existing[0]);
395 assert!(!is_keys_existing[1]);
396
397 let mut batch = Batch::new();
399 batch.put_key_value(b"foo".to_vec(), &32_u128)?;
400 batch.put_key_value(b"bar".to_vec(), &42_u128)?;
401 mock_store.write_batch(batch).await?;
402
403 let is_key_existing = mock_store.contains_key(b"foo").await?;
404 assert!(is_key_existing);
405
406 let value = mock_store.read_value(b"foo").await?;
407 assert_eq!(value, Some(32_u128));
408
409 let value = mock_store.read_value(b"bar").await?;
410 assert_eq!(value, Some(42_u128));
411
412 Ok(())
413 }
414}