use std::{collections::HashSet, sync::Arc, time::Duration}; 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, }; use super::{JobExecution, JobHandler}; #[derive(Debug, Deserialize)] struct AnalyzePayload { document_id: Uuid, document_version_id: Uuid, #[serde(default)] force: bool, } pub struct AnalyzeDocumentJob; impl AnalyzeDocumentJob { pub fn new() -> Self { Self } } #[async_trait] impl JobHandler for AnalyzeDocumentJob { fn job_type(&self) -> &'static str { JOB_ANALYZE_DOCUMENT } async fn handle( &self, state: Arc, job: crate::models::Job, _storage: TenantStorage, ) -> JobExecution { let payload: AnalyzePayload = match serde_json::from_value(job.payload.clone()) { Ok(payload) => payload, Err(err) => { return JobExecution::Failed { error: format!("invalid analyze payload: {err}"), } } }; 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}"), } } } } } fn analyze_document( 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()); } 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 ) })?; 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()); } } 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, ); if let Err(err) = enqueue_result { return Err(err.to_string()); } } Ok(JobExecution::Success) } pub(crate) fn determine_thumbnail_support(document: &Document) -> (bool, Option) { let supported_mimes: HashSet<&'static str> = [ "image/jpeg", "image/png", "image/gif", "image/tiff", "image/bmp", "image/webp", "application/pdf", ] .into_iter() .collect(); if let Some(ref content_type) = document.content_type { if supported_mimes.contains(content_type.as_str()) { return (true, None); } } if let Some(ext) = document .original_name .rsplit('.') .next() .map(|ext| ext.to_ascii_lowercase()) { let supported_exts = [ "jpg", "jpeg", "png", "gif", "tif", "tiff", "bmp", "webp", "pdf", ]; if supported_exts.contains(&ext.as_str()) { return (true, None); } } ( false, Some("content type not supported for thumbnails".into()), ) }