Allow reading tail across hosts

4 files changed, 24 insertions, 11 deletions

diff --git a/atuin-client/src/kv.rs b/atuin-client/src/kv.rs
index 021b73aa..de3aa2fc 100644
--- a/atuin-client/src/kv.rs
+++ b/atuin-client/src/kv.rs
@@ -101,7 +101,10 @@ impl KvStore {
 
         let bytes = record.serialize()?;
 
-        let parent = store.tail(host_id, KV_TAG).await?.map(|entry| entry.id);
+        let parent = store
+            .tail(Some(host_id), KV_TAG)
+            .await?
+            .map(|entry| entry.id);
 
         let record = atuin_common::record::Record::builder()
             .host(host_id)
@@ -127,15 +130,12 @@ impl KvStore {
         namespace: &str,
         key: &str,
     ) -> Result<Option<KvRecord>> {
-        // TODO: don't load this from disk so much
-        let host_id = Settings::host_id().expect("failed to get host_id");
-
         // Currently, this is O(n). When we have an actual KV store, it can be better
         // Just a poc for now!
 
         // iterate records to find the value we want
         // start at the end, so we get the most recent version
-        let Some(mut record) = store.tail(host_id, KV_TAG).await? else {
+        let Some(mut record) = store.tail(None, KV_TAG).await? else {
             return Ok(None);
         };
 
diff --git a/atuin-client/src/record/sqlite_store.rs b/atuin-client/src/record/sqlite_store.rs
index ebfd666e..c0d75282 100644
--- a/atuin-client/src/record/sqlite_store.rs
+++ b/atuin-client/src/record/sqlite_store.rs
@@ -171,15 +171,28 @@ impl Store for SqliteStore {
         Ok(res)
     }
 
-    async fn tail(&self, host: Uuid, tag: &str) -> Result<Option<Record<EncryptedData>>> {
-        let res = sqlx::query(
+    // Get the tail for a given tag
+    // If a host is provided, get the tail for a tag for that host
+    // Otherwise, the latest record across hosts
+    async fn tail(&self, host: Option<Uuid>, tag: &str) -> Result<Option<Record<EncryptedData>>> {
+        let res = if let Some(host) = host {
+            sqlx::query(
             "select * from records rp where tag=?1 and host=?2 and (select count(1) from records where parent=rp.id) = 0;",
         )
         .bind(tag)
         .bind(host.as_simple().to_string())
         .map(Self::query_row)
         .fetch_optional(&self.pool)
-        .await?;
+        .await?
+        } else {
+            sqlx::query(
+            "select * from records rp where tag=?1 and (select count(1) from records where parent=rp.id) = 0;",
+        )
+        .bind(tag)
+        .map(Self::query_row)
+        .fetch_optional(&self.pool)
+        .await?
+        };
 
         Ok(res)
     }
diff --git a/atuin-client/src/record/store.rs b/atuin-client/src/record/store.rs
index 029cdeb6..79f27cbd 100644
--- a/atuin-client/src/record/store.rs
+++ b/atuin-client/src/record/store.rs
@@ -30,7 +30,7 @@ pub trait Store {
     /// Get the first record for a given host and tag
     async fn head(&self, host: Uuid, tag: &str) -> Result<Option<Record<EncryptedData>>>;
     /// Get the last record for a given host and tag
-    async fn tail(&self, host: Uuid, tag: &str) -> Result<Option<Record<EncryptedData>>>;
+    async fn tail(&self, host: Option<Uuid>, tag: &str) -> Result<Option<Record<EncryptedData>>>;
 
     async fn tail_records(&self) -> Result<Vec<(Uuid, String, Uuid)>>;
 }
diff --git a/atuin-client/src/record/sync.rs b/atuin-client/src/record/sync.rs
index 1905fe26..8d0cbf6f 100644
--- a/atuin-client/src/record/sync.rs
+++ b/atuin-client/src/record/sync.rs
@@ -49,7 +49,7 @@ pub async fn operations(diff: Diff, store: &impl Store) -> Result<Vec<Operation>
             // if local has the ID, then we should find the actual tail of this
             // store, so we know what we need to update the remote to.
             let tail = store
-                .tail(host, tag.as_str())
+                .tail(Some(host), tag.as_str())
                 .await?
                 .expect("failed to fetch last record, expected tag/host to exist");
 
@@ -163,7 +163,7 @@ async fn sync_download(
     let remote_tail = remote_index
         .get(op.0, op.1.clone())
         .expect("remote index does not contain expected tail during download");
-    let local_tail = store.tail(op.0, op.1.as_str()).await?;
+    let local_tail = store.tail(Some(op.0), op.1.as_str()).await?;
     //
     // We expect that the operations diff will represent the desired state
     // In this case, that contains the remote tail.