This commit is contained in:
2025-10-11 14:39:58 +02:00
parent 6f6abef5d2
commit cd27f36e3d
12 changed files with 616 additions and 12 deletions
+68
View File
@@ -0,0 +1,68 @@
use std::env;
use anyhow::{Context, Result};
use diesel::prelude::*;
use paperless_backend::{
config::AppConfig,
db,
models::DocumentAsset,
s3,
schema::document_assets,
storage::{ObjectStorage, S3Storage},
};
#[tokio::main]
async fn main() -> Result<()> {
let mut args = env::args().skip(1);
match args.next().as_deref() {
Some("delete-assets") => delete_all_assets().await?,
Some(cmd) => {
eprintln!("Unknown command: {cmd}\nUsage: maintenance delete-assets");
std::process::exit(1);
}
None => {
eprintln!("Usage: maintenance delete-assets");
std::process::exit(1);
}
}
Ok(())
}
async fn delete_all_assets() -> Result<()> {
let config = AppConfig::from_env()?;
let pool = db::init_pool(&config.database_url)?;
let s3_client = s3::build_client(&config).await?;
let storage = S3Storage::new(s3_client, config.s3_bucket.clone());
let mut conn = pool.get().context("failed to get database connection")?;
let assets: Vec<DocumentAsset> = document_assets::table
.load(&mut conn)
.context("failed to load document assets")?;
if assets.is_empty() {
println!("No assets found.");
return Ok(());
}
println!("Deleting {} assets…", assets.len());
for asset in &assets {
if let Err(err) = storage.delete_object(&asset.s3_key).await {
eprintln!(
"Failed to delete object {} from storage: {err}",
asset.s3_key
);
}
}
diesel::delete(document_assets::table)
.execute(&mut conn)
.context("failed to remove asset records")?;
println!("Asset records deleted.");
Ok(())
}
+1
View File
@@ -17,6 +17,7 @@ pub const STATUS_FAILED: &str = "failed";
pub const JOB_ANALYZE_DOCUMENT: &str = "analyze-document";
pub const JOB_GENERATE_THUMBNAILS: &str = "generate-thumbnails";
pub const JOB_GENERATE_OCR_TEXT: &str = "generate-ocr-text";
#[derive(Debug, Error)]
pub enum JobQueueError {
+13
View File
@@ -18,6 +18,8 @@ pub trait ObjectStorage: Send + Sync + 'static {
async fn presign_get_object(&self, key: &str, expires_in: Duration) -> Result<String>;
async fn get_object(&self, key: &str) -> Result<Vec<u8>>;
async fn delete_object(&self, key: &str) -> Result<()>;
}
pub struct S3Storage {
@@ -99,4 +101,15 @@ impl ObjectStorage for S3Storage {
Ok(bytes)
}
async fn delete_object(&self, key: &str) -> Result<()> {
self.client
.delete_object()
.bucket(&self.bucket)
.key(key)
.send()
.await
.context("failed to delete object from S3")?;
Ok(())
}
}
+41 -3
View File
@@ -8,10 +8,11 @@ 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_THUMBNAILS},
models::{Document, DocumentVersion},
schema::{document_versions, documents},
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,
};
@@ -88,6 +89,16 @@ fn analyze_document(state: Arc<AppState>, payload: AnalyzePayload) -> Result<Job
.map_err(|err| format!("{err:?}"))?;
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))
.first(&mut conn)
.optional()
.map_err(|err| format!("{err:?}"))?;
let skip_ocr = existing_ocr.is_some() && !payload.force;
let mut summary_map = match version.operations_summary {
Value::Object(map) => map,
@@ -100,6 +111,16 @@ fn analyze_document(state: Arc<AppState>, payload: AnalyzePayload) -> Result<Job
summary_map.remove("thumbnail_reason");
}
summary_map.insert("ocr_supported".to_string(), Value::Bool(ocr_supported));
if ocr_supported {
summary_map.remove("ocr_reason");
} else {
summary_map.insert(
"ocr_reason".to_string(),
Value::String("document is not a PDF".into()),
);
}
diesel::update(document_versions::table.find(version.id))
.set(document_versions::operations_summary.eq(Value::Object(summary_map)))
.execute(&mut conn)
@@ -122,6 +143,23 @@ fn analyze_document(state: Arc<AppState>, payload: AnalyzePayload) -> Result<Job
}
}
if ocr_supported && !skip_ocr {
let enqueue_result = enqueue_job(
&mut conn,
JOB_GENERATE_OCR_TEXT,
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());
}
}
Ok(JobExecution::Success)
}
+2
View File
@@ -11,6 +11,7 @@ use crate::{
};
pub mod analyze;
pub mod ocr;
pub mod thumbnails;
#[derive(Debug)]
@@ -128,5 +129,6 @@ pub fn default_handlers() -> Vec<Arc<dyn JobHandler>> {
vec![
Arc::new(analyze::AnalyzeDocumentJob::new()),
Arc::new(thumbnails::GenerateThumbnailsJob::new()),
Arc::new(ocr::GenerateOcrTextJob::new()),
]
}
+414
View File
@@ -0,0 +1,414 @@
use std::{
fmt,
fs,
io::{ErrorKind, Write},
process::Command,
sync::Arc,
time::Duration,
};
use async_trait::async_trait;
use chrono::Utc;
use diesel::{pg::upsert::excluded, prelude::*};
use pdfium_render::prelude::*;
use serde::Deserialize;
use serde_json::json;
use tempfile::NamedTempFile;
use tokio::task;
use tracing::{error, info, warn};
use uuid::Uuid;
use crate::{
jobs::JOB_GENERATE_OCR_TEXT,
models::{Document, DocumentAsset, DocumentVersion, NewDocumentAsset},
schema::{document_assets, document_versions, documents},
state::AppState,
};
use super::{JobExecution, JobHandler};
pub const OCR_TEXT_ASSET_TYPE: &str = "ocr-text";
const MIN_TEXT_LENGTH: usize = 50;
#[derive(Debug, Deserialize)]
struct OcrPayload {
document_id: Uuid,
document_version_id: Uuid,
#[serde(default)]
force: bool,
}
pub struct GenerateOcrTextJob;
impl GenerateOcrTextJob {
pub fn new() -> Self {
Self
}
}
#[async_trait]
impl JobHandler for GenerateOcrTextJob {
fn job_type(&self) -> &'static str {
JOB_GENERATE_OCR_TEXT
}
async fn handle(&self, state: Arc<AppState>, job: crate::models::Job) -> JobExecution {
let payload: OcrPayload = match serde_json::from_value(job.payload.clone()) {
Ok(payload) => payload,
Err(err) => {
return JobExecution::Failed {
error: format!("invalid OCR payload: {err}"),
}
}
};
let state_clone = state.clone();
let context =
match task::spawn_blocking(move || load_ocr_context(state_clone, &payload)).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}"),
};
}
};
if context.skip {
info!(job_id = %job.id, "ocr already present; skipping");
return JobExecution::Success;
}
let bytes = match state.storage.get_object(&context.version.s3_key).await {
Ok(bytes) => bytes,
Err(err) => {
warn!(job_id = %job.id, error = %err, "failed to fetch document for ocr");
return JobExecution::Retry {
delay: Duration::from_secs(30),
error: err.to_string(),
};
}
};
let doc_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 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(),
};
};
let asset_id = context
.existing_asset
.as_ref()
.map(|asset| asset.id)
.unwrap_or_else(Uuid::new_v4);
let s3_key = format!(
"documents/{}/v{}/assets/{}/{}",
context.document.id, context.version.version_number, OCR_TEXT_ASSET_TYPE, asset_id
);
if let Err(err) = state
.storage
.put_object(
&s3_key,
generation.text.into_bytes(),
Some("text/plain".into()),
)
.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(),
};
}
let state_clone = state.clone();
match task::spawn_blocking(move || {
persist_ocr_metadata(state_clone, &context, asset_id, &s3_key, generation.source)
})
.await
{
Ok(Ok(())) => 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}"),
}
}
}
}
}
struct PdfDocumentMeta {
content_type: Option<String>,
original_name: String,
}
struct OcrContext {
document: Document,
version: DocumentVersion,
existing_asset: Option<DocumentAsset>,
skip: bool,
}
struct OcrGeneration {
text: String,
source: &'static str,
}
fn load_ocr_context(state: Arc<AppState>, payload: &OcrPayload) -> Result<OcrContext, String> {
let mut conn = state.db().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 existing: 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))
.first(&mut conn)
.optional()
.map_err(|err| format!("{err:?}"))?;
let is_pdf = document_is_pdf(&document);
if !is_pdf {
return Ok(OcrContext {
document,
version,
existing_asset: existing,
skip: true,
});
}
let skip = existing.is_some() && !payload.force;
Ok(OcrContext {
document,
version,
existing_asset: existing,
skip,
})
}
fn generate_ocr_text(meta: &PdfDocumentMeta, bytes: &[u8]) -> Option<OcrGeneration> {
if !document_meta_is_pdf(meta) {
return None;
}
if let Ok(text) = extract_pdf_text(bytes) {
if text.trim().chars().count() >= MIN_TEXT_LENGTH {
return Some(OcrGeneration {
text,
source: "pdf-text",
});
}
}
match run_ocr(bytes) {
Ok(Some(text)) => Some(OcrGeneration {
text,
source: "ocr",
}),
Ok(None) => None,
Err(OcrError::BinaryMissing) => {
warn!("ocrmypdf not installed; cannot perform OCR");
None
}
Err(err) => {
warn!(error = ?err, "ocr command failed");
None
}
}
}
fn extract_pdf_text(bytes: &[u8]) -> Result<String, String> {
let pdfium = Pdfium::default();
let document = pdfium
.load_pdf_from_byte_slice(bytes, None)
.map_err(|err| format!("load pdf: {err}"))?;
let mut combined = String::new();
let pages = document.pages();
for page_index in 0..pages.len() {
let page = pages
.get(page_index)
.map_err(|err| format!("load page {page_index}: {err}"))?;
if let Ok(page_text) = page.text() {
for segment in page_text.segments().iter() {
combined.push_str(&segment.text());
combined.push('\n');
}
};
}
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
.write_all(bytes)
.map_err(|err| OcrError::Failed(err.to_string()))?;
input
.flush()
.map_err(|err| OcrError::Failed(err.to_string()))?;
let output_pdf = NamedTempFile::new().map_err(|err| OcrError::Failed(err.to_string()))?;
let sidecar = NamedTempFile::new().map_err(|err| OcrError::Failed(err.to_string()))?;
let status = Command::new("ocrmypdf")
.arg("--sidecar")
.arg(sidecar.path())
.arg("--skip-text")
.arg(input.path())
.arg(output_pdf.path())
.output();
match status {
Ok(output) => {
if !output.status.success() {
return Err(OcrError::Failed(format!(
"ocrmypdf failed: exit={} stderr={}",
output.status,
String::from_utf8_lossy(&output.stderr)
)));
}
let text = fs::read_to_string(sidecar.path()).map_err(|err| OcrError::Failed(err.to_string()))?;
if text.trim().chars().count() >= MIN_TEXT_LENGTH {
Ok(Some(text))
} else {
Ok(None)
}
}
Err(err) => {
if err.kind() == ErrorKind::NotFound {
Err(OcrError::BinaryMissing)
} else {
Err(OcrError::Failed(err.to_string()))
}
}
}
}
fn persist_ocr_metadata(
state: Arc<AppState>,
context: &OcrContext,
asset_id: Uuid,
s3_key: &str,
source: &'static str,
) -> Result<(), String> {
let mut conn = state.db().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(),
s3_key: s3_key.to_string(),
mime_type: "text/plain".to_string(),
width: None,
height: None,
metadata: json!({
"generated_at": Utc::now().to_rfc3339(),
"source": source,
}),
};
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::s3_key.eq(excluded(document_assets::s3_key)),
document_assets::mime_type.eq(excluded(document_assets::mime_type)),
document_assets::metadata.eq(excluded(document_assets::metadata)),
))
.execute(&mut conn)
.map_err(|err| format!("{err:?}"))?;
Ok(())
}
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") {
return true;
}
}
meta.original_name
.rsplit('.')
.next()
.map(|ext| ext.eq_ignore_ascii_case("pdf"))
.unwrap_or(false)
}
+18 -8
View File
@@ -102,10 +102,20 @@ impl JobHandler for GenerateThumbnailsJob {
}
};
let asset_id = initial
.existing_asset
.as_ref()
.map(|asset| asset.id)
.unwrap_or_else(Uuid::new_v4);
let s3_key = format!(
"documents/{}/v{}/assets/{}/{}",
initial.document.id, initial.version.version_number, THUMBNAIL_ASSET_TYPE, asset_id
);
if let Err(err) = state
.storage
.put_object(
&generation.s3_key,
&s3_key,
generation.image_bytes.clone(),
Some("image/png".into()),
)
@@ -120,7 +130,7 @@ impl JobHandler for GenerateThumbnailsJob {
let state_clone = state.clone();
match task::spawn_blocking(move || {
persist_thumbnail_metadata(state_clone, &initial, &generation)
persist_thumbnail_metadata(state_clone, &initial, &generation, asset_id, &s3_key)
})
.await
{
@@ -148,6 +158,7 @@ impl JobHandler for GenerateThumbnailsJob {
struct ThumbnailContext {
document: Document,
version: DocumentVersion,
existing_asset: Option<DocumentAsset>,
skip: bool,
}
@@ -155,7 +166,6 @@ struct GeneratedThumbnail {
image_bytes: Vec<u8>,
width: Option<i32>,
height: Option<i32>,
s3_key: String,
}
fn load_thumbnail_context(
@@ -195,6 +205,7 @@ fn load_thumbnail_context(
Ok(ThumbnailContext {
document,
version,
existing_asset: existing,
skip,
})
}
@@ -219,13 +230,10 @@ fn generate_thumbnail(document: &Document, bytes: &[u8]) -> Result<GeneratedThum
generate_image_thumbnail(bytes)?
};
let s3_key = format!("thumbnails/{}/{}.png", document.id, Uuid::new_v4());
Ok(GeneratedThumbnail {
image_bytes: png_bytes,
width,
height,
s3_key,
})
}
@@ -284,14 +292,16 @@ fn persist_thumbnail_metadata(
state: Arc<AppState>,
context: &ThumbnailContext,
generated: &GeneratedThumbnail,
asset_id: Uuid,
s3_key: &str,
) -> Result<(), String> {
let mut conn = state.db().map_err(|err| format!("{err:?}"))?;
let new_asset = NewDocumentAsset {
id: Uuid::new_v4(),
id: asset_id,
document_version_id: context.version.id,
asset_type: THUMBNAIL_ASSET_TYPE.to_string(),
s3_key: generated.s3_key.clone(),
s3_key: s3_key.to_string(),
mime_type: "image/png".to_string(),
width: generated.width,
height: generated.height,