nils/papercrate
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

taskflow

Browse Source
This commit is contained in:
nils
2025-11-10 23:06:53 +01:00
parent 8dddb0b940
commit 8788cd3e25
11 changed files with 1336 additions and 1177 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/src/workers/ocr.rs
+302 -291
View File
@@ -14,13 +14,13 @@ use serde::Deserialize;
use serde_json::json;
use tempfile::NamedTempFile;
use tokio::task;
use tracing::{error, info, warn};
use tracing::{info, warn};
use uuid::Uuid;
use crate::{
documents::asset::delete_asset,
error::AppResult,
jobs::{enqueue_job, JOB_GENERATE_OCR_TEXT, JOB_INDEX_DOCUMENT_TEXT},
jobs::JOB_GENERATE_OCR_TEXT,
models::{
Document, DocumentAsset, DocumentAssetObject, DocumentVersion, NewDocumentAsset,
NewDocumentAssetObject,
@@ -32,8 +32,13 @@ use crate::{
};
use super::{
common::{load_document_version, load_version_assets},
fetch_version_object, handle_fetch_error, JobExecution, JobHandler,
index::IndexDocumentTask,
job_execution_from_task_error,
taskflow::{
document::DocumentVersionTaskContext, BoxedTask, Task, TaskContext, TaskError,
TaskExecutor, TaskPlanner, TaskResult,
},
JobExecution, JobHandler,
};
pub const OCR_TEXT_ASSET_TYPE: &str = "ocr-text";
@@ -72,104 +77,102 @@ impl JobHandler for GenerateOcrTextJob {
Err(err) => {
return JobExecution::Failed {
error: format!("invalid OCR payload: {err}"),
}
};
}
};
let state_clone = state.clone();
let payload_clone = payload.clone();
let tenant_id = job.tenant_id;
let context = match task::spawn_blocking(move || {
load_ocr_context(state_clone, tenant_id, payload_clone)
})
.await
{
Ok(Ok(ctx)) => ctx,
Ok(Err(err)) => {
warn!(job_id = %job.id, error = %err, "ocr job will retry");
return JobExecution::Retry {
delay: Duration::from_secs(30),
error: err,
};
}
Err(join_err) => {
error!(job_id = %job.id, error = %join_err, "ocr task panicked");
return JobExecution::Retry {
delay: Duration::from_secs(60),
error: format!("worker panicked: {join_err}"),
};
}
};
let mut context = DocumentVersionTaskContext::new(
job.id,
JOB_GENERATE_OCR_TEXT,
job.tenant_id,
payload.document_id,
payload.document_version_id,
payload.force,
state.config.worker_max_document_bytes,
state.clone(),
storage,
);
let planner = OcrPlanner::new(payload.force, state.clone());
match TaskExecutor::run(&planner, &mut context).await {
Ok(()) => JobExecution::Success,
Err(err) => job_execution_from_task_error(err),
}
}
}
struct OcrPlanner {
force: bool,
state: Arc<AppState>,
}
impl OcrPlanner {
fn new(force: bool, state: Arc<AppState>) -> Self {
Self { force, state }
}
}
#[async_trait]
impl TaskPlanner<DocumentVersionTaskContext> for OcrPlanner {
async fn plan(
&self,
_ctx: &mut DocumentVersionTaskContext,
) -> TaskResult<Vec<BoxedTask<DocumentVersionTaskContext>>> {
Ok(vec![
Box::new(GenerateOcrTask::new(self.force, self.state.clone())),
Box::new(IndexDocumentTask::new()),
])
}
}
pub struct GenerateOcrTask {
force: bool,
state: Arc<AppState>,
}
impl GenerateOcrTask {
pub fn new(force: bool, state: Arc<AppState>) -> Self {
Self { force, state }
}
}
#[async_trait]
impl Task<DocumentVersionTaskContext> for GenerateOcrTask {
fn name(&self) -> &'static str {
"generate-ocr-text"
}
async fn execute(&self, ctx: &mut DocumentVersionTaskContext) -> TaskResult<()> {
let context = build_ocr_context(ctx, self.force).await?;
if context.skip {
info!(job_id = %job.id, "ocr already present; skipping");
return JobExecution::Success;
info!(job_id = %ctx.job_id(), "ocr already present; skipping");
return Ok(());
}
let bytes = match fetch_version_object(
&context.version,
&storage,
&context.version.s3_key,
state.config.worker_max_document_bytes,
)
.await
{
Ok(bytes) => bytes,
Err(err) => return handle_fetch_error(&job, err, "failed to fetch document for ocr"),
};
let doc_meta = PdfDocumentMeta {
let bytes = ctx.buffered_object().await?.to_vec();
let meta = PdfDocumentMeta {
content_type: context.document.content_type.clone(),
original_name: context.document.original_name.clone(),
};
let generation =
match task::spawn_blocking(move || generate_ocr_text(&doc_meta, &bytes)).await {
Ok(result) => result,
Err(join_err) => {
error!(job_id = %job.id, error = %join_err, "ocr text task panicked");
return JobExecution::Retry {
delay: Duration::from_secs(60),
error: format!("worker panicked: {join_err}"),
};
}
};
let generation = task::spawn_blocking(move || generate_ocr_text(&meta, &bytes))
.await
.map_err(|err| {
TaskError::retry(
Duration::from_secs(60),
format!("ocr text task panicked: {err}"),
)
})?;
let Some(generation) = generation else {
warn!(job_id = %job.id, "no text extracted from document; failing job");
return JobExecution::Failed {
error: "no text extracted and OCR unavailable".into(),
};
warn!(job_id = %ctx.job_id(), "no text extracted from document; failing job");
return Err(TaskError::fail("no text extracted and OCR unavailable"));
};
if let Some(existing_asset) = &context.existing_asset {
for object in &context.existing_objects {
if let Err(err) = storage.delete_object(&object.s3_key).await {
warn!(job_id = %job.id, error = %err, s3_key = %object.s3_key, "failed to delete existing ocr asset object");
}
}
let tenant_id = context.document.tenant_id;
let asset_id = existing_asset.id;
let state_clone = state.clone();
match task::spawn_blocking(move || -> AppResult<()> {
let mut conn = state_clone.db_for_tenant(tenant_id)?;
delete_asset(&mut conn, tenant_id, asset_id)
})
.await
{
Ok(Ok(())) => {}
Ok(Err(err)) => {
warn!(job_id = %job.id, error = ?err, asset_id = %asset_id, "failed to remove ocr asset metadata after deletion");
}
Err(join_err) => {
warn!(job_id = %job.id, error = %join_err, asset_id = %asset_id, "failed to remove ocr asset metadata: task panicked");
}
}
}
remove_existing_ocr_asset(ctx, &context).await;
let asset_id = Uuid::new_v4();
let s3_key = document_asset_object_prefix(
context.document.id,
context.version.version_number,
@@ -177,7 +180,7 @@ impl JobHandler for GenerateOcrTextJob {
asset_id,
);
if let Err(err) = storage
ctx.storage()
.put_object(
&s3_key,
generation.text.into_bytes(),
@@ -185,47 +188,25 @@ impl JobHandler for GenerateOcrTextJob {
None,
)
.await
{
warn!(job_id = %job.id, error = %err, "failed to upload ocr text");
return JobExecution::Retry {
delay: Duration::from_secs(30),
error: err.to_string(),
};
}
.map_err(|err| TaskError::retry(Duration::from_secs(30), err.to_string()))?;
let state_clone = state.clone();
match task::spawn_blocking(move || {
persist_ocr_metadata(state_clone, &context, asset_id, &s3_key, generation.source)
let state = self.state.clone();
task::spawn_blocking(move || {
persist_ocr_metadata(state, &context, asset_id, &s3_key, generation.source)
})
.await
{
Ok(Ok(())) => {
if let Err(err) = enqueue_index_job(&state, job.tenant_id, &payload) {
warn!(job_id = %job.id, error = %err, "failed to enqueue index job");
}
JobExecution::Success
}
Ok(Err(err)) => {
warn!(job_id = %job.id, error = %err, "failed to persist ocr metadata");
JobExecution::Retry {
delay: Duration::from_secs(30),
error: err,
}
}
Err(join_err) => {
error!(job_id = %job.id, error = %join_err, "ocr metadata task panicked");
JobExecution::Retry {
delay: Duration::from_secs(30),
error: format!("metadata update panic: {join_err}"),
}
}
}
}
}
.map_err(|err| {
TaskError::retry(
Duration::from_secs(60),
format!("ocr metadata task panicked: {err}"),
)
})?
.map_err(|err| TaskError::retry(Duration::from_secs(30), err))?;
struct PdfDocumentMeta {
content_type: Option<String>,
original_name: String,
ctx.invalidate_asset_cache();
Ok(())
}
}
struct OcrContext {
@@ -236,61 +217,163 @@ struct OcrContext {
skip: bool,
}
struct OcrGeneration {
text: String,
source: &'static str,
}
async fn build_ocr_context(
ctx: &mut DocumentVersionTaskContext,
force: bool,
) -> TaskResult<OcrContext> {
let document = ctx.document().await?.clone();
let version = ctx.version().await?.clone();
let asset = ctx.asset(OCR_TEXT_ASSET_TYPE).await?;
fn load_ocr_context(
state: Arc<AppState>,
tenant_id: Uuid,
payload: OcrPayload,
) -> Result<OcrContext, String> {
let base = load_document_version(
state.as_ref(),
tenant_id,
payload.document_id,
payload.document_version_id,
)?;
let mut conn = state
.db_for_tenant(base.tenant_id)
.map_err(|err| format!("{err:?}"))?;
let mut assets = load_version_assets(
&mut conn,
base.tenant_id,
base.version.id,
&[OCR_TEXT_ASSET_TYPE],
)?;
let (existing_asset, existing_objects) = assets
.remove(OCR_TEXT_ASSET_TYPE)
.map(|entry| (Some(entry.asset), entry.objects))
let (existing_asset, existing_objects) = asset
.map(|asset| (Some(asset.asset.clone()), asset.objects.clone()))
.unwrap_or((None, Vec::new()));
let is_pdf = document_is_pdf(&base.document);
if !is_pdf {
if !document_is_pdf(&document) {
return Ok(OcrContext {
document: base.document,
version: base.version,
document,
version,
existing_asset,
existing_objects,
skip: true,
});
}
let skip = existing_asset.is_some() && !payload.force;
let skip = existing_asset.is_some() && !force;
Ok(OcrContext {
document: base.document,
version: base.version,
document,
version,
existing_asset,
existing_objects,
skip,
})
}
async fn remove_existing_ocr_asset(ctx: &DocumentVersionTaskContext, context: &OcrContext) {
if let Some(existing_asset) = &context.existing_asset {
for object in &context.existing_objects {
if let Err(err) = ctx.storage().delete_object(&object.s3_key).await {
warn!(
job_id = %ctx.job_id(),
error = %err,
s3_key = %object.s3_key,
"failed to delete existing ocr asset object"
);
}
}
let tenant_id = context.document.tenant_id;
let asset_id = existing_asset.id;
let state = ctx.state().clone();
match task::spawn_blocking(move || -> AppResult<()> {
let mut conn = state.db_for_tenant(tenant_id)?;
delete_asset(&mut conn, tenant_id, asset_id)
})
.await
{
Ok(Ok(())) => {}
Ok(Err(err)) => {
warn!(
job_id = %ctx.job_id(),
error = ?err,
asset_id = %asset_id,
"failed to remove ocr asset metadata after deletion"
);
}
Err(join_err) => {
warn!(
job_id = %ctx.job_id(),
error = %join_err,
asset_id = %asset_id,
"failed to remove ocr asset metadata: task panicked"
);
}
}
}
}
fn persist_ocr_metadata(
state: Arc<AppState>,
context: &OcrContext,
asset_id: Uuid,
s3_key: &str,
source: OcrSource,
) -> Result<(), String> {
let tenant_id = context.document.tenant_id;
let mut conn = state
.db_for_tenant(tenant_id)
.map_err(|err| format!("{err:?}"))?;
let document_version_id = context.version.id;
let existing_asset = context.existing_asset.as_ref().map(|asset| asset.id);
if let Some(existing_asset) = existing_asset {
diesel::delete(
document_assets::table
.filter(document_assets::id.eq(existing_asset))
.filter(document_assets::tenant_id.eq(tenant_id)),
)
.execute(&mut conn)
.map_err(|err| format!("{err:?}"))?;
}
let metadata = json!({
"source": source.to_string(),
"generated_at": Utc::now().to_rfc3339(),
});
let new_asset = NewDocumentAsset {
id: asset_id,
document_version_id,
asset_type: OCR_TEXT_ASSET_TYPE.to_string(),
mime_type: "text/plain".to_string(),
metadata,
cardinality: Some(1),
tenant_id,
};
diesel::insert_into(document_assets::table)
.values(&new_asset)
.on_conflict((
document_assets::document_version_id,
document_assets::asset_type,
))
.do_update()
.set((
document_assets::mime_type.eq(excluded(document_assets::mime_type)),
document_assets::metadata.eq(excluded(document_assets::metadata)),
document_assets::cardinality.eq(excluded(document_assets::cardinality)),
document_assets::id.eq(excluded(document_assets::id)),
))
.execute(&mut conn)
.map_err(|err| format!("{err:?}"))?;
diesel::delete(
document_asset_objects::table
.filter(document_asset_objects::asset_id.eq(asset_id))
.filter(document_asset_objects::tenant_id.eq(tenant_id)),
)
.execute(&mut conn)
.map_err(|err| format!("{err:?}"))?;
let object = NewDocumentAssetObject {
id: Uuid::new_v4(),
asset_id,
ordinal: 1,
s3_key: s3_key.to_string(),
metadata: json!({}),
tenant_id,
};
diesel::insert_into(document_asset_objects::table)
.values(&object)
.execute(&mut conn)
.map_err(|err| format!("{err:?}"))?;
Ok(())
}
fn generate_ocr_text(meta: &PdfDocumentMeta, bytes: &[u8]) -> Option<OcrGeneration> {
if !document_meta_is_pdf(meta) {
return None;
@@ -300,7 +383,7 @@ fn generate_ocr_text(meta: &PdfDocumentMeta, bytes: &[u8]) -> Option<OcrGenerati
if text.trim().chars().count() >= MIN_TEXT_LENGTH {
return Some(OcrGeneration {
text,
source: "pdf-text",
source: OcrSource::PdfText,
});
}
}
@@ -308,20 +391,60 @@ fn generate_ocr_text(meta: &PdfDocumentMeta, bytes: &[u8]) -> Option<OcrGenerati
match run_ocr(bytes) {
Ok(Some(text)) => Some(OcrGeneration {
text,
source: "ocr",
source: OcrSource::Ocr,
}),
Ok(None) => None,
Err(OcrError::BinaryMissing) => {
warn!("ocrmypdf not installed; cannot perform OCR");
warn!("ocrmypdf binary not found; OCR unavailable");
None
}
Err(err) => {
warn!(error = ?err, "ocr command failed");
warn!(error = %err, "ocr command failed");
None
}
}
}
struct PdfDocumentMeta {
content_type: Option<String>,
original_name: String,
}
struct OcrGeneration {
text: String,
source: OcrSource,
}
#[derive(Clone, Copy)]
enum OcrSource {
PdfText,
Ocr,
}
impl fmt::Display for OcrSource {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
OcrSource::PdfText => write!(f, "pdf-text"),
OcrSource::Ocr => write!(f, "ocr"),
}
}
}
#[derive(Debug)]
enum OcrError {
BinaryMissing,
Failed(String),
}
impl fmt::Display for OcrError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
OcrError::BinaryMissing => write!(f, "ocrmypdf binary not found"),
OcrError::Failed(msg) => write!(f, "ocr failed: {msg}"),
}
}
}
fn extract_pdf_text(bytes: &[u8]) -> Result<String, String> {
let pdfium = Pdfium::default();
let document = pdfium
@@ -345,21 +468,6 @@ fn extract_pdf_text(bytes: &[u8]) -> Result<String, String> {
Ok(combined)
}
#[derive(Debug)]
enum OcrError {
BinaryMissing,
Failed(String),
}
impl fmt::Display for OcrError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
OcrError::BinaryMissing => write!(f, "ocrmypdf binary not found"),
OcrError::Failed(msg) => write!(f, "ocr failed: {msg}"),
}
}
}
fn run_ocr(bytes: &[u8]) -> Result<Option<String>, OcrError> {
let mut input = NamedTempFile::new().map_err(|err| OcrError::Failed(err.to_string()))?;
input
@@ -408,118 +516,6 @@ fn run_ocr(bytes: &[u8]) -> Result<Option<String>, OcrError> {
}
}
fn persist_ocr_metadata(
state: Arc<AppState>,
context: &OcrContext,
asset_id: Uuid,
s3_key: &str,
source: &'static str,
) -> Result<(), String> {
let tenant_id = context.document.tenant_id;
let mut conn = state
.db_for_tenant(tenant_id)
.map_err(|err| format!("{err:?}"))?;
if let Some(existing_asset) = &context.existing_asset {
diesel::delete(document_assets::table.filter(document_assets::id.eq(existing_asset.id)))
.execute(&mut conn)
.map_err(|err| format!("{err:?}"))?;
}
let new_asset = NewDocumentAsset {
id: asset_id,
document_version_id: context.version.id,
asset_type: OCR_TEXT_ASSET_TYPE.to_string(),
mime_type: "text/plain".to_string(),
metadata: json!({
"generated_at": Utc::now().to_rfc3339(),
"source": source,
}),
cardinality: Some(1),
tenant_id,
};
diesel::insert_into(document_assets::table)
.values(&new_asset)
.on_conflict((
document_assets::document_version_id,
document_assets::asset_type,
))
.do_update()
.set((
document_assets::mime_type.eq(excluded(document_assets::mime_type)),
document_assets::metadata.eq(excluded(document_assets::metadata)),
document_assets::cardinality.eq(excluded(document_assets::cardinality)),
))
.execute(&mut conn)
.map_err(|err| format!("{err:?}"))?;
let existing_object_id: Option<Uuid> = document_asset_objects::table
.filter(document_asset_objects::asset_id.eq(asset_id))
.filter(document_asset_objects::ordinal.eq(1))
.filter(document_asset_objects::tenant_id.eq(tenant_id))
.select(document_asset_objects::id)
.first(&mut conn)
.optional()
.map_err(|err| format!("{err:?}"))?;
let object_id = existing_object_id.unwrap_or_else(Uuid::new_v4);
let new_object = NewDocumentAssetObject {
id: object_id,
asset_id,
ordinal: 1,
s3_key: s3_key.to_string(),
metadata: json!({}),
tenant_id,
};
diesel::insert_into(document_asset_objects::table)
.values(&new_object)
.on_conflict((
document_asset_objects::asset_id,
document_asset_objects::ordinal,
))
.do_update()
.set((
document_asset_objects::s3_key.eq(excluded(document_asset_objects::s3_key)),
document_asset_objects::metadata.eq(excluded(document_asset_objects::metadata)),
))
.execute(&mut conn)
.map_err(|err| format!("{err:?}"))?;
Ok(())
}
fn enqueue_index_job(
state: &AppState,
tenant_id: Uuid,
payload: &OcrPayload,
) -> Result<(), String> {
let mut conn = state
.db_for_tenant(tenant_id)
.map_err(|err| format!("{err:?}"))?;
enqueue_job(
&mut conn,
tenant_id,
JOB_INDEX_DOCUMENT_TEXT,
json!({
"document_id": payload.document_id,
"document_version_id": payload.document_version_id,
}),
None,
)
.map(|_| ())
.map_err(|err| err.to_string())
}
pub fn document_is_pdf(document: &Document) -> bool {
document_meta_is_pdf(&PdfDocumentMeta {
content_type: document.content_type.clone(),
original_name: document.original_name.clone(),
})
}
fn document_meta_is_pdf(meta: &PdfDocumentMeta) -> bool {
if let Some(content_type) = &meta.content_type {
if content_type.eq_ignore_ascii_case("application/pdf") {
@@ -533,3 +529,18 @@ fn document_meta_is_pdf(meta: &PdfDocumentMeta) -> bool {
.map(|ext| ext.eq_ignore_ascii_case("pdf"))
.unwrap_or(false)
}
fn document_is_pdf(document: &Document) -> bool {
document
.content_type
.as_deref()
.map(|mime| mime.eq_ignore_ascii_case("application/pdf"))
.unwrap_or_else(|| {
document
.original_name
.rsplit('.')
.next()
.map(|ext| ext.eq_ignore_ascii_case("pdf"))
.unwrap_or(false)
})
}