From 0c75cfbfda1f991c618fcd00921ac6be99fd03f7 Mon Sep 17 00:00:00 2001
From: Ellie Huxtable <ellie@elliehuxtable.com>
Date: Thu, 15 Jun 2023 08:17:30 +0100
Subject: [PATCH] Add namespaces to kv store (#1052)

---
 atuin-client/src/kv.rs         | 14 +++++++++++---
 atuin/src/command/client/kv.rs | 16 +++++++++++++---
 2 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/atuin-client/src/kv.rs b/atuin-client/src/kv.rs
index 87149275..35e8852e 100644
--- a/atuin-client/src/kv.rs
+++ b/atuin-client/src/kv.rs
@@ -9,6 +9,7 @@ const KV_TAG: &str = "kv";
 
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
 pub struct KvRecord {
+    pub namespace: String,
     pub key: String,
     pub value: String,
 }
@@ -38,12 +39,14 @@ impl KvStore {
     pub async fn set(
         &self,
         store: &mut (impl Store + Send + Sync),
+        namespace: &str,
         key: &str,
         value: &str,
     ) -> Result<()> {
         let host_id = Settings::host_id().expect("failed to get host_id");
 
         let record = KvRecord {
+            namespace: namespace.to_string(),
             key: key.to_string(),
             value: value.to_string(),
         };
@@ -70,7 +73,12 @@ impl KvStore {
 
     // TODO: setup an actual kv store, rebuild func, and do not pass the main store in here as
     // well.
-    pub async fn get(&self, store: &impl Store, key: &str) -> Result<Option<KvRecord>> {
+    pub async fn get(
+        &self,
+        store: &impl Store,
+        namespace: &str,
+        key: &str,
+    ) -> Result<Option<KvRecord>> {
         // TODO: don't load this from disk so much
         let host_id = Settings::host_id().expect("failed to get host_id");
 
@@ -84,7 +92,7 @@ impl KvStore {
         };
         let kv: KvRecord = rmp_serde::from_slice(&record.data)?;
 
-        if kv.key == key {
+        if kv.key == key && kv.namespace == namespace {
             return Ok(Some(kv));
         }
 
@@ -92,7 +100,7 @@ impl KvStore {
             record = store.get(parent.as_str()).await?;
             let kv: KvRecord = rmp_serde::from_slice(&record.data)?;
 
-            if kv.key == key {
+            if kv.key == key && kv.namespace == namespace {
                 return Ok(Some(kv));
             }
         }
diff --git a/atuin/src/command/client/kv.rs b/atuin/src/command/client/kv.rs
index f922b063..a3f642df 100644
--- a/atuin/src/command/client/kv.rs
+++ b/atuin/src/command/client/kv.rs
@@ -11,12 +11,18 @@ pub enum Cmd {
         #[arg(long, short)]
         key: String,
 
+        #[arg(long, short, default_value = "global")]
+        namespace: String,
+
         value: String,
     },
 
     // atuin kv get foo => bar baz
     Get {
         key: String,
+
+        #[arg(long, short, default_value = "global")]
+        namespace: String,
     },
 }
 
@@ -29,10 +35,14 @@ impl Cmd {
         let kv_store = KvStore::new();
 
         match self {
-            Self::Set { key, value } => kv_store.set(store, key, value).await,
+            Self::Set {
+                key,
+                value,
+                namespace,
+            } => kv_store.set(store, namespace, key, value).await,
 
-            Self::Get { key } => {
-                let val = kv_store.get(store, key).await?;
+            Self::Get { key, namespace } => {
+                let val = kv_store.get(store, namespace, key).await?;
 
                 if let Some(kv) = val {
                     println!("{}", kv.value);