taskflow

2025-11-10 23:06:53 +01:00
parent 8dddb0b940
commit 8788cd3e25
11 changed files with 1336 additions and 1177 deletions
@@ -2,24 +2,25 @@ 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,
+    job_execution_from_task_error,
+    ocr::OCR_TEXT_ASSET_TYPE,
+    taskflow::{
+        document::DocumentVersionTaskContext, BoxedTask, Task, TaskError, TaskExecutor,
+        TaskPlanner, TaskResult,
+    },
+    JobExecution, JobHandler,
 };

 #[derive(Debug, Deserialize)]
@@ -53,161 +54,116 @@ impl JobHandler for IndexDocumentTextJob {
            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 tenant_id = job.tenant_id;
-        let context =
-            match task::spawn_blocking(move || load_context(state_clone, tenant_id, 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,
+        let mut context = DocumentVersionTaskContext::new(
+            job.id,
+            JOB_INDEX_DOCUMENT_TEXT,
+            job.tenant_id,
+            payload.document_id,
+            payload.document_version_id,
+            false,
            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(),
-                };
-            }
-        };
+            state.clone(),
+            storage,
+        );
+
+        let planner = IndexPlanner::new();
+        match TaskExecutor::run(&planner, &mut context).await {
+            Ok(()) => JobExecution::Success,
+            Err(err) => job_execution_from_task_error(err),
+        }
+    }
+}
+
+struct IndexPlanner;
+
+impl IndexPlanner {
+    fn new() -> Self {
+        Self
+    }
+}
+
+#[async_trait]
+impl TaskPlanner<DocumentVersionTaskContext> for IndexPlanner {
+    async fn plan(
+        &self,
+        _ctx: &mut DocumentVersionTaskContext,
+    ) -> TaskResult<Vec<BoxedTask<DocumentVersionTaskContext>>> {
+        Ok(vec![Box::new(IndexDocumentTask::new())])
+    }
+}
+
+pub struct IndexDocumentTask;
+
+impl IndexDocumentTask {
+    pub fn new() -> Self {
+        Self
+    }
+}
+
+#[async_trait]
+impl Task<DocumentVersionTaskContext> for IndexDocumentTask {
+    fn name(&self) -> &'static str {
+        "index-document-text"
+    }
+
+    async fn execute(&self, ctx: &mut DocumentVersionTaskContext) -> TaskResult<()> {
+        let state = ctx.state().clone();
+        let quickwit_endpoint = state
+            .config
+            .quickwit_endpoint
+            .clone()
+            .ok_or_else(|| TaskError::fail("quickwit endpoint missing"))?;
+
+        let tenant = state.tenants.get_by_id(ctx.tenant_id()).map_err(|err| {
+            TaskError::retry(
+                Duration::from_secs(30),
+                format!("failed to load tenant: {err:?}"),
+            )
+        })?;
+
+        let quickwit_index = tenant
+            .quickwit_index
+            .clone()
+            .ok_or_else(|| TaskError::fail("tenant quickwit index not configured"))?;
+
+        let asset = ctx
+            .asset(OCR_TEXT_ASSET_TYPE)
+            .await?
+            .ok_or_else(|| TaskError::fail("missing OCR text asset"))?;
+        let object = asset
+            .objects
+            .first()
+            .ok_or_else(|| TaskError::fail("missing OCR text object"))?;
+
+        let s3_key = object.s3_key.clone();
+        let bytes = ctx.storage().get_object(&s3_key).await.map_err(|err| {
+            TaskError::retry(
+                Duration::from_secs(30),
+                format!("failed to download ocr text: {err}"),
+            )
+        })?;
+
+        let text = String::from_utf8(bytes)
+            .map_err(|err| TaskError::fail(format!("ocr text not valid UTF-8: {err}")))?;

        if text.trim().is_empty() {
-            warn!(job_id = %job.id, "ocr text empty; skipping");
-            return JobExecution::Failed {
-                error: "ocr text empty".into(),
-            };
+            return Err(TaskError::fail("ocr text empty"));
        }

-        let record =
-            build_quickwit_ingest_record(&context.document, &context.version, job.tenant_id, &text);
+        let document = ctx.document().await?.clone();
+        let version = ctx.version().await?.clone();

-        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(),
-                }
-            }
-        }
+        let record = build_quickwit_ingest_record(&document, &version, ctx.tenant_id(), &text);
+        let client = Client::new();
+
+        quickwit_ingest(&client, &quickwit_endpoint, &quickwit_index, &[record])
+            .await
+            .map_err(|err| TaskError::retry(Duration::from_secs(30), err.to_string()))?;
+
+        Ok(())
    }
 }
-
-struct IndexContext {
-    document: Document,
-    version: DocumentVersion,
-    text_s3_key: Option<String>,
-}
-
-fn load_context(
-    state: Arc<AppState>,
-    tenant_id: Uuid,
-    payload: IndexPayload,
-) -> Result<IndexContext, String> {
-    let mut conn = state
-        .db_for_tenant(tenant_id)
-        .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_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,
-    })
-}