cleanup

backend/src/workers/ocr.rs

@@ -24,13 +24,16 @@ use crate::{
         Document, DocumentAsset, DocumentAssetObject, DocumentVersion, NewDocumentAsset,
         NewDocumentAssetObject,
     },
-    schema::{document_asset_objects, document_assets, document_versions, documents},
+    schema::{document_asset_objects, document_assets, documents},
     state::AppState,
     storage::TenantStorage,
     utils::storage_paths::document_asset_object_prefix,
 };
 
-use super::{fetch_version_object, handle_fetch_error, JobExecution, JobHandler};
+use super::{
+    common::{load_document_version, load_version_assets},
+    fetch_version_object, handle_fetch_error, JobExecution, JobHandler,
+};
 
 pub const OCR_TEXT_ASSET_TYPE: &str = "ocr-text";
 const MIN_TEXT_LENGTH: usize = 50;
@@ -234,54 +237,34 @@ struct OcrGeneration {
 }
 
 fn load_ocr_context(state: Arc<AppState>, payload: &OcrPayload) -> Result<OcrContext, String> {
-    let mut base_conn = state.db_unscoped().map_err(|err| format!("{err:?}"))?;
-
-    let version: DocumentVersion = document_versions::table
-        .find(payload.document_version_id)
-        .first(&mut base_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 base_conn)
-        .map_err(|err| format!("{err:?}"))?;
-
-    let tenant_id = document.tenant_id;
-    drop(base_conn);
+    let base = load_document_version(
+        state.as_ref(),
+        payload.document_id,
+        payload.document_version_id,
+    )?;
 
     let mut conn = state
-        .db_for_tenant(tenant_id)
+        .db_for_tenant(base.tenant_id)
         .map_err(|err| format!("{err:?}"))?;
 
-    let existing_asset: Option<DocumentAsset> = 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 mut assets = load_version_assets(
+        &mut conn,
+        base.tenant_id,
+        base.version.id,
+        &[OCR_TEXT_ASSET_TYPE],
+    )?;
 
-    let existing_objects: Vec<DocumentAssetObject> = if let Some(asset) = &existing_asset {
-        document_asset_objects::table
-            .filter(document_asset_objects::asset_id.eq(asset.id))
-            .filter(document_asset_objects::tenant_id.eq(tenant_id))
-            .order(document_asset_objects::ordinal.asc())
-            .load(&mut conn)
-            .map_err(|err| format!("{err:?}"))?
-    } else {
-        Vec::new()
-    };
+    let (existing_asset, existing_objects) = assets
+        .remove(OCR_TEXT_ASSET_TYPE)
+        .map(|entry| (Some(entry.asset), entry.objects))
+        .unwrap_or((None, Vec::new()));
 
-    let is_pdf = document_is_pdf(&document);
+    let is_pdf = document_is_pdf(&base.document);
     if !is_pdf {
         return Ok(OcrContext {
-            document,
-            version,
-            existing_asset: existing_asset,
+            document: base.document,
+            version: base.version,
+            existing_asset,
             existing_objects,
             skip: true,
         });
@@ -290,8 +273,8 @@ fn load_ocr_context(state: Arc<AppState>, payload: &OcrPayload) -> Result<OcrCon
     let skip = existing_asset.is_some() && !payload.force;
 
     Ok(OcrContext {
-        document,
-        version,
+        document: base.document,
+        version: base.version,
         existing_asset,
         existing_objects,
         skip,