refactor: enhance batch processing for cipher operations

README.md

@@ -159,7 +159,7 @@ You can configure the following environment variables in `wrangler.toml` under t
 
 *   **`IMPORT_BATCH_SIZE`** (Optional, Default: `30`)
     
-    Number of records to process in each batch when importing data. This helps manage memory usage and processing time for large imports.
+    Number of records to process in each batch when importing and deleting data. This helps manage memory usage and processing time for large imports.
     
     *   Set to `0` to disable batching (all records imported in a single batch)
     *   Defaults to `30` records per batch if not specified

src/handlers/ciphers.rs

@@ -1,9 +1,10 @@
+use super::get_batch_size;
 use axum::{extract::State, Json};
 use chrono::Utc;
 use serde::Deserialize;
 use std::sync::Arc;
 use uuid::Uuid;
-use worker::{query, Env};
+use worker::{query, D1PreparedStatement, Env};
 
 use crate::auth::Claims;
 use crate::db;
@@ -11,6 +12,27 @@ use crate::error::AppError;
 use crate::models::cipher::{Cipher, CipherData, CipherRequestData, CreateCipherRequest};
 use axum::extract::Path;
 
+/// Execute D1 statements in batches, allowing batch_size 0 to run everything at once.
+async fn execute_in_batches(
+    db: &worker::D1Database,
+    statements: Vec<D1PreparedStatement>,
+    batch_size: usize,
+) -> Result<(), AppError> {
+    if statements.is_empty() {
+        return Ok(());
+    }
+
+    if batch_size == 0 {
+        db.batch(statements).await?;
+    } else {
+        for chunk in statements.chunks(batch_size) {
+            db.batch(chunk.to_vec()).await?;
+        }
+    }
+
+    Ok(())
+}
+
 #[worker::send]
 pub async fn create_cipher(
     claims: Claims,
@@ -198,20 +220,24 @@ pub async fn soft_delete_ciphers_bulk(
 ) -> Result<Json<()>, AppError> {
     let db = db::get_db(&env)?;
     let now = Utc::now().format("%Y-%m-%dT%H:%M:%S%.3fZ").to_string();
+    let batch_size = get_batch_size(&env);
+    let mut statements = Vec::with_capacity(payload.ids.len());
 
     for id in payload.ids {
-        query!(
+        let stmt = query!(
             &db,
             "UPDATE ciphers SET deleted_at = ?1, updated_at = ?1 WHERE id = ?2 AND user_id = ?3",
             now,
             id,
             claims.sub
         )
-        .map_err(|_| AppError::Database)?
-        .run()
-        .await?;
+        .map_err(|_| AppError::Database)?;
+
+        statements.push(stmt);
     }
 
+    execute_in_batches(&db, statements, batch_size).await?;
+
     Ok(Json(()))
 }
 
@@ -246,19 +272,23 @@ pub async fn hard_delete_ciphers_bulk(
     Json(payload): Json<CipherIdsData>,
 ) -> Result<Json<()>, AppError> {
     let db = db::get_db(&env)?;
+    let batch_size = get_batch_size(&env);
+    let mut statements = Vec::with_capacity(payload.ids.len());
 
     for id in payload.ids {
-        query!(
+        let stmt = query!(
             &db,
             "DELETE FROM ciphers WHERE id = ?1 AND user_id = ?2",
             id,
             claims.sub
         )
-        .map_err(|_| AppError::Database)?
-        .run()
-        .await?;
+        .map_err(|_| AppError::Database)?;
+
+        statements.push(stmt);
     }
 
+    execute_in_batches(&db, statements, batch_size).await?;
+
     Ok(Json(()))
 }
 
@@ -318,23 +348,28 @@ pub async fn restore_ciphers_bulk(
 ) -> Result<Json<BulkRestoreResponse>, AppError> {
     let db = db::get_db(&env)?;
     let now = Utc::now().format("%Y-%m-%dT%H:%M:%S%.3fZ").to_string();
+    let batch_size = get_batch_size(&env);
+    let ids = payload.ids;
+    let mut update_statements = Vec::with_capacity(ids.len());
 
-    let mut restored_ciphers = Vec::new();
-
-    for id in payload.ids {
-        // Update the cipher to clear deleted_at
-        query!(
+    for id in ids.iter() {
+        let stmt = query!(
             &db,
             "UPDATE ciphers SET deleted_at = NULL, updated_at = ?1 WHERE id = ?2 AND user_id = ?3",
             now,
-            id,
+            id.clone(),
             claims.sub
         )
-        .map_err(|_| AppError::Database)?
-        .run()
-        .await?;
+        .map_err(|_| AppError::Database)?;
 
-        // Fetch the restored cipher
+        update_statements.push(stmt);
+    }
+
+    execute_in_batches(&db, update_statements, batch_size).await?;
+
+    let mut restored_ciphers = Vec::with_capacity(ids.len());
+
+    for id in ids {
         let cipher_db: Option<crate::models::cipher::CipherDBModel> = query!(
             &db,
             "SELECT * FROM ciphers WHERE id = ?1 AND user_id = ?2",

src/handlers/import.rs

@@ -11,14 +11,7 @@ use crate::models::cipher::{Cipher, CipherData};
 use crate::models::folder::Folder;
 use crate::models::import::ImportRequest;
 
-/// Get the batch size from environment variable IMPORT_BATCH_SIZE.
-/// Defaults to 30 if not set or invalid.
-fn get_batch_size(env: &Env) -> usize {
-    env.var("IMPORT_BATCH_SIZE")
-        .ok()
-        .and_then(|v| v.to_string().parse::<usize>().ok())
-        .unwrap_or(30)
-}
+use super::{get_batch_size};
 
 /// Execute statements in batches. If batch_size is 0, execute all in one batch.
 async fn execute_in_batches(

src/handlers/mod.rs

@@ -6,3 +6,16 @@ pub mod sync;
 pub mod folders;
 pub mod import;
 pub mod purge;
+
+/// Shared helper for reading an environment variable into usize.
+pub(crate) fn get_env_usize(env: &worker::Env, var_name: &str, default: usize) -> usize {
+    env.var(var_name)
+        .ok()
+        .and_then(|value| value.to_string().parse::<usize>().ok())
+        .unwrap_or(default)
+}
+
+/// Convenience helper for cipher batch size using IMPORT_BATCH_SIZE.
+pub(crate) fn get_batch_size(env: &worker::Env) -> usize {
+    get_env_usize(env, "IMPORT_BATCH_SIZE", 30)
+}