diff --git a/crates/lib/mysql_queries/src/queries/batch_generations/insert_batch_generation_records.rs b/crates/lib/mysql_queries/src/queries/batch_generations/insert_batch_generation_records.rs index 4b9c6f4d56..b5007d1a22 100644 --- a/crates/lib/mysql_queries/src/queries/batch_generations/insert_batch_generation_records.rs +++ b/crates/lib/mysql_queries/src/queries/batch_generations/insert_batch_generation_records.rs @@ -13,6 +13,7 @@ pub struct BatchEntry { pub struct InsertBatchArgs<'a, 'b> { pub entries: Vec, + pub maybe_existing_batch_token: Option<&'a BatchGenerationToken>, pub transaction: &'a mut Transaction<'b, MySql>, } @@ -21,7 +22,11 @@ pub struct InsertBatchArgs<'a, 'b> { /// /// NB: Calling code is responsible for rolling back the transaction if this fails. pub async fn insert_batch_generation_records<'a, 'b>(args: InsertBatchArgs<'a, 'b>) -> AnyhowResult { - let batch_token = BatchGenerationToken::generate(); + + let batch_token = match args.maybe_existing_batch_token { + Some(existing_token) => existing_token.clone(), + None => BatchGenerationToken::generate(), + }; let mut query_builder = QueryBuilder::new(r#" INSERT INTO batch_generations (token, entity_type, entity_token) VALUES diff --git a/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_generic.rs b/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_generic.rs index f385dcb324..36bb1f6b9a 100644 --- a/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_generic.rs +++ b/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_generic.rs @@ -8,6 +8,7 @@ use enums::by_table::media_files::media_file_origin_model_type::MediaFileOriginM use enums::by_table::media_files::media_file_origin_product_category::MediaFileOriginProductCategory; use enums::by_table::media_files::media_file_type::MediaFileType; use errors::AnyhowResult; +use tokens::tokens::batch_generations::BatchGenerationToken; use tokens::tokens::media_files::MediaFileToken; use tokens::tokens::model_weights::ModelWeightToken; use tokens::tokens::users::UserToken; @@ -29,8 +30,8 @@ pub struct InsertArgs<'a> { pub maybe_origin_model_token: Option, pub maybe_origin_filename: Option, - // Generation flags and media details - pub is_batch_generated: bool, // Assuming you set this in your application logic + // If batch generated, this is the batch token. + pub maybe_batch_token: Option<&'a BatchGenerationToken>, pub maybe_mime_type: Option<&'a str>, @@ -67,6 +68,8 @@ pub async fn insert_media_file_generic( let mut maybe_creator_file_synthetic_id : Option = None; let mut maybe_creator_category_synthetic_id : Option = None; + let is_batch_generated = args.maybe_batch_token.is_some(); + let mut transaction = args.pool.begin().await?; // create a user token @@ -102,7 +105,8 @@ pub async fn insert_media_file_generic( maybe_origin_model_token = ?, maybe_origin_filename = ?, - is_batch_generated = ?, + is_batch_generated = ?, + maybe_batch_token = ?, maybe_mime_type = ?, file_size_bytes = ?, @@ -141,7 +145,8 @@ pub async fn insert_media_file_generic( args.maybe_origin_model_token.map(|e| e.to_string()), args.maybe_origin_filename, - args.is_batch_generated, + is_batch_generated, + args.maybe_batch_token.map(|e| e.to_string()), args.maybe_mime_type, args.file_size_bytes, diff --git a/crates/service/job/inference_job/src/job/job_types/bevy_to_workflow/bvh_to_workflow/process_job.rs b/crates/service/job/inference_job/src/job/job_types/bevy_to_workflow/bvh_to_workflow/process_job.rs index fa5d1b6cf9..88464ed4f5 100644 --- a/crates/service/job/inference_job/src/job/job_types/bevy_to_workflow/bvh_to_workflow/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/bevy_to_workflow/bvh_to_workflow/process_job.rs @@ -246,7 +246,7 @@ pub async fn process_job(args: BvhToWorkflowJobArgs<'_>) -> Result) -> Result { v.to_string() } Err(_err) => { return Err(