@@ -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,
@@ -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<Primitive
}
impl LlmDatabase {
- pub async fn create_subscription_usage(
- &self,
- user_id: UserId,
- period_start_at: DateTimeUtc,
- period_end_at: DateTimeUtc,
- plan: SubscriptionKind,
- model_requests: i32,
- edit_predictions: i32,
- ) -> Result<subscription_usage::Model> {
- 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<subscription_usage::Model> {
- // 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<SubscriptionKind>,
- new_subscription_status: StripeSubscriptionStatus,
- new_period_start_at: DateTimeUtc,
- new_period_end_at: DateTimeUtc,
- ) -> Result<Option<subscription_usage::Model>> {
- 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
- }
}
@@ -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"
- );
- }
-}