insert batch tokens on batch generation

.sqlx/query-7e223e284b992b6eaf31bec057bae00fc4e085d68b064e370f0e3d92b82049b2.json

@@ -0,0 +1,12 @@
+{
+  "db_name": "MySQL",
+  "query": "\nUPDATE generic_inference_jobs\nSET\n  status = ?,\n  on_success_result_entity_type = ?,\n  on_success_result_entity_token = ?,\n  on_success_result_batch_token = ?,\n  failure_reason = NULL,\n  internal_debugging_failure_reason = NULL,\n  success_execution_millis = NULL,\n  success_inference_execution_millis = NULL,\n  retry_at = NULL,\n  successfully_completed_at = NOW()\nWHERE token = ?\n        ",
+  "describe": {
+    "columns": [],
+    "parameters": {
+      "Right": 5
+    },
+    "nullable": []
+  },
+  "hash": "7e223e284b992b6eaf31bec057bae00fc4e085d68b064e370f0e3d92b82049b2"
+}

.sqlx/query-f68a8917e639418a0bb952bef606c96299001e143c745ee30df99b72d1235d35.json

@@ -1,12 +0,0 @@
-{
-  "db_name": "MySQL",
-  "query": "\nUPDATE generic_inference_jobs\nSET\n  status = ?,\n  on_success_result_entity_type = ?,\n  on_success_result_entity_token = ?,\n  failure_reason = NULL,\n  internal_debugging_failure_reason = NULL,\n  success_execution_millis = NULL,\n  success_inference_execution_millis = NULL,\n  retry_at = NULL,\n  successfully_completed_at = NOW()\nWHERE token = ?\n        ",
-  "describe": {
-    "columns": [],
-    "parameters": {
-      "Right": 4
-    },
-    "nullable": []
-  },
-  "hash": "f68a8917e639418a0bb952bef606c96299001e143c745ee30df99b72d1235d35"
-}

crates/schema/database/mysql_queries/src/queries/generic_inference/fal/mark_fal_generic_inference_job_successfully_done.rs

@@ -1,19 +1,27 @@
 use anyhow::anyhow;
-use sqlx::MySqlPool;
+use sqlx::{Executor, MySql};
+use std::marker::PhantomData;
 
 use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType;
 use enums::common::job_status_plus::JobStatusPlus;
 use errors::AnyhowResult;
+use tokens::tokens::batch_generations::BatchGenerationToken;
 use tokens::tokens::generic_inference_jobs::InferenceJobToken;
 use tokens::tokens::media_files::MediaFileToken;
 
-// TODO: Maybe in the future we set a batch_token instead of media token
-
-pub async fn mark_fal_generic_inference_job_successfully_done(
-  pool: &MySqlPool,
-  job_token: &InferenceJobToken,
-  media_file_token: MediaFileToken,
+pub struct MarkJobArgs<'e, 'c, E>
+  where E: 'e + Executor<'c, Database = MySql>
+{
+  pub job_token: &'e InferenceJobToken,
+  pub media_file_token: &'e MediaFileToken,
+  pub maybe_batch_token: Option<&'e BatchGenerationToken>,
+  pub mysql_executor: E,
+  pub phantom: PhantomData<&'c E>,
+}
+pub async fn mark_fal_generic_inference_job_successfully_done<'e, 'c : 'e, E>(
+  args: MarkJobArgs<'e, 'c, E>
 ) -> AnyhowResult<()>
+  where E: 'e + Executor<'c, Database = MySql>
 {
   const STATUS : JobStatusPlus = JobStatusPlus::CompleteSuccess;
   const RESULT_TYPE : InferenceResultType = InferenceResultType::MediaFile;
@@ -25,6 +33,7 @@ SET
   status = ?,
   on_success_result_entity_type = ?,
   on_success_result_entity_token = ?,
+  on_success_result_batch_token = ?,
   failure_reason = NULL,
   internal_debugging_failure_reason = NULL,
   success_execution_millis = NULL,
@@ -35,10 +44,11 @@ WHERE token = ?
         "#,
         STATUS.to_str(),
         RESULT_TYPE.to_str(),
-        media_file_token,
-        job_token,
+        args.media_file_token,
+        args.maybe_batch_token.map(|t| t.as_str()),
+        args.job_token,
     )
-      .execute(pool)
+      .execute(args.mysql_executor)
       .await;
 
   match query_result {

crates/schema/database/mysql_queries/src/queries/media_files/create/insert_builder/media_file_insert_builder.rs

@@ -10,6 +10,7 @@ use enums::by_table::media_files::media_file_type::MediaFileType;
 use enums::common::visibility::Visibility;
 use sqlx::MySqlPool;
 use tokens::tokens::anonymous_visitor_tracking::AnonymousVisitorTrackingToken;
+use tokens::tokens::batch_generations::BatchGenerationToken;
 use tokens::tokens::media_files::MediaFileToken;
 use tokens::tokens::prompts::PromptToken;
 use tokens::tokens::users::UserToken;
@@ -67,7 +68,7 @@ pub struct MediaFileInsertBuilder {
   // maybe_prompt_token: Option<&'a PromptToken>,
 
   // // If batch generated, this is the batch token.
-  // maybe_batch_token: Option<&'a BatchGenerationToken>,
+  maybe_batch_generation_token: Option<BatchGenerationToken>,
 
   // Storage details
   public_bucket_directory_hash: Option<MediaFileBucketPath>, // NB: Non-nullable field(s)
@@ -100,6 +101,7 @@ impl MediaFileInsertBuilder {
       maybe_frame_height: None,
       checksum_sha2: None,
       maybe_engine_category: None,
+      maybe_batch_generation_token: None,
       public_bucket_directory_hash: None,
       maybe_prompt_token: None,
     }
@@ -228,7 +230,12 @@ impl MediaFileInsertBuilder {
   // TODO: maybe_text_transcript
   // TODO: maybe_scene_source_media_file_token
   // TODO: maybe_prompt_token
-  // TODO: maybe_batch_token
+
+  pub fn maybe_batch_generation_token(mut self, maybe_batch_token: Option<&BatchGenerationToken>) -> Self {
+    self.maybe_batch_generation_token = maybe_batch_token
+        .map(|token| token.clone());
+    self
+  }
 
   pub fn public_bucket_directory_hash(mut self, public_bucket_directory_hash: &MediaFileBucketPath) -> Self {
     self.public_bucket_directory_hash = Some(public_bucket_directory_hash.clone());
@@ -281,7 +288,7 @@ impl MediaFileInsertBuilder {
       maybe_text_transcript: None, // TODO
       maybe_scene_source_media_file_token: None, // TODO
       maybe_prompt_token: self.maybe_prompt_token.as_ref(),
-      maybe_batch_token: None, // TODO
+      maybe_batch_token: self.maybe_batch_generation_token.as_ref(),
       public_bucket_directory_hash: bucket_path.get_object_hash(),
       maybe_public_bucket_prefix: bucket_path.get_optional_prefix(),
       maybe_public_bucket_extension: bucket_path.get_optional_extension(),

crates/service/web/storyteller_web/src/http_server/endpoints/webhooks/fal_webhook_handler.rs

@@ -15,7 +15,7 @@ use http_server_common::response::response_success_helpers::SimpleGenericJsonSuc
 use http_server_common::response::serialize_as_json_error::serialize_as_json_error;
 use log::{info, warn};
 use mysql_queries::queries::generic_inference::fal::get_inference_job_by_fal_id::get_inference_job_by_fal_id;
-use mysql_queries::queries::generic_inference::fal::mark_fal_generic_inference_job_successfully_done::mark_fal_generic_inference_job_successfully_done;
+use mysql_queries::queries::generic_inference::fal::mark_fal_generic_inference_job_successfully_done::{mark_fal_generic_inference_job_successfully_done, MarkJobArgs};
 use serde_json::Value;
 use utoipa::ToSchema;
 // 1. tauri --> hit endpoint to enqueue
@@ -95,7 +95,7 @@ impl From<anyhow::Error> for FalWebhookError {
 /// Fal webhook
 #[utoipa::path(
   post,
-  tag = "Webooks",
+  tag = "Webhooks",
   path = "/v1/webhooks/fal",
   responses(
     (status = 200, description = "Success", body = SimpleGenericJsonSuccess),
@@ -143,6 +143,7 @@ pub async fn fal_webhook_handler(
   info!("Fal webhook job record: {:?}", job);
 
  let mut maybe_media_token = None;
+ let mut maybe_batch_token = None;
 
   if let Some(payload_obj) = payload.as_object() {
     if payload_obj.contains_key("image") {
@@ -150,7 +151,7 @@ pub async fn fal_webhook_handler(
       let token = handle_image_payload(payload_obj, &job, &server_state).await?;
       maybe_media_token = Some(token);
     } else if payload_obj.contains_key("images") {
-      maybe_media_token = handle_images_payload(payload_obj, &job, &server_state).await?;
+      (maybe_media_token, maybe_batch_token) = handle_images_payload(payload_obj, &job, &server_state).await?;
     } else if payload_obj.contains_key("video") {
       info!("Handling video payload for job: {:?}", job.job_token);
       let token = handle_video_payload(payload_obj, &job, &server_state).await?;
@@ -165,11 +166,13 @@ pub async fn fal_webhook_handler(
   if let Some(media_token) = maybe_media_token {
     info!("Media file token: {:?}", media_token);
     // TODO: Update job metadata.
-    mark_fal_generic_inference_job_successfully_done(
-      &server_state.mysql_pool,
-      &job.job_token,
-      media_token,
-    ).await.map_err(|err| {
+    mark_fal_generic_inference_job_successfully_done(MarkJobArgs {
+      job_token: &job.job_token,
+      media_file_token: &media_token,
+      maybe_batch_token: maybe_batch_token.as_ref(),
+      mysql_executor: &server_state.mysql_pool,
+      phantom: Default::default(),
+    }).await.map_err(|err| {
       warn!("Error marking job as successfully done: {:?}", err);
       FalWebhookError::ServerError
     })?;

crates/service/web/storyteller_web/src/http_server/endpoints/webhooks/handle_images_payload.rs

@@ -15,6 +15,7 @@ use mimetypes::mimetype_info::mimetype_info::MimetypeInfo;
 use mysql_queries::queries::generic_inference::fal::get_inference_job_by_fal_id::FalJobDetails;
 use mysql_queries::queries::media_files::create::insert_builder::media_file_insert_builder::MediaFileInsertBuilder;
 use serde_json::{Map, Value};
+use tokens::tokens::batch_generations::BatchGenerationToken;
 use tokens::tokens::media_files::MediaFileToken;
 
 const PREFIX : Option<&str> = Some("artcraft_");
@@ -33,8 +34,8 @@ pub async fn handle_images_payload(
   payload: &Map<String, Value>,
   job: &FalJobDetails,
   server_state: &ServerState,
-) -> AnyhowResult<Option<MediaFileToken>> {
-  
+) -> AnyhowResult<(Option<MediaFileToken>, Option<BatchGenerationToken>)> {
+
   let images_value = payload.get("images")
       .ok_or_else(|| anyhow!("no `images` key in payload"))?;
 
@@ -43,22 +44,31 @@ pub async fn handle_images_payload(
   let images: Vec<FalWebhookImage> = serde_json::from_value(images_value.clone())?;
   
   let mut maybe_media_token = None;
-  
+
+  // NB: We are not going to create `batch_generations` table records. We don't need them.
+  // The foreign key in `media_files` is enough to look up the rest of the batch.
+  let mut maybe_batch_token = None;
+
+  if images.len() > 1 {
+    maybe_batch_token = Some(BatchGenerationToken::generate());
+  }
+
   for image in images {
-    let media_token = upload_image(job, server_state, &image).await?;
+    let media_token = upload_image(job, server_state, &image, maybe_batch_token.as_ref()).await?;
     
     if maybe_media_token.is_none() {
       maybe_media_token = Some(media_token); // Set the first media token
     }
   }
   
-  Ok(maybe_media_token)
+  Ok((maybe_media_token, maybe_batch_token))
 }
 
 async fn upload_image(
   job: &FalJobDetails,
   server_state: &ServerState,
   image: &FalWebhookImage,
+  maybe_batch_token: Option<&BatchGenerationToken>,
 ) -> AnyhowResult<MediaFileToken> {
   let image_url = image.url
       .as_deref()
@@ -98,6 +108,7 @@ async fn upload_image(
         image,
         &file_bytes,
         mimetype_info,
+        maybe_batch_token,
       ).await
     }
     _ => {
@@ -107,6 +118,7 @@ async fn upload_image(
         image,
         &file_bytes,
         mimetype_info,
+        maybe_batch_token,
       ).await
     }
   }
@@ -118,6 +130,7 @@ async fn upload_image_bytes(
   image: &FalWebhookImage,
   file_bytes: &[u8],
   mimetype_info: MimetypeInfo,
+  maybe_batch_token: Option<&BatchGenerationToken>,
 ) -> AnyhowResult<MediaFileToken> {
 
   let mime_type = mimetype_info.mime_type();
@@ -145,7 +158,7 @@ async fn upload_image_bytes(
     &mime_type)
       .await?;
 
-  let media_token = MediaFileInsertBuilder::new()
+  let mut query_builder = MediaFileInsertBuilder::new()
       .maybe_creator_user(job.maybe_creator_user_token.as_ref())
       .maybe_creator_anonymous_visitor(job.maybe_creator_anonymous_visitor_token.as_ref())
       .creator_ip_address(&job.creator_ip_address)
@@ -159,11 +172,18 @@ async fn upload_image_bytes(
       .frame_width(image_info.width())
       .frame_height(image_info.height())
       .maybe_prompt_token(job.maybe_prompt_token.as_ref())
-      .checksum_sha2(&file_hash)
+      .checksum_sha2(&file_hash);
+
+  if let Some(batch_token) = maybe_batch_token {
+    query_builder = query_builder.maybe_batch_generation_token(Some(batch_token));
+  }
+
+  let media_token = query_builder
       .insert_pool(&server_state.mysql_pool)
       .await?;
 
-  info!("Image media file uploaded with token: {}", media_token);
+  info!("Image media file uploaded with token: {} ; possible batch token: {:?}",
+    media_token, maybe_batch_token);
 
   Ok(media_token)
 }