enqueue vits uploads (in theory)

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<u64> {
+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))
 }

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

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(),