diff --git a/.github/workflows/test-build.yml b/.github/workflows/test-build.yml index a2c32f5f25e..3e3fa21f7f9 100644 --- a/.github/workflows/test-build.yml +++ b/.github/workflows/test-build.yml @@ -30,6 +30,19 @@ jobs: --health-interval 5s --health-timeout 5s --health-retries 10 + postgres-legacy: + image: postgres:16-alpine + env: + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + POSTGRES_DB: sim_billing_test + ports: + - 5433:5432 + options: >- + --health-cmd "pg_isready -U postgres -d sim_billing_test" + --health-interval 5s + --health-timeout 5s + --health-retries 10 env: DATABASE_URL: postgresql://postgres:postgres@127.0.0.1:5432/sim_auth_scim OAUTH_TOKEN_FAMILY_TEST_DATABASE_URL: postgresql://postgres:postgres@127.0.0.1:5432/sim_auth_scim @@ -83,6 +96,19 @@ jobs: ee/scim/lib/managed-membership.postgres.test.ts lib/auth/sso/application/admit-sso-user.postgres.test.ts + - name: Verify cumulative billing timeout recovery in PostgreSQL + working-directory: apps/sim + env: + BILLING_USAGE_TEST_DATABASE_URL: postgresql://postgres:postgres@127.0.0.1:5432/sim_auth_scim + run: bunx vitest run lib/billing/core/usage-log.postgres.test.ts + + - name: Verify cumulative billing timeout recovery on PostgreSQL 16 + if: matrix.provision == 'push' + working-directory: apps/sim + env: + BILLING_USAGE_TEST_DATABASE_URL: postgresql://postgres:postgres@127.0.0.1:5433/sim_billing_test + run: bunx vitest run lib/billing/core/usage-log.postgres.test.ts + - name: Verify SCIM and administration over real HTTP working-directory: apps/sim env: diff --git a/apps/sim/app/api/billing/update-cost/route.ts b/apps/sim/app/api/billing/update-cost/route.ts index f9537851d57..d9524b4a6f0 100644 --- a/apps/sim/app/api/billing/update-cost/route.ts +++ b/apps/sim/app/api/billing/update-cost/route.ts @@ -312,6 +312,7 @@ async function updateCostInner(req: NextRequest, span: Span): Promise { + const databaseUrl = process.env.BILLING_USAGE_TEST_DATABASE_URL + if (databaseUrl && !['localhost', '127.0.0.1', '[::1]'].includes(new URL(databaseUrl).hostname)) { + throw new Error('Billing usage integration tests require a disposable local database') + } + return { databaseUrl, transaction: vi.fn() } +}) + +vi.unmock('@sim/db/schema') +vi.unmock('drizzle-orm') +vi.mock('@sim/db', () => ({ db: { transaction }, dbReplica: {} })) +vi.mock('@/lib/billing/core/plan', () => ({ getHighestPrioritySubscription: vi.fn() })) +vi.mock('@/lib/billing/subscriptions/utils', () => ({ isOrgScopedSubscription: vi.fn() })) + +import { + CumulativeUsageContextMismatchError, + type RecordCumulativeUsageParams, + recordCumulativeUsage, +} from '@/lib/billing/core/usage-log' + +const require = createRequire(import.meta.url) +const commonJsPostgres = require('postgres') as typeof postgres + +const schemaName = `billing_usage_${generateId().replaceAll('-', '')}` +const connection = databaseUrl + ? postgres(databaseUrl, { + max: 8, + prepare: false, + fetch_types: false, + connection: { search_path: schemaName }, + onnotice: () => undefined, + }) + : undefined +const database = connection ? drizzle(connection) : undefined + +type Transaction = Parameters['transaction']>[0]>[0] + +function deferred() { + let resolve: () => void = () => undefined + const promise = new Promise((done) => { + resolve = done + }) + return { promise, resolve } +} + +interface PausedTransaction { + reached: ReturnType + release: ReturnType + lockOnly: boolean +} + +let nextPause: PausedTransaction | undefined +let holderTimeoutSetting = 'transaction_timeout' + +function pauseNextTransaction(lockOnly = false): PausedTransaction { + const pause = { reached: deferred(), release: deferred(), lockOnly } + nextPause = pause + return pause +} + +function usage(cost: number, eventKey = 'update-cost:shared-request'): RecordCumulativeUsageParams { + return { + userId: 'actor', + workspaceId: 'workspace', + billingEntity: { type: 'organization', id: 'payer' }, + billingPeriod: { + start: new Date('2026-09-01T00:00:00.000Z'), + end: new Date('2026-10-01T00:00:00.000Z'), + }, + source: 'workspace-chat', + model: 'test-model', + eventKey, + cost, + metadata: { inputTokens: 10, outputTokens: 5 }, + } +} + +async function ledgerRows() { + if (!connection) throw new Error('PostgreSQL fixture is unavailable') + return connection<{ event_key: string; cost: string }[]>` + select event_key, cost from usage_log order by event_key + ` +} + +afterAll(async () => { + nextPause?.release.resolve() + if (connection) { + await connection.unsafe(`DROP SCHEMA IF EXISTS "${schemaName}" CASCADE`) + await connection.end() + } +}) + +describe.skipIf(!databaseUrl)('Cumulative billing with PostgreSQL', () => { + beforeAll(async () => { + if (!connection || !database) throw new Error('PostgreSQL fixture is unavailable') + const [version] = await connection` + select current_setting('server_version_num')::integer as version, + current_setting('transaction_timeout', true) is not null as has_transaction_timeout + ` + expect(version.version).toBeGreaterThanOrEqual(150000) + holderTimeoutSetting = version.has_transaction_timeout + ? 'transaction_timeout' + : 'idle_in_transaction_session_timeout' + await connection.unsafe(`CREATE SCHEMA "${schemaName}"`) + await connection.unsafe(` + CREATE TABLE usage_log ( + id text PRIMARY KEY, user_id text NOT NULL, category text NOT NULL, + source text NOT NULL, description text NOT NULL, metadata jsonb, + cost numeric NOT NULL, event_key text, billing_entity_type text, + billing_entity_id text, billing_period_start timestamp, billing_period_end timestamp, + workspace_id text, workflow_id text, execution_id text, + created_at timestamp NOT NULL DEFAULT now() + ); + CREATE UNIQUE INDEX usage_log_event_key_unique ON usage_log(event_key) + WHERE event_key IS NOT NULL; + CREATE TABLE driver_probe (id text PRIMARY KEY) + `) + transaction.mockImplementation(async (callback: (tx: Transaction) => Promise) => { + const pause = nextPause + nextPause = undefined + return database.transaction(async (tx) => { + const result = await callback(tx) + if (pause) { + if (pause.lockOnly) { + /** Reproduce the old policy: lock_timeout does not expire an idle holder. */ + await tx.execute(sql`select set_config(${holderTimeoutSetting}, '0', true)`) + } + pause.reached.resolve() + await pause.release.promise + } + return result + }) + }) + }) + + beforeEach(async () => { + nextPause = undefined + if (!connection) throw new Error('PostgreSQL fixture is unavailable') + await connection`truncate usage_log` + }) + + it.each([ + { name: 'ESM', create: postgres }, + { name: 'CommonJS', create: commonJsPostgres }, + ])( + 'rejects resumed $name transaction queries after its connection is reused', + async ({ create }) => { + const pool = create(databaseUrl!, { + max: 1, + prepare: false, + fetch_types: false, + connection: { search_path: schemaName }, + }) + const release = deferred() + const resumed = deferred() + let resumedError: unknown + const holder = pool.begin(async (tx) => { + await tx`select set_config(${holderTimeoutSetting}, '150ms', true)` + await tx`select 1` + await release.promise + try { + await tx`insert into driver_probe (id) values (${generateId()})` + } catch (error) { + resumedError = error + } finally { + resumed.resolve() + } + }) + try { + await expect(holder).rejects.toMatchObject({ code: 'CONNECTION_CLOSED' }) + /** max: 1 forces the underlying connection object to serve a new session. */ + await pool`select 1` + release.resolve() + await resumed.promise + expect(getPostgresErrorCode(resumedError)).toBe('CONNECTION_CLOSED') + const [row] = await pool`select count(*)::integer as count from driver_probe` + expect(row.count).toBe(0) + } finally { + release.resolve() + await pool.end({ timeout: 0 }) + } + } + ) + + it('reproduces a retry timing out behind an idle holder when only lock waits are bounded', async () => { + const pause = pauseNextTransaction(true) + const holder = recordCumulativeUsage(usage(0.4)) + try { + await pause.reached.promise + const failure = await recordCumulativeUsage(usage(0.8)).catch((error: unknown) => error) + expect(getPostgresErrorCode(failure)).toBe('55P03') + expect(await ledgerRows()).toEqual([]) + } finally { + pause.release.resolve() + await holder + } + expect(await ledgerRows()).toEqual([{ event_key: usage(0).eventKey, cost: '0.4' }]) + }, 10_000) + + it.each([0, 0.2])( + 'expires the idle holder, rolls back its write, and recovers exactly once (initial %s)', + async (initial) => { + if (initial > 0) await recordCumulativeUsage(usage(initial)) + const pause = pauseNextTransaction() + const holder = recordCumulativeUsage(usage(0.4)).catch((error: unknown) => error) + try { + await pause.reached.promise + const failure = await recordCumulativeUsage(usage(0.8)).catch((error: unknown) => error) + expect(getPostgresErrorCode(failure)).toBe('55P03') + const recovered = await recordCumulativeUsage(usage(0.8)) + expect(recovered.billed).toBe(true) + expect(recovered.delta).toBeCloseTo(0.8 - initial, 9) + expect(recovered.total).toBe(0.8) + expect(await recordCumulativeUsage(usage(0.8))).toEqual({ + billed: false, + delta: 0, + total: 0.8, + }) + expect(await recordCumulativeUsage(usage(0.3))).toEqual({ + billed: false, + delta: 0, + total: 0.8, + }) + expect(await ledgerRows()).toEqual([{ event_key: usage(0).eventKey, cost: '0.8' }]) + } finally { + pause.release.resolve() + } + expect(getPostgresErrorCode(await holder)).toBe('CONNECTION_CLOSED') + }, + 12_000 + ) + + it('converges concurrent out-of-order callbacks and independent events to their exact totals', async () => { + const costs = [0.4, 0.1, 0.8, 0.3, 0.8, 0.6] + const results = await Promise.all(costs.map((cost) => recordCumulativeUsage(usage(cost)))) + expect(results.reduce((total, result) => total + result.delta, 0)).toBeCloseTo(0.8, 9) + await Promise.all( + Array.from({ length: 32 }, (_, index) => + recordCumulativeUsage(usage(0.25, `independent:${index}`)) + ) + ) + expect(await ledgerRows()).toHaveLength(33) + expect(await recordCumulativeUsage(usage(0.8))).toEqual({ billed: false, delta: 0, total: 0.8 }) + }) + + it.each([0.2, 0.8])( + 'rejects an actor mismatch even for a non-increasing callback (%s)', + async (cost) => { + await recordCumulativeUsage(usage(0.8)) + await expect( + recordCumulativeUsage({ ...usage(cost), userId: 'another-actor' }) + ).rejects.toBeInstanceOf(CumulativeUsageContextMismatchError) + expect(await ledgerRows()).toEqual([{ event_key: usage(0).eventKey, cost: '0.8' }]) + } + ) +}) diff --git a/apps/sim/lib/billing/core/usage-log.test.ts b/apps/sim/lib/billing/core/usage-log.test.ts index 1098fddc277..cd612889af2 100644 --- a/apps/sim/lib/billing/core/usage-log.test.ts +++ b/apps/sim/lib/billing/core/usage-log.test.ts @@ -429,7 +429,7 @@ describe('recordCumulativeUsage', () => { }) }) - it('bounds the advisory-lock wait and locks on the 64-bit event-key hash', async () => { + it('bounds the holder lifetime and lock wait before acquiring the event-key lock', async () => { const { tx } = setupTx({ id: 'row-1', cost: '0.3474447' }) await recordCumulativeUsage({ userId: 'user-1', @@ -438,7 +438,11 @@ describe('recordCumulativeUsage', () => { cost: 0.4662453, eventKey: 'update-cost:msg-1-billing', }) + expect(executedSqlContaining(tx, 'transaction_timeout')).toBe(true) + expect(executedSqlContaining(tx, 'idle_in_transaction_session_timeout')).toBe(true) + expect(executedSqlContaining(tx, 'statement_timeout')).toBe(true) expect(executedSqlContaining(tx, 'lock_timeout')).toBe(true) + expect(tx.execute.mock.calls[0][0]).toMatchObject({ values: ['4000ms', '3500ms', '3000ms'] }) expect(executedSqlContaining(tx, 'pg_advisory_xact_lock')).toBe(true) expect(executedSqlContaining(tx, 'hashtextextended')).toBe(true) }) diff --git a/apps/sim/lib/billing/core/usage-log.ts b/apps/sim/lib/billing/core/usage-log.ts index ed0609b2e01..07ed2580368 100644 --- a/apps/sim/lib/billing/core/usage-log.ts +++ b/apps/sim/lib/billing/core/usage-log.ts @@ -2,7 +2,7 @@ import { createHash } from 'node:crypto' import { db, dbReplica } from '@sim/db' import { usageLog, workflow } from '@sim/db/schema' import { createLogger } from '@sim/logger' -import { toError } from '@sim/utils/errors' +import { getPostgresErrorCode, toError } from '@sim/utils/errors' import { generateId } from '@sim/utils/id' import { and, desc, eq, gte, inArray, lt, lte, notInArray, or, sql } from 'drizzle-orm' import { defaultBillingPeriod } from '@/lib/billing/core/billing-period' @@ -635,14 +635,18 @@ function assertCumulativeUsageLedgerBinding( } /** - * Bounds the wait for the per-event-key advisory lock (and any row/index lock - * waits inside the critical section). The Go mothership gives each UpdateCost - * POST a 5s deadline, retries 3x with backoff, then dead-letters the charge - * keyed on the same idempotency key — so a stuck lock holder must surface as - * a fast, retryable failure (SQLSTATE 55P03) within that budget rather than - * an unbounded wait that pins pooled connections. + * PostgreSQL 17+ bounds the entire transaction below the callback's five-second + * deadline. Older supported servers instead bound each idle interval between + * statements, alongside the per-statement budget. Both policies release an idle + * lock holder without waiting for its application process to resume; only the + * newer policy also limits total elapsed transaction time. */ +const CUMULATIVE_FLUSH_TRANSACTION_TIMEOUT_MS = 4_000 +const CUMULATIVE_FLUSH_STATEMENT_TIMEOUT_MS = 3_500 const CUMULATIVE_FLUSH_LOCK_TIMEOUT_MS = 3_000 +const CUMULATIVE_FLUSH_SLOW_MS = 1_000 + +type CumulativeUsageStage = 'pool' | 'configure' | 'lock' | 'read' | 'write' | 'commit' /** * Record a request's CUMULATIVE cost idempotently with monotonic top-up. @@ -655,8 +659,9 @@ const CUMULATIVE_FLUSH_LOCK_TIMEOUT_MS = 3_000 * An existing row must match the incoming actor, workspace, payer, and billing * period before either a duplicate no-op or a top-up is accepted. * The billing context is resolved BEFORE the transaction and the lock wait is - * bounded by `lock_timeout`, keeping the critical section to one SELECT plus - * one INSERT/UPDATE on a single pooled connection. + * bounded by `lock_timeout`. A server-enforced transaction deadline, or idle + * transaction deadline on older PostgreSQL, releases a stalled holder. The + * critical section uses one SELECT plus one INSERT/UPDATE on a single connection. * * Because every leg flushes its cumulative and this converges to the max, * there is no under-billing if the request recovers after a partial flush, no @@ -684,78 +689,121 @@ export async function recordCumulativeUsage( const billingContext = await resolveBillingContext(userId, billingEntity, billingPeriod) - return db.transaction(async (tx) => { - // Serialize all flushes for this request (lock auto-releases at tx end), - // with a bounded wait so a pathological holder fails this flush fast and - // lets the caller retry instead of hanging the connection. - await tx.execute( - sql`select set_config('lock_timeout', ${`${CUMULATIVE_FLUSH_LOCK_TIMEOUT_MS}ms`}, true)` - ) - await tx.execute(sql`select pg_advisory_xact_lock(hashtextextended(${eventKey}, 0))`) + const startedAt = Date.now() + let stage: CumulativeUsageStage = 'pool' + let stageStartedAt = startedAt + const stageDurationsMs: Partial> = {} + let succeeded = false + let pgCode: string | undefined + const enterStage = (nextStage: CumulativeUsageStage) => { + const now = Date.now() + stageDurationsMs[stage] = now - stageStartedAt + stage = nextStage + stageStartedAt = now + } - const [existing] = await tx - .select({ - id: usageLog.id, - cost: usageLog.cost, - userId: usageLog.userId, - workspaceId: usageLog.workspaceId, - billingEntityType: usageLog.billingEntityType, - billingEntityId: usageLog.billingEntityId, - billingPeriodStart: usageLog.billingPeriodStart, - billingPeriodEnd: usageLog.billingPeriodEnd, - }) - .from(usageLog) - .where(eq(usageLog.eventKey, eventKey)) - .limit(1) + try { + const result = await db.transaction(async (tx) => { + enterStage('configure') + await tx.execute(sql` + select + set_config( + case when current_setting('transaction_timeout', true) is null + then 'idle_in_transaction_session_timeout' + else 'transaction_timeout' + end, + ${`${CUMULATIVE_FLUSH_TRANSACTION_TIMEOUT_MS}ms`}, + true + ), + set_config('statement_timeout', ${`${CUMULATIVE_FLUSH_STATEMENT_TIMEOUT_MS}ms`}, true), + set_config('lock_timeout', ${`${CUMULATIVE_FLUSH_LOCK_TIMEOUT_MS}ms`}, true) + `) + enterStage('lock') + await tx.execute(sql`select pg_advisory_xact_lock(hashtextextended(${eventKey}, 0))`) + + enterStage('read') + const [existing] = await tx + .select({ + id: usageLog.id, + cost: usageLog.cost, + userId: usageLog.userId, + workspaceId: usageLog.workspaceId, + billingEntityType: usageLog.billingEntityType, + billingEntityId: usageLog.billingEntityId, + billingPeriodStart: usageLog.billingPeriodStart, + billingPeriodEnd: usageLog.billingPeriodEnd, + }) + .from(usageLog) + .where(eq(usageLog.eventKey, eventKey)) + .limit(1) - if (existing) { - assertCumulativeUsageLedgerBinding(existing, { - userId, - workspaceId, - billingContext, - eventKey, - }) - } + if (existing) { + assertCumulativeUsageLedgerBinding(existing, { + userId, + workspaceId, + billingContext, + eventKey, + }) + } - const recorded = existing ? Number.parseFloat(existing.cost) : 0 - const { shouldBill, delta, newTotal } = resolveCumulativeTopUp(recorded, cost) + const recorded = existing ? Number.parseFloat(existing.cost) : 0 + const { shouldBill, delta, newTotal } = resolveCumulativeTopUp(recorded, cost) - if (!shouldBill) { - return { billed: false, delta: 0, total: recorded } - } + if (!shouldBill) { + enterStage('commit') + return { billed: false, delta: 0, total: recorded } + } - if (existing) { - // Top up the single row to the new (higher) cumulative; the - // period total is SUM(usage_log.cost), so this lifts it by the delta. - await tx - .update(usageLog) - .set({ cost: newTotal.toString(), metadata: metadata ?? null }) - .where(eq(usageLog.id, existing.id)) - } else { - // First flush for this request: insert the canonical row with the - // pre-resolved billing context. Runs in the same tx + advisory lock. - await recordUsage({ - userId, - workspaceId, - tx, - billingEntity: billingContext.billingEntity, - billingPeriod: billingContext.billingPeriod, - entries: [ - { - category: 'model', - source, - description: model, - cost: newTotal, - eventKey, - sourceReference: eventKey, - ...(metadata ? { metadata } : {}), - }, - ], + enterStage('write') + if (existing) { + await tx + .update(usageLog) + .set({ cost: newTotal.toString(), metadata: metadata ?? null }) + .where(eq(usageLog.id, existing.id)) + } else { + await recordUsage({ + userId, + workspaceId, + tx, + billingEntity: billingContext.billingEntity, + billingPeriod: billingContext.billingPeriod, + entries: [ + { + category: 'model', + source, + description: model, + cost: newTotal, + eventKey, + sourceReference: eventKey, + ...(metadata ? { metadata } : {}), + }, + ], + }) + } + + enterStage('commit') + return { billed: true, delta, total: newTotal } + }) + succeeded = true + return result + } catch (error) { + pgCode = getPostgresErrorCode(error) + throw error + } finally { + const now = Date.now() + stageDurationsMs[stage] = now - stageStartedAt + const durationMs = now - startedAt + if (!succeeded || durationMs >= CUMULATIVE_FLUSH_SLOW_MS) { + logger.warn('Cumulative usage transaction did not complete promptly', { + eventKey, + succeeded, + stage, + durationMs, + stageDurationsMs, + ...(pgCode ? { pgCode } : {}), }) } - - return { billed: true, delta, total: newTotal } - }) + } } interface UsageLogFilter { diff --git a/bun.lock b/bun.lock index ad1d60991e2..42dceec6923 100644 --- a/bun.lock +++ b/bun.lock @@ -806,6 +806,7 @@ ], "patchedDependencies": { "@better-auth/oauth-provider@1.6.27": "patches/@better-auth%2Foauth-provider@1.6.27.patch", + "postgres@3.4.9": "patches/postgres@3.4.9.patch", }, "overrides": { "@hono/node-server": "1.19.15", diff --git a/package.json b/package.json index 23995fc0f6d..5f7195c283d 100644 --- a/package.json +++ b/package.json @@ -186,6 +186,7 @@ "sharp" ], "patchedDependencies": { - "@better-auth/oauth-provider@1.6.27": "patches/@better-auth%2Foauth-provider@1.6.27.patch" + "@better-auth/oauth-provider@1.6.27": "patches/@better-auth%2Foauth-provider@1.6.27.patch", + "postgres@3.4.9": "patches/postgres@3.4.9.patch" } } diff --git a/patches/README.md b/patches/README.md index be3b325cec9..b4bacf0962a 100644 --- a/patches/README.md +++ b/patches/README.md @@ -18,3 +18,30 @@ Remove this patch when upgrading to a provider version with native authorization resource preservation. Review its persisted resource model and migrate Sim's opaque-token audience binding at the same time; preserving the query alone does not enforce an access token's audience. + +# PostgreSQL transaction closure + +`postgres@3.4.9` lets a transaction callback keep using its connection object after +PostgreSQL closes that transaction's session. Resuming the callback can crash on a +null socket or execute statements on a replacement session outside the transaction. + +The version-pinned patch carries the transaction-scope closure guard from +[upstream PR #1155](https://github.com/porsager/postgres/pull/1155). It records the +connection closure in `begin()` and rejects subsequent queries from that scope, +including its implicit COMMIT/ROLLBACK, before they reach the connection or pool. +The guard is applied to all published ESM, CommonJS, and Cloudflare entry points. +It does not change connection establishment, retries, or healthy transactions. + +This is required for cumulative billing's server-enforced holder deadline: +PostgreSQL 17+ uses `transaction_timeout`; older supported servers use +`idle_in_transaction_session_timeout` alongside the statement timeout. Both release +a stalled idle holder; the older fallback limits each idle interval and statement, +not the total elapsed transaction time. +`apps/sim/lib/billing/core/usage-log.postgres.test.ts` tests real ESM/CommonJS driver +closure and reconnection, rollback after billing INSERT/UPDATE, and exact retry +accounting. CI runs it against PostgreSQL 17 and 16. Set +`BILLING_USAGE_TEST_DATABASE_URL` to a disposable local PostgreSQL 15+ database to +run it manually. + +Remove this patch when the pinned driver includes equivalent transaction-scope +closure handling. Keep the reconnect regression tests when upgrading. diff --git a/patches/postgres@3.4.9.patch b/patches/postgres@3.4.9.patch new file mode 100644 index 00000000000..c59bab4f38d --- /dev/null +++ b/patches/postgres@3.4.9.patch @@ -0,0 +1,81 @@ +diff --git a/src/index.js b/src/index.js +--- a/src/index.js ++++ b/src/index.js +@@ -237,12 +237,13 @@ + let savepoints = 0 + , connection + , prepare = null ++ , closed = false + + try { + await sql.unsafe('begin ' + options.replace(/[^a-z ]/ig, ''), [], { onexecute }).execute() + return await Promise.race([ + scope(connection, fn), +- new Promise((_, reject) => connection.onclose = reject) ++ new Promise((_, reject) => connection.onclose = e => (closed = true, reject(e))) + ]) + } catch (error) { + throw error +@@ -289,6 +290,8 @@ + } + + function handler(q) { ++ if (closed) ++ return q.reject(Errors.connection('CONNECTION_CLOSED', options)) + q.catch(e => uncaughtError || (uncaughtError = e)) + c.queue === full + ? queries.push(q) +diff --git a/cjs/src/index.js b/cjs/src/index.js +--- a/cjs/src/index.js ++++ b/cjs/src/index.js +@@ -237,12 +237,13 @@ + let savepoints = 0 + , connection + , prepare = null ++ , closed = false + + try { + await sql.unsafe('begin ' + options.replace(/[^a-z ]/ig, ''), [], { onexecute }).execute() + return await Promise.race([ + scope(connection, fn), +- new Promise((_, reject) => connection.onclose = reject) ++ new Promise((_, reject) => connection.onclose = e => (closed = true, reject(e))) + ]) + } catch (error) { + throw error +@@ -289,6 +290,8 @@ + } + + function handler(q) { ++ if (closed) ++ return q.reject(Errors.connection('CONNECTION_CLOSED', options)) + q.catch(e => uncaughtError || (uncaughtError = e)) + c.queue === full + ? queries.push(q) +diff --git a/cf/src/index.js b/cf/src/index.js +--- a/cf/src/index.js ++++ b/cf/src/index.js +@@ -238,12 +238,13 @@ + let savepoints = 0 + , connection + , prepare = null ++ , closed = false + + try { + await sql.unsafe('begin ' + options.replace(/[^a-z ]/ig, ''), [], { onexecute }).execute() + return await Promise.race([ + scope(connection, fn), +- new Promise((_, reject) => connection.onclose = reject) ++ new Promise((_, reject) => connection.onclose = e => (closed = true, reject(e))) + ]) + } catch (error) { + throw error +@@ -290,6 +291,8 @@ + } + + function handler(q) { ++ if (closed) ++ return q.reject(Errors.connection('CONNECTION_CLOSED', options)) + q.catch(e => uncaughtError || (uncaughtError = e)) + c.queue === full + ? queries.push(q)