linera_sdk/views/
system_api.rs

1// Copyright (c) Zefchain Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Functions and types to interface with the system API available to application views.
5
6#[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
27/// We need to have a maximum key size that handles all possible underlying
28/// sizes. The constraint so far is DynamoDB which has a key length of 1024.
29/// That key length is decreased by 4 due to the use of a value splitting.
30/// Then the [`KeyValueStore`] needs to handle some base key and so we
31/// reduce to 900. Depending on the size, the error can occur in `system_api`
32/// or in the `KeyValueStoreView`.
33const MAX_KEY_SIZE: usize = 900;
34
35/// A type to interface with the key value storage provided to applications.
36#[derive(Clone)]
37pub struct KeyValueStore {
38    wit_api: WitInterface,
39}
40
41#[cfg_attr(with_testing, allow(dead_code))]
42impl KeyValueStore {
43    /// Returns a [`KeyValueStore`] that uses the contract WIT interface.
44    pub(crate) fn for_contracts() -> Self {
45        KeyValueStore {
46            wit_api: WitInterface::Contract,
47        }
48    }
49
50    /// Returns a [`KeyValueStore`] that uses the service WIT interface.
51    pub(crate) fn for_services() -> Self {
52        KeyValueStore {
53            wit_api: WitInterface::Service,
54        }
55    }
56
57    /// Returns a new [`KeyValueStore`] that just keeps the storage contents in memory.
58    #[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    /// Returns a mocked [`KeyValueStore`] that shares the memory storage with this instance but
69    /// allows write operations.
70    #[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/// The error type for [`KeyValueStore`] operations.
90#[derive(Error, Debug)]
91pub enum KeyValueStoreError {
92    /// Key too long
93    #[error("Key too long")]
94    KeyTooLong,
95
96    /// BCS serialization error.
97    #[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    // The KeyValueStore of the system_api does not have limits
107    // on the size of its values.
108    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/// Which system API should be used to interface with the storage.
204#[derive(Clone)]
205#[cfg_attr(with_testing, allow(dead_code))]
206enum WitInterface {
207    /// The contract system API.
208    Contract,
209    /// The service system API.
210    Service,
211    #[cfg(with_testing)]
212    /// A mock system API.
213    Mock {
214        store: Arc<MockKeyValueStore>,
215        read_only: bool,
216    },
217}
218
219impl WitInterface {
220    /// Creates a promise for testing if a key exist in the key-value store
221    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    /// Resolves a promise for testing if a key exist in the key-value store
231    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    /// Creates a promise for testing if multiple keys exist in the key-value store
241    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    /// Resolves a promise for testing if multiple keys exist in the key-value store
251    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    /// Creates a promise for reading multiple keys in the key-value store
261    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    /// Resolves a promise for reading multiple keys in the key-value store
271    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    /// Creates a promise for reading a key in the key-value store
281    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    /// Resolves a promise for reading a key in the key-value store
291    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    /// Creates a promise for finding keys having a specified prefix in the key-value store
301    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    /// Resolves a promise for finding keys having a specified prefix in the key-value store
311    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    /// Creates a promise for finding the key/values having a specified prefix in the key-value store
321    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    /// Resolves a promise for finding the key/values having a specified prefix in the key-value store
331    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    /// Calls the `write_batch` WIT function.
341    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
370/// Implementation of [`linera_views::context::Context`] to be used for data storage
371/// by Linera applications.
372pub 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        // Create a mock key-value store for testing
381        let store = KeyValueStore::mock();
382        let mock_store = store.to_mut();
383
384        // Check if key exists
385        let is_key_existing = mock_store.contains_key(b"foo").await?;
386        assert!(!is_key_existing);
387
388        // Check if keys exist
389        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        // Read and write values
396        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}