1use super::*;
 2
 3#[derive(Debug)]
 4pub struct CreateBillingPreferencesParams {
 5    pub max_monthly_llm_usage_spending_in_cents: i32,
 6}
 7
 8#[derive(Debug, Default)]
 9pub struct UpdateBillingPreferencesParams {
10    pub max_monthly_llm_usage_spending_in_cents: ActiveValue<i32>,
11}
12
13impl Database {
14    /// Returns the billing preferences for the given user, if they exist.
15    pub async fn get_billing_preferences(
16        &self,
17        user_id: UserId,
18    ) -> Result<Option<billing_preference::Model>> {
19        self.transaction(|tx| async move {
20            Ok(billing_preference::Entity::find()
21                .filter(billing_preference::Column::UserId.eq(user_id))
22                .one(&*tx)
23                .await?)
24        })
25        .await
26    }
27
28    /// Creates new billing preferences for the given user.
29    pub async fn create_billing_preferences(
30        &self,
31        user_id: UserId,
32        params: &CreateBillingPreferencesParams,
33    ) -> Result<billing_preference::Model> {
34        self.transaction(|tx| async move {
35            let preferences = billing_preference::Entity::insert(billing_preference::ActiveModel {
36                user_id: ActiveValue::set(user_id),
37                max_monthly_llm_usage_spending_in_cents: ActiveValue::set(
38                    params.max_monthly_llm_usage_spending_in_cents,
39                ),
40                ..Default::default()
41            })
42            .exec_with_returning(&*tx)
43            .await?;
44
45            Ok(preferences)
46        })
47        .await
48    }
49
50    /// Updates the billing preferences for the given user.
51    pub async fn update_billing_preferences(
52        &self,
53        user_id: UserId,
54        params: &UpdateBillingPreferencesParams,
55    ) -> Result<billing_preference::Model> {
56        self.transaction(|tx| async move {
57            let preferences = billing_preference::Entity::update_many()
58                .set(billing_preference::ActiveModel {
59                    max_monthly_llm_usage_spending_in_cents: params
60                        .max_monthly_llm_usage_spending_in_cents
61                        .clone(),
62                    ..Default::default()
63                })
64                .filter(billing_preference::Column::UserId.eq(user_id))
65                .exec_with_returning(&*tx)
66                .await?;
67
68            Ok(preferences
69                .into_iter()
70                .next()
71                .ok_or_else(|| anyhow!("billing preferences not found"))?)
72        })
73        .await
74    }
75}