Insert batch token in media_file records

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<BatchGenerationEntity>,
+  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<BatchGenerationToken> {
-  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

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<ModelWeightToken>,
     pub maybe_origin_filename: Option<String>,
 
-    // 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<u64> = None;
     let mut maybe_creator_category_synthetic_id : Option<u64> = 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, 

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<JobSuccessRes
     maybe_creator_file_synthetic_id_category: IdCategory::MediaFile,
     maybe_creator_category_synthetic_id_category: IdCategory::MocapResult,
     maybe_mod_user_token: None,
-    is_batch_generated: false,
+    maybe_batch_token: None,
     is_generated_on_prem: false,
     checksum_sha2: &file_checksum,
     generated_by_worker: Some(&args.job_dependencies.job.info.container.hostname),

crates/service/job/inference_job/src/job/job_types/format_conversion/fbx_to_gltf/process_job.rs

@@ -206,7 +206,7 @@ pub async fn process_job(args: FbxToGltfJobArgs<'_>) -> Result<JobSuccessResult,
     maybe_creator_file_synthetic_id_category: IdCategory::MediaFile,
     maybe_creator_category_synthetic_id_category: IdCategory::MocapResult,
     maybe_mod_user_token: None,
-    is_batch_generated: false,
+    maybe_batch_token: None,
     is_generated_on_prem: false,
     checksum_sha2: &upload_details.checksum,
     generated_by_worker: Some(&args.job_dependencies.job.info.container.hostname),

crates/service/job/inference_job/src/job/job_types/image_generation/sd/process_job_inference.rs

@@ -34,6 +34,7 @@ use mysql_queries::queries::model_weights::create::create_weight::{
 };
 use mysql_queries::queries::model_weights::get::get_weight::get_weight_by_token;
 use tokens::tokens::anonymous_visitor_tracking::AnonymousVisitorTrackingToken;
+use tokens::tokens::batch_generations::BatchGenerationToken;
 use tokens::tokens::model_weights::ModelWeightToken;
 use tokens::tokens::users::UserToken;
 
@@ -273,6 +274,8 @@ pub async fn process_job_inference(
     }
   };
 
+  let batch_token = BatchGenerationToken::generate();
+
   let mut maybe_first_media_file_token = None;
 
   for i in 0..sd_args.maybe_batch_count.unwrap_or(1) {
@@ -296,7 +299,6 @@ pub async fn process_job_inference(
       }
     };
 
-
     // extra_file_modification_info: todo!(), // JSON ENCODED STRUCT
     let (media_file_token, _id) = insert_media_file_generic(InsertArgs {
       pool: mysql_pool,
@@ -307,7 +309,7 @@ pub async fn process_job_inference(
       maybe_origin_model_type: Some(MediaFileOriginModelType::StableDiffusion15),
       maybe_origin_model_token: Some(sd_model_weight_token.clone()),
       maybe_origin_filename: Some(file_path),
-      is_batch_generated: true,
+      maybe_batch_token: Some(&batch_token),
       maybe_mime_type: Some(metadata.mimetype.as_ref()),
       file_size_bytes: metadata.file_size_bytes,
       maybe_duration_millis: Some(inference_duration.as_millis() as u64),
@@ -343,10 +345,11 @@ pub async fn process_job_inference(
 
   let batch_token_result = insert_batch_generation_records(InsertBatchArgs {
     entries,
+    maybe_existing_batch_token: Some(&batch_token),
     transaction: &mut transaction,
   }).await;
 
-  let batch_token = match batch_token_result {
+  let _batch_token = match batch_token_result {
     Ok(v) => { v.to_string() }
     Err(_err) => {
       return Err(