fix(db): give billing and auth their own connection pools

postgres.js pipelines new queries onto busy connections when none is idle,
so slow reads saturating the process pool made short usage-ledger
transactions wait seconds for BEGIN until Postgres killed them.

- Add `billing` (4) and `auth` (3) dbFor roles, with DATABASE_URL_BILLING
  and DATABASE_URL_AUTH overrides falling back to the process URL.
- Billing pool: recordCumulativeUsage, recordUsage without a caller tx,
  the execution logger's usage reconciliation, both threshold settlement
  transactions, and recordJobRunSpend. Each moved transaction only queries
  through its tx handle; readers stay on the process pool.
- Auth pool: getAuthPool() backs the Better Auth adapter, getAuthDatabase()
  outside a transaction, the OAuth adapter guard defaults, and the API-key
  hash lookup.
- isPrimaryPoolClient() in @sim/db replaces the `executor === db` checks in
  the security, session, and SSO policy caches, so reads through the auth
  pool keep hitting them instead of costing two queries per cached session.
- Integration test saturates the process pool and requires the billing
  write, an adapter session read, and an API-key lookup to finish anyway.

Author: waleedlatif1

apps/sim/app/api/billing/update-cost/replay.live.test.ts

@@ -47,7 +47,7 @@ vi.mock('@sim/db', async () => {
   })
   state.client = client
   const db = drizzle(client)
-  return { db, dbReplica: db }
+  return { db, dbReplica: db, dbFor: () => db }
 })
 
 /** Subscription and payment-provider fixtures; callback, ledger and replay code are real. */

apps/sim/app/api/billing/update-cost/worker.live.test.ts

@@ -47,7 +47,7 @@ vi.mock('@sim/db', async () => {
   })
   state.client = client
   const db = drizzle(client)
-  return { db, dbReplica: db }
+  return { db, dbReplica: db, dbFor: () => db }
 })
 
 /** Subscription and payment-provider fixtures; callback, ledger and replay code are real. */

apps/sim/lib/api-key/service.ts

@@ -3,6 +3,7 @@ import { apiKey as apiKeyTable, user as userTable } from '@sim/db/schema'
 import { createLogger, setRequestAuth } from '@sim/logger'
 import { and, eq, isNull, lt, or } from 'drizzle-orm'
 import { hashApiKey } from '@/lib/api-key/crypto'
+import { getAuthPool } from '@/lib/auth/database-context'
 import { createDetachedTouch } from '@/lib/core/utils/background'
 import { getUserEntityPermissions } from '@/lib/workspaces/permissions/utils'
 import { getWorkspaceBillingSettings, type WorkspaceBillingSettings } from '@/lib/workspaces/utils'
@@ -62,7 +63,7 @@ export async function authenticateApiKeyFromHeader(
     }
 
     const keyHash = hashApiKey(apiKeyHeader)
-    const rows: HashCandidate[] = await db
+    const rows: HashCandidate[] = await getAuthPool()
       .select({
         id: apiKeyTable.id,
         userId: apiKeyTable.userId,

apps/sim/lib/auth/database-context.ts

@@ -1,12 +1,21 @@
 import { AsyncLocalStorage } from 'node:async_hooks'
-import { db } from '@sim/db'
+import { type db, dbFor } from '@sim/db'
 import type { DbOrTx } from '@/lib/db/types'
 
 const authDatabase = new AsyncLocalStorage<DbOrTx>()
 
-/** The Drizzle executor backing the current Better Auth transaction, or the root database. */
+/**
+ * The auth pool's root client. Session, adapter, and API-key lookups gate every
+ * authenticated request, so they keep their own connections when slow work
+ * saturates the process pool.
+ */
+export function getAuthPool(): typeof db {
+  return dbFor('auth')
+}
+
+/** The Drizzle executor backing the current Better Auth transaction, or the auth pool. */
 export function getAuthDatabase(): DbOrTx {
-  return authDatabase.getStore() ?? db
+  return authDatabase.getStore() ?? getAuthPool()
 }
 
 /** Keeps database hooks on the adapter's connection so they can see uncommitted auth records. */

apps/sim/lib/auth/oauth-provider-adapter-guard.ts

@@ -1,10 +1,11 @@
 import { AsyncLocalStorage } from 'node:async_hooks'
-import { db } from '@sim/db'
+import type { db } from '@sim/db'
 import { oauthConsent, oauthTokenFamily } from '@sim/db/schema'
 import { generateId } from '@sim/utils/id'
 import type { drizzleAdapter } from 'better-auth/adapters/drizzle'
 import { APIError } from 'better-auth/api'
 import { inArray } from 'drizzle-orm'
+import { getAuthPool } from '@/lib/auth/database-context'
 import { getOAuthIssuedResource } from '@/lib/auth/oauth-resource'
 import { capabilityRefusal } from '@/lib/permission-groups/capabilities'
 import { isCapabilityWithheldForUser } from '@/lib/permission-groups/user-scope.server'
@@ -37,7 +38,7 @@ async function deleteTrackedFamilies(
  */
 export async function withOAuthProviderIssuanceCompensation<T extends Response>(
   work: () => Promise<T>,
-  database: AuthDatabase = db
+  database: AuthDatabase = getAuthPool()
 ): Promise<T> {
   const familyIds = new Set<string>()
   try {
@@ -87,7 +88,7 @@ function requireConsentInsert(data: Record<string, unknown>): OAuthConsentInsert
  */
 export function guardOAuthProviderWrites(
   adapter: BetterAuthAdapter,
-  database: AuthDatabase = db
+  database: AuthDatabase = getAuthPool()
 ): BetterAuthAdapter {
   return {
     ...adapter,

apps/sim/lib/auth/security-policy.test.ts

@@ -1,6 +1,6 @@
 import { db } from '@sim/db'
-import { dbChainMockFns, resetDbChainMock } from '@sim/testing'
-import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest'
+import { databaseMockFns, dbChainMockFns, resetDbChainMock } from '@sim/testing'
+import { afterAll, afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 import {
   getMemberOrganizationId,
   getSecurityPolicyVersion,
@@ -146,3 +146,37 @@ describe('getSessionCookieCacheVersion', () => {
     expect(limit).toHaveBeenCalledTimes(2)
   })
 })
+
+describe('shared caches behind a sub-pool client', () => {
+  beforeEach(() => {
+    resetDbChainMock()
+    invalidateMembershipCache('pool-user')
+    invalidateSecurityPolicyVersionCache('pool-org')
+  })
+
+  afterEach(() => {
+    databaseMockFns.mockIsPrimaryPoolClient.mockReset()
+  })
+
+  /**
+   * Better Auth reads the cookie-cache version on every cached-session request through the auth
+   * pool, not `db`; caching only reads made through `db` would cost two queries per request.
+   */
+  it('serves a repeat cookie-cache version from cache when read through the auth pool', async () => {
+    const authPool: typeof db = Object.create(db)
+    databaseMockFns.mockIsPrimaryPoolClient.mockImplementation(
+      (executor) => executor === db || executor === authPool
+    )
+    dbChainMockFns.limit
+      .mockResolvedValueOnce([{ organizationId: 'pool-org' }])
+      .mockResolvedValueOnce([{ version: 3 }])
+
+    await expect(getSessionCookieCacheVersion({ userId: 'pool-user' }, authPool)).resolves.toBe(
+      'pool-org:3'
+    )
+    await expect(getSessionCookieCacheVersion({ userId: 'pool-user' }, authPool)).resolves.toBe(
+      'pool-org:3'
+    )
+    expect(dbChainMockFns.limit).toHaveBeenCalledTimes(2)
+  })
+})

apps/sim/lib/auth/security-policy.ts

@@ -1,4 +1,4 @@
-import { db } from '@sim/db'
+import { db, isPrimaryPoolClient } from '@sim/db'
 import { member, organization } from '@sim/db/schema'
 import { createLogger } from '@sim/logger'
 import { eq } from 'drizzle-orm'
@@ -50,7 +50,7 @@ export async function getSecurityPolicyVersion(
 ): Promise<number> {
   if (!organizationId) return DEFAULT_VERSION
 
-  const cached = executor === db ? versionCache.get(organizationId) : undefined
+  const cached = isPrimaryPoolClient(executor) ? versionCache.get(organizationId) : undefined
   if (cached !== undefined) return cached
 
   try {
@@ -61,7 +61,7 @@ export async function getSecurityPolicyVersion(
       .limit(1)
 
     const version = row?.version ?? DEFAULT_VERSION
-    if (executor === db) versionCache.set(organizationId, version)
+    if (isPrimaryPoolClient(executor)) versionCache.set(organizationId, version)
     return version
   } catch (error) {
     logger.error('Failed to resolve security policy version; using default', {
@@ -133,7 +133,7 @@ export async function getMemberOrganizationId(
   if (!userId) return null
 
   /** A transaction may have just created or moved the membership; shared cache entries cannot see it. */
-  const cached = executor === db ? membershipCache.get(userId) : undefined
+  const cached = isPrimaryPoolClient(executor) ? membershipCache.get(userId) : undefined
   if (cached) return cached.organizationId
 
   try {
@@ -144,7 +144,7 @@ export async function getMemberOrganizationId(
       .limit(1)
 
     const organizationId = row?.organizationId ?? null
-    if (executor === db) {
+    if (isPrimaryPoolClient(executor)) {
       membershipCache.set(userId, { organizationId }, { ttl: membershipCacheTtlMs(organizationId) })
     }
     return organizationId

apps/sim/lib/auth/session-policy.ts

@@ -1,4 +1,4 @@
-import { db } from '@sim/db'
+import { db, isPrimaryPoolClient } from '@sim/db'
 import type { SessionPolicySettings } from '@sim/db/schema'
 import { organization } from '@sim/db/schema'
 import { createLogger } from '@sim/logger'
@@ -57,7 +57,7 @@ export async function getSessionPolicy(
   if (!organizationId) return NO_POLICY
 
   /** Uncommitted policy reads must neither consume nor populate the shared cache. */
-  const cached = executor === db ? policyCache.get(organizationId) : undefined
+  const cached = isPrimaryPoolClient(executor) ? policyCache.get(organizationId) : undefined
   if (cached) return cached
 
   try {
@@ -78,7 +78,7 @@ export async function getSessionPolicy(
           idleTimeoutHours: settings.idleTimeoutHours ?? null,
         }
       : NO_POLICY
-    if (executor === db) policyCache.set(organizationId, policy)
+    if (isPrimaryPoolClient(executor)) policyCache.set(organizationId, policy)
     return policy
   } catch (error) {
     logger.error('Failed to resolve session policy; applying no policy', {

apps/sim/lib/auth/sim-auth-adapter.ts

@@ -1,10 +1,9 @@
-import { db } from '@sim/db'
 import * as schema from '@sim/db/schema'
 import type { BetterAuthOptions } from 'better-auth'
 import { drizzleAdapter } from 'better-auth/adapters/drizzle'
 import { APIError } from 'better-auth/api'
 import { and, eq } from 'drizzle-orm'
-import { runWithAuthDatabase } from '@/lib/auth/database-context'
+import { getAuthPool, runWithAuthDatabase } from '@/lib/auth/database-context'
 import {
   type AuthDatabase,
   guardOAuthProviderWrites,
@@ -66,7 +65,7 @@ function createTransactionAdapter(
  */
 export function createSimAuthAdapter(
   options: BetterAuthOptions,
-  database: AuthDatabase = db
+  database: AuthDatabase = getAuthPool()
 ): BetterAuthAdapter {
   const guarded = createGuardedAdapter(options, database)
   const transaction: BetterAuthAdapter['transaction'] = (callback) =>

apps/sim/lib/auth/sso-policy.ts

@@ -1,4 +1,4 @@
-import { db } from '@sim/db'
+import { db, isPrimaryPoolClient } from '@sim/db'
 import { organization } from '@sim/db/schema'
 import { createLogger } from '@sim/logger'
 import { APIError } from 'better-auth/api'
@@ -43,7 +43,7 @@ export async function isSsoRequiredForOrganization(
   if (!organizationId) return false
 
   /** Uncommitted reads must neither consume nor populate the shared cache. */
-  const cached = executor === db ? requirementCache.get(organizationId) : undefined
+  const cached = isPrimaryPoolClient(executor) ? requirementCache.get(organizationId) : undefined
   if (cached !== undefined) return cached
 
   const [row] = await executor
@@ -64,7 +64,7 @@ export async function isSsoRequiredForOrganization(
       onError: 'throw',
     })) &&
     (await hasSignInCapableSsoProvider(organizationId, executor))
-  if (executor === db) requirementCache.set(organizationId, required)
+  if (isPrimaryPoolClient(executor)) requirementCache.set(organizationId, required)
   return required
 }