diff --git a/apps/sim/lib/billing/core/ledger-read.test.ts b/apps/sim/lib/billing/core/ledger-read.test.ts new file mode 100644 index 00000000000..0f7e77f5323 --- /dev/null +++ b/apps/sim/lib/billing/core/ledger-read.test.ts @@ -0,0 +1,35 @@ +/** + * @vitest-environment node + */ +import { beforeEach, describe, expect, it, vi } from 'vitest' +import { USAGE_LEDGER_STATEMENT_TIMEOUT_MS } from '@/lib/billing/constants' +import { readLedgerBounded } from '@/lib/billing/core/ledger-read' +import type { DbClient } from '@/lib/db/types' + +const renderedSql = (statement: unknown) => + (statement as { toSQL: () => { sql: string } }).toSQL().sql + +describe('readLedgerBounded', () => { + const execute = vi.fn().mockResolvedValue([]) + const tx = { execute } + const transaction = vi.fn((callback: (client: typeof tx) => Promise) => callback(tx)) + const executor = { transaction } as unknown as DbClient + + beforeEach(() => vi.clearAllMocks()) + + it('bounds the statement inside one transaction on the given client, before the read', async () => { + const read = vi.fn().mockResolvedValue([{ cost: '12.5' }]) + await expect(readLedgerBounded(executor, read)).resolves.toEqual([{ cost: '12.5' }]) + expect(transaction).toHaveBeenCalledTimes(1) + expect(execute.mock.calls.map(([statement]) => renderedSql(statement))).toEqual([ + `SET LOCAL statement_timeout = '${USAGE_LEDGER_STATEMENT_TIMEOUT_MS}ms'`, + ]) + expect(read).toHaveBeenCalledWith(tx) + expect(execute.mock.invocationCallOrder[0]).toBeLessThan(read.mock.invocationCallOrder[0]) + }) + + it('surfaces the read failure to the caller', async () => { + const failure = new Error('canceling statement due to statement timeout') + await expect(readLedgerBounded(executor, () => Promise.reject(failure))).rejects.toBe(failure) + }) +}) diff --git a/apps/sim/lib/billing/core/ledger-read.ts b/apps/sim/lib/billing/core/ledger-read.ts new file mode 100644 index 00000000000..577d1b7b5d0 --- /dev/null +++ b/apps/sim/lib/billing/core/ledger-read.ts @@ -0,0 +1,25 @@ +import { sql } from 'drizzle-orm' +import { USAGE_LEDGER_STATEMENT_TIMEOUT_MS } from '@/lib/billing/constants' +import type { DbClient, DbTransaction } from '@/lib/db/types' + +/** + * Runs one aggregate over a payer's usage ledger in a transaction of its own, bounded by + * {@link USAGE_LEDGER_STATEMENT_TIMEOUT_MS}. `SET LOCAL` scopes the bound to that transaction, + * so it ends with the read and never reaches the pool. Every sum over a payer's billing period + * reads through here, whether it admits a run, closes a cycle or previews a bill: a payer whose + * period has grown past what one statement can sum within the bound fails at the database + * instead of holding a connection without limit, and a caller that admits on the answer can + * size its own deadline from the bound. Reads keyed to one execution or one stamped period + * boundary, and the platform-wide admin analytics, are not period sums and read directly. + */ +export function readLedgerBounded( + executor: DbClient, + read: (tx: DbTransaction) => Promise +): Promise { + return executor.transaction(async (tx) => { + await tx.execute( + sql.raw(`SET LOCAL statement_timeout = '${USAGE_LEDGER_STATEMENT_TIMEOUT_MS}ms'`) + ) + return read(tx) + }) +} diff --git a/apps/sim/lib/billing/core/usage-gate-cache.ts b/apps/sim/lib/billing/core/usage-gate-cache.ts index dbcf326f3d4..a520220b3bf 100644 --- a/apps/sim/lib/billing/core/usage-gate-cache.ts +++ b/apps/sim/lib/billing/core/usage-gate-cache.ts @@ -23,14 +23,16 @@ export const USAGE_GATE_TTL_MS = 5 * 60 * 1000 /** * How long a coalesced usage read may take before its callers give up on it. The read's cost is - * the ledger sum, which the database ends at {@link USAGE_LEDGER_STATEMENT_TIMEOUT_MS}; the - * remainder is a few indexed lookups and the connection waits around them. The singleflight - * default of 30 s exists to bound a hung producer, and a slow sum is not a hung one: given up on - * early, it keeps running detached while every joined caller fails and the next caller starts a - * second sum alongside it. Derived from the statement bound so the database always ends the sum - * first, and the gate only gives up on a connection that never answers. + * its ledger aggregates, each of which the database ends at + * {@link USAGE_LEDGER_STATEMENT_TIMEOUT_MS}; at most two run in sequence (the payer's usage, + * then a member's cap), and the remainder is a few indexed lookups and the connection waits + * around them. The singleflight default of 30 s exists to bound a hung producer, and a slow + * aggregate is not a hung one: given up on early, it keeps running detached while every joined + * caller fails and the next caller starts a second one alongside it. Sized from the statement + * bound so the database always ends the aggregates first, and the gate only gives up on a + * connection that never answers. */ -export const USAGE_GATE_SETTLE_TIMEOUT_MS = USAGE_LEDGER_STATEMENT_TIMEOUT_MS + 15_000 +export const USAGE_GATE_SETTLE_TIMEOUT_MS = 2 * USAGE_LEDGER_STATEMENT_TIMEOUT_MS + 15_000 /** * Recent gate answers, admitted and refused, with `LRUCache` supplying the TTL diff --git a/apps/sim/lib/billing/core/usage-log.test.ts b/apps/sim/lib/billing/core/usage-log.test.ts index ccb88ae4c0b..f595e67cf2e 100644 --- a/apps/sim/lib/billing/core/usage-log.test.ts +++ b/apps/sim/lib/billing/core/usage-log.test.ts @@ -38,6 +38,10 @@ import { CUMULATIVE_COST_EPSILON, CumulativeUsageContextMismatchError, getBillingPeriodUsageCost, + getBillingPeriodUsageCostByUser, + getBillingPeriodUsageCostWithSourceSubset, + getBillingPeriodWorkflowRunCount, + getStampedPeriodRangeUsageCostByUser, getUserUsageLogs, getWorkspaceUsageLogs, recordCumulativeUsage, @@ -557,34 +561,83 @@ describe('usage-log query scopes', () => { }) }) -describe('getBillingPeriodUsageCost', () => { +describe('ledger aggregates', () => { + const billingEntity = { type: 'organization' as const, id: 'org-1' } + const billingPeriod = { + start: new Date('2026-05-01T00:00:00Z'), + end: new Date('2027-05-01T00:00:00Z'), + } + /** Every aggregate over the ledger, with the row the mocked read hands back and the value it yields. */ + const aggregates: Array<{ + name: string + read: () => Promise + rows: unknown[] + expected: unknown + }> = [ + { + name: 'getBillingPeriodUsageCost', + read: () => getBillingPeriodUsageCost(billingEntity, billingPeriod), + rows: [{ cost: '12.5' }], + expected: 12.5, + }, + { + name: 'getBillingPeriodWorkflowRunCount', + read: () => getBillingPeriodWorkflowRunCount(billingEntity, billingPeriod), + rows: [{ workflowRuns: 7 }], + expected: 7, + }, + { + name: 'getBillingPeriodUsageCostWithSourceSubset', + read: () => + getBillingPeriodUsageCostWithSourceSubset(billingEntity, billingPeriod, ['workflow']), + rows: [{ total: '20', subset: '5' }], + expected: { total: 20, subset: 5 }, + }, + { + name: 'getBillingPeriodUsageCostByUser', + read: () => getBillingPeriodUsageCostByUser(billingEntity, billingPeriod), + rows: [{ userId: 'user-1', cost: '3' }], + expected: new Map([['user-1', 3]]), + }, + { + name: 'getStampedPeriodRangeUsageCostByUser', + read: () => + getStampedPeriodRangeUsageCostByUser(billingEntity, { + from: billingPeriod.start, + to: billingPeriod.end, + }), + rows: [{ userId: 'user-2', cost: '4' }], + expected: new Map([['user-2', 4]]), + }, + ] + beforeEach(() => { vi.clearAllMocks() installSharedDbMocks() }) - it('bounds the ledger sum with its own statement timeout inside one transaction', async () => { - const execute = vi.fn().mockResolvedValue([]) - const where = vi.fn().mockResolvedValue([{ cost: '12.5' }]) - const tx = { execute, select: vi.fn(() => ({ from: vi.fn(() => ({ where })) })) } - mockTransaction.mockImplementation((callback: (client: typeof tx) => Promise) => - callback(tx) - ) - - const cost = await getBillingPeriodUsageCost( - { type: 'organization', id: 'org-1' }, - { start: new Date('2026-05-01T00:00:00Z'), end: new Date('2027-05-01T00:00:00Z') } - ) - - expect(cost).toBe(12.5) - expect(mockTransaction).toHaveBeenCalledTimes(1) - const executed = execute.mock.calls.map( - ([statement]) => (statement as { toSQL: () => { sql: string } }).toSQL().sql - ) - expect(executed).toContain( - `SET LOCAL statement_timeout = '${USAGE_LEDGER_STATEMENT_TIMEOUT_MS}ms'` - ) - /** The bound is set before the sum runs, not after. */ - expect(execute.mock.invocationCallOrder[0]).toBeLessThan(where.mock.invocationCallOrder[0]) - }) + for (const aggregate of aggregates) { + it(`${aggregate.name} reads through the bounded ledger transaction`, async () => { + const execute = vi.fn().mockResolvedValue([]) + const terminal = vi.fn().mockResolvedValue(aggregate.rows) + const chain: Record = {} + for (const step of ['select', 'from', 'where', 'leftJoin']) chain[step] = vi.fn(() => chain) + chain.groupBy = terminal + chain.then = (resolve: (rows: unknown[]) => unknown) => terminal().then(resolve) + const tx = { execute, select: chain.select } + mockTransaction.mockImplementation((callback: (client: typeof tx) => Promise) => + callback(tx) + ) + + await expect(aggregate.read()).resolves.toEqual(aggregate.expected) + expect(mockTransaction).toHaveBeenCalledTimes(1) + expect( + execute.mock.calls.map( + ([statement]) => (statement as { toSQL: () => { sql: string } }).toSQL().sql + ) + ).toEqual([`SET LOCAL statement_timeout = '${USAGE_LEDGER_STATEMENT_TIMEOUT_MS}ms'`]) + /** The bound is set before the aggregate runs, not after. */ + expect(execute.mock.invocationCallOrder[0]).toBeLessThan(terminal.mock.invocationCallOrder[0]) + }) + } }) diff --git a/apps/sim/lib/billing/core/usage-log.ts b/apps/sim/lib/billing/core/usage-log.ts index 0849352ef35..e44e1ba314f 100644 --- a/apps/sim/lib/billing/core/usage-log.ts +++ b/apps/sim/lib/billing/core/usage-log.ts @@ -15,8 +15,8 @@ import { textKey, timestampKey, } from '@/lib/api/list-query' -import { USAGE_LEDGER_STATEMENT_TIMEOUT_MS } from '@/lib/billing/constants' import { defaultBillingPeriod } from '@/lib/billing/core/billing-period' +import { readLedgerBounded } from '@/lib/billing/core/ledger-read' import { getHighestPrioritySubscription } from '@/lib/billing/core/plan' import { resolveSubscriptionUsagePeriod, @@ -216,10 +216,6 @@ async function resolveBillingContext( /** * Returns attributed ledger usage for a billing entity/period. The ledger is * the sole source of truth for usage — there is no userStats baseline. - * - * The sum runs in a transaction of its own on the given client so that it can - * be bounded by {@link USAGE_LEDGER_STATEMENT_TIMEOUT_MS} for that statement - * alone: `SET LOCAL` ends with the transaction and never reaches the pool. */ export async function getBillingPeriodUsageCost( billingEntity: BillingEntity, @@ -243,17 +239,14 @@ export async function getBillingPeriodUsageCost( ) } - const [row] = await executor.transaction(async (tx) => { - await tx.execute( - sql.raw(`SET LOCAL statement_timeout = '${USAGE_LEDGER_STATEMENT_TIMEOUT_MS}ms'`) - ) - return tx + const [row] = await readLedgerBounded(executor, (tx) => + tx .select({ cost: sql`COALESCE(SUM(${usageLog.cost}), 0)`, }) .from(usageLog) .where(and(...conditions)) - }) + ) return Number.parseFloat(row?.cost ?? '0') } @@ -274,36 +267,38 @@ export async function getBillingPeriodWorkflowRunCount( billingPeriod: UsageQueryPeriod, executor: DbClient = db ): Promise { - const [row] = await executor - .select({ - /** - * The exclusion goes through `notInArray`, not `<> ALL(${array})`. Interpolating - * a JavaScript array into a `sql` template emits parenthesized scalar binds — - * `ALL(($1))` — which Postgres rejects outright with "op ANY/ALL (array) - * requires array on right side". Unit tests cannot catch it, because `@sim/db` - * is mocked and no statement is ever rendered. - */ - workflowRuns: - sql`COUNT(DISTINCT ${usageLog.executionId}) FILTER (WHERE ${usageLog.source} = 'workflow' AND ${notInArray(usageLog.category, [...UNBILLED_USAGE_CATEGORIES])})`.mapWith( - Number - ), - }) - .from(usageLog) - .where( - and( - eq(usageLog.billingEntityType, billingEntity.type), - eq(usageLog.billingEntityId, billingEntity.id), - ...(billingPeriod.source === 'reporting' - ? [ - gte(usageLog.createdAt, billingPeriod.start), - lt(usageLog.createdAt, billingPeriod.end), - ] - : [ - eq(usageLog.billingPeriodStart, billingPeriod.start), - eq(usageLog.billingPeriodEnd, billingPeriod.end), - ]) + const [row] = await readLedgerBounded(executor, (tx) => + tx + .select({ + /** + * The exclusion goes through `notInArray`, not `<> ALL(${array})`. Interpolating + * a JavaScript array into a `sql` template emits parenthesized scalar binds — + * `ALL(($1))` — which Postgres rejects outright with "op ANY/ALL (array) + * requires array on right side". Unit tests cannot catch it, because `@sim/db` + * is mocked and no statement is ever rendered. + */ + workflowRuns: + sql`COUNT(DISTINCT ${usageLog.executionId}) FILTER (WHERE ${usageLog.source} = 'workflow' AND ${notInArray(usageLog.category, [...UNBILLED_USAGE_CATEGORIES])})`.mapWith( + Number + ), + }) + .from(usageLog) + .where( + and( + eq(usageLog.billingEntityType, billingEntity.type), + eq(usageLog.billingEntityId, billingEntity.id), + ...(billingPeriod.source === 'reporting' + ? [ + gte(usageLog.createdAt, billingPeriod.start), + lt(usageLog.createdAt, billingPeriod.end), + ] + : [ + eq(usageLog.billingPeriodStart, billingPeriod.start), + eq(usageLog.billingPeriodEnd, billingPeriod.end), + ]) + ) ) - ) + ) return row?.workflowRuns ?? 0 } @@ -321,27 +316,29 @@ export async function getBillingPeriodUsageCostWithSourceSubset( source: UsageLogSource[], executor: DbClient = db ): Promise<{ total: number; subset: number }> { - const [row] = await executor - .select({ - total: sql`COALESCE(SUM(${usageLog.cost}), 0)`, - subset: sql`COALESCE(SUM(${usageLog.cost}) FILTER (WHERE ${inArray(usageLog.source, source)}), 0)`, - }) - .from(usageLog) - .where( - and( - eq(usageLog.billingEntityType, billingEntity.type), - eq(usageLog.billingEntityId, billingEntity.id), - ...(billingPeriod.source === 'reporting' - ? [ - gte(usageLog.createdAt, billingPeriod.start), - lt(usageLog.createdAt, billingPeriod.end), - ] - : [ - eq(usageLog.billingPeriodStart, billingPeriod.start), - eq(usageLog.billingPeriodEnd, billingPeriod.end), - ]) + const [row] = await readLedgerBounded(executor, (tx) => + tx + .select({ + total: sql`COALESCE(SUM(${usageLog.cost}), 0)`, + subset: sql`COALESCE(SUM(${usageLog.cost}) FILTER (WHERE ${inArray(usageLog.source, source)}), 0)`, + }) + .from(usageLog) + .where( + and( + eq(usageLog.billingEntityType, billingEntity.type), + eq(usageLog.billingEntityId, billingEntity.id), + ...(billingPeriod.source === 'reporting' + ? [ + gte(usageLog.createdAt, billingPeriod.start), + lt(usageLog.createdAt, billingPeriod.end), + ] + : [ + eq(usageLog.billingPeriodStart, billingPeriod.start), + eq(usageLog.billingPeriodEnd, billingPeriod.end), + ]) + ) ) - ) + ) return { total: Number.parseFloat(row?.total ?? '0'), @@ -377,14 +374,16 @@ export async function getBillingPeriodUsageCostByUser( } if (userIds) conditions.push(inArray(usageLog.userId, [...userIds])) - const rows = await executor - .select({ - userId: usageLog.userId, - cost: sql`COALESCE(SUM(${usageLog.cost}), 0)`, - }) - .from(usageLog) - .where(and(...conditions)) - .groupBy(usageLog.userId) + const rows = await readLedgerBounded(executor, (tx) => + tx + .select({ + userId: usageLog.userId, + cost: sql`COALESCE(SUM(${usageLog.cost}), 0)`, + }) + .from(usageLog) + .where(and(...conditions)) + .groupBy(usageLog.userId) + ) return new Map(rows.map((row) => [row.userId, Number.parseFloat(row.cost ?? '0')])) } @@ -418,14 +417,16 @@ export async function getStampedPeriodRangeUsageCostByUser( ) } - const rows = await executor - .select({ - userId: usageLog.userId, - cost: sql`COALESCE(SUM(${usageLog.cost}), 0)`, - }) - .from(usageLog) - .where(and(...conditions)) - .groupBy(usageLog.userId) + const rows = await readLedgerBounded(executor, (tx) => + tx + .select({ + userId: usageLog.userId, + cost: sql`COALESCE(SUM(${usageLog.cost}), 0)`, + }) + .from(usageLog) + .where(and(...conditions)) + .groupBy(usageLog.userId) + ) return new Map(rows.map((row) => [row.userId, Number.parseFloat(row.cost ?? '0')])) } diff --git a/apps/sim/lib/billing/credits/weekly-refresh.test.ts b/apps/sim/lib/billing/credits/weekly-refresh.test.ts index 95344da6943..d1a796e73ab 100644 --- a/apps/sim/lib/billing/credits/weekly-refresh.test.ts +++ b/apps/sim/lib/billing/credits/weekly-refresh.test.ts @@ -3,6 +3,7 @@ */ import { dbChainMockFns, drizzleOrmMock, schemaMock } from '@sim/testing' import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' +import { USAGE_LEDGER_STATEMENT_TIMEOUT_MS } from '@/lib/billing/constants' vi.mock('drizzle-orm', () => { const sqlTag = () => { @@ -11,7 +12,9 @@ vi.mock('drizzle-orm', () => { } return { ...drizzleOrmMock, - sql: Object.assign(sqlTag, { raw: sqlTag }), + sql: Object.assign(sqlTag, { + raw: (rawSql: string) => ({ rawSql, toSQL: () => ({ sql: rawSql, params: [] }) }), + }), sum: () => ({ as: () => 'sum' }), } }) @@ -49,6 +52,14 @@ describe('computeBillingPeriodUsageWithWeeklyRefresh', () => { { ledgerTotal: '25.00', refreshWeekTotal: '12.00' }, { ledgerTotal: '25.00', refreshWeekTotal: '4.00' }, ]) + /** The scan is a ledger aggregate: it runs inside the bounded ledger transaction. */ + const boundBeforeScan = () => + dbChainMockFns.transaction.mock.calls.length === 1 && + dbChainMockFns.execute.mock.calls.some(([statement]) => + String((statement as { toSQL?: () => { sql: string } }).toSQL?.().sql).includes( + `SET LOCAL statement_timeout = '${USAGE_LEDGER_STATEMENT_TIMEOUT_MS}ms'` + ) + ) await expect( computeBillingPeriodUsageWithWeeklyRefresh({ @@ -65,6 +76,7 @@ describe('computeBillingPeriodUsageWithWeeklyRefresh', () => { periodStart ) expect(drizzleOrmMock.eq).toHaveBeenCalledWith(schemaMock.usageLog.billingPeriodEnd, periodEnd) + expect(boundBeforeScan()).toBe(true) }) it('uses the reporting time range for the ledger while retaining captured-period refresh', async () => { diff --git a/apps/sim/lib/billing/credits/weekly-refresh.ts b/apps/sim/lib/billing/credits/weekly-refresh.ts index 8779fc7389f..a12de029d63 100644 --- a/apps/sim/lib/billing/credits/weekly-refresh.ts +++ b/apps/sim/lib/billing/credits/weekly-refresh.ts @@ -24,6 +24,7 @@ import { db } from '@sim/db' import { usageLog } from '@sim/db/schema' import { createLogger } from '@sim/logger' import { and, eq, gte, lt, or, sql, sum } from 'drizzle-orm' +import { readLedgerBounded } from '@/lib/billing/core/ledger-read' import type { BillingEntity, UsageQueryPeriod } from '@/lib/billing/core/usage-log' import type { DbClient } from '@/lib/db/types' @@ -88,29 +89,31 @@ export async function computeBillingPeriodUsageWithWeeklyRefresh( const startEpoch = Math.floor(refreshPeriodStart.getTime() / 1000) const capEpoch = Math.floor(cap.getTime() / 1000) - const rows = await executor - .select({ - weekIndex: - sql`FLOOR((LEAST(GREATEST(EXTRACT(EPOCH FROM ${usageLog.createdAt}), ${startEpoch}), ${capEpoch - 1}) - ${startEpoch}) / 604800)`.as( - 'week_index' + const rows = await readLedgerBounded(executor, (tx) => + tx + .select({ + weekIndex: + sql`FLOOR((LEAST(GREATEST(EXTRACT(EPOCH FROM ${usageLog.createdAt}), ${startEpoch}), ${capEpoch - 1}) - ${startEpoch}) / 604800)`.as( + 'week_index' + ), + ledgerTotal: + sql`SUM(SUM(${usageLog.cost}) FILTER (WHERE ${ledgerPeriodFilter})) OVER ()`.as( + 'ledger_total' + ), + refreshWeekTotal: sql`SUM(${usageLog.cost}) FILTER (WHERE ${refreshFilter})`.as( + 'refresh_week_total' ), - ledgerTotal: - sql`SUM(SUM(${usageLog.cost}) FILTER (WHERE ${ledgerPeriodFilter})) OVER ()`.as( - 'ledger_total' - ), - refreshWeekTotal: sql`SUM(${usageLog.cost}) FILTER (WHERE ${refreshFilter})`.as( - 'refresh_week_total' - ), - }) - .from(usageLog) - .where( - and( - eq(usageLog.billingEntityType, billingEntity.type), - eq(usageLog.billingEntityId, billingEntity.id), - scanFilter + }) + .from(usageLog) + .where( + and( + eq(usageLog.billingEntityType, billingEntity.type), + eq(usageLog.billingEntityId, billingEntity.id), + scanFilter + ) ) - ) - .groupBy(sql`week_index`) + .groupBy(sql`week_index`) + ) let refreshConsumed = 0 for (const row of rows) { @@ -169,23 +172,25 @@ export async function computeWeeklyRefreshConsumed( // refresh in the period's final week rather than fall out of the deduction. const startEpoch = Math.floor(periodStart.getTime() / 1000) const capEpoch = Math.floor(cap.getTime() / 1000) - const rows = await executor - .select({ - weekIndex: - sql`FLOOR((LEAST(GREATEST(EXTRACT(EPOCH FROM ${usageLog.createdAt}), ${startEpoch}), ${capEpoch - 1}) - ${startEpoch}) / 604800)`.as( - 'week_index' - ), - weekTotal: sum(usageLog.cost).as('week_total'), - }) - .from(usageLog) - .where( - and( - eq(usageLog.billingEntityType, billingEntity.type), - eq(usageLog.billingEntityId, billingEntity.id), - eq(usageLog.billingPeriodStart, periodStart) + const rows = await readLedgerBounded(executor, (tx) => + tx + .select({ + weekIndex: + sql`FLOOR((LEAST(GREATEST(EXTRACT(EPOCH FROM ${usageLog.createdAt}), ${startEpoch}), ${capEpoch - 1}) - ${startEpoch}) / 604800)`.as( + 'week_index' + ), + weekTotal: sum(usageLog.cost).as('week_total'), + }) + .from(usageLog) + .where( + and( + eq(usageLog.billingEntityType, billingEntity.type), + eq(usageLog.billingEntityId, billingEntity.id), + eq(usageLog.billingPeriodStart, periodStart) + ) ) - ) - .groupBy(sql`week_index`) + .groupBy(sql`week_index`) + ) let totalConsumed = 0 for (const row of rows) { diff --git a/apps/sim/lib/billing/enterprise-provisioning.test.ts b/apps/sim/lib/billing/enterprise-provisioning.test.ts index b69eb60f452..1368760c0ee 100644 --- a/apps/sim/lib/billing/enterprise-provisioning.test.ts +++ b/apps/sim/lib/billing/enterprise-provisioning.test.ts @@ -199,8 +199,6 @@ describe('Enterprise issuance preflight', () => { queueTableRows(schemaMock.workspace, []) queueTableRows(schemaMock.workspace, []) queueTableRows(schemaMock.subscription, []) - /** The run count resolves first; the ledger sum opens its bounded transaction before it reads. */ - queueTableRows(schemaMock.usageLog, [{ workflowRuns: 0 }]) queueTableRows(schemaMock.usageLog, [{ cost: '150' }]) const result = await getEnterpriseIssuancePreflight({ diff --git a/apps/sim/lib/billing/organizations/member-limits.test.ts b/apps/sim/lib/billing/organizations/member-limits.test.ts index 8c1c9e60782..aed15bf0b53 100644 --- a/apps/sim/lib/billing/organizations/member-limits.test.ts +++ b/apps/sim/lib/billing/organizations/member-limits.test.ts @@ -3,6 +3,7 @@ */ import { dbChainMockFns, queueTableRows, resetDbChainMock } from '@sim/testing' import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' +import { USAGE_LEDGER_STATEMENT_TIMEOUT_MS } from '@/lib/billing/constants' const { schemaTables, @@ -57,7 +58,10 @@ vi.mock('drizzle-orm', () => ({ isNull: mockIsNull, lt: mockLt, or: mockOr, - sql: vi.fn((strings: TemplateStringsArray, ...values: unknown[]) => ({ strings, values })), + sql: Object.assign( + vi.fn((strings: TemplateStringsArray, ...values: unknown[]) => ({ strings, values })), + { raw: (rawSql: string) => ({ rawSql, toSQL: () => ({ sql: rawSql, params: [] }) }) } + ), })) vi.mock('@/lib/billing/core/billing', () => ({ @@ -109,6 +113,16 @@ describe('getOrgMemberUsageForBillingPeriod', () => { getOrgMemberUsageForBillingPeriod('snapshot-org', 'actor-2', billingPeriod) ).resolves.toBe(4.5) + /** The member sum is a ledger aggregate: it runs inside the bounded ledger transaction. */ + expect(dbChainMockFns.transaction).toHaveBeenCalledTimes(1) + expect( + dbChainMockFns.execute.mock.calls.some(([statement]) => + String((statement as { toSQL?: () => { sql: string } }).toSQL?.().sql).includes( + `SET LOCAL statement_timeout = '${USAGE_LEDGER_STATEMENT_TIMEOUT_MS}ms'` + ) + ) + ).toBe(true) + expect(mockEq).toHaveBeenCalledWith('usageLog.billingEntityType', 'organization') expect(mockEq).toHaveBeenCalledWith('usageLog.billingEntityId', 'snapshot-org') expect(mockEq).toHaveBeenCalledWith('usageLog.userId', 'actor-2') diff --git a/apps/sim/lib/billing/organizations/member-limits.ts b/apps/sim/lib/billing/organizations/member-limits.ts index 94b6c294ece..a78f8e98c93 100644 --- a/apps/sim/lib/billing/organizations/member-limits.ts +++ b/apps/sim/lib/billing/organizations/member-limits.ts @@ -11,6 +11,7 @@ import { generateId } from '@sim/utils/id' import { and, eq, gte, isNull, lt, or, sql } from 'drizzle-orm' import { getOrganizationSubscription } from '@/lib/billing/core/billing' import { defaultBillingPeriod } from '@/lib/billing/core/billing-period' +import { readLedgerBounded } from '@/lib/billing/core/ledger-read' import { resolveSubscriptionUsagePeriod } from '@/lib/billing/core/reporting-period' import type { UsageQueryPeriod } from '@/lib/billing/core/usage-log' import { toDecimal, toNumber } from '@/lib/billing/utils/decimal' @@ -143,43 +144,45 @@ export async function getOrgMemberUsageForBillingPeriod( userId: string, billingPeriod: UsageQueryPeriod ): Promise { - const [row] = await db - .select({ cost: sql`COALESCE(SUM(${usageLog.cost}), 0)` }) - .from(usageLog) - .leftJoin(workspace, eq(workspace.id, usageLog.workspaceId)) - .where( - and( - eq(usageLog.userId, userId), - ...(billingPeriod.source === 'reporting' - ? [ - eq(usageLog.billingEntityType, 'organization'), - eq(usageLog.billingEntityId, organizationId), - gte(usageLog.createdAt, billingPeriod.start), - lt(usageLog.createdAt, billingPeriod.end), - ] - : [ - or( - and( - eq(usageLog.billingEntityType, 'organization'), - eq(usageLog.billingEntityId, organizationId), - eq(usageLog.billingPeriodStart, billingPeriod.start), - eq(usageLog.billingPeriodEnd, billingPeriod.end) - ), - and( - isNull(usageLog.billingEntityType), - isNull(usageLog.billingEntityId), - eq(workspace.organizationId, organizationId), - or( - isNull(workspace.organizationAssignedAt), - gte(usageLog.createdAt, workspace.organizationAssignedAt) + const [row] = await readLedgerBounded(db, (tx) => + tx + .select({ cost: sql`COALESCE(SUM(${usageLog.cost}), 0)` }) + .from(usageLog) + .leftJoin(workspace, eq(workspace.id, usageLog.workspaceId)) + .where( + and( + eq(usageLog.userId, userId), + ...(billingPeriod.source === 'reporting' + ? [ + eq(usageLog.billingEntityType, 'organization'), + eq(usageLog.billingEntityId, organizationId), + gte(usageLog.createdAt, billingPeriod.start), + lt(usageLog.createdAt, billingPeriod.end), + ] + : [ + or( + and( + eq(usageLog.billingEntityType, 'organization'), + eq(usageLog.billingEntityId, organizationId), + eq(usageLog.billingPeriodStart, billingPeriod.start), + eq(usageLog.billingPeriodEnd, billingPeriod.end) ), - gte(usageLog.createdAt, billingPeriod.start), - lt(usageLog.createdAt, billingPeriod.end) - ) - ), - ]) + and( + isNull(usageLog.billingEntityType), + isNull(usageLog.billingEntityId), + eq(workspace.organizationId, organizationId), + or( + isNull(workspace.organizationAssignedAt), + gte(usageLog.createdAt, workspace.organizationAssignedAt) + ), + gte(usageLog.createdAt, billingPeriod.start), + lt(usageLog.createdAt, billingPeriod.end) + ) + ), + ]) + ) ) - ) + ) return Number.parseFloat(row?.cost ?? '0') }