diff --git a/Cargo.lock b/Cargo.lock index 49a020828b..e6d6b57ab8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1812,6 +1812,7 @@ name = "db-backfill" version = "0.0.1" dependencies = [ "anyhow", + "chrono", "clap", "dotenv", "easyenv", diff --git a/_sql/actions_and_analytics/tts/tts_results_by_user.sql b/_sql/actions_and_analytics/tts/tts_results_by_user.sql index 1a18383143..ce18ac6091 100644 --- a/_sql/actions_and_analytics/tts/tts_results_by_user.sql +++ b/_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 diff --git a/crates/cli/production/db-backfill/Cargo.toml b/crates/cli/production/db-backfill/Cargo.toml index 10311bc2fb..1762017036 100644 --- a/crates/cli/production/db-backfill/Cargo.toml +++ b/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" diff --git a/crates/cli/production/db-backfill/src/main.rs b/crates/cli/production/db-backfill/src/main.rs index 0cbba44702..6314b37fca 100644 --- a/crates/cli/production/db-backfill/src/main.rs +++ b/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(()) -} \ No newline at end of file +} + +async fn get_mysql(env_var_name: &str) -> AnyhowResult> { + 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) +} diff --git a/crates/cli/production/db-backfill/src/operations/calculate_model_weights_usages/run_migration.rs b/crates/cli/production/db-backfill/src/operations/calculate_model_weights_usages/run_migration.rs index 1411a027d7..cebf3b1237 100644 --- a/crates/cli/production/db-backfill/src/operations/calculate_model_weights_usages/run_migration.rs +++ b/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) -> 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(()) diff --git a/crates/lib/mysql_queries/src/queries/model_weights/count/count_model_use_using_media_files_on_date.rs b/crates/lib/mysql_queries/src/queries/model_weights/count/count_model_use_using_media_files_on_date.rs new file mode 100644 index 0000000000..3d21b5c17a --- /dev/null +++ b/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 { + 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), + }) +} diff --git a/crates/lib/mysql_queries/src/queries/model_weights/count/mod.rs b/crates/lib/mysql_queries/src/queries/model_weights/count/mod.rs new file mode 100644 index 0000000000..5ac01e5eb1 --- /dev/null +++ b/crates/lib/mysql_queries/src/queries/model_weights/count/mod.rs @@ -0,0 +1 @@ +pub mod count_model_use_using_media_files_on_date; diff --git a/crates/lib/mysql_queries/src/queries/model_weights/mod.rs b/crates/lib/mysql_queries/src/queries/model_weights/mod.rs index d72b5dfd00..eea74cf455 100644 --- a/crates/lib/mysql_queries/src/queries/model_weights/mod.rs +++ b/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;