From 767a2a38dad29a1fdcc8a54084008569e4dddca1 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sat, 28 Sep 2024 00:13:31 -0400 Subject: [PATCH] replace mysql_pool with connection reuse in style transfer --- .../get/batch_get_media_files_by_tokens.rs | 32 +++++++++++++++---- .../process_video_style_transfer_job.rs | 2 +- .../steps/download_input_videos.rs | 32 +++++++++++-------- 3 files changed, 44 insertions(+), 22 deletions(-) diff --git a/crates/lib/mysql_queries/src/queries/media_files/get/batch_get_media_files_by_tokens.rs b/crates/lib/mysql_queries/src/queries/media_files/get/batch_get_media_files_by_tokens.rs index 59cf89f3ab..485d41c666 100644 --- a/crates/lib/mysql_queries/src/queries/media_files/get/batch_get_media_files_by_tokens.rs +++ b/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> { + 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> { 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 = get_raw_media_files_by_tokens(&mut connection, media_file_tokens, can_see_deleted).await?; + let raw_media_files: Vec = 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> { - let connection = connection.acquire().await?; - let mut query_builder : QueryBuilder = 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::(); - 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) } diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/video_style_transfer/process_video_style_transfer_job.rs b/crates/service/job/inference_job/src/job/job_types/workflow/video_style_transfer/process_video_style_transfer_job.rs index 97e02893f7..8368f317ae 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/video_style_transfer/process_video_style_transfer_job.rs +++ b/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?; diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/video_style_transfer/steps/download_input_videos.rs b/crates/service/job/inference_job/src/job/job_types/workflow/video_style_transfer/steps/download_input_videos.rs index cceb71b731..fdbd31c743 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/video_style_transfer/steps/download_input_videos.rs +++ b/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 { - 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 { 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 { 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))