Support wasi-keyvalue for key-value store access

Signed-off-by: itowlson <[email protected]>

Author: itowlson

crates/factor-key-value/src/host.rs

@@ -2,6 +2,7 @@ use anyhow::{Context, Result};
 use spin_core::{async_trait, wasmtime::component::Resource};
 use spin_resource_table::Table;
 use spin_world::v2::key_value;
+use spin_world::wasi::keyvalue as wasi_keyvalue;
 use std::{collections::HashSet, sync::Arc};
 use tracing::{instrument, Level};
 
@@ -55,13 +56,22 @@ impl KeyValueDispatch {
         }
     }
 
-    pub fn get_store(&self, store: Resource<key_value::Store>) -> anyhow::Result<&Arc<dyn Store>> {
+    pub fn get_store<T: 'static>(&self, store: Resource<T>) -> anyhow::Result<&Arc<dyn Store>> {
         self.stores.get(store.rep()).context("invalid store")
     }
 
     pub fn allowed_stores(&self) -> &HashSet<String> {
         &self.allowed_stores
     }
+
+    pub fn get_store_wasi<T: 'static>(
+        &self,
+        store: Resource<T>,
+    ) -> Result<&Arc<dyn Store>, wasi_keyvalue::store::Error> {
+        self.stores
+            .get(store.rep())
+            .ok_or(wasi_keyvalue::store::Error::NoSuchStore)
+    }
 }
 
 #[async_trait]
@@ -141,6 +151,102 @@ impl key_value::HostStore for KeyValueDispatch {
     }
 }
 
+fn to_wasi_err(e: Error) -> wasi_keyvalue::store::Error {
+    match e {
+        Error::AccessDenied => wasi_keyvalue::store::Error::AccessDenied,
+        Error::NoSuchStore => wasi_keyvalue::store::Error::NoSuchStore,
+        Error::StoreTableFull => wasi_keyvalue::store::Error::Other("store table full".to_string()),
+        Error::Other(msg) => wasi_keyvalue::store::Error::Other(msg),
+    }
+}
+
+#[async_trait]
+impl wasi_keyvalue::store::Host for KeyValueDispatch {
+    async fn open(
+        &mut self,
+        identifier: String,
+    ) -> Result<Resource<wasi_keyvalue::store::Bucket>, wasi_keyvalue::store::Error> {
+        if self.allowed_stores.contains(&identifier) {
+            let store = self
+                .stores
+                .push(self.manager.get(&identifier).await.map_err(to_wasi_err)?)
+                .map_err(|()| wasi_keyvalue::store::Error::Other("store table full".to_string()))?;
+            Ok(Resource::new_own(store))
+        } else {
+            Err(wasi_keyvalue::store::Error::AccessDenied)
+        }
+    }
+
+    fn convert_error(
+        &mut self,
+        error: spin_world::wasi::keyvalue::store::Error,
+    ) -> std::result::Result<spin_world::wasi::keyvalue::store::Error, anyhow::Error> {
+        Ok(error)
+    }
+}
+
+use wasi_keyvalue::store::Bucket;
+#[async_trait]
+impl wasi_keyvalue::store::HostBucket for KeyValueDispatch {
+    async fn get(
+        &mut self,
+        self_: Resource<Bucket>,
+        key: String,
+    ) -> Result<Option<Vec<u8>>, wasi_keyvalue::store::Error> {
+        let store = self.get_store_wasi(self_)?;
+        store.get(&key).await.map_err(to_wasi_err)
+    }
+
+    async fn set(
+        &mut self,
+        self_: Resource<Bucket>,
+        key: String,
+        value: Vec<u8>,
+    ) -> Result<(), wasi_keyvalue::store::Error> {
+        let store = self.get_store_wasi(self_)?;
+        store.set(&key, &value).await.map_err(to_wasi_err)
+    }
+
+    async fn delete(
+        &mut self,
+        self_: Resource<Bucket>,
+        key: String,
+    ) -> Result<(), wasi_keyvalue::store::Error> {
+        let store = self.get_store_wasi(self_)?;
+        store.delete(&key).await.map_err(to_wasi_err)
+    }
+
+    async fn exists(
+        &mut self,
+        self_: Resource<Bucket>,
+        key: String,
+    ) -> Result<bool, wasi_keyvalue::store::Error> {
+        let store = self.get_store_wasi(self_)?;
+        store.exists(&key).await.map_err(to_wasi_err)
+    }
+
+    async fn list_keys(
+        &mut self,
+        self_: Resource<Bucket>,
+        cursor: Option<u64>,
+    ) -> Result<wasi_keyvalue::store::KeyResponse, wasi_keyvalue::store::Error> {
+        if cursor.unwrap_or_default() != 0 {
+            return Err(wasi_keyvalue::store::Error::Other(
+                "list_keys: cursor not supported".to_owned(),
+            ));
+        }
+
+        let store = self.get_store_wasi(self_)?;
+        let keys = store.get_keys().await.map_err(to_wasi_err)?;
+        Ok(wasi_keyvalue::store::KeyResponse { keys, cursor: None })
+    }
+
+    async fn drop(&mut self, rep: Resource<Bucket>) -> anyhow::Result<()> {
+        self.stores.remove(rep.rep());
+        Ok(())
+    }
+}
+
 pub fn log_error(err: impl std::fmt::Debug) -> Error {
     tracing::warn!("key-value error: {err:?}");
     Error::Other(format!("{err:?}"))

crates/factor-key-value/src/lib.rs

@@ -40,6 +40,7 @@ impl Factor for KeyValueFactor {
     fn init<T: Send + 'static>(&mut self, mut ctx: InitContext<T, Self>) -> anyhow::Result<()> {
         ctx.link_bindings(spin_world::v1::key_value::add_to_linker)?;
         ctx.link_bindings(spin_world::v2::key_value::add_to_linker)?;
+        ctx.link_bindings(spin_world::wasi::keyvalue::store::add_to_linker)?;
         Ok(())
     }
 

crates/world/src/lib.rs

@@ -27,6 +27,7 @@ wasmtime::component::bindgen!({
         "fermyon:spin/[email protected]/error" => v2::sqlite::Error,
         "fermyon:spin/sqlite/error" => v1::sqlite::Error,
         "fermyon:spin/[email protected]/error" => v2::variables::Error,
+        "wasi:keyvalue/store/error" => wasi::keyvalue::store::Error,
     },
     trappable_imports: true,
 });

tests/runtime-tests/tests/key-value-wasi/spin.toml

@@ -0,0 +1,14 @@
+spin_manifest_version = 2
+
+[application]
+name = "key-value-wasi"
+authors = ["Fermyon Engineering <[email protected]>"]
+version = "0.1.0"
+
+[[trigger.http]]
+route = "/"
+component = "test"
+
+[component.test]
+source = "%{source=key-value-wasi}"
+key_value_stores = ["default"]

tests/test-components/components/Cargo.lock

@@ -0,0 +1,11 @@
+[package]
+name = "key-value-wasi"
+version = "0.1.0"
+edition = "2021"
+
+[lib]
+crate-type = ["cdylib"]
+
+[dependencies]
+helper = { path = "../../helper" }
+wit-bindgen = "0.16.0"

tests/test-components/components/key-value-wasi/Cargo.toml

@@ -0,0 +1,10 @@
+# Key Value
+
+Tests the key/value interface.
+
+## Expectations
+
+This test component expects the following to be true:
+* It is given permission to open a connection to the "default" store.
+* It does not have permission to access a store named "forbidden".
+* It is empty

tests/test-components/components/key-value-wasi/README.md

@@ -0,0 +1,51 @@
+use helper::{ensure_matches, ensure_ok};
+
+use bindings::wasi::keyvalue::store::{Error, open, KeyResponse};
+
+helper::define_component!(Component);
+
+impl Component {
+    fn main() -> Result<(), String> {
+
+        ensure_matches!(open("forbidden"), Err(Error::AccessDenied));
+
+        let store = ensure_ok!(open("default"));
+
+        // Ensure nothing set in `bar` key
+        ensure_ok!(store.delete("bar"));
+        ensure_matches!(store.exists("bar"), Ok(false));
+        ensure_matches!(store.get("bar"), Ok(None));
+        ensure_matches!(keys(&store.list_keys(None)), Ok(&[]));
+
+        // Set `bar` key
+        ensure_ok!(store.set("bar", b"baz"));
+        ensure_matches!(store.exists("bar"), Ok(true));
+        ensure_matches!(store.get("bar"), Ok(Some(v)) if v == b"baz");
+        ensure_matches!(keys(&store.list_keys(None)), Ok([bar]) if bar == "bar");
+        ensure_matches!(keys(&store.list_keys(Some(0))), Ok([bar]) if bar == "bar");
+
+        // Override `bar` key
+        ensure_ok!(store.set("bar", b"wow"));
+        ensure_matches!(store.exists("bar"), Ok(true));
+        ensure_matches!(store.get("bar"), Ok(Some(wow)) if wow == b"wow");
+        ensure_matches!(keys(&store.list_keys(None)), Ok([bar]) if bar == "bar");
+
+        // Set another key
+        ensure_ok!(store.set("qux", b"yay"));
+        ensure_matches!(keys(&store.list_keys(None)), Ok(c) if c.len() == 2 && c.contains(&"bar".into()) && c.contains(&"qux".into()));
+
+        // Delete everything
+        ensure_ok!(store.delete("bar"));
+        ensure_ok!(store.delete("bar"));
+        ensure_ok!(store.delete("qux"));
+        ensure_matches!(store.exists("bar"), Ok(false));
+        ensure_matches!(store.get("qux"), Ok(None));
+        ensure_matches!(keys(&store.list_keys(None)), Ok(&[]));
+
+        Ok(())
+    }
+}
+
+fn keys<E>(res: &Result<KeyResponse, E>) -> Result<&[String], &E> {
+    res.as_ref().map(|kr| kr.keys.as_slice())
+}

tests/test-components/components/key-value-wasi/src/lib.rs

@@ -0,0 +1,22 @@
+/// A keyvalue interface that provides atomic operations.
+/// 
+/// Atomic operations are single, indivisible operations. When a fault causes an atomic operation to
+/// fail, it will appear to the invoker of the atomic operation that the action either completed
+/// successfully or did nothing at all.
+/// 
+/// Please note that this interface is bare functions that take a reference to a bucket. This is to
+/// get around the current lack of a way to "extend" a resource with additional methods inside of
+/// wit. Future version of the interface will instead extend these methods on the base `bucket`
+/// resource.
+interface atomics {
+  	use store.{bucket, error};
+
+  	/// Atomically increment the value associated with the key in the store by the given delta. It
+	/// returns the new value.
+	///
+	/// If the key does not exist in the store, it creates a new key-value pair with the value set
+	/// to the given delta. 
+	///
+	/// If any other error occurs, it returns an `Err(error)`.
+	increment: func(bucket: borrow<bucket>, key: string, delta: u64) -> result<u64, error>;
+}

wit/deps/keyvalue-2024-05-03/atomic.wit

@@ -0,0 +1,63 @@
+/// A keyvalue interface that provides batch operations.
+/// 
+/// A batch operation is an operation that operates on multiple keys at once.
+/// 
+/// Batch operations are useful for reducing network round-trip time. For example, if you want to
+/// get the values associated with 100 keys, you can either do 100 get operations or you can do 1
+/// batch get operation. The batch operation is faster because it only needs to make 1 network call
+/// instead of 100.
+/// 
+/// A batch operation does not guarantee atomicity, meaning that if the batch operation fails, some
+/// of the keys may have been modified and some may not. 
+/// 
+/// This interface does has the same consistency guarantees as the `store` interface, meaning that
+/// you should be able to "read your writes."
+/// 
+/// Please note that this interface is bare functions that take a reference to a bucket. This is to
+/// get around the current lack of a way to "extend" a resource with additional methods inside of
+/// wit. Future version of the interface will instead extend these methods on the base `bucket`
+/// resource.
+interface batch {
+    use store.{bucket, error};
+
+    /// Get the key-value pairs associated with the keys in the store. It returns a list of
+    /// key-value pairs.
+    ///
+    /// If any of the keys do not exist in the store, it returns a `none` value for that pair in the
+    /// list.
+    /// 
+    /// MAY show an out-of-date value if there are concurrent writes to the store.
+    /// 
+    /// If any other error occurs, it returns an `Err(error)`.
+    get-many: func(bucket: borrow<bucket>, keys: list<string>) -> result<list<option<tuple<string, list<u8>>>>, error>;
+
+    /// Set the values associated with the keys in the store. If the key already exists in the
+    /// store, it overwrites the value. 
+    /// 
+    /// Note that the key-value pairs are not guaranteed to be set in the order they are provided. 
+    ///
+    /// If any of the keys do not exist in the store, it creates a new key-value pair.
+    /// 
+    /// If any other error occurs, it returns an `Err(error)`. When an error occurs, it does not
+    /// rollback the key-value pairs that were already set. Thus, this batch operation does not
+    /// guarantee atomicity, implying that some key-value pairs could be set while others might
+    /// fail. 
+    /// 
+    /// Other concurrent operations may also be able to see the partial results.
+    set-many: func(bucket: borrow<bucket>, key-values: list<tuple<string, list<u8>>>) -> result<_, error>;
+
+    /// Delete the key-value pairs associated with the keys in the store.
+    /// 
+    /// Note that the key-value pairs are not guaranteed to be deleted in the order they are
+    /// provided.
+    /// 
+    /// If any of the keys do not exist in the store, it skips the key.
+    /// 
+    /// If any other error occurs, it returns an `Err(error)`. When an error occurs, it does not
+    /// rollback the key-value pairs that were already deleted. Thus, this batch operation does not
+    /// guarantee atomicity, implying that some key-value pairs could be deleted while others might
+    /// fail.
+    /// 
+    /// Other concurrent operations may also be able to see the partial results.
+    delete-many: func(bucket: borrow<bucket>, keys: list<string>) -> result<_, error>;
+}