refactor: move execute_in_batches function to db module for better reusability

This commit is contained in:
qaz741wsd856 2025-11-30 06:59:59 +00:00
parent e965276754
commit 6f842cc2ce
3 changed files with 31 additions and 54 deletions

View file

@ -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(())
}

View file

@ -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());

View file

@ -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(()))
}