diff --git a/crates/lib/mysql_queries/src/queries/model_weights/list_weights_query_builder.rs b/crates/lib/mysql_queries/src/queries/model_weights/list_weights_query_builder.rs index 355f60c0d8..fa0eda657c 100644 --- a/crates/lib/mysql_queries/src/queries/model_weights/list_weights_query_builder.rs +++ b/crates/lib/mysql_queries/src/queries/model_weights/list_weights_query_builder.rs @@ -2,7 +2,8 @@ #![forbid(unused_mut)] #![forbid(unused_variables)] -use chrono::{DateTime, Utc}; +use chrono::{ DateTime, Utc }; + use sqlx::MySqlPool; use config::shared_constants::DEFAULT_MYSQL_QUERY_RESULT_PAGE_SIZE; @@ -26,7 +27,6 @@ pub struct WeightsPage { pub last_id: Option, } - // TODO: figure out what you want to expose ? since we don't have specs really #[derive(Serialize)] pub struct WeightJoinUser { @@ -35,67 +35,64 @@ pub struct WeightJoinUser { pub weights_type: WeightsType, pub weights_category: WeightsCategory, - + pub title: String, - + pub maybe_thumbnail_token: Option, - + pub description_markdown: String, pub description_rendered_html: String, - + pub creator_user_token: UserToken, pub creator_ip_address: String, pub creator_set_visibility: Visibility, - + pub maybe_last_update_user_token: Option, - + pub original_download_url: Option, pub original_filename: Option, - + pub file_size_bytes: i32, pub file_checksum_sha2: String, - + pub private_bucket_hash: String, - + pub maybe_private_bucket_prefix: Option, pub maybe_private_bucket_extension: Option, - + pub cached_user_ratings_total_count: u32, pub cached_user_ratings_positive_count: u32, pub cached_user_ratings_negative_count: u32, pub maybe_cached_user_ratings_ratio: Option, pub cached_user_ratings_last_updated_at: DateTime, - + pub version: i32, - + pub created_at: DateTime, pub updated_at: DateTime, - + pub user_deleted_at: Option>, pub mod_deleted_at: Option>, - + pub creator_username: String, pub creator_display_name: String, pub creator_email_gravatar_hash: String, } - pub struct ListWeightsQueryBuilder { scope_creator_username: Option, include_mod_deleted_results: bool, include_user_deleted_results: bool, include_user_hidden: bool, sort_ascending: bool, - weights_type:Option, - weights_category:Option, + weights_type: Option, + weights_category: Option, offset: Option, limit: u16, cursor_is_reversed: bool, } - impl ListWeightsQueryBuilder { - pub fn new() -> Self { Self { scope_creator_username: None, @@ -115,7 +112,7 @@ impl ListWeightsQueryBuilder { self.scope_creator_username = scope_creator_username.map(|u| u.to_string()); self } - + pub fn include_user_hidden(mut self, include_user_hidden: bool) -> Self { self.include_user_hidden = include_user_hidden; self @@ -136,12 +133,12 @@ impl ListWeightsQueryBuilder { self } - pub fn weights_type(mut self,weights_type:WeightsType) -> Self { + pub fn weights_type(mut self, weights_type: WeightsType) -> Self { self.weights_type = Some(weights_type); self } - pub fn weights_category(mut self,weights_category:WeightsCategory) -> Self { + pub fn weights_category(mut self, weights_category: WeightsCategory) -> Self { self.weights_category = Some(weights_category); self } @@ -166,14 +163,13 @@ impl ListWeightsQueryBuilder { mysql_pool: &MySqlPool ) -> AnyhowResult> { let query = self.build_query_string(); + println!("HERE! query: {}", query); + //panic!("query: {:?}", query); + let mut query = sqlx::query_as::<_, RawWeightJoinUser>(&query); // NB: The following bindings must match the order of the query builder !! - if let Some(offset) = self.offset { - query = query.bind(offset); - } - if let Some(username) = self.scope_creator_username.as_deref() { query = query.bind(username); } @@ -188,71 +184,72 @@ impl ListWeightsQueryBuilder { query = query.bind(self.limit); + if let Some(offset) = self.offset { + query = query.bind(offset); + } - // debug!("query: {:?}", query); - // panic!("query: {:?}", query) - let mut results = query.fetch_all(mysql_pool) - .await?; + let mut results = query.fetch_all(mysql_pool).await?; if self.cursor_is_reversed { - results.reverse() + results.reverse(); } Ok(results) } - pub async fn perform_query_for_page( &self, mysql_pool: &MySqlPool ) -> AnyhowResult { - let weights = self.perform_internal_query(mysql_pool).await?; - let first_id = weights.first() - .map(|raw_result| raw_result.weight_id); + let first_id = weights.first().map(|raw_result| raw_result.weight_id); - let last_id = weights.last() - .map(|raw_result| raw_result.weight_id); + let last_id = weights.last().map(|raw_result| raw_result.weight_id); - let weights = weights.into_iter().map( - |record| { + let weights = weights + .into_iter() + .map(|record| { // map the raw result into a WeightJoinUser WeightJoinUser { - weight_id:record.weight_id, - token:record.token, - weights_type: - WeightsType::from_str(record.weights_type.as_str()).unwrap(), - weights_category:WeightsCategory::from_str(record.weights_category.as_str()).unwrap(), - title:record.title, - maybe_thumbnail_token:record.maybe_thumbnail_token, - description_markdown:record.description_markdown, - description_rendered_html:record.description_rendered_html, - creator_user_token: UserToken::new_from_str( &record.creator_user_token), - creator_ip_address:record.creator_ip_address, - creator_set_visibility: Visibility::from_str(&record.creator_set_visibility).unwrap(), - maybe_last_update_user_token: - record.maybe_last_update_user_token.map(|token| UserToken::new_from_str(&token)), - original_download_url:record.original_download_url, - original_filename:record.original_filename, - file_size_bytes:record.file_size_bytes, - file_checksum_sha2:record.file_checksum_sha2, - private_bucket_hash:record.private_bucket_hash, - maybe_private_bucket_prefix:record.maybe_private_bucket_prefix, - maybe_private_bucket_extension:record.maybe_private_bucket_extension, - cached_user_ratings_total_count:record.cached_user_ratings_total_count, - cached_user_ratings_positive_count:record.cached_user_ratings_positive_count, - cached_user_ratings_negative_count:record.cached_user_ratings_negative_count, - maybe_cached_user_ratings_ratio:record.maybe_cached_user_ratings_ratio, - cached_user_ratings_last_updated_at:record.cached_user_ratings_last_updated_at, - version:record.version, - created_at:record.created_at, - updated_at:record.updated_at, - user_deleted_at:record.user_deleted_at, - mod_deleted_at:record.mod_deleted_at, - creator_username:record.creator_username, - creator_display_name:record.creator_display_name, - creator_email_gravatar_hash:record.creator_email_gravatar_hash, + weight_id: record.weight_id, + token: record.token, + weights_type: WeightsType::from_str(record.weights_type.as_str()).unwrap(), + weights_category: WeightsCategory::from_str( + record.weights_category.as_str() + ).unwrap(), + title: record.title, + maybe_thumbnail_token: record.maybe_thumbnail_token, + description_markdown: record.description_markdown, + description_rendered_html: record.description_rendered_html, + creator_user_token: UserToken::new_from_str(&record.creator_user_token), + creator_ip_address: record.creator_ip_address, + creator_set_visibility: Visibility::from_str( + &record.creator_set_visibility + ).unwrap(), + maybe_last_update_user_token: record.maybe_last_update_user_token.map(|token| + UserToken::new_from_str(&token) + ), + original_download_url: record.original_download_url, + original_filename: record.original_filename, + file_size_bytes: record.file_size_bytes, + file_checksum_sha2: record.file_checksum_sha2, + private_bucket_hash: record.private_bucket_hash, + maybe_private_bucket_prefix: record.maybe_private_bucket_prefix, + maybe_private_bucket_extension: record.maybe_private_bucket_extension, + cached_user_ratings_total_count: record.cached_user_ratings_total_count, + cached_user_ratings_positive_count: record.cached_user_ratings_positive_count, + cached_user_ratings_negative_count: record.cached_user_ratings_negative_count, + maybe_cached_user_ratings_ratio: record.maybe_cached_user_ratings_ratio, + cached_user_ratings_last_updated_at: record.cached_user_ratings_last_updated_at, + version: record.version, + created_at: record.created_at, + updated_at: record.updated_at, + user_deleted_at: record.user_deleted_at, + mod_deleted_at: record.mod_deleted_at, + creator_username: record.creator_username, + creator_display_name: record.creator_display_name, + creator_email_gravatar_hash: record.creator_email_gravatar_hash, } }) .collect::>(); @@ -266,7 +263,8 @@ impl ListWeightsQueryBuilder { } pub fn build_query_string(&self) -> String { - let mut query = r#" + let mut query = + r#" SELECT model_weights.id AS weight_id, model_weights.token, @@ -328,18 +326,18 @@ impl ListWeightsQueryBuilder { if sort_ascending { if self.cursor_is_reversed { // NB: We're searching backwards. - query.push_str(" model_weights.id < ?"); + query.push_str(" model_weights.id"); sort_ascending = !sort_ascending; } else { - query.push_str(" model_weights.id > ?"); + query.push_str(" model_weights.id"); } } else { if self.cursor_is_reversed { // NB: We're searching backwards. - query.push_str(" model_weights.id > ?"); + query.push_str(" model_weights.id"); sort_ascending = !sort_ascending; } else { - query.push_str(" model_weights.id < ?"); + query.push_str(" model_weights.id"); } } } @@ -404,13 +402,18 @@ impl ListWeightsQueryBuilder { query.push_str(" ORDER BY model_weights.id DESC"); } - query.push_str(" LIMIT ?"); + if self.limit > 0 { + query.push_str(" LIMIT ?"); + } + + if let Some(_) = self.offset { + query.push_str(" OFFSET ?"); + } query } } - #[derive(Serialize)] #[derive(sqlx::FromRow)] struct RawWeightJoinUser { @@ -419,45 +422,45 @@ struct RawWeightJoinUser { pub weights_type: String, pub weights_category: String, - + pub title: String, - + pub maybe_thumbnail_token: Option, - + pub description_markdown: String, pub description_rendered_html: String, - + pub creator_user_token: String, pub creator_ip_address: String, pub creator_set_visibility: String, pub maybe_last_update_user_token: Option, - + pub original_download_url: Option, pub original_filename: Option, - + pub file_size_bytes: i32, pub file_checksum_sha2: String, - + pub private_bucket_hash: String, - + pub maybe_private_bucket_prefix: Option, pub maybe_private_bucket_extension: Option, - + pub cached_user_ratings_total_count: u32, pub cached_user_ratings_positive_count: u32, pub cached_user_ratings_negative_count: u32, pub maybe_cached_user_ratings_ratio: Option, pub cached_user_ratings_last_updated_at: DateTime, - + pub version: i32, - + pub created_at: DateTime, pub updated_at: DateTime, - + pub user_deleted_at: Option>, pub mod_deleted_at: Option>, - + pub creator_username: String, pub creator_display_name: String, pub creator_email_gravatar_hash: String, @@ -470,94 +473,104 @@ mod tests { weights_types::WeightsType, weights_category::WeightsCategory, }; - + #[test] fn predicates_without_scoping() { let query_builder = ListWeightsQueryBuilder::new(); - + println!("Query HERE! {}", &query_builder.build_predicates()); - assert_eq!(&query_builder.build_predicates(), - " WHERE model_weights.creator_set_visibility = 'public' \ + assert_eq!( + &query_builder.build_predicates(), + " WHERE model_weights.creator_set_visibility = 'public' \ AND model_weights.mod_deleted_at IS NULL \ AND model_weights.user_deleted_at IS NULL \ ORDER BY model_weights.id DESC \ - LIMIT ?"); + LIMIT ?" + ); } #[test] fn predicates_scoped_to_user() { - let query_builder = ListWeightsQueryBuilder::new() - .scope_creator_username(Some("echelon")); - assert_eq!(&query_builder.build_predicates(), - " WHERE users.username = ? \ + let query_builder = ListWeightsQueryBuilder::new().scope_creator_username(Some("echelon")); + assert_eq!( + &query_builder.build_predicates(), + " WHERE users.username = ? \ AND model_weights.creator_set_visibility = 'public' \ AND model_weights.mod_deleted_at IS NULL \ AND model_weights.user_deleted_at IS NULL \ ORDER BY model_weights.id DESC \ - LIMIT ?"); + LIMIT ?" + ); } #[test] fn predicates_scoped_to_weights_type() { - let query_builder = ListWeightsQueryBuilder::new() - .weights_type(WeightsType::RvcV2); - assert_eq!(&query_builder.build_predicates(), - " WHERE model_weights.weights_type = ? \ + let query_builder = ListWeightsQueryBuilder::new().weights_type(WeightsType::RvcV2); + assert_eq!( + &query_builder.build_predicates(), + " WHERE model_weights.weights_type = ? \ AND model_weights.creator_set_visibility = 'public' \ AND model_weights.mod_deleted_at IS NULL \ AND model_weights.user_deleted_at IS NULL \ ORDER BY model_weights.id DESC \ - LIMIT ?"); + LIMIT ?" + ); } #[test] fn predicates_scoped_to_weights_category() { - let query_builder = ListWeightsQueryBuilder::new() - .weights_category(WeightsCategory::VoiceConversion); + let query_builder = ListWeightsQueryBuilder::new().weights_category( + WeightsCategory::VoiceConversion + ); println!("Query HERE! {}", &query_builder.build_predicates()); - assert_eq!(&query_builder.build_predicates(), - " WHERE model_weights.weights_category = ? \ + assert_eq!( + &query_builder.build_predicates(), + " WHERE model_weights.weights_category = ? \ AND model_weights.creator_set_visibility = 'public' \ AND model_weights.mod_deleted_at IS NULL \ AND model_weights.user_deleted_at IS NULL \ ORDER BY model_weights.id DESC \ - LIMIT ?"); + LIMIT ?" + ); } #[test] fn predicates_including_user_hidden() { - let query_builder = ListWeightsQueryBuilder::new() - .include_user_hidden(true); + let query_builder = ListWeightsQueryBuilder::new().include_user_hidden(true); - assert_eq!(&query_builder.build_predicates(), - " WHERE model_weights.mod_deleted_at IS NULL \ + assert_eq!( + &query_builder.build_predicates(), + " WHERE model_weights.mod_deleted_at IS NULL \ AND model_weights.user_deleted_at IS NULL \ ORDER BY model_weights.id DESC \ - LIMIT ?"); + LIMIT ?" + ); } #[test] fn predicates_including_mod_deleted() { - let query_builder = ListWeightsQueryBuilder::new() - .include_mod_deleted_results(true); + let query_builder = ListWeightsQueryBuilder::new().include_mod_deleted_results(true); - assert_eq!(&query_builder.build_predicates(), - " WHERE model_weights.creator_set_visibility = 'public' \ + assert_eq!( + &query_builder.build_predicates(), + " WHERE model_weights.creator_set_visibility = 'public' \ AND model_weights.user_deleted_at IS NULL \ ORDER BY model_weights.id DESC \ - LIMIT ?"); + LIMIT ?" + ); } #[test] fn predicates_including_user_deleted() { - let query_builder = ListWeightsQueryBuilder::new() - .include_user_deleted_results(true); + let query_builder = ListWeightsQueryBuilder::new().include_user_deleted_results(true); - assert_eq!(&query_builder.build_predicates(), - " WHERE model_weights.creator_set_visibility = 'public' \ + assert_eq!( + &query_builder.build_predicates(), + " WHERE model_weights.creator_set_visibility = 'public' \ AND model_weights.mod_deleted_at IS NULL \ ORDER BY model_weights.id DESC \ - LIMIT ?"); + LIMIT ?" + ); } #[test] @@ -566,9 +579,11 @@ mod tests { .include_mod_deleted_results(true) .include_user_deleted_results(true); - assert_eq!(&query_builder.build_predicates(), - " WHERE model_weights.creator_set_visibility = 'public' \ + assert_eq!( + &query_builder.build_predicates(), + " WHERE model_weights.creator_set_visibility = 'public' \ ORDER BY model_weights.id DESC \ - LIMIT ?"); + LIMIT ?" + ); } -} \ No newline at end of file +} diff --git a/crates/lib/mysql_queries/src/queries/model_weights/tests.rs b/crates/lib/mysql_queries/src/queries/model_weights/tests.rs index 91615d3e11..23f706328f 100644 --- a/crates/lib/mysql_queries/src/queries/model_weights/tests.rs +++ b/crates/lib/mysql_queries/src/queries/model_weights/tests.rs @@ -2,6 +2,8 @@ mod tests { use std::result; + use anyhow::Ok; + use enums::by_table::model_weights::{weights_types, weights_category}; use sqlx::{MySqlPool, Acquire}; use rand::Rng; @@ -290,6 +292,154 @@ mod tests { Ok(()) } + + async fn seed_weights_with_category_and_type(weights_type:WeightsType,weights_category:WeightsCategory,number_of_items:u32,creator_username:&str) -> AnyhowResult<()> { + + let pool = setup().await; + let creator_token = get_user_token_by_username(&creator_username, &pool).await?.unwrap_or_else( + || panic!("Could not find user with username {}", creator_username) + ); + + for i in 0..number_of_items { + + let mut rng = rand::thread_rng(); + let number: u32 = rng.gen(); + + let model_weight_token = ModelWeightToken(number.to_string()); + + let args = CreateModelWeightsArgs { + token: &model_weight_token, // replace with actual ModelWeightToken + weights_type:weights_type, // replace with actual WeightsType + weights_category: weights_category, // replace with actual WeightsCategory + title: format!("Title {}", i), + maybe_thumbnail_token: Some(format!("Thumbnail {}", i)), + description_markdown: format!("Description {}", i), + description_rendered_html: format!("

Description {}

", i), + creator_user_token: Some(&creator_token), // replace with actual UserToken + creator_ip_address: "192.168.1.1", + creator_set_visibility: Visibility::Public, + maybe_last_update_user_token: Some(format!("Last Update User Token {}", i)), + original_download_url: Some(format!("http://example.com/download{}", i)), + original_filename: Some(format!("filename {}.txt", i)), + file_size_bytes: 1024, + file_checksum_sha2: format!("checksum{}", i), + private_bucket_hash: format!("bucket_hash{}", i), + maybe_private_bucket_prefix: Some("_fake".to_string()), + maybe_private_bucket_extension: Some("rvc".to_string()), + cached_user_ratings_total_count: 10, + cached_user_ratings_positive_count: 9, + cached_user_ratings_negative_count: 1, + maybe_cached_user_ratings_ratio: Some(0.9), + version: 1, + mysql_pool: &pool, // replace with actual MySqlPool + }; + + create_weight(args).await?; + } + + Ok(()) + + } + + // Tests paging for the list weights query builder + #[tokio::test] + async fn list_weights_query_build_test_paging() -> AnyhowResult<()> { + let pool = setup().await; + let creator_username = "hanashi".to_string(); + + seed_weights_with_category_and_type(WeightsType::LoRA,WeightsCategory::ImageGeneration,20,&creator_username).await?; + + let qb = ListWeightsQueryBuilder::new() + .offset(Some(0)) + .weights_type(WeightsType::LoRA) + .weights_category(WeightsCategory::ImageGeneration) + .limit(10) + .scope_creator_username(Some("hanashi")) + .include_user_deleted_results(false); + + let result = qb.perform_query_for_page(&pool).await?; + + // let result = qb.perform_query_for_page(&pool).await?; + assert_eq!(result.weights.len(), 10); + + let qb = ListWeightsQueryBuilder::new() + .weights_type(WeightsType::LoRA) + .weights_category(WeightsCategory::ImageGeneration) + .scope_creator_username(Some("hanashi")) + .include_user_deleted_results(false) + .offset(Some(10)) + .limit(10); + + let result = qb.perform_query_for_page(&pool).await?; + + assert_eq!(result.weights.len(), 10); + + + + let qb = ListWeightsQueryBuilder::new() + .weights_type(WeightsType::LoRA) + .weights_category(WeightsCategory::ImageGeneration) + .scope_creator_username(Some("hanashi")) + .include_user_deleted_results(false) + .offset(Some(10)) + .limit(10); + + let result = qb.perform_query_for_page(&pool).await?; + + assert_eq!(result.weights.len(), 10); + + Ok(()) + } + + #[tokio::test] + // asc and desc tests + async fn list_weights_query_build_test_asc_desc() -> AnyhowResult<()> { + let pool = setup().await; + let creator_username = "hanashi".to_string(); + + seed_weights_with_category_and_type(WeightsType::LoRA,WeightsCategory::ImageGeneration,10,&creator_username).await?; + + let qb = ListWeightsQueryBuilder::new() + .offset(Some(0)) + .weights_type(WeightsType::LoRA) + .weights_category(WeightsCategory::ImageGeneration) + .limit(10) + .scope_creator_username(Some("hanashi")) + .include_user_deleted_results(false) + .sort_ascending(true); + let result = qb.perform_query_for_page(&pool).await?; + + // write code that will loop through the result and check if the ids are in ascending order + let mut previous_id = 0; + for weight in result.weights.iter() { + let current_id = weight.weight_id; + assert!(current_id > previous_id); + previous_id = current_id; + } + + let qb = ListWeightsQueryBuilder::new() + .offset(Some(0)) + .weights_type(WeightsType::LoRA) + .weights_category(WeightsCategory::ImageGeneration) + .limit(10) + .scope_creator_username(Some("hanashi")) + .include_user_deleted_results(false) + .sort_ascending(false); + let result = qb.perform_query_for_page(&pool).await?; + + // write code that will loop through the result and check if the ids are in descending order + + let mut previous_id + = 10000; + for weight in result.weights.iter() { + let current_id = weight.weight_id; + assert!(current_id < previous_id); + previous_id = current_id; + } + + Ok(()) + } + #[tokio::test] async fn list_weights_query_build_test() -> AnyhowResult<()> { let pool = setup().await; @@ -374,8 +524,20 @@ mod tests { let result = qb.perform_query_for_page(&pool).await?; + assert_eq!(result.weights.len(), 2); + + let qb = ListWeightsQueryBuilder::new() + .weights_type(WeightsType::LoRA) + .weights_category(WeightsCategory::ImageGeneration) + .scope_creator_username(Some("hanashi")) + .include_user_deleted_results(false); + + let result = qb.perform_query_for_page(&pool).await?; assert_eq!(result.weights.len(), 1); + assert_eq!(result.weights[0].weights_category, WeightsCategory::ImageGeneration); + assert_eq!(result.weights[0].weights_type, WeightsType::LoRA); Ok(()) } } +