collab: Don't transfer existing usage when upgrading to Zed Pro (#30045)
This PR makes it so we don't transfer existing usage over when upgrading from a trial to Zed Pro. Release Notes: - N/A
This commit is contained in:
Marshall Bowers
GitHub
parent
3cdf5ce947
commit
cec1d2584b
4 changed files with 1 additions and 252 deletions
|
|
@ -1023,29 +1023,6 @@ async fn handle_customer_subscription_event(
|
||||||
.get_billing_subscription_by_stripe_subscription_id(&subscription.id)
|
.get_billing_subscription_by_stripe_subscription_id(&subscription.id)
|
||||||
.await?
|
.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
|
app.db
|
||||||
.update_billing_subscription(
|
.update_billing_subscription(
|
||||||
existing_subscription.id,
|
existing_subscription.id,
|
||||||
|
|
|
||||||
|
|
@ -1,8 +1,6 @@
|
||||||
use chrono::Timelike;
|
|
||||||
use time::PrimitiveDateTime;
|
use time::PrimitiveDateTime;
|
||||||
|
|
||||||
use crate::db::billing_subscription::{StripeSubscriptionStatus, SubscriptionKind};
|
use crate::db::UserId;
|
||||||
use crate::db::{UserId, billing_subscription};
|
|
||||||
|
|
||||||
use super::*;
|
use super::*;
|
||||||
|
|
||||||
|
|
@ -26,62 +24,6 @@ pub fn convert_chrono_to_time(datetime: DateTimeUtc) -> anyhow::Result<Primitive
|
||||||
}
|
}
|
||||||
|
|
||||||
impl LlmDatabase {
|
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(
|
pub async fn get_subscription_usage_for_period(
|
||||||
&self,
|
&self,
|
||||||
user_id: UserId,
|
user_id: UserId,
|
||||||
|
|
@ -114,54 +56,4 @@ impl LlmDatabase {
|
||||||
.one(tx)
|
.one(tx)
|
||||||
.await?)
|
.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,5 +1,4 @@
|
||||||
mod provider_tests;
|
mod provider_tests;
|
||||||
mod subscription_usage_tests;
|
|
||||||
|
|
||||||
use gpui::BackgroundExecutor;
|
use gpui::BackgroundExecutor;
|
||||||
use parking_lot::Mutex;
|
use parking_lot::Mutex;
|
||||||
|
|
|
||||||
|
|
@ -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"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue