reduce to 1 object per asset

This commit is contained in:
2025-11-20 16:54:12 +01:00
parent 42f3917881
commit 3ce6b0ce93
24 changed files with 366 additions and 607 deletions
+3 -21
View File
@@ -3,8 +3,8 @@ use std::collections::HashMap;
use diesel::{prelude::*, PgConnection};
use uuid::Uuid;
use crate::models::{Document, DocumentAsset, DocumentAssetObject, DocumentVersion};
use crate::schema::{document_asset_objects, document_assets, document_versions, documents};
use crate::models::{Document, DocumentAsset, DocumentVersion};
use crate::schema::{document_assets, document_versions, documents};
use crate::state::AppState;
pub(crate) struct LoadedDocumentVersion {
@@ -41,7 +41,6 @@ pub(crate) fn load_document_version(
pub struct LoadedAsset {
pub asset: DocumentAsset,
pub objects: Vec<DocumentAssetObject>,
}
pub(crate) fn load_version_assets(
@@ -65,26 +64,9 @@ pub(crate) fn load_version_assets(
.load(conn)
.map_err(|err| format!("{err:?}"))?;
let asset_ids: Vec<Uuid> = assets.iter().map(|asset| asset.id).collect();
let mut object_map: HashMap<Uuid, Vec<DocumentAssetObject>> = HashMap::new();
if !asset_ids.is_empty() {
let objects: Vec<DocumentAssetObject> = document_asset_objects::table
.filter(document_asset_objects::asset_id.eq_any(&asset_ids))
.filter(document_asset_objects::tenant_id.eq(tenant_id))
.order(document_asset_objects::ordinal.asc())
.load(conn)
.map_err(|err| format!("{err:?}"))?;
for object in objects {
object_map.entry(object.asset_id).or_default().push(object);
}
}
let mut result = HashMap::with_capacity(assets.len());
for asset in assets {
let objects = object_map.remove(&asset.id).unwrap_or_default();
result.insert(asset.asset_type.clone(), LoadedAsset { asset, objects });
result.insert(asset.asset_type.clone(), LoadedAsset { asset });
}
Ok(result)
+1 -5
View File
@@ -48,12 +48,8 @@ impl Task<DocumentVersionTaskContext> for IndexDocumentTask {
.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 s3_key = asset.asset.s3_key.clone();
let bytes = ctx.storage().get_object(&s3_key).await.map_err(|err| {
TaskError::retry(
Duration::from_secs(30),
+1 -2
View File
@@ -430,8 +430,7 @@ async fn load_document_text(ctx: &mut DocumentVersionTaskContext) -> TaskResult<
let assets = ctx.assets().await?;
assets
.get(OCR_TEXT_ASSET_TYPE)
.and_then(|asset| asset.objects.first())
.map(|object| object.s3_key.clone())
.map(|asset| asset.asset.s3_key.clone())
};
let Some(key) = object_key else {
+14 -46
View File
@@ -19,13 +19,10 @@ use uuid::Uuid;
use crate::{
documents::asset::delete_asset,
error::AppResult,
models::{
Document, DocumentAsset, DocumentAssetObject, DocumentVersion, NewDocumentAsset,
NewDocumentAssetObject,
},
schema::{document_asset_objects, document_assets},
models::{Document, DocumentAsset, DocumentVersion, NewDocumentAsset},
schema::document_assets,
state::AppState,
utils::storage_paths::document_asset_object_prefix,
utils::storage_paths::document_asset_key,
};
use super::taskflow::{
@@ -83,7 +80,7 @@ impl Task<DocumentVersionTaskContext> for GenerateOcrTask {
remove_existing_ocr_asset(ctx, &context).await;
let asset_id = Uuid::new_v4();
let s3_key = document_asset_object_prefix(
let s3_key = document_asset_key(
context.document.id,
context.version.version_number,
OCR_TEXT_ASSET_TYPE,
@@ -123,7 +120,6 @@ struct OcrContext {
document: Document,
version: DocumentVersion,
existing_asset: Option<DocumentAsset>,
existing_objects: Vec<DocumentAssetObject>,
skip: bool,
}
@@ -135,16 +131,13 @@ async fn build_ocr_context(
let version = ctx.version().await?.clone();
let asset = ctx.asset(OCR_TEXT_ASSET_TYPE).await?;
let (existing_asset, existing_objects) = asset
.map(|asset| (Some(asset.asset.clone()), asset.objects.clone()))
.unwrap_or((None, Vec::new()));
let existing_asset = asset.map(|asset| asset.asset.clone());
if !document_is_pdf(&document) {
return Ok(OcrContext {
document,
version,
existing_asset,
existing_objects,
skip: true,
});
}
@@ -155,22 +148,19 @@ async fn build_ocr_context(
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"
);
}
if let Err(err) = ctx.storage().delete_object(&existing_asset.s3_key).await {
warn!(
job_id = %ctx.job_id(),
error = %err,
s3_key = %existing_asset.s3_key,
"failed to delete existing ocr asset object"
);
}
let tenant_id = context.document.tenant_id;
@@ -239,7 +229,7 @@ fn persist_ocr_metadata(
asset_type: OCR_TEXT_ASSET_TYPE.to_string(),
mime_type: "text/plain".to_string(),
metadata,
cardinality: Some(1),
s3_key: s3_key.to_string(),
tenant_id,
};
@@ -253,34 +243,12 @@ fn persist_ocr_metadata(
.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::s3_key.eq(excluded(document_assets::s3_key)),
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(())
}
+4 -14
View File
@@ -11,7 +11,7 @@ use uuid::Uuid;
use crate::auth::ensure_active_tenant;
use crate::jobs::JOB_PURGE_DOCUMENT;
use crate::models::{Document, DocumentVersion};
use crate::schema::{document_asset_objects, document_assets, document_versions};
use crate::schema::{document_assets, document_versions};
use crate::state::AppState;
use crate::storage::TenantStorage;
@@ -235,21 +235,11 @@ fn prepare_purge_context(
let asset_keys = if version_ids.is_empty() {
Vec::new()
} else {
let asset_ids: Vec<Uuid> = document_assets::table
document_assets::table
.filter(document_assets::document_version_id.eq_any(&version_ids))
.filter(document_assets::tenant_id.eq(tenant_id))
.select(document_assets::id)
.load(conn)?;
if asset_ids.is_empty() {
Vec::new()
} else {
document_asset_objects::table
.filter(document_asset_objects::asset_id.eq_any(&asset_ids))
.filter(document_asset_objects::tenant_id.eq(tenant_id))
.select(document_asset_objects::s3_key)
.load(conn)?
}
.select(document_assets::s3_key)
.load(conn)?
};
Ok(Some(PurgeContext {
+4 -9
View File
@@ -22,9 +22,8 @@ use crate::documents::search::{delete_quickwit_index, ensure_quickwit_index};
use crate::jobs::{JOB_DELETE_TENANT, JOB_PROVISION_TENANT};
use crate::models::{NewUserMembership, Tenant, TenantStatus};
use crate::schema::{
api_tokens, correspondents, document_asset_objects, document_assets, document_correspondents,
document_tags, document_versions, documents, folders, tags, tenants, user_memberships,
user_sessions,
api_tokens, correspondents, document_assets, document_correspondents, document_tags,
document_versions, documents, folders, tags, tenants, user_memberships, user_sessions,
};
use crate::state::AppState;
use crate::tenants::TenantRepository;
@@ -401,8 +400,8 @@ fn collect_object_keys(conn: &mut PgConnection) -> Result<TenantObjectKeys, dies
let version_keys = document_versions::table
.select(document_versions::s3_key)
.load::<String>(conn)?;
let asset_keys = document_asset_objects::table
.select(document_asset_objects::s3_key)
let asset_keys = document_assets::table
.select(document_assets::s3_key)
.load::<String>(conn)?;
Ok(TenantObjectKeys {
@@ -459,10 +458,6 @@ fn delete_tenant_rows(
remove_memberships: bool,
) -> Result<(), diesel::result::Error> {
conn.transaction(|conn| {
diesel::delete(
document_asset_objects::table.filter(document_asset_objects::tenant_id.eq(tenant_id)),
)
.execute(conn)?;
diesel::delete(document_assets::table.filter(document_assets::tenant_id.eq(tenant_id)))
.execute(conn)?;
diesel::delete(
+78 -179
View File
@@ -5,7 +5,7 @@ use chrono::Utc;
use diesel::{pg::upsert::excluded, prelude::*};
use image::{GenericImageView, ImageFormat, ImageReader};
use pdfium_render::prelude::*;
use serde_json::{json, Map, Value};
use serde_json::{Map, Value};
use tokio::task;
use tracing::{info, warn};
use uuid::Uuid;
@@ -13,13 +13,10 @@ use uuid::Uuid;
use crate::{
documents::asset::delete_asset,
error::AppResult,
models::{
Document, DocumentAsset, DocumentAssetObject, DocumentVersion, NewDocumentAsset,
NewDocumentAssetObject,
},
schema::{document_asset_objects, document_assets, document_versions},
models::{Document, DocumentAsset, DocumentVersion, NewDocumentAsset},
schema::{document_assets, document_versions},
state::AppState,
utils::storage_paths::document_asset_object_key,
utils::storage_paths::document_asset_key,
};
use super::{
@@ -82,7 +79,7 @@ impl Task<DocumentVersionTaskContext> for GenerateThumbnailsTask {
let thumbnail_asset_id = Uuid::new_v4();
let thumbnail_objects = upload_generated_objects(
let thumbnail_persistence = upload_generated_asset(
ctx,
&context,
THUMBNAIL_ASSET_TYPE,
@@ -91,11 +88,7 @@ impl Task<DocumentVersionTaskContext> for GenerateThumbnailsTask {
)
.await?;
let asset_persistences = vec![AssetPersistence {
asset_type: THUMBNAIL_ASSET_TYPE,
asset_id: thumbnail_asset_id,
objects: thumbnail_objects,
}];
let asset_persistences = vec![thumbnail_persistence];
let state = ctx.state().clone();
let tenant_id = context.document.tenant_id;
@@ -131,32 +124,22 @@ async fn build_thumbnail_context(
document,
version,
existing_thumbnail: None,
existing_thumbnail_objects: Vec::new(),
skip: true,
tenant_id,
});
}
let assets = ctx.assets().await?;
let (existing_thumbnail, existing_thumbnail_objects) = assets
let existing_thumbnail = assets
.get(THUMBNAIL_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 thumbnail_cardinality = existing_thumbnail
.as_ref()
.and_then(|asset| asset.cardinality)
.unwrap_or_else(|| existing_thumbnail_objects.len() as i32);
let needs_regeneration = thumbnail_cardinality < expected_cardinality
|| (existing_thumbnail_objects.len() as i32) < expected_cardinality;
.map(|entry| entry.asset.clone());
let skip = existing_thumbnail.is_some() && !force && !needs_regeneration;
let skip = existing_thumbnail.is_some() && !force;
Ok(ThumbnailContext {
document,
version,
existing_thumbnail,
existing_thumbnail_objects,
skip,
tenant_id,
})
@@ -167,25 +150,18 @@ async fn remove_existing_thumbnail_assets(
context: &ThumbnailContext,
) {
if let Some(existing_thumbnail) = &context.existing_thumbnail {
delete_asset_with_objects(ctx, existing_thumbnail, &context.existing_thumbnail_objects)
.await;
delete_asset_object(ctx, existing_thumbnail).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"
);
}
async fn delete_asset_object(ctx: &DocumentVersionTaskContext, asset: &DocumentAsset) {
if let Err(err) = ctx.storage().delete_object(&asset.s3_key).await {
warn!(
job_id = %ctx.job_id(),
error = %err,
s3_key = %asset.s3_key,
"failed to delete existing asset object"
);
}
let tenant_id = ctx.tenant_id();
@@ -217,54 +193,45 @@ async fn delete_asset_with_objects(
}
}
async fn upload_generated_objects(
async fn upload_generated_asset(
ctx: &DocumentVersionTaskContext,
context: &ThumbnailContext,
asset_type: &str,
asset_id: Uuid,
asset: &GeneratedAsset,
) -> TaskResult<Vec<AssetObjectPersistence>> {
let mut objects = Vec::with_capacity(asset.objects.len());
) -> TaskResult<AssetPersistence> {
let image = &asset.image;
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,
);
let s3_key = document_asset_key(
context.document.id,
context.version.version_number,
asset_type,
asset_id,
);
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()))?;
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)
Ok(AssetPersistence {
asset_type: asset_type.to_string(),
asset_id,
s3_key,
width: image.width,
height: image.height,
})
}
struct ThumbnailContext {
document: Document,
version: DocumentVersion,
existing_thumbnail: Option<DocumentAsset>,
existing_thumbnail_objects: Vec<DocumentAssetObject>,
skip: bool,
tenant_id: Uuid,
}
@@ -276,7 +243,7 @@ struct GeneratedImage {
}
struct GeneratedAsset {
objects: Vec<GeneratedImage>,
image: GeneratedImage,
}
struct GeneratedAssets {
@@ -284,19 +251,14 @@ struct GeneratedAssets {
page_count: Option<u32>,
}
struct AssetObjectPersistence {
ordinal: i32,
struct AssetPersistence {
asset_type: String,
asset_id: Uuid,
s3_key: String,
width: Option<i32>,
height: Option<i32>,
}
struct AssetPersistence {
asset_type: &'static str,
asset_id: Uuid,
objects: Vec<AssetObjectPersistence>,
}
fn generate_thumbnails(document: &Document, bytes: &[u8]) -> Result<GeneratedAssets, String> {
if document_is_pdf(document) {
let pdf_assets = generate_pdf_assets(bytes)?;
@@ -334,9 +296,7 @@ fn generate_image_assets(bytes: &[u8]) -> Result<GeneratedAsset, String> {
let thumbnail = encode_dynamic_image(thumbnail_image)?;
Ok(GeneratedAsset {
objects: vec![thumbnail],
})
Ok(GeneratedAsset { image: thumbnail })
}
struct PdfGeneratedAssets {
@@ -354,6 +314,9 @@ fn generate_pdf_assets(bytes: &[u8]) -> Result<PdfGeneratedAssets, String> {
let pages = document.pages();
let total_pages = pages.len() as usize;
if total_pages == 0 {
return Err("pdf has no pages".to_string());
}
let render_config = PdfRenderConfig::new()
.set_target_width(RENDER_WIDTH as i32)
@@ -361,29 +324,21 @@ fn generate_pdf_assets(bytes: &[u8]) -> Result<PdfGeneratedAssets, String> {
.render_form_data(true)
.rotate_if_landscape(PdfPageRenderRotation::None, true);
let mut thumbnail_objects: Vec<GeneratedImage> = Vec::with_capacity(total_pages);
let first_page = pages.get(0).map_err(|err| format!("load page 0: {err}"))?;
for page_index in 0..total_pages {
let page = pages
.get(u16::try_from(page_index).map_err(|_| "page index overflow".to_string())?)
.map_err(|err| format!("load page {page_index}: {err}"))?;
let bitmap = first_page
.render_with_config(&render_config)
.map_err(|err| format!("render pdf page 0: {err}"))?;
let bitmap = page
.render_with_config(&render_config)
.map_err(|err| format!("render pdf page {page_index}: {err}"))?;
let render_buffer = bitmap.as_image().to_rgb8();
let render_image = image::DynamicImage::ImageRgb8(render_buffer);
let render_buffer = bitmap.as_image().to_rgb8();
let render_image = image::DynamicImage::ImageRgb8(render_buffer);
let thumbnail_image =
if render_image.width() > THUMBNAIL_WIDTH || render_image.height() > THUMBNAIL_HEIGHT {
render_image.thumbnail(THUMBNAIL_WIDTH, THUMBNAIL_HEIGHT)
} else {
render_image.clone()
};
thumbnail_objects.push(encode_dynamic_image(thumbnail_image)?);
}
let thumbnail_image =
if render_image.width() > THUMBNAIL_WIDTH || render_image.height() > THUMBNAIL_HEIGHT {
render_image.thumbnail(THUMBNAIL_WIDTH, THUMBNAIL_HEIGHT)
} else {
render_image.clone()
};
let page_count: u32 = total_pages
.try_into()
@@ -391,7 +346,7 @@ fn generate_pdf_assets(bytes: &[u8]) -> Result<PdfGeneratedAssets, String> {
Ok(PdfGeneratedAssets {
thumbnail: GeneratedAsset {
objects: thumbnail_objects,
image: encode_dynamic_image(thumbnail_image)?,
},
page_count,
})
@@ -421,28 +376,25 @@ fn persist_assets_metadata(
.map_err(|err| format!("{err:?}"))?;
for asset in assets {
if asset.objects.is_empty() {
return Err(format!(
"asset {} has no generated objects",
asset.asset_type
));
let mut metadata_map = Map::new();
if let Some(width) = asset.width {
metadata_map.insert("width".to_string(), Value::from(width));
}
let object_count: i32 = asset
.objects
.len()
.try_into()
.map_err(|_| "asset contains too many objects".to_string())?;
if let Some(height) = asset.height {
metadata_map.insert("height".to_string(), Value::from(height));
}
metadata_map.insert(
"generated_at".to_string(),
Value::from(Utc::now().to_rfc3339()),
);
let new_asset = NewDocumentAsset {
id: asset.asset_id,
document_version_id: version_id,
asset_type: asset.asset_type.to_string(),
asset_type: asset.asset_type.clone(),
mime_type: "image/webp".to_string(),
metadata: json!({
"generated_at": Utc::now().to_rfc3339(),
}),
cardinality: Some(object_count),
metadata: Value::Object(metadata_map),
s3_key: asset.s3_key.clone(),
tenant_id,
};
@@ -456,44 +408,10 @@ fn persist_assets_metadata(
.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::s3_key.eq(excluded(document_assets::s3_key)),
))
.execute(&mut conn)
.map_err(|err| format!("{err:?}"))?;
diesel::delete(
document_asset_objects::table
.filter(document_asset_objects::asset_id.eq(new_asset.id))
.filter(document_asset_objects::tenant_id.eq(tenant_id)),
)
.execute(&mut conn)
.map_err(|err| format!("{err:?}"))?;
for object in &asset.objects {
let mut metadata_map = Map::new();
if let Some(width) = object.width {
metadata_map.insert("width".to_string(), Value::from(width));
}
if let Some(height) = object.height {
metadata_map.insert("height".to_string(), Value::from(height));
}
let object_metadata = Value::Object(metadata_map);
let new_object = NewDocumentAssetObject {
id: Uuid::new_v4(),
asset_id: new_asset.id,
ordinal: object.ordinal,
s3_key: object.s3_key.clone(),
metadata: object_metadata,
tenant_id,
};
diesel::insert_into(document_asset_objects::table)
.values(&new_object)
.execute(&mut conn)
.map_err(|err| format!("{err:?}"))?;
}
}
Ok(())
@@ -557,22 +475,3 @@ fn document_is_pdf(document: &Document) -> bool {
.unwrap_or(false)
})
}
fn expected_asset_cardinality(document: &Document, version: &DocumentVersion) -> i32 {
if let Value::Object(map) = &version.metadata {
if let Some(count) = map.get("page_count").and_then(|v| v.as_i64()) {
if count > 0 {
return count
.min(i64::from(i32::MAX))
.try_into()
.unwrap_or(i32::MAX);
}
}
}
if document_is_pdf(document) {
1
} else {
1
}
}