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
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/// Which system API should be used to interface with the storage.
206#[derive(Clone)]
207#[cfg_attr(with_testing, allow(dead_code))]
208enum WitInterface {
209    /// The contract system API.
210    Contract,
211    /// The service system API.
212    Service,
213    #[cfg(with_testing)]
214    /// A mock system API.
215    Mock {
216        store: Arc<MockKeyValueStore>,
217        read_only: bool,
218    },
219}
220
221impl WitInterface {
222    /// Creates a promise for testing if a key exist in the key-value store
223    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    /// Resolves a promise for testing if a key exist in the key-value store
233    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    /// Creates a promise for testing if multiple keys exist in the key-value store
243    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    /// Resolves a promise for testing if multiple keys exist in the key-value store
253    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    /// Creates a promise for reading multiple keys in the key-value store
263    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    /// Resolves a promise for reading multiple keys in the key-value store
273    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    /// Creates a promise for reading a key in the key-value store
283    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    /// Resolves a promise for reading a key in the key-value store
293    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    /// Creates a promise for finding keys having a specified prefix in the key-value store
303    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    /// Resolves a promise for finding keys having a specified prefix in the key-value store
313    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    /// Creates a promise for finding the key/values having a specified prefix in the key-value store
323    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    /// Resolves a promise for finding the key/values having a specified prefix in the key-value store
333    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    /// Calls the `write_batch` WIT function.
343    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
372/// Implementation of [`linera_views::context::Context`] to be used for data storage
373/// by Linera applications.
374pub 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        // Create a mock key-value store for testing
383        let store = KeyValueStore::mock();
384        let mock_store = store.to_mut();
385
386        // Check if key exists
387        let is_key_existing = mock_store.contains_key(b"foo").await?;
388        assert!(!is_key_existing);
389
390        // Check if keys exist
391        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        // Read and write values
398        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}