feat: Implement KV as a write-through cache (#2732)

3 files changed, 381 insertions, 0 deletions

diff --git a/crates/atuin-kv/src/store.rs b/crates/atuin-kv/src/store.rs
new file mode 100644
index 00000000..3394b8c0
--- /dev/null
+++ b/crates/atuin-kv/src/store.rs
@@ -0,0 +1,214 @@
+use std::collections::HashSet;
+
+use eyre::{Result, bail};
+
+use atuin_client::record::sqlite_store::SqliteStore;
+use atuin_client::record::{encryption::PASETO_V4, store::Store};
+use atuin_common::record::{Host, HostId, Record, RecordId, RecordIdx};
+use entry::KvEntry;
+use record::{KV_TAG, KV_VERSION, KvRecord};
+
+use crate::database::Database;
+
+pub mod entry;
+pub mod record;
+
+#[derive(Debug, Clone)]
+pub struct KvStore {
+    pub record_store: SqliteStore,
+    pub kv_db: Database,
+    pub host_id: HostId,
+    pub encryption_key: [u8; 32],
+}
+
+impl KvStore {
+    pub fn new(
+        record_store: SqliteStore,
+        kv_db: Database,
+        host_id: HostId,
+        encryption_key: [u8; 32],
+    ) -> Self {
+        KvStore {
+            record_store,
+            kv_db,
+            host_id,
+            encryption_key,
+        }
+    }
+
+    pub async fn set(&self, namespace: &str, key: &str, value: &str) -> Result<()> {
+        let kv_record = KvRecord::builder()
+            .namespace(namespace.to_string())
+            .key(key.to_string())
+            .value(Some(value.to_string()))
+            .build();
+
+        self.push_record(kv_record).await?;
+
+        let kv = KvEntry::builder()
+            .namespace(namespace.to_string())
+            .key(key.to_string())
+            .value(value.to_string())
+            .build();
+
+        self.kv_db.save(&kv).await?;
+
+        Ok(())
+    }
+
+    pub async fn get(&self, namespace: &str, key: &str) -> Result<Option<String>> {
+        let kv = self.kv_db.load(namespace, key).await?;
+        Ok(kv.map(|kv| kv.value))
+    }
+
+    pub async fn delete(&self, namespace: &str, keys: &[String]) -> Result<()> {
+        for key in keys {
+            let record = KvRecord::builder()
+                .namespace(namespace.to_string())
+                .key(key.to_string())
+                .value(None)
+                .build();
+
+            self.push_record(record).await?;
+            self.kv_db.delete(namespace, key).await?;
+        }
+
+        Ok(())
+    }
+
+    pub async fn list(&self, namespace: Option<&str>) -> Result<Vec<KvEntry>> {
+        let entries = self.kv_db.list(namespace).await?;
+
+        Ok(entries)
+    }
+
+    async fn push_record(&self, record: KvRecord) -> Result<(RecordId, RecordIdx)> {
+        let bytes = record.serialize()?;
+        let idx = self
+            .record_store
+            .last(self.host_id, KV_TAG)
+            .await?
+            .map_or(0, |p| p.idx + 1);
+
+        let record = Record::builder()
+            .host(Host::new(self.host_id))
+            .version(KV_VERSION.to_string())
+            .tag(KV_TAG.to_string())
+            .idx(idx)
+            .data(bytes)
+            .build();
+
+        let id = record.id;
+
+        self.record_store
+            .push(&record.encrypt::<PASETO_V4>(&self.encryption_key))
+            .await?;
+
+        Ok((id, idx))
+    }
+
+    pub async fn build(&self) -> Result<()> {
+        let mut tagged = self.record_store.all_tagged(KV_TAG).await?;
+        tagged.reverse();
+
+        let cached = self.kv_db.list(None).await?;
+
+        let mut visited = HashSet::new();
+
+        // Iterate through all KV records from newest to oldest;
+        // only visit each KV once, inserting or deleting based on the first time we see it
+        for record in tagged {
+            let decrypted = match record.version.as_str() {
+                "v0" | KV_VERSION => record.decrypt::<PASETO_V4>(&self.encryption_key)?,
+                version => bail!("unknown version {version:?}"),
+            };
+
+            let kv = KvRecord::deserialize(&decrypted.data, &decrypted.version)?;
+            let uniq_id = format!("{}.{}", kv.namespace, kv.key);
+
+            if visited.insert(uniq_id) {
+                match kv.value {
+                    Some(value) => {
+                        self.kv_db
+                            .save(
+                                &KvEntry::builder()
+                                    .namespace(kv.namespace.clone())
+                                    .key(kv.key.clone())
+                                    .value(value)
+                                    .build(),
+                            )
+                            .await?;
+                    }
+                    None => {
+                        self.kv_db
+                            .delete(kv.namespace.as_str(), kv.key.as_str())
+                            .await?;
+                    }
+                }
+            }
+        }
+
+        // Any KVs that were in the cache but not in the tagged list should be deleted;
+        // this should never happen in practice since the cache is always built from the tagged list,
+        // but just in case because ** S O F T W A R E **
+        for kv in cached {
+            if !visited.contains(&format!("{}.{}", kv.namespace, kv.key)) {
+                self.kv_db
+                    .delete(kv.namespace.as_str(), kv.key.as_str())
+                    .await?;
+            }
+        }
+
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    async fn setup() -> Result<KvStore> {
+        let record_store = SqliteStore::new("sqlite::memory:", 1.0).await.unwrap();
+        let kv_db = Database::new("sqlite::memory:", 1.0).await.unwrap();
+        let host_id = atuin_common::record::HostId(atuin_common::utils::uuid_v7());
+        let encryption_key = [0; 32];
+        Ok(KvStore::new(record_store, kv_db, host_id, encryption_key))
+    }
+
+    #[tokio::test]
+    async fn test_kv_store() -> Result<()> {
+        let store = setup().await?;
+
+        store.set("test", "key", "value").await.unwrap();
+        let value = store.get("test", "key").await.unwrap();
+        assert_eq!(value, Some("value".to_string()));
+
+        let records = store.record_store.all_tagged(KV_TAG).await?;
+        assert_eq!(records.len(), 1);
+
+        let list = store.list(Some("test")).await.unwrap();
+        let expected = vec![
+            KvEntry::builder()
+                .namespace("test".to_string())
+                .key("key".to_string())
+                .value("value".to_string())
+                .build(),
+        ];
+        assert_eq!(list, expected);
+
+        let ns_list = store.list(None).await.unwrap();
+        assert_eq!(ns_list, expected);
+
+        store
+            .delete("test", &vec!["key".to_string()])
+            .await
+            .unwrap();
+        let value = store.get("test", "key").await.unwrap();
+        assert_eq!(value, None);
+
+        let records = store.record_store.all_tagged(KV_TAG).await?;
+        assert_eq!(records.len(), 2);
+
+        Ok(())
+    }
+}
diff --git a/crates/atuin-kv/src/store/entry.rs b/crates/atuin-kv/src/store/entry.rs
new file mode 100644
index 00000000..1d6a1ef8
--- /dev/null
+++ b/crates/atuin-kv/src/store/entry.rs
@@ -0,0 +1,8 @@
+use typed_builder::TypedBuilder;
+
+#[derive(Debug, Clone, PartialEq, Eq, TypedBuilder)]
+pub struct KvEntry {
+    pub namespace: String,
+    pub key: String,
+    pub value: String,
+}
diff --git a/crates/atuin-kv/src/store/record.rs b/crates/atuin-kv/src/store/record.rs
new file mode 100644
index 00000000..37254176
--- /dev/null
+++ b/crates/atuin-kv/src/store/record.rs
@@ -0,0 +1,159 @@
+use atuin_common::record::DecryptedData;
+use eyre::{Result, bail, ensure, eyre};
+use typed_builder::TypedBuilder;
+
+pub const KV_VERSION: &str = "v1";
+pub const KV_TAG: &str = "kv";
+pub const KV_VAL_MAX_LEN: usize = 100 * 1024;
+
+#[derive(Debug, Clone, PartialEq, Eq, TypedBuilder)]
+pub struct KvRecord {
+    pub namespace: String,
+    pub key: String,
+    pub value: Option<String>,
+}
+
+impl KvRecord {
+    pub fn serialize(&self) -> Result<DecryptedData> {
+        use rmp::encode;
+
+        let mut output = vec![];
+
+        // INFO: ensure this is updated when adding new fields
+        encode::write_array_len(&mut output, 4)?;
+
+        encode::write_str(&mut output, &self.namespace)?;
+        encode::write_str(&mut output, &self.key)?;
+        encode::write_bool(&mut output, self.value.is_some())?;
+
+        if let Some(value) = &self.value {
+            encode::write_str(&mut output, value)?;
+        }
+
+        Ok(DecryptedData(output))
+    }
+
+    pub fn deserialize(data: &DecryptedData, version: &str) -> Result<Self> {
+        use rmp::decode;
+
+        fn error_report<E: std::fmt::Debug>(err: E) -> eyre::Report {
+            eyre!("{err:?}")
+        }
+
+        match version {
+            "v0" => {
+                let mut bytes = decode::Bytes::new(&data.0);
+
+                let nfields = decode::read_array_len(&mut bytes).map_err(error_report)?;
+                ensure!(nfields == 3, "too many entries in v0 kv record");
+
+                let bytes = bytes.remaining_slice();
+
+                let (namespace, bytes) =
+                    decode::read_str_from_slice(bytes).map_err(error_report)?;
+                let (key, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?;
+                let (value, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?;
+
+                if !bytes.is_empty() {
+                    bail!("trailing bytes in encoded kvrecord. malformed")
+                }
+
+                Ok(KvRecord {
+                    namespace: namespace.to_owned(),
+                    key: key.to_owned(),
+                    value: Some(value.to_owned()),
+                })
+            }
+            KV_VERSION => {
+                let mut bytes = decode::Bytes::new(&data.0);
+
+                let nfields = decode::read_array_len(&mut bytes).map_err(error_report)?;
+                ensure!(nfields == 4, "too many entries in v1 kv record");
+
+                let bytes = bytes.remaining_slice();
+
+                let (namespace, bytes) =
+                    decode::read_str_from_slice(bytes).map_err(error_report)?;
+                let (key, mut bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?;
+                let has_value = decode::read_bool(&mut bytes).map_err(error_report)?;
+
+                let (value, bytes) = if has_value {
+                    let (value, bytes) =
+                        decode::read_str_from_slice(bytes).map_err(error_report)?;
+                    (Some(value.to_owned()), bytes)
+                } else {
+                    (None, bytes)
+                };
+
+                if !bytes.is_empty() {
+                    bail!("trailing bytes in encoded kvrecord. malformed")
+                }
+
+                Ok(KvRecord {
+                    namespace: namespace.to_owned(),
+                    key: key.to_owned(),
+                    value,
+                })
+            }
+            _ => {
+                bail!("unknown version {version:?}")
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::{DecryptedData, KV_VERSION, KvRecord};
+
+    #[test]
+    fn encode_decode_some() {
+        let kv = KvRecord {
+            namespace: "foo".to_owned(),
+            key: "bar".to_owned(),
+            value: Some("baz".to_owned()),
+        };
+        let snapshot = [
+            0x94, 0xa3, b'f', b'o', b'o', 0xa3, b'b', b'a', b'r', 0xc3, 0xa3, b'b', b'a', b'z',
+        ];
+
+        let encoded = kv.serialize().unwrap();
+        let decoded = KvRecord::deserialize(&encoded, KV_VERSION).unwrap();
+
+        assert_eq!(encoded.0, &snapshot);
+        assert_eq!(decoded, kv);
+    }
+
+    #[test]
+    fn encode_decode_none() {
+        let kv = KvRecord {
+            namespace: "foo".to_owned(),
+            key: "bar".to_owned(),
+            value: None,
+        };
+        let snapshot = [0x94, 0xa3, b'f', b'o', b'o', 0xa3, b'b', b'a', b'r', 0xc2];
+
+        let encoded = kv.serialize().unwrap();
+        let decoded = KvRecord::deserialize(&encoded, KV_VERSION).unwrap();
+
+        assert_eq!(encoded.0, &snapshot);
+        assert_eq!(decoded, kv);
+    }
+
+    #[test]
+    fn decode_v0() {
+        let kv = KvRecord {
+            namespace: "foo".to_owned(),
+            key: "bar".to_owned(),
+            value: Some("baz".to_owned()),
+        };
+
+        let snapshot = vec![
+            0x93, 0xa3, b'f', b'o', b'o', 0xa3, b'b', b'a', b'r', 0xa3, b'b', b'a', b'z',
+        ];
+
+        let decoded = KvRecord::deserialize(&DecryptedData(snapshot), "v0").unwrap();
+
+        assert_eq!(decoded, kv);
+    }
+}