feat(kv): support setting key/value pairs

Signed-off-by: Kate Goldenring <kate.goldenring@fermyon.com>

crates/cloud/src/client.rs

@@ -240,15 +240,15 @@ impl CloudClientInterface for Client {
     // Key value API methods
     async fn add_key_value_pair(
         &self,
-        app_id: Uuid,
+        app_id: Option<Uuid>,
         store_name: String,
         key: String,
         value: String,
     ) -> anyhow::Result<()> {
         api_key_value_pairs_post(
             &self.configuration,
             CreateKeyValuePairCommand {
-                app_id: Some(app_id),
+                app_id,
                 store_name: Some(store_name),
                 key,
                 value,

crates/cloud/src/client_interface.rs

@@ -49,7 +49,7 @@ pub trait CloudClientInterface: Send + Sync {
 
     async fn add_key_value_pair(
         &self,
-        app_id: Uuid,
+        app_id: Option<Uuid>,
         store_name: String,
         key: String,
         value: String,

src/commands/deploy.rs

@@ -224,7 +224,12 @@ impl DeployCommand {
 
                 for kv in self.key_values {
                     client
-                        .add_key_value_pair(app_id, SPIN_DEFAULT_KV_STORE.to_string(), kv.0, kv.1)
+                        .add_key_value_pair(
+                            Some(app_id),
+                            SPIN_DEFAULT_KV_STORE.to_string(),
+                            kv.0,
+                            kv.1,
+                        )
                         .await
                         .context("Problem creating key/value")?;
                 }
@@ -258,7 +263,12 @@ impl DeployCommand {
 
                 for kv in self.key_values {
                     client
-                        .add_key_value_pair(app_id, SPIN_DEFAULT_KV_STORE.to_string(), kv.0, kv.1)
+                        .add_key_value_pair(
+                            Some(app_id),
+                            SPIN_DEFAULT_KV_STORE.to_string(),
+                            kv.0,
+                            kv.1,
+                        )
                         .await
                         .context("Problem creating key/value")?;
                 }

src/commands/key_value.rs

@@ -2,10 +2,13 @@ use crate::commands::links_output::{
     print_json, print_table, prompt_delete_resource, ListFormat, ResourceGroupBy, ResourceLinks,
     ResourceType,
 };
-use crate::commands::{create_cloud_client, CommonArgs};
+use crate::commands::links_target::ResourceTarget;
+use crate::commands::{create_cloud_client, disallow_empty, CommonArgs};
 use anyhow::{bail, Context, Result};
 use clap::{Parser, ValueEnum};
 use cloud::CloudClientInterface;
+use cloud_openapi::models::KeyValueStoreItem;
+use spin_common::arg_parser::parse_kv;
 
 #[derive(Parser, Debug)]
 #[clap(about = "Manage Fermyon Cloud key value stores")]
@@ -16,6 +19,8 @@ pub enum KeyValueCommand {
     Delete(DeleteCommand),
     /// List key value stores
     List(ListCommand),
+    /// Set a key value pair in a store
+    Set(SetCommand),
 }
 
 #[derive(Parser, Debug)]
@@ -74,6 +79,29 @@ impl From<GroupBy> for ResourceGroupBy {
     }
 }
 
+#[derive(Parser, Debug)]
+pub struct SetCommand {
+    /// The name of the key value store
+    #[clap(name = "STORE", short = 's', long = "store", value_parser = clap::builder::ValueParser::new(disallow_empty), required_unless_present_all = ["LABEL", "APP"], conflicts_with_all = &["LABEL", "APP"])]
+    pub store: Option<String>,
+
+    /// Label of the key value store to set pairs in
+    #[clap(name = "LABEL", short = 'l', long = "label", value_parser = clap::builder::ValueParser::new(disallow_empty), requires = "APP", required_unless_present = "STORE")]
+    pub label: Option<String>,
+
+    /// App to which label relates
+    #[clap(name = "APP", short = 'a', long = "app", value_parser = clap::builder::ValueParser::new(disallow_empty), requires = "LABEL", required_unless_present = "STORE")]
+    pub app: Option<String>,
+
+    /// A key/value pair (key=value) to set in the store. Any existing value will be overwritten.
+    /// Can be used multiple times.
+    #[clap(parse(try_from_str = parse_kv))]
+    pub key_values: Vec<(String, String)>,
+
+    #[clap(flatten)]
+    common: CommonArgs,
+}
+
 impl KeyValueCommand {
     pub async fn run(&self) -> Result<()> {
         match self {
@@ -89,6 +117,10 @@ impl KeyValueCommand {
                 let client = create_cloud_client(cmd.common.deployment_env_id.as_deref()).await?;
                 cmd.run(client).await
             }
+            KeyValueCommand::Set(cmd) => {
+                let client = create_cloud_client(cmd.common.deployment_env_id.as_deref()).await?;
+                cmd.run(client).await
+            }
         }
     }
 }
@@ -165,6 +197,39 @@ impl ListCommand {
         }
     }
 }
+
+impl SetCommand {
+    pub async fn run(&self, client: impl CloudClientInterface) -> Result<()> {
+        let target = ResourceTarget::from_inputs(&self.store, &self.label, &self.app)?;
+        let stores = client
+            .get_key_value_stores(None)
+            .await
+            .context("Problem fetching key value stores")?;
+        let store = target
+            .find_in(to_resource_links(stores), ResourceType::KeyValueStore)?
+            .name;
+        for (key, value) in &self.key_values {
+            client
+                .add_key_value_pair(None, store.clone(), key.clone(), value.clone())
+                .await
+                .with_context(|| {
+                    format!(
+                        "Error adding key value pair '{key}={value}' to store '{}'",
+                        store
+                    )
+                })?;
+        }
+        Ok(())
+    }
+}
+
+fn to_resource_links(stores: Vec<KeyValueStoreItem>) -> Vec<ResourceLinks> {
+    stores
+        .into_iter()
+        .map(|s| ResourceLinks::new(s.name, s.links))
+        .collect()
+}
+
 #[cfg(test)]
 mod key_value_tests {
     use super::*;

src/commands/links_output.rs

@@ -15,6 +15,7 @@ pub enum ListFormat {
     Json,
 }
 
+#[derive(PartialEq, Debug, Clone)]
 pub struct ResourceLinks {
     pub name: String,
     pub links: Vec<ResourceLabel>,
@@ -24,6 +25,12 @@ impl ResourceLinks {
     pub fn new(name: String, links: Vec<ResourceLabel>) -> Self {
         Self { name, links }
     }
+
+    pub fn has_link(&self, label: &str, app: Option<&str>) -> bool {
+        self.links
+            .iter()
+            .any(|l| l.label == label && l.app_name.as_deref() == app)
+    }
 }
 
 #[derive(Debug, Clone, Copy)]

src/commands/links_target.rs

@@ -0,0 +1,115 @@
+/// Module for determining the linkable resource to target for a command
+use crate::commands::links_output::{ResourceLinks, ResourceType};
+
+#[derive(Debug, PartialEq)]
+pub enum ResourceTarget {
+    ByName(String),
+    ByLabel { label: String, app: String },
+}
+
+impl ResourceTarget {
+    pub fn find_in(
+        &self,
+        resources: Vec<ResourceLinks>,
+        resource_type: ResourceType,
+    ) -> anyhow::Result<ResourceLinks> {
+        match self {
+            Self::ByName(resource) => resources
+                .into_iter()
+                .find(|r| &r.name == resource)
+                .ok_or_else(|| {
+                    anyhow::anyhow!("No {resource_type} found with name \"{resource}\"")
+                }),
+            Self::ByLabel { label, app } => resources
+                .into_iter()
+                .find(|r| r.has_link(label, Some(app.as_str())))
+                .ok_or_else(|| {
+                    anyhow::anyhow!(
+                        r#"No {resource_type} found with label "{label}" for app "{app}""#
+                    )
+                }),
+        }
+    }
+
+    pub fn from_inputs(
+        resource: &Option<String>,
+        label: &Option<String>,
+        app: &Option<String>,
+    ) -> anyhow::Result<ResourceTarget> {
+        match (resource, label, app) {
+            (Some(r), None, None) => Ok(ResourceTarget::ByName(r.to_owned())),
+            (None, Some(l), Some(a)) => Ok(ResourceTarget::ByLabel {
+                label: l.to_owned(),
+                app: a.to_owned(),
+            }),
+            _ => Err(anyhow::anyhow!("Invalid combination of arguments")), // Should be prevented by clap
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use cloud_openapi::models::ResourceLabel;
+    use uuid::Uuid;
+
+    use super::*;
+
+    #[test]
+    fn test_execute_target_from_inputs() {
+        assert_eq!(
+            ResourceTarget::from_inputs(&Some("mykv".to_owned()), &None, &None).unwrap(),
+            ResourceTarget::ByName("mykv".to_owned())
+        );
+        assert_eq!(
+            ResourceTarget::from_inputs(&None, &Some("label".to_owned()), &Some("app".to_owned()))
+                .unwrap(),
+            ResourceTarget::ByLabel {
+                label: "label".to_owned(),
+                app: "app".to_owned(),
+            }
+        );
+        assert!(ResourceTarget::from_inputs(&None, &None, &None).is_err());
+        assert!(ResourceTarget::from_inputs(
+            &Some("mykv".to_owned()),
+            &Some("label".to_owned()),
+            &Some("app".to_owned())
+        )
+        .is_err());
+    }
+
+    #[test]
+    fn test_execute_target_find_in() {
+        let links = vec![ResourceLabel {
+            app_id: Uuid::new_v4(),
+            app_name: Some("app".to_owned()),
+            label: "label".to_owned(),
+        }];
+        let rl1 = ResourceLinks::new("mykv".to_owned(), vec![]);
+        let rl2 = ResourceLinks::new("mykv2".to_owned(), links);
+        let resources = vec![rl1.clone(), rl2.clone()];
+        assert_eq!(
+            ResourceTarget::ByName("mykv".to_owned())
+                .find_in(resources.clone(), ResourceType::KeyValueStore)
+                .unwrap(),
+            rl1
+        );
+        assert_eq!(
+            ResourceTarget::ByLabel {
+                label: "label".to_owned(),
+                app: "app".to_owned(),
+            }
+            .find_in(resources.clone(), ResourceType::KeyValueStore)
+            .unwrap(),
+            rl2
+        );
+        assert!(ResourceTarget::ByName("foo".to_owned())
+            .find_in(resources.clone(), ResourceType::KeyValueStore)
+            .is_err());
+        assert!(ResourceTarget::ByLabel {
+            label: "foo".to_owned(),
+            app: "app".to_owned(),
+        }
+        .find_in(resources.clone(), ResourceType::KeyValueStore)
+        .is_err());
+    }
+}

src/commands/mod.rs

@@ -3,6 +3,7 @@ pub mod deploy;
 pub mod key_value;
 pub mod link;
 pub mod links_output;
+pub mod links_target;
 pub mod login;
 pub mod logs;
 pub mod sqlite;
@@ -54,3 +55,10 @@ struct CommonArgs {
     )]
     pub deployment_env_id: Option<String>,
 }
+
+fn disallow_empty(statement: &str) -> anyhow::Result<String> {
+    if statement.trim().is_empty() {
+        anyhow::bail!("cannot be empty");
+    }
+    return Ok(statement.trim().to_owned());
+}