quickwit

2025-10-11 16:42:25 +02:00
parent cd27f36e3d
commit 0a669064dd
13 changed files with 521 additions and 12 deletions
@@ -0,0 +1,204 @@
+use std::sync::Arc;
+use std::time::Duration;
+
+use async_trait::async_trait;
+use diesel::prelude::*;
+use reqwest::Client;
+use serde::Deserialize;
+use serde_json::json;
+use tokio::task;
+use tracing::{error, warn};
+use uuid::Uuid;
+
+use crate::{
+    jobs::JOB_INDEX_DOCUMENT_TEXT,
+    models::{Document, DocumentAsset, DocumentVersion},
+    schema::{document_assets, document_versions, documents},
+    state::AppState,
+};
+
+use super::{ocr::OCR_TEXT_ASSET_TYPE, JobExecution, JobHandler};
+
+#[derive(Debug, Deserialize)]
+struct IndexPayload {
+    document_id: Uuid,
+    document_version_id: Uuid,
+}
+
+pub struct IndexDocumentTextJob;
+
+impl IndexDocumentTextJob {
+    pub fn new() -> Self {
+        Self
+    }
+}
+
+#[async_trait]
+impl JobHandler for IndexDocumentTextJob {
+    fn job_type(&self) -> &'static str {
+        JOB_INDEX_DOCUMENT_TEXT
+    }
+
+    async fn handle(&self, state: Arc<AppState>, job: crate::models::Job) -> JobExecution {
+        let payload: IndexPayload = match serde_json::from_value(job.payload.clone()) {
+            Ok(payload) => payload,
+            Err(err) => {
+                return JobExecution::Failed {
+                    error: format!("invalid index payload: {err}"),
+                }
+            }
+        };
+
+        let quickwit_endpoint = match &state.config.quickwit_endpoint {
+            Some(endpoint) => endpoint.clone(),
+            None => {
+                warn!("quickwit endpoint missing; skipping indexing");
+                return JobExecution::Success;
+            }
+        };
+
+        let quickwit_index = match &state.config.quickwit_index {
+            Some(index) => index.clone(),
+            None => {
+                warn!("quickwit index missing; skipping indexing");
+                return JobExecution::Success;
+            }
+        };
+
+        let client = Client::new();
+
+        let state_clone = state.clone();
+        let context = match task::spawn_blocking(move || load_context(state_clone, &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_asset.is_none() {
+            warn!(job_id = %job.id, "missing OCR text asset; failing indexing job");
+            return JobExecution::Failed {
+                error: "missing OCR text asset".into(),
+            };
+        }
+
+        let asset = context.text_asset.unwrap();
+        let text = match state.storage.get_object(&asset.s3_key).await {
+            Ok(bytes) => 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(),
+                    };
+                }
+            },
+            Err(err) => {
+                warn!(job_id = %job.id, error = %err, "failed to download ocr text");
+                return JobExecution::Retry {
+                    delay: Duration::from_secs(30),
+                    error: err.to_string(),
+                };
+            }
+        };
+
+        if text.trim().is_empty() {
+            warn!(job_id = %job.id, "ocr text empty; skipping");
+            return JobExecution::Failed {
+                error: "ocr text empty".into(),
+            };
+        }
+
+        let client = client;
+        let url = format!(
+            "{}/api/v1/{}/ingest?commit=auto",
+            quickwit_endpoint, quickwit_index
+        );
+        let payload = json!({
+            "document_id": context.document.id,
+            "version_id": context.version.id,
+            "title": context.document.title,
+            "text": text,
+        });
+
+        let body = serde_json::to_string(&payload).unwrap();
+
+        match client
+            .post(&url)
+            .header("content-type", "application/x-ndjson")
+            .body(format!("{}\n", body))
+            .send()
+            .await
+        {
+            Ok(response) => {
+                if response.status().is_success() {
+                    JobExecution::Success
+                } else {
+                    let status = response.status();
+                    let body = response.text().await.unwrap_or_default();
+                    warn!(job_id = %job.id, %status, %body, "quickwit ingest failed");
+                    JobExecution::Retry {
+                        delay: Duration::from_secs(30),
+                        error: format!("quickwit ingest failed with status {status}"),
+                    }
+                }
+            }
+            Err(err) => {
+                warn!(job_id = %job.id, error = %err, "quickwit request failed");
+                JobExecution::Retry {
+                    delay: Duration::from_secs(30),
+                    error: err.to_string(),
+                }
+            }
+        }
+    }
+}
+
+struct IndexContext {
+    document: Document,
+    version: DocumentVersion,
+    text_asset: Option<DocumentAsset>,
+}
+
+fn load_context(state: Arc<AppState>, payload: &IndexPayload) -> Result<IndexContext, 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 text_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))
+        .first(&mut conn)
+        .optional()
+        .map_err(|err| format!("{err:?}"))?;
+
+    Ok(IndexContext {
+        document,
+        version,
+        text_asset,
+    })
+}
@@ -11,6 +11,7 @@ use crate::{
 };

 pub mod analyze;
+pub mod index;
 pub mod ocr;
 pub mod thumbnails;

@@ -130,5 +131,6 @@ pub fn default_handlers() -> Vec<Arc<dyn JobHandler>> {
        Arc::new(analyze::AnalyzeDocumentJob::new()),
        Arc::new(thumbnails::GenerateThumbnailsJob::new()),
        Arc::new(ocr::GenerateOcrTextJob::new()),
+        Arc::new(index::IndexDocumentTextJob::new()),
    ]
 }
@@ -1,6 +1,5 @@
 use std::{
-    fmt,
-    fs,
+    fmt, fs,
    io::{ErrorKind, Write},
    process::Command,
    sync::Arc,
@@ -19,7 +18,7 @@ use tracing::{error, info, warn};
 use uuid::Uuid;

 use crate::{
-    jobs::JOB_GENERATE_OCR_TEXT,
+    jobs::{enqueue_job, JOB_GENERATE_OCR_TEXT, JOB_INDEX_DOCUMENT_TEXT},
    models::{Document, DocumentAsset, DocumentVersion, NewDocumentAsset},
    schema::{document_assets, document_versions, documents},
    state::AppState,
@@ -30,7 +29,7 @@ use super::{JobExecution, JobHandler};
 pub const OCR_TEXT_ASSET_TYPE: &str = "ocr-text";
 const MIN_TEXT_LENGTH: usize = 50;

-#[derive(Debug, Deserialize)]
+#[derive(Clone, Debug, Deserialize)]
 struct OcrPayload {
    document_id: Uuid,
    document_version_id: Uuid,
@@ -63,8 +62,10 @@ impl JobHandler for GenerateOcrTextJob {
        };

        let state_clone = state.clone();
+        let payload_clone = payload.clone();
        let context =
-            match task::spawn_blocking(move || load_ocr_context(state_clone, &payload)).await {
+            match task::spawn_blocking(move || load_ocr_context(state_clone, &payload_clone)).await
+            {
                Ok(Ok(ctx)) => ctx,
                Ok(Err(err)) => {
                    warn!(job_id = %job.id, error = %err, "ocr job will retry");
@@ -155,7 +156,15 @@ impl JobHandler for GenerateOcrTextJob {
        })
        .await
        {
-            Ok(Ok(())) => JobExecution::Success,
+            Ok(Ok(())) => {
+                if state.config.quickwit_endpoint.is_some() && state.config.quickwit_index.is_some()
+                {
+                    if let Err(err) = enqueue_index_job(&state, &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 {
@@ -334,7 +343,8 @@ fn run_ocr(bytes: &[u8]) -> Result<Option<String>, OcrError> {
                )));
            }

-            let text = fs::read_to_string(sidecar.path()).map_err(|err| OcrError::Failed(err.to_string()))?;
+            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 {
@@ -392,6 +402,21 @@ fn persist_ocr_metadata(
    Ok(())
 }

+fn enqueue_index_job(state: &AppState, payload: &OcrPayload) -> Result<(), String> {
+    let mut conn = state.db().map_err(|err| format!("{err:?}"))?;
+    enqueue_job(
+        &mut conn,
+        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(),