collab: Add billing preferences for maximum LLM monthly spend (#18948)

This PR adds a new `billing_preferences` table.

Right now there is a single preference: the maximum monthly spend for
LLM usage.

Release Notes:

- N/A

---------

Co-authored-by: Richard <richard@zed.dev>

crates/collab/src/db/ids.rs

@@ -72,6 +72,7 @@ macro_rules! id_type {
 id_type!(AccessTokenId);
 id_type!(BillingCustomerId);
 id_type!(BillingSubscriptionId);
+id_type!(BillingPreferencesId);
 id_type!(BufferId);
 id_type!(ChannelBufferCollaboratorId);
 id_type!(ChannelChatParticipantId);

crates/collab/src/db/queries.rs

@@ -2,6 +2,7 @@ use super::*;
 
 pub mod access_tokens;
 pub mod billing_customers;
+pub mod billing_preferences;
 pub mod billing_subscriptions;
 pub mod buffers;
 pub mod channels;

crates/collab/src/db/queries/billing_preferences.rs

@@ -0,0 +1,75 @@
+use super::*;
+
+#[derive(Debug)]
+pub struct CreateBillingPreferencesParams {
+    pub max_monthly_llm_usage_spending_in_cents: i32,
+}
+
+#[derive(Debug, Default)]
+pub struct UpdateBillingPreferencesParams {
+    pub max_monthly_llm_usage_spending_in_cents: ActiveValue<i32>,
+}
+
+impl Database {
+    /// Returns the billing preferences for the given user, if they exist.
+    pub async fn get_billing_preferences(
+        &self,
+        user_id: UserId,
+    ) -> Result<Option<billing_preference::Model>> {
+        self.transaction(|tx| async move {
+            Ok(billing_preference::Entity::find()
+                .filter(billing_preference::Column::UserId.eq(user_id))
+                .one(&*tx)
+                .await?)
+        })
+        .await
+    }
+
+    /// Creates new billing preferences for the given user.
+    pub async fn create_billing_preferences(
+        &self,
+        user_id: UserId,
+        params: &CreateBillingPreferencesParams,
+    ) -> Result<billing_preference::Model> {
+        self.transaction(|tx| async move {
+            let preferences = billing_preference::Entity::insert(billing_preference::ActiveModel {
+                user_id: ActiveValue::set(user_id),
+                max_monthly_llm_usage_spending_in_cents: ActiveValue::set(
+                    params.max_monthly_llm_usage_spending_in_cents,
+                ),
+                ..Default::default()
+            })
+            .exec_with_returning(&*tx)
+            .await?;
+
+            Ok(preferences)
+        })
+        .await
+    }
+
+    /// Updates the billing preferences for the given user.
+    pub async fn update_billing_preferences(
+        &self,
+        user_id: UserId,
+        params: &UpdateBillingPreferencesParams,
+    ) -> Result<billing_preference::Model> {
+        self.transaction(|tx| async move {
+            let preferences = billing_preference::Entity::update_many()
+                .set(billing_preference::ActiveModel {
+                    max_monthly_llm_usage_spending_in_cents: params
+                        .max_monthly_llm_usage_spending_in_cents
+                        .clone(),
+                    ..Default::default()
+                })
+                .filter(billing_preference::Column::UserId.eq(user_id))
+                .exec_with_returning(&*tx)
+                .await?;
+
+            Ok(preferences
+                .into_iter()
+                .next()
+                .ok_or_else(|| anyhow!("billing preferences not found"))?)
+        })
+        .await
+    }
+}

crates/collab/src/db/tables.rs

@@ -1,5 +1,6 @@
 pub mod access_token;
 pub mod billing_customer;
+pub mod billing_preference;
 pub mod billing_subscription;
 pub mod buffer;
 pub mod buffer_operation;

crates/collab/src/db/tables/billing_preference.rs

@@ -0,0 +1,30 @@
+use crate::db::{BillingPreferencesId, UserId};
+use sea_orm::entity::prelude::*;
+
+#[derive(Clone, Debug, Default, PartialEq, Eq, DeriveEntityModel)]
+#[sea_orm(table_name = "billing_preferences")]
+pub struct Model {
+    #[sea_orm(primary_key)]
+    pub id: BillingPreferencesId,
+    pub created_at: DateTime,
+    pub user_id: UserId,
+    pub max_monthly_llm_usage_spending_in_cents: i32,
+}
+
+#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
+pub enum Relation {
+    #[sea_orm(
+        belongs_to = "super::user::Entity",
+        from = "Column::UserId",
+        to = "super::user::Column::Id"
+    )]
+    User,
+}
+
+impl Related<super::user::Entity> for Entity {
+    fn to() -> RelationDef {
+        Relation::User.def()
+    }
+}
+
+impl ActiveModelBehavior for ActiveModel {}