From 08c05cbcf20db167f40e90975eb3eff519618428 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sat, 6 Aug 2022 00:41:18 -0400 Subject: [PATCH] Schema changes, endpoints, routes, queries, etc. --- Cargo.lock | 2 + .../up.sql | 7 +- crates/lib/database_queries/Cargo.toml | 1 + .../src/mediators/firehose_publisher.rs | 25 +++ .../get_generic_download_job_status.rs | 59 ++++++ .../insert_generic_download_job.rs | 58 ++++++ .../src/queries/generic_download/mod.rs | 2 + .../lib/database_queries/src/queries/mod.rs | 1 + crates/lib/redis_common/src/redis_keys.rs | 6 + ...pload_type.rs => generic_download_type.rs} | 10 +- crates/lib/reusable_types/src/lib.rs | 4 +- crates/service/storyteller_web/Cargo.toml | 1 + .../download_job/enqueue_generic_download.rs | 178 ++++++------------ .../get_generic_upload_job_status.rs | 121 ++++++------ .../http_server/endpoints/download_job/mod.rs | 1 - crates/service/storyteller_web/src/routes.rs | 31 +++ .../storyteller_web/src/validations/mod.rs | 1 + .../validate_idempotency_token_format.rs | 43 +++++ 18 files changed, 360 insertions(+), 191 deletions(-) create mode 100644 crates/lib/database_queries/src/queries/generic_download/get_generic_download_job_status.rs create mode 100644 crates/lib/database_queries/src/queries/generic_download/insert_generic_download_job.rs create mode 100644 crates/lib/database_queries/src/queries/generic_download/mod.rs rename crates/lib/reusable_types/src/{generic_upload_type.rs => generic_download_type.rs} (56%) create mode 100644 crates/service/storyteller_web/src/validations/validate_idempotency_token_format.rs diff --git a/Cargo.lock b/Cargo.lock index 9f70f7d31e..cf3ee9a5c0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1297,6 +1297,7 @@ dependencies = [ "log", "once_cell", "regex", + "reusable_types", "serde", "serde_derive", "serde_json", @@ -3893,6 +3894,7 @@ dependencies = [ "memory_caching", "native-tls", "newrelic-telemetry", + "once_cell", "pulldown-cmark", "r2d2_redis", "rand 0.8.4", diff --git a/_migrations/2022-08-05-202017_generic_download_jobs/up.sql b/_migrations/2022-08-05-202017_generic_download_jobs/up.sql index 84c377c6a4..7aef66c9f8 100644 --- a/_migrations/2022-08-05-202017_generic_download_jobs/up.sql +++ b/_migrations/2022-08-05-202017_generic_download_jobs/up.sql @@ -18,7 +18,11 @@ CREATE TABLE generic_download_jobs ( -- If the job completes successfully, this is the result token. -- This is only populated on a successful result. - on_success_result_token VARCHAR(32) DEFAULT NULL, + on_success_downloaded_entity_token VARCHAR(32) DEFAULT NULL, + + -- The type of the object will vary based on the type of the upload, + -- and we may include heuristics that auto-detect types in the future + on_success_downloaded_entity_type VARCHAR(32) DEFAULT NULL, -- ========== UPLOAD DETAILS ========== @@ -98,7 +102,6 @@ CREATE TABLE generic_download_jobs ( PRIMARY KEY (id), UNIQUE KEY (token), UNIQUE KEY (uuid_idempotency_token), - KEY fk_on_success_result_token (on_success_result_token), KEY fk_on_download_type (download_type), KEY fk_creator_user_token (creator_user_token), KEY index_status (status), diff --git a/crates/lib/database_queries/Cargo.toml b/crates/lib/database_queries/Cargo.toml index ab092d51d2..1b2bd703bc 100644 --- a/crates/lib/database_queries/Cargo.toml +++ b/crates/lib/database_queries/Cargo.toml @@ -17,6 +17,7 @@ path = "src/lib.rs" # Internal config = { path = "../config" } container_common = { path = "../container_common" } +reusable_types = { path = "../reusable_types" } # External anyhow = "1.0.40" diff --git a/crates/lib/database_queries/src/mediators/firehose_publisher.rs b/crates/lib/database_queries/src/mediators/firehose_publisher.rs index 2499a2b144..3b15833fdd 100644 --- a/crates/lib/database_queries/src/mediators/firehose_publisher.rs +++ b/crates/lib/database_queries/src/mediators/firehose_publisher.rs @@ -24,6 +24,9 @@ enum FirehoseEvent { W2lInferenceStarted, W2lInferenceCompleted, + GenericDownloadStarted, + GenericDownloadCompleted, + TwitterMention, TwitterRetweet, DiscordJoin, @@ -51,6 +54,8 @@ impl FirehoseEvent { FirehoseEvent::DiscordMessage => "discord_message", FirehoseEvent::TwitchSubscribe => "twitch_subscribe", FirehoseEvent::TwitchFollow => "twitch_follow", + FirehoseEvent::GenericDownloadStarted => "generic_download_started", + FirehoseEvent::GenericDownloadCompleted => "generic_download_completed", } } } @@ -172,6 +177,26 @@ impl FirehosePublisher { Ok(()) } + pub async fn enqueue_generic_download(&self, user_token: &str, job_token: &str) -> AnyhowResult<()> { + let _record_id = self.insert( + FirehoseEvent::GenericDownloadStarted, + Some(user_token), + None, + Some(job_token) + ).await?; + Ok(()) + } + + pub async fn publish_generic_download_finished(&self, user_token: &str, template_token: &str) -> AnyhowResult<()> { + let _record_id = self.insert( + FirehoseEvent::GenericDownloadCompleted, + Some(user_token), + Some(template_token), + Some(template_token) + ).await?; + Ok(()) + } + // ======================================================================= async fn insert( diff --git a/crates/lib/database_queries/src/queries/generic_download/get_generic_download_job_status.rs b/crates/lib/database_queries/src/queries/generic_download/get_generic_download_job_status.rs new file mode 100644 index 0000000000..ae5eba70f2 --- /dev/null +++ b/crates/lib/database_queries/src/queries/generic_download/get_generic_download_job_status.rs @@ -0,0 +1,59 @@ +use anyhow::anyhow; +use chrono::{DateTime, Utc}; +use log::warn; +use sqlx::MySqlPool; +use container_common::anyhow_result::AnyhowResult; +use reusable_types::entity_visibility::EntityVisibility; +use reusable_types::generic_download_type::GenericDownloadType; + +pub struct GenericDownloadJobStatus { + pub job_token: String, + + pub status: String, + pub attempt_count: i32, + pub maybe_downloaded_entity_type: Option, + pub maybe_downloaded_entity_token: Option, + + pub created_at: DateTime, + pub updated_at: DateTime, +} + +/// Look up job status. +/// Returns Ok(None) when the record cannot be found. +pub async fn get_generic_download_job_status(job_token: &str, mysql_pool: &MySqlPool) + -> AnyhowResult> +{ + let maybe_status = sqlx::query_as!( + GenericDownloadJobStatus, + r#" +SELECT + jobs.token as job_token, + + jobs.status, + jobs.attempt_count, + jobs.on_success_downloaded_entity_token as maybe_downloaded_entity_token, + jobs.on_success_downloaded_entity_type as maybe_downloaded_entity_type, + + jobs.created_at, + jobs.updated_at + +FROM generic_download_jobs as jobs + +WHERE jobs.token = ? + "#, + job_token + ) + .fetch_one(mysql_pool) + .await; + + match maybe_status { + Ok(record) => Ok(Some(record)), + Err(err) => match err { + sqlx::Error::RowNotFound => Ok(None), + _ => { + warn!("error querying job record: {:?}", err); + Err(anyhow!("error querying job record: {:?}", err)) + } + } + } +} diff --git a/crates/lib/database_queries/src/queries/generic_download/insert_generic_download_job.rs b/crates/lib/database_queries/src/queries/generic_download/insert_generic_download_job.rs new file mode 100644 index 0000000000..cd3761a3f8 --- /dev/null +++ b/crates/lib/database_queries/src/queries/generic_download/insert_generic_download_job.rs @@ -0,0 +1,58 @@ +use anyhow::anyhow; +use sqlx::MySqlPool; +use container_common::anyhow_result::AnyhowResult; +use reusable_types::entity_visibility::EntityVisibility; +use reusable_types::generic_download_type::GenericDownloadType; + +pub struct Args <'a> { + pub job_token: &'a str, + pub uuid_idempotency_token: &'a str, + pub download_type: GenericDownloadType, + pub download_url: &'a str, + pub title: &'a str, + pub creator_user_token: &'a str, + pub creator_ip_address: &'a str, + pub creator_set_visibility: EntityVisibility, + pub mysql_pool: &'a MySqlPool, +} + +pub async fn insert_generic_download_job(args: Args<'_>) -> AnyhowResult { + + let query = sqlx::query!( + r#" +INSERT INTO generic_download_jobs +SET + token = ?, + uuid_idempotency_token = ?, + download_type = ?, + download_url = ?, + title = ?, + creator_user_token = ?, + creator_ip_address = ?, + creator_set_visibility = ?, + status = "pending" + "#, + args.job_token, + args.uuid_idempotency_token, + args.download_type, + args.download_url, + args.title, + args.creator_user_token, + args.creator_ip_address, + args.creator_set_visibility, + ); + + let query_result = query.execute(args.mysql_pool) + .await; + + let record_id = match query_result { + Ok(res) => { + res.last_insert_id() + }, + Err(err) => { + return Err(anyhow!("error inserting new generic download job: {:?}", err)); + } + }; + + Ok(record_id) +} diff --git a/crates/lib/database_queries/src/queries/generic_download/mod.rs b/crates/lib/database_queries/src/queries/generic_download/mod.rs new file mode 100644 index 0000000000..4def6e37b1 --- /dev/null +++ b/crates/lib/database_queries/src/queries/generic_download/mod.rs @@ -0,0 +1,2 @@ +pub mod get_generic_download_job_status; +pub mod insert_generic_download_job; diff --git a/crates/lib/database_queries/src/queries/mod.rs b/crates/lib/database_queries/src/queries/mod.rs index eb9ea3bf70..b3e2ed6161 100644 --- a/crates/lib/database_queries/src/queries/mod.rs +++ b/crates/lib/database_queries/src/queries/mod.rs @@ -1,4 +1,5 @@ pub mod api_tokens; +pub mod generic_download; pub mod health_check; pub mod ip_bans; pub mod model_categories; diff --git a/crates/lib/redis_common/src/redis_keys.rs b/crates/lib/redis_common/src/redis_keys.rs index 4970bf9ad7..c151d73c91 100644 --- a/crates/lib/redis_common/src/redis_keys.rs +++ b/crates/lib/redis_common/src/redis_keys.rs @@ -25,6 +25,12 @@ impl RedisKeys { format!("ttsDownloadExtraStatus:{}", inference_job_token) } + /// We write extra in-progress status information to keys. + /// These keys should have a TTL. + pub fn generic_download_extra_status_info(inference_job_token: &str) -> String { + format!("genericDownloadExtraStatus:{}", inference_job_token) + } + /// We write extra in-progress status information to keys. /// These keys should have a TTL. pub fn tts_inference_extra_status_info(inference_job_token: &str) -> String { diff --git a/crates/lib/reusable_types/src/generic_upload_type.rs b/crates/lib/reusable_types/src/generic_download_type.rs similarity index 56% rename from crates/lib/reusable_types/src/generic_upload_type.rs rename to crates/lib/reusable_types/src/generic_download_type.rs index 3cb962f880..23c179bd6a 100644 --- a/crates/lib/reusable_types/src/generic_upload_type.rs +++ b/crates/lib/reusable_types/src/generic_download_type.rs @@ -1,10 +1,10 @@ -/// Our "generic uploads" pipeline supports a wide variety of ML models and other media. -/// Each type of upload is identified by the following enum variants. +/// Our "generic downloads" pipeline supports a wide variety of ML models and other media. +/// Each type of download is identified by the following enum variants. /// These types are present in the HTTP API and database columns as serialized here. /// /// DO NOT CHANGE VALUES WITHOUT A MIGRATION STRATEGY. #[derive(Clone, Copy, Eq, PartialEq, Debug, Deserialize, Serialize, sqlx::Type)] -pub enum GenericUploadType { +pub enum GenericDownloadType { #[serde(rename = "hifigan")] #[sqlx(rename = "hifigan")] HifiGan, @@ -12,11 +12,11 @@ pub enum GenericUploadType { #[cfg(test)] mod tests { - use crate::generic_upload_type::GenericUploadType; + use crate::generic_download_type::GenericDownloadType; use crate::test_helpers::assert_serialization; #[test] fn test_serialization() { - assert_serialization(GenericUploadType::HifiGan, "hifigan"); + assert_serialization(GenericDownloadType::HifiGan, "hifigan"); } } diff --git a/crates/lib/reusable_types/src/lib.rs b/crates/lib/reusable_types/src/lib.rs index 3ba21fd256..35a956b57d 100644 --- a/crates/lib/reusable_types/src/lib.rs +++ b/crates/lib/reusable_types/src/lib.rs @@ -3,7 +3,7 @@ #![forbid(unused_must_use)] // NB: It's unsafe to not close/check some things // Okay to toggle -//#![forbid(unused_imports)] +#![forbid(unused_imports)] #![forbid(unused_mut)] #![forbid(unused_variables)] @@ -14,6 +14,6 @@ #[macro_use] extern crate serde_derive; pub mod entity_visibility; -pub mod generic_upload_type; +pub mod generic_download_type; #[cfg(test)] pub mod test_helpers; \ No newline at end of file diff --git a/crates/service/storyteller_web/Cargo.toml b/crates/service/storyteller_web/Cargo.toml index 1872b73058..3bb8eb3e6b 100644 --- a/crates/service/storyteller_web/Cargo.toml +++ b/crates/service/storyteller_web/Cargo.toml @@ -46,6 +46,7 @@ lazy_static = "1.4.0" lexical-sort = "0.3.1" lfu = "0.2.5" log = "0.4.14" +once_cell = "1.9.0" rand = "0.8.3" regex = "1.5.4" time = "0.2.26" diff --git a/crates/service/storyteller_web/src/http_server/endpoints/download_job/enqueue_generic_download.rs b/crates/service/storyteller_web/src/http_server/endpoints/download_job/enqueue_generic_download.rs index 62d33a8ce9..a7b1d0e5ec 100644 --- a/crates/service/storyteller_web/src/http_server/endpoints/download_job/enqueue_generic_download.rs +++ b/crates/service/storyteller_web/src/http_server/endpoints/download_job/enqueue_generic_download.rs @@ -12,71 +12,63 @@ use crate::server_state::ServerState; use crate::validations::model_uploads::validate_model_title; use crate::validations::passwords::validate_passwords; use crate::validations::username::validate_username; +use crate::validations::validate_idempotency_token_format::validate_idempotency_token_format; +use database_queries::queries::generic_download::insert_generic_download_job::{Args, insert_generic_download_job}; use database_queries::tokens::Tokens; use log::{info, warn, log}; use regex::Regex; -use reusable_types::generic_upload_type::GenericUploadType; +use reusable_types::entity_visibility::EntityVisibility; +use reusable_types::generic_download_type::GenericDownloadType; use sqlx::error::DatabaseError; use sqlx::error::Error::Database; use sqlx::mysql::MySqlDatabaseError; use std::fmt; use std::sync::Arc; -#[deprecated(note = "Use `RecordVisibility` instead!")] #[derive(Deserialize)] -pub enum CreatorSetVisibility { - /// public - Public, - /// hidden - Hidden, - /// private - Private, -} - -#[derive(Deserialize)] -pub struct UploadTtsModelRequest { +pub struct EnqueueGenericDownloadRequest { idempotency_token: String, title: String, download_url: String, - generic_upload_type: GenericUploadType, + generic_download_type: GenericDownloadType, /// Not all upload types will have an associated visibility, so supplying this is optional. /// If an upload type supports it and it is not supplied, assume a reasonable default or user /// setting. - creator_set_visibility: Option, + creator_set_visibility: Option, } #[derive(Serialize)] -pub struct UploadTtsModelSuccessResponse { +pub struct EnqueueGenericDownloadSuccessResponse { pub success: bool, /// This is how frontend clients can request the job execution status. pub job_token: String, } #[derive(Debug)] -pub enum UploadTtsModelError { +pub enum EnqueueGenericDownloadError { BadInput(String), MustBeLoggedIn, ServerError, RateLimited, } -impl ResponseError for UploadTtsModelError { +impl ResponseError for EnqueueGenericDownloadError { fn status_code(&self) -> StatusCode { match *self { - UploadTtsModelError::BadInput(_) => StatusCode::BAD_REQUEST, - UploadTtsModelError::MustBeLoggedIn => StatusCode::UNAUTHORIZED, - UploadTtsModelError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, - UploadTtsModelError::RateLimited => StatusCode::TOO_MANY_REQUESTS, + EnqueueGenericDownloadError::BadInput(_) => StatusCode::BAD_REQUEST, + EnqueueGenericDownloadError::MustBeLoggedIn => StatusCode::UNAUTHORIZED, + EnqueueGenericDownloadError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + EnqueueGenericDownloadError::RateLimited => StatusCode::TOO_MANY_REQUESTS, } } fn error_response(&self) -> HttpResponse { let error_reason = match self { - UploadTtsModelError::BadInput(reason) => reason.to_string(), - UploadTtsModelError::MustBeLoggedIn => "user must be logged in".to_string(), - UploadTtsModelError::ServerError => "server error".to_string(), - UploadTtsModelError::RateLimited => "rate limited".to_string(), + EnqueueGenericDownloadError::BadInput(reason) => reason.to_string(), + EnqueueGenericDownloadError::MustBeLoggedIn => "user must be logged in".to_string(), + EnqueueGenericDownloadError::ServerError => "server error".to_string(), + EnqueueGenericDownloadError::RateLimited => "rate limited".to_string(), }; to_simple_json_error(&error_reason, self.status_code()) @@ -84,19 +76,19 @@ impl ResponseError for UploadTtsModelError { } // NB: Not using derive_more::Display since Clion doesn't understand it. -impl fmt::Display for UploadTtsModelError { +impl fmt::Display for EnqueueGenericDownloadError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "{:?}", self) } } -pub async fn upload_tts_model_handler( +pub async fn enqueue_generic_download_handler( http_request: HttpRequest, - request: web::Json, - server_state: web::Data>) -> Result + request: web::Json, + server_state: web::Data>) -> Result { if let Err(_err) = server_state.redis_rate_limiters.model_upload.rate_limit_request(&http_request) { - return Err(UploadTtsModelError::RateLimited); + return Err(EnqueueGenericDownloadError::RateLimited); } let maybe_user_session = server_state @@ -105,133 +97,87 @@ pub async fn upload_tts_model_handler( .await .map_err(|e| { warn!("Session checker error: {:?}", e); - UploadTtsModelError::ServerError + EnqueueGenericDownloadError::ServerError })?; let user_session = match maybe_user_session { Some(session) => session, None => { warn!("not logged in"); - return Err(UploadTtsModelError::MustBeLoggedIn); + return Err(EnqueueGenericDownloadError::MustBeLoggedIn); } }; - if let Err(reason) = validate_idempotency_token(&request.idempotency_token) { - return Err(UploadTtsModelError::BadInput(reason)); - } - - if let Err(reason) = validate_model_title(&request.title) { - return Err(UploadTtsModelError::BadInput(reason)); - } - let ip_address = get_request_ip(&http_request); - let uuid = request.idempotency_token.to_string(); - let title = request.title.to_string(); - let download_url = request.download_url.to_string(); + let title = request.title.trim().to_string(); + let download_url = request.download_url.trim().to_string(); + + let creator_set_visibility = request.creator_set_visibility + .unwrap_or(EntityVisibility::Public); + + if let Err(reason) = validate_idempotency_token_format(&uuid) { + return Err(EnqueueGenericDownloadError::BadInput(reason)); + } + + if let Err(reason) = validate_model_title(&title) { + return Err(EnqueueGenericDownloadError::BadInput(reason)); + } match is_bad_tts_model_download_url(&download_url) { Ok(false) => {} // Ok case Ok(true) => { - return Err(UploadTtsModelError::BadInput("Bad model download URL".to_string())); + return Err(EnqueueGenericDownloadError::BadInput("Bad model download URL".to_string())); } Err(err) => { warn!("Error parsing url: {:?}", err); - return Err(UploadTtsModelError::BadInput("Bad model download URL".to_string())); + return Err(EnqueueGenericDownloadError::BadInput("Bad model download URL".to_string())); } } - let tts_model_type = "tacotron2".to_string(); - let creator_set_visibility = "public".to_string(); - // This token is returned to the client. let job_token = Tokens::new_tts_model_upload_job() .map_err(|e| { warn!("Error creating token"); - UploadTtsModelError::ServerError + EnqueueGenericDownloadError::ServerError })?; - let query_result = sqlx::query!( - r#" -INSERT INTO tts_model_upload_jobs -SET - token = ?, - uuid_idempotency_token = ?, - creator_user_token = ?, - creator_ip_address = ?, - creator_set_visibility = ?, - title = ?, - tts_model_type = ?, - download_url = ?, - status = "pending" - "#, - job_token.to_string(), - uuid.to_string(), - user_session.user_token.to_string(), - ip_address.to_string(), - creator_set_visibility.to_string(), - title.to_string(), - tts_model_type, - download_url, - ) - .execute(&server_state.mysql_pool) - .await; - let record_id = match query_result { - Ok(res) => { - res.last_insert_id() - }, - Err(err) => { - warn!("New tts model upload creation DB error: {:?}", err); + let record_id = insert_generic_download_job(Args { + job_token: &job_token, + uuid_idempotency_token: &uuid, + download_type: request.generic_download_type, + download_url: &download_url, + title: &title, + creator_user_token: &user_session.user_token, + creator_ip_address: &ip_address, + creator_set_visibility, + mysql_pool: &server_state.mysql_pool, + }) + .await + .map_err(|err| { + warn!("New generic download creation DB error: {:?}", err); + EnqueueGenericDownloadError::ServerError + })?; - // NB: SQLSTATE[23000]: Integrity constraint violation - // NB: MySQL Error Code 1062: Duplicate key insertion (this is harder to access) - match err { - Database(err) => { - let maybe_code = err.code().map(|c| c.into_owned()); - /*match maybe_code.as_deref() { - Some("23000") => { - if err.message().contains("username") { - return Err(UsernameTaken); - } else if err.message().contains("email_address") { - return Err(EmailTaken); - } - } - _ => {}, - }*/ - }, - _ => {}, - } - return Err(UploadTtsModelError::ServerError); - } - }; + info!("new generic download job id: {}", record_id); - info!("new model upload job id: {}", record_id); - - server_state.firehose_publisher.enqueue_tts_model_upload(&user_session.user_token, &job_token) + server_state.firehose_publisher.enqueue_generic_download(&user_session.user_token, &job_token) .await .map_err(|e| { warn!("error publishing event: {:?}", e); - UploadTtsModelError::ServerError + EnqueueGenericDownloadError::ServerError })?; - let response = UploadTtsModelSuccessResponse { + let response = EnqueueGenericDownloadSuccessResponse { success: true, job_token: job_token.to_string(), }; let body = serde_json::to_string(&response) - .map_err(|e| UploadTtsModelError::ServerError)?; + .map_err(|e| EnqueueGenericDownloadError::ServerError)?; Ok(HttpResponse::Ok() .content_type("application/json") .body(body)) } - -fn validate_idempotency_token(token: &str) -> Result<(), String> { - if token.len() != 36 { - return Err("idempotency token should be 36 characters".to_string()); - } - - Ok(()) -} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/download_job/get_generic_upload_job_status.rs b/crates/service/storyteller_web/src/http_server/endpoints/download_job/get_generic_upload_job_status.rs index 17d24a204a..ec2eacc840 100644 --- a/crates/service/storyteller_web/src/http_server/endpoints/download_job/get_generic_upload_job_status.rs +++ b/crates/service/storyteller_web/src/http_server/endpoints/download_job/get_generic_upload_job_status.rs @@ -1,16 +1,17 @@ use actix_http::Error; use actix_http::http::header; -use actix_web::cookie::Cookie; use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; use actix_web::error::ResponseError; use actix_web::http::StatusCode; use actix_web::web::Path; use actix_web::{Responder, web, HttpResponse, error, HttpRequest, HttpMessage}; +use anyhow::anyhow; use chrono::{DateTime, Utc}; use crate::AnyhowResult; use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; use crate::server_state::ServerState; -use derive_more::{Display, Error}; +use database_queries::queries::generic_download::get_generic_download_job_status::{GenericDownloadJobStatus, get_generic_download_job_status}; use log::{info, warn, log}; use r2d2_redis::redis::Commands; use redis_common::redis_keys::RedisKeys; @@ -19,46 +20,46 @@ use sqlx::MySqlPool; use sqlx::error::DatabaseError; use sqlx::error::Error::Database; use sqlx::mysql::MySqlDatabaseError; +use std::fmt; use std::sync::Arc; /// For the URL PathInfo #[derive(Deserialize)] -pub struct GetTtsUploadModelStatusPathInfo { +pub struct GetGenericDownloadJobStatusPathInfo { token: String, } #[derive(Serialize)] -pub struct TtsUploadModelJobStatusForResponse { +pub struct GetGenericDownloadJobStatusForResponse { pub job_token: String, - - /// Primary status from the database (a state machine). pub status: String, - // TODO: Not yet used. /// Extra, temporary status from Redis. /// This can denote inference progress, and the Python code can write to it. pub maybe_extra_status_description: Option, pub attempt_count: u8, - pub maybe_model_token: Option, + pub maybe_downloaded_entity_type: Option, + pub maybe_downloaded_entity_token: Option, pub created_at: DateTime, pub updated_at: DateTime, } #[derive(Serialize)] -pub struct GetTtsUploadModelStatusSuccessResponse { +pub struct GetGenericDownloadJobStatusSuccessResponse { pub success: bool, - pub state: TtsUploadModelJobStatusForResponse, + pub state: GetGenericDownloadJobStatusForResponse, } -#[derive(Debug, Display)] -pub enum GetTtsUploadModelStatusError { +#[derive(Debug)] +pub enum GetGenericDownloadJobStatusError { + NotFoundError, ServerError, } #[derive(Serialize)] -pub struct TtsUploadModelJobStatusRecord { +pub struct GetGenericDownloadJobStatusRecord { pub job_token: String, pub status: String, @@ -69,63 +70,52 @@ pub struct TtsUploadModelJobStatusRecord { pub updated_at: DateTime, } -impl ResponseError for GetTtsUploadModelStatusError { +impl ResponseError for GetGenericDownloadJobStatusError { fn status_code(&self) -> StatusCode { match *self { - GetTtsUploadModelStatusError::ServerError=> StatusCode::INTERNAL_SERVER_ERROR, + GetGenericDownloadJobStatusError::NotFoundError => StatusCode::NOT_FOUND, + GetGenericDownloadJobStatusError::ServerError=> StatusCode::INTERNAL_SERVER_ERROR, } } fn error_response(&self) -> HttpResponse { let error_reason = match self { - GetTtsUploadModelStatusError::ServerError => "server error".to_string(), + GetGenericDownloadJobStatusError::NotFoundError=> "not found".to_string(), + GetGenericDownloadJobStatusError::ServerError => "server error".to_string(), }; to_simple_json_error(&error_reason, self.status_code()) } } -pub async fn get_tts_upload_model_job_status_handler( +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for GetGenericDownloadJobStatusError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + + +pub async fn get_generic_download_job_status_handler( http_request: HttpRequest, - path: Path, - server_state: web::Data>) -> Result + path: Path, + server_state: web::Data>) -> Result { - // NB: Lookup failure is Err(RowNotFound). - // NB: Since this is publicly exposed, we don't query sensitive data. - let maybe_status = sqlx::query_as!( - TtsUploadModelJobStatusRecord, - r#" -SELECT - jobs.token as job_token, + let maybe_job_status = get_generic_download_job_status( + &path.token, + &server_state.mysql_pool + ) + .await + .map_err(|err| { + warn!("error querying job record: {:?}", err); + return GetGenericDownloadJobStatusError::ServerError + })?; - jobs.status, - jobs.attempt_count, - jobs.on_success_result_token as maybe_model_token, - - jobs.created_at, - jobs.updated_at - -FROM tts_model_upload_jobs as jobs - -WHERE jobs.token = ? - "#, - &path.token - ) - .fetch_one(&server_state.mysql_pool) - .await; // TODO: This will return error if it doesn't exist - - let record : TtsUploadModelJobStatusRecord = match maybe_status { - Ok(record) => record, - Err(err) => { - match err { - sqlx::Error::RowNotFound => { - return Err(GetTtsUploadModelStatusError::ServerError); - }, - _ => { - warn!("tts template query error: {:?}", err); - return Err(GetTtsUploadModelStatusError::ServerError); - } - } + let job_status = match maybe_job_status { + Some(record) => record, + None => { + warn!("job record not found"); + return Err(GetGenericDownloadJobStatusError::NotFoundError); } }; @@ -133,10 +123,10 @@ WHERE jobs.token = ? .get() .map_err(|e| { warn!("redis error: {:?}", e); - GetTtsUploadModelStatusError::ServerError + GetGenericDownloadJobStatusError::ServerError })?; - let extra_status_key = RedisKeys::tts_download_extra_status_info(&path.token); + let extra_status_key = RedisKeys::generic_download_extra_status_info(&path.token); let maybe_extra_status_description : Option = match redis.get(&extra_status_key) { Ok(Some(status)) => { Some(status) @@ -148,23 +138,24 @@ WHERE jobs.token = ? }, }; - let model_for_response = TtsUploadModelJobStatusForResponse { - job_token: record.job_token.clone(), - status: record.status.clone(), + let model_for_response = GetGenericDownloadJobStatusForResponse { + job_token: job_status.job_token, + status: job_status.status, maybe_extra_status_description, - attempt_count: record.attempt_count as u8, - maybe_model_token: record.maybe_model_token.clone(), - created_at: record.created_at.clone(), - updated_at: record.updated_at.clone(), + attempt_count: job_status.attempt_count as u8, + maybe_downloaded_entity_token: job_status.maybe_downloaded_entity_token, + maybe_downloaded_entity_type: job_status.maybe_downloaded_entity_type, + created_at: job_status.created_at.clone(), + updated_at: job_status.updated_at.clone(), }; - let response = GetTtsUploadModelStatusSuccessResponse { + let response = GetGenericDownloadJobStatusSuccessResponse { success: true, state: model_for_response, }; let body = serde_json::to_string(&response) - .map_err(|e| GetTtsUploadModelStatusError::ServerError)?; + .map_err(|e| GetGenericDownloadJobStatusError::ServerError)?; Ok(HttpResponse::Ok() .content_type("application/json") diff --git a/crates/service/storyteller_web/src/http_server/endpoints/download_job/mod.rs b/crates/service/storyteller_web/src/http_server/endpoints/download_job/mod.rs index e273d91520..83ff514ad4 100644 --- a/crates/service/storyteller_web/src/http_server/endpoints/download_job/mod.rs +++ b/crates/service/storyteller_web/src/http_server/endpoints/download_job/mod.rs @@ -1,3 +1,2 @@ pub mod enqueue_generic_download; -pub mod generic_upload_type; pub mod get_generic_upload_job_status; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/routes.rs b/crates/service/storyteller_web/src/routes.rs index 11146bcbd0..4c5f6fe5a2 100644 --- a/crates/service/storyteller_web/src/routes.rs +++ b/crates/service/storyteller_web/src/routes.rs @@ -12,6 +12,8 @@ use crate::http_server::endpoints::categories::create_category::create_category_ use crate::http_server::endpoints::categories::get_category::get_category_handler; use crate::http_server::endpoints::categories::list_tts_categories::list_tts_categories_handler; use crate::http_server::endpoints::categories::list_tts_model_assigned_categories::list_tts_model_assigned_categories_handler; +use crate::http_server::endpoints::download_job::enqueue_generic_download::enqueue_generic_download_handler; +use crate::http_server::endpoints::download_job::get_generic_upload_job_status::get_generic_download_job_status_handler; use crate::http_server::endpoints::events::list_events::list_events_handler; use crate::http_server::endpoints::flags::design_refresh_flag::disable_design_refresh_flag_handler::disable_design_refresh_flag_handler; use crate::http_server::endpoints::flags::design_refresh_flag::enable_design_refresh_flag_handler::enable_design_refresh_flag_handler; @@ -106,6 +108,7 @@ pub fn add_routes (app: App) -> App let mut app = add_moderator_routes(app); /* /moderation */ app = add_tts_routes(app); /* /tts */ app = add_w2l_routes(app); /* /w2l */ + app = add_retrieval_routes(app); /* /retrieval (aka. "generic_download_jobs") */ app = add_category_routes(app); /* /category */ app = add_user_profile_routes(app); /* /user */ app = add_api_token_routes(app); /* /api_tokens */ @@ -469,6 +472,34 @@ fn add_w2l_routes (app: App) -> App ) } +// ==================== RETRIEVAL ROUTES ("GENERIC_DOWNLOAD_JOBS") ==================== + +fn add_retrieval_routes (app: App) -> App + where + B: MessageBody, + T: ServiceFactory< + ServiceRequest, + Config = (), + Response = ServiceResponse, + Error = Error, + InitError = (), + >, +{ + app.service( + web::scope("/retrieval") + .service( + web::resource("/enqueue") + .route(web::post().to(enqueue_generic_download_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/job_status/{token}") + .route(web::get().to(get_generic_download_job_status_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + ) +} + // ==================== CATEGORY ROUTES ==================== fn add_category_routes (app: App) -> App diff --git a/crates/service/storyteller_web/src/validations/mod.rs b/crates/service/storyteller_web/src/validations/mod.rs index e881b761a1..073fcfcec1 100644 --- a/crates/service/storyteller_web/src/validations/mod.rs +++ b/crates/service/storyteller_web/src/validations/mod.rs @@ -8,4 +8,5 @@ pub mod twitch_username; pub mod twitter_username; pub mod username; pub mod username_reservations; +pub mod validate_idempotency_token_format; pub mod website_url; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/validations/validate_idempotency_token_format.rs b/crates/service/storyteller_web/src/validations/validate_idempotency_token_format.rs new file mode 100644 index 0000000000..46e8e9adba --- /dev/null +++ b/crates/service/storyteller_web/src/validations/validate_idempotency_token_format.rs @@ -0,0 +1,43 @@ +use once_cell::sync::Lazy; +use regex::Regex; + +static EXPECTED_FORMAT : Lazy = Lazy::new(|| { + // Support UUID with or without the dashes. + Regex::new(r"^(([a-zA-Z0-9]{32})|([a-zA-Z0-9\-]{36}))$").expect("should be valid regex") +}); + +/// Validate the idempotency token +/// Idempotency tokens are supposed to be random entropy from the frontend to prevent duplicate +/// creation/insertion requests from being sent and acknowledged by the database. +/// We expect them to be in UUID format. +pub fn validate_idempotency_token_format(token: &str) -> Result<(), String> { + if !EXPECTED_FORMAT.is_match(token) { + return Err("idempotency token does not match the expected format".to_string()); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use crate::validations::validate_idempotency_token_format::validate_idempotency_token_format; + + #[test] + fn valid_idempotency_tokens() { + // UUID with the dashes + assert!(validate_idempotency_token_format("bf12126a-2de0-4527-bbbf-156c4c45ae21").is_ok()); + assert!(validate_idempotency_token_format("220ace1b-db32-42c4-bb87-49176c85a907").is_ok()); + + // UUID without the dashes + assert!(validate_idempotency_token_format("220ace1bdb3242c4bb8749176c85a907").is_ok()); + } + + #[test] + fn invalid_idempotency_tokens() { + assert!(!validate_idempotency_token_format("").is_ok()); + assert!(!validate_idempotency_token_format(" ").is_ok()); + assert!(!validate_idempotency_token_format("\n\n\n\n").is_ok()); + assert!(!validate_idempotency_token_format("bob").is_ok()); + assert!(!validate_idempotency_token_format("220ace1bdb324").is_ok()); + assert!(!validate_idempotency_token_format("220ace1bdb3242c4bb8749176c85a907ace0b0def0ff0").is_ok()); + } +}