diff --git a/crates/lib/mysql_queries/src/queries/model_weights/batch_get/batch_get_model_weights_for_elastic_search_backfill.rs b/crates/lib/mysql_queries/src/queries/model_weights/batch_get/batch_get_model_weights_for_elastic_search_backfill.rs new file mode 100644 index 0000000000..d5aa4f8c86 --- /dev/null +++ b/crates/lib/mysql_queries/src/queries/model_weights/batch_get/batch_get_model_weights_for_elastic_search_backfill.rs @@ -0,0 +1,296 @@ +use anyhow::anyhow; +use chrono::{DateTime, Utc}; +use log::warn; +use sqlx::{Executor, FromRow, MySql, QueryBuilder, Row}; +use sqlx::mysql::MySqlRow; + +use enums::by_table::model_weights::weights_category::WeightsCategory; +use enums::by_table::model_weights::weights_types::WeightsType; +use enums::common::visibility::Visibility; +use enums::traits::mysql_from_row::MySqlFromRow; +use errors::AnyhowResult; +use tokens::tokens::media_files::MediaFileToken; +use tokens::tokens::model_weights::ModelWeightToken; +use tokens::tokens::users::UserToken; + +/// This is meant to be the entire table +#[derive(Debug)] +pub struct ModelWeightForElasticsearchRecord { + pub id: i64, + pub token: ModelWeightToken, + + pub creator_set_visibility: Visibility, + + pub weights_type: WeightsType, + pub weights_category: WeightsCategory, + + pub title: String, + + // Cover images + pub maybe_cover_image_media_file_token: Option, + pub maybe_cover_image_public_bucket_hash: Option, + pub maybe_cover_image_public_bucket_prefix: Option, + pub maybe_cover_image_public_bucket_extension: Option, + + // Creator + pub creator_user_token: UserToken, + pub creator_username: String, + pub creator_display_name: String, + pub creator_gravatar_hash: String, + + // Statistics + pub maybe_ratings_positive_count: Option, + pub maybe_ratings_negative_count: Option, + pub maybe_bookmark_count: Option, + + // TTS extensions + pub maybe_tts_ietf_language_tag: Option, + pub maybe_tts_ietf_primary_language_subtag: Option, + + // VC extensions + pub maybe_voice_conversion_ietf_language_tag: Option, + pub maybe_voice_conversion_ietf_primary_language_subtag: Option, + + pub created_at: DateTime, + pub updated_at: DateTime, + + pub user_deleted_at: Option>, + pub mod_deleted_at: Option>, +} + +pub async fn batch_get_model_weights_for_elastic_search_backfill<'e, 'c, E>( + mysql_executor: E, + tokens: &Vec, +) -> AnyhowResult> + where E: 'e + Executor<'c, Database=MySql> +{ + + let maybe_models + = list_model_weights(mysql_executor, tokens) + .await; + + let models : Vec = match maybe_models { + Ok(models) => models, + Err(sqlx::error::Error::RowNotFound) => return Ok(Vec::new()), + Err(err) => { + warn!("model weight list query error: {:?}", err); + return Err(anyhow!("model weight list query error")); + } + }; + + Ok(models.into_iter() + .map(|model| { + ModelWeightForElasticsearchRecord { + id: model.id, + token: model.token, + title: model.title, + weights_type: model.weights_type, + weights_category: model.weights_category, + maybe_cover_image_media_file_token: model.maybe_cover_image_media_file_token, + maybe_cover_image_public_bucket_hash: model.maybe_cover_image_public_bucket_hash, + maybe_cover_image_public_bucket_prefix: model.maybe_cover_image_public_bucket_prefix, + maybe_cover_image_public_bucket_extension: model.maybe_cover_image_public_bucket_extension, + creator_user_token: model.creator_user_token, + creator_username: model.creator_username, + creator_display_name: model.creator_display_name, + creator_gravatar_hash: model.creator_gravatar_hash, + maybe_tts_ietf_language_tag: model.maybe_tts_ietf_language_tag, + maybe_tts_ietf_primary_language_subtag: model.maybe_tts_ietf_primary_language_subtag, + maybe_voice_conversion_ietf_language_tag: model.maybe_vc_ietf_language_tag, + maybe_voice_conversion_ietf_primary_language_subtag: model.maybe_vc_ietf_primary_language_subtag, + maybe_ratings_positive_count: model.maybe_ratings_positive_count, + maybe_ratings_negative_count: model.maybe_ratings_negative_count, + maybe_bookmark_count: model.maybe_bookmark_count, + creator_set_visibility: model.creator_set_visibility, + created_at: model.created_at, + updated_at: model.updated_at, + user_deleted_at: model.user_deleted_at, + mod_deleted_at: model.mod_deleted_at, + } + }) + .collect::>()) +} + +async fn list_model_weights<'e, 'c, E>( + mysql_executor: E, + tokens: &Vec, +) -> Result, sqlx::Error> + where E: 'e + Executor<'c, Database=MySql> +{ + if tokens.len() == 0 { + return Ok(Vec::new()); + } + + let mut query_builder: QueryBuilder = QueryBuilder::new( + r#" +SELECT + w.id, + w.token + + w.creator_set_visibility + + w.weights_type, + w.weights_category, + + w.title, + + cover_image.token as maybe_cover_image_media_file_token, + cover_image.public_bucket_directory_hash as maybe_cover_image_public_bucket_hash, + cover_image.maybe_public_bucket_prefix as maybe_cover_image_public_bucket_prefix, + cover_image.maybe_public_bucket_extension as maybe_cover_image_public_bucket_extension, + + w.creator_user_token, + users.username as creator_username, + users.display_name as creator_display_name, + users.email_gravatar_hash as creator_gravatar_hash, + + entity_stats.ratings_positive_count as maybe_ratings_positive_count, + entity_stats.ratings_negative_count as maybe_ratings_negative_count, + entity_stats.bookmark_count as maybe_bookmark_count, + + extension_tts.ietf_language_tag as maybe_tts_ietf_language_tag, + extension_tts.ietf_primary_language_subtag as maybe_tts_ietf_primary_language_subtag, + + extension_vc.ietf_language_tag as maybe_vc_ietf_language_tag, + extension_vc.ietf_primary_language_subtag as maybe_vc_ietf_primary_language_subtag, + + w.created_at, + w.updated_at, + w.user_deleted_at, + w.mod_deleted_at + +FROM model_weights as w + +JOIN users + ON users.token = w.creator_user_token + +LEFT OUTER JOIN model_weights_extension_tts_details as extension_tts + ON extension_tts.model_weights_token = w.token +LEFT OUTER JOIN model_weights_extension_voice_conversion_details as extension_vc + ON extension_vc.model_weights_token = w.token +LEFT OUTER JOIN entity_stats + ON entity_stats.entity_type = "model_weights" + AND entity_stats.entity_token = w.token +LEFT OUTER JOIN media_files as cover_image + ON cover_image.token = w.maybe_cover_image_media_file_token + +WHERE token IN ( + "#); + + // NB: Syntax will be wrong if list has zero length + let mut separated = query_builder.separated(", "); + + for token in tokens.iter() { + separated.push_bind(token.as_str()); + } + + separated.push_unseparated(") "); + + // NB: We should return the most recent update timestamp last (for cursoring) + query_builder.push(" ORDER BY w.updated_at ASC, w.id ASC "); + + let query = query_builder.build_query_as::(); + + let maybe_results = query.fetch_all(mysql_executor).await?; + + Ok(maybe_results) +} + +struct RawRecord { + pub id: i64, + pub token: ModelWeightToken, + + pub creator_set_visibility: Visibility, + + pub weights_type: WeightsType, + pub weights_category: WeightsCategory, + + pub title: String, + + // Cover images + pub maybe_cover_image_media_file_token: Option, + pub maybe_cover_image_public_bucket_hash: Option, + pub maybe_cover_image_public_bucket_prefix: Option, + pub maybe_cover_image_public_bucket_extension: Option, + + // Creator + pub creator_user_token: UserToken, + pub creator_username: String, + pub creator_display_name: String, + pub creator_gravatar_hash: String, + + // Statistics + pub maybe_ratings_positive_count: Option, + pub maybe_ratings_negative_count: Option, + pub maybe_bookmark_count: Option, + + // TTS extensions + pub maybe_tts_ietf_language_tag: Option, + pub maybe_tts_ietf_primary_language_subtag: Option, + + // Voice conversion extensions + pub maybe_vc_ietf_language_tag: Option, + pub maybe_vc_ietf_primary_language_subtag: Option, + + pub created_at: DateTime, + pub updated_at: DateTime, + + pub user_deleted_at: Option>, + pub mod_deleted_at: Option>, +} + +// NB(bt,2023-12-05): There's an issue with type hinting in the `as` clauses with QueryBuilder (or +// raw query strings) and sqlx::FromRow, regardless of whether it is derived of manually +// implemented. Perhaps this will improve in the future, but for now manually constructed queries +// cannot have type hints, eg. the following: +// +// m.token as `token: tokens::tokens::media_files::MediaFileToken`, +// m.origin_category as `origin_category: enums::by_table::media_files::media_file_origin_category::MediaFileOriginCategory`, +// m.creator_set_visibility as `creator_set_visibility: enums::common::visibility::Visibility`, +// +// This results in the automatic mapping not being able to be found by name (for macro derive), and +// in the manual case `row.try_get()` etc. won't have the correct column name (since the name is the +// full "as" clause). +impl FromRow<'_, MySqlRow> for RawRecord { + fn from_row(row: &MySqlRow) -> Result { + // NB(bt): Not sure why IDE is complaining about inlining this with the call to map() below. + let maybe_cover_image_media_file_token: Option = row.try_get("maybe_cover_image_media_file_token")?; + + Ok(Self { + id: row.try_get("id")?, + token: ModelWeightToken::new(row.try_get("token")?), + + creator_set_visibility: Visibility::try_from_mysql_row(row, "creator_set_visibility")?, + + weights_type: WeightsType::try_from_mysql_row(row, "weights_type")?, + weights_category: WeightsCategory::try_from_mysql_row(row, "weights_category")?, + + title: row.try_get("title")?, + + maybe_cover_image_media_file_token: maybe_cover_image_media_file_token.map(|token| MediaFileToken::new(token)), + maybe_cover_image_public_bucket_hash: row.try_get("maybe_cover_image_public_bucket_hash")?, + maybe_cover_image_public_bucket_prefix: row.try_get("maybe_cover_image_public_bucket_prefix")?, + maybe_cover_image_public_bucket_extension: row.try_get("maybe_cover_image_public_bucket_extension")?, + + creator_user_token: UserToken::new(row.try_get("creator_user_token")?), + creator_username: row.try_get("creator_username")?, + creator_display_name: row.try_get("creator_display_name")?, + creator_gravatar_hash: row.try_get("creator_gravatar_hash")?, + + maybe_ratings_positive_count: row.try_get("maybe_ratings_positive_count")?, + maybe_ratings_negative_count: row.try_get("maybe_ratings_negative_count")?, + maybe_bookmark_count: row.try_get("maybe_bookmark_count")?, + + maybe_tts_ietf_language_tag: row.try_get("maybe_tts_ietf_language_tag")?, + maybe_tts_ietf_primary_language_subtag: row.try_get("maybe_tts_ietf_primary_language_subtag")?, + + maybe_vc_ietf_language_tag: row.try_get("maybe_vc_ietf_language_tag")?, + maybe_vc_ietf_primary_language_subtag: row.try_get("maybe_vc_ietf_primary_language_subtag")?, + + created_at: row.try_get("created_at")?, + updated_at: row.try_get("updated_at")?, + user_deleted_at: row.try_get("user_deleted_at")?, + mod_deleted_at: row.try_get("mod_deleted_at")?, + }) + } +} diff --git a/crates/lib/mysql_queries/src/queries/model_weights/batch_get/mod.rs b/crates/lib/mysql_queries/src/queries/model_weights/batch_get/mod.rs new file mode 100644 index 0000000000..4e5e9dc900 --- /dev/null +++ b/crates/lib/mysql_queries/src/queries/model_weights/batch_get/mod.rs @@ -0,0 +1 @@ +pub mod batch_get_model_weights_for_elastic_search_backfill; \ No newline at end of file diff --git a/crates/lib/mysql_queries/src/queries/model_weights/list/list_model_weight_tokens_updated_since.rs b/crates/lib/mysql_queries/src/queries/model_weights/list/list_model_weight_tokens_updated_since.rs new file mode 100644 index 0000000000..87e92399f3 --- /dev/null +++ b/crates/lib/mysql_queries/src/queries/model_weights/list/list_model_weight_tokens_updated_since.rs @@ -0,0 +1,55 @@ +use chrono::{DateTime, Utc}; +use sqlx::{Executor, MySql, Row}; + +use enums::by_table::model_weights::weights_types::WeightsType; +use errors::AnyhowResult; +use tokens::tokens::model_weights::ModelWeightToken; + +pub struct UpdatedModelWeight { + pub token: ModelWeightToken, + pub weights_type: WeightsType, + pub updated_at: DateTime, +} + +/// Get records updated since a particular date. +/// NB: This doesn't batch yet at our scale, but in the future we'll likely need/want to batch, +/// and then we'll need another way to cursor in the event that batch_size < # records updated +/// within a given time quantum. +pub async fn list_model_weight_tokens_updated_since<'e, 'c, E>( + mysql_executor: E, + since_date: &DateTime, +) -> AnyhowResult> + where E: 'e + Executor<'c, Database=MySql> +{ + let query = sqlx::query_as!( + UpdatedModelWeightRaw, + r#" + SELECT + token as `token: tokens::tokens::model_weights::ModelWeightToken`, + weights_type as `weights_type: enums::by_table::model_weights::weights_types::WeightsType`, + updated_at + FROM model_weights + WHERE + updated_at > ? + "#, + since_date + ); + + let results = query.fetch_all(mysql_executor).await?; + + let results = results.into_iter() + .map(|record| UpdatedModelWeight { + token: record.token, + weights_type: record.weights_type, + updated_at: record.updated_at, + }) + .collect(); + + Ok(results) +} + +pub struct UpdatedModelWeightRaw { + pub token: ModelWeightToken, + pub weights_type: WeightsType, + pub updated_at: DateTime, +} diff --git a/crates/lib/mysql_queries/src/queries/model_weights/list/mod.rs b/crates/lib/mysql_queries/src/queries/model_weights/list/mod.rs index f56b2896be..c13e6a3e34 100644 --- a/crates/lib/mysql_queries/src/queries/model_weights/list/mod.rs +++ b/crates/lib/mysql_queries/src/queries/model_weights/list/mod.rs @@ -1,6 +1,7 @@ +pub mod list_model_weight_tokens_updated_since; pub mod list_model_weights_for_elastic_search_backfill_using_cursor; pub mod list_model_weights_for_text_to_speech; pub mod list_model_weights_for_voice_conversion; pub mod list_weights_by_tokens; pub mod list_weights_by_user; -pub mod list_weights_query_builder; \ No newline at end of file +pub mod list_weights_query_builder; 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 8c32897033..d72b5dfd00 100644 --- a/crates/lib/mysql_queries/src/queries/model_weights/mod.rs +++ b/crates/lib/mysql_queries/src/queries/model_weights/mod.rs @@ -1,3 +1,4 @@ +pub mod batch_get; pub mod create; pub mod delete_weights; pub mod edit; diff --git a/crates/service/job/es-update-job/src/job_state.rs b/crates/service/job/es-update-job/src/job_state.rs index 23d9ab4d5f..ce80365dce 100644 --- a/crates/service/job/es-update-job/src/job_state.rs +++ b/crates/service/job/es-update-job/src/job_state.rs @@ -10,8 +10,8 @@ pub struct JobState { /// Use sleep to not overload the database. pub struct SleepConfigs { - // How long to wait between individual "jobs". - pub between_job_wait_millis: u64, + // How long to wait between individual ES writes. + pub between_es_writes_wait_millis: u64, // How long to wait between individual job batches. pub between_job_batch_wait_millis: u64, diff --git a/crates/service/job/es-update-job/src/jobs/calculate_model_analytics.rs b/crates/service/job/es-update-job/src/jobs/calculate_model_analytics.rs deleted file mode 100644 index d0b6be90e1..0000000000 --- a/crates/service/job/es-update-job/src/jobs/calculate_model_analytics.rs +++ /dev/null @@ -1,121 +0,0 @@ -use std::time::Duration; - -use log::info; -use rand::seq::SliceRandom; -use rand::thread_rng; -use sqlx::MySql; -use sqlx::pool::PoolConnection; - -use enums::by_table::trending_model_analytics::window_name::WindowName; -use errors::AnyhowResult; -use mysql_queries::queries::trending_model_analytics::upsert_trending_model_analytics::{Args, ModelToken, upsert_trending_model_analytics}; -use mysql_queries::queries::tts::tts_models::count_tts_model_uses::count_tts_model_uses; -use mysql_queries::queries::tts::tts_models::count_tts_model_uses_total::count_tts_model_uses_total; -use mysql_queries::queries::tts::tts_models::list_all_tts_model_tokens::{list_all_tts_model_tokens, TtsModelTokens}; -use tokens::tokens::tts_models::TtsModelToken; - -use crate::job_state::JobState; - -// TODO(bt, 2023-01-16): Make a job trait and make this an instance. -// Jobs can store state, batch progress, and not starve one another. -// It'll be easier to interleave jobs. - -pub async fn calculate_model_analytics(job_state: &JobState) -> AnyhowResult<()> { - // TODO(bt, 2023-01-16): It's conceivable the pool connection dies mid-workload and we starve - // our tokens. An ideal fix would be a self-healing pool connection, but for now we'll use - // ugly hacks. - let mut mysql_connection = job_state.mysql_pool.acquire().await?; - - let tokens = query_all_model_tokens(&mut mysql_connection).await?; - - std::thread::sleep(Duration::from_millis(100)); - - for token in tokens.tokens { - info!("Running analytics for TTS model {}", token.token); - - let _r = query_single_model_statistics( - job_state, - &mut mysql_connection, - &token.token).await; - - std::thread::sleep(Duration::from_millis(job_state.sleep_config.between_job_wait_millis)); - } - - Ok(()) -} - -async fn query_all_model_tokens( - mysql_connection: &mut PoolConnection, -) -> AnyhowResult { - let mut tokens = list_all_tts_model_tokens(mysql_connection).await?; - - // NB: This is just a hack to not starve the work queue. We should replace this with a - // self-healing connection and batch cursors. - let mut rng = thread_rng(); - tokens.tokens.shuffle(&mut rng); - - Ok(tokens) -} - -async fn query_single_model_statistics( - job_state: &JobState, - mysql_connection: &mut PoolConnection, - model_token: &TtsModelToken, -) -> AnyhowResult<()> { - { - let three_hours_in_minutes = 60 * 3; - - let result = count_tts_model_uses( - model_token, - three_hours_in_minutes, - mysql_connection).await?; - - info!("TTS model {} uses {} times (three hours)", model_token, result.use_count); - - let _r = upsert_trending_model_analytics(Args { - model_token: ModelToken::Tts(model_token), - window_name: WindowName::Last3Hours, - numeric_value: result.use_count, - mysql_connection, - }).await?; - } - - std::thread::sleep(Duration::from_millis(job_state.sleep_config.between_query_wait_millis)); - - { - let three_days_in_minutes = 60 * 24 * 3; - - let result = count_tts_model_uses( - model_token, - three_days_in_minutes, - mysql_connection).await?; - - info!("TTS model {} uses {} times (three days)", model_token, result.use_count); - - let _r = upsert_trending_model_analytics(Args { - model_token: ModelToken::Tts(model_token), - window_name: WindowName::Last3Days, - numeric_value: result.use_count, - mysql_connection, - }).await?; - } - - std::thread::sleep(Duration::from_millis(job_state.sleep_config.between_query_wait_millis)); - - { - let result = count_tts_model_uses_total( - model_token, - mysql_connection).await?; - - info!("TTS model {} uses {} times (total)", model_token, result.total_use_count); - - let _r = upsert_trending_model_analytics(Args { - model_token: ModelToken::Tts(model_token), - window_name: WindowName::AllTime, - numeric_value: result.total_use_count, - mysql_connection, - }).await?; - } - - Ok(()) -} diff --git a/crates/service/job/es-update-job/src/jobs/mod.rs b/crates/service/job/es-update-job/src/jobs/mod.rs deleted file mode 100644 index 3b2fdb8c68..0000000000 --- a/crates/service/job/es-update-job/src/jobs/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod calculate_model_analytics; \ No newline at end of file diff --git a/crates/service/job/es-update-job/src/main.rs b/crates/service/job/es-update-job/src/main.rs index 9e135f4620..75afcbb0bf 100644 --- a/crates/service/job/es-update-job/src/main.rs +++ b/crates/service/job/es-update-job/src/main.rs @@ -31,7 +31,6 @@ use crate::job_state::{JobState, SleepConfigs}; use crate::main_loop::main_loop; pub mod job_state; -pub mod jobs; pub mod main_loop; #[tokio::main] @@ -53,7 +52,7 @@ async fn main() -> AnyhowResult<()> { mysql_pool, elasticsearch, sleep_config: SleepConfigs { - between_job_wait_millis: easyenv::get_env_num("BETWEEN_JOB_WAIT_MILLIS", 100)?, + between_es_writes_wait_millis: easyenv::get_env_num("BETWEEN_JOB_WAIT_MILLIS", 100)?, between_job_batch_wait_millis: easyenv::get_env_num("BETWEEN_JOB_BATCH_WAIT_MILLIS", 5000)?, between_query_wait_millis: easyenv::get_env_num("BETWEEN_QUERY_WAIT_MILLIS", 100)?, between_error_wait_millis: easyenv::get_env_num("BETWEEN_ERROR_WAIT_MILLIS", 10_000)?, diff --git a/crates/service/job/es-update-job/src/main_loop.rs b/crates/service/job/es-update-job/src/main_loop.rs index 17fc8c5b14..801c808723 100644 --- a/crates/service/job/es-update-job/src/main_loop.rs +++ b/crates/service/job/es-update-job/src/main_loop.rs @@ -1,24 +1,156 @@ use std::time::Duration; +use anyhow::anyhow; +use chrono::{DateTime, Utc}; +use elasticsearch::{BulkOperation, BulkParts, Elasticsearch}; use log::{error, info}; +use serde_json::Value; +use elasticsearch_schema::documents::model_weight_document::{MODEL_WEIGHT_INDEX, ModelWeightDocument}; +use elasticsearch_schema::traits::document::Document; +use enums::by_table::model_weights::weights_category::WeightsCategory; use errors::AnyhowResult; +use mysql_queries::queries::model_weights::batch_get::batch_get_model_weights_for_elastic_search_backfill::{batch_get_model_weights_for_elastic_search_backfill, ModelWeightForElasticsearchRecord}; +use mysql_queries::queries::model_weights::list::list_model_weight_tokens_updated_since::list_model_weight_tokens_updated_since; use crate::job_state::JobState; -use crate::jobs::calculate_model_analytics::calculate_model_analytics; -pub async fn main_loop(job_state: JobState) -> AnyhowResult<()> { +pub async fn main_loop(job_state: JobState) { + let cursor = DateTime::UNIX_EPOCH; + loop { - info!("Begin elasticsearch job batch."); + info!("Main loop; cursor @ {:?}", &cursor); - match calculate_model_analytics(&job_state).await { - Ok(_) => { - std::thread::sleep(Duration::from_millis(job_state.sleep_config.between_job_batch_wait_millis)); - } - Err(e) => { - error!("Error: {:?}", e); - std::thread::sleep(Duration::from_millis(job_state.sleep_config.between_error_wait_millis)); - } + let result = with_database_main_loop(cursor, &job_state).await; + + if let Err(err) = result { + error!("Error in main loop: {:?}", err); + std::thread::sleep(Duration::from_millis(job_state.sleep_config.between_error_wait_millis)); } + + // NB: Function should never return in success case, but we'll sleep just in case. } } + +pub async fn with_database_main_loop(mut updated_at_cursor: DateTime, job_state: &JobState) -> AnyhowResult<()> { + let mut mysql_connection = job_state.mysql_pool.acquire().await?; + + let mut last_successful_update_at = updated_at_cursor.clone(); + + loop { + info!("Querying tokens updated since: {:?}", &updated_at_cursor); + + let mut maybe_tokens = + list_model_weight_tokens_updated_since(&mut *mysql_connection, &updated_at_cursor).await?; + + while !maybe_tokens.is_empty() { + // NB: This list might be very large if we query from (1) the epoch, or (2) there was a large series of updates + let drained_tokens = maybe_tokens.drain(0..50) + .into_iter() + .map(|record| record.token) + .collect::>(); + + let records + = batch_get_model_weights_for_elastic_search_backfill(&mut *mysql_connection, &drained_tokens).await?; + + for record in records { + let updated_at = record.updated_at.clone(); + + create_document_from_record(&job_state.elasticsearch, record).await?; + + // NB: We don't want to advance the cursor to the current second just yet, because we might be processing + // several records with the exact same "updated_at" timestamp. We do know that due to ordering, we can update + // the timestamp cursor to the second before this timestamp, however. + // + // If we change the cursoring or batch sizes, we'll need to reconsider all of this logic. + if let Some(cursor_at_least) = updated_at.checked_sub_signed(chrono::Duration::seconds(1)) { + // NB: We don't want the cursor to slide backwards (it shouldn't, but the max will save us). + last_successful_update_at = cursor_at_least.max(last_successful_update_at); + } + + updated_at_cursor = last_successful_update_at.max(updated_at_cursor); + + std::thread::sleep(Duration::from_millis(job_state.sleep_config.between_es_writes_wait_millis)); + } + } + + info!("Up to date at cursor = {:?}", &updated_at_cursor); + + std::thread::sleep(Duration::from_millis(job_state.sleep_config.between_job_batch_wait_millis)); + } +} + +async fn create_document_from_record(elasticsearch: &Elasticsearch, record: ModelWeightForElasticsearchRecord) -> AnyhowResult<()> { + info!("Create record for {:?} - {:?}", record.token, record.title); + + let document = ModelWeightDocument { + token: record.token, + + creator_set_visibility: record.creator_set_visibility, + + weights_type: record.weights_type, + weights_category: record.weights_category, + + title: record.title, + + maybe_cover_image_media_file_token: record.maybe_cover_image_media_file_token, + maybe_cover_image_public_bucket_hash: record.maybe_cover_image_public_bucket_hash, + maybe_cover_image_public_bucket_prefix: record.maybe_cover_image_public_bucket_prefix, + maybe_cover_image_public_bucket_extension: record.maybe_cover_image_public_bucket_extension, + + creator_user_token: record.creator_user_token, + creator_username: record.creator_username, + creator_display_name: record.creator_display_name, + creator_gravatar_hash: record.creator_gravatar_hash, + + ratings_positive_count: record.maybe_ratings_positive_count.unwrap_or(0), + ratings_negative_count: record.maybe_ratings_negative_count.unwrap_or(0), + bookmark_count: record.maybe_bookmark_count.unwrap_or(0), + + maybe_ietf_language_tag: match record.weights_category { + WeightsCategory::TextToSpeech => record.maybe_tts_ietf_language_tag, + WeightsCategory::VoiceConversion => record.maybe_voice_conversion_ietf_language_tag, + _ => None, + }, + + maybe_ietf_primary_language_subtag: match record.weights_category { + WeightsCategory::TextToSpeech => record.maybe_tts_ietf_primary_language_subtag, + WeightsCategory::VoiceConversion => record.maybe_voice_conversion_ietf_primary_language_subtag, + _ => None, + }, + + created_at: record.created_at, + updated_at: record.updated_at, + user_deleted_at: record.user_deleted_at, + mod_deleted_at: record.mod_deleted_at, + }; + + let op : BulkOperation<_> = BulkOperation::index(&document) + .id(document.get_document_id()) + .into(); + + let response = elasticsearch + .bulk(BulkParts::Index(MODEL_WEIGHT_INDEX)) + .body(vec![op]) + .send() + .await?; + + let json: Value = response.json().await?; + + let had_errors = json["errors"].as_bool().unwrap_or(false); + + if had_errors { + let failed: Vec<&Value> = json["items"] + .as_array() + .unwrap() + .iter() + .filter(|v| !v["error"].is_null()) + .collect(); + + error!("Errors during indexing. Failures: {}", failed.len()); + + return Err(anyhow!("Errors during indexing. Failures: {}", failed.len())); + } + + Ok(()) +}