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/error.rs
+4 -3
View File
@@ -6,6 +6,7 @@ use axum::{
use serde::Serialize;
use serde_json::Value;
use std::fmt::Display;
use utoipa::ToSchema;
pub type AppResult<T> = Result<T, AppError>;
@@ -61,7 +62,7 @@ impl AppError {
impl IntoResponse for AppError {
fn into_response(self) -> Response {
let status = self.status;
let body = Json(ErrorResponse {
let body = Json(ApiErrorResponse {
error: self.message,
code: self.code,
details: self.details,
@@ -70,8 +71,8 @@ impl IntoResponse for AppError {
}
}
#[derive(Serialize)]
struct ErrorResponse {
#[derive(Serialize, ToSchema)]
pub struct ApiErrorResponse {
error: String,
#[serde(skip_serializing_if = "Option::is_none")]
code: Option<String>,
backend/src/workers/analyze.rs
+79 -107
View File
@@ -1,23 +1,24 @@
use std::{collections::HashSet, sync::Arc, time::Duration};
use std::collections::HashSet;
use std::sync::Arc;
use async_trait::async_trait;
use diesel::prelude::*;
use serde::Deserialize;
use serde_json::json;
use tokio::task;
use tracing::{error, warn};
use uuid::Uuid;
use super::ocr::{document_is_pdf, OCR_TEXT_ASSET_TYPE};
use crate::{
jobs::{enqueue_job, JOB_ANALYZE_DOCUMENT, JOB_GENERATE_OCR_TEXT, JOB_GENERATE_THUMBNAILS},
models::{Document, DocumentAsset, DocumentVersion},
schema::{document_assets, document_versions, documents},
state::AppState,
storage::TenantStorage,
jobs::JOB_ANALYZE_DOCUMENT, models::Document, state::AppState, storage::TenantStorage,
};
use super::{JobExecution, JobHandler};
use super::{
index::IndexDocumentTask,
job_execution_from_task_error,
ocr::{GenerateOcrTask, OCR_TEXT_ASSET_TYPE},
taskflow::{
document::DocumentVersionTaskContext, BoxedTask, TaskExecutor, TaskPlanner, TaskResult,
},
thumbnails::GenerateThumbnailsTask,
JobExecution, JobHandler,
};
#[derive(Debug, Deserialize)]
struct AnalyzePayload {
@@ -45,7 +46,7 @@ impl JobHandler for AnalyzeDocumentJob {
&self,
state: Arc<AppState>,
job: crate::models::Job,
_storage: TenantStorage,
storage: TenantStorage,
) -> JobExecution {
let payload: AnalyzePayload = match serde_json::from_value(job.payload.clone()) {
Ok(payload) => payload,
@@ -56,114 +57,70 @@ impl JobHandler for AnalyzeDocumentJob {
}
};
let state_clone = state.clone();
let tenant_id = job.tenant_id;
match task::spawn_blocking(move || analyze_document(state_clone, tenant_id, payload)).await
{
Ok(Ok(execution)) => execution,
Ok(Err(err)) => {
warn!(job_id = %job.id, error = %err, "analyze job will retry");
JobExecution::Retry {
delay: Duration::from_secs(30),
error: err,
}
}
Err(join_err) => {
error!(job_id = %job.id, error = %join_err, "analyze task panicked");
JobExecution::Retry {
delay: Duration::from_secs(60),
error: format!("worker panicked: {join_err}"),
}
}
let mut context = DocumentVersionTaskContext::new(
job.id,
JOB_ANALYZE_DOCUMENT,
job.tenant_id,
payload.document_id,
payload.document_version_id,
payload.force,
state.config.worker_max_document_bytes,
state.clone(),
storage,
);
let planner = AnalyzePlanner::new(payload.force, state.clone());
match TaskExecutor::run(&planner, &mut context).await {
Ok(()) => JobExecution::Success,
Err(err) => job_execution_from_task_error(err),
}
}
}
fn analyze_document(
struct AnalyzePlanner {
force: bool,
state: Arc<AppState>,
tenant_id: Uuid,
payload: AnalyzePayload,
) -> Result<JobExecution, 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!(
"failed to load document_version {} for tenant {}: {err:?}",
payload.document_version_id, tenant_id
)
})?;
if version.document_id != payload.document_id {
return Err("document/version mismatch".into());
impl AnalyzePlanner {
fn new(force: bool, state: Arc<AppState>) -> Self {
Self { force, state }
}
}
let document: Document = documents::table
.find(payload.document_id)
.first(&mut conn)
.map_err(|err| {
format!(
"failed to load document {} for tenant {}: {err:?}",
payload.document_id, tenant_id
)
})?;
#[async_trait]
impl TaskPlanner<DocumentVersionTaskContext> for AnalyzePlanner {
async fn plan(
&self,
ctx: &mut DocumentVersionTaskContext,
) -> TaskResult<Vec<BoxedTask<DocumentVersionTaskContext>>> {
let document = ctx.document().await?.clone();
let mut tasks: Vec<BoxedTask<DocumentVersionTaskContext>> = Vec::new();
let tenant_id = document.tenant_id;
let (supported, _reason) = determine_thumbnail_support(&document);
let ocr_supported = document_is_pdf(&document);
let existing_ocr: Option<DocumentAsset> = document_assets::table
.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))
.first(&mut conn)
.optional()
.map_err(|err| format!("{err:?}"))?;
let skip_ocr = existing_ocr.is_some() && !payload.force;
if supported {
let enqueue_result = enqueue_job(
&mut conn,
tenant_id,
JOB_GENERATE_THUMBNAILS,
json!({
"document_id": payload.document_id,
"document_version_id": payload.document_version_id,
"force": payload.force,
}),
None,
);
if let Err(err) = enqueue_result {
return Err(err.to_string());
let (thumbnail_supported, _) = determine_thumbnail_support(&document);
if thumbnail_supported {
tasks.push(Box::new(GenerateThumbnailsTask::new(self.force)));
}
}
if ocr_supported && !skip_ocr {
let enqueue_result = enqueue_job(
&mut conn,
tenant_id,
JOB_GENERATE_OCR_TEXT,
json!({
"document_id": payload.document_id,
"document_version_id": payload.document_version_id,
"force": payload.force,
}),
None,
);
let existing_ocr = ctx.asset(OCR_TEXT_ASSET_TYPE).await?.is_some();
let mut should_index = existing_ocr;
if let Err(err) = enqueue_result {
return Err(err.to_string());
if document_supports_ocr(&document) {
if self.force || !existing_ocr {
tasks.push(Box::new(GenerateOcrTask::new(
self.force,
self.state.clone(),
)));
should_index = true;
}
}
}
Ok(JobExecution::Success)
if should_index {
tasks.push(Box::new(IndexDocumentTask::new()));
}
Ok(tasks)
}
}
pub(crate) fn determine_thumbnail_support(document: &Document) -> (bool, Option<String>) {
@@ -204,3 +161,18 @@ pub(crate) fn determine_thumbnail_support(document: &Document) -> (bool, Option<
Some("content type not supported for thumbnails".into()),
)
}
fn document_supports_ocr(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)
})
}
backend/src/workers/common.rs
+4 -9
View File
@@ -7,10 +7,9 @@ use crate::models::{Document, DocumentAsset, DocumentAssetObject, DocumentVersio
use crate::schema::{document_asset_objects, document_assets, document_versions, documents};
use crate::state::AppState;
pub(crate) struct DocumentVersionContext {
pub(crate) struct LoadedDocumentVersion {
pub document: Document,
pub version: DocumentVersion,
pub tenant_id: Uuid,
}
pub(crate) fn load_document_version(
@@ -18,7 +17,7 @@ pub(crate) fn load_document_version(
tenant_id: Uuid,
document_id: Uuid,
version_id: Uuid,
) -> Result<DocumentVersionContext, String> {
) -> Result<LoadedDocumentVersion, String> {
let mut conn = state
.db_for_tenant(tenant_id)
.map_err(|err| format!("{err:?}"))?;
@@ -37,14 +36,10 @@ pub(crate) fn load_document_version(
.first(&mut conn)
.map_err(|err| format!("{err:?}"))?;
Ok(DocumentVersionContext {
document,
version,
tenant_id,
})
Ok(LoadedDocumentVersion { document, version })
}
pub(crate) struct LoadedAsset {
pub struct LoadedAsset {
pub asset: DocumentAsset,
pub objects: Vec<DocumentAssetObject>,
}
backend/src/workers/index.rs
+108 -152
View File
@@ -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,
})
}
backend/src/workers/mod.rs
+9 -32
View File
@@ -11,12 +11,14 @@ use crate::{
state::AppState,
storage::TenantStorage,
};
use taskflow::TaskError;
pub mod analyze;
pub mod common;
pub mod index;
pub mod ocr;
pub mod purge;
pub mod taskflow;
pub mod tenants;
pub mod thumbnails;
@@ -156,6 +158,13 @@ pub fn default_handlers() -> Vec<Arc<dyn JobHandler>> {
]
}
pub(crate) fn job_execution_from_task_error(error: TaskError) -> JobExecution {
match error {
TaskError::Fail { error } => JobExecution::Failed { error },
TaskError::Retry { delay, error } => JobExecution::Retry { delay, error },
}
}
pub(crate) fn check_worker_document_limit(
size_bytes: i64,
limit_bytes: u64,
@@ -187,35 +196,3 @@ pub(crate) async fn fetch_version_object(
.await
.map_err(FetchVersionError::Storage)
}
pub(crate) fn handle_fetch_error(
job: &crate::models::Job,
err: FetchVersionError,
message: &str,
) -> JobExecution {
match err {
FetchVersionError::TooLarge { size, limit } => {
warn!(
job_id = %job.id,
size_bytes = size,
limit_bytes = limit,
"document exceeds worker size limit"
);
JobExecution::Failed {
error: format!("document size {size} bytes exceeds worker limit of {limit} bytes"),
}
}
FetchVersionError::Storage(err) => {
warn!(
job_id = %job.id,
error = %err,
context = message,
"failed to fetch object for worker"
);
JobExecution::Retry {
delay: Duration::from_secs(30),
error: err.to_string(),
}
}
}
}
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)
})
}
backend/src/workers/purge.rs
+111 -58
View File
@@ -6,7 +6,6 @@ use async_trait::async_trait;
use diesel::prelude::*;
use diesel::result::Error as DieselError;
use serde::Deserialize;
use tracing::{error, warn};
use uuid::Uuid;
use crate::jobs::JOB_PURGE_DOCUMENT;
@@ -15,7 +14,11 @@ use crate::schema::{document_asset_objects, document_assets, document_versions};
use crate::state::AppState;
use crate::storage::TenantStorage;
use super::{JobExecution, JobHandler};
use super::{
job_execution_from_task_error,
taskflow::{BoxedTask, Task, TaskContext, TaskError, TaskExecutor, TaskPlanner, TaskResult},
JobExecution, JobHandler,
};
#[derive(Debug, Deserialize)]
struct PurgeDocumentPayload {
@@ -58,70 +61,120 @@ impl JobHandler for PurgeDocumentJob {
}
};
let tenant_id = job.tenant_id;
let document_id = payload.document_id;
let state_for_prepare = state.clone();
let mut context = PurgeTaskContext::new(
job.id,
JOB_PURGE_DOCUMENT,
job.tenant_id,
payload.document_id,
state.clone(),
storage,
);
let planner = PurgePlanner;
match TaskExecutor::run(&planner, &mut context).await {
Ok(()) => JobExecution::Success,
Err(err) => job_execution_from_task_error(err),
}
}
}
struct PurgeTaskContext {
job_id: Uuid,
job_type: &'static str,
tenant_id: Uuid,
document_id: Uuid,
state: Arc<AppState>,
storage: TenantStorage,
}
impl PurgeTaskContext {
fn new(
job_id: Uuid,
job_type: &'static str,
tenant_id: Uuid,
document_id: Uuid,
state: Arc<AppState>,
storage: TenantStorage,
) -> Self {
Self {
job_id,
job_type,
tenant_id,
document_id,
state,
storage,
}
}
}
impl TaskContext for PurgeTaskContext {
fn job_id(&self) -> Uuid {
self.job_id
}
fn job_type(&self) -> &'static str {
self.job_type
}
}
struct PurgePlanner;
#[async_trait]
impl TaskPlanner<PurgeTaskContext> for PurgePlanner {
async fn plan(
&self,
_ctx: &mut PurgeTaskContext,
) -> TaskResult<Vec<BoxedTask<PurgeTaskContext>>> {
Ok(vec![Box::new(PurgeTask)])
}
}
struct PurgeTask;
#[async_trait]
impl Task<PurgeTaskContext> for PurgeTask {
fn name(&self) -> &'static str {
"purge-document"
}
async fn execute(&self, ctx: &mut PurgeTaskContext) -> TaskResult<()> {
let tenant_id = ctx.tenant_id;
let document_id = ctx.document_id;
let state = ctx.state.clone();
let preparation = tokio::task::spawn_blocking(move || {
prepare_purge_context(state_for_prepare, tenant_id, document_id)
prepare_purge_context(state, tenant_id, document_id)
})
.await;
.await
.map_err(|err| {
TaskError::retry(
Duration::from_secs(60),
format!("purge preparation panicked: {err}"),
)
})?;
let context = match preparation {
Ok(Ok(Some(ctx))) => ctx,
Ok(Ok(None)) => {
// Document already gone or restored; nothing to do.
return JobExecution::Success;
}
Ok(Err(err)) => {
warn!(job_id = %job.id, error = %err, "purge preparation failed");
return JobExecution::Retry {
delay: Duration::from_secs(30),
error: err,
};
}
Err(join_err) => {
error!(job_id = %job.id, error = %join_err, "purge preparation task panicked");
return JobExecution::Retry {
delay: Duration::from_secs(60),
error: format!("purge preparation panicked: {join_err}"),
};
}
let Some(context) =
preparation.map_err(|err| TaskError::retry(Duration::from_secs(30), err))?
else {
return Ok(());
};
if let Err(err) = delete_storage_objects(&storage, &context).await {
warn!(job_id = %job.id, error = %err, "failed to delete storage objects for purge");
return JobExecution::Retry {
delay: Duration::from_secs(30),
error: err,
};
}
delete_storage_objects(&ctx.storage, &context)
.await
.map_err(|err| TaskError::retry(Duration::from_secs(30), err))?;
let PurgeContext { document_id, .. } = context;
let state_for_finalize = state.clone();
let state = ctx.state.clone();
tokio::task::spawn_blocking(move || finalize_purge(state, tenant_id, context.document_id))
.await
.map_err(|err| {
TaskError::retry(
Duration::from_secs(60),
format!("purge finalize panicked: {err}"),
)
})?
.map_err(|err| TaskError::retry(Duration::from_secs(30), err))?;
let finalize = tokio::task::spawn_blocking(move || {
finalize_purge(state_for_finalize, tenant_id, document_id)
})
.await;
match finalize {
Ok(Ok(())) => JobExecution::Success,
Ok(Err(err)) => {
warn!(job_id = %job.id, error = %err, "failed to finalize purge");
JobExecution::Retry {
delay: Duration::from_secs(30),
error: err,
}
}
Err(join_err) => {
error!(job_id = %job.id, error = %join_err, "purge finalize task panicked");
JobExecution::Retry {
delay: Duration::from_secs(60),
error: format!("purge finalize panicked: {join_err}"),
}
}
}
Ok(())
}
}
backend/src/workers/taskflow/document.rs
+197
View File
@@ -0,0 +1,197 @@
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio::task;
use uuid::Uuid;
use crate::models::{Document, DocumentVersion};
use crate::state::AppState;
use crate::storage::TenantStorage;
use crate::workers::common::{load_document_version, load_version_assets, LoadedAsset};
use crate::workers::{fetch_version_object, FetchVersionError};
use super::{TaskContext, TaskError, TaskResult};
const BLOCKING_RETRY_DELAY: Duration = Duration::from_secs(60);
const DEFAULT_RETRY_DELAY: Duration = Duration::from_secs(30);
pub struct DocumentVersionTaskContext {
job_id: Uuid,
job_type: &'static str,
tenant_id: Uuid,
document_id: Uuid,
document_version_id: Uuid,
force: bool,
max_document_bytes: u64,
state: Arc<AppState>,
storage: TenantStorage,
document: Option<Document>,
version: Option<DocumentVersion>,
assets: Option<HashMap<String, LoadedAsset>>, // keyed by asset_type
object_bytes: Option<Vec<u8>>,
}
impl DocumentVersionTaskContext {
#[allow(clippy::too_many_arguments)]
pub fn new(
job_id: Uuid,
job_type: &'static str,
tenant_id: Uuid,
document_id: Uuid,
document_version_id: Uuid,
force: bool,
max_document_bytes: u64,
state: Arc<AppState>,
storage: TenantStorage,
) -> Self {
Self {
job_id,
job_type,
tenant_id,
document_id,
document_version_id,
force,
max_document_bytes,
state,
storage,
document: None,
version: None,
assets: None,
object_bytes: None,
}
}
pub fn tenant_id(&self) -> Uuid {
self.tenant_id
}
pub fn document_id(&self) -> Uuid {
self.document_id
}
pub fn version_id(&self) -> Uuid {
self.document_version_id
}
pub fn force(&self) -> bool {
self.force
}
pub fn storage(&self) -> &TenantStorage {
&self.storage
}
pub fn invalidate_asset_cache(&mut self) {
self.assets = None;
}
pub fn state(&self) -> &Arc<AppState> {
&self.state
}
pub fn max_document_bytes(&self) -> u64 {
self.max_document_bytes
}
pub async fn document(&mut self) -> TaskResult<&Document> {
self.ensure_document_loaded().await?;
Ok(self.document.as_ref().expect("document hydrated"))
}
pub async fn version(&mut self) -> TaskResult<&DocumentVersion> {
self.ensure_document_loaded().await?;
Ok(self.version.as_ref().expect("version hydrated"))
}
pub async fn assets(&mut self) -> TaskResult<&HashMap<String, LoadedAsset>> {
if self.assets.is_none() {
let tenant_id = self.tenant_id;
let version_id = self.document_version_id;
let state = self.state.clone();
let result = task::spawn_blocking(move || {
let mut conn = state
.db_for_tenant(tenant_id)
.map_err(|err| format!("failed to scope tenant connection: {err:?}"))?;
load_version_assets(&mut conn, tenant_id, version_id, &[])
})
.await
.map_err(|err| {
TaskError::retry(
BLOCKING_RETRY_DELAY,
format!("asset load task panicked: {err}"),
)
})?
.map_err(|err| TaskError::retry(DEFAULT_RETRY_DELAY, err))?;
self.assets = Some(result);
}
Ok(self.assets.as_ref().expect("asset map hydrated"))
}
pub async fn asset(&mut self, asset_type: &str) -> TaskResult<Option<&LoadedAsset>> {
let assets = self.assets().await?;
Ok(assets.get(asset_type))
}
pub async fn buffered_object(&mut self) -> TaskResult<&[u8]> {
if self.object_bytes.is_some() {
return Ok(self.object_bytes.as_deref().expect("bytes present"));
}
let version = self.version().await?.clone();
let bytes = fetch_version_object(
&version,
&self.storage,
&version.s3_key,
self.max_document_bytes,
)
.await
.map_err(|err| match err {
FetchVersionError::TooLarge { size, limit } => TaskError::fail(format!(
"document size {size} bytes exceeds worker limit of {limit} bytes"
)),
FetchVersionError::Storage(err) => TaskError::retry(
DEFAULT_RETRY_DELAY,
format!("failed to fetch object: {err}"),
),
})?;
self.object_bytes = Some(bytes);
Ok(self.object_bytes.as_deref().expect("bytes hydrated"))
}
async fn ensure_document_loaded(&mut self) -> TaskResult<()> {
if self.document.is_some() && self.version.is_some() {
return Ok(());
}
let tenant_id = self.tenant_id;
let document_id = self.document_id;
let version_id = self.document_version_id;
let state = self.state.clone();
let loaded = task::spawn_blocking(move || {
load_document_version(state.as_ref(), tenant_id, document_id, version_id)
})
.await
.map_err(|err| {
TaskError::retry(
BLOCKING_RETRY_DELAY,
format!("document load task panicked: {err}"),
)
})?
.map_err(|err| TaskError::fail(format!("failed to load document context: {err}")))?;
self.document = Some(loaded.document);
self.version = Some(loaded.version);
Ok(())
}
}
impl TaskContext for DocumentVersionTaskContext {
fn job_id(&self) -> Uuid {
self.job_id
}
fn job_type(&self) -> &'static str {
self.job_type
}
}
backend/src/workers/taskflow/mod.rs
+85
View File
@@ -0,0 +1,85 @@
use std::time::Duration;
use async_trait::async_trait;
use thiserror::Error;
use tracing::info;
use uuid::Uuid;
pub mod document;
pub type TaskResult<T> = Result<T, TaskError>;
#[derive(Debug, Error)]
pub enum TaskError {
#[error("{error}")]
Fail { error: String },
#[error("{error}")]
Retry { delay: Duration, error: String },
}
impl TaskError {
pub fn fail(error: impl Into<String>) -> Self {
Self::Fail {
error: error.into(),
}
}
pub fn retry(delay: Duration, error: impl Into<String>) -> Self {
Self::Retry {
delay,
error: error.into(),
}
}
}
pub trait TaskContext: Send + Sync {
fn job_id(&self) -> Uuid;
fn job_type(&self) -> &'static str;
}
#[async_trait]
pub trait Task<Ctx>: Send + Sync
where
Ctx: TaskContext,
{
fn name(&self) -> &'static str;
async fn execute(&self, ctx: &mut Ctx) -> TaskResult<()>;
}
pub type BoxedTask<Ctx> = Box<dyn Task<Ctx> + Send + Sync>;
#[async_trait]
pub trait TaskPlanner<Ctx>: Send + Sync
where
Ctx: TaskContext,
{
async fn plan(&self, ctx: &mut Ctx) -> TaskResult<Vec<BoxedTask<Ctx>>>;
}
pub struct TaskExecutor;
impl TaskExecutor {
pub async fn run<P, C>(planner: &P, ctx: &mut C) -> TaskResult<()>
where
P: TaskPlanner<C>,
C: TaskContext,
{
let tasks = planner.plan(ctx).await?;
for task in tasks {
info!(
job_id = %ctx.job_id(),
job_type = ctx.job_type(),
task = task.name(),
"starting job task"
);
task.execute(ctx).await?;
info!(
job_id = %ctx.job_id(),
job_type = ctx.job_type(),
task = task.name(),
"finished job task"
);
}
Ok(())
}
}
backend/src/workers/tenants.rs
+155 -155
View File
@@ -1,4 +1,5 @@
use std::sync::Arc;
use std::time::Duration;
use async_trait::async_trait;
use chrono::Utc;
@@ -18,7 +19,11 @@ use crate::models::{NewUserMembership, TenantStatus};
use crate::schema::{tenants, user_memberships};
use crate::state::AppState;
use crate::tenants::TenantRepository;
use crate::workers::{JobExecution, JobHandler};
use crate::workers::{
job_execution_from_task_error,
taskflow::{BoxedTask, Task, TaskContext, TaskError, TaskExecutor, TaskPlanner, TaskResult},
JobExecution, JobHandler,
};
pub struct ProvisionTenantJob;
@@ -40,77 +45,120 @@ impl JobHandler for ProvisionTenantJob {
job: crate::models::Job,
_storage: crate::storage::TenantStorage,
) -> JobExecution {
let mut conn = match state.db_unscoped() {
Ok(conn) => conn,
Err(err) => {
warn!(job_id = %job.id, error = ?err, "failed to get connection for tenant provisioning");
return JobExecution::Retry {
delay: std::time::Duration::from_secs(30),
error: "database connection unavailable".into(),
};
}
};
let members = ProvisionPayload::from_job(&job).unwrap_or_default();
let mut context = ProvisionContext::new(
job.id,
JOB_PROVISION_TENANT,
job.tenant_id,
state.clone(),
members,
);
let tenant = match TenantRepository::get_by_id(&mut conn, job.tenant_id) {
Ok(tenant) => tenant,
Err(err) => {
warn!(job_id = %job.id, error = ?err, "tenant not found for provisioning");
return JobExecution::Failed {
error: "tenant not found".into(),
};
}
};
let planner = ProvisionPlanner;
match TaskExecutor::run(&planner, &mut context).await {
Ok(()) => JobExecution::Success,
Err(err) => job_execution_from_task_error(err),
}
}
}
struct ProvisionContext {
job_id: Uuid,
job_type: &'static str,
tenant_id: Uuid,
state: Arc<AppState>,
members: Vec<Uuid>,
}
impl ProvisionContext {
fn new(
job_id: Uuid,
job_type: &'static str,
tenant_id: Uuid,
state: Arc<AppState>,
members: Vec<Uuid>,
) -> Self {
Self {
job_id,
job_type,
tenant_id,
state,
members,
}
}
}
impl TaskContext for ProvisionContext {
fn job_id(&self) -> Uuid {
self.job_id
}
fn job_type(&self) -> &'static str {
self.job_type
}
}
struct ProvisionPlanner;
#[async_trait]
impl TaskPlanner<ProvisionContext> for ProvisionPlanner {
async fn plan(
&self,
_ctx: &mut ProvisionContext,
) -> TaskResult<Vec<BoxedTask<ProvisionContext>>> {
Ok(vec![Box::new(ProvisionTask)])
}
}
struct ProvisionTask;
#[async_trait]
impl Task<ProvisionContext> for ProvisionTask {
fn name(&self) -> &'static str {
"provision-tenant"
}
async fn execute(&self, ctx: &mut ProvisionContext) -> TaskResult<()> {
let mut conn = ctx
.state
.db_unscoped()
.map_err(|err| TaskError::retry(Duration::from_secs(30), format!("{err:?}")))?;
let tenant = TenantRepository::get_by_id(&mut conn, ctx.tenant_id).map_err(|err| {
TaskError::fail(format!("tenant not found for provisioning: {err:?}"))
})?;
drop(conn);
let mut conn = match state.db_for_tenant(tenant.id) {
Ok(conn) => conn,
Err(err) => {
warn!(job_id = %job.id, error = ?err, "failed to scope connection for tenant provisioning");
return JobExecution::Retry {
delay: std::time::Duration::from_secs(30),
error: "tenant connection unavailable".into(),
};
}
};
let mut conn = ctx
.state
.db_for_tenant(tenant.id)
.map_err(|err| TaskError::retry(Duration::from_secs(30), format!("{err:?}")))?;
if tenant.status == TenantStatus::Active {
warn!(
job_id = %job.id,
job_id = %ctx.job_id(),
tenant_id = %tenant.id,
"tenant already active; skipping provisioning"
);
return JobExecution::Success;
return Ok(());
}
if tenant.status != TenantStatus::Creating {
warn!(
job_id = %job.id,
tenant_id = %tenant.id,
status = %tenant.status.as_str(),
"tenant not in creating state; skipping provisioning"
);
return JobExecution::Failed {
error: format!(
"tenant status '{}' not eligible for provisioning",
tenant.status.as_str()
),
};
return Err(TaskError::fail(format!(
"tenant status '{}' not eligible for provisioning",
tenant.status.as_str()
)));
}
let endpoint = match &state.config.quickwit_endpoint {
Some(endpoint) => endpoint.trim_end_matches('/').to_owned(),
None => {
warn!(
job_id = %job.id,
tenant_id = %tenant.id,
"quickwit endpoint not configured; retrying"
);
return JobExecution::Retry {
delay: std::time::Duration::from_secs(30),
error: "quickwit endpoint not configured".into(),
};
}
};
let endpoint = ctx
.state
.config
.quickwit_endpoint
.as_ref()
.map(|value| value.trim_end_matches('/').to_owned())
.ok_or_else(|| {
TaskError::retry(Duration::from_secs(30), "quickwit endpoint not configured")
})?;
let index_id = tenant
.quickwit_index
@@ -119,117 +167,69 @@ impl JobHandler for ProvisionTenantJob {
.unwrap_or_else(|| format!("documents-{}", tenant.id));
let client = Client::new();
if let Err(err) = ensure_quickwit_index(&client, &endpoint, &index_id).await {
warn!(
job_id = %job.id,
tenant_id = %tenant.id,
error = %err,
"failed to ensure quickwit index"
);
return JobExecution::Retry {
delay: std::time::Duration::from_secs(30),
error: err.to_string(),
};
}
ensure_quickwit_index(&client, &endpoint, &index_id)
.await
.map_err(|err| TaskError::retry(Duration::from_secs(30), err.to_string()))?;
let owner_capability_set_id =
match ensure_capability_set(&mut conn, tenant.id, owner_capabilities()) {
Ok(set) => set.id,
Err(err) => {
warn!(
job_id = %job.id,
tenant_id = %tenant.id,
error = ?err,
"failed to ensure owner capability set during provisioning"
);
return JobExecution::Retry {
delay: std::time::Duration::from_secs(30),
error: "owner capability set unavailable".into(),
};
}
ensure_capability_set(&mut conn, tenant.id, owner_capabilities())
.map_err(|_| {
TaskError::retry(Duration::from_secs(30), "owner capability set unavailable")
})?
.id;
ensure_capability_set(&mut conn, tenant.id, user_capabilities()).map_err(|_| {
TaskError::retry(Duration::from_secs(30), "user capability set unavailable")
})?;
ensure_capability_set(&mut conn, tenant.id, readonly_capabilities()).map_err(|_| {
TaskError::retry(
Duration::from_secs(30),
"readonly capability set unavailable",
)
})?;
ensure_capability_set(&mut conn, tenant.id, webdav_capabilities()).map_err(|_| {
TaskError::retry(Duration::from_secs(30), "webdav capability set unavailable")
})?;
for member in &ctx.members {
let new_membership = NewUserMembership {
id: Uuid::new_v4(),
user_id: *member,
tenant_id: tenant.id,
capability_set_id: Some(owner_capability_set_id),
};
if let Err(err) = ensure_capability_set(&mut conn, tenant.id, user_capabilities()) {
warn!(
job_id = %job.id,
tenant_id = %tenant.id,
error = ?err,
"failed to ensure user capability set during provisioning"
);
return JobExecution::Retry {
delay: std::time::Duration::from_secs(30),
error: "user capability set unavailable".into(),
};
}
if let Err(err) = ensure_capability_set(&mut conn, tenant.id, readonly_capabilities()) {
warn!(
job_id = %job.id,
tenant_id = %tenant.id,
error = ?err,
"failed to ensure readonly capability set during provisioning"
);
return JobExecution::Retry {
delay: std::time::Duration::from_secs(30),
error: "readonly capability set unavailable".into(),
};
}
if let Err(err) = ensure_capability_set(&mut conn, tenant.id, webdav_capabilities()) {
warn!(
job_id = %job.id,
tenant_id = %tenant.id,
error = ?err,
"failed to ensure webdav capability set during provisioning"
);
return JobExecution::Retry {
delay: std::time::Duration::from_secs(30),
error: "webdav capability set unavailable".into(),
};
}
if let Some(members) = ProvisionPayload::from_job(&job) {
for member in members {
let new_membership = NewUserMembership {
id: Uuid::new_v4(),
user_id: member,
tenant_id: tenant.id,
capability_set_id: Some(owner_capability_set_id),
};
if let Err(err) = diesel::insert_into(user_memberships::table)
.values(&new_membership)
.on_conflict((user_memberships::user_id, user_memberships::tenant_id))
.do_nothing()
.execute(&mut conn)
{
warn!(
job_id = %job.id,
tenant_id = %tenant.id,
user_id = %member,
error = %err,
"failed to assign initial membership"
);
}
if let Err(err) = diesel::insert_into(user_memberships::table)
.values(&new_membership)
.on_conflict((user_memberships::user_id, user_memberships::tenant_id))
.do_nothing()
.execute(&mut conn)
{
warn!(
job_id = %ctx.job_id(),
tenant_id = %tenant.id,
user_id = %member,
error = %err,
"failed to assign initial membership"
);
}
}
if let Err(err) = diesel::update(tenants::table.find(tenant.id))
diesel::update(tenants::table.find(tenant.id))
.set((
tenants::status.eq(TenantStatus::Active),
tenants::quickwit_index.eq(Some(index_id)),
tenants::updated_at.eq(Utc::now().naive_utc()),
))
.execute(&mut conn)
{
warn!(job_id = %job.id, error = %err, "failed to activate tenant");
return JobExecution::Retry {
delay: std::time::Duration::from_secs(30),
error: format!("failed to update tenant status: {err}"),
};
}
.map_err(|err| {
TaskError::retry(
Duration::from_secs(30),
format!("failed to update tenant status: {err}"),
)
})?;
JobExecution::Success
Ok(())
}
}
backend/src/workers/thumbnails.rs
+282 -370
View File
@@ -8,7 +8,7 @@ use pdfium_render::prelude::*;
use serde::Deserialize;
use serde_json::{json, Map, Value};
use tokio::task;
use tracing::{error, info, warn};
use tracing::{info, warn};
use uuid::Uuid;
use crate::{
@@ -27,16 +27,20 @@ use crate::{
use super::{
analyze::determine_thumbnail_support,
common::{load_document_version, load_version_assets},
fetch_version_object, handle_fetch_error, JobExecution, JobHandler,
job_execution_from_task_error,
taskflow::{
document::DocumentVersionTaskContext, BoxedTask, Task, TaskContext, TaskError,
TaskExecutor, TaskPlanner, TaskResult,
},
JobExecution, JobHandler,
};
const THUMBNAIL_WIDTH: u32 = 512;
const THUMBNAIL_HEIGHT: u32 = 512;
pub const THUMBNAIL_WIDTH: u32 = 512;
pub const THUMBNAIL_HEIGHT: u32 = 512;
const PREVIEW_WIDTH: u32 = THUMBNAIL_WIDTH * 4;
const PREVIEW_HEIGHT: u32 = THUMBNAIL_HEIGHT * 4;
const THUMBNAIL_ASSET_TYPE: &str = "thumbnail";
const PREVIEW_ASSET_TYPE: &str = "preview";
pub const THUMBNAIL_ASSET_TYPE: &str = "thumbnail";
pub const PREVIEW_ASSET_TYPE: &str = "preview";
#[derive(Debug, Deserialize, Clone)]
struct ThumbnailPayload {
@@ -75,268 +79,109 @@ impl JobHandler for GenerateThumbnailsJob {
}
};
let state_clone = state.clone();
let tenant_id = job.tenant_id;
let payload_for_context = payload.clone();
let initial = match task::spawn_blocking(move || {
load_thumbnail_context(state_clone, tenant_id, payload_for_context)
})
.await
{
Ok(Ok(ctx)) => ctx,
Ok(Err(err)) => {
warn!(job_id = %job.id, error = %err, "thumbnail job will retry");
return JobExecution::Retry {
delay: Duration::from_secs(30),
error: err,
};
}
Err(join_err) => {
error!(job_id = %job.id, error = %join_err, "thumbnail 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_THUMBNAILS,
job.tenant_id,
payload.document_id,
payload.document_version_id,
payload.force,
state.config.worker_max_document_bytes,
state.clone(),
storage,
);
if initial.skip {
info!(job_id = %job.id, "thumbnails already exist; skipping");
return JobExecution::Success;
let planner = ThumbnailPlanner {
force: payload.force,
};
match TaskExecutor::run(&planner, &mut context).await {
Ok(()) => JobExecution::Success,
Err(err) => job_execution_from_task_error(err),
}
}
}
struct ThumbnailPlanner {
force: bool,
}
#[async_trait]
impl TaskPlanner<DocumentVersionTaskContext> for ThumbnailPlanner {
async fn plan(
&self,
_ctx: &mut DocumentVersionTaskContext,
) -> TaskResult<Vec<BoxedTask<DocumentVersionTaskContext>>> {
Ok(vec![Box::new(GenerateThumbnailsTask::new(self.force))])
}
}
pub struct GenerateThumbnailsTask {
force: bool,
}
impl GenerateThumbnailsTask {
pub fn new(force: bool) -> Self {
Self { force }
}
}
#[async_trait]
impl Task<DocumentVersionTaskContext> for GenerateThumbnailsTask {
fn name(&self) -> &'static str {
"generate-thumbnails"
}
async fn execute(&self, ctx: &mut DocumentVersionTaskContext) -> TaskResult<()> {
let context = build_thumbnail_context(ctx, self.force).await?;
if context.skip {
info!(job_id = %ctx.job_id(), "thumbnails already exist; skipping");
return Ok(());
}
let bytes = match fetch_version_object(
&initial.version,
&storage,
&initial.version.s3_key,
state.config.worker_max_document_bytes,
)
.await
{
Ok(bytes) => bytes,
Err(err) => return handle_fetch_error(&job, err, "thumbnail fetch failed; will retry"),
};
let generation = match generate_preview_and_thumbnail(&initial.document, &bytes) {
Ok(result) => result,
Err(err) => {
return JobExecution::Failed { error: err };
}
};
let bytes = ctx.buffered_object().await?;
let generation =
generate_preview_and_thumbnail(&context.document, bytes).map_err(TaskError::fail)?;
if let Some(page_count) = generation.page_count {
let state_clone = state.clone();
let document_id = initial.document.id;
let version_id = initial.version.id;
let tenant_id = initial.tenant_id;
match task::spawn_blocking(move || {
persist_document_page_count(
state_clone,
tenant_id,
document_id,
version_id,
page_count,
let state = ctx.state().clone();
let tenant_id = context.tenant_id;
let document_id = context.document.id;
let version_id = context.version.id;
task::spawn_blocking(move || {
persist_document_page_count(state, tenant_id, document_id, version_id, page_count)
})
.await
.map_err(|err| {
TaskError::retry(
Duration::from_secs(60),
format!("page count task panicked: {err}"),
)
})
.await
{
Ok(Ok(())) => {}
Ok(Err(err)) => {
warn!(
job_id = %job.id,
document_id = %document_id,
version_id = %version_id,
error = %err,
"failed to update document page count metadata; retrying"
);
return JobExecution::Retry {
delay: Duration::from_secs(30),
error: err,
};
}
Err(join_err) => {
error!(
job_id = %job.id,
document_id = %document_id,
version_id = %version_id,
error = %join_err,
"page count metadata task panicked"
);
return JobExecution::Retry {
delay: Duration::from_secs(60),
error: format!("metadata panic: {join_err}"),
};
}
}
})?
.map_err(|err| TaskError::retry(Duration::from_secs(30), err))?;
}
if let Some(existing_preview) = &initial.existing_preview {
for object in &initial.existing_preview_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 preview object"
);
}
}
let tenant_id = initial.document.tenant_id;
let asset_id = existing_preview.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 preview metadata after deletion"
);
}
Err(join_err) => {
warn!(
job_id = %job.id,
error = %join_err,
asset_id = %asset_id,
"failed to remove preview metadata: task panicked"
);
}
}
}
if let Some(existing_thumbnail) = &initial.existing_thumbnail {
for object in &initial.existing_thumbnail_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 thumbnail object"
);
}
}
let tenant_id = initial.document.tenant_id;
let asset_id = existing_thumbnail.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 thumbnail metadata after deletion"
);
}
Err(join_err) => {
warn!(
job_id = %job.id,
error = %join_err,
asset_id = %asset_id,
"failed to remove thumbnail metadata: task panicked"
);
}
}
}
remove_existing_assets(ctx, &context).await;
let preview_asset_id = Uuid::new_v4();
let thumbnail_asset_id = Uuid::new_v4();
let mut preview_objects: Vec<AssetObjectPersistence> =
Vec::with_capacity(generation.preview.objects.len());
for (index, image) in generation.preview.objects.iter().enumerate() {
if index + 1 > i32::MAX as usize {
return JobExecution::Failed {
error: "too many preview objects".to_string(),
};
}
let ordinal = (index + 1) as i32;
let s3_key = document_asset_object_key(
initial.document.id,
initial.version.version_number,
PREVIEW_ASSET_TYPE,
preview_asset_id,
ordinal,
);
if let Err(err) = storage
.put_object(
&s3_key,
image.image_bytes.clone(),
Some("image/webp".into()),
None,
)
.await
{
warn!(job_id = %job.id, error = %err, ordinal, "failed to upload preview; retrying");
return JobExecution::Retry {
delay: Duration::from_secs(30),
error: err.to_string(),
};
}
preview_objects.push(AssetObjectPersistence {
ordinal,
s3_key,
width: image.width,
height: image.height,
});
}
let mut thumbnail_objects: Vec<AssetObjectPersistence> =
Vec::with_capacity(generation.thumbnail.objects.len());
for (index, image) in generation.thumbnail.objects.iter().enumerate() {
if index + 1 > i32::MAX as usize {
return JobExecution::Failed {
error: "too many thumbnail objects".to_string(),
};
}
let ordinal = (index + 1) as i32;
let s3_key = document_asset_object_key(
initial.document.id,
initial.version.version_number,
THUMBNAIL_ASSET_TYPE,
thumbnail_asset_id,
ordinal,
);
if let Err(err) = storage
.put_object(
&s3_key,
image.image_bytes.clone(),
Some("image/webp".into()),
None,
)
.await
{
warn!(job_id = %job.id, error = %err, ordinal, "failed to upload thumbnail; retrying");
return JobExecution::Retry {
delay: Duration::from_secs(30),
error: err.to_string(),
};
}
thumbnail_objects.push(AssetObjectPersistence {
ordinal,
s3_key,
width: image.width,
height: image.height,
});
}
let preview_objects = upload_generated_objects(
ctx,
&context,
PREVIEW_ASSET_TYPE,
preview_asset_id,
&generation.preview,
)
.await?;
let thumbnail_objects = upload_generated_objects(
ctx,
&context,
THUMBNAIL_ASSET_TYPE,
thumbnail_asset_id,
&generation.thumbnail,
)
.await?;
let asset_persistences = vec![
AssetPersistence {
@@ -351,33 +196,186 @@ impl JobHandler for GenerateThumbnailsJob {
},
];
let state_clone = state.clone();
match task::spawn_blocking(move || {
persist_assets_metadata(state_clone, &initial, &asset_persistences)
let state = ctx.state().clone();
let tenant_id = context.document.tenant_id;
let version_id = context.version.id;
task::spawn_blocking(move || {
persist_assets_metadata(state, tenant_id, version_id, &asset_persistences)
})
.await
{
Ok(Ok(())) => {}
Ok(Err(err)) => {
warn!(job_id = %job.id, error = %err, "failed to persist thumbnail metadata; retrying");
return JobExecution::Retry {
delay: Duration::from_secs(30),
error: err,
};
}
Err(join_err) => {
error!(job_id = %job.id, error = %join_err, "thumbnail metadata update panicked");
return JobExecution::Retry {
delay: Duration::from_secs(30),
error: format!("metadata update panic: {join_err}"),
};
}
}
.map_err(|err| {
TaskError::retry(
Duration::from_secs(60),
format!("thumbnail metadata task panicked: {err}"),
)
})?
.map_err(|err| TaskError::retry(Duration::from_secs(30), err))?;
JobExecution::Success
ctx.invalidate_asset_cache();
Ok(())
}
}
async fn build_thumbnail_context(
ctx: &mut DocumentVersionTaskContext,
force: bool,
) -> TaskResult<ThumbnailContext> {
let document = ctx.document().await?.clone();
let version = ctx.version().await?.clone();
let tenant_id = ctx.tenant_id();
let (supported, _) = determine_thumbnail_support(&document);
if !supported {
return Ok(ThumbnailContext {
document,
version,
existing_thumbnail: None,
existing_thumbnail_objects: Vec::new(),
existing_preview: None,
existing_preview_objects: Vec::new(),
skip: true,
tenant_id,
});
}
let assets = ctx.assets().await?;
let (existing_thumbnail, existing_thumbnail_objects) = assets
.get(THUMBNAIL_ASSET_TYPE)
.map(|entry| (Some(entry.asset.clone()), entry.objects.clone()))
.unwrap_or((None, Vec::new()));
let (existing_preview, existing_preview_objects) = assets
.get(PREVIEW_ASSET_TYPE)
.map(|entry| (Some(entry.asset.clone()), entry.objects.clone()))
.unwrap_or((None, Vec::new()));
let expected_cardinality = expected_asset_cardinality(&document, &version);
let preview_cardinality = existing_preview
.as_ref()
.and_then(|asset| asset.cardinality)
.unwrap_or_else(|| existing_preview_objects.len() as i32);
let thumbnail_cardinality = existing_thumbnail
.as_ref()
.and_then(|asset| asset.cardinality)
.unwrap_or_else(|| existing_thumbnail_objects.len() as i32);
let needs_regeneration = preview_cardinality < expected_cardinality
|| thumbnail_cardinality < expected_cardinality
|| (existing_preview_objects.len() as i32) < expected_cardinality
|| (existing_thumbnail_objects.len() as i32) < expected_cardinality;
let skip =
existing_thumbnail.is_some() && existing_preview.is_some() && !force && !needs_regeneration;
Ok(ThumbnailContext {
document,
version,
existing_thumbnail,
existing_thumbnail_objects,
existing_preview,
existing_preview_objects,
skip,
tenant_id,
})
}
async fn remove_existing_assets(ctx: &DocumentVersionTaskContext, context: &ThumbnailContext) {
if let Some(existing_preview) = &context.existing_preview {
delete_asset_with_objects(ctx, existing_preview, &context.existing_preview_objects).await;
}
if let Some(existing_thumbnail) = &context.existing_thumbnail {
delete_asset_with_objects(ctx, existing_thumbnail, &context.existing_thumbnail_objects)
.await;
}
}
async fn delete_asset_with_objects(
ctx: &DocumentVersionTaskContext,
asset: &DocumentAsset,
objects: &[DocumentAssetObject],
) {
for object in 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 asset object"
);
}
}
let tenant_id = ctx.tenant_id();
let asset_id = 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(),
asset_id = %asset_id,
error = ?err,
"failed to delete asset metadata"
);
}
Err(join_err) => {
warn!(
job_id = %ctx.job_id(),
asset_id = %asset_id,
error = %join_err,
"failed to delete asset metadata task panicked"
);
}
}
}
async fn upload_generated_objects(
ctx: &DocumentVersionTaskContext,
context: &ThumbnailContext,
asset_type: &str,
asset_id: Uuid,
asset: &GeneratedAsset,
) -> TaskResult<Vec<AssetObjectPersistence>> {
let mut objects = Vec::with_capacity(asset.objects.len());
for (index, image) in asset.objects.iter().enumerate() {
if index + 1 > i32::MAX as usize {
return Err(TaskError::fail("too many generated asset objects"));
}
let ordinal = (index + 1) as i32;
let s3_key = document_asset_object_key(
context.document.id,
context.version.version_number,
asset_type,
asset_id,
ordinal,
);
ctx.storage()
.put_object(
&s3_key,
image.image_bytes.clone(),
Some("image/webp".into()),
None,
)
.await
.map_err(|err| TaskError::retry(Duration::from_secs(30), err.to_string()))?;
objects.push(AssetObjectPersistence {
ordinal,
s3_key,
width: image.width,
height: image.height,
});
}
Ok(objects)
}
struct ThumbnailContext {
document: Document,
version: DocumentVersion,
@@ -418,83 +416,11 @@ struct AssetPersistence {
objects: Vec<AssetObjectPersistence>,
}
fn load_thumbnail_context(
state: Arc<AppState>,
tenant_id: Uuid,
payload: ThumbnailPayload,
) -> Result<ThumbnailContext, 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,
&[THUMBNAIL_ASSET_TYPE, PREVIEW_ASSET_TYPE],
)?;
let (existing_thumbnail, existing_thumbnail_objects) = assets
.remove(THUMBNAIL_ASSET_TYPE)
.map(|entry| (Some(entry.asset), entry.objects))
.unwrap_or((None, Vec::new()));
let (existing_preview, existing_preview_objects) = assets
.remove(PREVIEW_ASSET_TYPE)
.map(|entry| (Some(entry.asset), entry.objects))
.unwrap_or((None, Vec::new()));
let (supported, _) = determine_thumbnail_support(&base.document);
if !supported {
return Err("thumbnail generation not supported for this document".into());
}
let expected_cardinality = expected_asset_cardinality(&base.document, &base.version);
let preview_cardinality = existing_preview
.as_ref()
.and_then(|asset| asset.cardinality)
.unwrap_or_else(|| existing_preview_objects.len() as i32);
let thumbnail_cardinality = existing_thumbnail
.as_ref()
.and_then(|asset| asset.cardinality)
.unwrap_or_else(|| existing_thumbnail_objects.len() as i32);
let needs_regeneration = preview_cardinality < expected_cardinality
|| thumbnail_cardinality < expected_cardinality
|| (existing_preview_objects.len() as i32) < expected_cardinality
|| (existing_thumbnail_objects.len() as i32) < expected_cardinality;
let skip = existing_thumbnail.is_some()
&& existing_preview.is_some()
&& !payload.force
&& !needs_regeneration;
Ok(ThumbnailContext {
document: base.document,
version: base.version,
existing_thumbnail,
existing_thumbnail_objects,
existing_preview,
existing_preview_objects,
skip,
tenant_id: base.tenant_id,
})
}
fn generate_preview_and_thumbnail(
document: &Document,
bytes: &[u8],
) -> Result<GeneratedAssets, String> {
let is_pdf = document_is_pdf(document);
if is_pdf {
if document_is_pdf(document) {
let pdf_assets = generate_pdf_assets(bytes)?;
Ok(GeneratedAssets {
preview: pdf_assets.preview,
@@ -623,28 +549,14 @@ fn encode_dynamic_image(image: image::DynamicImage) -> Result<GeneratedImage, St
fn persist_assets_metadata(
state: Arc<AppState>,
context: &ThumbnailContext,
tenant_id: Uuid,
version_id: Uuid,
assets: &[AssetPersistence],
) -> 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_preview) = &context.existing_preview {
diesel::delete(document_assets::table.filter(document_assets::id.eq(existing_preview.id)))
.execute(&mut conn)
.map_err(|err| format!("{err:?}"))?;
}
if let Some(existing_thumbnail) = &context.existing_thumbnail {
diesel::delete(
document_assets::table.filter(document_assets::id.eq(existing_thumbnail.id)),
)
.execute(&mut conn)
.map_err(|err| format!("{err:?}"))?;
}
for asset in assets {
if asset.objects.is_empty() {
return Err(format!(
@@ -661,7 +573,7 @@ fn persist_assets_metadata(
let new_asset = NewDocumentAsset {
id: asset.asset_id,
document_version_id: context.version.id,
document_version_id: version_id,
asset_type: asset.asset_type.to_string(),
mime_type: "image/webp".to_string(),
metadata: json!({
@@ -688,7 +600,7 @@ fn persist_assets_metadata(
diesel::delete(
document_asset_objects::table
.filter(document_asset_objects::asset_id.eq(asset.asset_id))
.filter(document_asset_objects::asset_id.eq(new_asset.id))
.filter(document_asset_objects::tenant_id.eq(tenant_id)),
)
.execute(&mut conn)
@@ -707,7 +619,7 @@ fn persist_assets_metadata(
let new_object = NewDocumentAssetObject {
id: Uuid::new_v4(),
asset_id: asset.asset_id,
asset_id: new_asset.id,
ordinal: object.ordinal,
s3_key: object.s3_key.clone(),
metadata: object_metadata,