Add kv map builder and list function

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"
+        );
+    }
 }

atuin/src/command/client/kv.rs

@@ -11,7 +11,7 @@ pub enum Cmd {
         #[arg(long, short)]
         key: String,
 
-        #[arg(long, short, default_value = "global")]
+        #[arg(long, short, default_value = "default")]
         namespace: String,
 
         value: String,
@@ -21,9 +21,17 @@ pub enum Cmd {
     Get {
         key: String,
 
-        #[arg(long, short, default_value = "global")]
+        #[arg(long, short, default_value = "default")]
         namespace: String,
     },
+
+    List {
+        #[arg(long, short, default_value = "default")]
+        namespace: String,
+
+        #[arg(long, short)]
+        all_namespaces: bool,
+    },
 }
 
 impl Cmd {
@@ -38,14 +46,16 @@ impl Cmd {
             .context("could not load encryption key")?
             .into();
 
+        let host_id = Settings::host_id().expect("failed to get host_id");
+
         match self {
             Self::Set {
                 key,
                 value,
                 namespace,
             } => {
                 kv_store
-                    .set(store, &encryption_key, namespace, key, value)
+                    .set(store, &encryption_key, host_id, namespace, key, value)
                     .await
             }
 
@@ -58,6 +68,33 @@ impl Cmd {
 
                 Ok(())
             }
+
+            Self::List {
+                namespace,
+                all_namespaces,
+            } => {
+                // TODO: don't rebuild this every time lol
+                let map = kv_store.build_kv(store, &encryption_key).await?;
+
+                // slower, but sorting is probably useful
+                if *all_namespaces {
+                    for (ns, kv) in &map {
+                        for k in kv.keys() {
+                            println!("{ns}.{k}");
+                        }
+                    }
+                } else {
+                    let ns = map.get(namespace);
+
+                    if let Some(ns) = ns {
+                        for k in ns.keys() {
+                            println!("{k}");
+                        }
+                    }
+                }
+
+                Ok(())
+            }
         }
     }
 }