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};
25
26const MAX_KEY_SIZE: usize = 900;
33
34#[derive(Clone)]
36pub struct KeyValueStore {
37 wit_api: WitInterface,
38}
39
40#[cfg_attr(with_testing, allow(dead_code))]
41impl KeyValueStore {
42 pub(crate) fn for_contracts() -> Self {
44 KeyValueStore {
45 wit_api: WitInterface::Contract,
46 }
47 }
48
49 pub(crate) fn for_services() -> Self {
51 KeyValueStore {
52 wit_api: WitInterface::Service,
53 }
54 }
55
56 #[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 #[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#[derive(Error, Debug)]
90pub enum KeyValueStoreError {
91 #[error("Key too long")]
93 KeyTooLong,
94
95 #[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 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#[derive(Clone)]
200#[cfg_attr(with_testing, allow(dead_code))]
201enum WitInterface {
202 Contract,
204 Service,
206 #[cfg(with_testing)]
207 Mock {
209 store: Arc<MockKeyValueStore>,
210 read_only: bool,
211 },
212}
213
214impl WitInterface {
215 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 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 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 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 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 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 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 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 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 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 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 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 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
365pub 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 let store = KeyValueStore::mock();
377 let mock_store = store.to_mut();
378
379 let is_key_existing = mock_store.contains_key(b"foo").await?;
381 assert!(!is_key_existing);
382
383 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 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}