feat: Add a parameter to the sync to specify the download/upload page (#2408)

This adds a parameter to change the upload/download page size, this
could help if one gets a '413' in a selfhosted scenario.

## Checks
- [x] I am happy for maintainers to push small adjustments to this PR,
to speed up the review cycle
- [x] I have checked that there are no existing pull requests for the
same thing

---------

Co-authored-by: Ellie Huxtable <ellie@elliehuxtable.com>

3 files changed, 25 insertions, 9 deletions

diff --git a/crates/atuin-client/src/record/sync.rs b/crates/atuin-client/src/record/sync.rs
index a47f100b..bd357b79 100644
--- 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);
@@ -184,7 +184,7 @@ async fn sync_upload(
 
     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:?}");
@@ -222,10 +222,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();
 
@@ -244,7 +244,7 @@ async fn sync_download(
 
     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() })?;
 
@@ -276,6 +276,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,
@@ -299,7 +300,10 @@ 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,
@@ -307,7 +311,9 @@ pub async fn sync_remote(
                 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)
             }
 
@@ -324,7 +330,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
index df6e7346..af7ca00e 100644
--- 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 download 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
index 78d119bc..b5e13092 100644
--- 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");