From 407182dfbd9c7f28806e44a6764e2059567001f5 Mon Sep 17 00:00:00 2001
From: Darwin Boersma <darwin@sadlark.com>
Date: Wed, 30 Oct 2024 09:00:46 -0600
Subject: [PATCH] Update to DynamoKeyValueStore to make other AWS KV store
 implementations easier

Signed-off-by: Darwin Boersma <darwin@sadlark.com>
---
 crates/key-value-aws/src/lib.rs   |  6 +--
 crates/key-value-aws/src/store.rs | 62 ++++++++++++++++---------------
 crates/runtime-config/src/lib.rs  |  2 +-
 3 files changed, 36 insertions(+), 34 deletions(-)

diff --git a/crates/key-value-aws/src/lib.rs b/crates/key-value-aws/src/lib.rs
index aaf3ecca7..a0305be66 100644
--- a/crates/key-value-aws/src/lib.rs
+++ b/crates/key-value-aws/src/lib.rs
@@ -8,11 +8,11 @@ use store::{
 
 /// A key-value store that uses AWS Dynamo as the backend.
 #[derive(Default)]
-pub struct AwsKeyValueStore {
+pub struct AwsDynamoKeyValueStore {
     _priv: (),
 }
 
-impl AwsKeyValueStore {
+impl AwsDynamoKeyValueStore {
     /// Creates a new `AwsKeyValueStore`.
     pub fn new() -> Self {
         Self::default()
@@ -34,7 +34,7 @@ pub struct AwsDynamoKeyValueRuntimeConfig {
     table: String,
 }
 
-impl MakeKeyValueStore for AwsKeyValueStore {
+impl MakeKeyValueStore for AwsDynamoKeyValueStore {
     const RUNTIME_CONFIG_TYPE: &'static str = "aws_dynamo";
 
     type RuntimeConfig = AwsDynamoKeyValueRuntimeConfig;
diff --git a/crates/key-value-aws/src/store.rs b/crates/key-value-aws/src/store.rs
index 7a1b8548b..ca9d0d573 100644
--- a/crates/key-value-aws/src/store.rs
+++ b/crates/key-value-aws/src/store.rs
@@ -77,27 +77,24 @@ impl KeyValueAwsDynamo {
         auth_options: KeyValueAwsDynamoAuthOptions,
     ) -> Result<Self> {
         let region_clone = region.clone();
-        let client_fut: std::pin::Pin<Box<dyn std::future::Future<Output = Client> + Send>> =
-            Box::pin(async move {
-                let config = match auth_options {
-                    KeyValueAwsDynamoAuthOptions::RuntimeConfigValues(config) => {
-                        SdkConfig::builder()
-                            .credentials_provider(SharedCredentialsProvider::new(config))
-                            .region(Region::new(region_clone))
-                            .behavior_version(BehaviorVersion::latest())
-                            .build()
-                    }
-                    KeyValueAwsDynamoAuthOptions::Environmental => {
-                        aws_config::load_defaults(BehaviorVersion::latest()).await
-                    }
-                };
-                Client::new(&config)
-            });
+        let client_fut = Box::pin(async move {
+            let sdk_config = match auth_options {
+                KeyValueAwsDynamoAuthOptions::RuntimeConfigValues(config) => SdkConfig::builder()
+                    .credentials_provider(SharedCredentialsProvider::new(config))
+                    .region(Region::new(region_clone))
+                    .behavior_version(BehaviorVersion::latest())
+                    .build(),
+                KeyValueAwsDynamoAuthOptions::Environmental => {
+                    aws_config::load_defaults(BehaviorVersion::latest()).await
+                }
+            };
+            Client::new(&sdk_config)
+        });
 
         Ok(Self {
-            client: async_once_cell::Lazy::from_future(client_fut),
             table,
             region,
+            client: async_once_cell::Lazy::from_future(client_fut),
         })
     }
 }
@@ -143,7 +140,7 @@ impl Store for AwsDynamoStore {
     async fn set(&self, key: &str, value: &[u8]) -> Result<(), Error> {
         self.client
             .put_item()
-            .table_name(self.table.clone())
+            .table_name(&self.table)
             .item(PK, AttributeValue::S(key.to_string()))
             .item(VAL, AttributeValue::B(Blob::new(value)))
             .send()
@@ -156,7 +153,7 @@ impl Store for AwsDynamoStore {
         if self.exists(key).await? {
             self.client
                 .delete_item()
-                .table_name(self.table.clone())
+                .table_name(&self.table)
                 .key(PK, AttributeValue::S(key.to_string()))
                 .send()
                 .await
@@ -176,22 +173,27 @@ impl Store for AwsDynamoStore {
 
 impl AwsDynamoStore {
     async fn get_item(&self, key: &str) -> Result<Option<Vec<u8>>, Error> {
-        let query = self
+        let response = self
             .client
             .get_item()
-            .table_name(self.table.clone())
-            .key(PK, aws_sdk_dynamodb::types::AttributeValue::S(key.into()))
+            .table_name(&self.table)
+            .key(
+                PK,
+                aws_sdk_dynamodb::types::AttributeValue::S(key.to_string()),
+            )
             .send()
             .await
             .map_err(log_error)?;
 
-        Ok(query.item.and_then(|item| {
-            if let Some(AttributeValue::B(val)) = item.get(VAL) {
-                Some(val.clone().into_inner())
+        let val = response.item.and_then(|mut item| {
+            if let Some(AttributeValue::B(val)) = item.remove(VAL) {
+                Some(val.into_inner())
             } else {
                 None
             }
-        }))
+        });
+
+        Ok(val)
     }
 
     async fn get_keys(&self) -> Result<Vec<String>, Error> {
@@ -202,7 +204,7 @@ impl AwsDynamoStore {
             let mut scan_builder = self
                 .client
                 .scan()
-                .table_name(self.table.clone())
+                .table_name(&self.table)
                 .projection_expression(PK);
 
             if let Some(keys) = last_evaluated_key {
@@ -214,9 +216,9 @@ impl AwsDynamoStore {
             let scan_output = scan_builder.send().await.map_err(log_error)?;
 
             if let Some(items) = scan_output.items {
-                for item in items {
-                    if let Some(AttributeValue::S(pk)) = item.get(PK) {
-                        primary_keys.push(pk.clone());
+                for mut item in items {
+                    if let Some(AttributeValue::S(pk)) = item.remove(PK) {
+                        primary_keys.push(pk);
                     }
                 }
             }
diff --git a/crates/runtime-config/src/lib.rs b/crates/runtime-config/src/lib.rs
index aa14c5dc4..963924363 100644
--- a/crates/runtime-config/src/lib.rs
+++ b/crates/runtime-config/src/lib.rs
@@ -401,7 +401,7 @@ pub fn key_value_config_resolver(
         .register_store_type(spin_key_value_azure::AzureKeyValueStore::new())
         .unwrap();
     key_value
-        .register_store_type(spin_key_value_aws::AwsKeyValueStore::new())
+        .register_store_type(spin_key_value_aws::AwsDynamoKeyValueStore::new())
         .unwrap();
 
     // Add handling of "default" store.