aboutsummaryrefslogtreecommitdiffstats
path: root/atuin-client/src/history/store.rs
diff options
context:
space:
mode:
authorEllie Huxtable <ellie@elliehuxtable.com>2024-04-18 16:41:28 +0100
committerGitHub <noreply@github.com>2024-04-18 16:41:28 +0100
commit95cc472037fcb3207b510e67f1a44af4e2a2cae9 (patch)
treefc1d3e71d8e0bdb806370e4144fd6f373bcc9c5e /atuin-client/src/history/store.rs
parentfeat: show preview auto (#1804) (diff)
downloadatuin-95cc472037fcb3207b510e67f1a44af4e2a2cae9.zip
chore: move crates into crates/ dir (#1958)
I'd like to tidy up the root a little, and it's nice to have all the rust crates in one place
Diffstat (limited to 'atuin-client/src/history/store.rs')
-rw-r--r--atuin-client/src/history/store.rs410
1 files changed, 0 insertions, 410 deletions
diff --git a/atuin-client/src/history/store.rs b/atuin-client/src/history/store.rs
deleted file mode 100644
index fe2b7b92..00000000
--- a/atuin-client/src/history/store.rs
+++ /dev/null
@@ -1,410 +0,0 @@
-use std::{collections::HashSet, fmt::Write, time::Duration};
-
-use eyre::{bail, eyre, Result};
-use indicatif::{ProgressBar, ProgressState, ProgressStyle};
-use rmp::decode::Bytes;
-
-use crate::{
- database::{current_context, Database},
- record::{encryption::PASETO_V4, sqlite_store::SqliteStore, store::Store},
-};
-use atuin_common::record::{DecryptedData, Host, HostId, Record, RecordId, RecordIdx};
-
-use super::{History, HistoryId, HISTORY_TAG, HISTORY_VERSION};
-
-#[derive(Debug)]
-pub struct HistoryStore {
- pub store: SqliteStore,
- pub host_id: HostId,
- pub encryption_key: [u8; 32],
-}
-
-#[derive(Debug, Eq, PartialEq, Clone)]
-pub enum HistoryRecord {
- Create(History), // Create a history record
- Delete(HistoryId), // Delete a history record, identified by ID
-}
-
-impl HistoryRecord {
- /// Serialize a history record, returning DecryptedData
- /// The record will be of a certain type
- /// We map those like so:
- ///
- /// HistoryRecord::Create -> 0
- /// HistoryRecord::Delete-> 1
- ///
- /// This numeric identifier is then written as the first byte to the buffer. For history, we
- /// append the serialized history right afterwards, to avoid having to handle serialization
- /// twice.
- ///
- /// Deletion simply refers to the history by ID
- pub fn serialize(&self) -> Result<DecryptedData> {
- // probably don't actually need to use rmp here, but if we ever need to extend it, it's a
- // nice wrapper around raw byte stuff
- use rmp::encode;
-
- let mut output = vec![];
-
- match self {
- HistoryRecord::Create(history) => {
- // 0 -> a history create
- encode::write_u8(&mut output, 0)?;
-
- let bytes = history.serialize()?;
-
- encode::write_bin(&mut output, &bytes.0)?;
- }
- HistoryRecord::Delete(id) => {
- // 1 -> a history delete
- encode::write_u8(&mut output, 1)?;
- encode::write_str(&mut output, id.0.as_str())?;
- }
- };
-
- Ok(DecryptedData(output))
- }
-
- pub fn deserialize(bytes: &DecryptedData, version: &str) -> Result<Self> {
- use rmp::decode;
-
- fn error_report<E: std::fmt::Debug>(err: E) -> eyre::Report {
- eyre!("{err:?}")
- }
-
- let mut bytes = Bytes::new(&bytes.0);
-
- let record_type = decode::read_u8(&mut bytes).map_err(error_report)?;
-
- match record_type {
- // 0 -> HistoryRecord::Create
- 0 => {
- // not super useful to us atm, but perhaps in the future
- // written by write_bin above
- let _ = decode::read_bin_len(&mut bytes).map_err(error_report)?;
-
- let record = History::deserialize(bytes.remaining_slice(), version)?;
-
- Ok(HistoryRecord::Create(record))
- }
-
- // 1 -> HistoryRecord::Delete
- 1 => {
- let bytes = bytes.remaining_slice();
- let (id, bytes) = decode::read_str_from_slice(bytes).map_err(error_report)?;
-
- if !bytes.is_empty() {
- bail!(
- "trailing bytes decoding HistoryRecord::Delete - malformed? got {bytes:?}"
- );
- }
-
- Ok(HistoryRecord::Delete(id.to_string().into()))
- }
-
- n => {
- bail!("unknown HistoryRecord type {n}")
- }
- }
- }
-}
-
-impl HistoryStore {
- pub fn new(store: SqliteStore, host_id: HostId, encryption_key: [u8; 32]) -> Self {
- HistoryStore {
- store,
- host_id,
- encryption_key,
- }
- }
-
- async fn push_record(&self, record: HistoryRecord) -> Result<(RecordId, RecordIdx)> {
- let bytes = record.serialize()?;
- let idx = self
- .store
- .last(self.host_id, HISTORY_TAG)
- .await?
- .map_or(0, |p| p.idx + 1);
-
- let record = Record::builder()
- .host(Host::new(self.host_id))
- .version(HISTORY_VERSION.to_string())
- .tag(HISTORY_TAG.to_string())
- .idx(idx)
- .data(bytes)
- .build();
-
- let id = record.id;
-
- self.store
- .push(&record.encrypt::<PASETO_V4>(&self.encryption_key))
- .await?;
-
- Ok((id, idx))
- }
-
- async fn push_batch(&self, records: impl Iterator<Item = HistoryRecord>) -> Result<()> {
- let mut ret = Vec::new();
-
- let idx = self
- .store
- .last(self.host_id, HISTORY_TAG)
- .await?
- .map_or(0, |p| p.idx + 1);
-
- // Could probably _also_ do this as an iterator, but let's see how this is for now.
- // optimizing for minimal sqlite transactions, this code can be optimised later
- for (n, record) in records.enumerate() {
- let bytes = record.serialize()?;
-
- let record = Record::builder()
- .host(Host::new(self.host_id))
- .version(HISTORY_VERSION.to_string())
- .tag(HISTORY_TAG.to_string())
- .idx(idx + n as u64)
- .data(bytes)
- .build();
-
- let record = record.encrypt::<PASETO_V4>(&self.encryption_key);
-
- ret.push(record);
- }
-
- self.store.push_batch(ret.iter()).await?;
-
- Ok(())
- }
-
- pub async fn delete(&self, id: HistoryId) -> Result<(RecordId, RecordIdx)> {
- let record = HistoryRecord::Delete(id);
-
- self.push_record(record).await
- }
-
- pub async fn push(&self, history: History) -> Result<(RecordId, RecordIdx)> {
- // TODO(ellie): move the history store to its own file
- // it's tiny rn so fine as is
- let record = HistoryRecord::Create(history);
-
- self.push_record(record).await
- }
-
- pub async fn history(&self) -> Result<Vec<HistoryRecord>> {
- // Atm this loads all history into memory
- // Not ideal as that is potentially quite a lot, although history will be small.
- let records = self.store.all_tagged(HISTORY_TAG).await?;
- let mut ret = Vec::with_capacity(records.len());
-
- for record in records.into_iter() {
- let hist = match record.version.as_str() {
- HISTORY_VERSION => {
- let decrypted = record.decrypt::<PASETO_V4>(&self.encryption_key)?;
-
- HistoryRecord::deserialize(&decrypted.data, HISTORY_VERSION)
- }
- version => bail!("unknown history version {version:?}"),
- }?;
-
- ret.push(hist);
- }
-
- Ok(ret)
- }
-
- pub async fn build(&self, database: &dyn Database) -> Result<()> {
- // I'd like to change how we rebuild and not couple this with the database, but need to
- // consider the structure more deeply. This will be easy to change.
-
- // TODO(ellie): page or iterate this
- let history = self.history().await?;
-
- // In theory we could flatten this here
- // The current issue is that the database may have history in it already, from the old sync
- // This didn't actually delete old history
- // If we're sure we have a DB only maintained by the new store, we can flatten
- // create/delete before we even get to sqlite
- let mut creates = Vec::new();
- let mut deletes = Vec::new();
-
- for i in history {
- match i {
- HistoryRecord::Create(h) => {
- creates.push(h);
- }
- HistoryRecord::Delete(id) => {
- deletes.push(id);
- }
- }
- }
-
- database.save_bulk(&creates).await?;
- database.delete_rows(&deletes).await?;
-
- Ok(())
- }
-
- pub async fn incremental_build(&self, database: &dyn Database, ids: &[RecordId]) -> Result<()> {
- for id in ids {
- let record = self.store.get(*id).await;
-
- let record = if let Ok(record) = record {
- record
- } else {
- continue;
- };
-
- if record.tag != HISTORY_TAG {
- continue;
- }
-
- let decrypted = record.decrypt::<PASETO_V4>(&self.encryption_key)?;
- let record = HistoryRecord::deserialize(&decrypted.data, HISTORY_VERSION)?;
-
- match record {
- HistoryRecord::Create(h) => {
- // TODO: benchmark CPU time/memory tradeoff of batch commit vs one at a time
- database.save(&h).await?;
- }
- HistoryRecord::Delete(id) => {
- database.delete_rows(&[id]).await?;
- }
- }
- }
-
- Ok(())
- }
-
- /// Get a list of history IDs that exist in the store
- /// Note: This currently involves loading all history into memory. This is not going to be a
- /// large amount in absolute terms, but do not all it in a hot loop.
- pub async fn history_ids(&self) -> Result<HashSet<HistoryId>> {
- let history = self.history().await?;
-
- let ret = HashSet::from_iter(history.iter().map(|h| match h {
- HistoryRecord::Create(h) => h.id.clone(),
- HistoryRecord::Delete(id) => id.clone(),
- }));
-
- Ok(ret)
- }
-
- pub async fn init_store(&self, db: &impl Database) -> Result<()> {
- let pb = ProgressBar::new_spinner();
- pb.set_style(
- ProgressStyle::with_template("{spinner:.blue} {msg}")
- .unwrap()
- .with_key("eta", |state: &ProgressState, w: &mut dyn Write| {
- write!(w, "{:.1}s", state.eta().as_secs_f64()).unwrap()
- })
- .progress_chars("#>-"),
- );
- pb.enable_steady_tick(Duration::from_millis(500));
-
- pb.set_message("Fetching history from old database");
-
- let context = current_context();
- let history = db.list(&[], &context, None, false, true).await?;
-
- pb.set_message("Fetching history already in store");
- let store_ids = self.history_ids().await?;
-
- pb.set_message("Converting old history to new store");
- let mut records = Vec::new();
-
- for i in history {
- debug!("loaded {}", i.id);
-
- if store_ids.contains(&i.id) {
- debug!("skipping {} - already exists", i.id);
- continue;
- }
-
- if i.deleted_at.is_some() {
- records.push(HistoryRecord::Delete(i.id));
- } else {
- records.push(HistoryRecord::Create(i));
- }
- }
-
- pb.set_message("Writing to db");
-
- if !records.is_empty() {
- self.push_batch(records.into_iter()).await?;
- }
-
- pb.finish_with_message("Import complete");
-
- Ok(())
- }
-}
-
-#[cfg(test)]
-mod tests {
- use atuin_common::record::DecryptedData;
- use time::macros::datetime;
-
- use crate::history::{store::HistoryRecord, HISTORY_VERSION};
-
- use super::History;
-
- #[test]
- fn test_serialize_deserialize_create() {
- let bytes = [
- 204, 0, 196, 141, 205, 0, 0, 153, 217, 32, 48, 49, 56, 99, 100, 52, 102, 101, 56, 49,
- 55, 53, 55, 99, 100, 50, 97, 101, 101, 54, 53, 99, 100, 55, 56, 54, 49, 102, 57, 99,
- 56, 49, 207, 23, 166, 251, 212, 181, 82, 0, 0, 100, 0, 162, 108, 115, 217, 41, 47, 85,
- 115, 101, 114, 115, 47, 101, 108, 108, 105, 101, 47, 115, 114, 99, 47, 103, 105, 116,
- 104, 117, 98, 46, 99, 111, 109, 47, 97, 116, 117, 105, 110, 115, 104, 47, 97, 116, 117,
- 105, 110, 217, 32, 48, 49, 56, 99, 100, 52, 102, 101, 97, 100, 56, 57, 55, 53, 57, 55,
- 56, 53, 50, 53, 50, 55, 97, 51, 49, 99, 57, 57, 56, 48, 53, 57, 170, 98, 111, 111, 112,
- 58, 101, 108, 108, 105, 101, 192,
- ];
-
- let history = History {
- id: "018cd4fe81757cd2aee65cd7861f9c81".to_owned().into(),
- timestamp: datetime!(2024-01-04 00:00:00.000000 +00:00),
- duration: 100,
- exit: 0,
- command: "ls".to_owned(),
- cwd: "/Users/ellie/src/github.com/atuinsh/atuin".to_owned(),
- session: "018cd4fead897597852527a31c998059".to_owned(),
- hostname: "boop:ellie".to_owned(),
- deleted_at: None,
- };
-
- let record = HistoryRecord::Create(history);
-
- let serialized = record.serialize().expect("failed to serialize history");
- assert_eq!(serialized.0, bytes);
-
- let deserialized = HistoryRecord::deserialize(&serialized, HISTORY_VERSION)
- .expect("failed to deserialize HistoryRecord");
- assert_eq!(deserialized, record);
-
- // check the snapshot too
- let deserialized =
- HistoryRecord::deserialize(&DecryptedData(Vec::from(bytes)), HISTORY_VERSION)
- .expect("failed to deserialize HistoryRecord");
- assert_eq!(deserialized, record);
- }
-
- #[test]
- fn test_serialize_deserialize_delete() {
- let bytes = [
- 204, 1, 217, 32, 48, 49, 56, 99, 100, 52, 102, 101, 56, 49, 55, 53, 55, 99, 100, 50,
- 97, 101, 101, 54, 53, 99, 100, 55, 56, 54, 49, 102, 57, 99, 56, 49,
- ];
- let record = HistoryRecord::Delete("018cd4fe81757cd2aee65cd7861f9c81".to_string().into());
-
- let serialized = record.serialize().expect("failed to serialize history");
- assert_eq!(serialized.0, bytes);
-
- let deserialized = HistoryRecord::deserialize(&serialized, HISTORY_VERSION)
- .expect("failed to deserialize HistoryRecord");
- assert_eq!(deserialized, record);
-
- let deserialized =
- HistoryRecord::deserialize(&DecryptedData(Vec::from(bytes)), HISTORY_VERSION)
- .expect("failed to deserialize HistoryRecord");
- assert_eq!(deserialized, record);
- }
-}