From 6f842cc2ce4d157132711a972f8e2d26adf415ac Mon Sep 17 00:00:00 2001
From: qaz741wsd856 <hysth@163.com>
Date: Sun, 30 Nov 2025 06:59:59 +0000
Subject: [PATCH] refactor: move execute_in_batches function to db module for
 better reusability

---
 src/db.rs               | 23 ++++++++++++++++++++++-
 src/handlers/ciphers.rs | 33 ++++++---------------------------
 src/handlers/import.rs  | 29 +++--------------------------
 3 files changed, 31 insertions(+), 54 deletions(-)

diff --git a/src/db.rs b/src/db.rs
index 7416df1..603ec99 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -1,7 +1,28 @@
 use crate::error::AppError;
 use std::sync::Arc;
-use worker::{D1Database, Env};
+use worker::{D1Database, D1PreparedStatement, Env};
 
 pub fn get_db(env: &Arc<Env>) -> Result<D1Database, AppError> {
     env.d1("vault1").map_err(AppError::Worker)
 }
+
+/// Execute D1 statements in batches, allowing batch_size 0 to run everything at once.
+pub async fn execute_in_batches(
+    db: &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(())
+}
diff --git a/src/handlers/ciphers.rs b/src/handlers/ciphers.rs
index 79044f2..60b9e86 100644
--- a/src/handlers/ciphers.rs
+++ b/src/handlers/ciphers.rs
@@ -12,27 +12,6 @@ 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,
@@ -221,7 +200,7 @@ pub async fn soft_delete_ciphers_bulk(
     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());
+    let mut statements: Vec<D1PreparedStatement> = Vec::with_capacity(payload.ids.len());
 
     for id in payload.ids {
         let stmt = query!(
@@ -236,7 +215,7 @@ pub async fn soft_delete_ciphers_bulk(
         statements.push(stmt);
     }
 
-    execute_in_batches(&db, statements, batch_size).await?;
+    db::execute_in_batches(&db, statements, batch_size).await?;
 
     Ok(Json(()))
 }
@@ -273,7 +252,7 @@ pub async fn hard_delete_ciphers_bulk(
 ) -> 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());
+    let mut statements: Vec<D1PreparedStatement> = Vec::with_capacity(payload.ids.len());
 
     for id in payload.ids {
         let stmt = query!(
@@ -287,7 +266,7 @@ pub async fn hard_delete_ciphers_bulk(
         statements.push(stmt);
     }
 
-    execute_in_batches(&db, statements, batch_size).await?;
+    db::execute_in_batches(&db, statements, batch_size).await?;
 
     Ok(Json(()))
 }
@@ -350,7 +329,7 @@ pub async fn restore_ciphers_bulk(
     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 update_statements: Vec<D1PreparedStatement> = Vec::with_capacity(ids.len());
 
     for id in ids.iter() {
         let stmt = query!(
@@ -365,7 +344,7 @@ pub async fn restore_ciphers_bulk(
         update_statements.push(stmt);
     }
 
-    execute_in_batches(&db, update_statements, batch_size).await?;
+    db::execute_in_batches(&db, update_statements, batch_size).await?;
 
     let mut restored_ciphers = Vec::with_capacity(ids.len());
 
diff --git a/src/handlers/import.rs b/src/handlers/import.rs
index e5eb23b..a15d6db 100644
--- a/src/handlers/import.rs
+++ b/src/handlers/import.rs
@@ -11,30 +11,7 @@ use crate::models::cipher::{Cipher, CipherData};
 use crate::models::folder::Folder;
 use crate::models::import::ImportRequest;
 
-use super::{get_batch_size};
-
-/// Execute statements in batches. If batch_size is 0, execute all in one batch.
-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 {
-        // Execute all statements in one batch
-        db.batch(statements).await?;
-    } else {
-        // Execute in chunks of batch_size
-        for chunk in statements.chunks(batch_size) {
-            db.batch(chunk.to_vec()).await?;
-        }
-    }
-
-    Ok(())
-}
+use super::get_batch_size;
 
 #[worker::send]
 pub async fn import_data(
@@ -74,7 +51,7 @@ pub async fn import_data(
     }
 
     // Execute folder inserts in batches
-    execute_in_batches(&db, folder_statements, batch_size).await?;
+    db::execute_in_batches(&db, folder_statements, batch_size).await?;
 
     // Process folder relationships
     for relationship in payload.folder_relationships {
@@ -146,7 +123,7 @@ pub async fn import_data(
     }
 
     // Execute cipher inserts in batches
-    execute_in_batches(&db, cipher_statements, batch_size).await?;
+    db::execute_in_batches(&db, cipher_statements, batch_size).await?;
 
     Ok(Json(()))
 }
\ No newline at end of file