From 09831ca3998254f0c05cf3e1e31a60256d34b273 Mon Sep 17 00:00:00 2001
From: Callum-A <callumanderson745@gmail.com>
Date: Wed, 11 Sep 2024 17:49:28 +0100
Subject: [PATCH] Feat : Get or create Key Value store

---
 async-nats/src/jetstream/context.rs | 40 ++++++++++++++++++++++
 async-nats/tests/kv_tests.rs        | 51 +++++++++++++++++++++++++++++
 2 files changed, 91 insertions(+)

diff --git a/async-nats/src/jetstream/context.rs b/async-nats/src/jetstream/context.rs
index ae13e8429..57c85562e 100644
--- a/async-nats/src/jetstream/context.rs
+++ b/async-nats/src/jetstream/context.rs
@@ -821,6 +821,46 @@ impl Context {
         Ok(store)
     }
 
+    /// Tries to get an existing key-value bucket, if one cannot be found it will create a new key-value bucket.
+    ///
+    /// Note: This does not validate if the key-value on the server is compatible with the configuration passed in.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # #[tokio::main]
+    /// # async fn main() -> Result<(), async_nats::Error> {
+    /// let client = async_nats::connect("demo.nats.io:4222").await?;
+    /// let jetstream = async_nats::jetstream::new(client);
+    /// let kv = jetstream
+    ///     .get_or_create_key_value("kv".to_string(), async_nats::jetstream::kv::Config {
+    ///         bucket: "kv".to_string(),
+    ///         history: 10,
+    ///         ..Default::default()
+    ///     })
+    ///     .await?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub async fn get_or_create_key_value<T: Into<String>>(
+        &self,
+        bucket: T,
+        config: crate::jetstream::kv::Config,
+    ) -> Result<Store, CreateKeyValueError> {
+        match self.get_key_value(bucket).await {
+            Ok(kv) => Ok(kv),
+            Err(e) => match e.kind() {
+                KeyValueErrorKind::GetBucket => self.create_key_value(config).await,
+                KeyValueErrorKind::InvalidStoreName => Err(CreateKeyValueError::new(
+                    CreateKeyValueErrorKind::InvalidStoreName,
+                )),
+                KeyValueErrorKind::JetStream => {
+                    Err(CreateKeyValueError::new(CreateKeyValueErrorKind::JetStream))
+                }
+            },
+        }
+    }
+
     /// Deletes given key-value bucket.
     ///
     /// # Examples
diff --git a/async-nats/tests/kv_tests.rs b/async-nats/tests/kv_tests.rs
index 1901ca092..3fd3c7674 100644
--- a/async-nats/tests/kv_tests.rs
+++ b/async-nats/tests/kv_tests.rs
@@ -52,6 +52,57 @@ mod kv {
         assert!(info.config.allow_direct);
     }
 
+    #[tokio::test]
+    async fn get_or_create_bucket() {
+        let server = nats_server::run_server("tests/configs/jetstream.conf");
+        let client = ConnectOptions::new()
+            .event_callback(|event| async move { println!("event: {event:?}") })
+            .connect(server.client_url())
+            .await
+            .unwrap();
+
+        let context = async_nats::jetstream::new(client);
+
+        let mut kv = context
+            .get_or_create_key_value(
+                "test",
+                async_nats::jetstream::kv::Config {
+                    bucket: "test".into(),
+                    description: "test_description".into(),
+                    history: 10,
+                    storage: StorageType::File,
+                    num_replicas: 1,
+                    ..Default::default()
+                },
+            )
+            .await
+            .unwrap();
+        let info = kv.stream.info().await.unwrap();
+        assert_eq!("KV_test", kv.stream_name);
+        assert_eq!(info.config.discard, DiscardPolicy::New);
+        assert!(info.config.allow_direct);
+
+        let mut kv2 = context
+            .get_or_create_key_value(
+                "test",
+                async_nats::jetstream::kv::Config {
+                    bucket: "test".into(),
+                    description: "test_description".into(),
+                    history: 10,
+                    storage: StorageType::File,
+                    num_replicas: 1,
+                    ..Default::default()
+                },
+            )
+            .await
+            .unwrap();
+
+        let info2 = kv2.stream.info().await.unwrap();
+        assert_eq!(kv2.stream_name, kv.stream_name);
+        assert_eq!(info2.config.discard, info.config.discard);
+        assert_eq!(info2.config.allow_direct, info.config.allow_direct);
+    }
+
     #[tokio::test]
     async fn create() {
         let server = nats_server::run_server("tests/configs/jetstream.conf");