use std::sync::Arc; use std::time::Duration; use async_trait::async_trait; use diesel::prelude::*; use reqwest::Client; use serde::Deserialize; use tokio::task; use tracing::{error, warn}; use uuid::Uuid; use crate::{ documents::search::{build_quickwit_ingest_record, quickwit_ingest}, jobs::JOB_INDEX_DOCUMENT_TEXT, models::{Document, DocumentVersion}, schema::{document_asset_objects, document_assets, document_versions, documents}, state::AppState, storage::TenantStorage, }; use super::{ fetch_version_object, handle_fetch_error, 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, storage: TenantStorage, ) -> 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 => { return JobExecution::Failed { error: "quickwit endpoint missing".into(), }; } }; let tenant = match state.tenants.get_by_id(job.tenant_id) { Ok(tenant) => tenant, Err(err) => { warn!(job_id = %job.id, error = ?err, "failed to load tenant for indexing"); return JobExecution::Retry { delay: Duration::from_secs(30), error: format!("failed to load tenant: {err:?}"), }; } }; let quickwit_index = match tenant.quickwit_index.clone() { Some(index) => index, None => { return JobExecution::Failed { error: "tenant quickwit index not configured".into(), }; } }; 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_s3_key.is_none() { warn!(job_id = %job.id, "missing OCR text asset; failing indexing job"); return JobExecution::Failed { error: "missing OCR text asset".into(), }; } let s3_key = context.text_s3_key.unwrap(); let bytes = match fetch_version_object( &context.version, &storage, &s3_key, state.config.worker_max_document_bytes, ) .await { Ok(bytes) => bytes, Err(err) => return handle_fetch_error(&job, err, "failed to download ocr text"), }; let text = 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(), }; } }; if text.trim().is_empty() { warn!(job_id = %job.id, "ocr text empty; skipping"); return JobExecution::Failed { error: "ocr text empty".into(), }; } let record = build_quickwit_ingest_record(&context.document, &context.version, job.tenant_id, &text); match quickwit_ingest(&client, &quickwit_endpoint, &quickwit_index, &[record]).await { Ok(()) => JobExecution::Success, Err(err) => { warn!(job_id = %job.id, error = %err, "quickwit ingest failed"); JobExecution::Retry { delay: Duration::from_secs(30), error: err.to_string(), } } } } } struct IndexContext { document: Document, version: DocumentVersion, text_s3_key: Option, } fn load_context(state: Arc, payload: &IndexPayload) -> Result { 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 mut conn = state .db_for_tenant(tenant_id) .map_err(|err| format!("{err:?}"))?; let text_s3_key: Option = document_asset_objects::table .inner_join( document_assets::table.on(document_asset_objects::asset_id.eq(document_assets::id)), ) .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)) .filter(document_asset_objects::ordinal.eq(1)) .filter(document_asset_objects::tenant_id.eq(tenant_id)) .select(document_asset_objects::s3_key) .first(&mut conn) .optional() .map_err(|err| format!("{err:?}"))?; Ok(IndexContext { document, version, text_s3_key, }) }