get count of usages while supporting old token formats

Cargo.lock

@@ -1812,6 +1812,7 @@ name = "db-backfill"
 version = "0.0.1"
 dependencies = [
  "anyhow",
+ "chrono",
  "clap",
  "dotenv",
  "easyenv",

_sql/actions_and_analytics/tts/tts_results_by_user.sql

@@ -43,7 +43,7 @@ left outer join model_weights as w_migrated
   on f.maybe_origin_model_token = w_migrated.maybe_migration_old_model_token
 join users as u
   on f.maybe_creator_user_token = u.token
-where u.username = lower('powerchat')
+where u.username = lower('peepostream')
 order by f.id desc
 limit 100
 

crates/cli/production/db-backfill/Cargo.toml

@@ -28,6 +28,7 @@ tokens = { path = "../../../_types/tokens" }
 
 # External
 anyhow = "1.0.86"
+chrono = { version = "0.4.22", features = ["serde"] }
 clap = "4.4.8"
 dotenv = "0.15.0"
 log = "0.4.14"

crates/cli/production/db-backfill/src/main.rs

@@ -1,7 +1,10 @@
 use log::info;
+use sqlx::{MySql, Pool};
+use sqlx::mysql::MySqlPoolOptions;
 
 use easyenv::init_all_with_default_logging;
 use errors::AnyhowResult;
+
 use crate::args::{Command, parse_cli_args};
 use crate::operations::calculate_model_weights_usages::run_migration::run_migration;
 
@@ -16,8 +19,14 @@ async fn main() -> AnyhowResult<()> {
 
     let command = parse_cli_args()?;
 
+    // NB: This secrets file differs from the rest because we might want to backfill production from local dev.
+    // (Hopefully this isn't getting out of hand at this point.)
+    easyenv::from_filename(".env-db-backfill-secrets")?;
+
+    let mysql = get_mysql("MYSQL_PRODUCTION_URL").await?;
+
     match command.sub_command {
-        Command::CalculateModelWeightsUsages => run_migration().await?,
+        Command::CalculateModelWeightsUsages => run_migration(mysql).await?,
         Command::CalculateTtsResultsUsages => {}
     }
     info!("command: {:?}", command);
@@ -25,4 +34,15 @@ async fn main() -> AnyhowResult<()> {
     info!("TODO...");
 
     Ok(())
-}
+}
+
+async fn get_mysql(env_var_name: &str) -> AnyhowResult<Pool<MySql>> {
+    info!("Connecting to MySQL {env_var_name}...");
+
+    let pool = MySqlPoolOptions::new()
+        .max_connections(easyenv::get_env_num("MYSQL_MAX_CONNECTIONS", 3)?)
+        .connect(&easyenv::get_env_string_required(env_var_name)?)
+        .await?;
+
+    Ok(pool)
+}

crates/cli/production/db-backfill/src/operations/calculate_model_weights_usages/run_migration.rs

@@ -1,13 +1,28 @@
+use anyhow::anyhow;
+use chrono::NaiveDate;
 use log::info;
+use sqlx::{MySql, Pool};
 use errors::AnyhowResult;
-
+use mysql_queries::queries::model_weights::count::count_model_use_using_media_files_on_date::count_model_use_using_media_files_on_date;
+use tokens::tokens::media_files::MediaFileToken;
 use crate::operations::calculate_model_weights_usages::sub_args::parse_cli_sub_args;
 
-pub async fn run_migration() -> AnyhowResult<()> {
+pub async fn run_migration(mysql: Pool<MySql>) -> AnyhowResult<()> {
   let args = parse_cli_sub_args()?;
 
   info!("args: {:?}", args);
 
+  let model_token = args.model_token
+      .map(|token| MediaFileToken::new_from_str(&token))
+      .ok_or_else(|| anyhow!("no token provided"))?;
+
+  let date = NaiveDate::from_ymd_opt(2024, 9, 1)
+      .ok_or_else(|| anyhow!("invalid start date"))?;
+
+  let count = count_model_use_using_media_files_on_date(&mysql, &model_token, date).await?;
+
+  info!("count: {:?}", count.record_count);
+
   println!("calculate_model_weights_usages: run migration");
 
   Ok(())

crates/lib/mysql_queries/src/queries/model_weights/count/count_model_use_using_media_files_on_date.rs

@@ -0,0 +1,56 @@
+use anyhow::anyhow;
+use chrono::{Days, NaiveDate};
+use sqlx::MySqlPool;
+
+use errors::AnyhowResult;
+use tokens::tokens::media_files::MediaFileToken;
+
+use crate::helpers::numeric_converters::try_i64_to_u64_or_min;
+
+#[derive(Clone)]
+pub struct ModelUseCount {
+  pub record_count: u64,
+}
+
+struct RawResult {
+  record_count: i64,
+}
+
+pub async fn count_model_use_using_media_files_on_date(
+  pool: &MySqlPool,
+  token: &MediaFileToken,
+  date: NaiveDate,
+) -> AnyhowResult<ModelUseCount> {
+  let start_date = date;
+  //let end_date = date.checked_add_days(Days::new(1))
+  //    .ok_or_else(|| anyhow!("Failed to add days to date"))?;
+  // Avoid full table scan on date-to-timestamp: https://stackoverflow.com/a/14769096
+  //  AND f.created_at >= CURDATE()
+  //  AND f.created_at < CURDATE() + INTERVAL 1 DAY
+  // WHERE (w.token = ? OR w.maybe_migration_old_model_token = ?)
+  let result = sqlx::query_as!(
+      RawResult ,
+        r#"
+SELECT
+  COUNT(*) as record_count
+  FROM media_files as f
+  left outer join model_weights as w
+     on f.maybe_origin_model_token = w.token
+  left outer join model_weights as w_old
+     on f.maybe_origin_model_token = w_old.maybe_migration_old_model_token
+  WHERE (w.token = ? OR w_old.token = ?)
+  AND f.created_at >= ?
+  AND f.created_at < ? + INTERVAL 1 DAY
+        "#,
+      token.as_str(),
+      token.as_str(),
+      start_date,
+      start_date,
+    )
+      .fetch_one(pool)
+      .await?;
+
+  Ok(ModelUseCount {
+    record_count: try_i64_to_u64_or_min(result.record_count),
+  })
+}

crates/lib/mysql_queries/src/queries/model_weights/count/mod.rs

@@ -0,0 +1 @@
+pub mod count_model_use_using_media_files_on_date;

crates/lib/mysql_queries/src/queries/model_weights/mod.rs

@@ -1,4 +1,5 @@
 pub mod batch_get;
+pub mod count;
 pub mod create;
 pub mod delete_weights;
 pub mod edit;