From f22ab5ffaf66d904de3332eb19659d7fa29f0513 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Gerg=C5=91=20M=C3=B3ricz?= Date: Wed, 25 Sep 2024 20:57:45 +0200 Subject: [PATCH] feat(db): implement bill_team RPC --- apps/api/src/controllers/v0/scrape.ts | 2 +- apps/api/src/controllers/v0/search.ts | 4 +- apps/api/src/controllers/v1/map.ts | 2 +- apps/api/src/controllers/v1/scrape.ts | 2 +- apps/api/src/controllers/v1/types.ts | 8 +- apps/api/src/main/runWebScraper.ts | 2 +- .../src/services/billing/credit_billing.ts | 185 ++---------------- 7 files changed, 23 insertions(+), 182 deletions(-) diff --git a/apps/api/src/controllers/v0/scrape.ts b/apps/api/src/controllers/v0/scrape.ts index 696f8f74..db304fed 100644 --- a/apps/api/src/controllers/v0/scrape.ts +++ b/apps/api/src/controllers/v0/scrape.ts @@ -244,7 +244,7 @@ export async function scrapeController(req: Request, res: Response) { } if (creditsToBeBilled > 0) { // billing for doc done on queue end, bill only for llm extraction - billTeam(team_id, creditsToBeBilled).catch(error => { + billTeam(team_id, chunk.sub_id, creditsToBeBilled).catch(error => { Logger.error(`Failed to bill team ${team_id} for ${creditsToBeBilled} credits: ${error}`); // Optionally, you could notify an admin or add to a retry queue here }); diff --git a/apps/api/src/controllers/v0/search.ts b/apps/api/src/controllers/v0/search.ts index 970acc25..2d8c69de 100644 --- a/apps/api/src/controllers/v0/search.ts +++ b/apps/api/src/controllers/v0/search.ts @@ -18,6 +18,7 @@ export async function searchHelper( jobId: string, req: Request, team_id: string, + subscription_id: string, crawlerOptions: any, pageOptions: PageOptions, searchOptions: SearchOptions, @@ -54,7 +55,7 @@ export async function searchHelper( if (justSearch) { - billTeam(team_id, res.length).catch(error => { + billTeam(team_id, subscription_id, res.length).catch(error => { Logger.error(`Failed to bill team ${team_id} for ${res.length} credits: ${error}`); // Optionally, you could notify an admin or add to a retry queue here }); @@ -169,6 +170,7 @@ export async function searchController(req: Request, res: Response) { jobId, req, team_id, + chunk.sub_id, crawlerOptions, pageOptions, searchOptions, diff --git a/apps/api/src/controllers/v1/map.ts b/apps/api/src/controllers/v1/map.ts index 6b13f762..aca03ced 100644 --- a/apps/api/src/controllers/v1/map.ts +++ b/apps/api/src/controllers/v1/map.ts @@ -152,7 +152,7 @@ export async function mapController( // remove duplicates that could be due to http/https or www links = removeDuplicateUrls(links); - billTeam(req.auth.team_id, 1).catch((error) => { + billTeam(req.auth.team_id, req.acuc.sub_id, 1).catch((error) => { Logger.error( `Failed to bill team ${req.auth.team_id} for 1 credit: ${error}` ); diff --git a/apps/api/src/controllers/v1/scrape.ts b/apps/api/src/controllers/v1/scrape.ts index ebbabc00..41974917 100644 --- a/apps/api/src/controllers/v1/scrape.ts +++ b/apps/api/src/controllers/v1/scrape.ts @@ -108,7 +108,7 @@ export async function scrapeController( creditsToBeBilled = 5; } - billTeam(req.auth.team_id, creditsToBeBilled).catch(error => { + billTeam(req.auth.team_id, req.acuc.sub_id, creditsToBeBilled).catch(error => { Logger.error(`Failed to bill team ${req.auth.team_id} for ${creditsToBeBilled} credits: ${error}`); // Optionally, you could notify an admin or add to a retry queue here }); diff --git a/apps/api/src/controllers/v1/types.ts b/apps/api/src/controllers/v1/types.ts index c09a29f2..5e99fdab 100644 --- a/apps/api/src/controllers/v1/types.ts +++ b/apps/api/src/controllers/v1/types.ts @@ -318,10 +318,10 @@ type Account = { export type AuthCreditUsageChunk = { api_key: string; team_id: string; - sub_id: string; - sub_current_period_start: string; - sub_current_period_end: string; - price_id: string; + sub_id: string | null; + sub_current_period_start: string | null; + sub_current_period_end: string | null; + price_id: string | null; price_credits: number; // credit limit with assoicated price, or free_credits (500) if free plan credits_used: number; coupon_credits: number; diff --git a/apps/api/src/main/runWebScraper.ts b/apps/api/src/main/runWebScraper.ts index f67a1cd0..571122f9 100644 --- a/apps/api/src/main/runWebScraper.ts +++ b/apps/api/src/main/runWebScraper.ts @@ -120,7 +120,7 @@ export async function runWebScraper({ : docs; if(is_scrape === false) { - billTeam(team_id, filteredDocs.length).catch(error => { + billTeam(team_id, undefined, filteredDocs.length).catch(error => { Logger.error(`Failed to bill team ${team_id} for ${filteredDocs.length} credits: ${error}`); // Optionally, you could notify an admin or add to a retry queue here }); diff --git a/apps/api/src/services/billing/credit_billing.ts b/apps/api/src/services/billing/credit_billing.ts index 48dd3a17..9d0d1f09 100644 --- a/apps/api/src/services/billing/credit_billing.ts +++ b/apps/api/src/services/billing/credit_billing.ts @@ -3,168 +3,30 @@ import { withAuth } from "../../lib/withAuth"; import { sendNotification } from "../notification/email_notification"; import { supabase_service } from "../supabase"; import { Logger } from "../../lib/logger"; -import { getValue, setValue } from "../redis"; -import { redlock } from "../redlock"; import * as Sentry from "@sentry/node"; import { AuthCreditUsageChunk } from "../../controllers/v1/types"; const FREE_CREDITS = 500; - -export async function billTeam(team_id: string, credits: number) { - return withAuth(supaBillTeam)(team_id, credits); +/** + * If you do not know the subscription_id in the current context, pass subscription_id as undefined. + */ +export async function billTeam(team_id: string, subscription_id: string | null | undefined, credits: number) { + return withAuth(supaBillTeam)(team_id, subscription_id, credits); } -export async function supaBillTeam(team_id: string, credits: number) { +export async function supaBillTeam(team_id: string, subscription_id: string, credits: number) { if (team_id === "preview") { return { success: true, message: "Preview team, no credits used" }; } Logger.info(`Billing team ${team_id} for ${credits} credits`); - // When the API is used, you can log the credit usage in the credit_usage table: - // team_id: The ID of the team using the API. - // subscription_id: The ID of the team's active subscription. - // credits_used: The number of credits consumed by the API call. - // created_at: The timestamp of the API usage. - // 1. get the subscription and check for available coupons concurrently - const [{ data: subscription }, { data: coupons }] = await Promise.all([ - supabase_service - .from("subscriptions") - .select("*") - .eq("team_id", team_id) - .eq("status", "active") - .single(), - supabase_service - .from("coupons") - .select("id, credits") - .eq("team_id", team_id) - .eq("status", "active"), - ]); - - let couponCredits = 0; - let sortedCoupons = []; - - if (coupons && coupons.length > 0) { - couponCredits = coupons.reduce( - (total, coupon) => total + coupon.credits, - 0 - ); - sortedCoupons = [...coupons].sort((a, b) => b.credits - a.credits); + const { error } = + await supabase_service.rpc("bill_team", { _team_id: team_id, sub_id: subscription_id ?? null, fetch_subscription: subscription_id === undefined, credits }); + + if (error) { + Sentry.captureException(error); + Logger.error("Failed to bill team: " + JSON.stringify(error)); } - // using coupon credits: - if (couponCredits > 0) { - // if there is no subscription and they have enough coupon credits - if (!subscription) { - // using only coupon credits: - // if there are enough coupon credits - if (couponCredits >= credits) { - // remove credits from coupon credits - let usedCredits = credits; - while (usedCredits > 0) { - // update coupons - if (sortedCoupons[0].credits < usedCredits) { - usedCredits = usedCredits - sortedCoupons[0].credits; - // update coupon credits - await supabase_service - .from("coupons") - .update({ - credits: 0, - }) - .eq("id", sortedCoupons[0].id); - sortedCoupons.shift(); - } else { - // update coupon credits - await supabase_service - .from("coupons") - .update({ - credits: sortedCoupons[0].credits - usedCredits, - }) - .eq("id", sortedCoupons[0].id); - usedCredits = 0; - } - } - - return await createCreditUsage({ team_id, credits: 0 }); - - // not enough coupon credits and no subscription - } else { - // update coupon credits - const usedCredits = credits - couponCredits; - for (let i = 0; i < sortedCoupons.length; i++) { - await supabase_service - .from("coupons") - .update({ - credits: 0, - }) - .eq("id", sortedCoupons[i].id); - } - - return await createCreditUsage({ team_id, credits: usedCredits }); - } - } - - // with subscription - // using coupon + subscription credits: - if (credits > couponCredits) { - // update coupon credits - for (let i = 0; i < sortedCoupons.length; i++) { - await supabase_service - .from("coupons") - .update({ - credits: 0, - }) - .eq("id", sortedCoupons[i].id); - } - const usedCredits = credits - couponCredits; - return await createCreditUsage({ - team_id, - subscription_id: subscription.id, - credits: usedCredits, - }); - } else { - // using only coupon credits - let usedCredits = credits; - while (usedCredits > 0) { - // update coupons - if (sortedCoupons[0].credits < usedCredits) { - usedCredits = usedCredits - sortedCoupons[0].credits; - // update coupon credits - await supabase_service - .from("coupons") - .update({ - credits: 0, - }) - .eq("id", sortedCoupons[0].id); - sortedCoupons.shift(); - } else { - // update coupon credits - await supabase_service - .from("coupons") - .update({ - credits: sortedCoupons[0].credits - usedCredits, - }) - .eq("id", sortedCoupons[0].id); - usedCredits = 0; - } - } - - return await createCreditUsage({ - team_id, - subscription_id: subscription.id, - credits: 0, - }); - } - } - - // not using coupon credits - if (!subscription) { - return await createCreditUsage({ team_id, credits }); - } - - return await createCreditUsage({ - team_id, - subscription_id: subscription.id, - credits, - }); } export async function checkTeamCredits(chunk: AuthCreditUsageChunk, team_id: string, credits: number) { @@ -297,26 +159,3 @@ export async function countCreditsAndRemainingForCurrentBillingPeriod( totalCredits: price.credits, }; } - -async function createCreditUsage({ - team_id, - subscription_id, - credits, -}: { - team_id: string; - subscription_id?: string; - credits: number; -}) { - await supabase_service - .from("credit_usage") - .insert([ - { - team_id, - credits_used: credits, - subscription_id: subscription_id || null, - created_at: new Date(), - }, - ]); - - return { success: true }; -}