From df3ec24d2799e2dfac1b574d39c76379f4ec4181 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Wed, 15 May 2024 13:32:52 -0700 Subject: [PATCH] insert beta keys should work --- Cargo.lock | 1 + .../by_table/beta_keys/beta_key_product.rs | 109 +++++++++++++++++ .../enums/src/by_table/beta_keys/mod.rs | 1 + crates/_types/enums/src/by_table/mod.rs | 1 + crates/_types/tokens/src/prefixes.rs | 9 +- crates/_types/tokens/src/tokens/beta_keys.rs | 16 +++ crates/_types/tokens/src/tokens/mod.rs | 1 + .../beta_keys/insert_batch_beta_keys.rs | 110 ++++++++++++++++++ .../src/queries/beta_keys/mod.rs | 1 + crates/lib/mysql_queries/src/queries/mod.rs | 1 + crates/service/web/storyteller_web/Cargo.toml | 1 + .../beta_keys/create_beta_keys_handler.rs | 23 +++- 12 files changed, 269 insertions(+), 5 deletions(-) create mode 100644 crates/_types/enums/src/by_table/beta_keys/beta_key_product.rs create mode 100644 crates/_types/enums/src/by_table/beta_keys/mod.rs create mode 100644 crates/_types/tokens/src/tokens/beta_keys.rs create mode 100644 crates/lib/mysql_queries/src/queries/beta_keys/insert_batch_beta_keys.rs create mode 100644 crates/lib/mysql_queries/src/queries/beta_keys/mod.rs diff --git a/Cargo.lock b/Cargo.lock index 490350f97a..40e02972fd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6355,6 +6355,7 @@ dependencies = [ "config", "container_common", "cookies", + "crockford", "data-encoding", "datetimes", "derive_more", diff --git a/crates/_types/enums/src/by_table/beta_keys/beta_key_product.rs b/crates/_types/enums/src/by_table/beta_keys/beta_key_product.rs new file mode 100644 index 0000000000..3c9f00b514 --- /dev/null +++ b/crates/_types/enums/src/by_table/beta_keys/beta_key_product.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 BetaKeyProduct { + /// Media files + /// This will probably be the only type supported, but we'll leave the possibility of future types. + #[serde(rename = "studio")] + Studio, +} + +// TODO(bt, 2023-01-17): This desperately needs MySQL integration tests! +impl_enum_display_and_debug_using_to_str!(BetaKeyProduct); +impl_mysql_enum_coders!(BetaKeyProduct); +impl_mysql_from_row!(BetaKeyProduct); + +/// NB: Legacy API for older code. +impl BetaKeyProduct { + pub fn to_str(&self) -> &'static str { + match self { + Self::Studio => "studio", + } + } + + pub fn from_str(value: &str) -> Result { + match value { + "studio" => Ok(Self::Studio), + _ => 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::Studio, + ]) + } +} + +#[cfg(test)] +mod tests { + use crate::by_table::beta_keys::beta_key_product::BetaKeyProduct; + use crate::test_helpers::assert_serialization; + + mod serde { + use super::*; + + #[test] + fn test_serialization() { + assert_serialization(BetaKeyProduct::Studio, "studio"); + } + } + + mod impl_methods { + use super::*; + + #[test] + fn to_str() { + assert_eq!(BetaKeyProduct::Studio.to_str(), "studio"); + } + + #[test] + fn from_str() { + assert_eq!(BetaKeyProduct::from_str("studio").unwrap(), BetaKeyProduct::Studio); + assert!(BetaKeyProduct::from_str("foo").is_err()); + } + } + + mod manual_variant_checks { + use super::*; + + #[test] + fn all_variants() { + let mut variants = BetaKeyProduct::all_variants(); + assert_eq!(variants.len(), 1); + assert_eq!(variants.pop_first(), Some(BetaKeyProduct::Studio)); + assert_eq!(variants.pop_first(), None); + } + } + + mod mechanical_checks { + use super::*; + + #[test] + fn variant_length() { + use strum::IntoEnumIterator; + assert_eq!(BetaKeyProduct::all_variants().len(), BetaKeyProduct::iter().len()); + } + + #[test] + fn round_trip() { + for variant in BetaKeyProduct::all_variants() { + assert_eq!(variant, BetaKeyProduct::from_str(variant.to_str()).unwrap()); + assert_eq!(variant, BetaKeyProduct::from_str(&format!("{}", variant)).unwrap()); + assert_eq!(variant, BetaKeyProduct::from_str(&format!("{:?}", variant)).unwrap()); + } + } + } +} diff --git a/crates/_types/enums/src/by_table/beta_keys/mod.rs b/crates/_types/enums/src/by_table/beta_keys/mod.rs new file mode 100644 index 0000000000..a3484beb5d --- /dev/null +++ b/crates/_types/enums/src/by_table/beta_keys/mod.rs @@ -0,0 +1 @@ +pub mod beta_key_product; \ 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 0ea73ef4df..e31bf821c2 100644 --- a/crates/_types/enums/src/by_table/mod.rs +++ b/crates/_types/enums/src/by_table/mod.rs @@ -1,6 +1,7 @@ // ===== MySql ===== pub mod audit_logs; pub mod batch_generations; +pub mod beta_keys; pub mod comments; pub mod email_sender_jobs; pub mod entity_stats; diff --git a/crates/_types/tokens/src/prefixes.rs b/crates/_types/tokens/src/prefixes.rs index 7bbae30010..2fc57c5c88 100644 --- a/crates/_types/tokens/src/prefixes.rs +++ b/crates/_types/tokens/src/prefixes.rs @@ -25,6 +25,7 @@ pub(crate) enum TokenPrefix { AnonymousVisitorTracking, // AVTs are not stored as primary keys in any table, but an index in many tables. AuditLog, BatchGeneration, + BetaKey, Comment, DownloadJob, EmailSenderJob, @@ -93,6 +94,7 @@ impl PrefixGenerator for TokenPrefix { Self::AnonymousVisitorTracking => "avt_", Self::AuditLog => "audit_", Self::BatchGeneration => "batch_g_", + Self::BetaKey => "beta_key_", Self::Comment => "comment_", Self::DownloadJob => "jdown_", // NB: Previously "JGUP:" Self::EmailSenderJob => "email_job_", @@ -302,11 +304,12 @@ mod tests { 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_" || prefix == "batch_g_" + || prefix == "beta_key_" + || prefix == "email_job_" + || prefix == "pw_reset_" + || prefix == "tts_task_" { // 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. diff --git a/crates/_types/tokens/src/tokens/beta_keys.rs b/crates/_types/tokens/src/tokens/beta_keys.rs new file mode 100644 index 0000000000..c9d8078b0a --- /dev/null +++ b/crates/_types/tokens/src/tokens/beta_keys.rs @@ -0,0 +1,16 @@ +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 BetaKeyToken(pub String); + +impl_crockford_generator!(BetaKeyToken, 32usize, TokenPrefix::BetaKey, CrockfordLower); +impl_mysql_token_from_row!(BetaKeyToken); +impl_string_token!(BetaKeyToken); diff --git a/crates/_types/tokens/src/tokens/mod.rs b/crates/_types/tokens/src/tokens/mod.rs index 977a0bf577..d756c77bd8 100644 --- a/crates/_types/tokens/src/tokens/mod.rs +++ b/crates/_types/tokens/src/tokens/mod.rs @@ -3,6 +3,7 @@ pub mod api_tokens_external; pub mod api_tokens_internal; pub mod audit_logs; pub mod batch_generations; +pub mod beta_keys; pub mod comments; pub mod email_sender_job_token; pub mod firehose_entries; diff --git a/crates/lib/mysql_queries/src/queries/beta_keys/insert_batch_beta_keys.rs b/crates/lib/mysql_queries/src/queries/beta_keys/insert_batch_beta_keys.rs new file mode 100644 index 0000000000..ed3aaef8c4 --- /dev/null +++ b/crates/lib/mysql_queries/src/queries/beta_keys/insert_batch_beta_keys.rs @@ -0,0 +1,110 @@ +use anyhow::anyhow; +use log::error; +use sqlx::{MySql, MySqlPool, QueryBuilder, Transaction}; + +use composite_identifiers::by_table::batch_generations::batch_generation_entity::BatchGenerationEntity; +use enums::by_table::beta_keys::beta_key_product::BetaKeyProduct; +use errors::AnyhowResult; +use tokens::tokens::batch_generations::BatchGenerationToken; +use tokens::tokens::beta_keys::BetaKeyToken; +use tokens::tokens::users::UserToken; + +pub struct BatchEntry { + pub batch_token: BatchGenerationToken, + pub entity: BatchGenerationEntity, +} + +pub struct InsertBatchArgs<'a> { + pub product: BetaKeyProduct, + pub maybe_referrer_user_token: Option<&'a UserToken>, + pub beta_keys: &'a Vec, + pub mysql_pool: &'a MySqlPool, +} + +pub async fn insert_batch_beta_keys<'a, 'b>(args: InsertBatchArgs<'a>) -> AnyhowResult<()> { + + let mut query_builder = QueryBuilder::new(r#" + INSERT INTO beta_keys ( + token, + product, + key_value, + maybe_referrer_user_token + ) VALUES + "#); + + for (i, beta_key) in args.beta_keys.iter().enumerate() { + let token = BetaKeyToken::generate(); + + query_builder.push("("); + + query_builder.push_bind(token.to_string()); + query_builder.push(","); + + query_builder.push_bind(args.product.to_str()); + query_builder.push(","); + + query_builder.push_bind(beta_key); + query_builder.push(","); + + query_builder.push_bind(args.maybe_referrer_user_token.map(|t| t.as_str())); + + query_builder.push(")"); + + if i < args.beta_keys.len() - 1 { + query_builder.push(","); + } + } + + let query = query_builder.build(); + + let query_result = query.execute(args.mysql_pool).await; + + match query_result { + Ok(_) => Ok(()), + Err(err) => { + error!("Error with batch generation query: {:?}", &err); + Err(anyhow!("model category creation error: {:?}", &err)) + }, + } +} + +// #[ignore] +// #[cfg(test)] +// mod tests { +// use sqlx::mysql::MySqlPoolOptions; +// +// use composite_identifiers::by_table::batch_generations::batch_generation_entity::BatchGenerationEntity; +// use config::shared_constants::DEFAULT_MYSQL_CONNECTION_STRING; +// use tokens::tokens::media_files::MediaFileToken; +// +// use crate::queries::batch_generations::insert_batch_generation_records::{insert_batch_generation_records, InsertBatchArgs}; +// +// async fn setup() -> sqlx::Pool { +// MySqlPoolOptions::new() +// .max_connections(3) +// .connect(&DEFAULT_MYSQL_CONNECTION_STRING).await +// .unwrap() +// } +// +// #[ignore] +// #[tokio::test] +// async fn test() { +// let pool = setup().await; +// +// let mut transaction = pool.begin().await.unwrap(); +// +// let entries = vec![ +// BatchGenerationEntity::MediaFile(MediaFileToken::new_from_str("media_foo")), +// BatchGenerationEntity::MediaFile(MediaFileToken::new_from_str("media_bar")), +// BatchGenerationEntity::MediaFile(MediaFileToken::new_from_str("media_baz")), +// ]; +// +// let r = insert_batch_generation_records(InsertBatchArgs { +// entries, +// maybe_existing_batch_token: None, +// transaction: &mut transaction, +// }).await; +// +// transaction.commit().await.unwrap(); +// } +// } \ No newline at end of file diff --git a/crates/lib/mysql_queries/src/queries/beta_keys/mod.rs b/crates/lib/mysql_queries/src/queries/beta_keys/mod.rs new file mode 100644 index 0000000000..10bfbf6f73 --- /dev/null +++ b/crates/lib/mysql_queries/src/queries/beta_keys/mod.rs @@ -0,0 +1 @@ +pub mod insert_batch_beta_keys; \ No newline at end of file diff --git a/crates/lib/mysql_queries/src/queries/mod.rs b/crates/lib/mysql_queries/src/queries/mod.rs index aa5e17aa0f..39f968345e 100644 --- a/crates/lib/mysql_queries/src/queries/mod.rs +++ b/crates/lib/mysql_queries/src/queries/mod.rs @@ -1,6 +1,7 @@ pub mod api_tokens; pub mod audit_logs; pub mod batch_generations; +pub mod beta_keys; pub mod billing; pub mod comments; pub mod email_sender_jobs; diff --git a/crates/service/web/storyteller_web/Cargo.toml b/crates/service/web/storyteller_web/Cargo.toml index 6e894386bb..9e97d8205b 100644 --- a/crates/service/web/storyteller_web/Cargo.toml +++ b/crates/service/web/storyteller_web/Cargo.toml @@ -34,6 +34,7 @@ composite_identifiers = { path = "../../../_types/composite_identifiers" } config = { path = "../../../lib/deprecated/config" } container_common = { path = "../../../lib/deprecated/container_common" } cookies = { path = "../../../lib/cookies" } +crockford = { path = "../../../lib/data/crockford" } datetimes = { path = "../../../lib/data/datetimes" } easyenv = { path = "../../../lib/easyenv" } elasticsearch_schema = { path = "../../../lib/elasticsearch_schema" } diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/beta_keys/create_beta_keys_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/beta_keys/create_beta_keys_handler.rs index bac159ce1a..494ac7c5b9 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/beta_keys/create_beta_keys_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/beta_keys/create_beta_keys_handler.rs @@ -7,8 +7,11 @@ use actix_web::http::StatusCode; use log::warn; use utoipa::ToSchema; +use crockford::crockford_entropy_lower; +use enums::by_table::beta_keys::beta_key_product::BetaKeyProduct; use enums::by_table::comments::comment_entity_type::CommentEntityType; use http_server_common::request::get_request_ip::get_request_ip; +use mysql_queries::queries::beta_keys::insert_batch_beta_keys::{insert_batch_beta_keys, InsertBatchArgs}; use mysql_queries::queries::comments::comment_entity_token::CommentEntityToken; use mysql_queries::queries::comments::insert_comment::{insert_comment, InsertCommentArgs}; use tokens::tokens::comments::CommentToken; @@ -27,6 +30,8 @@ use crate::http_server::web_utils::require_moderator::{require_moderator, Requir use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; use crate::server_state::ServerState; +const MAXIMUM_KEYS : u32 = 100; + #[derive(Deserialize, ToSchema)] pub struct CreateBetaKeysRequest { uuid_idempotency_token: String, @@ -102,11 +107,25 @@ pub async fn create_beta_keys_handler( RequireModeratorError::NotAuthorized => CreateBetaKeysError::NotAuthorized, })?; - // TODO(bt,2024-05-13): Create beta key records + let number_of_keys = request.number_of_keys.min(MAXIMUM_KEYS); + + let beta_keys = (0..number_of_keys).map(|_| { + crockford_entropy_lower(8) + }).collect::>(); + + insert_batch_beta_keys(InsertBatchArgs { + product: BetaKeyProduct::Studio, + maybe_referrer_user_token: None, + beta_keys: &beta_keys, + mysql_pool: &server_state.mysql_pool, + }).await.map_err(|err| { + warn!("Error inserting beta keys: {:?}", err); + CreateBetaKeysError::ServerError + })?; let response = CreateBetaKeysSuccessResponse { success: true, - beta_keys: Vec::new(), + beta_keys, }; let body = serde_json::to_string(&response)