multi tenancy part 1

This commit is contained in:
2025-10-22 22:52:59 +02:00
parent 8e3f09774a
commit e175d28c2c
35 changed files with 1288 additions and 314 deletions
+16 -3
View File
@@ -51,7 +51,9 @@ impl JobHandler for AnalyzeDocumentJob {
};
let state_clone = state.clone();
match task::spawn_blocking(move || analyze_document(state_clone, payload)).await {
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");
@@ -71,8 +73,14 @@ impl JobHandler for AnalyzeDocumentJob {
}
}
fn analyze_document(state: Arc<AppState>, payload: AnalyzePayload) -> Result<JobExecution, String> {
let mut conn = state.db().map_err(|err| format!("{err:?}"))?;
fn analyze_document(
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)
@@ -88,12 +96,15 @@ fn analyze_document(state: Arc<AppState>, payload: AnalyzePayload) -> Result<Job
.first(&mut conn)
.map_err(|err| format!("{err:?}"))?;
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:?}"))?;
@@ -129,6 +140,7 @@ fn analyze_document(state: Arc<AppState>, payload: AnalyzePayload) -> Result<Job
if supported {
let enqueue_result = enqueue_job(
&mut conn,
tenant_id,
JOB_GENERATE_THUMBNAILS,
json!({
"document_id": payload.document_id,
@@ -146,6 +158,7 @@ fn analyze_document(state: Arc<AppState>, payload: AnalyzePayload) -> Result<Job
if ocr_supported && !skip_ocr {
let enqueue_result = enqueue_job(
&mut conn,
tenant_id,
JOB_GENERATE_OCR_TEXT,
json!({
"document_id": payload.document_id,
+12 -3
View File
@@ -173,11 +173,11 @@ struct IndexContext {
}
fn load_context(state: Arc<AppState>, payload: &IndexPayload) -> Result<IndexContext, String> {
let mut conn = state.db().map_err(|err| format!("{err:?}"))?;
let mut base_conn = state.db_unscoped().map_err(|err| format!("{err:?}"))?;
let version: DocumentVersion = document_versions::table
.find(payload.document_version_id)
.first(&mut conn)
.first(&mut base_conn)
.map_err(|err| format!("{err:?}"))?;
if version.document_id != payload.document_id {
@@ -186,7 +186,14 @@ fn load_context(state: Arc<AppState>, payload: &IndexPayload) -> Result<IndexCon
let document: Document = documents::table
.find(payload.document_id)
.first(&mut conn)
.first(&mut base_conn)
.map_err(|err| format!("{err:?}"))?;
let tenant_id = document.tenant_id;
drop(base_conn);
let mut conn = state
.db_for_tenant(tenant_id)
.map_err(|err| format!("{err:?}"))?;
let text_s3_key: Option<String> = document_asset_objects::table
@@ -195,7 +202,9 @@ fn load_context(state: Arc<AppState>, payload: &IndexPayload) -> Result<IndexCon
)
.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()
+5 -5
View File
@@ -71,7 +71,7 @@ impl Worker {
return Ok(false);
}
let mut conn = match self.state.db() {
let mut conn = match self.state.db_unscoped() {
Ok(conn) => conn,
Err(err) => {
error!(?err, "failed to obtain database connection in worker");
@@ -87,7 +87,7 @@ impl Worker {
let result = handler.handle(self.state.clone(), job.clone()).await;
match result {
JobExecution::Success => {
if let Ok(mut conn) = self.state.db() {
if let Ok(mut conn) = self.state.db_unscoped() {
mark_job_succeeded(&mut conn, job.id)?;
info!(job_id = %job.id, job_type = %job.job_type, "job completed successfully");
} else {
@@ -96,7 +96,7 @@ impl Worker {
}
JobExecution::Retry { delay, error } => {
warn!(job_id = %job.id, job_type = %job.job_type, %error, "job will retry");
if let Ok(mut conn) = self.state.db() {
if let Ok(mut conn) = self.state.db_unscoped() {
retry_job_after(&mut conn, job.id, delay, &error)?;
} else {
error!("failed to requeue job for retry due to pool error");
@@ -104,7 +104,7 @@ impl Worker {
}
JobExecution::Failed { error } => {
error!(job_id = %job.id, job_type = %job.job_type, %error, "job failed");
if let Ok(mut conn) = self.state.db() {
if let Ok(mut conn) = self.state.db_unscoped() {
mark_job_failed(&mut conn, job.id, &error)?;
} else {
error!("failed to mark job failed due to pool error");
@@ -113,7 +113,7 @@ impl Worker {
}
} else {
error!(job_type = %job.job_type, "no handler registered for job type");
if let Ok(mut conn) = self.state.db() {
if let Ok(mut conn) = self.state.db_unscoped() {
mark_job_failed(&mut conn, job.id, "no handler registered")?;
} else {
error!("failed to mark job failed for missing handler due to pool error");
+31 -5
View File
@@ -213,11 +213,11 @@ struct OcrGeneration {
}
fn load_ocr_context(state: Arc<AppState>, payload: &OcrPayload) -> Result<OcrContext, String> {
let mut conn = state.db().map_err(|err| format!("{err:?}"))?;
let mut base_conn = state.db_unscoped().map_err(|err| format!("{err:?}"))?;
let version: DocumentVersion = document_versions::table
.find(payload.document_version_id)
.first(&mut conn)
.first(&mut base_conn)
.map_err(|err| format!("{err:?}"))?;
if version.document_id != payload.document_id {
@@ -226,12 +226,20 @@ fn load_ocr_context(state: Arc<AppState>, payload: &OcrPayload) -> Result<OcrCon
let document: Document = documents::table
.find(payload.document_id)
.first(&mut conn)
.first(&mut base_conn)
.map_err(|err| format!("{err:?}"))?;
let tenant_id = document.tenant_id;
drop(base_conn);
let mut conn = state
.db_for_tenant(tenant_id)
.map_err(|err| format!("{err:?}"))?;
let existing_asset: 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:?}"))?;
@@ -239,6 +247,7 @@ fn load_ocr_context(state: Arc<AppState>, payload: &OcrPayload) -> Result<OcrCon
let existing_objects: Vec<DocumentAssetObject> = if let Some(asset) = &existing_asset {
document_asset_objects::table
.filter(document_asset_objects::asset_id.eq(asset.id))
.filter(document_asset_objects::tenant_id.eq(tenant_id))
.order(document_asset_objects::ordinal.asc())
.load(&mut conn)
.map_err(|err| format!("{err:?}"))?
@@ -392,7 +401,10 @@ fn persist_ocr_metadata(
s3_key: &str,
source: &'static str,
) -> Result<(), String> {
let mut conn = state.db().map_err(|err| format!("{err:?}"))?;
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)))
@@ -410,6 +422,7 @@ fn persist_ocr_metadata(
"source": source,
}),
cardinality: Some(1),
tenant_id,
};
diesel::insert_into(document_assets::table)
@@ -430,6 +443,7 @@ fn persist_ocr_metadata(
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()
@@ -443,6 +457,7 @@ fn persist_ocr_metadata(
ordinal: 1,
s3_key: s3_key.to_string(),
metadata: json!({}),
tenant_id,
};
diesel::insert_into(document_asset_objects::table)
@@ -463,9 +478,20 @@ fn persist_ocr_metadata(
}
fn enqueue_index_job(state: &AppState, payload: &OcrPayload) -> Result<(), String> {
let mut conn = state.db().map_err(|err| format!("{err:?}"))?;
let mut base_conn = state.db_unscoped().map_err(|err| format!("{err:?}"))?;
let tenant_id: Uuid = documents::table
.find(payload.document_id)
.select(documents::tenant_id)
.first(&mut base_conn)
.map_err(|err| format!("{err:?}"))?;
drop(base_conn);
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,
+23 -5
View File
@@ -343,7 +343,7 @@ fn load_thumbnail_context(
state: Arc<AppState>,
payload: &ThumbnailPayload,
) -> Result<ThumbnailContext, String> {
let mut conn = state.db().map_err(|err| format!("{err:?}"))?;
let mut conn = state.db_unscoped().map_err(|err| format!("{err:?}"))?;
let version: DocumentVersion = document_versions::table
.find(payload.document_version_id)
@@ -359,12 +359,15 @@ fn load_thumbnail_context(
.first(&mut conn)
.map_err(|err| format!("{err:?}"))?;
let tenant_id = document.tenant_id;
let existing_assets: Vec<DocumentAsset> = document_assets::table
.filter(document_assets::document_version_id.eq(payload.document_version_id))
.filter(document_assets::asset_type.eq_any(vec![
THUMBNAIL_ASSET_TYPE.to_string(),
PREVIEW_ASSET_TYPE.to_string(),
]))
.filter(document_assets::tenant_id.eq(tenant_id))
.load(&mut conn)
.map_err(|err| format!("{err:?}"))?;
@@ -377,6 +380,7 @@ fn load_thumbnail_context(
THUMBNAIL_ASSET_TYPE => {
existing_thumbnail_objects = document_asset_objects::table
.filter(document_asset_objects::asset_id.eq(asset.id))
.filter(document_asset_objects::tenant_id.eq(tenant_id))
.order(document_asset_objects::ordinal.asc())
.load(&mut conn)
.map_err(|err| format!("{err:?}"))?;
@@ -385,6 +389,7 @@ fn load_thumbnail_context(
PREVIEW_ASSET_TYPE => {
existing_preview_objects = document_asset_objects::table
.filter(document_asset_objects::asset_id.eq(asset.id))
.filter(document_asset_objects::tenant_id.eq(tenant_id))
.order(document_asset_objects::ordinal.asc())
.load(&mut conn)
.map_err(|err| format!("{err:?}"))?;
@@ -568,7 +573,10 @@ fn persist_assets_metadata(
context: &ThumbnailContext,
assets: &[AssetPersistence],
) -> Result<(), String> {
let mut conn = state.db().map_err(|err| format!("{err:?}"))?;
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)))
@@ -607,6 +615,7 @@ fn persist_assets_metadata(
"generated_at": Utc::now().to_rfc3339(),
}),
cardinality: Some(object_count),
tenant_id,
};
diesel::insert_into(document_assets::table)
@@ -626,7 +635,8 @@ 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(asset.asset_id))
.filter(document_asset_objects::tenant_id.eq(tenant_id)),
)
.execute(&mut conn)
.map_err(|err| format!("{err:?}"))?;
@@ -648,6 +658,7 @@ fn persist_assets_metadata(
ordinal: object.ordinal,
s3_key: object.s3_key.clone(),
metadata: object_metadata,
tenant_id,
};
diesel::insert_into(document_asset_objects::table)
@@ -666,11 +677,17 @@ fn persist_document_page_count(
document_version_id: Uuid,
page_count: u32,
) -> Result<(), String> {
let mut conn = state.db().map_err(|err| format!("{err:?}"))?;
let mut conn = state.db_unscoped().map_err(|err| format!("{err:?}"))?;
let tenant_id: Uuid = documents::table
.find(document_id)
.select(documents::tenant_id)
.first(&mut conn)
.map_err(|err| format!("{err:?}"))?;
let existing_metadata: Value = document_versions::table
.filter(document_versions::id.eq(document_version_id))
.filter(document_versions::document_id.eq(document_id))
.filter(document_versions::tenant_id.eq(tenant_id))
.select(document_versions::metadata)
.first(&mut conn)
.map_err(|err| format!("{err:?}"))?;
@@ -690,7 +707,8 @@ fn persist_document_page_count(
diesel::update(
document_versions::table
.filter(document_versions::id.eq(document_version_id))
.filter(document_versions::document_id.eq(document_id)),
.filter(document_versions::document_id.eq(document_id))
.filter(document_versions::tenant_id.eq(tenant_id)),
)
.set(document_versions::metadata.eq(updated))
.execute(&mut conn)