Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

collab: Remove unused billing queries (#35679)

Browse source 
This PR removes some billing-related queries that are no longer used.

Release Notes:

- N/A
This commit is contained in:
Marshall Bowers 2025-08-05 18:42:45 -04:00 • committed by GitHub
parent f10ffc2a72
commit 142efbac0d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 0 additions and 317 deletions
Download patch file Download diff file Expand all files Collapse all files

126
crates/collab/src/db/queries/billing_subscriptions.rs
View file

@ -85,19 +85,6 @@ impl Database {
.await
}
/// Returns the billing subscription with the specified ID.
pub async fn get_billing_subscription_by_id(
&self,
id: BillingSubscriptionId,
) -> Result<Option<billing_subscription::Model>> {
self.transaction(|tx| async move {
Ok(billing_subscription::Entity::find_by_id(id)
.one(&*tx)
.await?)
})
.await
}
/// Returns the billing subscription with the specified Stripe subscription ID.
pub async fn get_billing_subscription_by_stripe_subscription_id(
&self,
@ -143,119 +130,6 @@ impl Database {
.await
}
/// Returns all of the billing subscriptions for the user with the specified ID.
///
/// Note that this returns the subscriptions regardless of their status.
/// If you're wanting to check if a use has an active billing subscription,
/// use `get_active_billing_subscriptions` instead.
pub async fn get_billing_subscriptions(
&self,
user_id: UserId,
) -> Result<Vec<billing_subscription::Model>> {
self.transaction(|tx| async move {
let subscriptions = billing_subscription::Entity::find()
.inner_join(billing_customer::Entity)
.filter(billing_customer::Column::UserId.eq(user_id))
.order_by_asc(billing_subscription::Column::Id)
.all(&*tx)
.await?;
Ok(subscriptions)
})
.await
}
pub async fn get_active_billing_subscriptions(
&self,
user_ids: HashSet<UserId>,
) -> Result<HashMap<UserId, (billing_customer::Model, billing_subscription::Model)>> {
self.transaction(|tx| {
let user_ids = user_ids.clone();
async move {
let mut rows = billing_subscription::Entity::find()
.inner_join(billing_customer::Entity)
.select_also(billing_customer::Entity)
.filter(billing_customer::Column::UserId.is_in(user_ids))
.filter(
billing_subscription::Column::StripeSubscriptionStatus
.eq(StripeSubscriptionStatus::Active),
)
.filter(billing_subscription::Column::Kind.is_null())
.order_by_asc(billing_subscription::Column::Id)
.stream(&*tx)
.await?;
let mut subscriptions = HashMap::default();
while let Some(row) = rows.next().await {
if let (subscription, Some(customer)) = row? {
subscriptions.insert(customer.user_id, (customer, subscription));
}
}
Ok(subscriptions)
}
})
.await
}
pub async fn get_active_zed_pro_billing_subscriptions(
&self,
) -> Result<HashMap<UserId, (billing_customer::Model, billing_subscription::Model)>> {
self.transaction(|tx| async move {
let mut rows = billing_subscription::Entity::find()
.inner_join(billing_customer::Entity)
.select_also(billing_customer::Entity)
.filter(
billing_subscription::Column::StripeSubscriptionStatus
.eq(StripeSubscriptionStatus::Active),
)
.filter(billing_subscription::Column::Kind.eq(SubscriptionKind::ZedPro))
.order_by_asc(billing_subscription::Column::Id)
.stream(&*tx)
.await?;
let mut subscriptions = HashMap::default();
while let Some(row) = rows.next().await {
if let (subscription, Some(customer)) = row? {
subscriptions.insert(customer.user_id, (customer, subscription));
}
}
Ok(subscriptions)
})
.await
}
pub async fn get_active_zed_pro_billing_subscriptions_for_users(
&self,
user_ids: HashSet<UserId>,
) -> Result<HashMap<UserId, (billing_customer::Model, billing_subscription::Model)>> {
self.transaction(|tx| {
let user_ids = user_ids.clone();
async move {
let mut rows = billing_subscription::Entity::find()
.inner_join(billing_customer::Entity)
.select_also(billing_customer::Entity)
.filter(billing_customer::Column::UserId.is_in(user_ids))
.filter(
billing_subscription::Column::StripeSubscriptionStatus
.eq(StripeSubscriptionStatus::Active),
)
.filter(billing_subscription::Column::Kind.eq(SubscriptionKind::ZedPro))
.order_by_asc(billing_subscription::Column::Id)
.stream(&*tx)
.await?;
let mut subscriptions = HashMap::default();
while let Some(row) = rows.next().await {
if let (subscription, Some(customer)) = row? {
subscriptions.insert(customer.user_id, (customer, subscription));
}
}
Ok(subscriptions)
}
})
.await
}
/// Returns whether the user has an active billing subscription.
pub async fn has_active_billing_subscription(&self, user_id: UserId) -> Result<bool> {
Ok(self.count_active_billing_subscriptions(user_id).await? > 0)

1
crates/collab/src/db/tests.rs
View file

@ -1,4 +1,3 @@
mod billing_subscription_tests;
mod buffer_tests;
mod channel_tests;
mod contributor_tests;

96
crates/collab/src/db/tests/billing_subscription_tests.rs
View file

@ -1,96 +0,0 @@
use std::sync::Arc;
use crate::db::billing_subscription::StripeSubscriptionStatus;
use crate::db::tests::new_test_user;
use crate::db::{CreateBillingCustomerParams, CreateBillingSubscriptionParams};
use crate::test_both_dbs;
use super::Database;
test_both_dbs!(
test_get_active_billing_subscriptions,
test_get_active_billing_subscriptions_postgres,
test_get_active_billing_subscriptions_sqlite
);
async fn test_get_active_billing_subscriptions(db: &Arc<Database>) {
// A user with no subscription has no active billing subscriptions.
{
let user_id = new_test_user(db, "no-subscription-user@example.com").await;
let subscription_count = db
.count_active_billing_subscriptions(user_id)
.await
.unwrap();
assert_eq!(subscription_count, 0);
}
// A user with an active subscription has one active billing subscription.
{
let user_id = new_test_user(db, "active-user@example.com").await;
let customer = db
.create_billing_customer(&CreateBillingCustomerParams {
user_id,
stripe_customer_id: "cus_active_user".into(),
})
.await
.unwrap();
assert_eq!(customer.stripe_customer_id, "cus_active_user".to_string());
db.create_billing_subscription(&CreateBillingSubscriptionParams {
billing_customer_id: customer.id,
kind: None,
stripe_subscription_id: "sub_active_user".into(),
stripe_subscription_status: StripeSubscriptionStatus::Active,
stripe_cancellation_reason: None,
stripe_current_period_start: None,
stripe_current_period_end: None,
})
.await
.unwrap();
let subscriptions = db.get_billing_subscriptions(user_id).await.unwrap();
assert_eq!(subscriptions.len(), 1);
let subscription = &subscriptions[0];
assert_eq!(
subscription.stripe_subscription_id,
"sub_active_user".to_string()
);
assert_eq!(
subscription.stripe_subscription_status,
StripeSubscriptionStatus::Active
);
}
// A user with a past-due subscription has no active billing subscriptions.
{
let user_id = new_test_user(db, "past-due-user@example.com").await;
let customer = db
.create_billing_customer(&CreateBillingCustomerParams {
user_id,
stripe_customer_id: "cus_past_due_user".into(),
})
.await
.unwrap();
assert_eq!(customer.stripe_customer_id, "cus_past_due_user".to_string());
db.create_billing_subscription(&CreateBillingSubscriptionParams {
billing_customer_id: customer.id,
kind: None,
stripe_subscription_id: "sub_past_due_user".into(),
stripe_subscription_status: StripeSubscriptionStatus::PastDue,
stripe_cancellation_reason: None,
stripe_current_period_start: None,
stripe_current_period_end: None,
})
.await
.unwrap();
let subscription_count = db
.count_active_billing_subscriptions(user_id)
.await
.unwrap();
assert_eq!(subscription_count, 0);
}
}

1
crates/collab/src/llm/db/queries.rs
View file

@ -1,6 +1,5 @@
use super::*;
pub mod providers;
pub mod subscription_usage_meters;
pub mod subscription_usages;
pub mod usages;

72
crates/collab/src/llm/db/queries/subscription_usage_meters.rs
View file

@ -1,72 +0,0 @@
use crate::db::UserId;
use crate::llm::db::queries::subscription_usages::convert_chrono_to_time;
use super::*;
impl LlmDatabase {
/// Returns all current subscription usage meters as of the given timestamp.
pub async fn get_current_subscription_usage_meters(
&self,
now: DateTimeUtc,
) -> Result<Vec<(subscription_usage_meter::Model, subscription_usage::Model)>> {
let now = convert_chrono_to_time(now)?;
self.transaction(|tx| async move {
let result = subscription_usage_meter::Entity::find()
.inner_join(subscription_usage::Entity)
.filter(
subscription_usage::Column::PeriodStartAt
.lte(now)
.and(subscription_usage::Column::PeriodEndAt.gte(now)),
)
.select_also(subscription_usage::Entity)
.all(&*tx)
.await?;
let result = result
.into_iter()
.filter_map(|(meter, usage)| {
let usage = usage?;
Some((meter, usage))
})
.collect();
Ok(result)
})
.await
}
/// Returns all current subscription usage meters for the given user as of the given timestamp.
pub async fn get_current_subscription_usage_meters_for_user(
&self,
user_id: UserId,
now: DateTimeUtc,
) -> Result<Vec<(subscription_usage_meter::Model, subscription_usage::Model)>> {
let now = convert_chrono_to_time(now)?;
self.transaction(|tx| async move {
let result = subscription_usage_meter::Entity::find()
.inner_join(subscription_usage::Entity)
.filter(subscription_usage::Column::UserId.eq(user_id))
.filter(
subscription_usage::Column::PeriodStartAt
.lte(now)
.and(subscription_usage::Column::PeriodEndAt.gte(now)),
)
.select_also(subscription_usage::Entity)
.all(&*tx)
.await?;
let result = result
.into_iter()
.filter_map(|(meter, usage)| {
let usage = usage?;
Some((meter, usage))
})
.collect();
Ok(result)
})
.await
}
}

21
crates/collab/src/llm/db/queries/subscription_usages.rs
View file

@ -1,28 +1,7 @@
use time::PrimitiveDateTime;
use crate::db::UserId;
use super::*;
pub fn convert_chrono_to_time(datetime: DateTimeUtc) -> anyhow::Result<PrimitiveDateTime> {
use chrono::{Datelike as _, Timelike as _};
let date = time::Date::from_calendar_date(
datetime.year(),
time::Month::try_from(datetime.month() as u8).unwrap(),
datetime.day() as u8,
)?;
let time = time::Time::from_hms_nano(
datetime.hour() as u8,
datetime.minute() as u8,
datetime.second() as u8,
datetime.nanosecond(),
)?;
Ok(PrimitiveDateTime::new(date, time))
}
impl LlmDatabase {
pub async fn get_subscription_usage_for_period(
&self,