From e5e93364aebc07754b97d9f176d2d50f96b3caa0 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sun, 21 Jul 2024 09:25:04 -0400 Subject: [PATCH] usage tracker wip --- Cargo.lock | 2 + crates/_types/enums/src/no_table/mod.rs | 1 + .../enums/src/no_table/premium_product/mod.rs | 1 + .../premium_product/premium_product_name.rs | 134 ++++++++++++++++++ crates/lib/redis_schema/Cargo.toml | 4 + .../lib/redis_schema/src/keys/premium/mod.rs | 2 +- .../src/keys/premium/premium_avt_redis_key.rs | 7 +- .../premium/premium_ip_address_redis_key.rs | 7 +- ...redis_key.rs => premium_user_redis_key.rs} | 15 +- crates/lib/redis_schema/src/payloads/mod.rs | 1 + .../redis_schema/src/payloads/premium/mod.rs | 5 + .../payloads/premium/premium_user_payload.rs | 76 ++++++++++ .../payloads/premium/product_by_week_store.rs | 57 ++++++++ .../premium/product_by_week_subkey.rs | 67 +++++++++ 14 files changed, 365 insertions(+), 14 deletions(-) create mode 100644 crates/_types/enums/src/no_table/premium_product/mod.rs create mode 100644 crates/_types/enums/src/no_table/premium_product/premium_product_name.rs rename crates/lib/redis_schema/src/keys/premium/{premium_user_token_redis_key.rs => premium_user_redis_key.rs} (60%) create mode 100644 crates/lib/redis_schema/src/payloads/premium/mod.rs create mode 100644 crates/lib/redis_schema/src/payloads/premium/premium_user_payload.rs create mode 100644 crates/lib/redis_schema/src/payloads/premium/product_by_week_store.rs create mode 100644 crates/lib/redis_schema/src/payloads/premium/product_by_week_subkey.rs diff --git a/Cargo.lock b/Cargo.lock index 683c13ad67..14f16158e2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5563,8 +5563,10 @@ name = "redis_schema" version = "0.0.1" dependencies = [ "chrono", + "enums", "errors", "my-workspace-hack", + "r2d2_redis", "tokens", ] diff --git a/crates/_types/enums/src/no_table/mod.rs b/crates/_types/enums/src/no_table/mod.rs index 51ac556c25..7bb3fd4399 100644 --- a/crates/_types/enums/src/no_table/mod.rs +++ b/crates/_types/enums/src/no_table/mod.rs @@ -1 +1,2 @@ +pub mod premium_product; pub mod style_transfer; \ No newline at end of file diff --git a/crates/_types/enums/src/no_table/premium_product/mod.rs b/crates/_types/enums/src/no_table/premium_product/mod.rs new file mode 100644 index 0000000000..c6890ee97f --- /dev/null +++ b/crates/_types/enums/src/no_table/premium_product/mod.rs @@ -0,0 +1 @@ +pub mod premium_product_name; \ No newline at end of file diff --git a/crates/_types/enums/src/no_table/premium_product/premium_product_name.rs b/crates/_types/enums/src/no_table/premium_product/premium_product_name.rs new file mode 100644 index 0000000000..ca3d2c463f --- /dev/null +++ b/crates/_types/enums/src/no_table/premium_product/premium_product_name.rs @@ -0,0 +1,134 @@ +use std::collections::BTreeSet; + +#[cfg(test)] +use strum::EnumCount; +#[cfg(test)] +use strum::EnumIter; +use utoipa::ToSchema; + +/// This enum is not backed by a particular database table. +/// This is used to count premium product uses for free and paid users. +/// +/// DO NOT CHANGE VALUES WITHOUT A MIGRATION STRATEGY. +#[cfg_attr(test, derive(EnumIter, EnumCount))] +#[derive(Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Hash, Deserialize, Serialize, ToSchema)] +#[serde(rename_all = "snake_case")] +pub enum PremiumProductName { + // NB: These keys are kept short to preserve space + #[serde(rename = "fa")] + FaceAnimator, + // NB: These keys are kept short to preserve space + #[serde(rename = "fm")] + FaceMirror, + // NB: These keys are kept short to preserve space + #[serde(rename = "lip")] + Lipsync, + // NB: These keys are kept short to preserve space + #[serde(rename = "vst")] + VideoStyleTransfer, +} + +// TODO(bt, 2022-12-21): This desperately needs MySQL integration tests! +impl_enum_display_and_debug_using_to_str!(PremiumProductName); +//impl_mysql_enum_coders!(PremiumProductName); +//impl_mysql_from_row!(PremiumProductName); + +/// NB: Legacy API for older code. +impl PremiumProductName { + pub fn to_str(&self) -> &'static str { + match self { + Self::FaceAnimator => "fa", + Self::FaceMirror => "fm", + Self::Lipsync => "lip", + Self::VideoStyleTransfer => "vst", + } + } + + pub fn from_str(value: &str) -> Result { + match value { + "fa" => Ok(Self::FaceAnimator), + "fm" => Ok(Self::FaceMirror), + "lip" => Ok(Self::Lipsync), + "vst" => Ok(Self::VideoStyleTransfer), + _ => Err(format!("Unknown PremiumProductName: {}", 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::FaceAnimator, + Self::FaceMirror, + Self::Lipsync, + Self::VideoStyleTransfer, + ]) + } +} + +#[cfg(test)] +mod tests { + use crate::no_table::premium_product::premium_product_name::PremiumProductName; + use crate::test_helpers::assert_serialization; + + mod serde { + use super::*; + + #[test] + fn test_serialization() { + assert_serialization(PremiumProductName::FaceAnimator, "fa"); + assert_serialization(PremiumProductName::FaceMirror, "fm"); + assert_serialization(PremiumProductName::Lipsync, "lip"); + assert_serialization(PremiumProductName::VideoStyleTransfer, "vst"); + } + + mod impl_methods { + use super::*; + + #[test] + fn to_str() { + assert_eq!(PremiumProductName::FaceAnimator.to_str(), "fa"); + assert_eq!(PremiumProductName::FaceMirror.to_str(), "fm"); + assert_eq!(PremiumProductName::Lipsync.to_str(), "lip"); + assert_eq!(PremiumProductName::VideoStyleTransfer.to_str(), "vst"); + } + + #[test] + fn from_str() { + assert_eq!(PremiumProductName::from_str("fa").unwrap(), PremiumProductName::FaceAnimator); + assert_eq!(PremiumProductName::from_str("fm").unwrap(), PremiumProductName::FaceMirror); + assert_eq!(PremiumProductName::from_str("lip").unwrap(), PremiumProductName::Lipsync); + assert_eq!(PremiumProductName::from_str("vst").unwrap(), PremiumProductName::VideoStyleTransfer); + } + } + + mod manual_variant_checks { + use super::*; + + #[test] + fn all_variants() { + let variants = PremiumProductName::all_variants(); + assert_eq!(variants.len(), 4); + } + } + + mod mechanical_checks { + use super::*; + + #[test] + fn variant_length() { + use strum::IntoEnumIterator; + assert_eq!(PremiumProductName::all_variants().len(), PremiumProductName::iter().len()); + } + + #[test] + fn round_trip() { + for variant in PremiumProductName::all_variants() { + assert_eq!(variant, PremiumProductName::from_str(variant.to_str()).unwrap()); + assert_eq!(variant, PremiumProductName::from_str(&format!("{}", variant)).unwrap()); + assert_eq!(variant, PremiumProductName::from_str(&format!("{:?}", variant)).unwrap()); + } + } + } + } +} diff --git a/crates/lib/redis_schema/Cargo.toml b/crates/lib/redis_schema/Cargo.toml index 7a57d4c195..2a620619be 100644 --- a/crates/lib/redis_schema/Cargo.toml +++ b/crates/lib/redis_schema/Cargo.toml @@ -17,6 +17,7 @@ normal = ["my-workspace-hack"] [dependencies] # Internal +enums = { path = "../../_types/enums" } errors = { path = "../errors" } tokens = { path = "../../_types/tokens" } @@ -26,5 +27,8 @@ my-workspace-hack = { version = "0.1", path = "../../my-workspace-hack" } # External chrono = { version = "0.4.22", features = ["serde"] } +# Redis +r2d2_redis = "0.14.0" + [dev-dependencies] # None yet diff --git a/crates/lib/redis_schema/src/keys/premium/mod.rs b/crates/lib/redis_schema/src/keys/premium/mod.rs index 8f3f3920b5..faf8d9ee7a 100644 --- a/crates/lib/redis_schema/src/keys/premium/mod.rs +++ b/crates/lib/redis_schema/src/keys/premium/mod.rs @@ -1,3 +1,3 @@ pub mod premium_avt_redis_key; pub mod premium_ip_address_redis_key; -pub mod premium_user_token_redis_key; \ No newline at end of file +pub mod premium_user_redis_key; \ No newline at end of file diff --git a/crates/lib/redis_schema/src/keys/premium/premium_avt_redis_key.rs b/crates/lib/redis_schema/src/keys/premium/premium_avt_redis_key.rs index 0eb4defbbd..b65f9cfdae 100644 --- a/crates/lib/redis_schema/src/keys/premium/premium_avt_redis_key.rs +++ b/crates/lib/redis_schema/src/keys/premium/premium_avt_redis_key.rs @@ -5,7 +5,8 @@ pub struct PremiumAvtRedisKey(pub String); impl_string_key!(PremiumAvtRedisKey); -const DURATION : Duration = Duration::milliseconds(1000 * 60 * 60 * 24 * 32); // 32 days +// NB: 62 days to last over a month with enough time for debugging. +const REDIS_KEY_TTL_DURATION: Duration = Duration::milliseconds(1000 * 60 * 60 * 24 * 62); impl PremiumAvtRedisKey { pub fn new_for_user(avt: &AnonymousVisitorTrackingToken, time: DateTime) -> Self { @@ -15,7 +16,7 @@ impl PremiumAvtRedisKey { } pub fn get_redis_ttl() -> Duration { - DURATION + REDIS_KEY_TTL_DURATION } } @@ -35,6 +36,6 @@ mod tests { #[test] fn test_duration() { - assert_eq!(PremiumAvtRedisKey::get_redis_ttl().num_days(), 32); + assert_eq!(PremiumAvtRedisKey::get_redis_ttl().num_days(), 62); } } diff --git a/crates/lib/redis_schema/src/keys/premium/premium_ip_address_redis_key.rs b/crates/lib/redis_schema/src/keys/premium/premium_ip_address_redis_key.rs index bef49d3768..6136924114 100644 --- a/crates/lib/redis_schema/src/keys/premium/premium_ip_address_redis_key.rs +++ b/crates/lib/redis_schema/src/keys/premium/premium_ip_address_redis_key.rs @@ -4,7 +4,8 @@ pub struct PremiumIpAddressRedisKey(pub String); impl_string_key!(PremiumIpAddressRedisKey); -const DURATION : Duration = Duration::milliseconds(1000 * 60 * 60 * 24 * 32); // 32 days +// NB: 62 days to last over a month with enough time for debugging. +const REDIS_KEY_TTL_DURATION: Duration = Duration::milliseconds(1000 * 60 * 60 * 24 * 62); impl PremiumIpAddressRedisKey { pub fn new_for_user(ip_addr: &str, time: DateTime) -> Self { @@ -14,7 +15,7 @@ impl PremiumIpAddressRedisKey { } pub fn get_redis_ttl() -> Duration { - DURATION + REDIS_KEY_TTL_DURATION } } @@ -33,6 +34,6 @@ mod tests { #[test] fn test_duration() { - assert_eq!(PremiumIpAddressRedisKey::get_redis_ttl().num_days(), 32); + assert_eq!(PremiumIpAddressRedisKey::get_redis_ttl().num_days(), 62); } } diff --git a/crates/lib/redis_schema/src/keys/premium/premium_user_token_redis_key.rs b/crates/lib/redis_schema/src/keys/premium/premium_user_redis_key.rs similarity index 60% rename from crates/lib/redis_schema/src/keys/premium/premium_user_token_redis_key.rs rename to crates/lib/redis_schema/src/keys/premium/premium_user_redis_key.rs index e8618e03fc..406a198bab 100644 --- a/crates/lib/redis_schema/src/keys/premium/premium_user_token_redis_key.rs +++ b/crates/lib/redis_schema/src/keys/premium/premium_user_redis_key.rs @@ -2,13 +2,14 @@ use chrono::{Datelike, DateTime, Duration, Utc}; use tokens::tokens::users::UserToken; -pub struct PremiumUserTokenRedisKey(pub String); +pub struct PremiumUserRedisKey(pub String); -impl_string_key!(PremiumUserTokenRedisKey); +impl_string_key!(PremiumUserRedisKey); -const DURATION : Duration = Duration::milliseconds(1000 * 60 * 60 * 24 * 32); // 32 days +// NB: 62 days to last over a month with enough time for debugging. +const REDIS_KEY_TTL_DURATION: Duration = Duration::milliseconds(1000 * 60 * 60 * 24 * 62); -impl PremiumUserTokenRedisKey { +impl PremiumUserRedisKey { pub fn new_for_user(user_token: &UserToken, time: DateTime) -> Self { let month = time.month(); let key = format!("premium:user:{}:{}", user_token.as_str(), month); @@ -16,7 +17,7 @@ impl PremiumUserTokenRedisKey { } pub fn get_redis_ttl() -> Duration { - DURATION + REDIS_KEY_TTL_DURATION } } @@ -30,12 +31,12 @@ mod tests { fn test_new_for_user() { let user_token = UserToken::new_from_str("token"); let time = Utc.with_ymd_and_hms(2021, 1, 1, 0, 0, 0).unwrap(); - let key = PremiumUserTokenRedisKey::new_for_user(&user_token, time); + let key = PremiumUserRedisKey::new_for_user(&user_token, time); assert_eq!(key.as_str(), "premium:user:token:1"); } #[test] fn test_duration() { - assert_eq!(PremiumUserTokenRedisKey::get_redis_ttl().num_days(), 32); + assert_eq!(PremiumUserRedisKey::get_redis_ttl().num_days(), 62); } } diff --git a/crates/lib/redis_schema/src/payloads/mod.rs b/crates/lib/redis_schema/src/payloads/mod.rs index e69de29bb2..7be8597f74 100644 --- a/crates/lib/redis_schema/src/payloads/mod.rs +++ b/crates/lib/redis_schema/src/payloads/mod.rs @@ -0,0 +1 @@ +pub mod premium; \ No newline at end of file diff --git a/crates/lib/redis_schema/src/payloads/premium/mod.rs b/crates/lib/redis_schema/src/payloads/premium/mod.rs new file mode 100644 index 0000000000..289c9f1b64 --- /dev/null +++ b/crates/lib/redis_schema/src/payloads/premium/mod.rs @@ -0,0 +1,5 @@ +pub mod premium_user_payload; + +// Private +pub (crate) mod product_by_week_store; +pub (crate) mod product_by_week_subkey; diff --git a/crates/lib/redis_schema/src/payloads/premium/premium_user_payload.rs b/crates/lib/redis_schema/src/payloads/premium/premium_user_payload.rs new file mode 100644 index 0000000000..6313176cf5 --- /dev/null +++ b/crates/lib/redis_schema/src/payloads/premium/premium_user_payload.rs @@ -0,0 +1,76 @@ +use std::collections::HashMap; +use std::ops::Add; +use chrono::Utc; +use r2d2_redis::r2d2::PooledConnection; +use r2d2_redis::redis::{Commands}; +use r2d2_redis::RedisConnectionManager; +use errors::AnyhowResult; +use crate::keys::premium::premium_user_redis_key::PremiumUserRedisKey; + +pub struct PremiumUserPayload { + key: PremiumUserRedisKey, + + credits_used: Option, + +} + +impl PremiumUserPayload { + + pub fn read_from_redis(key: PremiumUserRedisKey, mut redis: PooledConnection) -> AnyhowResult { + let _values : HashMap = redis.hgetall(key.as_str())?; + + let mut free_uses_per_product_map = HashMap::new(); + + let keys: Vec = redis.hkeys(key.as_str())?; + + for key in keys { + if key == "credits" { + continue; + } + + let count: u64 = redis.hget(key.as_str(), "count")?; + + free_uses_per_product_map.insert(key, count); + } + + Ok(Self { + key, + credits_used: Some(0), // TODO + //free_uses_per_product_map, + }) + } + +// pub fn persist_to_redis(&self, mut redis: PooledConnection) -> AnyhowResult<()> { +// let mut values = Vec::with_capacity(self.free_uses_per_product_map.len() + 1); +// +// values.push(("credits", self.credits_used.unwrap_or(0).to_string())); +// +// for (product_by_week_key, count) in &self.free_uses_per_product_map { +// values.push((product_by_week_key, count.to_string())); +// } +// +// redis.hset_multiple(self.key.as_str(), &values)?; +// +// self.set_key_expiry(redis)?; +// +// Ok(()) +// } + + pub fn set_key_expiry(&self, mut redis: PooledConnection) -> AnyhowResult<()> { + let expire_at= Utc::now() + .add(PremiumUserRedisKey::get_redis_ttl()) + .timestamp() as usize; + + redis.expire_at(self.key.as_str(), expire_at)?; + + Ok(()) + } +} + +#[cfg(test)] +mod tests { + + #[test] + fn todo() { + } +} diff --git a/crates/lib/redis_schema/src/payloads/premium/product_by_week_store.rs b/crates/lib/redis_schema/src/payloads/premium/product_by_week_store.rs new file mode 100644 index 0000000000..d89a273e63 --- /dev/null +++ b/crates/lib/redis_schema/src/payloads/premium/product_by_week_store.rs @@ -0,0 +1,57 @@ +use std::collections::HashMap; +use enums::no_table::premium_product::premium_product_name::PremiumProductName; +use crate::payloads::premium::product_by_week_subkey::ProductByWeekSubkey; + +pub struct ProductByWeekStore { + + /// Key: "{product_id}:{iso_week_zero_index}" + /// Value: count of uses + free_uses_per_product_map: HashMap, +} + +impl ProductByWeekStore { + pub fn new() -> Self { + Self { + free_uses_per_product_map: HashMap::new(), + } + } + + pub fn increment_use(&mut self, name: PremiumProductName, week: u32) { + let key = ProductByWeekSubkey::new(name, week); + let count = self.free_uses_per_product_map.entry(key).or_insert(0); + *count += 1; + } + + pub fn get_use_count(&self, name: PremiumProductName, week: u32) -> u64 { + let key = ProductByWeekSubkey::new(name, week); + *self.free_uses_per_product_map.get(&key).unwrap_or(&0) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_increment_use() { + let mut store = ProductByWeekStore::new(); + + store.increment_use(PremiumProductName::FaceAnimator, 10); + store.increment_use(PremiumProductName::FaceAnimator, 10); + store.increment_use(PremiumProductName::FaceAnimator, 10); + store.increment_use(PremiumProductName::FaceMirror, 10); + store.increment_use(PremiumProductName::Lipsync, 11); + + // Week 10 + assert_eq!(store.get_use_count(PremiumProductName::FaceAnimator, 10), 3); + assert_eq!(store.get_use_count(PremiumProductName::FaceMirror, 10), 1); + assert_eq!(store.get_use_count(PremiumProductName::Lipsync, 10), 0); + assert_eq!(store.get_use_count(PremiumProductName::VideoStyleTransfer, 10), 0); + + // Week 11 + assert_eq!(store.get_use_count(PremiumProductName::FaceAnimator, 11), 0); + assert_eq!(store.get_use_count(PremiumProductName::FaceMirror, 11), 0); + assert_eq!(store.get_use_count(PremiumProductName::Lipsync, 11), 1); + assert_eq!(store.get_use_count(PremiumProductName::VideoStyleTransfer, 11), 0); + } +} \ No newline at end of file diff --git a/crates/lib/redis_schema/src/payloads/premium/product_by_week_subkey.rs b/crates/lib/redis_schema/src/payloads/premium/product_by_week_subkey.rs new file mode 100644 index 0000000000..8aed0a3fb7 --- /dev/null +++ b/crates/lib/redis_schema/src/payloads/premium/product_by_week_subkey.rs @@ -0,0 +1,67 @@ +use std::num::ParseIntError; +use chrono::{Datelike, Utc}; +use enums::no_table::premium_product::premium_product_name::PremiumProductName; +use errors::{anyhow, AnyhowResult}; + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct ProductByWeekSubkey { + pub product_name: PremiumProductName, + pub week: u32, +} + +impl ProductByWeekSubkey { + pub fn new(product_name: PremiumProductName, week: u32) -> Self { + Self { product_name, week } + } + + pub fn new_this_week(product_name: PremiumProductName) -> Self { + let time = Utc::now(); + let week = time.iso_week(); + let week = week.week0(); + Self { product_name, week } + } + + pub fn from_string(key: &str) -> AnyhowResult { + let mut parts = key.split(':'); + let product_name = parts.next() + .ok_or_else(|| anyhow!("No product name"))?; + let product_name = PremiumProductName::from_str(product_name) + .map_err(|err| anyhow!(err))?; + let week = parts.next() + .ok_or_else(|| anyhow!("No week"))? + .parse() + .map_err(|err: ParseIntError| anyhow!(err))?; + Ok(Self { + product_name, + week, + }) + } + + pub fn to_string(&self) -> String { + format!("{}:{}", self.product_name.to_str(), self.week) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_new_this_week() { + let key = ProductByWeekSubkey::new_this_week(PremiumProductName::FaceAnimator); + assert_eq!(key.product_name, PremiumProductName::FaceAnimator); + } + + #[test] + fn test_from_string() { + let key = ProductByWeekSubkey::from_string("fa:1").unwrap(); + assert_eq!(key.product_name, PremiumProductName::FaceAnimator); + assert_eq!(key.week, 1); + } + + #[test] + fn test_to_string() { + let key = ProductByWeekSubkey::new(PremiumProductName::FaceMirror, 2); + assert_eq!(key.to_string(), "fm:2"); + } +} \ No newline at end of file