From d107ba8a2811ae321ff54781a8a0d416ca22a2d3 Mon Sep 17 00:00:00 2001 From: qaz741wsd856 Date: Sun, 30 Nov 2025 04:31:45 +0000 Subject: [PATCH] refactor: enhance batch processing for cipher operations --- README.md | 2 +- src/handlers/ciphers.rs | 73 ++++++++++++++++++++++++++++++----------- src/handlers/import.rs | 9 +---- src/handlers/mod.rs | 13 ++++++++ 4 files changed, 69 insertions(+), 28 deletions(-) diff --git a/README.md b/README.md index 8219cfd..05b5a24 100644 --- a/README.md +++ b/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 diff --git a/src/handlers/ciphers.rs b/src/handlers/ciphers.rs index 5944ef4..79044f2 100644 --- a/src/handlers/ciphers.rs +++ b/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, + 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, 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, ) -> Result, 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, 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 = query!( &db, "SELECT * FROM ciphers WHERE id = ?1 AND user_id = ?2", diff --git a/src/handlers/import.rs b/src/handlers/import.rs index 2a40b6d..e5eb23b 100644 --- a/src/handlers/import.rs +++ b/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::().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( diff --git a/src/handlers/mod.rs b/src/handlers/mod.rs index 5644148..2bf4465 100644 --- a/src/handlers/mod.rs +++ b/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::().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) +}