Add a parameter to the sync to specify the download/upload page

atuinsh · Oct 4, 2024 · 8e4280a · 8e4280a
1 parent 80e950c
commit 8e4280a
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 9 deletions.
diff --git a/crates/atuin-client/src/record/sync.rs b/crates/atuin-client/src/record/sync.rs
@@ -163,10 +163,10 @@ async fn sync_upload(
     tag: String,
     local: RecordIdx,
     remote: Option<RecordIdx>,
+    page_size: u64,
 ) -> Result<i64, SyncError> {
     let remote = remote.unwrap_or(0);
     let expected = local - remote;
-    let upload_page_size = 100;
     let mut progress = 0;
 
     let pb = ProgressBar::new(expected);
@@ -185,7 +185,7 @@ async fn sync_upload(
     // preload with the first entry if remote does not know of this store
     loop {
         let page = store
-            .next(host, tag.as_str(), remote + progress, upload_page_size)
+            .next(host, tag.as_str(), remote + progress, page_size)
             .await
             .map_err(|e| {
                 error!("failed to read upload page: {e:?}");
@@ -219,10 +219,10 @@ async fn sync_download(
     tag: String,
     local: Option<RecordIdx>,
     remote: RecordIdx,
+    page_size: u64,
 ) -> Result<Vec<RecordId>, SyncError> {
     let local = local.unwrap_or(0);
     let expected = remote - local;
-    let download_page_size = 100;
     let mut progress = 0;
     let mut ret = Vec::new();
 
@@ -242,7 +242,7 @@ async fn sync_download(
     // preload with the first entry if remote does not know of this store
     loop {
         let page = client
-            .next_records(host, tag.clone(), local + progress, download_page_size)
+            .next_records(host, tag.clone(), local + progress, page_size)
             .await
             .map_err(|e| SyncError::RemoteRequestError { msg: e.to_string() })?;
 
@@ -270,6 +270,7 @@ pub async fn sync_remote(
     operations: Vec<Operation>,
     local_store: &impl Store,
     settings: &Settings,
+    page_size: u64,
 ) -> Result<(i64, Vec<RecordId>), SyncError> {
     let client = Client::new(
         &settings.sync_address,
@@ -293,15 +294,20 @@ pub async fn sync_remote(
                 tag,
                 local,
                 remote,
-            } => uploaded += sync_upload(local_store, &client, host, tag, local, remote).await?,
+            } => {
+                uploaded += 
+                    sync_upload(local_store, &client, host, tag, local, remote, page_size).await?
+            }
 
             Operation::Download {
                 host,
                 tag,
                 local,
                 remote,
             } => {
-                let mut d = sync_download(local_store, &client, host, tag, local, remote).await?;
+                let mut d = 
+                    sync_download(local_store, &client, host, tag, local, remote, page_size)
+                        .await?;
                 downloaded.append(&mut d)
             }
 
@@ -318,7 +324,7 @@ pub async fn sync(
 ) -> Result<(i64, Vec<RecordId>), SyncError> {
     let (diff, _) = diff(settings, store).await?;
     let operations = operations(diff, store).await?;
-    let (uploaded, downloaded) = sync_remote(operations, store, settings).await?;
+    let (uploaded, downloaded) = sync_remote(operations, store, settings, 100).await?;
 
     Ok((uploaded, downloaded))
 }

diff --git a/crates/atuin/src/command/client/store/pull.rs b/crates/atuin/src/command/client/store/pull.rs
@@ -19,6 +19,11 @@ pub struct Pull {
     /// This will first wipe the local store, and then download all records from the remote
     #[arg(long, default_value = "false")]
     pub force: bool,
+
+    /// Page Size
+    /// How many records to upload at once. Defaults to 100
+    #[arg(long, default_value = "100")]
+    pub page: u64,
 }
 
 impl Pull {
@@ -67,7 +72,7 @@ impl Pull {
             })
             .collect();
 
-        let (_, downloaded) = sync::sync_remote(operations, &store, settings).await?;
+        let (_, downloaded) = sync::sync_remote(operations, &store, settings, self.page).await?;
 
         println!("Downloaded {} records", downloaded.len());
 

diff --git a/crates/atuin/src/command/client/store/push.rs b/crates/atuin/src/command/client/store/push.rs
@@ -25,6 +25,11 @@ pub struct Push {
     /// local store
     #[arg(long, default_value = "false")]
     pub force: bool,
+
+    /// Page Size
+    /// How many records to upload at once. Defaults to 100
+    #[arg(long, default_value = "100")]
+    pub page: u64,
 }
 
 impl Push {
@@ -87,7 +92,7 @@ impl Push {
             })
             .collect();
 
-        let (uploaded, _) = sync::sync_remote(operations, &store, settings).await?;
+        let (uploaded, _) = sync::sync_remote(operations, &store, settings, self.page).await?;
 
         println!("Uploaded {uploaded} records");