papercrate/backend/src/workers/index.rs

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<AppState>,
        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<String>,
}

fn load_context(state: Arc<AppState>, payload: &IndexPayload) -> Result<IndexContext, 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 mut conn = state
        .db_for_tenant(tenant_id)
        .map_err(|err| format!("{err:?}"))?;

    let text_s3_key: Option<String> = 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,
    })
}