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
This commit is contained in:
Ellie Huxtable 2023-08-18 08:36:55 +01:00 committed by GitHub
parent fbbe24da75
commit 69a772d1ca
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 125 additions and 8 deletions

View file

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

View file

@ -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(())
}
}
}
}