diff --git a/crates/cli/production/elasticsearch-cli/src/cli_args.rs b/crates/cli/production/elasticsearch-cli/src/cli_args.rs index aece674700..df7cddabf4 100644 --- a/crates/cli/production/elasticsearch-cli/src/cli_args.rs +++ b/crates/cli/production/elasticsearch-cli/src/cli_args.rs @@ -23,10 +23,16 @@ pub enum Environment { pub enum Action { ReindexTts, SearchTts, + /// Reindex model_weights into ElasticSearch ReindexModelWeights, + /// Search the model weights SearchModelWeights, + + /// Reindex media_files into ElasticSearch + /// For now, this is just media files that are relevant to the engine. + ReindexMediaFiles, } #[derive(Parser, Debug)] diff --git a/crates/cli/production/elasticsearch-cli/src/main.rs b/crates/cli/production/elasticsearch-cli/src/main.rs index 3edd5d8cba..9c14e1a259 100644 --- a/crates/cli/production/elasticsearch-cli/src/main.rs +++ b/crates/cli/production/elasticsearch-cli/src/main.rs @@ -20,6 +20,7 @@ use errors::AnyhowResult; use crate::cli_args::{Action, Environment, parse_cli_args}; use crate::plans::create_all_model_weight_documents::create_all_model_weight_documents; use crate::plans::create_all_tts_documents::create_all_tts_documents; +use crate::plans::create_dimensional_media_file_documents::create_dimensional_media_file_documents; pub mod cli_args; pub mod plans; @@ -62,6 +63,10 @@ pub async fn main() -> AnyhowResult<()> { println!("Result: {:#?}", result); } } + Action::ReindexMediaFiles => { + info!("Reindexing media files..."); + create_dimensional_media_file_documents(&mysql, &elasticsearch).await?; + } } info!("Done!"); diff --git a/crates/cli/production/elasticsearch-cli/src/plans/create_dimensional_media_file_documents.rs b/crates/cli/production/elasticsearch-cli/src/plans/create_dimensional_media_file_documents.rs index d1187239ad..32590eca14 100644 --- a/crates/cli/production/elasticsearch-cli/src/plans/create_dimensional_media_file_documents.rs +++ b/crates/cli/production/elasticsearch-cli/src/plans/create_dimensional_media_file_documents.rs @@ -1,14 +1,20 @@ +use std::collections::HashSet; use std::fs::read_to_string; +use std::iter::FromIterator; use elasticsearch::{BulkOperation, BulkParts, Elasticsearch}; use log::{error, info, warn}; use serde_json::Value; use sqlx::{MySql, Pool}; -use elasticsearch_schema::documents::media_file_document::MEDIA_FILE_INDEX; +use elasticsearch_schema::documents::media_file_document::{MEDIA_FILE_INDEX, MediaFileDocument}; use elasticsearch_schema::traits::document::Document; use elasticsearch_schema::utils::create_index_if_not_exists::{create_index_if_not_exists, CreateIndexArgs}; +use enums::by_table::media_files::media_file_class::MediaFileClass; +use enums::by_table::media_files::media_file_engine_category::MediaFileEngineCategory; +use enums::by_table::media_files::media_file_type::MediaFileType; use errors::AnyhowResult; +use mysql_queries::queries::media_files::list::list_media_files_for_elastic_search_backfill_using_cursor::{list_media_files_for_elastic_search_backfill_using_cursor, ListArgs, MediaFileForElasticsearchRecord}; use mysql_queries::queries::model_weights::list::list_model_weights_for_elastic_search_backfill_using_cursor::{list_model_weights_for_elastic_search_backfill_using_cursor, ModelWeightForElasticsearchRecord}; use storyteller_root::get_storyteller_rust_root; @@ -21,14 +27,41 @@ pub async fn create_dimensional_media_file_documents( // TODO(bt, 2024-01-10): expose this as a CLI flag const DELETE_EXISTING_INDEX : bool = true; - const PAGE_SIZE : u64 = 1000; + const PAGE_SIZE : usize = 1000; create_media_file_index(&elasticsearch, DELETE_EXISTING_INDEX).await?; let mut cursor = 0; loop { - let results = list_model_weights_for_elastic_search_backfill_using_cursor(mysql, PAGE_SIZE, cursor).await?; + let results = list_media_files_for_elastic_search_backfill_using_cursor(ListArgs { + mysql_pool: mysql, + page_size: PAGE_SIZE, + cursor, + maybe_filter_engine_categories: Some(&HashSet::from_iter(vec![ + MediaFileEngineCategory::Animation, + MediaFileEngineCategory::Character, + MediaFileEngineCategory::Expression, + MediaFileEngineCategory::ImagePlane, + MediaFileEngineCategory::Object, + MediaFileEngineCategory::Scene, + ])), + maybe_filter_media_types: Some(&HashSet::from_iter(vec![ + // Engine types + MediaFileType::Csv, + MediaFileType::Fbx, + MediaFileType::Glb, + MediaFileType::SceneJson, + // Image types + MediaFileType::Gif, + MediaFileType::Jpg, + MediaFileType::Png, + ])), + maybe_filter_media_classes: Some(&HashSet::from_iter(vec![ + MediaFileClass::Image, + MediaFileClass::Dimensional, + ])), + }).await?; if results.is_empty() { info!("No more results at cursor {cursor}"); @@ -38,7 +71,7 @@ pub async fn create_dimensional_media_file_documents( let maybe_last_id = results.last().map(|result| result.id); match maybe_last_id { - Some(last_id) => cursor = last_id as u64, + Some(last_id) => cursor = last_id as usize, None => { warn!("No final ID at cursor {cursor}"); break; @@ -53,83 +86,73 @@ pub async fn create_dimensional_media_file_documents( Ok(()) } -//async fn create_document_from_record(elasticsearch: &Elasticsearch, record: ModelWeightForElasticsearchRecord) -> AnyhowResult<()> { -// info!("Create record for {:?} - {:?}", record.token, record.title); -// -// let is_deleted = record.user_deleted_at.is_some() || record.mod_deleted_at.is_some(); -// -// 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.clone(), -// title_as_keyword: 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, -// -// is_deleted, -// }; -// -// 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(); -// -// // TODO: retry failures -// error!("Errors during indexing. Failures: {}", failed.len()); -// } -// -// Ok(()) -//} +async fn create_document_from_record(elasticsearch: &Elasticsearch, record: MediaFileForElasticsearchRecord) -> AnyhowResult<()> { + info!("Create record for {:?}", record.token); + + let is_deleted = record.user_deleted_at.is_some() || record.mod_deleted_at.is_some(); + + let document = MediaFileDocument { + token: record.token, + + media_class: record.media_class, + media_type: record.media_type, + maybe_media_subtype: record.maybe_media_subtype, + maybe_mime_type: record.maybe_mime_type, + public_bucket_directory_hash: record.public_bucket_directory_hash, + maybe_public_bucket_prefix: record.maybe_public_bucket_prefix, + maybe_public_bucket_extension: record.maybe_public_bucket_extension, + maybe_animation_type: record.maybe_animation_type, + creator_set_visibility: record.creator_set_visibility, + + maybe_title: record.maybe_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, + + maybe_creator_user_token: record.maybe_creator_user_token, + maybe_creator_username: record.maybe_creator_username, + maybe_creator_display_name: record.maybe_creator_display_name, + maybe_creator_gravatar_hash: record.maybe_creator_gravatar_hash, + + created_at: record.created_at, + updated_at: record.updated_at, + user_deleted_at: record.user_deleted_at, + mod_deleted_at: record.mod_deleted_at, + + is_deleted, + maybe_title_as_keyword: None, + }; + + let op : BulkOperation<_> = BulkOperation::index(&document) + .id(document.get_document_id()) + .into(); + + let response = elasticsearch + .bulk(BulkParts::Index(MEDIA_FILE_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(); + + // TODO: retry failures + error!("Errors during indexing. Failures: {}", failed.len()); + } + + Ok(()) +} // NB: Adapted from elasticsearch crate examples source async fn create_media_file_index(client: &Elasticsearch, delete_existing: bool) -> AnyhowResult<()> { diff --git a/crates/lib/elasticsearch_schema/src/documents/media_file_document.rs b/crates/lib/elasticsearch_schema/src/documents/media_file_document.rs index 1bdfa3b92e..c48acc4ac8 100644 --- a/crates/lib/elasticsearch_schema/src/documents/media_file_document.rs +++ b/crates/lib/elasticsearch_schema/src/documents/media_file_document.rs @@ -31,9 +31,9 @@ pub struct MediaFileDocument { pub maybe_mime_type: Option, - pub public_bucket_hash: Option, - pub public_bucket_prefix: Option, - pub public_bucket_extension: Option, + pub public_bucket_directory_hash: String, + pub maybe_public_bucket_prefix: Option, + pub maybe_public_bucket_extension: Option, pub maybe_animation_type: Option, diff --git a/crates/lib/mysql_queries/src/queries/media_files/list/list_media_files_for_elastic_search_backfill_using_cursor.rs b/crates/lib/mysql_queries/src/queries/media_files/list/list_media_files_for_elastic_search_backfill_using_cursor.rs index f112d6b1b0..5f3c347f2e 100644 --- a/crates/lib/mysql_queries/src/queries/media_files/list/list_media_files_for_elastic_search_backfill_using_cursor.rs +++ b/crates/lib/mysql_queries/src/queries/media_files/list/list_media_files_for_elastic_search_backfill_using_cursor.rs @@ -2,8 +2,8 @@ use std::collections::HashSet; use anyhow::anyhow; use chrono::{DateTime, Utc}; -use log::warn; -use sqlx::{FromRow, MySql, MySqlPool, QueryBuilder, Row}; +use log::{info, warn}; +use sqlx::{Execute, FromRow, MySql, MySqlPool, QueryBuilder, Row}; use sqlx::mysql::MySqlRow; use sqlx::pool::PoolConnection; @@ -15,8 +15,6 @@ use enums::by_table::media_files::media_file_origin_model_type::MediaFileOriginM use enums::by_table::media_files::media_file_origin_product_category::MediaFileOriginProductCategory; use enums::by_table::media_files::media_file_subtype::MediaFileSubtype; use enums::by_table::media_files::media_file_type::MediaFileType; -use enums::by_table::model_weights::weights_category::WeightsCategory; -use enums::by_table::model_weights::weights_types::WeightsType; use enums::common::view_as::ViewAs; use enums::common::visibility::Visibility; use enums::traits::mysql_from_row::MySqlFromRow; @@ -127,13 +125,16 @@ pub struct ListArgs<'a> { pub mysql_pool: &'a MySqlPool, } -pub async fn list_model_weights_for_elastic_search_backfill_using_cursor( +pub async fn list_media_files_for_elastic_search_backfill_using_cursor( args: ListArgs<'_> ) -> AnyhowResult> { let mut query = query_builder(&args); + let query = query.build_query_as::(); + info!("Query: {:?}", query.sql()); + let maybe_media_files = query.fetch_all(args.mysql_pool).await; let media_files : Vec = match maybe_media_files {