1use chrono::Timelike;
2use time::PrimitiveDateTime;
3
4use crate::db::billing_subscription::{StripeSubscriptionStatus, SubscriptionKind};
5use crate::db::{UserId, billing_subscription};
6
7use super::*;
8
9pub fn convert_chrono_to_time(datetime: DateTimeUtc) -> anyhow::Result<PrimitiveDateTime> {
10 use chrono::{Datelike as _, Timelike as _};
11
12 let date = time::Date::from_calendar_date(
13 datetime.year(),
14 time::Month::try_from(datetime.month() as u8).unwrap(),
15 datetime.day() as u8,
16 )?;
17
18 let time = time::Time::from_hms_nano(
19 datetime.hour() as u8,
20 datetime.minute() as u8,
21 datetime.second() as u8,
22 datetime.nanosecond(),
23 )?;
24
25 Ok(PrimitiveDateTime::new(date, time))
26}
27
28impl LlmDatabase {
29 pub async fn create_subscription_usage(
30 &self,
31 user_id: UserId,
32 period_start_at: DateTimeUtc,
33 period_end_at: DateTimeUtc,
34 plan: SubscriptionKind,
35 model_requests: i32,
36 edit_predictions: i32,
37 ) -> Result<subscription_usage::Model> {
38 self.transaction(|tx| async move {
39 self.create_subscription_usage_in_tx(
40 user_id,
41 period_start_at,
42 period_end_at,
43 plan,
44 model_requests,
45 edit_predictions,
46 &tx,
47 )
48 .await
49 })
50 .await
51 }
52
53 async fn create_subscription_usage_in_tx(
54 &self,
55 user_id: UserId,
56 period_start_at: DateTimeUtc,
57 period_end_at: DateTimeUtc,
58 plan: SubscriptionKind,
59 model_requests: i32,
60 edit_predictions: i32,
61 tx: &DatabaseTransaction,
62 ) -> Result<subscription_usage::Model> {
63 // Clear out the nanoseconds so that these timestamps are comparable with Unix timestamps.
64 let period_start_at = period_start_at.with_nanosecond(0).unwrap();
65 let period_end_at = period_end_at.with_nanosecond(0).unwrap();
66
67 let period_start_at = convert_chrono_to_time(period_start_at)?;
68 let period_end_at = convert_chrono_to_time(period_end_at)?;
69
70 Ok(
71 subscription_usage::Entity::insert(subscription_usage::ActiveModel {
72 id: ActiveValue::set(Uuid::now_v7()),
73 user_id: ActiveValue::set(user_id),
74 period_start_at: ActiveValue::set(period_start_at),
75 period_end_at: ActiveValue::set(period_end_at),
76 plan: ActiveValue::set(plan),
77 model_requests: ActiveValue::set(model_requests),
78 edit_predictions: ActiveValue::set(edit_predictions),
79 })
80 .exec_with_returning(tx)
81 .await?,
82 )
83 }
84
85 pub async fn get_subscription_usage_for_period(
86 &self,
87 user_id: UserId,
88 period_start_at: DateTimeUtc,
89 period_end_at: DateTimeUtc,
90 ) -> Result<Option<subscription_usage::Model>> {
91 self.transaction(|tx| async move {
92 self.get_subscription_usage_for_period_in_tx(
93 user_id,
94 period_start_at,
95 period_end_at,
96 &tx,
97 )
98 .await
99 })
100 .await
101 }
102
103 async fn get_subscription_usage_for_period_in_tx(
104 &self,
105 user_id: UserId,
106 period_start_at: DateTimeUtc,
107 period_end_at: DateTimeUtc,
108 tx: &DatabaseTransaction,
109 ) -> Result<Option<subscription_usage::Model>> {
110 Ok(subscription_usage::Entity::find()
111 .filter(subscription_usage::Column::UserId.eq(user_id))
112 .filter(subscription_usage::Column::PeriodStartAt.eq(period_start_at))
113 .filter(subscription_usage::Column::PeriodEndAt.eq(period_end_at))
114 .one(tx)
115 .await?)
116 }
117
118 pub async fn transfer_existing_subscription_usage(
119 &self,
120 user_id: UserId,
121 existing_subscription: &billing_subscription::Model,
122 new_subscription_kind: Option<SubscriptionKind>,
123 new_subscription_status: StripeSubscriptionStatus,
124 new_period_start_at: DateTimeUtc,
125 new_period_end_at: DateTimeUtc,
126 ) -> Result<Option<subscription_usage::Model>> {
127 self.transaction(|tx| async move {
128 match (existing_subscription.kind, new_subscription_status) {
129 (Some(SubscriptionKind::ZedProTrial), StripeSubscriptionStatus::Active) => {
130 let trial_period_start_at = existing_subscription
131 .current_period_start_at()
132 .ok_or_else(|| anyhow!("No trial subscription period start"))?;
133 let trial_period_end_at = existing_subscription
134 .current_period_end_at()
135 .ok_or_else(|| anyhow!("No trial subscription period end"))?;
136
137 let existing_usage = self
138 .get_subscription_usage_for_period_in_tx(
139 user_id,
140 trial_period_start_at,
141 trial_period_end_at,
142 &tx,
143 )
144 .await?;
145 if let Some(existing_usage) = existing_usage {
146 return Ok(Some(
147 self.create_subscription_usage_in_tx(
148 user_id,
149 new_period_start_at,
150 new_period_end_at,
151 new_subscription_kind.unwrap_or(existing_usage.plan),
152 existing_usage.model_requests,
153 existing_usage.edit_predictions,
154 &tx,
155 )
156 .await?,
157 ));
158 }
159 }
160 _ => {}
161 }
162
163 Ok(None)
164 })
165 .await
166 }
167}