diff --git a/.sqlx/query-19c1458e12d614bfface303102898b6a34f6b7a56eb906f5df3568172d51778b.json b/.sqlx/query-19c1458e12d614bfface303102898b6a34f6b7a56eb906f5df3568172d51778b.json new file mode 100644 index 0000000000..b4b424bdac --- /dev/null +++ b/.sqlx/query-19c1458e12d614bfface303102898b6a34f6b7a56eb906f5df3568172d51778b.json @@ -0,0 +1,12 @@ +{ + "db_name": "MySQL", + "query": "\nUPDATE users\nSET\n can_access_studio = ?,\n version = version + 1\n\nWHERE users.token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 2 + }, + "nullable": [] + }, + "hash": "19c1458e12d614bfface303102898b6a34f6b7a56eb906f5df3568172d51778b" +} diff --git a/crates/lib/mysql_queries/src/queries/users/user/mod.rs b/crates/lib/mysql_queries/src/queries/users/user/mod.rs index 2e2fdf6a73..e4a5293707 100644 --- a/crates/lib/mysql_queries/src/queries/users/user/mod.rs +++ b/crates/lib/mysql_queries/src/queries/users/user/mod.rs @@ -4,6 +4,7 @@ pub mod list_users_query_builder; pub mod lookup_user_for_login_by_email; pub mod lookup_user_for_login_by_username; pub mod lookup_user_for_login_result; +pub mod set_can_access_studio_transactional; pub mod set_user_ban_status; pub mod set_user_feature_flags; pub mod set_user_feature_flags_transactional; diff --git a/crates/lib/mysql_queries/src/queries/users/user/set_can_access_studio_transactional.rs b/crates/lib/mysql_queries/src/queries/users/user/set_can_access_studio_transactional.rs new file mode 100644 index 0000000000..4b33315fa0 --- /dev/null +++ b/crates/lib/mysql_queries/src/queries/users/user/set_can_access_studio_transactional.rs @@ -0,0 +1,38 @@ +use anyhow::anyhow; +use sqlx::{MySql, MySqlPool, Transaction}; + +use errors::AnyhowResult; +use tokens::tokens::users::UserToken; + +pub struct SetCanAccessStudioArgs<'a, 'b> { + // The action's target user token. + pub subject_user_token: &'a UserToken, + + pub can_access_studio: bool, + + pub transaction: &'a mut Transaction<'b, MySql>, +} + +pub async fn set_can_access_studio_transactional(args: SetCanAccessStudioArgs<'_, '_>) -> AnyhowResult<()> { + let query_result = sqlx::query!( + r#" +UPDATE users +SET + can_access_studio = ?, + version = version + 1 + +WHERE users.token = ? +LIMIT 1 + "#, + args.can_access_studio, + args.subject_user_token, + ) + .execute(&mut **args.transaction) + .await; + + if let Err(err) = query_result { + return Err(anyhow!("error with query: {:?}", err)); + } + + Ok(()) +} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/beta_keys/redeem_beta_key_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/beta_keys/redeem_beta_key_handler.rs index 3fd5a51e30..4a766e0c15 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/beta_keys/redeem_beta_key_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/beta_keys/redeem_beta_key_handler.rs @@ -7,6 +7,7 @@ use actix_web::http::StatusCode; use log::warn; use r2d2_redis::redis::transaction; use utoipa::ToSchema; +use enums::by_table::beta_keys::beta_key_product::BetaKeyProduct; use enums::by_table::comments::comment_entity_type::CommentEntityType; use enums::by_table::users::user_feature_flag::UserFeatureFlag; @@ -15,7 +16,9 @@ use mysql_queries::queries::beta_keys::get_beta_key_by_value::get_beta_key_by_va use mysql_queries::queries::beta_keys::redeem_beta_key::redeem_beta_key; use mysql_queries::queries::comments::comment_entity_token::CommentEntityToken; use mysql_queries::queries::comments::insert_comment::{insert_comment, InsertCommentArgs}; +use mysql_queries::queries::users::user::set_can_access_studio_transactional::{set_can_access_studio_transactional, SetCanAccessStudioArgs}; use mysql_queries::queries::users::user::set_user_feature_flags::{set_user_feature_flags, SetUserFeatureFlagArgs}; +use mysql_queries::queries::users::user_sessions::get_user_session_by_token::SessionUserRecord; use tokens::tokens::comments::CommentToken; use tokens::tokens::media_files::MediaFileToken; use tokens::tokens::model_weights::ModelWeightToken; @@ -138,6 +141,37 @@ pub async fn redeem_beta_key_handler( return Err(RedeemBetaKeyError::BadInput("beta key already redeemed".to_string())); } + let ip_address = get_request_ip(&http_request); + + match beta_key.product { + BetaKeyProduct::Studio => { + enroll_in_studio(&request, &server_state, &user_session, &ip_address) + .await + .map_err(|err| { + warn!("Error enrolling in studio: {:?}", &err); + RedeemBetaKeyError::ServerError + })?; + } + } + + let response = RedeemBetaKeySuccessResponse { + success: true, + }; + + let body = serde_json::to_string(&response) + .map_err(|_e| RedeemBetaKeyError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} + +async fn enroll_in_studio( + request: &RedeemBetaKeyRequest, + server_state: &ServerState, + user_session: &SessionUserRecord, + ip_address: &str, +) -> Result<(), RedeemBetaKeyError> { let mut user_feature_flags = UserSessionFeatureFlags::new(user_session.maybe_feature_flags.as_deref()); @@ -146,8 +180,6 @@ pub async fn redeem_beta_key_handler( UserFeatureFlag::VideoStyleTransfer, ]); - let ip_address = get_request_ip(&http_request); - let mut transaction = server_state.mysql_pool.begin() .await .map_err(|e| { @@ -167,6 +199,17 @@ pub async fn redeem_beta_key_handler( RedeemBetaKeyError::ServerError })?; + // NB: This isn't a necessary field, but can be useful for analytics. + set_can_access_studio_transactional(SetCanAccessStudioArgs { + subject_user_token: &user_session.user_token, + can_access_studio: true, + transaction: &mut transaction, + }).await + .map_err(|e| { + warn!("Could not set can_access_studio: {:?}", e); + RedeemBetaKeyError::ServerError + })?; + redeem_beta_key(&request.beta_key, &user_session.user_token, &mut transaction) .await .map_err(|e| { @@ -181,14 +224,5 @@ pub async fn redeem_beta_key_handler( RedeemBetaKeyError::ServerError })?; - let response = RedeemBetaKeySuccessResponse { - success: true, - }; - - let body = serde_json::to_string(&response) - .map_err(|_e| RedeemBetaKeyError::ServerError)?; - - Ok(HttpResponse::Ok() - .content_type("application/json") - .body(body)) + Ok(()) }