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