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 type Keys = Vec<Vec<u8>>;
110 type KeyValues = Vec<(Vec<u8>, Vec<u8>)>;
111
112 fn max_stream_queries(&self) -> usize {
113 1
114 }
115
116 async fn contains_key(&self, key: &[u8]) -> Result<bool, KeyValueStoreError> {
117 ensure!(
118 key.len() <= Self::MAX_KEY_SIZE,
119 KeyValueStoreError::KeyTooLong
120 );
121 let promise = self.wit_api.contains_key_new(key);
122 yield_once().await;
123 Ok(self.wit_api.contains_key_wait(promise))
124 }
125
126 async fn contains_keys(&self, keys: Vec<Vec<u8>>) -> Result<Vec<bool>, KeyValueStoreError> {
127 for key in &keys {
128 ensure!(
129 key.len() <= Self::MAX_KEY_SIZE,
130 KeyValueStoreError::KeyTooLong
131 );
132 }
133 let promise = self.wit_api.contains_keys_new(&keys);
134 yield_once().await;
135 Ok(self.wit_api.contains_keys_wait(promise))
136 }
137
138 async fn read_multi_values_bytes(
139 &self,
140 keys: Vec<Vec<u8>>,
141 ) -> Result<Vec<Option<Vec<u8>>>, KeyValueStoreError> {
142 for key in &keys {
143 ensure!(
144 key.len() <= Self::MAX_KEY_SIZE,
145 KeyValueStoreError::KeyTooLong
146 );
147 }
148 let promise = self.wit_api.read_multi_values_bytes_new(&keys);
149 yield_once().await;
150 Ok(self.wit_api.read_multi_values_bytes_wait(promise))
151 }
152
153 async fn read_value_bytes(&self, key: &[u8]) -> Result<Option<Vec<u8>>, KeyValueStoreError> {
154 ensure!(
155 key.len() <= Self::MAX_KEY_SIZE,
156 KeyValueStoreError::KeyTooLong
157 );
158 let promise = self.wit_api.read_value_bytes_new(key);
159 yield_once().await;
160 Ok(self.wit_api.read_value_bytes_wait(promise))
161 }
162
163 async fn find_keys_by_prefix(
164 &self,
165 key_prefix: &[u8],
166 ) -> Result<Self::Keys, KeyValueStoreError> {
167 ensure!(
168 key_prefix.len() <= Self::MAX_KEY_SIZE,
169 KeyValueStoreError::KeyTooLong
170 );
171 let promise = self.wit_api.find_keys_new(key_prefix);
172 yield_once().await;
173 Ok(self.wit_api.find_keys_wait(promise))
174 }
175
176 async fn find_key_values_by_prefix(
177 &self,
178 key_prefix: &[u8],
179 ) -> Result<Self::KeyValues, KeyValueStoreError> {
180 ensure!(
181 key_prefix.len() <= Self::MAX_KEY_SIZE,
182 KeyValueStoreError::KeyTooLong
183 );
184 let promise = self.wit_api.find_key_values_new(key_prefix);
185 yield_once().await;
186 Ok(self.wit_api.find_key_values_wait(promise))
187 }
188}
189
190impl WritableKeyValueStore for KeyValueStore {
191 const MAX_VALUE_SIZE: usize = usize::MAX;
192
193 async fn write_batch(&self, batch: Batch) -> Result<(), KeyValueStoreError> {
194 self.wit_api.write_batch(batch);
195 Ok(())
196 }
197
198 async fn clear_journal(&self) -> Result<(), KeyValueStoreError> {
199 Ok(())
200 }
201}
202
203#[derive(Clone)]
205#[cfg_attr(with_testing, allow(dead_code))]
206enum WitInterface {
207 Contract,
209 Service,
211 #[cfg(with_testing)]
212 Mock {
214 store: Arc<MockKeyValueStore>,
215 read_only: bool,
216 },
217}
218
219impl WitInterface {
220 fn contains_key_new(&self, key: &[u8]) -> u32 {
222 match self {
223 WitInterface::Contract => contract_wit::contains_key_new(key),
224 WitInterface::Service => service_wit::contains_key_new(key),
225 #[cfg(with_testing)]
226 WitInterface::Mock { store, .. } => store.contains_key_new(key),
227 }
228 }
229
230 fn contains_key_wait(&self, promise: u32) -> bool {
232 match self {
233 WitInterface::Contract => contract_wit::contains_key_wait(promise),
234 WitInterface::Service => service_wit::contains_key_wait(promise),
235 #[cfg(with_testing)]
236 WitInterface::Mock { store, .. } => store.contains_key_wait(promise),
237 }
238 }
239
240 fn contains_keys_new(&self, keys: &[Vec<u8>]) -> u32 {
242 match self {
243 WitInterface::Contract => contract_wit::contains_keys_new(keys),
244 WitInterface::Service => service_wit::contains_keys_new(keys),
245 #[cfg(with_testing)]
246 WitInterface::Mock { store, .. } => store.contains_keys_new(keys),
247 }
248 }
249
250 fn contains_keys_wait(&self, promise: u32) -> Vec<bool> {
252 match self {
253 WitInterface::Contract => contract_wit::contains_keys_wait(promise),
254 WitInterface::Service => service_wit::contains_keys_wait(promise),
255 #[cfg(with_testing)]
256 WitInterface::Mock { store, .. } => store.contains_keys_wait(promise),
257 }
258 }
259
260 fn read_multi_values_bytes_new(&self, keys: &[Vec<u8>]) -> u32 {
262 match self {
263 WitInterface::Contract => contract_wit::read_multi_values_bytes_new(keys),
264 WitInterface::Service => service_wit::read_multi_values_bytes_new(keys),
265 #[cfg(with_testing)]
266 WitInterface::Mock { store, .. } => store.read_multi_values_bytes_new(keys),
267 }
268 }
269
270 fn read_multi_values_bytes_wait(&self, promise: u32) -> Vec<Option<Vec<u8>>> {
272 match self {
273 WitInterface::Contract => contract_wit::read_multi_values_bytes_wait(promise),
274 WitInterface::Service => service_wit::read_multi_values_bytes_wait(promise),
275 #[cfg(with_testing)]
276 WitInterface::Mock { store, .. } => store.read_multi_values_bytes_wait(promise),
277 }
278 }
279
280 fn read_value_bytes_new(&self, key: &[u8]) -> u32 {
282 match self {
283 WitInterface::Contract => contract_wit::read_value_bytes_new(key),
284 WitInterface::Service => service_wit::read_value_bytes_new(key),
285 #[cfg(with_testing)]
286 WitInterface::Mock { store, .. } => store.read_value_bytes_new(key),
287 }
288 }
289
290 fn read_value_bytes_wait(&self, promise: u32) -> Option<Vec<u8>> {
292 match self {
293 WitInterface::Contract => contract_wit::read_value_bytes_wait(promise),
294 WitInterface::Service => service_wit::read_value_bytes_wait(promise),
295 #[cfg(with_testing)]
296 WitInterface::Mock { store, .. } => store.read_value_bytes_wait(promise),
297 }
298 }
299
300 fn find_keys_new(&self, key_prefix: &[u8]) -> u32 {
302 match self {
303 WitInterface::Contract => contract_wit::find_keys_new(key_prefix),
304 WitInterface::Service => service_wit::find_keys_new(key_prefix),
305 #[cfg(with_testing)]
306 WitInterface::Mock { store, .. } => store.find_keys_new(key_prefix),
307 }
308 }
309
310 fn find_keys_wait(&self, promise: u32) -> Vec<Vec<u8>> {
312 match self {
313 WitInterface::Contract => contract_wit::find_keys_wait(promise),
314 WitInterface::Service => service_wit::find_keys_wait(promise),
315 #[cfg(with_testing)]
316 WitInterface::Mock { store, .. } => store.find_keys_wait(promise),
317 }
318 }
319
320 fn find_key_values_new(&self, key_prefix: &[u8]) -> u32 {
322 match self {
323 WitInterface::Contract => contract_wit::find_key_values_new(key_prefix),
324 WitInterface::Service => service_wit::find_key_values_new(key_prefix),
325 #[cfg(with_testing)]
326 WitInterface::Mock { store, .. } => store.find_key_values_new(key_prefix),
327 }
328 }
329
330 fn find_key_values_wait(&self, promise: u32) -> Vec<(Vec<u8>, Vec<u8>)> {
332 match self {
333 WitInterface::Contract => contract_wit::find_key_values_wait(promise),
334 WitInterface::Service => service_wit::find_key_values_wait(promise),
335 #[cfg(with_testing)]
336 WitInterface::Mock { store, .. } => store.find_key_values_wait(promise),
337 }
338 }
339
340 fn write_batch(&self, batch: Batch) {
342 match self {
343 WitInterface::Contract => {
344 let batch_operations = batch
345 .operations
346 .into_iter()
347 .map(WriteOperation::from)
348 .collect::<Vec<_>>();
349
350 contract_runtime_api::write_batch(&batch_operations);
351 }
352 WitInterface::Service => panic!("Attempt to modify storage from a service"),
353 #[cfg(with_testing)]
354 WitInterface::Mock {
355 store,
356 read_only: false,
357 } => {
358 store.write_batch(batch);
359 }
360 #[cfg(with_testing)]
361 WitInterface::Mock {
362 read_only: true, ..
363 } => {
364 panic!("Attempt to modify storage from a service")
365 }
366 }
367 }
368}
369
370pub type ViewStorageContext = linera_views::context::ViewContext<(), KeyValueStore>;
373
374#[cfg(all(test, not(target_arch = "wasm32")))]
375mod tests {
376 use super::*;
377
378 #[tokio::test]
379 async fn test_key_value_store_mock() -> anyhow::Result<()> {
380 let store = KeyValueStore::mock();
382 let mock_store = store.to_mut();
383
384 let is_key_existing = mock_store.contains_key(b"foo").await?;
386 assert!(!is_key_existing);
387
388 let is_keys_existing = mock_store
390 .contains_keys(vec![b"foo".to_vec(), b"bar".to_vec()])
391 .await?;
392 assert!(!is_keys_existing[0]);
393 assert!(!is_keys_existing[1]);
394
395 let mut batch = Batch::new();
397 batch.put_key_value(b"foo".to_vec(), &32_u128)?;
398 batch.put_key_value(b"bar".to_vec(), &42_u128)?;
399 mock_store.write_batch(batch).await?;
400
401 let is_key_existing = mock_store.contains_key(b"foo").await?;
402 assert!(is_key_existing);
403
404 let value = mock_store.read_value(b"foo").await?;
405 assert_eq!(value, Some(32_u128));
406
407 let value = mock_store.read_value(b"bar").await?;
408 assert_eq!(value, Some(42_u128));
409
410 Ok(())
411 }
412}