multi tenancy part 1
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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");
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user