appks/service/workspace/billing.rs

use serde::{Deserialize, Serialize};
use uuid::Uuid;

use crate::error::AppError;
use crate::models::common::Role;
use crate::models::workspaces::{Workspace, WorkspaceBilling};
use crate::service::WorkspaceService;
use crate::session::Session;

use super::util::{merge_optional_text, set_local_user_id};

#[derive(Deserialize, Serialize, Clone, Debug, utoipa::ToSchema)]
pub struct UpdateBillingParams {
    pub plan: Option<String>,
    pub billing_email: Option<String>,
    pub seats: Option<i32>,
}

impl WorkspaceService {
    pub async fn workspace_billing(
        &self,
        ctx: &Session,
        ws: &Workspace,
    ) -> Result<WorkspaceBilling, AppError> {
        let user_uid = ctx.user().ok_or(AppError::Unauthorized)?;
        self.ensure_workspace_role_at_least(user_uid, &ws, Role::Owner)
            .await?;
        self.ensure_workspace_billing(ws.id).await
    }

    pub async fn workspace_update_billing(
        &self,
        ctx: &Session,
        ws: &Workspace,
        params: UpdateBillingParams,
    ) -> Result<WorkspaceBilling, AppError> {
        let user_uid = ctx.user().ok_or(AppError::Unauthorized)?;
        self.ensure_workspace_role_at_least(user_uid, &ws, Role::Owner)
            .await?;

        let current = self.ensure_workspace_billing(ws.id).await?;
        let plan = params.plan.unwrap_or(current.plan.clone());
        let billing_email = merge_optional_text(params.billing_email, current.billing_email);
        let seats = params.seats.unwrap_or(current.seats);
        let now = chrono::Utc::now();

        let mut txn = self
            .ctx
            .db
            .writer()
            .begin()
            .await
            .map_err(|_| AppError::TxnError)?;
        sqlx::query(set_local_user_id(user_uid))
            .execute(&mut *txn)
            .await
            .map_err(AppError::Database)?;

        let result = sqlx::query_as::<_, WorkspaceBilling>(
            "UPDATE workspace_billing SET plan = $1, billing_email = $2, seats = $3, updated_at = $4 \
             WHERE workspace_id = $5 \
             RETURNING workspace_id, customer_id, subscription_id, plan, billing_email, status, \
             seats, trial_ends_at, current_period_start, current_period_end, canceled_at, \
             created_at, updated_at",
        )
        .bind(&plan)
        .bind(&billing_email)
        .bind(seats)
        .bind(now)
        .bind(ws.id)
        .fetch_one(&mut *txn)
        .await
        .map_err(AppError::Database)?;

        txn.commit().await.map_err(|_| AppError::TxnError)?;
        Ok(result)
    }

    async fn ensure_workspace_billing(
        &self,
        workspace_id: Uuid,
    ) -> Result<WorkspaceBilling, AppError> {
        if let Some(billing) = self.find_workspace_billing(workspace_id).await? {
            return Ok(billing);
        }
        let now = chrono::Utc::now();
        sqlx::query(
            "INSERT INTO workspace_billing (workspace_id, plan, status, seats, created_at, updated_at) \
             VALUES ($1, 'free', 'active', 1, $2, $2) ON CONFLICT (workspace_id) DO NOTHING",
        )
        .bind(workspace_id)
        .bind(now)
        .execute(self.ctx.db.writer())
        .await
        .map_err(AppError::Database)?;
        self.find_workspace_billing(workspace_id)
            .await?
            .ok_or(AppError::NotFound("workspace billing not found".into()))
    }

    async fn find_workspace_billing(
        &self,
        workspace_id: Uuid,
    ) -> Result<Option<WorkspaceBilling>, AppError> {
        sqlx::query_as::<_, WorkspaceBilling>(
            "SELECT workspace_id, customer_id, subscription_id, plan, billing_email, status, \
             seats, trial_ends_at, current_period_start, current_period_end, canceled_at, \
             created_at, updated_at FROM workspace_billing WHERE workspace_id = $1",
        )
        .bind(workspace_id)
        .fetch_optional(self.ctx.db.reader())
        .await
        .map_err(AppError::Database)
    }
}