diff --git a/backend/src/error.rs b/backend/src/error.rs index 62dfbfc..7a75b8c 100644 --- a/backend/src/error.rs +++ b/backend/src/error.rs @@ -6,6 +6,7 @@ use axum::{ use serde::Serialize; use serde_json::Value; use std::fmt::Display; +use utoipa::ToSchema; pub type AppResult = Result; @@ -61,7 +62,7 @@ impl AppError { impl IntoResponse for AppError { fn into_response(self) -> Response { let status = self.status; - let body = Json(ErrorResponse { + let body = Json(ApiErrorResponse { error: self.message, code: self.code, details: self.details, @@ -70,8 +71,8 @@ impl IntoResponse for AppError { } } -#[derive(Serialize)] -struct ErrorResponse { +#[derive(Serialize, ToSchema)] +pub struct ApiErrorResponse { error: String, #[serde(skip_serializing_if = "Option::is_none")] code: Option, diff --git a/backend/src/workers/analyze.rs b/backend/src/workers/analyze.rs index 3678644..d997c9a 100644 --- a/backend/src/workers/analyze.rs +++ b/backend/src/workers/analyze.rs @@ -1,23 +1,24 @@ -use std::{collections::HashSet, sync::Arc, time::Duration}; +use std::collections::HashSet; +use std::sync::Arc; use async_trait::async_trait; -use diesel::prelude::*; use serde::Deserialize; -use serde_json::json; -use tokio::task; -use tracing::{error, warn}; use uuid::Uuid; -use super::ocr::{document_is_pdf, OCR_TEXT_ASSET_TYPE}; use crate::{ - jobs::{enqueue_job, JOB_ANALYZE_DOCUMENT, JOB_GENERATE_OCR_TEXT, JOB_GENERATE_THUMBNAILS}, - models::{Document, DocumentAsset, DocumentVersion}, - schema::{document_assets, document_versions, documents}, - state::AppState, - storage::TenantStorage, + jobs::JOB_ANALYZE_DOCUMENT, models::Document, state::AppState, storage::TenantStorage, }; -use super::{JobExecution, JobHandler}; +use super::{ + index::IndexDocumentTask, + job_execution_from_task_error, + ocr::{GenerateOcrTask, OCR_TEXT_ASSET_TYPE}, + taskflow::{ + document::DocumentVersionTaskContext, BoxedTask, TaskExecutor, TaskPlanner, TaskResult, + }, + thumbnails::GenerateThumbnailsTask, + JobExecution, JobHandler, +}; #[derive(Debug, Deserialize)] struct AnalyzePayload { @@ -45,7 +46,7 @@ impl JobHandler for AnalyzeDocumentJob { &self, state: Arc, job: crate::models::Job, - _storage: TenantStorage, + storage: TenantStorage, ) -> JobExecution { let payload: AnalyzePayload = match serde_json::from_value(job.payload.clone()) { Ok(payload) => payload, @@ -56,114 +57,70 @@ impl JobHandler for AnalyzeDocumentJob { } }; - let state_clone = state.clone(); - let tenant_id = job.tenant_id; - match task::spawn_blocking(move || analyze_document(state_clone, tenant_id, payload)).await - { - Ok(Ok(execution)) => execution, - Ok(Err(err)) => { - warn!(job_id = %job.id, error = %err, "analyze job will retry"); - JobExecution::Retry { - delay: Duration::from_secs(30), - error: err, - } - } - Err(join_err) => { - error!(job_id = %job.id, error = %join_err, "analyze task panicked"); - JobExecution::Retry { - delay: Duration::from_secs(60), - error: format!("worker panicked: {join_err}"), - } - } + let mut context = DocumentVersionTaskContext::new( + job.id, + JOB_ANALYZE_DOCUMENT, + job.tenant_id, + payload.document_id, + payload.document_version_id, + payload.force, + state.config.worker_max_document_bytes, + state.clone(), + storage, + ); + + let planner = AnalyzePlanner::new(payload.force, state.clone()); + match TaskExecutor::run(&planner, &mut context).await { + Ok(()) => JobExecution::Success, + Err(err) => job_execution_from_task_error(err), } } } -fn analyze_document( +struct AnalyzePlanner { + force: bool, state: Arc, - tenant_id: Uuid, - payload: AnalyzePayload, -) -> Result { - let mut conn = state - .db_for_tenant(tenant_id) - .map_err(|err| format!("{err:?}"))?; +} - let version: DocumentVersion = document_versions::table - .find(payload.document_version_id) - .first(&mut conn) - .map_err(|err| { - format!( - "failed to load document_version {} for tenant {}: {err:?}", - payload.document_version_id, tenant_id - ) - })?; - - if version.document_id != payload.document_id { - return Err("document/version mismatch".into()); +impl AnalyzePlanner { + fn new(force: bool, state: Arc) -> Self { + Self { force, state } } +} - let document: Document = documents::table - .find(payload.document_id) - .first(&mut conn) - .map_err(|err| { - format!( - "failed to load document {} for tenant {}: {err:?}", - payload.document_id, tenant_id - ) - })?; +#[async_trait] +impl TaskPlanner for AnalyzePlanner { + async fn plan( + &self, + ctx: &mut DocumentVersionTaskContext, + ) -> TaskResult>> { + let document = ctx.document().await?.clone(); + let mut tasks: Vec> = Vec::new(); - let tenant_id = document.tenant_id; - - let (supported, _reason) = determine_thumbnail_support(&document); - let ocr_supported = document_is_pdf(&document); - - let existing_ocr: Option = document_assets::table - .filter(document_assets::document_version_id.eq(payload.document_version_id)) - .filter(document_assets::asset_type.eq(OCR_TEXT_ASSET_TYPE)) - .filter(document_assets::tenant_id.eq(tenant_id)) - .first(&mut conn) - .optional() - .map_err(|err| format!("{err:?}"))?; - - let skip_ocr = existing_ocr.is_some() && !payload.force; - - if supported { - let enqueue_result = enqueue_job( - &mut conn, - tenant_id, - JOB_GENERATE_THUMBNAILS, - json!({ - "document_id": payload.document_id, - "document_version_id": payload.document_version_id, - "force": payload.force, - }), - None, - ); - - if let Err(err) = enqueue_result { - return Err(err.to_string()); + let (thumbnail_supported, _) = determine_thumbnail_support(&document); + if thumbnail_supported { + tasks.push(Box::new(GenerateThumbnailsTask::new(self.force))); } - } - if ocr_supported && !skip_ocr { - let enqueue_result = enqueue_job( - &mut conn, - tenant_id, - JOB_GENERATE_OCR_TEXT, - json!({ - "document_id": payload.document_id, - "document_version_id": payload.document_version_id, - "force": payload.force, - }), - None, - ); + let existing_ocr = ctx.asset(OCR_TEXT_ASSET_TYPE).await?.is_some(); + let mut should_index = existing_ocr; - if let Err(err) = enqueue_result { - return Err(err.to_string()); + if document_supports_ocr(&document) { + if self.force || !existing_ocr { + tasks.push(Box::new(GenerateOcrTask::new( + self.force, + self.state.clone(), + ))); + should_index = true; + } } - } - Ok(JobExecution::Success) + if should_index { + tasks.push(Box::new(IndexDocumentTask::new())); + } + + Ok(tasks) + } } pub(crate) fn determine_thumbnail_support(document: &Document) -> (bool, Option) { @@ -204,3 +161,18 @@ pub(crate) fn determine_thumbnail_support(document: &Document) -> (bool, Option< Some("content type not supported for thumbnails".into()), ) } + +fn document_supports_ocr(document: &Document) -> bool { + document + .content_type + .as_deref() + .map(|mime| mime.eq_ignore_ascii_case("application/pdf")) + .unwrap_or_else(|| { + document + .original_name + .rsplit('.') + .next() + .map(|ext| ext.eq_ignore_ascii_case("pdf")) + .unwrap_or(false) + }) +} diff --git a/backend/src/workers/common.rs b/backend/src/workers/common.rs index bccb92e..4428663 100644 --- a/backend/src/workers/common.rs +++ b/backend/src/workers/common.rs @@ -7,10 +7,9 @@ use crate::models::{Document, DocumentAsset, DocumentAssetObject, DocumentVersio use crate::schema::{document_asset_objects, document_assets, document_versions, documents}; use crate::state::AppState; -pub(crate) struct DocumentVersionContext { +pub(crate) struct LoadedDocumentVersion { pub document: Document, pub version: DocumentVersion, - pub tenant_id: Uuid, } pub(crate) fn load_document_version( @@ -18,7 +17,7 @@ pub(crate) fn load_document_version( tenant_id: Uuid, document_id: Uuid, version_id: Uuid, -) -> Result { +) -> Result { let mut conn = state .db_for_tenant(tenant_id) .map_err(|err| format!("{err:?}"))?; @@ -37,14 +36,10 @@ pub(crate) fn load_document_version( .first(&mut conn) .map_err(|err| format!("{err:?}"))?; - Ok(DocumentVersionContext { - document, - version, - tenant_id, - }) + Ok(LoadedDocumentVersion { document, version }) } -pub(crate) struct LoadedAsset { +pub struct LoadedAsset { pub asset: DocumentAsset, pub objects: Vec, } diff --git a/backend/src/workers/index.rs b/backend/src/workers/index.rs index e3a472b..18d235e 100644 --- a/backend/src/workers/index.rs +++ b/backend/src/workers/index.rs @@ -2,24 +2,25 @@ use std::sync::Arc; use std::time::Duration; use async_trait::async_trait; -use diesel::prelude::*; use reqwest::Client; use serde::Deserialize; -use tokio::task; -use tracing::{error, warn}; use uuid::Uuid; use crate::{ documents::search::{build_quickwit_ingest_record, quickwit_ingest}, jobs::JOB_INDEX_DOCUMENT_TEXT, - models::{Document, DocumentVersion}, - schema::{document_asset_objects, document_assets, document_versions, documents}, state::AppState, storage::TenantStorage, }; use super::{ - fetch_version_object, handle_fetch_error, ocr::OCR_TEXT_ASSET_TYPE, JobExecution, JobHandler, + job_execution_from_task_error, + ocr::OCR_TEXT_ASSET_TYPE, + taskflow::{ + document::DocumentVersionTaskContext, BoxedTask, Task, TaskError, TaskExecutor, + TaskPlanner, TaskResult, + }, + JobExecution, JobHandler, }; #[derive(Debug, Deserialize)] @@ -53,161 +54,116 @@ impl JobHandler for IndexDocumentTextJob { Err(err) => { return JobExecution::Failed { error: format!("invalid index payload: {err}"), - } - } - }; - - let quickwit_endpoint = match &state.config.quickwit_endpoint { - Some(endpoint) => endpoint.clone(), - None => { - return JobExecution::Failed { - error: "quickwit endpoint missing".into(), }; } }; - let tenant = match state.tenants.get_by_id(job.tenant_id) { - Ok(tenant) => tenant, - Err(err) => { - warn!(job_id = %job.id, error = ?err, "failed to load tenant for indexing"); - return JobExecution::Retry { - delay: Duration::from_secs(30), - error: format!("failed to load tenant: {err:?}"), - }; - } - }; - - let quickwit_index = match tenant.quickwit_index.clone() { - Some(index) => index, - None => { - return JobExecution::Failed { - error: "tenant quickwit index not configured".into(), - }; - } - }; - - let client = Client::new(); - - let state_clone = state.clone(); - let tenant_id = job.tenant_id; - let context = - match task::spawn_blocking(move || load_context(state_clone, tenant_id, payload)).await - { - Ok(Ok(ctx)) => ctx, - Ok(Err(err)) => { - warn!(job_id = %job.id, error = %err, "index job will retry"); - return JobExecution::Retry { - delay: Duration::from_secs(30), - error: err, - }; - } - Err(join_err) => { - error!(job_id = %job.id, error = %join_err, "index task panicked"); - return JobExecution::Retry { - delay: Duration::from_secs(60), - error: format!("worker panicked: {join_err}"), - }; - } - }; - - if context.text_s3_key.is_none() { - warn!(job_id = %job.id, "missing OCR text asset; failing indexing job"); - return JobExecution::Failed { - error: "missing OCR text asset".into(), - }; - } - - let s3_key = context.text_s3_key.unwrap(); - let bytes = match fetch_version_object( - &context.version, - &storage, - &s3_key, + let mut context = DocumentVersionTaskContext::new( + job.id, + JOB_INDEX_DOCUMENT_TEXT, + job.tenant_id, + payload.document_id, + payload.document_version_id, + false, state.config.worker_max_document_bytes, - ) - .await - { - Ok(bytes) => bytes, - Err(err) => return handle_fetch_error(&job, err, "failed to download ocr text"), - }; - let text = match String::from_utf8(bytes) { - Ok(text) => text, - Err(err) => { - warn!(job_id = %job.id, error = %err, "ocr text not valid UTF-8"); - return JobExecution::Failed { - error: "ocr text not valid UTF-8".into(), - }; - } - }; + state.clone(), + storage, + ); + + let planner = IndexPlanner::new(); + match TaskExecutor::run(&planner, &mut context).await { + Ok(()) => JobExecution::Success, + Err(err) => job_execution_from_task_error(err), + } + } +} + +struct IndexPlanner; + +impl IndexPlanner { + fn new() -> Self { + Self + } +} + +#[async_trait] +impl TaskPlanner for IndexPlanner { + async fn plan( + &self, + _ctx: &mut DocumentVersionTaskContext, + ) -> TaskResult>> { + Ok(vec![Box::new(IndexDocumentTask::new())]) + } +} + +pub struct IndexDocumentTask; + +impl IndexDocumentTask { + pub fn new() -> Self { + Self + } +} + +#[async_trait] +impl Task for IndexDocumentTask { + fn name(&self) -> &'static str { + "index-document-text" + } + + async fn execute(&self, ctx: &mut DocumentVersionTaskContext) -> TaskResult<()> { + let state = ctx.state().clone(); + let quickwit_endpoint = state + .config + .quickwit_endpoint + .clone() + .ok_or_else(|| TaskError::fail("quickwit endpoint missing"))?; + + let tenant = state.tenants.get_by_id(ctx.tenant_id()).map_err(|err| { + TaskError::retry( + Duration::from_secs(30), + format!("failed to load tenant: {err:?}"), + ) + })?; + + let quickwit_index = tenant + .quickwit_index + .clone() + .ok_or_else(|| TaskError::fail("tenant quickwit index not configured"))?; + + let asset = ctx + .asset(OCR_TEXT_ASSET_TYPE) + .await? + .ok_or_else(|| TaskError::fail("missing OCR text asset"))?; + let object = asset + .objects + .first() + .ok_or_else(|| TaskError::fail("missing OCR text object"))?; + + let s3_key = object.s3_key.clone(); + let bytes = ctx.storage().get_object(&s3_key).await.map_err(|err| { + TaskError::retry( + Duration::from_secs(30), + format!("failed to download ocr text: {err}"), + ) + })?; + + let text = String::from_utf8(bytes) + .map_err(|err| TaskError::fail(format!("ocr text not valid UTF-8: {err}")))?; if text.trim().is_empty() { - warn!(job_id = %job.id, "ocr text empty; skipping"); - return JobExecution::Failed { - error: "ocr text empty".into(), - }; + return Err(TaskError::fail("ocr text empty")); } - let record = - build_quickwit_ingest_record(&context.document, &context.version, job.tenant_id, &text); + let document = ctx.document().await?.clone(); + let version = ctx.version().await?.clone(); - match quickwit_ingest(&client, &quickwit_endpoint, &quickwit_index, &[record]).await { - Ok(()) => JobExecution::Success, - Err(err) => { - warn!(job_id = %job.id, error = %err, "quickwit ingest failed"); - JobExecution::Retry { - delay: Duration::from_secs(30), - error: err.to_string(), - } - } - } + let record = build_quickwit_ingest_record(&document, &version, ctx.tenant_id(), &text); + let client = Client::new(); + + quickwit_ingest(&client, &quickwit_endpoint, &quickwit_index, &[record]) + .await + .map_err(|err| TaskError::retry(Duration::from_secs(30), err.to_string()))?; + + Ok(()) } } - -struct IndexContext { - document: Document, - version: DocumentVersion, - text_s3_key: Option, -} - -fn load_context( - state: Arc, - tenant_id: Uuid, - payload: IndexPayload, -) -> Result { - let mut conn = state - .db_for_tenant(tenant_id) - .map_err(|err| format!("{err:?}"))?; - - let version: DocumentVersion = document_versions::table - .find(payload.document_version_id) - .first(&mut conn) - .map_err(|err| format!("{err:?}"))?; - - if version.document_id != payload.document_id { - return Err("document/version mismatch".into()); - } - - let document: Document = documents::table - .find(payload.document_id) - .first(&mut conn) - .map_err(|err| format!("{err:?}"))?; - - let text_s3_key: Option = document_asset_objects::table - .inner_join( - document_assets::table.on(document_asset_objects::asset_id.eq(document_assets::id)), - ) - .filter(document_assets::document_version_id.eq(payload.document_version_id)) - .filter(document_assets::asset_type.eq(OCR_TEXT_ASSET_TYPE)) - .filter(document_assets::tenant_id.eq(tenant_id)) - .filter(document_asset_objects::ordinal.eq(1)) - .filter(document_asset_objects::tenant_id.eq(tenant_id)) - .select(document_asset_objects::s3_key) - .first(&mut conn) - .optional() - .map_err(|err| format!("{err:?}"))?; - - Ok(IndexContext { - document, - version, - text_s3_key, - }) -} diff --git a/backend/src/workers/mod.rs b/backend/src/workers/mod.rs index dfaf438..333046e 100644 --- a/backend/src/workers/mod.rs +++ b/backend/src/workers/mod.rs @@ -11,12 +11,14 @@ use crate::{ state::AppState, storage::TenantStorage, }; +use taskflow::TaskError; pub mod analyze; pub mod common; pub mod index; pub mod ocr; pub mod purge; +pub mod taskflow; pub mod tenants; pub mod thumbnails; @@ -156,6 +158,13 @@ pub fn default_handlers() -> Vec> { ] } +pub(crate) fn job_execution_from_task_error(error: TaskError) -> JobExecution { + match error { + TaskError::Fail { error } => JobExecution::Failed { error }, + TaskError::Retry { delay, error } => JobExecution::Retry { delay, error }, + } +} + pub(crate) fn check_worker_document_limit( size_bytes: i64, limit_bytes: u64, @@ -187,35 +196,3 @@ pub(crate) async fn fetch_version_object( .await .map_err(FetchVersionError::Storage) } - -pub(crate) fn handle_fetch_error( - job: &crate::models::Job, - err: FetchVersionError, - message: &str, -) -> JobExecution { - match err { - FetchVersionError::TooLarge { size, limit } => { - warn!( - job_id = %job.id, - size_bytes = size, - limit_bytes = limit, - "document exceeds worker size limit" - ); - JobExecution::Failed { - error: format!("document size {size} bytes exceeds worker limit of {limit} bytes"), - } - } - FetchVersionError::Storage(err) => { - warn!( - job_id = %job.id, - error = %err, - context = message, - "failed to fetch object for worker" - ); - JobExecution::Retry { - delay: Duration::from_secs(30), - error: err.to_string(), - } - } - } -} diff --git a/backend/src/workers/ocr.rs b/backend/src/workers/ocr.rs index 5aa493f..3bf3ee9 100644 --- a/backend/src/workers/ocr.rs +++ b/backend/src/workers/ocr.rs @@ -14,13 +14,13 @@ use serde::Deserialize; use serde_json::json; use tempfile::NamedTempFile; use tokio::task; -use tracing::{error, info, warn}; +use tracing::{info, warn}; use uuid::Uuid; use crate::{ documents::asset::delete_asset, error::AppResult, - jobs::{enqueue_job, JOB_GENERATE_OCR_TEXT, JOB_INDEX_DOCUMENT_TEXT}, + jobs::JOB_GENERATE_OCR_TEXT, models::{ Document, DocumentAsset, DocumentAssetObject, DocumentVersion, NewDocumentAsset, NewDocumentAssetObject, @@ -32,8 +32,13 @@ use crate::{ }; use super::{ - common::{load_document_version, load_version_assets}, - fetch_version_object, handle_fetch_error, JobExecution, JobHandler, + index::IndexDocumentTask, + job_execution_from_task_error, + taskflow::{ + document::DocumentVersionTaskContext, BoxedTask, Task, TaskContext, TaskError, + TaskExecutor, TaskPlanner, TaskResult, + }, + JobExecution, JobHandler, }; pub const OCR_TEXT_ASSET_TYPE: &str = "ocr-text"; @@ -72,104 +77,102 @@ impl JobHandler for GenerateOcrTextJob { Err(err) => { return JobExecution::Failed { error: format!("invalid OCR payload: {err}"), - } + }; } }; - let state_clone = state.clone(); - let payload_clone = payload.clone(); - let tenant_id = job.tenant_id; - let context = match task::spawn_blocking(move || { - load_ocr_context(state_clone, tenant_id, payload_clone) - }) - .await - { - Ok(Ok(ctx)) => ctx, - Ok(Err(err)) => { - warn!(job_id = %job.id, error = %err, "ocr job will retry"); - return JobExecution::Retry { - delay: Duration::from_secs(30), - error: err, - }; - } - Err(join_err) => { - error!(job_id = %job.id, error = %join_err, "ocr task panicked"); - return JobExecution::Retry { - delay: Duration::from_secs(60), - error: format!("worker panicked: {join_err}"), - }; - } - }; + let mut context = DocumentVersionTaskContext::new( + job.id, + JOB_GENERATE_OCR_TEXT, + job.tenant_id, + payload.document_id, + payload.document_version_id, + payload.force, + state.config.worker_max_document_bytes, + state.clone(), + storage, + ); + + let planner = OcrPlanner::new(payload.force, state.clone()); + match TaskExecutor::run(&planner, &mut context).await { + Ok(()) => JobExecution::Success, + Err(err) => job_execution_from_task_error(err), + } + } +} + +struct OcrPlanner { + force: bool, + state: Arc, +} + +impl OcrPlanner { + fn new(force: bool, state: Arc) -> Self { + Self { force, state } + } +} + +#[async_trait] +impl TaskPlanner for OcrPlanner { + async fn plan( + &self, + _ctx: &mut DocumentVersionTaskContext, + ) -> TaskResult>> { + Ok(vec![ + Box::new(GenerateOcrTask::new(self.force, self.state.clone())), + Box::new(IndexDocumentTask::new()), + ]) + } +} + +pub struct GenerateOcrTask { + force: bool, + state: Arc, +} + +impl GenerateOcrTask { + pub fn new(force: bool, state: Arc) -> Self { + Self { force, state } + } +} + +#[async_trait] +impl Task for GenerateOcrTask { + fn name(&self) -> &'static str { + "generate-ocr-text" + } + + async fn execute(&self, ctx: &mut DocumentVersionTaskContext) -> TaskResult<()> { + let context = build_ocr_context(ctx, self.force).await?; if context.skip { - info!(job_id = %job.id, "ocr already present; skipping"); - return JobExecution::Success; + info!(job_id = %ctx.job_id(), "ocr already present; skipping"); + return Ok(()); } - let bytes = match fetch_version_object( - &context.version, - &storage, - &context.version.s3_key, - state.config.worker_max_document_bytes, - ) - .await - { - Ok(bytes) => bytes, - Err(err) => return handle_fetch_error(&job, err, "failed to fetch document for ocr"), - }; - - let doc_meta = PdfDocumentMeta { + let bytes = ctx.buffered_object().await?.to_vec(); + let meta = PdfDocumentMeta { content_type: context.document.content_type.clone(), original_name: context.document.original_name.clone(), }; - let generation = - match task::spawn_blocking(move || generate_ocr_text(&doc_meta, &bytes)).await { - Ok(result) => result, - Err(join_err) => { - error!(job_id = %job.id, error = %join_err, "ocr text task panicked"); - return JobExecution::Retry { - delay: Duration::from_secs(60), - error: format!("worker panicked: {join_err}"), - }; - } - }; + let generation = task::spawn_blocking(move || generate_ocr_text(&meta, &bytes)) + .await + .map_err(|err| { + TaskError::retry( + Duration::from_secs(60), + format!("ocr text task panicked: {err}"), + ) + })?; let Some(generation) = generation else { - warn!(job_id = %job.id, "no text extracted from document; failing job"); - return JobExecution::Failed { - error: "no text extracted and OCR unavailable".into(), - }; + warn!(job_id = %ctx.job_id(), "no text extracted from document; failing job"); + return Err(TaskError::fail("no text extracted and OCR unavailable")); }; - if let Some(existing_asset) = &context.existing_asset { - for object in &context.existing_objects { - if let Err(err) = storage.delete_object(&object.s3_key).await { - warn!(job_id = %job.id, error = %err, s3_key = %object.s3_key, "failed to delete existing ocr asset object"); - } - } - - let tenant_id = context.document.tenant_id; - let asset_id = existing_asset.id; - let state_clone = state.clone(); - match task::spawn_blocking(move || -> AppResult<()> { - let mut conn = state_clone.db_for_tenant(tenant_id)?; - delete_asset(&mut conn, tenant_id, asset_id) - }) - .await - { - Ok(Ok(())) => {} - Ok(Err(err)) => { - warn!(job_id = %job.id, error = ?err, asset_id = %asset_id, "failed to remove ocr asset metadata after deletion"); - } - Err(join_err) => { - warn!(job_id = %job.id, error = %join_err, asset_id = %asset_id, "failed to remove ocr asset metadata: task panicked"); - } - } - } + remove_existing_ocr_asset(ctx, &context).await; let asset_id = Uuid::new_v4(); - let s3_key = document_asset_object_prefix( context.document.id, context.version.version_number, @@ -177,7 +180,7 @@ impl JobHandler for GenerateOcrTextJob { asset_id, ); - if let Err(err) = storage + ctx.storage() .put_object( &s3_key, generation.text.into_bytes(), @@ -185,47 +188,25 @@ impl JobHandler for GenerateOcrTextJob { None, ) .await - { - warn!(job_id = %job.id, error = %err, "failed to upload ocr text"); - return JobExecution::Retry { - delay: Duration::from_secs(30), - error: err.to_string(), - }; - } + .map_err(|err| TaskError::retry(Duration::from_secs(30), err.to_string()))?; - let state_clone = state.clone(); - match task::spawn_blocking(move || { - persist_ocr_metadata(state_clone, &context, asset_id, &s3_key, generation.source) + let state = self.state.clone(); + task::spawn_blocking(move || { + persist_ocr_metadata(state, &context, asset_id, &s3_key, generation.source) }) .await - { - Ok(Ok(())) => { - if let Err(err) = enqueue_index_job(&state, job.tenant_id, &payload) { - warn!(job_id = %job.id, error = %err, "failed to enqueue index job"); - } - JobExecution::Success - } - Ok(Err(err)) => { - warn!(job_id = %job.id, error = %err, "failed to persist ocr metadata"); - JobExecution::Retry { - delay: Duration::from_secs(30), - error: err, - } - } - Err(join_err) => { - error!(job_id = %job.id, error = %join_err, "ocr metadata task panicked"); - JobExecution::Retry { - delay: Duration::from_secs(30), - error: format!("metadata update panic: {join_err}"), - } - } - } - } -} + .map_err(|err| { + TaskError::retry( + Duration::from_secs(60), + format!("ocr metadata task panicked: {err}"), + ) + })? + .map_err(|err| TaskError::retry(Duration::from_secs(30), err))?; -struct PdfDocumentMeta { - content_type: Option, - original_name: String, + ctx.invalidate_asset_cache(); + + Ok(()) + } } struct OcrContext { @@ -236,61 +217,163 @@ struct OcrContext { skip: bool, } -struct OcrGeneration { - text: String, - source: &'static str, -} +async fn build_ocr_context( + ctx: &mut DocumentVersionTaskContext, + force: bool, +) -> TaskResult { + let document = ctx.document().await?.clone(); + let version = ctx.version().await?.clone(); + let asset = ctx.asset(OCR_TEXT_ASSET_TYPE).await?; -fn load_ocr_context( - state: Arc, - tenant_id: Uuid, - payload: OcrPayload, -) -> Result { - let base = load_document_version( - state.as_ref(), - tenant_id, - payload.document_id, - payload.document_version_id, - )?; - - let mut conn = state - .db_for_tenant(base.tenant_id) - .map_err(|err| format!("{err:?}"))?; - - let mut assets = load_version_assets( - &mut conn, - base.tenant_id, - base.version.id, - &[OCR_TEXT_ASSET_TYPE], - )?; - - let (existing_asset, existing_objects) = assets - .remove(OCR_TEXT_ASSET_TYPE) - .map(|entry| (Some(entry.asset), entry.objects)) + let (existing_asset, existing_objects) = asset + .map(|asset| (Some(asset.asset.clone()), asset.objects.clone())) .unwrap_or((None, Vec::new())); - let is_pdf = document_is_pdf(&base.document); - if !is_pdf { + if !document_is_pdf(&document) { return Ok(OcrContext { - document: base.document, - version: base.version, + document, + version, existing_asset, existing_objects, skip: true, }); } - let skip = existing_asset.is_some() && !payload.force; + let skip = existing_asset.is_some() && !force; Ok(OcrContext { - document: base.document, - version: base.version, + document, + version, existing_asset, existing_objects, skip, }) } +async fn remove_existing_ocr_asset(ctx: &DocumentVersionTaskContext, context: &OcrContext) { + if let Some(existing_asset) = &context.existing_asset { + for object in &context.existing_objects { + if let Err(err) = ctx.storage().delete_object(&object.s3_key).await { + warn!( + job_id = %ctx.job_id(), + error = %err, + s3_key = %object.s3_key, + "failed to delete existing ocr asset object" + ); + } + } + + let tenant_id = context.document.tenant_id; + let asset_id = existing_asset.id; + let state = ctx.state().clone(); + match task::spawn_blocking(move || -> AppResult<()> { + let mut conn = state.db_for_tenant(tenant_id)?; + delete_asset(&mut conn, tenant_id, asset_id) + }) + .await + { + Ok(Ok(())) => {} + Ok(Err(err)) => { + warn!( + job_id = %ctx.job_id(), + error = ?err, + asset_id = %asset_id, + "failed to remove ocr asset metadata after deletion" + ); + } + Err(join_err) => { + warn!( + job_id = %ctx.job_id(), + error = %join_err, + asset_id = %asset_id, + "failed to remove ocr asset metadata: task panicked" + ); + } + } + } +} + +fn persist_ocr_metadata( + state: Arc, + context: &OcrContext, + asset_id: Uuid, + s3_key: &str, + source: OcrSource, +) -> Result<(), String> { + let tenant_id = context.document.tenant_id; + let mut conn = state + .db_for_tenant(tenant_id) + .map_err(|err| format!("{err:?}"))?; + + let document_version_id = context.version.id; + let existing_asset = context.existing_asset.as_ref().map(|asset| asset.id); + + if let Some(existing_asset) = existing_asset { + diesel::delete( + document_assets::table + .filter(document_assets::id.eq(existing_asset)) + .filter(document_assets::tenant_id.eq(tenant_id)), + ) + .execute(&mut conn) + .map_err(|err| format!("{err:?}"))?; + } + + let metadata = json!({ + "source": source.to_string(), + "generated_at": Utc::now().to_rfc3339(), + }); + + let new_asset = NewDocumentAsset { + id: asset_id, + document_version_id, + asset_type: OCR_TEXT_ASSET_TYPE.to_string(), + mime_type: "text/plain".to_string(), + metadata, + cardinality: Some(1), + tenant_id, + }; + + diesel::insert_into(document_assets::table) + .values(&new_asset) + .on_conflict(( + document_assets::document_version_id, + document_assets::asset_type, + )) + .do_update() + .set(( + document_assets::mime_type.eq(excluded(document_assets::mime_type)), + document_assets::metadata.eq(excluded(document_assets::metadata)), + document_assets::cardinality.eq(excluded(document_assets::cardinality)), + document_assets::id.eq(excluded(document_assets::id)), + )) + .execute(&mut conn) + .map_err(|err| format!("{err:?}"))?; + + diesel::delete( + document_asset_objects::table + .filter(document_asset_objects::asset_id.eq(asset_id)) + .filter(document_asset_objects::tenant_id.eq(tenant_id)), + ) + .execute(&mut conn) + .map_err(|err| format!("{err:?}"))?; + + let object = NewDocumentAssetObject { + id: Uuid::new_v4(), + asset_id, + ordinal: 1, + s3_key: s3_key.to_string(), + metadata: json!({}), + tenant_id, + }; + + diesel::insert_into(document_asset_objects::table) + .values(&object) + .execute(&mut conn) + .map_err(|err| format!("{err:?}"))?; + + Ok(()) +} + fn generate_ocr_text(meta: &PdfDocumentMeta, bytes: &[u8]) -> Option { if !document_meta_is_pdf(meta) { return None; @@ -300,7 +383,7 @@ fn generate_ocr_text(meta: &PdfDocumentMeta, bytes: &[u8]) -> Option= MIN_TEXT_LENGTH { return Some(OcrGeneration { text, - source: "pdf-text", + source: OcrSource::PdfText, }); } } @@ -308,20 +391,60 @@ fn generate_ocr_text(meta: &PdfDocumentMeta, bytes: &[u8]) -> Option Some(OcrGeneration { text, - source: "ocr", + source: OcrSource::Ocr, }), Ok(None) => None, Err(OcrError::BinaryMissing) => { - warn!("ocrmypdf not installed; cannot perform OCR"); + warn!("ocrmypdf binary not found; OCR unavailable"); None } Err(err) => { - warn!(error = ?err, "ocr command failed"); + warn!(error = %err, "ocr command failed"); None } } } +struct PdfDocumentMeta { + content_type: Option, + original_name: String, +} + +struct OcrGeneration { + text: String, + source: OcrSource, +} + +#[derive(Clone, Copy)] +enum OcrSource { + PdfText, + Ocr, +} + +impl fmt::Display for OcrSource { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + OcrSource::PdfText => write!(f, "pdf-text"), + OcrSource::Ocr => write!(f, "ocr"), + } + } +} + +#[derive(Debug)] +enum OcrError { + BinaryMissing, + Failed(String), +} + +impl fmt::Display for OcrError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + OcrError::BinaryMissing => write!(f, "ocrmypdf binary not found"), + OcrError::Failed(msg) => write!(f, "ocr failed: {msg}"), + } + } +} + fn extract_pdf_text(bytes: &[u8]) -> Result { let pdfium = Pdfium::default(); let document = pdfium @@ -345,21 +468,6 @@ fn extract_pdf_text(bytes: &[u8]) -> Result { Ok(combined) } -#[derive(Debug)] -enum OcrError { - BinaryMissing, - Failed(String), -} - -impl fmt::Display for OcrError { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - OcrError::BinaryMissing => write!(f, "ocrmypdf binary not found"), - OcrError::Failed(msg) => write!(f, "ocr failed: {msg}"), - } - } -} - fn run_ocr(bytes: &[u8]) -> Result, OcrError> { let mut input = NamedTempFile::new().map_err(|err| OcrError::Failed(err.to_string()))?; input @@ -408,118 +516,6 @@ fn run_ocr(bytes: &[u8]) -> Result, OcrError> { } } -fn persist_ocr_metadata( - state: Arc, - context: &OcrContext, - asset_id: Uuid, - s3_key: &str, - source: &'static str, -) -> Result<(), String> { - let tenant_id = context.document.tenant_id; - let mut conn = state - .db_for_tenant(tenant_id) - .map_err(|err| format!("{err:?}"))?; - - if let Some(existing_asset) = &context.existing_asset { - diesel::delete(document_assets::table.filter(document_assets::id.eq(existing_asset.id))) - .execute(&mut conn) - .map_err(|err| format!("{err:?}"))?; - } - - let new_asset = NewDocumentAsset { - id: asset_id, - document_version_id: context.version.id, - asset_type: OCR_TEXT_ASSET_TYPE.to_string(), - mime_type: "text/plain".to_string(), - metadata: json!({ - "generated_at": Utc::now().to_rfc3339(), - "source": source, - }), - cardinality: Some(1), - tenant_id, - }; - - diesel::insert_into(document_assets::table) - .values(&new_asset) - .on_conflict(( - document_assets::document_version_id, - document_assets::asset_type, - )) - .do_update() - .set(( - document_assets::mime_type.eq(excluded(document_assets::mime_type)), - document_assets::metadata.eq(excluded(document_assets::metadata)), - document_assets::cardinality.eq(excluded(document_assets::cardinality)), - )) - .execute(&mut conn) - .map_err(|err| format!("{err:?}"))?; - - let existing_object_id: Option = document_asset_objects::table - .filter(document_asset_objects::asset_id.eq(asset_id)) - .filter(document_asset_objects::ordinal.eq(1)) - .filter(document_asset_objects::tenant_id.eq(tenant_id)) - .select(document_asset_objects::id) - .first(&mut conn) - .optional() - .map_err(|err| format!("{err:?}"))?; - - let object_id = existing_object_id.unwrap_or_else(Uuid::new_v4); - - let new_object = NewDocumentAssetObject { - id: object_id, - asset_id, - ordinal: 1, - s3_key: s3_key.to_string(), - metadata: json!({}), - tenant_id, - }; - - diesel::insert_into(document_asset_objects::table) - .values(&new_object) - .on_conflict(( - document_asset_objects::asset_id, - document_asset_objects::ordinal, - )) - .do_update() - .set(( - document_asset_objects::s3_key.eq(excluded(document_asset_objects::s3_key)), - document_asset_objects::metadata.eq(excluded(document_asset_objects::metadata)), - )) - .execute(&mut conn) - .map_err(|err| format!("{err:?}"))?; - - Ok(()) -} - -fn enqueue_index_job( - state: &AppState, - tenant_id: Uuid, - payload: &OcrPayload, -) -> Result<(), String> { - let mut conn = state - .db_for_tenant(tenant_id) - .map_err(|err| format!("{err:?}"))?; - enqueue_job( - &mut conn, - tenant_id, - JOB_INDEX_DOCUMENT_TEXT, - json!({ - "document_id": payload.document_id, - "document_version_id": payload.document_version_id, - }), - None, - ) - .map(|_| ()) - .map_err(|err| err.to_string()) -} - -pub fn document_is_pdf(document: &Document) -> bool { - document_meta_is_pdf(&PdfDocumentMeta { - content_type: document.content_type.clone(), - original_name: document.original_name.clone(), - }) -} - fn document_meta_is_pdf(meta: &PdfDocumentMeta) -> bool { if let Some(content_type) = &meta.content_type { if content_type.eq_ignore_ascii_case("application/pdf") { @@ -533,3 +529,18 @@ fn document_meta_is_pdf(meta: &PdfDocumentMeta) -> bool { .map(|ext| ext.eq_ignore_ascii_case("pdf")) .unwrap_or(false) } + +fn document_is_pdf(document: &Document) -> bool { + document + .content_type + .as_deref() + .map(|mime| mime.eq_ignore_ascii_case("application/pdf")) + .unwrap_or_else(|| { + document + .original_name + .rsplit('.') + .next() + .map(|ext| ext.eq_ignore_ascii_case("pdf")) + .unwrap_or(false) + }) +} diff --git a/backend/src/workers/purge.rs b/backend/src/workers/purge.rs index 0bce8b7..3244635 100644 --- a/backend/src/workers/purge.rs +++ b/backend/src/workers/purge.rs @@ -6,7 +6,6 @@ use async_trait::async_trait; use diesel::prelude::*; use diesel::result::Error as DieselError; use serde::Deserialize; -use tracing::{error, warn}; use uuid::Uuid; use crate::jobs::JOB_PURGE_DOCUMENT; @@ -15,7 +14,11 @@ use crate::schema::{document_asset_objects, document_assets, document_versions}; use crate::state::AppState; use crate::storage::TenantStorage; -use super::{JobExecution, JobHandler}; +use super::{ + job_execution_from_task_error, + taskflow::{BoxedTask, Task, TaskContext, TaskError, TaskExecutor, TaskPlanner, TaskResult}, + JobExecution, JobHandler, +}; #[derive(Debug, Deserialize)] struct PurgeDocumentPayload { @@ -58,70 +61,120 @@ impl JobHandler for PurgeDocumentJob { } }; - let tenant_id = job.tenant_id; - let document_id = payload.document_id; - let state_for_prepare = state.clone(); + let mut context = PurgeTaskContext::new( + job.id, + JOB_PURGE_DOCUMENT, + job.tenant_id, + payload.document_id, + state.clone(), + storage, + ); + + let planner = PurgePlanner; + match TaskExecutor::run(&planner, &mut context).await { + Ok(()) => JobExecution::Success, + Err(err) => job_execution_from_task_error(err), + } + } +} + +struct PurgeTaskContext { + job_id: Uuid, + job_type: &'static str, + tenant_id: Uuid, + document_id: Uuid, + state: Arc, + storage: TenantStorage, +} + +impl PurgeTaskContext { + fn new( + job_id: Uuid, + job_type: &'static str, + tenant_id: Uuid, + document_id: Uuid, + state: Arc, + storage: TenantStorage, + ) -> Self { + Self { + job_id, + job_type, + tenant_id, + document_id, + state, + storage, + } + } +} + +impl TaskContext for PurgeTaskContext { + fn job_id(&self) -> Uuid { + self.job_id + } + + fn job_type(&self) -> &'static str { + self.job_type + } +} + +struct PurgePlanner; + +#[async_trait] +impl TaskPlanner for PurgePlanner { + async fn plan( + &self, + _ctx: &mut PurgeTaskContext, + ) -> TaskResult>> { + Ok(vec![Box::new(PurgeTask)]) + } +} + +struct PurgeTask; + +#[async_trait] +impl Task for PurgeTask { + fn name(&self) -> &'static str { + "purge-document" + } + + async fn execute(&self, ctx: &mut PurgeTaskContext) -> TaskResult<()> { + let tenant_id = ctx.tenant_id; + let document_id = ctx.document_id; + let state = ctx.state.clone(); let preparation = tokio::task::spawn_blocking(move || { - prepare_purge_context(state_for_prepare, tenant_id, document_id) + prepare_purge_context(state, tenant_id, document_id) }) - .await; + .await + .map_err(|err| { + TaskError::retry( + Duration::from_secs(60), + format!("purge preparation panicked: {err}"), + ) + })?; - let context = match preparation { - Ok(Ok(Some(ctx))) => ctx, - Ok(Ok(None)) => { - // Document already gone or restored; nothing to do. - return JobExecution::Success; - } - Ok(Err(err)) => { - warn!(job_id = %job.id, error = %err, "purge preparation failed"); - return JobExecution::Retry { - delay: Duration::from_secs(30), - error: err, - }; - } - Err(join_err) => { - error!(job_id = %job.id, error = %join_err, "purge preparation task panicked"); - return JobExecution::Retry { - delay: Duration::from_secs(60), - error: format!("purge preparation panicked: {join_err}"), - }; - } + let Some(context) = + preparation.map_err(|err| TaskError::retry(Duration::from_secs(30), err))? + else { + return Ok(()); }; - if let Err(err) = delete_storage_objects(&storage, &context).await { - warn!(job_id = %job.id, error = %err, "failed to delete storage objects for purge"); - return JobExecution::Retry { - delay: Duration::from_secs(30), - error: err, - }; - } + delete_storage_objects(&ctx.storage, &context) + .await + .map_err(|err| TaskError::retry(Duration::from_secs(30), err))?; - let PurgeContext { document_id, .. } = context; - let state_for_finalize = state.clone(); + let state = ctx.state.clone(); + tokio::task::spawn_blocking(move || finalize_purge(state, tenant_id, context.document_id)) + .await + .map_err(|err| { + TaskError::retry( + Duration::from_secs(60), + format!("purge finalize panicked: {err}"), + ) + })? + .map_err(|err| TaskError::retry(Duration::from_secs(30), err))?; - let finalize = tokio::task::spawn_blocking(move || { - finalize_purge(state_for_finalize, tenant_id, document_id) - }) - .await; - - match finalize { - Ok(Ok(())) => JobExecution::Success, - Ok(Err(err)) => { - warn!(job_id = %job.id, error = %err, "failed to finalize purge"); - JobExecution::Retry { - delay: Duration::from_secs(30), - error: err, - } - } - Err(join_err) => { - error!(job_id = %job.id, error = %join_err, "purge finalize task panicked"); - JobExecution::Retry { - delay: Duration::from_secs(60), - error: format!("purge finalize panicked: {join_err}"), - } - } - } + Ok(()) } } diff --git a/backend/src/workers/taskflow/document.rs b/backend/src/workers/taskflow/document.rs new file mode 100644 index 0000000..772b62f --- /dev/null +++ b/backend/src/workers/taskflow/document.rs @@ -0,0 +1,197 @@ +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Duration; + +use tokio::task; +use uuid::Uuid; + +use crate::models::{Document, DocumentVersion}; +use crate::state::AppState; +use crate::storage::TenantStorage; +use crate::workers::common::{load_document_version, load_version_assets, LoadedAsset}; +use crate::workers::{fetch_version_object, FetchVersionError}; + +use super::{TaskContext, TaskError, TaskResult}; + +const BLOCKING_RETRY_DELAY: Duration = Duration::from_secs(60); +const DEFAULT_RETRY_DELAY: Duration = Duration::from_secs(30); + +pub struct DocumentVersionTaskContext { + job_id: Uuid, + job_type: &'static str, + tenant_id: Uuid, + document_id: Uuid, + document_version_id: Uuid, + force: bool, + max_document_bytes: u64, + state: Arc, + storage: TenantStorage, + document: Option, + version: Option, + assets: Option>, // keyed by asset_type + object_bytes: Option>, +} + +impl DocumentVersionTaskContext { + #[allow(clippy::too_many_arguments)] + pub fn new( + job_id: Uuid, + job_type: &'static str, + tenant_id: Uuid, + document_id: Uuid, + document_version_id: Uuid, + force: bool, + max_document_bytes: u64, + state: Arc, + storage: TenantStorage, + ) -> Self { + Self { + job_id, + job_type, + tenant_id, + document_id, + document_version_id, + force, + max_document_bytes, + state, + storage, + document: None, + version: None, + assets: None, + object_bytes: None, + } + } + + pub fn tenant_id(&self) -> Uuid { + self.tenant_id + } + + pub fn document_id(&self) -> Uuid { + self.document_id + } + + pub fn version_id(&self) -> Uuid { + self.document_version_id + } + + pub fn force(&self) -> bool { + self.force + } + + pub fn storage(&self) -> &TenantStorage { + &self.storage + } + + pub fn invalidate_asset_cache(&mut self) { + self.assets = None; + } + + pub fn state(&self) -> &Arc { + &self.state + } + + pub fn max_document_bytes(&self) -> u64 { + self.max_document_bytes + } + + pub async fn document(&mut self) -> TaskResult<&Document> { + self.ensure_document_loaded().await?; + Ok(self.document.as_ref().expect("document hydrated")) + } + + pub async fn version(&mut self) -> TaskResult<&DocumentVersion> { + self.ensure_document_loaded().await?; + Ok(self.version.as_ref().expect("version hydrated")) + } + + pub async fn assets(&mut self) -> TaskResult<&HashMap> { + if self.assets.is_none() { + let tenant_id = self.tenant_id; + let version_id = self.document_version_id; + let state = self.state.clone(); + let result = task::spawn_blocking(move || { + let mut conn = state + .db_for_tenant(tenant_id) + .map_err(|err| format!("failed to scope tenant connection: {err:?}"))?; + load_version_assets(&mut conn, tenant_id, version_id, &[]) + }) + .await + .map_err(|err| { + TaskError::retry( + BLOCKING_RETRY_DELAY, + format!("asset load task panicked: {err}"), + ) + })? + .map_err(|err| TaskError::retry(DEFAULT_RETRY_DELAY, err))?; + self.assets = Some(result); + } + Ok(self.assets.as_ref().expect("asset map hydrated")) + } + + pub async fn asset(&mut self, asset_type: &str) -> TaskResult> { + let assets = self.assets().await?; + Ok(assets.get(asset_type)) + } + + pub async fn buffered_object(&mut self) -> TaskResult<&[u8]> { + if self.object_bytes.is_some() { + return Ok(self.object_bytes.as_deref().expect("bytes present")); + } + let version = self.version().await?.clone(); + let bytes = fetch_version_object( + &version, + &self.storage, + &version.s3_key, + self.max_document_bytes, + ) + .await + .map_err(|err| match err { + FetchVersionError::TooLarge { size, limit } => TaskError::fail(format!( + "document size {size} bytes exceeds worker limit of {limit} bytes" + )), + FetchVersionError::Storage(err) => TaskError::retry( + DEFAULT_RETRY_DELAY, + format!("failed to fetch object: {err}"), + ), + })?; + self.object_bytes = Some(bytes); + Ok(self.object_bytes.as_deref().expect("bytes hydrated")) + } + + async fn ensure_document_loaded(&mut self) -> TaskResult<()> { + if self.document.is_some() && self.version.is_some() { + return Ok(()); + } + + let tenant_id = self.tenant_id; + let document_id = self.document_id; + let version_id = self.document_version_id; + let state = self.state.clone(); + + let loaded = task::spawn_blocking(move || { + load_document_version(state.as_ref(), tenant_id, document_id, version_id) + }) + .await + .map_err(|err| { + TaskError::retry( + BLOCKING_RETRY_DELAY, + format!("document load task panicked: {err}"), + ) + })? + .map_err(|err| TaskError::fail(format!("failed to load document context: {err}")))?; + + self.document = Some(loaded.document); + self.version = Some(loaded.version); + Ok(()) + } +} + +impl TaskContext for DocumentVersionTaskContext { + fn job_id(&self) -> Uuid { + self.job_id + } + + fn job_type(&self) -> &'static str { + self.job_type + } +} diff --git a/backend/src/workers/taskflow/mod.rs b/backend/src/workers/taskflow/mod.rs new file mode 100644 index 0000000..9fc3831 --- /dev/null +++ b/backend/src/workers/taskflow/mod.rs @@ -0,0 +1,85 @@ +use std::time::Duration; + +use async_trait::async_trait; +use thiserror::Error; +use tracing::info; +use uuid::Uuid; + +pub mod document; + +pub type TaskResult = Result; + +#[derive(Debug, Error)] +pub enum TaskError { + #[error("{error}")] + Fail { error: String }, + #[error("{error}")] + Retry { delay: Duration, error: String }, +} + +impl TaskError { + pub fn fail(error: impl Into) -> Self { + Self::Fail { + error: error.into(), + } + } + + pub fn retry(delay: Duration, error: impl Into) -> Self { + Self::Retry { + delay, + error: error.into(), + } + } +} + +pub trait TaskContext: Send + Sync { + fn job_id(&self) -> Uuid; + fn job_type(&self) -> &'static str; +} + +#[async_trait] +pub trait Task: Send + Sync +where + Ctx: TaskContext, +{ + fn name(&self) -> &'static str; + async fn execute(&self, ctx: &mut Ctx) -> TaskResult<()>; +} + +pub type BoxedTask = Box + Send + Sync>; + +#[async_trait] +pub trait TaskPlanner: Send + Sync +where + Ctx: TaskContext, +{ + async fn plan(&self, ctx: &mut Ctx) -> TaskResult>>; +} + +pub struct TaskExecutor; + +impl TaskExecutor { + pub async fn run(planner: &P, ctx: &mut C) -> TaskResult<()> + where + P: TaskPlanner, + C: TaskContext, + { + let tasks = planner.plan(ctx).await?; + for task in tasks { + info!( + job_id = %ctx.job_id(), + job_type = ctx.job_type(), + task = task.name(), + "starting job task" + ); + task.execute(ctx).await?; + info!( + job_id = %ctx.job_id(), + job_type = ctx.job_type(), + task = task.name(), + "finished job task" + ); + } + Ok(()) + } +} diff --git a/backend/src/workers/tenants.rs b/backend/src/workers/tenants.rs index d8b542b..6d8a8c7 100644 --- a/backend/src/workers/tenants.rs +++ b/backend/src/workers/tenants.rs @@ -1,4 +1,5 @@ use std::sync::Arc; +use std::time::Duration; use async_trait::async_trait; use chrono::Utc; @@ -18,7 +19,11 @@ use crate::models::{NewUserMembership, TenantStatus}; use crate::schema::{tenants, user_memberships}; use crate::state::AppState; use crate::tenants::TenantRepository; -use crate::workers::{JobExecution, JobHandler}; +use crate::workers::{ + job_execution_from_task_error, + taskflow::{BoxedTask, Task, TaskContext, TaskError, TaskExecutor, TaskPlanner, TaskResult}, + JobExecution, JobHandler, +}; pub struct ProvisionTenantJob; @@ -40,77 +45,120 @@ impl JobHandler for ProvisionTenantJob { job: crate::models::Job, _storage: crate::storage::TenantStorage, ) -> JobExecution { - let mut conn = match state.db_unscoped() { - Ok(conn) => conn, - Err(err) => { - warn!(job_id = %job.id, error = ?err, "failed to get connection for tenant provisioning"); - return JobExecution::Retry { - delay: std::time::Duration::from_secs(30), - error: "database connection unavailable".into(), - }; - } - }; + let members = ProvisionPayload::from_job(&job).unwrap_or_default(); + let mut context = ProvisionContext::new( + job.id, + JOB_PROVISION_TENANT, + job.tenant_id, + state.clone(), + members, + ); - let tenant = match TenantRepository::get_by_id(&mut conn, job.tenant_id) { - Ok(tenant) => tenant, - Err(err) => { - warn!(job_id = %job.id, error = ?err, "tenant not found for provisioning"); - return JobExecution::Failed { - error: "tenant not found".into(), - }; - } - }; + let planner = ProvisionPlanner; + match TaskExecutor::run(&planner, &mut context).await { + Ok(()) => JobExecution::Success, + Err(err) => job_execution_from_task_error(err), + } + } +} +struct ProvisionContext { + job_id: Uuid, + job_type: &'static str, + tenant_id: Uuid, + state: Arc, + members: Vec, +} + +impl ProvisionContext { + fn new( + job_id: Uuid, + job_type: &'static str, + tenant_id: Uuid, + state: Arc, + members: Vec, + ) -> Self { + Self { + job_id, + job_type, + tenant_id, + state, + members, + } + } +} + +impl TaskContext for ProvisionContext { + fn job_id(&self) -> Uuid { + self.job_id + } + + fn job_type(&self) -> &'static str { + self.job_type + } +} + +struct ProvisionPlanner; + +#[async_trait] +impl TaskPlanner for ProvisionPlanner { + async fn plan( + &self, + _ctx: &mut ProvisionContext, + ) -> TaskResult>> { + Ok(vec![Box::new(ProvisionTask)]) + } +} + +struct ProvisionTask; + +#[async_trait] +impl Task for ProvisionTask { + fn name(&self) -> &'static str { + "provision-tenant" + } + + async fn execute(&self, ctx: &mut ProvisionContext) -> TaskResult<()> { + let mut conn = ctx + .state + .db_unscoped() + .map_err(|err| TaskError::retry(Duration::from_secs(30), format!("{err:?}")))?; + + let tenant = TenantRepository::get_by_id(&mut conn, ctx.tenant_id).map_err(|err| { + TaskError::fail(format!("tenant not found for provisioning: {err:?}")) + })?; drop(conn); - let mut conn = match state.db_for_tenant(tenant.id) { - Ok(conn) => conn, - Err(err) => { - warn!(job_id = %job.id, error = ?err, "failed to scope connection for tenant provisioning"); - return JobExecution::Retry { - delay: std::time::Duration::from_secs(30), - error: "tenant connection unavailable".into(), - }; - } - }; + let mut conn = ctx + .state + .db_for_tenant(tenant.id) + .map_err(|err| TaskError::retry(Duration::from_secs(30), format!("{err:?}")))?; if tenant.status == TenantStatus::Active { warn!( - job_id = %job.id, + job_id = %ctx.job_id(), tenant_id = %tenant.id, "tenant already active; skipping provisioning" ); - return JobExecution::Success; + return Ok(()); } if tenant.status != TenantStatus::Creating { - warn!( - job_id = %job.id, - tenant_id = %tenant.id, - status = %tenant.status.as_str(), - "tenant not in creating state; skipping provisioning" - ); - return JobExecution::Failed { - error: format!( - "tenant status '{}' not eligible for provisioning", - tenant.status.as_str() - ), - }; + return Err(TaskError::fail(format!( + "tenant status '{}' not eligible for provisioning", + tenant.status.as_str() + ))); } - let endpoint = match &state.config.quickwit_endpoint { - Some(endpoint) => endpoint.trim_end_matches('/').to_owned(), - None => { - warn!( - job_id = %job.id, - tenant_id = %tenant.id, - "quickwit endpoint not configured; retrying" - ); - return JobExecution::Retry { - delay: std::time::Duration::from_secs(30), - error: "quickwit endpoint not configured".into(), - }; - } - }; + + let endpoint = ctx + .state + .config + .quickwit_endpoint + .as_ref() + .map(|value| value.trim_end_matches('/').to_owned()) + .ok_or_else(|| { + TaskError::retry(Duration::from_secs(30), "quickwit endpoint not configured") + })?; let index_id = tenant .quickwit_index @@ -119,117 +167,69 @@ impl JobHandler for ProvisionTenantJob { .unwrap_or_else(|| format!("documents-{}", tenant.id)); let client = Client::new(); - if let Err(err) = ensure_quickwit_index(&client, &endpoint, &index_id).await { - warn!( - job_id = %job.id, - tenant_id = %tenant.id, - error = %err, - "failed to ensure quickwit index" - ); - return JobExecution::Retry { - delay: std::time::Duration::from_secs(30), - error: err.to_string(), - }; - } + ensure_quickwit_index(&client, &endpoint, &index_id) + .await + .map_err(|err| TaskError::retry(Duration::from_secs(30), err.to_string()))?; let owner_capability_set_id = - match ensure_capability_set(&mut conn, tenant.id, owner_capabilities()) { - Ok(set) => set.id, - Err(err) => { - warn!( - job_id = %job.id, - tenant_id = %tenant.id, - error = ?err, - "failed to ensure owner capability set during provisioning" - ); - return JobExecution::Retry { - delay: std::time::Duration::from_secs(30), - error: "owner capability set unavailable".into(), - }; - } + ensure_capability_set(&mut conn, tenant.id, owner_capabilities()) + .map_err(|_| { + TaskError::retry(Duration::from_secs(30), "owner capability set unavailable") + })? + .id; + + ensure_capability_set(&mut conn, tenant.id, user_capabilities()).map_err(|_| { + TaskError::retry(Duration::from_secs(30), "user capability set unavailable") + })?; + ensure_capability_set(&mut conn, tenant.id, readonly_capabilities()).map_err(|_| { + TaskError::retry( + Duration::from_secs(30), + "readonly capability set unavailable", + ) + })?; + ensure_capability_set(&mut conn, tenant.id, webdav_capabilities()).map_err(|_| { + TaskError::retry(Duration::from_secs(30), "webdav capability set unavailable") + })?; + + for member in &ctx.members { + let new_membership = NewUserMembership { + id: Uuid::new_v4(), + user_id: *member, + tenant_id: tenant.id, + capability_set_id: Some(owner_capability_set_id), }; - if let Err(err) = ensure_capability_set(&mut conn, tenant.id, user_capabilities()) { - warn!( - job_id = %job.id, - tenant_id = %tenant.id, - error = ?err, - "failed to ensure user capability set during provisioning" - ); - return JobExecution::Retry { - delay: std::time::Duration::from_secs(30), - error: "user capability set unavailable".into(), - }; - } - - if let Err(err) = ensure_capability_set(&mut conn, tenant.id, readonly_capabilities()) { - warn!( - job_id = %job.id, - tenant_id = %tenant.id, - error = ?err, - "failed to ensure readonly capability set during provisioning" - ); - return JobExecution::Retry { - delay: std::time::Duration::from_secs(30), - error: "readonly capability set unavailable".into(), - }; - } - - if let Err(err) = ensure_capability_set(&mut conn, tenant.id, webdav_capabilities()) { - warn!( - job_id = %job.id, - tenant_id = %tenant.id, - error = ?err, - "failed to ensure webdav capability set during provisioning" - ); - return JobExecution::Retry { - delay: std::time::Duration::from_secs(30), - error: "webdav capability set unavailable".into(), - }; - } - - if let Some(members) = ProvisionPayload::from_job(&job) { - for member in members { - let new_membership = NewUserMembership { - id: Uuid::new_v4(), - user_id: member, - tenant_id: tenant.id, - capability_set_id: Some(owner_capability_set_id), - }; - - if let Err(err) = diesel::insert_into(user_memberships::table) - .values(&new_membership) - .on_conflict((user_memberships::user_id, user_memberships::tenant_id)) - .do_nothing() - .execute(&mut conn) - { - warn!( - job_id = %job.id, - tenant_id = %tenant.id, - user_id = %member, - error = %err, - "failed to assign initial membership" - ); - } + if let Err(err) = diesel::insert_into(user_memberships::table) + .values(&new_membership) + .on_conflict((user_memberships::user_id, user_memberships::tenant_id)) + .do_nothing() + .execute(&mut conn) + { + warn!( + job_id = %ctx.job_id(), + tenant_id = %tenant.id, + user_id = %member, + error = %err, + "failed to assign initial membership" + ); } } - if let Err(err) = diesel::update(tenants::table.find(tenant.id)) + diesel::update(tenants::table.find(tenant.id)) .set(( tenants::status.eq(TenantStatus::Active), tenants::quickwit_index.eq(Some(index_id)), tenants::updated_at.eq(Utc::now().naive_utc()), )) .execute(&mut conn) - { - warn!(job_id = %job.id, error = %err, "failed to activate tenant"); - return JobExecution::Retry { - delay: std::time::Duration::from_secs(30), - error: format!("failed to update tenant status: {err}"), - }; - } + .map_err(|err| { + TaskError::retry( + Duration::from_secs(30), + format!("failed to update tenant status: {err}"), + ) + })?; - JobExecution::Success + Ok(()) } } diff --git a/backend/src/workers/thumbnails.rs b/backend/src/workers/thumbnails.rs index 766d546..deae483 100644 --- a/backend/src/workers/thumbnails.rs +++ b/backend/src/workers/thumbnails.rs @@ -8,7 +8,7 @@ use pdfium_render::prelude::*; use serde::Deserialize; use serde_json::{json, Map, Value}; use tokio::task; -use tracing::{error, info, warn}; +use tracing::{info, warn}; use uuid::Uuid; use crate::{ @@ -27,16 +27,20 @@ use crate::{ use super::{ analyze::determine_thumbnail_support, - common::{load_document_version, load_version_assets}, - fetch_version_object, handle_fetch_error, JobExecution, JobHandler, + job_execution_from_task_error, + taskflow::{ + document::DocumentVersionTaskContext, BoxedTask, Task, TaskContext, TaskError, + TaskExecutor, TaskPlanner, TaskResult, + }, + JobExecution, JobHandler, }; -const THUMBNAIL_WIDTH: u32 = 512; -const THUMBNAIL_HEIGHT: u32 = 512; +pub const THUMBNAIL_WIDTH: u32 = 512; +pub const THUMBNAIL_HEIGHT: u32 = 512; const PREVIEW_WIDTH: u32 = THUMBNAIL_WIDTH * 4; const PREVIEW_HEIGHT: u32 = THUMBNAIL_HEIGHT * 4; -const THUMBNAIL_ASSET_TYPE: &str = "thumbnail"; -const PREVIEW_ASSET_TYPE: &str = "preview"; +pub const THUMBNAIL_ASSET_TYPE: &str = "thumbnail"; +pub const PREVIEW_ASSET_TYPE: &str = "preview"; #[derive(Debug, Deserialize, Clone)] struct ThumbnailPayload { @@ -75,268 +79,109 @@ impl JobHandler for GenerateThumbnailsJob { } }; - let state_clone = state.clone(); - let tenant_id = job.tenant_id; - let payload_for_context = payload.clone(); - let initial = match task::spawn_blocking(move || { - load_thumbnail_context(state_clone, tenant_id, payload_for_context) - }) - .await - { - Ok(Ok(ctx)) => ctx, - Ok(Err(err)) => { - warn!(job_id = %job.id, error = %err, "thumbnail job will retry"); - return JobExecution::Retry { - delay: Duration::from_secs(30), - error: err, - }; - } - Err(join_err) => { - error!(job_id = %job.id, error = %join_err, "thumbnail task panicked"); - return JobExecution::Retry { - delay: Duration::from_secs(60), - error: format!("worker panicked: {join_err}"), - }; - } - }; + let mut context = DocumentVersionTaskContext::new( + job.id, + JOB_GENERATE_THUMBNAILS, + job.tenant_id, + payload.document_id, + payload.document_version_id, + payload.force, + state.config.worker_max_document_bytes, + state.clone(), + storage, + ); - if initial.skip { - info!(job_id = %job.id, "thumbnails already exist; skipping"); - return JobExecution::Success; + let planner = ThumbnailPlanner { + force: payload.force, + }; + match TaskExecutor::run(&planner, &mut context).await { + Ok(()) => JobExecution::Success, + Err(err) => job_execution_from_task_error(err), + } + } +} + +struct ThumbnailPlanner { + force: bool, +} + +#[async_trait] +impl TaskPlanner for ThumbnailPlanner { + async fn plan( + &self, + _ctx: &mut DocumentVersionTaskContext, + ) -> TaskResult>> { + Ok(vec![Box::new(GenerateThumbnailsTask::new(self.force))]) + } +} + +pub struct GenerateThumbnailsTask { + force: bool, +} + +impl GenerateThumbnailsTask { + pub fn new(force: bool) -> Self { + Self { force } + } +} + +#[async_trait] +impl Task for GenerateThumbnailsTask { + fn name(&self) -> &'static str { + "generate-thumbnails" + } + + async fn execute(&self, ctx: &mut DocumentVersionTaskContext) -> TaskResult<()> { + let context = build_thumbnail_context(ctx, self.force).await?; + + if context.skip { + info!(job_id = %ctx.job_id(), "thumbnails already exist; skipping"); + return Ok(()); } - let bytes = match fetch_version_object( - &initial.version, - &storage, - &initial.version.s3_key, - state.config.worker_max_document_bytes, - ) - .await - { - Ok(bytes) => bytes, - Err(err) => return handle_fetch_error(&job, err, "thumbnail fetch failed; will retry"), - }; - - let generation = match generate_preview_and_thumbnail(&initial.document, &bytes) { - Ok(result) => result, - Err(err) => { - return JobExecution::Failed { error: err }; - } - }; + let bytes = ctx.buffered_object().await?; + let generation = + generate_preview_and_thumbnail(&context.document, bytes).map_err(TaskError::fail)?; if let Some(page_count) = generation.page_count { - let state_clone = state.clone(); - let document_id = initial.document.id; - let version_id = initial.version.id; - let tenant_id = initial.tenant_id; - match task::spawn_blocking(move || { - persist_document_page_count( - state_clone, - tenant_id, - document_id, - version_id, - page_count, + let state = ctx.state().clone(); + let tenant_id = context.tenant_id; + let document_id = context.document.id; + let version_id = context.version.id; + task::spawn_blocking(move || { + persist_document_page_count(state, tenant_id, document_id, version_id, page_count) + }) + .await + .map_err(|err| { + TaskError::retry( + Duration::from_secs(60), + format!("page count task panicked: {err}"), ) - }) - .await - { - Ok(Ok(())) => {} - Ok(Err(err)) => { - warn!( - job_id = %job.id, - document_id = %document_id, - version_id = %version_id, - error = %err, - "failed to update document page count metadata; retrying" - ); - return JobExecution::Retry { - delay: Duration::from_secs(30), - error: err, - }; - } - Err(join_err) => { - error!( - job_id = %job.id, - document_id = %document_id, - version_id = %version_id, - error = %join_err, - "page count metadata task panicked" - ); - return JobExecution::Retry { - delay: Duration::from_secs(60), - error: format!("metadata panic: {join_err}"), - }; - } - } + })? + .map_err(|err| TaskError::retry(Duration::from_secs(30), err))?; } - if let Some(existing_preview) = &initial.existing_preview { - for object in &initial.existing_preview_objects { - if let Err(err) = storage.delete_object(&object.s3_key).await { - warn!( - job_id = %job.id, - error = %err, - s3_key = %object.s3_key, - "failed to delete existing preview object" - ); - } - } - - let tenant_id = initial.document.tenant_id; - let asset_id = existing_preview.id; - let state_clone = state.clone(); - match task::spawn_blocking(move || -> AppResult<()> { - let mut conn = state_clone.db_for_tenant(tenant_id)?; - delete_asset(&mut conn, tenant_id, asset_id) - }) - .await - { - Ok(Ok(())) => {} - Ok(Err(err)) => { - warn!( - job_id = %job.id, - error = ?err, - asset_id = %asset_id, - "failed to remove preview metadata after deletion" - ); - } - Err(join_err) => { - warn!( - job_id = %job.id, - error = %join_err, - asset_id = %asset_id, - "failed to remove preview metadata: task panicked" - ); - } - } - } - - if let Some(existing_thumbnail) = &initial.existing_thumbnail { - for object in &initial.existing_thumbnail_objects { - if let Err(err) = storage.delete_object(&object.s3_key).await { - warn!( - job_id = %job.id, - error = %err, - s3_key = %object.s3_key, - "failed to delete existing thumbnail object" - ); - } - } - - let tenant_id = initial.document.tenant_id; - let asset_id = existing_thumbnail.id; - let state_clone = state.clone(); - match task::spawn_blocking(move || -> AppResult<()> { - let mut conn = state_clone.db_for_tenant(tenant_id)?; - delete_asset(&mut conn, tenant_id, asset_id) - }) - .await - { - Ok(Ok(())) => {} - Ok(Err(err)) => { - warn!( - job_id = %job.id, - error = ?err, - asset_id = %asset_id, - "failed to remove thumbnail metadata after deletion" - ); - } - Err(join_err) => { - warn!( - job_id = %job.id, - error = %join_err, - asset_id = %asset_id, - "failed to remove thumbnail metadata: task panicked" - ); - } - } - } + remove_existing_assets(ctx, &context).await; let preview_asset_id = Uuid::new_v4(); - let thumbnail_asset_id = Uuid::new_v4(); - let mut preview_objects: Vec = - Vec::with_capacity(generation.preview.objects.len()); - for (index, image) in generation.preview.objects.iter().enumerate() { - if index + 1 > i32::MAX as usize { - return JobExecution::Failed { - error: "too many preview objects".to_string(), - }; - } - let ordinal = (index + 1) as i32; - let s3_key = document_asset_object_key( - initial.document.id, - initial.version.version_number, - PREVIEW_ASSET_TYPE, - preview_asset_id, - ordinal, - ); - - if let Err(err) = storage - .put_object( - &s3_key, - image.image_bytes.clone(), - Some("image/webp".into()), - None, - ) - .await - { - warn!(job_id = %job.id, error = %err, ordinal, "failed to upload preview; retrying"); - return JobExecution::Retry { - delay: Duration::from_secs(30), - error: err.to_string(), - }; - } - - preview_objects.push(AssetObjectPersistence { - ordinal, - s3_key, - width: image.width, - height: image.height, - }); - } - - let mut thumbnail_objects: Vec = - Vec::with_capacity(generation.thumbnail.objects.len()); - for (index, image) in generation.thumbnail.objects.iter().enumerate() { - if index + 1 > i32::MAX as usize { - return JobExecution::Failed { - error: "too many thumbnail objects".to_string(), - }; - } - let ordinal = (index + 1) as i32; - let s3_key = document_asset_object_key( - initial.document.id, - initial.version.version_number, - THUMBNAIL_ASSET_TYPE, - thumbnail_asset_id, - ordinal, - ); - - if let Err(err) = storage - .put_object( - &s3_key, - image.image_bytes.clone(), - Some("image/webp".into()), - None, - ) - .await - { - warn!(job_id = %job.id, error = %err, ordinal, "failed to upload thumbnail; retrying"); - return JobExecution::Retry { - delay: Duration::from_secs(30), - error: err.to_string(), - }; - } - - thumbnail_objects.push(AssetObjectPersistence { - ordinal, - s3_key, - width: image.width, - height: image.height, - }); - } + let preview_objects = upload_generated_objects( + ctx, + &context, + PREVIEW_ASSET_TYPE, + preview_asset_id, + &generation.preview, + ) + .await?; + let thumbnail_objects = upload_generated_objects( + ctx, + &context, + THUMBNAIL_ASSET_TYPE, + thumbnail_asset_id, + &generation.thumbnail, + ) + .await?; let asset_persistences = vec![ AssetPersistence { @@ -351,33 +196,186 @@ impl JobHandler for GenerateThumbnailsJob { }, ]; - let state_clone = state.clone(); - match task::spawn_blocking(move || { - persist_assets_metadata(state_clone, &initial, &asset_persistences) + let state = ctx.state().clone(); + let tenant_id = context.document.tenant_id; + let version_id = context.version.id; + task::spawn_blocking(move || { + persist_assets_metadata(state, tenant_id, version_id, &asset_persistences) }) .await - { - Ok(Ok(())) => {} - Ok(Err(err)) => { - warn!(job_id = %job.id, error = %err, "failed to persist thumbnail metadata; retrying"); - return JobExecution::Retry { - delay: Duration::from_secs(30), - error: err, - }; - } - Err(join_err) => { - error!(job_id = %job.id, error = %join_err, "thumbnail metadata update panicked"); - return JobExecution::Retry { - delay: Duration::from_secs(30), - error: format!("metadata update panic: {join_err}"), - }; - } - } + .map_err(|err| { + TaskError::retry( + Duration::from_secs(60), + format!("thumbnail metadata task panicked: {err}"), + ) + })? + .map_err(|err| TaskError::retry(Duration::from_secs(30), err))?; - JobExecution::Success + ctx.invalidate_asset_cache(); + Ok(()) } } +async fn build_thumbnail_context( + ctx: &mut DocumentVersionTaskContext, + force: bool, +) -> TaskResult { + let document = ctx.document().await?.clone(); + let version = ctx.version().await?.clone(); + let tenant_id = ctx.tenant_id(); + + let (supported, _) = determine_thumbnail_support(&document); + if !supported { + return Ok(ThumbnailContext { + document, + version, + existing_thumbnail: None, + existing_thumbnail_objects: Vec::new(), + existing_preview: None, + existing_preview_objects: Vec::new(), + skip: true, + tenant_id, + }); + } + + let assets = ctx.assets().await?; + let (existing_thumbnail, existing_thumbnail_objects) = assets + .get(THUMBNAIL_ASSET_TYPE) + .map(|entry| (Some(entry.asset.clone()), entry.objects.clone())) + .unwrap_or((None, Vec::new())); + let (existing_preview, existing_preview_objects) = assets + .get(PREVIEW_ASSET_TYPE) + .map(|entry| (Some(entry.asset.clone()), entry.objects.clone())) + .unwrap_or((None, Vec::new())); + + let expected_cardinality = expected_asset_cardinality(&document, &version); + let preview_cardinality = existing_preview + .as_ref() + .and_then(|asset| asset.cardinality) + .unwrap_or_else(|| existing_preview_objects.len() as i32); + let thumbnail_cardinality = existing_thumbnail + .as_ref() + .and_then(|asset| asset.cardinality) + .unwrap_or_else(|| existing_thumbnail_objects.len() as i32); + let needs_regeneration = preview_cardinality < expected_cardinality + || thumbnail_cardinality < expected_cardinality + || (existing_preview_objects.len() as i32) < expected_cardinality + || (existing_thumbnail_objects.len() as i32) < expected_cardinality; + + let skip = + existing_thumbnail.is_some() && existing_preview.is_some() && !force && !needs_regeneration; + + Ok(ThumbnailContext { + document, + version, + existing_thumbnail, + existing_thumbnail_objects, + existing_preview, + existing_preview_objects, + skip, + tenant_id, + }) +} + +async fn remove_existing_assets(ctx: &DocumentVersionTaskContext, context: &ThumbnailContext) { + if let Some(existing_preview) = &context.existing_preview { + delete_asset_with_objects(ctx, existing_preview, &context.existing_preview_objects).await; + } + + if let Some(existing_thumbnail) = &context.existing_thumbnail { + delete_asset_with_objects(ctx, existing_thumbnail, &context.existing_thumbnail_objects) + .await; + } +} + +async fn delete_asset_with_objects( + ctx: &DocumentVersionTaskContext, + asset: &DocumentAsset, + objects: &[DocumentAssetObject], +) { + for object in objects { + if let Err(err) = ctx.storage().delete_object(&object.s3_key).await { + warn!( + job_id = %ctx.job_id(), + error = %err, + s3_key = %object.s3_key, + "failed to delete existing asset object" + ); + } + } + + let tenant_id = ctx.tenant_id(); + let asset_id = asset.id; + let state = ctx.state().clone(); + match task::spawn_blocking(move || -> AppResult<()> { + let mut conn = state.db_for_tenant(tenant_id)?; + delete_asset(&mut conn, tenant_id, asset_id) + }) + .await + { + Ok(Ok(())) => {} + Ok(Err(err)) => { + warn!( + job_id = %ctx.job_id(), + asset_id = %asset_id, + error = ?err, + "failed to delete asset metadata" + ); + } + Err(join_err) => { + warn!( + job_id = %ctx.job_id(), + asset_id = %asset_id, + error = %join_err, + "failed to delete asset metadata task panicked" + ); + } + } +} + +async fn upload_generated_objects( + ctx: &DocumentVersionTaskContext, + context: &ThumbnailContext, + asset_type: &str, + asset_id: Uuid, + asset: &GeneratedAsset, +) -> TaskResult> { + let mut objects = Vec::with_capacity(asset.objects.len()); + + for (index, image) in asset.objects.iter().enumerate() { + if index + 1 > i32::MAX as usize { + return Err(TaskError::fail("too many generated asset objects")); + } + let ordinal = (index + 1) as i32; + let s3_key = document_asset_object_key( + context.document.id, + context.version.version_number, + asset_type, + asset_id, + ordinal, + ); + + ctx.storage() + .put_object( + &s3_key, + image.image_bytes.clone(), + Some("image/webp".into()), + None, + ) + .await + .map_err(|err| TaskError::retry(Duration::from_secs(30), err.to_string()))?; + + objects.push(AssetObjectPersistence { + ordinal, + s3_key, + width: image.width, + height: image.height, + }); + } + + Ok(objects) +} + struct ThumbnailContext { document: Document, version: DocumentVersion, @@ -418,83 +416,11 @@ struct AssetPersistence { objects: Vec, } -fn load_thumbnail_context( - state: Arc, - tenant_id: Uuid, - payload: ThumbnailPayload, -) -> Result { - let base = load_document_version( - state.as_ref(), - tenant_id, - payload.document_id, - payload.document_version_id, - )?; - - let mut conn = state - .db_for_tenant(base.tenant_id) - .map_err(|err| format!("{err:?}"))?; - - let mut assets = load_version_assets( - &mut conn, - base.tenant_id, - base.version.id, - &[THUMBNAIL_ASSET_TYPE, PREVIEW_ASSET_TYPE], - )?; - - let (existing_thumbnail, existing_thumbnail_objects) = assets - .remove(THUMBNAIL_ASSET_TYPE) - .map(|entry| (Some(entry.asset), entry.objects)) - .unwrap_or((None, Vec::new())); - - let (existing_preview, existing_preview_objects) = assets - .remove(PREVIEW_ASSET_TYPE) - .map(|entry| (Some(entry.asset), entry.objects)) - .unwrap_or((None, Vec::new())); - - let (supported, _) = determine_thumbnail_support(&base.document); - if !supported { - return Err("thumbnail generation not supported for this document".into()); - } - - let expected_cardinality = expected_asset_cardinality(&base.document, &base.version); - let preview_cardinality = existing_preview - .as_ref() - .and_then(|asset| asset.cardinality) - .unwrap_or_else(|| existing_preview_objects.len() as i32); - let thumbnail_cardinality = existing_thumbnail - .as_ref() - .and_then(|asset| asset.cardinality) - .unwrap_or_else(|| existing_thumbnail_objects.len() as i32); - - let needs_regeneration = preview_cardinality < expected_cardinality - || thumbnail_cardinality < expected_cardinality - || (existing_preview_objects.len() as i32) < expected_cardinality - || (existing_thumbnail_objects.len() as i32) < expected_cardinality; - - let skip = existing_thumbnail.is_some() - && existing_preview.is_some() - && !payload.force - && !needs_regeneration; - - Ok(ThumbnailContext { - document: base.document, - version: base.version, - existing_thumbnail, - existing_thumbnail_objects, - existing_preview, - existing_preview_objects, - skip, - tenant_id: base.tenant_id, - }) -} - fn generate_preview_and_thumbnail( document: &Document, bytes: &[u8], ) -> Result { - let is_pdf = document_is_pdf(document); - - if is_pdf { + if document_is_pdf(document) { let pdf_assets = generate_pdf_assets(bytes)?; Ok(GeneratedAssets { preview: pdf_assets.preview, @@ -623,28 +549,14 @@ fn encode_dynamic_image(image: image::DynamicImage) -> Result, - context: &ThumbnailContext, + tenant_id: Uuid, + version_id: Uuid, assets: &[AssetPersistence], ) -> Result<(), String> { - let tenant_id = context.document.tenant_id; let mut conn = state .db_for_tenant(tenant_id) .map_err(|err| format!("{err:?}"))?; - if let Some(existing_preview) = &context.existing_preview { - diesel::delete(document_assets::table.filter(document_assets::id.eq(existing_preview.id))) - .execute(&mut conn) - .map_err(|err| format!("{err:?}"))?; - } - - if let Some(existing_thumbnail) = &context.existing_thumbnail { - diesel::delete( - document_assets::table.filter(document_assets::id.eq(existing_thumbnail.id)), - ) - .execute(&mut conn) - .map_err(|err| format!("{err:?}"))?; - } - for asset in assets { if asset.objects.is_empty() { return Err(format!( @@ -661,7 +573,7 @@ fn persist_assets_metadata( let new_asset = NewDocumentAsset { id: asset.asset_id, - document_version_id: context.version.id, + document_version_id: version_id, asset_type: asset.asset_type.to_string(), mime_type: "image/webp".to_string(), metadata: json!({ @@ -688,7 +600,7 @@ fn persist_assets_metadata( diesel::delete( document_asset_objects::table - .filter(document_asset_objects::asset_id.eq(asset.asset_id)) + .filter(document_asset_objects::asset_id.eq(new_asset.id)) .filter(document_asset_objects::tenant_id.eq(tenant_id)), ) .execute(&mut conn) @@ -707,7 +619,7 @@ fn persist_assets_metadata( let new_object = NewDocumentAssetObject { id: Uuid::new_v4(), - asset_id: asset.asset_id, + asset_id: new_asset.id, ordinal: object.ordinal, s3_key: object.s3_key.clone(), metadata: object_metadata,