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 serde_json::json;
use tokio::task;
use tracing::{error, warn};
use uuid::Uuid;

use crate::{
    jobs::JOB_INDEX_DOCUMENT_TEXT,
    models::{Document, DocumentVersion},
    schema::{document_asset_objects, document_assets, document_versions, documents},
    state::AppState,
    storage::TenantStorage,
};

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,
        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 text = match storage.get_object(&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,
            "tenant_id": job.tenant_id,
            "title": context.document.title.to_lowercase(),
            "text": text.to_lowercase()
        });

        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_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,
    })
}