Add kv map builder and list function (#1179)

* Add kv map builder and list function

1. BREAKING - default namespace is now called "default"
2. Build an in-memory hashmap from the kv store
3. Allow listing

I need to cache the hashmap next, probs with a write-through to avoid
constant rebuilds.

Also check if BTreeMap is suitable. Sorted is useful for listing but
there's probs a better ds to use.

* Allow pure kv set, no filesystem

1 files changed, 85 insertions, 5 deletions

diff --git a/atuin-client/src/kv.rs b/atuin-client/src/kv.rs
index ad53a3d1..5710424f 100644
--- a/atuin-client/src/kv.rs
+++ b/atuin-client/src/kv.rs
@@ -1,9 +1,11 @@
-use atuin_common::record::DecryptedData;
+use std::collections::BTreeMap;
+
+use atuin_common::record::{DecryptedData, HostId};
 use eyre::{bail, ensure, eyre, Result};
+use serde::Deserialize;
 
 use crate::record::encryption::PASETO_V4;
 use crate::record::store::Store;
-use crate::settings::Settings;
 
 const KV_VERSION: &str = "v0";
 const KV_TAG: &str = "kv";
@@ -70,6 +72,7 @@ impl KvRecord {
     }
 }
 
+#[derive(Debug, Clone, Deserialize)]
 pub struct KvStore;
 
 impl Default for KvStore {
@@ -88,6 +91,7 @@ impl KvStore {
         &self,
         store: &mut (impl Store + Send + Sync),
         encryption_key: &[u8; 32],
+        host_id: HostId,
         namespace: &str,
         key: &str,
         value: &str,
@@ -99,8 +103,6 @@ impl KvStore {
             ));
         }
 
-        let host_id = Settings::host_id().expect("failed to get host_id");
-
         let record = KvRecord {
             namespace: namespace.to_string(),
             key: key.to_string(),
@@ -173,11 +175,55 @@ impl KvStore {
         // if we get here, then... we didn't find the record with that key :(
         Ok(None)
     }
+
+    // Build a kv map out of the linked list kv store
+    // Map is Namespace -> Key -> Value
+    // TODO(ellie): "cache" this into a real kv structure, which we can
+    // use as a write-through cache to avoid constant rebuilds.
+    pub async fn build_kv(
+        &self,
+        store: &impl Store,
+        encryption_key: &[u8; 32],
+    ) -> Result<BTreeMap<String, BTreeMap<String, String>>> {
+        let mut map = BTreeMap::new();
+        let tails = store.tag_tails(KV_TAG).await?;
+
+        if tails.is_empty() {
+            return Ok(map);
+        }
+
+        let mut record = tails.iter().max_by_key(|r| r.timestamp).unwrap().clone();
+
+        loop {
+            let decrypted = match record.version.as_str() {
+                KV_VERSION => record.decrypt::<PASETO_V4>(encryption_key)?,
+                version => bail!("unknown version {version:?}"),
+            };
+
+            let kv = KvRecord::deserialize(&decrypted.data, &decrypted.version)?;
+
+            let ns = map.entry(kv.namespace).or_insert_with(BTreeMap::new);
+            ns.entry(kv.key).or_insert_with(|| kv.value);
+
+            if let Some(parent) = decrypted.parent {
+                record = store.get(parent).await?;
+            } else {
+                break;
+            }
+        }
+
+        Ok(map)
+    }
 }
 
 #[cfg(test)]
 mod tests {
-    use super::{KvRecord, KV_VERSION};
+    use rand::rngs::OsRng;
+    use xsalsa20poly1305::{KeyInit, XSalsa20Poly1305};
+
+    use crate::record::sqlite_store::SqliteStore;
+
+    use super::{KvRecord, KvStore, KV_VERSION};
 
     #[test]
     fn encode_decode() {
@@ -196,4 +242,38 @@ mod tests {
         assert_eq!(encoded.0, &snapshot);
         assert_eq!(decoded, kv);
     }
+
+    #[tokio::test]
+    async fn build_kv() {
+        let mut store = SqliteStore::new(":memory:").await.unwrap();
+        let kv = KvStore::new();
+        let key: [u8; 32] = XSalsa20Poly1305::generate_key(&mut OsRng).into();
+        let host_id = atuin_common::record::HostId(atuin_common::utils::uuid_v7());
+
+        kv.set(&mut store, &key, host_id, "test-kv", "foo", "bar")
+            .await
+            .unwrap();
+
+        kv.set(&mut store, &key, host_id, "test-kv", "1", "2")
+            .await
+            .unwrap();
+
+        let map = kv.build_kv(&store, &key).await.unwrap();
+
+        assert_eq!(
+            map.get("test-kv")
+                .expect("map namespace not set")
+                .get("foo")
+                .expect("map key not set"),
+            "bar"
+        );
+
+        assert_eq!(
+            map.get("test-kv")
+                .expect("map namespace not set")
+                .get("1")
+                .expect("map key not set"),
+            "2"
+        );
+    }
 }