diff --git a/crates/lib/database_queries/src/queries/mod.rs b/crates/lib/database_queries/src/queries/mod.rs index caf2868bd8..57ffb02a8e 100644 --- a/crates/lib/database_queries/src/queries/mod.rs +++ b/crates/lib/database_queries/src/queries/mod.rs @@ -4,6 +4,7 @@ pub mod generic_download; pub mod health_check; pub mod ip_bans; pub mod model_categories; +pub mod stats; pub mod tts; pub mod twitch; pub mod users; diff --git a/crates/lib/database_queries/src/queries/stats/get_on_prem_worker_stats.rs b/crates/lib/database_queries/src/queries/stats/get_on_prem_worker_stats.rs new file mode 100644 index 0000000000..49db723a37 --- /dev/null +++ b/crates/lib/database_queries/src/queries/stats/get_on_prem_worker_stats.rs @@ -0,0 +1,63 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use log::warn; +use sqlx::MySqlPool; + +#[derive(Serialize, Clone)] +pub struct OnPremWorkerStats { + pub total_records_sampled: i64, + pub on_prem_count: i64, + pub cloud_count: i64, +} + +/// Return this when there are no results. +const NO_RESULTS_SENTINEL : OnPremWorkerStats = OnPremWorkerStats { + total_records_sampled: -1, + on_prem_count: -1, + cloud_count: -1, +}; + +pub async fn get_on_prem_worker_stats( + mysql_pool: &MySqlPool, + sample_size: u32, +) -> AnyhowResult { + + let sample_size_wide = sample_size as i64; + + // NB: Lookup failure is Err(RowNotFound). + let maybe_result = sqlx::query_as!( + OnPremRaw, + r#" +select count(*) as on_prem_count +from ( + select is_generated_on_prem + from tts_results + order by id desc limit ? +) as sample +where sample.is_generated_on_prem IS TRUE; + "#, + sample_size + ) + .fetch_one(mysql_pool) + .await; + + let result : OnPremWorkerStats = match maybe_result { + Ok(result) => OnPremWorkerStats { + total_records_sampled: sample_size_wide, + on_prem_count: result.on_prem_count, + cloud_count: sample_size_wide.saturating_sub(result.on_prem_count), + }, + Err(sqlx::Error::RowNotFound) => NO_RESULTS_SENTINEL.clone(), + Err(_) => { + warn!("get on prem worker stats error: {:?}", err); + return Err(anyhow!("couldn't fetch on prem worker stats")); + } + }; + + Ok(result) +} + + +struct OnPremRaw { + on_prem_count: i64, +} \ No newline at end of file diff --git a/crates/lib/database_queries/src/queries/stats/get_voice_count_stats.rs b/crates/lib/database_queries/src/queries/stats/get_voice_count_stats.rs new file mode 100644 index 0000000000..5d75c5e515 --- /dev/null +++ b/crates/lib/database_queries/src/queries/stats/get_voice_count_stats.rs @@ -0,0 +1,49 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use log::warn; +use sqlx::MySqlPool; + +#[derive(Serialize, Clone)] +pub struct VoiceStats { + pub public_count: i64, + pub all_count: i64, +} + +/// Return this when there are no results. +const NO_RESULTS_SENTINEL : VoiceStats = VoiceStats { + public_count: -1, + all_count: -1, +}; + +pub async fn get_voice_count_stats( + mysql_pool: &MySqlPool +) -> AnyhowResult { + + // NB: Lookup failure is Err(RowNotFound). + let maybe_result = sqlx::query_as!( + VoiceStats, + r#" +SELECT public_count, all_count FROM ( + (SELECT count(*) as public_count + FROM tts_models + WHERE user_deleted_at IS NULL + AND mod_deleted_at IS NULL) as pc, + (SELECT count(*) as all_count + FROM tts_models) as ac +) + "#, + ) + .fetch_one(mysql_pool) + .await; + + let result : VoiceStats = match maybe_result { + Ok(result) => result, + Err(sqlx::Error::RowNotFound) => NO_RESULTS_SENTINEL.clone(), + Err(_) => { + warn!("get voice count stats error: {:?}", err); + return Err(anyhow!("couldn't fetch voice count stats")); + } + }; + + Ok(result) +} diff --git a/crates/lib/database_queries/src/queries/stats/mod.rs b/crates/lib/database_queries/src/queries/stats/mod.rs new file mode 100644 index 0000000000..9116b30a7b --- /dev/null +++ b/crates/lib/database_queries/src/queries/stats/mod.rs @@ -0,0 +1,2 @@ +pub mod get_on_prem_worker_stats; +pub mod get_voice_count_stats; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/http_server/endpoints/moderation/stats/get_on_prem_worker_stats.rs b/crates/service/storyteller_web/src/http_server/endpoints/moderation/stats/get_on_prem_worker_stats.rs new file mode 100644 index 0000000000..3abed00ba5 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/moderation/stats/get_on_prem_worker_stats.rs @@ -0,0 +1,117 @@ +use actix_web::web::Query; +use actix_web::{HttpRequest, HttpResponse, ResponseError, web}; +use crate::ServerState; +use crate::http_server::web_utils::serialize_as_json_error::serialize_as_json_error; +use database_queries::queries::stats::get_on_prem_worker_stats::get_on_prem_worker_stats; +use hyper::StatusCode; +use log::warn; +use std::sync::Arc; + +const MIN_SAMPLE_SIZE : u32 = 100; +const MAX_SAMPLE_SIZE : u32 = 10_000; +const DEFAULT_SAMPLE_SIZE : u32 = 1000; + +#[derive(Deserialize)] +pub struct QueryFields { + sample_size: Option, +} + +#[derive(Serialize)] +pub struct GetOnPremWorkerStatsResponse { + pub success: bool, + pub worker_mix_stats: WorkerMixStats, +} + +#[derive(Serialize)] +pub struct WorkerMixStats { + // NB: Numbers are negative when sample could not be queried. + pub total_records_sampled: i64, + pub on_prem_count: i64, + pub cloud_count: i64, +} + +#[derive(Debug, Serialize)] +pub enum GetOnPremWorkerStatsError { + ServerError, + Unauthorized, +} + +impl ResponseError for GetOnPremWorkerStatsError { + fn status_code(&self) -> StatusCode { + match *self { + GetOnPremWorkerStatsError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + GetOnPremWorkerStatsError::Unauthorized => StatusCode::UNAUTHORIZED, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using DeriveMore since Clion doesn't understand it. +impl std::fmt::Display for GetOnPremWorkerStatsError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn get_on_prem_worker_stats_handler( + http_request: HttpRequest, + query: Query, + server_state: web::Data>, +) -> Result { + + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + GetOnPremWorkerStatsError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(GetOnPremWorkerStatsError::Unauthorized); + } + }; + + // TODO: Not a good fit for this permission. + if !user_session.can_edit_other_users_tts_models { + warn!("user is not allowed to edit user tts: {}", user_session.user_token); + return Err(GetOnPremWorkerStatsError::Unauthorized); + } + + let sample_size = query.sample_size + .unwrap_or(DEFAULT_SAMPLE_SIZE) + .max(MIN_SAMPLE_SIZE) + .min(MAX_SAMPLE_SIZE); + + let result = get_on_prem_worker_stats( + &server_state.mysql_pool, sample_size + ) + .await + .map_err(|e| { + return Err(GetOnPremWorkerStatsError::ServerError) + })?; + + let response = GetOnPremWorkerStatsResponse { + success: true, + worker_mix_stats: WorkerMixStats { + total_records_sampled: result.total_records_sampled, + on_prem_count: result.on_prem_count, + cloud_count: result.cloud_count, + } + }; + + let body = serde_json::to_string(&response) + .map_err(|e| GetOnPremWorkerStatsError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} + diff --git a/crates/service/storyteller_web/src/http_server/endpoints/moderation/stats/get_voice_count_stats.rs b/crates/service/storyteller_web/src/http_server/endpoints/moderation/stats/get_voice_count_stats.rs index b1937d3fb5..a9c6a3b60d 100644 --- a/crates/service/storyteller_web/src/http_server/endpoints/moderation/stats/get_voice_count_stats.rs +++ b/crates/service/storyteller_web/src/http_server/endpoints/moderation/stats/get_voice_count_stats.rs @@ -1,7 +1,7 @@ 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; @@ -10,12 +10,12 @@ use chrono::{DateTime, Utc}; use crate::http_server::web_utils::response_success_helpers::simple_json_success; use crate::http_server::web_utils::serialize_as_json_error::serialize_as_json_error; use crate::server_state::ServerState; +use database_queries::queries::stats::get_voice_count_stats::get_voice_count_stats; use log::{info, warn, log}; use regex::Regex; use sqlx::error::DatabaseError; use sqlx::error::Error::Database; use sqlx::mysql::MySqlDatabaseError; -use std::fmt; use std::sync::Arc; #[derive(Serialize)] @@ -45,8 +45,8 @@ impl ResponseError for GetVoiceCountStatsError { } // NB: Not using DeriveMore since Clion doesn't understand it. -impl fmt::Display for GetVoiceCountStatsError { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { +impl std::fmt::Display for GetVoiceCountStatsError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "{:?}", self) } } @@ -79,41 +79,11 @@ pub async fn get_voice_count_stats_handler( return Err(GetVoiceCountStatsError::Unauthorized); } - // NB: Lookup failure is Err(RowNotFound). - let maybe_result = sqlx::query_as!( - VoiceStats, - r#" -SELECT public_count, all_count FROM ( - (SELECT count(*) as public_count - FROM tts_models - WHERE user_deleted_at IS NULL - AND mod_deleted_at IS NULL) as pc, - (SELECT count(*) as all_count - FROM tts_models) as ac -) - "#, - ) - .fetch_one(&server_state.mysql_pool) - .await; - - let result : VoiceStats = match maybe_result { - Ok(result) => result, - Err(err) => { - match err { - sqlx::Error::RowNotFound => { - // NB: Not Found for null results means nothing is pending in the queue - VoiceStats { - all_count: -1, - public_count: -1, - } - }, - _ => { - warn!("get tts pending count error: {:?}", err); - return Err(GetVoiceCountStatsError::ServerError) - } - } - }, - }; + let result = get_voice_count_stats(&server_state.mysql_pool) + .await + .map_err(|e| { + return Err(GetVoiceCountStatsError::ServerError) + })?; let response = GetVoiceCountStatsResponse { success: true, @@ -129,8 +99,3 @@ SELECT public_count, all_count FROM ( .body(body)) } -#[derive(Serialize)] -pub struct VoiceStats { - pub public_count: i64, - pub all_count: i64, -} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/moderation/stats/mod.rs b/crates/service/storyteller_web/src/http_server/endpoints/moderation/stats/mod.rs index a2441d1780..9116b30a7b 100644 --- a/crates/service/storyteller_web/src/http_server/endpoints/moderation/stats/mod.rs +++ b/crates/service/storyteller_web/src/http_server/endpoints/moderation/stats/mod.rs @@ -1 +1,2 @@ +pub mod get_on_prem_worker_stats; pub mod get_voice_count_stats; \ 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 819cd3cb4d..80698ac7fb 100644 --- a/crates/service/storyteller_web/src/routes.rs +++ b/crates/service/storyteller_web/src/routes.rs @@ -97,6 +97,7 @@ use crate::http_server::endpoints::w2l::set_w2l_template_mod_approval::set_w2l_t use users_component::default_routes::add_suggested_api_v1_account_creation_and_session_routes; use users_component::endpoints::edit_profile_handler::edit_profile_handler; use users_component::endpoints::get_profile_handler::get_profile_handler; +use crate::http_server::endpoints::moderation::stats::get_on_prem_worker_stats::get_on_prem_worker_stats_handler; pub fn add_routes (app: App) -> App where @@ -273,6 +274,11 @@ fn add_moderator_routes (app: App) -> App .route(web::get().to(get_voice_count_stats_handler)) .route(web::head().to(|| HttpResponse::Ok())) ) + .service( + web::resource("/on_prem_workers") + .route(web::get().to(get_on_prem_worker_stats_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) ) .service( web::scope("/categories")