replace mysql_pool with connection reuse in style transfer

crates/lib/mysql_queries/src/queries/media_files/get/batch_get_media_files_by_tokens.rs

@@ -16,6 +16,7 @@ use tokens::tokens::users::UserToken;
 
 use crate::helpers::boolean_converters::i8_to_bool;
 use crate::payloads::prompt_args::prompt_inner_payload::PromptInnerPayload;
+use crate::utils::transactor::Transactor;
 
 #[derive(Serialize, Clone)]
 pub struct MediaFilesByTokensRecord {
@@ -73,28 +74,35 @@ pub async fn batch_get_media_files_by_tokens(
   media_file_tokens: &[MediaFileToken],
   can_see_deleted: bool
 ) -> AnyhowResult<Vec<MediaFilesByTokensRecord>> {
+  batch_get_media_files_by_tokens_with_transactor(
+    Transactor::Pool { pool: mysql_pool },
+    media_file_tokens,
+    can_see_deleted
+  ).await
+}
 
+pub async fn batch_get_media_files_by_tokens_with_transactor(
+  transactor: Transactor<'_, '_>,
+  media_file_tokens: &[MediaFileToken],
+  can_see_deleted: bool
+) -> AnyhowResult<Vec<MediaFilesByTokensRecord>> {
   if media_file_tokens.is_empty() {
     // NB: We should always eagerly return, but if we don't, the query builder will build an
     // invalid query.
     return Ok(vec![]);
   }
 
-  let mut connection = mysql_pool.acquire().await?;
-
-  let raw_media_files: Vec<RawMediaFileJoinUser> = get_raw_media_files_by_tokens(&mut connection, media_file_tokens, can_see_deleted).await?;
+  let raw_media_files: Vec<RawMediaFileJoinUser> = get_raw_media_files_by_tokens(transactor, media_file_tokens, can_see_deleted).await?;
 
   Ok(map_to_media_files(raw_media_files))
 }
 
 async fn get_raw_media_files_by_tokens(
-  connection: &mut MySqlConnection,
+  mut transactor: Transactor<'_, '_>,
   media_file_tokens: &[MediaFileToken],
   can_see_deleted: bool
 ) -> AnyhowResult<Vec<RawMediaFileJoinUser>> {
 
-  let connection = connection.acquire().await?;
-
   let mut query_builder : QueryBuilder<MySql> = if can_see_deleted {
     QueryBuilder::new(r#"
       SELECT
@@ -234,7 +242,17 @@ async fn get_raw_media_files_by_tokens(
 
   let query = query_builder.build_query_as::<RawMediaFileJoinUser>();
 
-  let results = query.fetch_all(connection).await?;
+  let results = match transactor {
+    Transactor::Pool { pool } => {
+      query.fetch_all(pool).await?
+    },
+    Transactor::Connection { connection } => {
+      query.fetch_all(connection).await?
+    },
+    Transactor::Transaction { transaction } => {
+      query.fetch_all(&mut **transaction).await?
+    },
+  };
 
   Ok(results)
 }

crates/service/job/inference_job/src/job/job_types/workflow/video_style_transfer/process_video_style_transfer_job.rs

@@ -265,7 +265,7 @@ pub async fn process_video_style_transfer_job(deps: &JobDependencies, job: &Avai
     let mut videos = download_input_videos(DownloadInputVideoArgs {
         job_args: &job_args,
         comfy_dirs: &comfy_dirs,
-        mysql_pool: &deps.db.mysql_pool,
+        mysql_connection: &mut mysql_connection,
         remote_cloud_file_client: &remote_cloud_file_client,
     }).await?;
 

crates/service/job/inference_job/src/job/job_types/workflow/video_style_transfer/steps/download_input_videos.rs

@@ -2,14 +2,15 @@ use std::path::{Path, PathBuf};
 
 use anyhow::anyhow;
 use log::{error, info};
-use sqlx::MySqlPool;
+use sqlx::{MySqlConnection, MySqlPool};
 
 use bucket_paths::legacy::typified_paths::public::media_files::bucket_file_path::MediaFileBucketPath;
 use cloud_storage::remote_file_manager::remote_cloud_file_manager::RemoteCloudFileClient;
 use errors::AnyhowResult;
 use filesys::path_to_string::path_to_string;
-use mysql_queries::queries::media_files::get::batch_get_media_files_by_tokens::{batch_get_media_files_by_tokens, MediaFilesByTokensRecord};
-use mysql_queries::queries::media_files::get::get_media_file::{get_media_file, MediaFile};
+use mysql_queries::queries::media_files::get::batch_get_media_files_by_tokens::{batch_get_media_files_by_tokens, batch_get_media_files_by_tokens_with_transactor, MediaFilesByTokensRecord};
+use mysql_queries::queries::media_files::get::get_media_file::{get_media_file, get_media_file_with_transactor, MediaFile};
+use mysql_queries::utils::transactor::Transactor;
 use tokens::tokens::media_files::MediaFileToken;
 use videos::ffprobe_get_dimensions::ffprobe_get_dimensions;
 
@@ -22,21 +23,21 @@ use crate::job::job_types::workflow::video_style_transfer::util::video_pathing::
 pub struct DownloadInputVideoArgs<'a> {
   pub job_args: &'a JobArgs<'a>,
   pub comfy_dirs: &'a ComfyDirs,
-  pub mysql_pool: &'a MySqlPool,
+  pub mysql_connection: &'a mut MySqlConnection,
   pub remote_cloud_file_client: &'a RemoteCloudFileClient,
 }
 
 pub async fn download_input_videos(
-  args: DownloadInputVideoArgs<'_>
+  mut args: DownloadInputVideoArgs<'_>,
 ) -> Result<VideoPathing, ProcessSingleJobError> {
-  let video_downloads = download_primary_video(&args).await?;
-  let video_downloads = maybe_download_secondary_videos(video_downloads, &args).await?;
+  let video_downloads = download_primary_video(&mut args).await?;
+  let video_downloads = maybe_download_secondary_videos(video_downloads, &mut args).await?;
   Ok(video_downloads)
 }
 
 // TODO: Consolidate primary video download with secondary download logic.
 async fn download_primary_video(
-  args: &DownloadInputVideoArgs<'_>,
+  args: &mut DownloadInputVideoArgs<'_>,
 ) -> Result<VideoPathing, ProcessSingleJobError> {
   let input_media_file_token = match args.job_args.maybe_input_file {
     None => return Err(ProcessSingleJobError::InvalidJob(anyhow!("No primary input video file provided"))),
@@ -45,10 +46,10 @@ async fn download_primary_video(
 
   info!("Querying primary input media file by token: {:?} ...", &input_media_file_token);
 
-  let mut input_media_file =  get_media_file(
+  let mut input_media_file =  get_media_file_with_transactor(
     &input_media_file_token,
     false,
-    args.mysql_pool
+    Transactor::for_connection(args.mysql_connection),
   ).await?.ok_or_else(|| {
     error!("primary input media_file not found: {:?}", &input_media_file_token);
     ProcessSingleJobError::Other(anyhow!("primary input media_file not found: {:?}", &input_media_file_token))
@@ -62,10 +63,10 @@ async fn download_primary_video(
     // TODO(bt,2024-05-14): Perhaps fail open and use the first media file if the original
     //  isn't found?
     info!("Looking up original style transfer source media file...");
-    input_media_file =  get_media_file(
+    input_media_file =  get_media_file_with_transactor(
       &source_media_file_token,
       true, // NB: In case the original was deleted, allow this to continue.
-      args.mysql_pool
+      Transactor::for_connection(args.mysql_connection),
     ).await?.ok_or_else(|| {
       error!("source input media_file not found: {:?}", &input_media_file_token);
       ProcessSingleJobError::Other(anyhow!("source input media_file not found: {:?}",
@@ -115,7 +116,7 @@ enum SecondaryVideoType {
 
 async fn maybe_download_secondary_videos(
   mut video_downloads: VideoPathing,
-  args: &DownloadInputVideoArgs<'_>
+  args: &mut DownloadInputVideoArgs<'_>,
 ) -> Result<VideoPathing, ProcessSingleJobError> {
   const CAN_SEE_DELETED : bool = true; // We don't need to care about the deleted flag.
 
@@ -137,7 +138,10 @@ async fn maybe_download_secondary_videos(
     return Ok(video_downloads);
   }
 
-  let results = batch_get_media_files_by_tokens(args.mysql_pool, &tokens, CAN_SEE_DELETED)
+  let results = batch_get_media_files_by_tokens_with_transactor(
+    Transactor::for_connection(args.mysql_connection),
+    &tokens,
+    CAN_SEE_DELETED)
       .await
       .map_err(|err| {
         ProcessSingleJobError::Other(anyhow!("error querying secondary videos: {:?}", &err))