mirror of
https://github.com/ellie/atuin.git
synced 2024-09-11 21:18:22 +03:00
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:
parent
fbbe24da75
commit
69a772d1ca
@ -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"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
@ -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(())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
Loading…
Reference in New Issue
Block a user