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, 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, } fn load_context(state: Arc, payload: &IndexPayload) -> Result { 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 = 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, }) }