ocr

backend/src/workers/analyze.rs

@@ -8,10 +8,11 @@ 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_THUMBNAILS},
-    models::{Document, DocumentVersion},
-    schema::{document_versions, documents},
+    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,
 };
 
@@ -88,6 +89,16 @@ fn analyze_document(state: Arc<AppState>, payload: AnalyzePayload) -> Result<Job
         .map_err(|err| format!("{err:?}"))?;
 
     let (supported, reason) = determine_thumbnail_support(&document);
+    let ocr_supported = document_is_pdf(&document);
+
+    let existing_ocr: 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))
+        .first(&mut conn)
+        .optional()
+        .map_err(|err| format!("{err:?}"))?;
+
+    let skip_ocr = existing_ocr.is_some() && !payload.force;
 
     let mut summary_map = match version.operations_summary {
         Value::Object(map) => map,
@@ -100,6 +111,16 @@ fn analyze_document(state: Arc<AppState>, payload: AnalyzePayload) -> Result<Job
         summary_map.remove("thumbnail_reason");
     }
 
+    summary_map.insert("ocr_supported".to_string(), Value::Bool(ocr_supported));
+    if ocr_supported {
+        summary_map.remove("ocr_reason");
+    } else {
+        summary_map.insert(
+            "ocr_reason".to_string(),
+            Value::String("document is not a PDF".into()),
+        );
+    }
+
     diesel::update(document_versions::table.find(version.id))
         .set(document_versions::operations_summary.eq(Value::Object(summary_map)))
         .execute(&mut conn)
@@ -122,6 +143,23 @@ fn analyze_document(state: Arc<AppState>, payload: AnalyzePayload) -> Result<Job
         }
     }
 
+    if ocr_supported && !skip_ocr {
+        let enqueue_result = enqueue_job(
+            &mut conn,
+            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)
 }
 

backend/src/workers/mod.rs

@@ -11,6 +11,7 @@ use crate::{
 };
 
 pub mod analyze;
+pub mod ocr;
 pub mod thumbnails;
 
 #[derive(Debug)]
@@ -128,5 +129,6 @@ pub fn default_handlers() -> Vec<Arc<dyn JobHandler>> {
     vec![
         Arc::new(analyze::AnalyzeDocumentJob::new()),
         Arc::new(thumbnails::GenerateThumbnailsJob::new()),
+        Arc::new(ocr::GenerateOcrTextJob::new()),
     ]
 }

backend/src/workers/ocr.rs

@@ -0,0 +1,414 @@
+use std::{
+    fmt,
+    fs,
+    io::{ErrorKind, Write},
+    process::Command,
+    sync::Arc,
+    time::Duration,
+};
+
+use async_trait::async_trait;
+use chrono::Utc;
+use diesel::{pg::upsert::excluded, prelude::*};
+use pdfium_render::prelude::*;
+use serde::Deserialize;
+use serde_json::json;
+use tempfile::NamedTempFile;
+use tokio::task;
+use tracing::{error, info, warn};
+use uuid::Uuid;
+
+use crate::{
+    jobs::JOB_GENERATE_OCR_TEXT,
+    models::{Document, DocumentAsset, DocumentVersion, NewDocumentAsset},
+    schema::{document_assets, document_versions, documents},
+    state::AppState,
+};
+
+use super::{JobExecution, JobHandler};
+
+pub const OCR_TEXT_ASSET_TYPE: &str = "ocr-text";
+const MIN_TEXT_LENGTH: usize = 50;
+
+#[derive(Debug, Deserialize)]
+struct OcrPayload {
+    document_id: Uuid,
+    document_version_id: Uuid,
+    #[serde(default)]
+    force: bool,
+}
+
+pub struct GenerateOcrTextJob;
+
+impl GenerateOcrTextJob {
+    pub fn new() -> Self {
+        Self
+    }
+}
+
+#[async_trait]
+impl JobHandler for GenerateOcrTextJob {
+    fn job_type(&self) -> &'static str {
+        JOB_GENERATE_OCR_TEXT
+    }
+
+    async fn handle(&self, state: Arc<AppState>, job: crate::models::Job) -> JobExecution {
+        let payload: OcrPayload = match serde_json::from_value(job.payload.clone()) {
+            Ok(payload) => payload,
+            Err(err) => {
+                return JobExecution::Failed {
+                    error: format!("invalid OCR payload: {err}"),
+                }
+            }
+        };
+
+        let state_clone = state.clone();
+        let context =
+            match task::spawn_blocking(move || load_ocr_context(state_clone, &payload)).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}"),
+                    };
+                }
+            };
+
+        if context.skip {
+            info!(job_id = %job.id, "ocr already present; skipping");
+            return JobExecution::Success;
+        }
+
+        let bytes = match state.storage.get_object(&context.version.s3_key).await {
+            Ok(bytes) => bytes,
+            Err(err) => {
+                warn!(job_id = %job.id, error = %err, "failed to fetch document for ocr");
+                return JobExecution::Retry {
+                    delay: Duration::from_secs(30),
+                    error: err.to_string(),
+                };
+            }
+        };
+
+        let doc_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 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(),
+            };
+        };
+
+        let asset_id = context
+            .existing_asset
+            .as_ref()
+            .map(|asset| asset.id)
+            .unwrap_or_else(Uuid::new_v4);
+
+        let s3_key = format!(
+            "documents/{}/v{}/assets/{}/{}",
+            context.document.id, context.version.version_number, OCR_TEXT_ASSET_TYPE, asset_id
+        );
+
+        if let Err(err) = state
+            .storage
+            .put_object(
+                &s3_key,
+                generation.text.into_bytes(),
+                Some("text/plain".into()),
+            )
+            .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(),
+            };
+        }
+
+        let state_clone = state.clone();
+        match task::spawn_blocking(move || {
+            persist_ocr_metadata(state_clone, &context, asset_id, &s3_key, generation.source)
+        })
+        .await
+        {
+            Ok(Ok(())) => 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}"),
+                }
+            }
+        }
+    }
+}
+
+struct PdfDocumentMeta {
+    content_type: Option<String>,
+    original_name: String,
+}
+
+struct OcrContext {
+    document: Document,
+    version: DocumentVersion,
+    existing_asset: Option<DocumentAsset>,
+    skip: bool,
+}
+
+struct OcrGeneration {
+    text: String,
+    source: &'static str,
+}
+
+fn load_ocr_context(state: Arc<AppState>, payload: &OcrPayload) -> Result<OcrContext, String> {
+    let mut conn = state.db().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 existing: 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))
+        .first(&mut conn)
+        .optional()
+        .map_err(|err| format!("{err:?}"))?;
+
+    let is_pdf = document_is_pdf(&document);
+    if !is_pdf {
+        return Ok(OcrContext {
+            document,
+            version,
+            existing_asset: existing,
+            skip: true,
+        });
+    }
+
+    let skip = existing.is_some() && !payload.force;
+
+    Ok(OcrContext {
+        document,
+        version,
+        existing_asset: existing,
+        skip,
+    })
+}
+
+fn generate_ocr_text(meta: &PdfDocumentMeta, bytes: &[u8]) -> Option<OcrGeneration> {
+    if !document_meta_is_pdf(meta) {
+        return None;
+    }
+
+    if let Ok(text) = extract_pdf_text(bytes) {
+        if text.trim().chars().count() >= MIN_TEXT_LENGTH {
+            return Some(OcrGeneration {
+                text,
+                source: "pdf-text",
+            });
+        }
+    }
+
+    match run_ocr(bytes) {
+        Ok(Some(text)) => Some(OcrGeneration {
+            text,
+            source: "ocr",
+        }),
+        Ok(None) => None,
+        Err(OcrError::BinaryMissing) => {
+            warn!("ocrmypdf not installed; cannot perform OCR");
+            None
+        }
+        Err(err) => {
+            warn!(error = ?err, "ocr command failed");
+            None
+        }
+    }
+}
+
+fn extract_pdf_text(bytes: &[u8]) -> Result<String, String> {
+    let pdfium = Pdfium::default();
+    let document = pdfium
+        .load_pdf_from_byte_slice(bytes, None)
+        .map_err(|err| format!("load pdf: {err}"))?;
+
+    let mut combined = String::new();
+    let pages = document.pages();
+    for page_index in 0..pages.len() {
+        let page = pages
+            .get(page_index)
+            .map_err(|err| format!("load page {page_index}: {err}"))?;
+        if let Ok(page_text) = page.text() {
+            for segment in page_text.segments().iter() {
+                combined.push_str(&segment.text());
+                combined.push('\n');
+            }
+        };
+    }
+
+    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<Option<String>, OcrError> {
+    let mut input = NamedTempFile::new().map_err(|err| OcrError::Failed(err.to_string()))?;
+    input
+        .write_all(bytes)
+        .map_err(|err| OcrError::Failed(err.to_string()))?;
+    input
+        .flush()
+        .map_err(|err| OcrError::Failed(err.to_string()))?;
+
+    let output_pdf = NamedTempFile::new().map_err(|err| OcrError::Failed(err.to_string()))?;
+    let sidecar = NamedTempFile::new().map_err(|err| OcrError::Failed(err.to_string()))?;
+
+    let status = Command::new("ocrmypdf")
+        .arg("--sidecar")
+        .arg(sidecar.path())
+        .arg("--skip-text")
+        .arg(input.path())
+        .arg(output_pdf.path())
+        .output();
+
+    match status {
+        Ok(output) => {
+            if !output.status.success() {
+                return Err(OcrError::Failed(format!(
+                    "ocrmypdf failed: exit={} stderr={}",
+                    output.status,
+                    String::from_utf8_lossy(&output.stderr)
+                )));
+            }
+
+            let text = fs::read_to_string(sidecar.path()).map_err(|err| OcrError::Failed(err.to_string()))?;
+            if text.trim().chars().count() >= MIN_TEXT_LENGTH {
+                Ok(Some(text))
+            } else {
+                Ok(None)
+            }
+        }
+        Err(err) => {
+            if err.kind() == ErrorKind::NotFound {
+                Err(OcrError::BinaryMissing)
+            } else {
+                Err(OcrError::Failed(err.to_string()))
+            }
+        }
+    }
+}
+
+fn persist_ocr_metadata(
+    state: Arc<AppState>,
+    context: &OcrContext,
+    asset_id: Uuid,
+    s3_key: &str,
+    source: &'static str,
+) -> Result<(), String> {
+    let mut conn = state.db().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(),
+        s3_key: s3_key.to_string(),
+        mime_type: "text/plain".to_string(),
+        width: None,
+        height: None,
+        metadata: json!({
+            "generated_at": Utc::now().to_rfc3339(),
+            "source": source,
+        }),
+    };
+
+    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::s3_key.eq(excluded(document_assets::s3_key)),
+            document_assets::mime_type.eq(excluded(document_assets::mime_type)),
+            document_assets::metadata.eq(excluded(document_assets::metadata)),
+        ))
+        .execute(&mut conn)
+        .map_err(|err| format!("{err:?}"))?;
+
+    Ok(())
+}
+
+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") {
+            return true;
+        }
+    }
+
+    meta.original_name
+        .rsplit('.')
+        .next()
+        .map(|ext| ext.eq_ignore_ascii_case("pdf"))
+        .unwrap_or(false)
+}

backend/src/workers/thumbnails.rs

@@ -102,10 +102,20 @@ impl JobHandler for GenerateThumbnailsJob {
             }
         };
 
+        let asset_id = initial
+            .existing_asset
+            .as_ref()
+            .map(|asset| asset.id)
+            .unwrap_or_else(Uuid::new_v4);
+        let s3_key = format!(
+            "documents/{}/v{}/assets/{}/{}",
+            initial.document.id, initial.version.version_number, THUMBNAIL_ASSET_TYPE, asset_id
+        );
+
         if let Err(err) = state
             .storage
             .put_object(
-                &generation.s3_key,
+                &s3_key,
                 generation.image_bytes.clone(),
                 Some("image/png".into()),
             )
@@ -120,7 +130,7 @@ impl JobHandler for GenerateThumbnailsJob {
 
         let state_clone = state.clone();
         match task::spawn_blocking(move || {
-            persist_thumbnail_metadata(state_clone, &initial, &generation)
+            persist_thumbnail_metadata(state_clone, &initial, &generation, asset_id, &s3_key)
         })
         .await
         {
@@ -148,6 +158,7 @@ impl JobHandler for GenerateThumbnailsJob {
 struct ThumbnailContext {
     document: Document,
     version: DocumentVersion,
+    existing_asset: Option<DocumentAsset>,
     skip: bool,
 }
 
@@ -155,7 +166,6 @@ struct GeneratedThumbnail {
     image_bytes: Vec<u8>,
     width: Option<i32>,
     height: Option<i32>,
-    s3_key: String,
 }
 
 fn load_thumbnail_context(
@@ -195,6 +205,7 @@ fn load_thumbnail_context(
     Ok(ThumbnailContext {
         document,
         version,
+        existing_asset: existing,
         skip,
     })
 }
@@ -219,13 +230,10 @@ fn generate_thumbnail(document: &Document, bytes: &[u8]) -> Result<GeneratedThum
         generate_image_thumbnail(bytes)?
     };
 
-    let s3_key = format!("thumbnails/{}/{}.png", document.id, Uuid::new_v4());
-
     Ok(GeneratedThumbnail {
         image_bytes: png_bytes,
         width,
         height,
-        s3_key,
     })
 }
 
@@ -284,14 +292,16 @@ fn persist_thumbnail_metadata(
     state: Arc<AppState>,
     context: &ThumbnailContext,
     generated: &GeneratedThumbnail,
+    asset_id: Uuid,
+    s3_key: &str,
 ) -> Result<(), String> {
     let mut conn = state.db().map_err(|err| format!("{err:?}"))?;
 
     let new_asset = NewDocumentAsset {
-        id: Uuid::new_v4(),
+        id: asset_id,
         document_version_id: context.version.id,
         asset_type: THUMBNAIL_ASSET_TYPE.to_string(),
-        s3_key: generated.s3_key.clone(),
+        s3_key: s3_key.to_string(),
         mime_type: "image/png".to_string(),
         width: generated.width,
         height: generated.height,