From 33613cb5d8b6e6679533cd04c0c737d66f3b0779 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Tue, 16 Jan 2024 23:00:17 -0500 Subject: [PATCH] insert batch generation records --- Cargo.lock | 21 ++++ Cargo.toml | 1 + .../_types/composite_identifiers/Cargo.toml | 40 +++++++ .../batch_generation_entity.rs | 28 +++++ .../src/by_table/batch_generations/mod.rs | 1 + .../composite_identifiers/src/by_table/mod.rs | 1 + .../_types/composite_identifiers/src/lib.rs | 22 ++++ .../batch_generation_entity_type.rs | 109 ++++++++++++++++++ .../src/by_table/batch_generations/mod.rs | 1 + crates/_types/enums/src/by_table/mod.rs | 2 + crates/_types/tokens/src/prefixes.rs | 10 +- .../tokens/src/tokens/batch_generations.rs | 15 +++ crates/_types/tokens/src/tokens/mod.rs | 3 +- crates/lib/mysql_queries/Cargo.toml | 1 + .../insert_batch_generation_records.rs | 54 +++++++++ .../src/queries/batch_generations/mod.rs | 1 + crates/lib/mysql_queries/src/queries/mod.rs | 3 +- crates/service/job/inference_job/Cargo.toml | 1 + crates/service/web/storyteller_web/Cargo.toml | 1 + 19 files changed, 311 insertions(+), 4 deletions(-) create mode 100644 crates/_types/composite_identifiers/Cargo.toml create mode 100644 crates/_types/composite_identifiers/src/by_table/batch_generations/batch_generation_entity.rs create mode 100644 crates/_types/composite_identifiers/src/by_table/batch_generations/mod.rs create mode 100644 crates/_types/composite_identifiers/src/by_table/mod.rs create mode 100644 crates/_types/composite_identifiers/src/lib.rs create mode 100644 crates/_types/enums/src/by_table/batch_generations/batch_generation_entity_type.rs create mode 100644 crates/_types/enums/src/by_table/batch_generations/mod.rs create mode 100644 crates/_types/tokens/src/tokens/batch_generations.rs create mode 100644 crates/lib/mysql_queries/src/queries/batch_generations/insert_batch_generation_records.rs create mode 100644 crates/lib/mysql_queries/src/queries/batch_generations/mod.rs diff --git a/Cargo.lock b/Cargo.lock index 5e29620c49..e0b6036e71 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1259,6 +1259,24 @@ dependencies = [ "my-workspace-hack", ] +[[package]] +name = "composite_identifiers" +version = "0.0.1" +dependencies = [ + "enums", + "my-workspace-hack", + "serde", + "serde_derive", + "serde_json", + "sqlx", + "sqlx-core", + "sqlx-mysql", + "strum", + "tokens", + "toml 0.5.11", + "utoipa", +] + [[package]] name = "concurrency" version = "0.0.1" @@ -3052,6 +3070,7 @@ dependencies = [ "chrono", "cloud_storage", "collections", + "composite_identifiers", "concurrency", "config", "dotenv", @@ -3857,6 +3876,7 @@ dependencies = [ "anyhow", "buckets", "chrono", + "composite_identifiers", "config", "container_common", "dotenv", @@ -6006,6 +6026,7 @@ dependencies = [ "buckets", "chrono", "cloud_storage", + "composite_identifiers", "config", "container_common", "cookies", diff --git a/Cargo.toml b/Cargo.toml index 0bd2f6e824..4b04d5e36c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,6 +5,7 @@ members = [ # Types: common strongly-typed tokens, ids, payloads, etc. # Also common cloud buckets and other "strongly typed" configs. "crates/_types/buckets", + "crates/_types/composite_identifiers", "crates/_types/enums", "crates/_types/reusable_types", # TODO(bt, 2022-12-21): Deprecate "crates/_types/tokens", diff --git a/crates/_types/composite_identifiers/Cargo.toml b/crates/_types/composite_identifiers/Cargo.toml new file mode 100644 index 0000000000..f9aa2156b2 --- /dev/null +++ b/crates/_types/composite_identifiers/Cargo.toml @@ -0,0 +1,40 @@ +[package] +name = "composite_identifiers" +edition = "2018" +version = "0.0.1" +authors = [ + "Brandon Thomas ", + "Brandon Thomas ", +] +publish = false + +[lib] +name = "composite_identifiers" +path = "src/lib.rs" + +[dependencies] +# Internal +enums = { path = "../enums" } +tokens = { path = "../tokens" } + +# Workspace hack (faster compile times) +my-workspace-hack = { version = "0.1", path = "../../my-workspace-hack" } + +# Serialization +serde = "1.0.125" +serde_derive = "1.0.125" + +sqlx = { version = "0.7.2", features = [ "mysql", "sqlite", "runtime-tokio-native-tls", "chrono" ] } +sqlx-core = { version = "0.7.2" } +sqlx-mysql= { version = "0.7.2" } + +utoipa = { version = "4", features = ["chrono","actix_extras"] } + +[dev-dependencies] + +# Enum utilities +strum = { version = "0.25", features = ["derive"] } + +# Test serialization +serde_json = "1.0.108" +toml = "0.5.8" diff --git a/crates/_types/composite_identifiers/src/by_table/batch_generations/batch_generation_entity.rs b/crates/_types/composite_identifiers/src/by_table/batch_generations/batch_generation_entity.rs new file mode 100644 index 0000000000..0c9a261314 --- /dev/null +++ b/crates/_types/composite_identifiers/src/by_table/batch_generations/batch_generation_entity.rs @@ -0,0 +1,28 @@ +use enums::by_table::batch_generations::batch_generation_entity_type::BatchGenerationEntityType; +use tokens::tokens::media_files::MediaFileToken; + +// TODO(bt, 2024-01-16): A ref type with <'a> lifetime of inner data instead of ownership? +pub enum BatchGenerationEntity { + MediaFile(MediaFileToken), +} + +// TODO(bt, 2024-01-16): Codegen? Make traits for these? Maybe overkill. +impl BatchGenerationEntity { + pub fn get_entity_type(&self) -> BatchGenerationEntityType { + match self { + Self::MediaFile(_) => BatchGenerationEntityType::MediaFile, + } + } + + pub fn get_entity_token_str(&self) -> &str { + match self { + Self::MediaFile(token) => token.as_str(), + } + } + + pub fn get_composite_keys(&self) -> (BatchGenerationEntityType, &str) { + match self { + Self::MediaFile(token) => (BatchGenerationEntityType::MediaFile, token.as_str()), + } + } +} diff --git a/crates/_types/composite_identifiers/src/by_table/batch_generations/mod.rs b/crates/_types/composite_identifiers/src/by_table/batch_generations/mod.rs new file mode 100644 index 0000000000..0f4d8b5201 --- /dev/null +++ b/crates/_types/composite_identifiers/src/by_table/batch_generations/mod.rs @@ -0,0 +1 @@ +pub mod batch_generation_entity; \ No newline at end of file diff --git a/crates/_types/composite_identifiers/src/by_table/mod.rs b/crates/_types/composite_identifiers/src/by_table/mod.rs new file mode 100644 index 0000000000..4bbf977cdb --- /dev/null +++ b/crates/_types/composite_identifiers/src/by_table/mod.rs @@ -0,0 +1 @@ +pub mod batch_generations; diff --git a/crates/_types/composite_identifiers/src/lib.rs b/crates/_types/composite_identifiers/src/lib.rs new file mode 100644 index 0000000000..08d90188da --- /dev/null +++ b/crates/_types/composite_identifiers/src/lib.rs @@ -0,0 +1,22 @@ +//! composite_identifiers +//! +//! The purpose of this library is to provide easy to construct, **typesafe**, composite identifiers. +//! This is mostly a combination of `tokens` and `enums`. +//! + +// Never allow these +#![forbid(private_bounds)] +#![forbid(private_interfaces)] +#![forbid(unused_must_use)] // NB: It's unsafe to not close/check some things + +// Okay to toggle +#![forbid(unreachable_patterns)] +#![forbid(unused_imports)] +#![forbid(unused_mut)] +#![forbid(unused_variables)] + +// Always allow +#![allow(dead_code)] +#![allow(non_snake_case)] + +pub mod by_table; diff --git a/crates/_types/enums/src/by_table/batch_generations/batch_generation_entity_type.rs b/crates/_types/enums/src/by_table/batch_generations/batch_generation_entity_type.rs new file mode 100644 index 0000000000..41633225f9 --- /dev/null +++ b/crates/_types/enums/src/by_table/batch_generations/batch_generation_entity_type.rs @@ -0,0 +1,109 @@ +use std::collections::BTreeSet; + +use serde::Deserialize; +use serde::Serialize; +#[cfg(test)] +use strum::EnumCount; +#[cfg(test)] +use strum::EnumIter; +use utoipa::ToSchema; + +/// Used in the `batch_generations` table in a `VARCHAR(32)` field named `entity_type`. +#[cfg_attr(test, derive(EnumIter, EnumCount))] +#[derive(Clone, Copy, Eq, PartialEq, Hash, Deserialize, Serialize, Ord, PartialOrd, ToSchema)] +pub enum BatchGenerationEntityType { + /// Media files + /// This will probably be the only type supported, but we'll leave the possibility of future types. + #[serde(rename = "media_file")] + MediaFile, +} + +// TODO(bt, 2023-01-17): This desperately needs MySQL integration tests! +impl_enum_display_and_debug_using_to_str!(BatchGenerationEntityType); +impl_mysql_enum_coders!(BatchGenerationEntityType); +impl_mysql_from_row!(BatchGenerationEntityType); + +/// NB: Legacy API for older code. +impl BatchGenerationEntityType { + pub fn to_str(&self) -> &'static str { + match self { + Self::MediaFile => "media_file", + } + } + + pub fn from_str(value: &str) -> Result { + match value { + "media_file" => Ok(Self::MediaFile), + _ => Err(format!("invalid value: {:?}", value)), + } + } + + pub fn all_variants() -> BTreeSet { + // NB: BTreeSet is sorted + // NB: BTreeSet::from() isn't const, but not worth using LazyStatic, etc. + BTreeSet::from([ + Self::MediaFile, + ]) + } +} + +#[cfg(test)] +mod tests { + use crate::by_table::batch_generations::batch_generation_entity_type::BatchGenerationEntityType; + use crate::test_helpers::assert_serialization; + + mod serde { + use super::*; + + #[test] + fn test_serialization() { + assert_serialization(BatchGenerationEntityType::MediaFile, "media_file"); + } + } + + mod impl_methods { + use super::*; + + #[test] + fn to_str() { + assert_eq!(BatchGenerationEntityType::MediaFile.to_str(), "media_file"); + } + + #[test] + fn from_str() { + assert_eq!(BatchGenerationEntityType::from_str("media_file").unwrap(), BatchGenerationEntityType::MediaFile); + assert!(BatchGenerationEntityType::from_str("foo").is_err()); + } + } + + mod manual_variant_checks { + use super::*; + + #[test] + fn all_variants() { + let mut variants = BatchGenerationEntityType::all_variants(); + assert_eq!(variants.len(), 1); + assert_eq!(variants.pop_first(), Some(BatchGenerationEntityType::MediaFile)); + assert_eq!(variants.pop_first(), None); + } + } + + mod mechanical_checks { + use super::*; + + #[test] + fn variant_length() { + use strum::IntoEnumIterator; + assert_eq!(BatchGenerationEntityType::all_variants().len(), BatchGenerationEntityType::iter().len()); + } + + #[test] + fn round_trip() { + for variant in BatchGenerationEntityType::all_variants() { + assert_eq!(variant, BatchGenerationEntityType::from_str(variant.to_str()).unwrap()); + assert_eq!(variant, BatchGenerationEntityType::from_str(&format!("{}", variant)).unwrap()); + assert_eq!(variant, BatchGenerationEntityType::from_str(&format!("{:?}", variant)).unwrap()); + } + } + } +} diff --git a/crates/_types/enums/src/by_table/batch_generations/mod.rs b/crates/_types/enums/src/by_table/batch_generations/mod.rs new file mode 100644 index 0000000000..2babf4fabc --- /dev/null +++ b/crates/_types/enums/src/by_table/batch_generations/mod.rs @@ -0,0 +1 @@ +pub mod batch_generation_entity_type; \ No newline at end of file diff --git a/crates/_types/enums/src/by_table/mod.rs b/crates/_types/enums/src/by_table/mod.rs index b3134ad59d..0c0169b2ac 100644 --- a/crates/_types/enums/src/by_table/mod.rs +++ b/crates/_types/enums/src/by_table/mod.rs @@ -1,5 +1,6 @@ // ===== MySql ===== pub mod audit_logs; +pub mod batch_generations; pub mod comments; pub mod email_sender_jobs; pub mod entity_stats; @@ -17,6 +18,7 @@ pub mod user_ratings; pub mod voice_conversion_models; pub mod voice_conversion_results; pub mod zs_voices; + // ===== Sqlite ===== pub mod tts_render_tasks; pub mod web_rendition_targets; diff --git a/crates/_types/tokens/src/prefixes.rs b/crates/_types/tokens/src/prefixes.rs index 382492bae4..e4263e21bb 100644 --- a/crates/_types/tokens/src/prefixes.rs +++ b/crates/_types/tokens/src/prefixes.rs @@ -24,6 +24,7 @@ use strum::EnumIter; pub(crate) enum TokenPrefix { AnonymousVisitorTracking, // AVTs are not stored as primary keys in any table, but an index in many tables. AuditLog, + BatchGeneration, Comment, DownloadJob, EmailSenderJob, @@ -90,6 +91,7 @@ impl PrefixGenerator for TokenPrefix { match self { Self::AnonymousVisitorTracking => "avt_", Self::AuditLog => "audit_", + Self::BatchGeneration => "batch_g_", Self::Comment => "comment_", Self::DownloadJob => "jdown_", // NB: Previously "JGUP:" Self::EmailSenderJob => "email_job_", @@ -297,7 +299,13 @@ mod tests { #[test] fn test_all_prefixes_end_with_separator_length_one() { for prefix in get_all_prefixes().iter().map(|s| *s) { - if prefix == "news_story_" || prefix == "tts_task_" || prefix == "INT_API:" || prefix == "pw_reset_" || prefix == "email_job_" { + if prefix == "news_story_" + || prefix == "tts_task_" + || prefix == "INT_API:" + || prefix == "pw_reset_" + || prefix == "email_job_" + || prefix == "batch_g_" + { // TODO/FIXME: I'm too tired at 5AM to replacen from the left. Make this test valid. // These tokens are from the AIChatBot sidecar, so asserting their validity is less important. continue; diff --git a/crates/_types/tokens/src/tokens/batch_generations.rs b/crates/_types/tokens/src/tokens/batch_generations.rs new file mode 100644 index 0000000000..052681f21a --- /dev/null +++ b/crates/_types/tokens/src/tokens/batch_generations.rs @@ -0,0 +1,15 @@ +use std::fmt::Debug; + +use serde::Deserialize; +use serde::Serialize; +use utoipa::ToSchema; + +use crate::prefixes::TokenPrefix; + +/// The primary key for Audit Logs. +#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, sqlx::Type, Debug, Serialize, Deserialize, ToSchema)] +#[sqlx(transparent)] +pub struct BatchGenerationToken(pub String); + +impl_string_token!(BatchGenerationToken); +impl_crockford_generator!(BatchGenerationToken, 32usize, TokenPrefix::BatchGeneration, CrockfordLower); diff --git a/crates/_types/tokens/src/tokens/mod.rs b/crates/_types/tokens/src/tokens/mod.rs index b7c4eed592..447cbedf43 100644 --- a/crates/_types/tokens/src/tokens/mod.rs +++ b/crates/_types/tokens/src/tokens/mod.rs @@ -2,6 +2,7 @@ pub mod anonymous_visitor_tracking; pub mod api_tokens_external; pub mod api_tokens_internal; pub mod audit_logs; +pub mod batch_generations; pub mod comments; pub mod email_sender_job_token; pub mod firehose_entries; @@ -20,6 +21,7 @@ pub mod tts_results; pub mod twitch_event_rule; pub mod twitch_oauth_tokens_grouping; pub mod twitch_oauth_tokens_internal; +pub mod user_bookmarks; pub mod user_sessions; pub mod user_subscriptions; pub mod users; @@ -34,4 +36,3 @@ pub mod w2l_templates; pub mod zs_voice_dataset_samples; pub mod zs_voice_datasets; pub mod zs_voices; -pub mod user_bookmarks; diff --git a/crates/lib/mysql_queries/Cargo.toml b/crates/lib/mysql_queries/Cargo.toml index 228c472322..2dcf665e29 100644 --- a/crates/lib/mysql_queries/Cargo.toml +++ b/crates/lib/mysql_queries/Cargo.toml @@ -16,6 +16,7 @@ path = "src/lib.rs" # Internal buckets = { path = "../../_types/buckets" } +composite_identifiers = { path = "../../_types/composite_identifiers" } config = { path = "../deprecated/config" } container_common = { path = "../deprecated/container_common" } enums = { path = "../../_types/enums" } diff --git a/crates/lib/mysql_queries/src/queries/batch_generations/insert_batch_generation_records.rs b/crates/lib/mysql_queries/src/queries/batch_generations/insert_batch_generation_records.rs new file mode 100644 index 0000000000..a572c42257 --- /dev/null +++ b/crates/lib/mysql_queries/src/queries/batch_generations/insert_batch_generation_records.rs @@ -0,0 +1,54 @@ +use anyhow::anyhow; +use log::error; +use sqlx::{Executor, MySql, QueryBuilder, Transaction}; + +use composite_identifiers::by_table::batch_generations::batch_generation_entity::BatchGenerationEntity; +use errors::AnyhowResult; +use tokens::tokens::batch_generations::BatchGenerationToken; + +pub struct BatchEntry { + pub batch_token: BatchGenerationToken, + pub entity: BatchGenerationEntity, +} + +pub struct InsertBatchArgs<'a> { + pub entries: Vec, + pub transaction: &'a mut Transaction<'a, MySql>, +} + +/// NB: Caller is responsible for rolling back the transaction if this fails. +pub async fn insert_batch_generation_records(args: InsertBatchArgs<'_>) -> AnyhowResult<()> { + let mut query_builder = QueryBuilder::new(r#" + INSERT INTO batch_generations (token, entity_type, entity_token) VALUES + "#); + + for (i, entry) in args.entries.iter().enumerate() { + let (entity_type, entity_token) = entry.entity.get_composite_keys(); + + query_builder.push("("); + query_builder.push_bind(&entry.batch_token); + query_builder.push(","); + + query_builder.push_bind(entity_type); + query_builder.push(","); + + query_builder.push_bind(entity_token); + query_builder.push(")"); + + if i < args.entries.len() - 1 { + query_builder.push(","); + } + } + + let query = query_builder.build(); + + let query_result = query.execute(&mut **args.transaction).await; + + match query_result { + Ok(_) => Ok(()), + Err(err) => { + error!("Error with batch generation query: {:?}", &err); + Err(anyhow!("model category creation error: {:?}", &err)) + }, + } +} diff --git a/crates/lib/mysql_queries/src/queries/batch_generations/mod.rs b/crates/lib/mysql_queries/src/queries/batch_generations/mod.rs new file mode 100644 index 0000000000..5f34676f55 --- /dev/null +++ b/crates/lib/mysql_queries/src/queries/batch_generations/mod.rs @@ -0,0 +1 @@ +pub mod insert_batch_generation_records; diff --git a/crates/lib/mysql_queries/src/queries/mod.rs b/crates/lib/mysql_queries/src/queries/mod.rs index 55f1039245..86079b7494 100644 --- a/crates/lib/mysql_queries/src/queries/mod.rs +++ b/crates/lib/mysql_queries/src/queries/mod.rs @@ -1,5 +1,6 @@ pub mod api_tokens; pub mod audit_logs; +pub mod batch_generations; pub mod billing; pub mod comments; pub mod email_sender_jobs; @@ -24,5 +25,3 @@ pub mod voice_clone_requests; pub mod voice_conversion; pub mod voice_designer; pub mod w2l; - - diff --git a/crates/service/job/inference_job/Cargo.toml b/crates/service/job/inference_job/Cargo.toml index 19e17e422b..5644d6e950 100644 --- a/crates/service/job/inference_job/Cargo.toml +++ b/crates/service/job/inference_job/Cargo.toml @@ -20,6 +20,7 @@ bootstrap = { path = "../../../lib/bootstrap" } buckets = { path = "../../../_types/buckets" } cloud_storage = { path = "../../../lib/api_clients/cloud_storage" } collections = { path = "../../../lib/collections" } +composite_identifiers = { path = "../../../_types/composite_identifiers" } concurrency = { path = "../../../lib/concurrency" } config = { path = "../../../lib/deprecated/config" } easyenv = { path = "../../../lib/easyenv" } diff --git a/crates/service/web/storyteller_web/Cargo.toml b/crates/service/web/storyteller_web/Cargo.toml index d8224d0269..5c959b2966 100644 --- a/crates/service/web/storyteller_web/Cargo.toml +++ b/crates/service/web/storyteller_web/Cargo.toml @@ -24,6 +24,7 @@ actix_helpers = { path = "../../../lib/actix_helpers" } bootstrap = { path = "../../../lib/bootstrap" } buckets = { path = "../../../_types/buckets" } cloud_storage = { path = "../../../lib/api_clients/cloud_storage" } +composite_identifiers = { path = "../../../_types/composite_identifiers" } config = { path = "../../../lib/deprecated/config" } container_common = { path = "../../../lib/deprecated/container_common" } cookies = { path = "../../../lib/cookies" }