diff --git a/crates/collab/src/api/billing.rs b/crates/collab/src/api/billing.rs index ce6b8dfaa8b2b2f2d1ddfbfdedceca6d571ed94e..39d124dd964c07bb08f2bc1386f2945815e10ad1 100644 --- a/crates/collab/src/api/billing.rs +++ b/crates/collab/src/api/billing.rs @@ -1023,29 +1023,6 @@ async fn handle_customer_subscription_event( .get_billing_subscription_by_stripe_subscription_id(&subscription.id) .await? { - let llm_db = app - .llm_db - .clone() - .ok_or_else(|| anyhow!("LLM DB not initialized"))?; - - let new_period_start_at = - chrono::DateTime::from_timestamp(subscription.current_period_start, 0) - .ok_or_else(|| anyhow!("No subscription period start"))?; - let new_period_end_at = - chrono::DateTime::from_timestamp(subscription.current_period_end, 0) - .ok_or_else(|| anyhow!("No subscription period end"))?; - - llm_db - .transfer_existing_subscription_usage( - billing_customer.user_id, - &existing_subscription, - subscription_kind, - subscription.status.into(), - new_period_start_at, - new_period_end_at, - ) - .await?; - app.db .update_billing_subscription( existing_subscription.id, diff --git a/crates/collab/src/llm/db/queries/subscription_usages.rs b/crates/collab/src/llm/db/queries/subscription_usages.rs index c772e8c5341a5dbc100d64bdd4d2c99cf93c775b..ee1ebf59b8cfc6bf3f4b74a81061389a160ccb4e 100644 --- a/crates/collab/src/llm/db/queries/subscription_usages.rs +++ b/crates/collab/src/llm/db/queries/subscription_usages.rs @@ -1,8 +1,6 @@ -use chrono::Timelike; use time::PrimitiveDateTime; -use crate::db::billing_subscription::{StripeSubscriptionStatus, SubscriptionKind}; -use crate::db::{UserId, billing_subscription}; +use crate::db::UserId; use super::*; @@ -26,62 +24,6 @@ pub fn convert_chrono_to_time(datetime: DateTimeUtc) -> anyhow::Result Result { - self.transaction(|tx| async move { - self.create_subscription_usage_in_tx( - user_id, - period_start_at, - period_end_at, - plan, - model_requests, - edit_predictions, - &tx, - ) - .await - }) - .await - } - - async fn create_subscription_usage_in_tx( - &self, - user_id: UserId, - period_start_at: DateTimeUtc, - period_end_at: DateTimeUtc, - plan: SubscriptionKind, - model_requests: i32, - edit_predictions: i32, - tx: &DatabaseTransaction, - ) -> Result { - // Clear out the nanoseconds so that these timestamps are comparable with Unix timestamps. - let period_start_at = period_start_at.with_nanosecond(0).unwrap(); - let period_end_at = period_end_at.with_nanosecond(0).unwrap(); - - let period_start_at = convert_chrono_to_time(period_start_at)?; - let period_end_at = convert_chrono_to_time(period_end_at)?; - - Ok( - subscription_usage::Entity::insert(subscription_usage::ActiveModel { - id: ActiveValue::set(Uuid::now_v7()), - user_id: ActiveValue::set(user_id), - period_start_at: ActiveValue::set(period_start_at), - period_end_at: ActiveValue::set(period_end_at), - plan: ActiveValue::set(plan), - model_requests: ActiveValue::set(model_requests), - edit_predictions: ActiveValue::set(edit_predictions), - }) - .exec_with_returning(tx) - .await?, - ) - } - pub async fn get_subscription_usage_for_period( &self, user_id: UserId, @@ -114,54 +56,4 @@ impl LlmDatabase { .one(tx) .await?) } - - pub async fn transfer_existing_subscription_usage( - &self, - user_id: UserId, - existing_subscription: &billing_subscription::Model, - new_subscription_kind: Option, - new_subscription_status: StripeSubscriptionStatus, - new_period_start_at: DateTimeUtc, - new_period_end_at: DateTimeUtc, - ) -> Result> { - self.transaction(|tx| async move { - match (existing_subscription.kind, new_subscription_status) { - (Some(SubscriptionKind::ZedProTrial), StripeSubscriptionStatus::Active) => { - let trial_period_start_at = existing_subscription - .current_period_start_at() - .ok_or_else(|| anyhow!("No trial subscription period start"))?; - let trial_period_end_at = existing_subscription - .current_period_end_at() - .ok_or_else(|| anyhow!("No trial subscription period end"))?; - - let existing_usage = self - .get_subscription_usage_for_period_in_tx( - user_id, - trial_period_start_at, - trial_period_end_at, - &tx, - ) - .await?; - if let Some(existing_usage) = existing_usage { - return Ok(Some( - self.create_subscription_usage_in_tx( - user_id, - new_period_start_at, - new_period_end_at, - new_subscription_kind.unwrap_or(existing_usage.plan), - existing_usage.model_requests, - existing_usage.edit_predictions, - &tx, - ) - .await?, - )); - } - } - _ => {} - } - - Ok(None) - }) - .await - } } diff --git a/crates/collab/src/llm/db/tests.rs b/crates/collab/src/llm/db/tests.rs index f77aedf3c8fcada0f3167cdf0246e8abe5cd4ccd..43a1b8b0d457817d1e94d72d0cad094011424c83 100644 --- a/crates/collab/src/llm/db/tests.rs +++ b/crates/collab/src/llm/db/tests.rs @@ -1,5 +1,4 @@ mod provider_tests; -mod subscription_usage_tests; use gpui::BackgroundExecutor; use parking_lot::Mutex; diff --git a/crates/collab/src/llm/db/tests/subscription_usage_tests.rs b/crates/collab/src/llm/db/tests/subscription_usage_tests.rs deleted file mode 100644 index 5ae49515bef6640f9f4b46e4e85227f9f7d68894..0000000000000000000000000000000000000000 --- a/crates/collab/src/llm/db/tests/subscription_usage_tests.rs +++ /dev/null @@ -1,119 +0,0 @@ -use chrono::{Duration, Utc}; -use pretty_assertions::assert_eq; - -use crate::db::billing_subscription::{StripeSubscriptionStatus, SubscriptionKind}; -use crate::db::{UserId, billing_subscription}; -use crate::llm::db::LlmDatabase; -use crate::test_llm_db; - -test_llm_db!( - test_transfer_existing_subscription_usage, - test_transfer_existing_subscription_usage_postgres -); - -async fn test_transfer_existing_subscription_usage(db: &mut LlmDatabase) { - // Test when an existing Zed Pro trial subscription is upgraded to Zed Pro. - { - let user_id = UserId(1); - - let now = Utc::now(); - - let trial_period_start_at = now - Duration::days(14); - let trial_period_end_at = now; - - let new_period_start_at = now; - let new_period_end_at = now + Duration::days(30); - - let existing_subscription = billing_subscription::Model { - kind: Some(SubscriptionKind::ZedProTrial), - stripe_current_period_start: Some(trial_period_start_at.timestamp()), - stripe_current_period_end: Some(trial_period_end_at.timestamp()), - ..Default::default() - }; - - let existing_usage = db - .create_subscription_usage( - user_id, - trial_period_start_at, - trial_period_end_at, - SubscriptionKind::ZedProTrial, - 25, - 1_000, - ) - .await - .unwrap(); - - let transferred_usage = db - .transfer_existing_subscription_usage( - user_id, - &existing_subscription, - Some(SubscriptionKind::ZedPro), - StripeSubscriptionStatus::Active, - new_period_start_at, - new_period_end_at, - ) - .await - .unwrap(); - - assert!( - transferred_usage.is_some(), - "subscription usage not transferred successfully" - ); - let transferred_usage = transferred_usage.unwrap(); - - assert_eq!( - transferred_usage.model_requests, - existing_usage.model_requests - ); - assert_eq!( - transferred_usage.edit_predictions, - existing_usage.edit_predictions - ); - } - - // Test when an existing Zed Pro trial subscription is canceled. - { - let user_id = UserId(2); - - let now = Utc::now(); - - let trial_period_start_at = now - Duration::days(14); - let trial_period_end_at = now; - - let existing_subscription = billing_subscription::Model { - kind: Some(SubscriptionKind::ZedProTrial), - stripe_current_period_start: Some(trial_period_start_at.timestamp()), - stripe_current_period_end: Some(trial_period_end_at.timestamp()), - ..Default::default() - }; - - let _existing_usage = db - .create_subscription_usage( - user_id, - trial_period_start_at, - trial_period_end_at, - SubscriptionKind::ZedProTrial, - 25, - 1_000, - ) - .await - .unwrap(); - - let transferred_usage = db - .transfer_existing_subscription_usage( - user_id, - &existing_subscription, - Some(SubscriptionKind::ZedPro), - StripeSubscriptionStatus::Canceled, - trial_period_start_at, - trial_period_end_at, - ) - .await - .unwrap(); - - assert!( - transferred_usage.is_none(), - "subscription usage was transferred when it should not have been" - ); - } -}