diff options
| -rw-r--r-- | atuin-client/src/kv.rs | 90 | ||||
| -rw-r--r-- | atuin/src/command/client/kv.rs | 43 |
2 files changed, 125 insertions, 8 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" + ); + } } diff --git a/atuin/src/command/client/kv.rs b/atuin/src/command/client/kv.rs index 694ee67e..d9a22760 100644 --- a/atuin/src/command/client/kv.rs +++ b/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,6 +46,8 @@ 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, @@ -45,7 +55,7 @@ impl Cmd { 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(()) + } } } } |