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};
25
26/// We need to have a maximum key size that handles all possible underlying
27/// sizes. The constraint so far is DynamoDB which has a key length of 1024.
28/// That key length is decreased by 4 due to the use of a value splitting.
29/// Then the [`KeyValueStore`] needs to handle some base key and so we
30/// reduce to 900. Depending on the size, the error can occur in `system_api`
31/// or in the `KeyValueStoreView`.
32const MAX_KEY_SIZE: usize = 900;
33
34/// A type to interface with the key value storage provided to applications.
35#[derive(Clone)]
36pub struct KeyValueStore {
37    wit_api: WitInterface,
38}
39
40#[cfg_attr(with_testing, allow(dead_code))]
41impl KeyValueStore {
42    /// Returns a [`KeyValueStore`] that uses the contract WIT interface.
43    pub(crate) fn for_contracts() -> Self {
44        KeyValueStore {
45            wit_api: WitInterface::Contract,
46        }
47    }
48
49    /// Returns a [`KeyValueStore`] that uses the service WIT interface.
50    pub(crate) fn for_services() -> Self {
51        KeyValueStore {
52            wit_api: WitInterface::Service,
53        }
54    }
55
56    /// Returns a new [`KeyValueStore`] that just keeps the storage contents in memory.
57    #[cfg(with_testing)]
58    pub fn mock() -> Self {
59        KeyValueStore {
60            wit_api: WitInterface::Mock {
61                store: Arc::new(MockKeyValueStore::default()),
62                read_only: true,
63            },
64        }
65    }
66
67    /// Returns a mocked [`KeyValueStore`] that shares the memory storage with this instance but
68    /// allows write operations.
69    #[cfg(with_testing)]
70    pub fn to_mut(&self) -> Self {
71        let WitInterface::Mock { store, .. } = &self.wit_api else {
72            panic!("Real `KeyValueStore` should not be used in unit tests");
73        };
74
75        KeyValueStore {
76            wit_api: WitInterface::Mock {
77                store: store.clone(),
78                read_only: false,
79            },
80        }
81    }
82}
83
84impl WithError for KeyValueStore {
85    type Error = KeyValueStoreError;
86}
87
88/// The error type for [`KeyValueStore`] operations.
89#[derive(Error, Debug)]
90pub enum KeyValueStoreError {
91    /// Key too long
92    #[error("Key too long")]
93    KeyTooLong,
94
95    /// BCS serialization error.
96    #[error(transparent)]
97    BcsError(#[from] bcs::Error),
98}
99
100impl linera_views::store::KeyValueStoreError for KeyValueStoreError {
101    const BACKEND: &'static str = "key_value_store";
102}
103
104impl ReadableKeyValueStore for KeyValueStore {
105    // The KeyValueStore of the system_api does not have limits
106    // on the size of its values.
107    const MAX_KEY_SIZE: usize = MAX_KEY_SIZE;
108
109    fn max_stream_queries(&self) -> usize {
110        1
111    }
112
113    fn root_key(&self) -> Result<Vec<u8>, KeyValueStoreError> {
114        Ok(Vec::new())
115    }
116
117    async fn contains_key(&self, key: &[u8]) -> Result<bool, KeyValueStoreError> {
118        ensure!(
119            key.len() <= Self::MAX_KEY_SIZE,
120            KeyValueStoreError::KeyTooLong
121        );
122        let promise = self.wit_api.contains_key_new(key);
123        Ok(self.wit_api.contains_key_wait(promise))
124    }
125
126    async fn contains_keys(&self, keys: &[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        Ok(self.wit_api.contains_keys_wait(promise))
135    }
136
137    async fn read_multi_values_bytes(
138        &self,
139        keys: &[Vec<u8>],
140    ) -> Result<Vec<Option<Vec<u8>>>, KeyValueStoreError> {
141        for key in keys {
142            ensure!(
143                key.len() <= Self::MAX_KEY_SIZE,
144                KeyValueStoreError::KeyTooLong
145            );
146        }
147        let promise = self.wit_api.read_multi_values_bytes_new(keys);
148        Ok(self.wit_api.read_multi_values_bytes_wait(promise))
149    }
150
151    async fn read_value_bytes(&self, key: &[u8]) -> Result<Option<Vec<u8>>, KeyValueStoreError> {
152        ensure!(
153            key.len() <= Self::MAX_KEY_SIZE,
154            KeyValueStoreError::KeyTooLong
155        );
156        let promise = self.wit_api.read_value_bytes_new(key);
157        Ok(self.wit_api.read_value_bytes_wait(promise))
158    }
159
160    async fn find_keys_by_prefix(
161        &self,
162        key_prefix: &[u8],
163    ) -> Result<Vec<Vec<u8>>, KeyValueStoreError> {
164        ensure!(
165            key_prefix.len() <= Self::MAX_KEY_SIZE,
166            KeyValueStoreError::KeyTooLong
167        );
168        let promise = self.wit_api.find_keys_new(key_prefix);
169        Ok(self.wit_api.find_keys_wait(promise))
170    }
171
172    async fn find_key_values_by_prefix(
173        &self,
174        key_prefix: &[u8],
175    ) -> Result<Vec<(Vec<u8>, Vec<u8>)>, KeyValueStoreError> {
176        ensure!(
177            key_prefix.len() <= Self::MAX_KEY_SIZE,
178            KeyValueStoreError::KeyTooLong
179        );
180        let promise = self.wit_api.find_key_values_new(key_prefix);
181        Ok(self.wit_api.find_key_values_wait(promise))
182    }
183}
184
185impl WritableKeyValueStore for KeyValueStore {
186    const MAX_VALUE_SIZE: usize = usize::MAX;
187
188    async fn write_batch(&self, batch: Batch) -> Result<(), KeyValueStoreError> {
189        self.wit_api.write_batch(batch);
190        Ok(())
191    }
192
193    async fn clear_journal(&self) -> Result<(), KeyValueStoreError> {
194        Ok(())
195    }
196}
197
198/// Which system API should be used to interface with the storage.
199#[derive(Clone)]
200#[cfg_attr(with_testing, allow(dead_code))]
201enum WitInterface {
202    /// The contract system API.
203    Contract,
204    /// The service system API.
205    Service,
206    #[cfg(with_testing)]
207    /// A mock system API.
208    Mock {
209        store: Arc<MockKeyValueStore>,
210        read_only: bool,
211    },
212}
213
214impl WitInterface {
215    /// Creates a promise for testing if a key exist in the key-value store
216    fn contains_key_new(&self, key: &[u8]) -> u32 {
217        match self {
218            WitInterface::Contract => contract_wit::contains_key_new(key),
219            WitInterface::Service => service_wit::contains_key_new(key),
220            #[cfg(with_testing)]
221            WitInterface::Mock { store, .. } => store.contains_key_new(key),
222        }
223    }
224
225    /// Resolves a promise for testing if a key exist in the key-value store
226    fn contains_key_wait(&self, promise: u32) -> bool {
227        match self {
228            WitInterface::Contract => contract_wit::contains_key_wait(promise),
229            WitInterface::Service => service_wit::contains_key_wait(promise),
230            #[cfg(with_testing)]
231            WitInterface::Mock { store, .. } => store.contains_key_wait(promise),
232        }
233    }
234
235    /// Creates a promise for testing if multiple keys exist in the key-value store
236    fn contains_keys_new(&self, keys: &[Vec<u8>]) -> u32 {
237        match self {
238            WitInterface::Contract => contract_wit::contains_keys_new(keys),
239            WitInterface::Service => service_wit::contains_keys_new(keys),
240            #[cfg(with_testing)]
241            WitInterface::Mock { store, .. } => store.contains_keys_new(keys),
242        }
243    }
244
245    /// Resolves a promise for testing if multiple keys exist in the key-value store
246    fn contains_keys_wait(&self, promise: u32) -> Vec<bool> {
247        match self {
248            WitInterface::Contract => contract_wit::contains_keys_wait(promise),
249            WitInterface::Service => service_wit::contains_keys_wait(promise),
250            #[cfg(with_testing)]
251            WitInterface::Mock { store, .. } => store.contains_keys_wait(promise),
252        }
253    }
254
255    /// Creates a promise for reading multiple keys in the key-value store
256    fn read_multi_values_bytes_new(&self, keys: &[Vec<u8>]) -> u32 {
257        match self {
258            WitInterface::Contract => contract_wit::read_multi_values_bytes_new(keys),
259            WitInterface::Service => service_wit::read_multi_values_bytes_new(keys),
260            #[cfg(with_testing)]
261            WitInterface::Mock { store, .. } => store.read_multi_values_bytes_new(keys),
262        }
263    }
264
265    /// Resolves a promise for reading multiple keys in the key-value store
266    fn read_multi_values_bytes_wait(&self, promise: u32) -> Vec<Option<Vec<u8>>> {
267        match self {
268            WitInterface::Contract => contract_wit::read_multi_values_bytes_wait(promise),
269            WitInterface::Service => service_wit::read_multi_values_bytes_wait(promise),
270            #[cfg(with_testing)]
271            WitInterface::Mock { store, .. } => store.read_multi_values_bytes_wait(promise),
272        }
273    }
274
275    /// Creates a promise for reading a key in the key-value store
276    fn read_value_bytes_new(&self, key: &[u8]) -> u32 {
277        match self {
278            WitInterface::Contract => contract_wit::read_value_bytes_new(key),
279            WitInterface::Service => service_wit::read_value_bytes_new(key),
280            #[cfg(with_testing)]
281            WitInterface::Mock { store, .. } => store.read_value_bytes_new(key),
282        }
283    }
284
285    /// Resolves a promise for reading a key in the key-value store
286    fn read_value_bytes_wait(&self, promise: u32) -> Option<Vec<u8>> {
287        match self {
288            WitInterface::Contract => contract_wit::read_value_bytes_wait(promise),
289            WitInterface::Service => service_wit::read_value_bytes_wait(promise),
290            #[cfg(with_testing)]
291            WitInterface::Mock { store, .. } => store.read_value_bytes_wait(promise),
292        }
293    }
294
295    /// Creates a promise for finding keys having a specified prefix in the key-value store
296    fn find_keys_new(&self, key_prefix: &[u8]) -> u32 {
297        match self {
298            WitInterface::Contract => contract_wit::find_keys_new(key_prefix),
299            WitInterface::Service => service_wit::find_keys_new(key_prefix),
300            #[cfg(with_testing)]
301            WitInterface::Mock { store, .. } => store.find_keys_new(key_prefix),
302        }
303    }
304
305    /// Resolves a promise for finding keys having a specified prefix in the key-value store
306    fn find_keys_wait(&self, promise: u32) -> Vec<Vec<u8>> {
307        match self {
308            WitInterface::Contract => contract_wit::find_keys_wait(promise),
309            WitInterface::Service => service_wit::find_keys_wait(promise),
310            #[cfg(with_testing)]
311            WitInterface::Mock { store, .. } => store.find_keys_wait(promise),
312        }
313    }
314
315    /// Creates a promise for finding the key/values having a specified prefix in the key-value store
316    fn find_key_values_new(&self, key_prefix: &[u8]) -> u32 {
317        match self {
318            WitInterface::Contract => contract_wit::find_key_values_new(key_prefix),
319            WitInterface::Service => service_wit::find_key_values_new(key_prefix),
320            #[cfg(with_testing)]
321            WitInterface::Mock { store, .. } => store.find_key_values_new(key_prefix),
322        }
323    }
324
325    /// Resolves a promise for finding the key/values having a specified prefix in the key-value store
326    fn find_key_values_wait(&self, promise: u32) -> Vec<(Vec<u8>, Vec<u8>)> {
327        match self {
328            WitInterface::Contract => contract_wit::find_key_values_wait(promise),
329            WitInterface::Service => service_wit::find_key_values_wait(promise),
330            #[cfg(with_testing)]
331            WitInterface::Mock { store, .. } => store.find_key_values_wait(promise),
332        }
333    }
334
335    /// Calls the `write_batch` WIT function.
336    fn write_batch(&self, batch: Batch) {
337        match self {
338            WitInterface::Contract => {
339                let batch_operations = batch
340                    .operations
341                    .into_iter()
342                    .map(WriteOperation::from)
343                    .collect::<Vec<_>>();
344
345                contract_runtime_api::write_batch(&batch_operations);
346            }
347            WitInterface::Service => panic!("Attempt to modify storage from a service"),
348            #[cfg(with_testing)]
349            WitInterface::Mock {
350                store,
351                read_only: false,
352            } => {
353                store.write_batch(batch);
354            }
355            #[cfg(with_testing)]
356            WitInterface::Mock {
357                read_only: true, ..
358            } => {
359                panic!("Attempt to modify storage from a service")
360            }
361        }
362    }
363}
364
365/// Implementation of [`linera_views::context::Context`] to be used for data storage
366/// by Linera applications.
367pub type ViewStorageContext = linera_views::context::ViewContext<(), KeyValueStore>;
368
369#[cfg(all(test, not(target_arch = "wasm32")))]
370mod tests {
371    use super::*;
372
373    #[tokio::test]
374    async fn test_key_value_store_mock() -> anyhow::Result<()> {
375        // Create a mock key-value store for testing
376        let store = KeyValueStore::mock();
377        let mock_store = store.to_mut();
378
379        // Check if key exists
380        let is_key_existing = mock_store.contains_key(b"foo").await?;
381        assert!(!is_key_existing);
382
383        // Check if keys exist
384        let is_keys_existing = mock_store
385            .contains_keys(&[b"foo".to_vec(), b"bar".to_vec()])
386            .await?;
387        assert!(!is_keys_existing[0]);
388        assert!(!is_keys_existing[1]);
389
390        // Read and write values
391        let mut batch = Batch::new();
392        batch.put_key_value(b"foo".to_vec(), &32_u128)?;
393        batch.put_key_value(b"bar".to_vec(), &42_u128)?;
394        mock_store.write_batch(batch).await?;
395
396        let is_key_existing = mock_store.contains_key(b"foo").await?;
397        assert!(is_key_existing);
398
399        let value = mock_store.read_value(b"foo").await?;
400        assert_eq!(value, Some(32_u128));
401
402        let value = mock_store.read_value(b"bar").await?;
403        assert_eq!(value, Some(42_u128));
404
405        Ok(())
406    }
407}