on-prem worker stats

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;

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<OnPremWorkerStats> {
+
+  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,
+}

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

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;

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<u32>,
+}
+
+#[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<QueryFields>,
+  server_state: web::Data<Arc<ServerState>>,
+) -> Result<HttpResponse, GetOnPremWorkerStatsError> {
+
+  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))
+}
+

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,
-}

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;

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<T, B> (app: App<T, B>) -> App<T, B>
   where
@@ -273,6 +274,11 @@ fn add_moderator_routes<T, B> (app: App<T, B>) -> App<T, B>
                   .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")