From 9fbecd55683856c8de87f32cbeebbd7c19a01411 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sun, 2 Apr 2023 21:21:16 -0700 Subject: [PATCH] enqueue vits uploads (in theory) --- .../web/insert_generic_download_job.rs | 11 ++-- .../download_job/enqueue_generic_download.rs | 12 ++--- .../endpoints/tts/enqueue_upload_tts_model.rs | 50 +++++++++++++------ 3 files changed, 44 insertions(+), 29 deletions(-) diff --git a/crates/lib/mysql_queries/src/queries/generic_download/web/insert_generic_download_job.rs b/crates/lib/mysql_queries/src/queries/generic_download/web/insert_generic_download_job.rs index fe8da42451..abf1238b60 100644 --- a/crates/lib/mysql_queries/src/queries/generic_download/web/insert_generic_download_job.rs +++ b/crates/lib/mysql_queries/src/queries/generic_download/web/insert_generic_download_job.rs @@ -5,8 +5,7 @@ use enums::workers::generic_download_type::GenericDownloadType; use sqlx::MySqlPool; use tokens::jobs::download::DownloadJobToken; -pub struct Args <'a> { - pub job_token: &'a DownloadJobToken, +pub struct InsertGenericDownloadJobArgs<'a> { pub uuid_idempotency_token: &'a str, pub download_type: GenericDownloadType, pub download_url: &'a str, @@ -17,7 +16,9 @@ pub struct Args <'a> { pub mysql_pool: &'a MySqlPool, } -pub async fn insert_generic_download_job(args: Args<'_>) -> AnyhowResult { +pub async fn insert_generic_download_job(args: InsertGenericDownloadJobArgs<'_>) -> AnyhowResult<(DownloadJobToken, u64)> { + // This token is returned to the client. + let job_token = DownloadJobToken::generate(); let query = sqlx::query!( r#" @@ -33,7 +34,7 @@ SET creator_set_visibility = ?, status = "pending" "#, - args.job_token, + &job_token, args.uuid_idempotency_token, args.download_type.to_str(), args.download_url, @@ -55,5 +56,5 @@ SET } }; - Ok(record_id) + Ok((job_token, record_id)) } diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/download_job/enqueue_generic_download.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/download_job/enqueue_generic_download.rs index 74dc33c2e6..2cd1159c69 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/download_job/enqueue_generic_download.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/download_job/enqueue_generic_download.rs @@ -10,13 +10,13 @@ use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; use crate::server_state::ServerState; use crate::validations::model_uploads::validate_model_title; use crate::validations::validate_idempotency_token_format::validate_idempotency_token_format; -use mysql_queries::queries::generic_download::web::insert_generic_download_job::{Args, insert_generic_download_job}; -use mysql_queries::tokens::Tokens; use enums::common::visibility::Visibility; use enums::workers::generic_download_type::GenericDownloadType; use http_server_common::request::get_request_ip::get_request_ip; use http_server_common::response::serialize_as_json_error::serialize_as_json_error; use log::{info, warn, log}; +use mysql_queries::queries::generic_download::web::insert_generic_download_job::{InsertGenericDownloadJobArgs, insert_generic_download_job}; +use mysql_queries::tokens::Tokens; use regex::Regex; use sqlx::error::DatabaseError; use sqlx::error::Error::Database; @@ -128,11 +128,7 @@ pub async fn enqueue_generic_download_handler( } } - // This token is returned to the client. - let job_token = DownloadJobToken::generate(); - - let record_id = insert_generic_download_job(Args { - job_token: &job_token, + let (job_token, record_id) = insert_generic_download_job(InsertGenericDownloadJobArgs { uuid_idempotency_token: &uuid, download_type: request.generic_download_type, download_url: &download_url, @@ -148,7 +144,7 @@ pub async fn enqueue_generic_download_handler( EnqueueGenericDownloadError::ServerError })?; - info!("new generic download job id: {}", record_id); + info!("new generic download job id: {}, token: {}", record_id, job_token); server_state.firehose_publisher.enqueue_generic_download(&user_session.user_token, job_token.as_str()) .await diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/tts/enqueue_upload_tts_model.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/tts/enqueue_upload_tts_model.rs index 50b346da9f..49528eb8ba 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/tts/enqueue_upload_tts_model.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/tts/enqueue_upload_tts_model.rs @@ -17,8 +17,9 @@ use log::warn; use std::fmt; use std::sync::Arc; use enums::workers::generic_download_type::GenericDownloadType; +use mysql_queries::queries::generic_download::web::insert_generic_download_job::{insert_generic_download_job, InsertGenericDownloadJobArgs}; -#[derive(Deserialize)] +#[derive(Deserialize, Copy, Clone)] pub enum SupportedTtsModelType { /// tacotron2 Tacotron2, @@ -139,50 +140,67 @@ pub async fn upload_tts_model_handler( } } - let mut supported_tts_model_type = request.tts_model_type + let supported_tts_model_type = request.tts_model_type .clone() .unwrap_or(SupportedTtsModelType::Tacotron2); - let mut job_token = ""; + let mut job_token = "".to_string(); // TODO/FIXME: This is messy. match supported_tts_model_type { SupportedTtsModelType::Tacotron2 => { - - let tts_model_type = "tacotron2".to_string(); - let creator_set_visibility = "public".to_string(); + // NB: This is the legacy upload path. In the future, all TT2 models should be generic uploads. // This token is returned to the client. let maybe_new_job_token = insert_tts_model_upload_job(InsertTtsModelUploadJobArgs { uuid: &uuid, creator_user_token: &user_session.user_token, creator_ip_address: &ip_address, - creator_set_visibility: &creator_set_visibility, + creator_set_visibility: "public", // TODO: Creator set preference. title: &title, - tts_model_type: &tts_model_type, + tts_model_type: "tacotron2", download_url: &download_url, mysql_pool: &server_state.mysql_pool, }).await; - let job_token = match maybe_new_job_token { + job_token = match maybe_new_job_token { Ok(token) => token, Err(err) => { warn!("Error inserting new job record: {:?}", err); return Err(UploadTtsModelError::ServerError); } }; - - server_state.firehose_publisher.enqueue_tts_model_upload(&user_session.user_token, &job_token) - .await - .map_err(|e| { - warn!("error publishing event: {:?}", e); - UploadTtsModelError::ServerError - })?; } SupportedTtsModelType::Vits => { + // NB: This is the new upload path. In the future, all models should use this path. + // When that's done, the code can be simplified. + // **ACTUALLY**, all future downloads should use the generic download job endpoint instead of this one! + let (download_job_token, _record_id)= insert_generic_download_job(InsertGenericDownloadJobArgs { + uuid_idempotency_token: &uuid, + download_type: GenericDownloadType::Vits, + download_url: &download_url, + title: &title, + creator_user_token: &user_session.user_token, + creator_ip_address: &ip_address, + creator_set_visibility: Visibility::Public, // TODO: Creator set preference + mysql_pool: &server_state.mysql_pool, + }) + .await + .map_err(|err| { + warn!("New generic download creation DB error (for TTS model): {:?}", err); + UploadTtsModelError::ServerError + })?; + job_token = download_job_token.to_string(); } } + server_state.firehose_publisher.enqueue_tts_model_upload(&user_session.user_token, &job_token) + .await + .map_err(|e| { + warn!("error publishing event: {:?}", e); + UploadTtsModelError::ServerError + })?; + let response = UploadTtsModelSuccessResponse { success: true, job_token: job_token.to_string(),