From e516f85e5fd8561044c35d8e72c1a198074cb751 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 17 Mar 2023 07:54:57 -0400 Subject: [PATCH] filter out models that are private/hidden --- .../queries/tts/tts_models/list_tts_models.rs | 12 ++++- .../endpoints/tts/list_tts_models.rs | 45 ++++++++++++++++++- 2 files changed, 55 insertions(+), 2 deletions(-) diff --git a/crates/lib/database_queries/src/queries/tts/tts_models/list_tts_models.rs b/crates/lib/database_queries/src/queries/tts/tts_models/list_tts_models.rs index 42c3808fc5..4533d0e2bd 100644 --- a/crates/lib/database_queries/src/queries/tts/tts_models/list_tts_models.rs +++ b/crates/lib/database_queries/src/queries/tts/tts_models/list_tts_models.rs @@ -1,7 +1,8 @@ use anyhow::anyhow; use chrono::{DateTime, Utc}; -use errors::AnyhowResult; use crate::helpers::boolean_converters::i8_to_bool; +use enums::common::visibility::Visibility; +use errors::AnyhowResult; use log::{warn, info}; use sqlx::pool::PoolConnection; use sqlx::{MySql, MySqlPool}; @@ -35,6 +36,8 @@ pub struct TtsModelRecordForList { pub user_ratings_negative_count: u32, pub user_ratings_total_count: u32, // NB: Does not include "neutral" ratings. + pub creator_set_visibility: Visibility, + pub created_at: DateTime, pub updated_at: DateTime, } @@ -101,6 +104,7 @@ pub async fn list_tts_models_with_connection( user_ratings_positive_count: model.user_ratings_positive_count, user_ratings_negative_count: model.user_ratings_negative_count, user_ratings_total_count: model.user_ratings_total_count, + creator_set_visibility: model.creator_set_visibility, created_at: model.created_at, updated_at: model.updated_at, } @@ -136,6 +140,7 @@ SELECT tts.user_ratings_positive_count, tts.user_ratings_negative_count, tts.user_ratings_total_count, + tts.creator_set_visibility as `creator_set_visibility: enums::common::visibility::Visibility`, tts.created_at, tts.updated_at FROM tts_models as tts @@ -170,6 +175,7 @@ SELECT tts.user_ratings_positive_count, tts.user_ratings_negative_count, tts.user_ratings_total_count, + tts.creator_set_visibility as `creator_set_visibility: enums::common::visibility::Visibility`, tts.created_at, tts.updated_at FROM tts_models as tts @@ -215,6 +221,7 @@ SELECT tts.user_ratings_positive_count, tts.user_ratings_negative_count, tts.user_ratings_total_count, + tts.creator_set_visibility as `creator_set_visibility: enums::common::visibility::Visibility`, tts.created_at, tts.updated_at FROM tts_models as tts @@ -252,6 +259,7 @@ SELECT tts.user_ratings_positive_count, tts.user_ratings_negative_count, tts.user_ratings_total_count, + tts.creator_set_visibility as `creator_set_visibility: enums::common::visibility::Visibility`, tts.created_at, tts.updated_at FROM tts_models as tts @@ -296,6 +304,8 @@ struct InternalRawTtsModelRecordForList { pub user_ratings_negative_count: u32, pub user_ratings_total_count: u32, // NB: Does not include "neutral" ratings. + pub creator_set_visibility: Visibility, + pub created_at: DateTime, pub updated_at: DateTime, } diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/tts/list_tts_models.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/tts/list_tts_models.rs index 67273410cd..c34ca66a4c 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/tts/list_tts_models.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/tts/list_tts_models.rs @@ -19,6 +19,8 @@ use sqlx::pool::PoolConnection; use std::collections::HashSet; use std::fmt; use std::sync::Arc; +use database_queries::queries::users::user_sessions::get_user_session_by_token::SessionUserRecord; +use enums::common::visibility::Visibility; #[derive(Serialize, Clone)] pub struct TtsModelRecordForResponse { @@ -35,6 +37,8 @@ pub struct TtsModelRecordForResponse { pub is_twitch_featured: bool, pub maybe_suggested_unique_bot_command: Option, + pub creator_set_visibility: Visibility, + pub user_ratings: UserRatingsStats, /// Category assignments @@ -88,7 +92,7 @@ impl fmt::Display for ListTtsModelsError { } pub async fn list_tts_models_handler( - _http_request: HttpRequest, + http_request: HttpRequest, server_state: web::Data>) -> Result { if server_state.flags.disable_tts_model_list_endpoint { @@ -107,6 +111,9 @@ pub async fn list_tts_models_handler( ListTtsModelsError::ServerError })?; + // NB: We don't know if we need a MySQL connection, so don't grab one unless we do. + let mut maybe_mysql_connection = None; + let models = match maybe_models { Some(models) => { info!("Serving TTS models from cache"); @@ -122,6 +129,7 @@ pub async fn list_tts_models_handler( ListTtsModelsError::ServerError })?; + // TODO: Fail open in case the DB is down. Pull from expired cache if query fails. let models = get_all_models(&mut mysql_connection) .await .map_err(|e| { @@ -129,15 +137,49 @@ pub async fn list_tts_models_handler( ListTtsModelsError::ServerError })?; + maybe_mysql_connection = Some(mysql_connection); + server_state.caches.voice_list.store_copy(&models) .map_err(|e| { error!("Error storing cache: {:?}", e); ListTtsModelsError::ServerError })?; + models }, }; + let maybe_user_session : Option = match maybe_mysql_connection { + None => { + server_state.session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + } + Some(mut mysql_connection) => { + server_state.session_checker + .maybe_get_user_session_from_connection(&http_request, &mut mysql_connection) + .await + } + }.map_err(|e| { + warn!("Session checker error: {:?}", e); + ListTtsModelsError::ServerError + })?; + + let maybe_session_user_token = maybe_user_session + .as_ref() + .map(|s| s.user_token.as_str()); + + let models = models.into_iter() + .filter(|model| { + match model.creator_set_visibility { + Visibility::Public => true, + Visibility::Hidden | Visibility::Private => maybe_session_user_token + .map(|token| token == model.creator_user_token.as_str()) + .unwrap_or(false), + } + }) + .collect(); + render_response_ok(ListTtsModelsSuccessResponse { success: true, models, @@ -197,6 +239,7 @@ async fn get_all_models(mysql_connection: &mut PoolConnection) -> AnyhowR is_front_page_featured: model.is_front_page_featured, is_twitch_featured: model.is_twitch_featured, maybe_suggested_unique_bot_command: model.maybe_suggested_unique_bot_command, + creator_set_visibility: model.creator_set_visibility, user_ratings: UserRatingsStats { positive_count: model.user_ratings_positive_count, negative_count: model.user_ratings_negative_count,