Revert "Revert "feat(web): serve /usage analytics from Postgres replica""

apps/web/src/lib/drizzle.test.ts

@@ -1,4 +1,10 @@
-import { pool, db, selectReplicaUrl, shouldExitOnPoolError } from '@/lib/drizzle';
+import {
+  pool,
+  db,
+  selectReplicaUrl,
+  selectUsageReplicaUrl,
+  shouldExitOnPoolError,
+} from '@/lib/drizzle';
 
 describe('drizzle', () => {
   describe('pool', () => {
@@ -106,6 +112,39 @@ describe('drizzle', () => {
       ).toBe('postgres://eu-2');
     });
 
+    it('uses the dedicated usage replica when configured', () => {
+      expect(
+        selectUsageReplicaUrl({
+          primaryUrl,
+          nodeEnv: 'production',
+          usageReplicaUrl: 'postgres://eu-2',
+          fallbackReplicaUrl: 'postgres://eu-1',
+        })
+      ).toBe('postgres://eu-2');
+    });
+
+    it('falls back to the standard replica when the usage replica is unset', () => {
+      expect(
+        selectUsageReplicaUrl({
+          primaryUrl,
+          nodeEnv: 'production',
+          usageReplicaUrl: undefined,
+          fallbackReplicaUrl: 'postgres://eu-1',
+        })
+      ).toBe('postgres://eu-1');
+    });
+
+    it('uses the primary for usage reads in local development', () => {
+      expect(
+        selectUsageReplicaUrl({
+          primaryUrl,
+          nodeEnv: 'development',
+          usageReplicaUrl: 'postgres://eu-2',
+          fallbackReplicaUrl: 'postgres://eu-1',
+        })
+      ).toBe(primaryUrl);
+    });
+
     it('falls back to the primary when the regional replica is unavailable', () => {
       expect(
         selectReplicaUrl({

apps/web/src/lib/drizzle.ts

@@ -68,21 +68,48 @@ export function selectReplicaUrl({
 /**
  * Get the read replica URL based on deployment region.
  * - US deployments use the US replica (POSTGRES_REPLICA_US_URL) for lower latency
- * - EU deployments randomly select one of two EU replicas to split read traffic
- *   across ~2,200 concurrent Vercel instances (~50/50 statistical distribution)
+ * - EU deployments use POSTGRES_REPLICA_EU_URL. POSTGRES_REPLICA_EU_URL_2 is
+ *   reserved for usage-analytics scans so those queries stay off this pool.
  * - Falls back to primary if no replica URL is configured for the region
  */
 function getReplicaUrl(): string {
   if (NODE_ENV === 'development') return postgresUrl;
+  const euReplicaUrl = getEnvVariable('POSTGRES_REPLICA_EU_URL');
   return selectReplicaUrl({
     primaryUrl: postgresUrl,
     nodeEnv: NODE_ENV,
     vercelRegion: VERCEL_REGION,
     usReplicaUrl: getEnvVariable('POSTGRES_REPLICA_US_URL'),
-    euReplicaUrls: [
-      getEnvVariable('POSTGRES_REPLICA_EU_URL'),
-      getEnvVariable('POSTGRES_REPLICA_EU_URL_2'),
-    ].filter(Boolean) as string[],
+    euReplicaUrls: euReplicaUrl ? [euReplicaUrl] : [],
+  });
+}
+
+/**
+ * Replica used by /usage analytics. Always POSTGRES_REPLICA_EU_URL_2 in
+ * production so heavy microdollar_usage scans do not compete with ordinary
+ * readDb traffic. Falls back to the standard replica, then primary.
+ */
+export function selectUsageReplicaUrl({
+  primaryUrl,
+  nodeEnv,
+  usageReplicaUrl,
+  fallbackReplicaUrl,
+}: {
+  primaryUrl: string;
+  nodeEnv: string | undefined;
+  usageReplicaUrl: string | undefined;
+  fallbackReplicaUrl: string;
+}): string {
+  if (nodeEnv === 'development') return primaryUrl;
+  return usageReplicaUrl || fallbackReplicaUrl;
+}
+
+function getUsageReplicaUrl(): string {
+  return selectUsageReplicaUrl({
+    primaryUrl: postgresUrl,
+    nodeEnv: NODE_ENV,
+    usageReplicaUrl: getEnvVariable('POSTGRES_REPLICA_EU_URL_2'),
+    fallbackReplicaUrl: getReplicaUrl(),
   });
 }
 
@@ -121,13 +148,35 @@ const replica = usesSeparateReplica
   : primary;
 const replicaPool = replica.pool;
 
+const usageReplicaUrl = getUsageReplicaUrl();
+export const usesDedicatedUsageReplica =
+  usageReplicaUrl !== postgresUrl && usageReplicaUrl !== replicaUrl;
+
+const usageReplica = usesDedicatedUsageReplica
+  ? createDrizzleClient({
+      connectionString: usageReplicaUrl,
+      poolConfig: {
+        ...sharedPoolConfig,
+        max: 2,
+        application_name: `${appName}-usage-replica`,
+      },
+      logger: !!DEBUG_QUERY_LOGGING,
+    })
+  : usageReplicaUrl === replicaUrl
+    ? replica
+    : primary;
+const usageReplicaPool = usageReplica.pool;
+
 // Attach pools to ensure idle connections close before suspension
 // Skip in test environment as it interferes with Jest's cleanup
 if (process.env.NODE_ENV !== 'test') {
   attachDatabasePool(pool);
   if (usesSeparateReplica) {
     attachDatabasePool(replicaPool);
   }
+  if (usesDedicatedUsageReplica) {
+    attachDatabasePool(usageReplicaPool);
+  }
 }
 
 /**
@@ -174,6 +223,15 @@ if (usesSeparateReplica) {
   });
 }
 
+if (usesDedicatedUsageReplica) {
+  usageReplicaPool.on('error', (err: Error) => {
+    console.error('Unexpected error on idle client (usage-replica)', err);
+    if (shouldExitOnPoolError('replica')) {
+      process.exit(-1);
+    }
+  });
+}
+
 // Pool observability is handled centrally by /api/cron/db-pool-metrics,
 // which scrapes PgBouncer metrics from the Supabase Prometheus endpoint
 // for all databases (primary + replicas) every minute.
@@ -196,6 +254,13 @@ const primaryDb = primary.db;
  */
 export const readDb = replica.db;
 
+/**
+ * Read replica reserved for /usage analytics scans.
+ * Points at POSTGRES_REPLICA_EU_URL_2 in production so those queries do not
+ * share the standard readDb pool. Falls back to readDb, then primary.
+ */
+export const usageReadDb = usageReplica.db;
+
 /**
  * Default database instance - connects to the primary database.
  * Use this for writes and for reads that need strong consistency.
@@ -225,6 +290,9 @@ export async function closeAllDrizzleConnections(): Promise<void> {
   if (usesSeparateReplica) {
     await replicaPool.end();
   }
+  if (usesDedicatedUsageReplica) {
+    await usageReplicaPool.end();
+  }
 }
 
 export type DrizzleTransaction = Parameters<Parameters<typeof db.transaction>[0]>[0];

apps/web/src/routers/usage-analytics-router.test.ts

@@ -1,5 +1,6 @@
 jest.mock('@/lib/redis', () => ({ redisClient: {} }));
 
+import { PgDialect } from 'drizzle-orm/pg-core';
 import {
   BreakdownInputSchema,
   CostSourceSchema,
@@ -25,26 +26,38 @@ const PARENT_ORG = '11111111-1111-4111-8111-111111111111';
 const CHILD_ORG_A = '22222222-2222-4222-8222-222222222222';
 const CHILD_ORG_B = '33333333-3333-4333-8333-333333333333';
 
+const dialect = new PgDialect();
+
+function compile(builder: WhereBuilder) {
+  const sql = builder.toSQL();
+  if (!sql) return { sql: '', params: [] as unknown[] };
+  const compiled = dialect.sqlToQuery(sql);
+  return { sql: compiled.sql, params: compiled.params };
+}
+
 function scopeSql(rawFilters: Record<string, unknown>) {
   const filters = UsageAnalyticsFiltersSchema.parse({ ...baseFilters, ...rawFilters });
   const where = new WhereBuilder();
   buildScopeConditions(where, filters, CTX_USER);
-  return { sql: where.sql(), bindings: where.bindings.map(b => b.value) };
+  return compile(where);
 }
 
 describe('usage analytics cost source', () => {
   it('defaults to billable cost for existing clients', () => {
     expect(UsageAnalyticsFiltersSchema.parse(baseFilters).costSource).toBe('cost');
-    expect(costColumnFor('cost')).toBe('total_cost_microdollars');
-    expect(costSumExprSql('cost')).toBe('COALESCE(SUM(total_cost_microdollars), 0)');
+    expect(dialect.sqlToQuery(costColumnFor('cost')).sql).toContain('cost');
+    expect(dialect.sqlToQuery(costSumExprSql('cost')).sql).toMatch(/COALESCE\(SUM\(/);
+    expect(dialect.sqlToQuery(costSumExprSql('cost')).sql).toContain('cost');
+    expect(dialect.sqlToQuery(costSumExprSql('cost')).sql).not.toContain('market_cost');
   });
 
-  it('uses the estimated market cost rollup when selected', () => {
+  it('uses the estimated market cost when selected', () => {
     expect(
       UsageAnalyticsFiltersSchema.parse({ ...baseFilters, costSource: 'market' }).costSource
     ).toBe('market');
-    expect(costColumnFor('market')).toBe('total_market_cost_microdollars');
-    expect(costSumExprSql('market')).toBe('COALESCE(SUM(total_market_cost_microdollars), 0)');
+    expect(dialect.sqlToQuery(costColumnFor('market')).sql).toContain('market_cost');
+    expect(dialect.sqlToQuery(costSumExprSql('market')).sql).toMatch(/COALESCE\(SUM\(/);
+    expect(dialect.sqlToQuery(costSumExprSql('market')).sql).toContain('market_cost');
   });
 
   it('rejects arbitrary cost source values', () => {
@@ -56,53 +69,63 @@ describe('usage analytics cost source', () => {
 
 describe('usage analytics scope conditions', () => {
   it('pins a single org to the caller in self view', () => {
-    const { sql, bindings } = scopeSql({ organizationId: PARENT_ORG, viewAs: 'self' });
-    expect(sql).toContain('organization_id = ?');
-    expect(sql).toContain('kilo_user_id = ?');
-    expect(bindings).toEqual([PARENT_ORG, CTX_USER]);
+    const { sql, params } = scopeSql({ organizationId: PARENT_ORG, viewAs: 'self' });
+    expect(sql).toContain('organization_id');
+    expect(sql).toContain('kilo_user_id');
+    expect(sql).not.toMatch(/is null/i);
+    expect(params).toEqual([PARENT_ORG, CTX_USER]);
   });
 
   it('does not pin to the caller in org-wide view', () => {
-    const { sql, bindings } = scopeSql({ organizationId: PARENT_ORG, viewAs: 'org-wide' });
-    expect(sql).toContain('organization_id = ?');
+    const { sql, params } = scopeSql({ organizationId: PARENT_ORG, viewAs: 'org-wide' });
+    expect(sql).toContain('organization_id');
     expect(sql).not.toContain('kilo_user_id');
-    expect(bindings).toEqual([PARENT_ORG]);
+    expect(params).toEqual([PARENT_ORG]);
   });
 
   it('aggregates org-wide across all orgs when organizationIds is set', () => {
-    const { sql, bindings } = scopeSql({
+    const { sql, params } = scopeSql({
       organizationIds: [PARENT_ORG, CHILD_ORG_A, CHILD_ORG_B],
     });
-    expect(sql).toContain('organization_id IN (?, ?, ?)');
+    expect(sql).toContain('organization_id');
+    expect(sql).toMatch(/in/i);
     expect(sql).not.toContain('kilo_user_id');
-    expect(bindings).toEqual([PARENT_ORG, CHILD_ORG_A, CHILD_ORG_B]);
+    expect(params).toEqual([PARENT_ORG, CHILD_ORG_A, CHILD_ORG_B]);
   });
 
   it('honors explicit user filters in the all-orgs aggregate', () => {
-    const { sql, bindings } = scopeSql({
+    const { sql, params } = scopeSql({
       organizationIds: [PARENT_ORG, CHILD_ORG_A],
       userIds: [CTX_USER],
     });
-    expect(sql).toContain('organization_id IN (?, ?)');
-    expect(sql).toContain('kilo_user_id IN (?)');
-    expect(bindings).toEqual([PARENT_ORG, CHILD_ORG_A, CTX_USER]);
+    expect(sql).toContain('organization_id');
+    expect(sql).toContain('kilo_user_id');
+    expect(params).toEqual([PARENT_ORG, CHILD_ORG_A, CTX_USER]);
   });
 
   it('takes precedence over a single organizationId', () => {
-    const { sql, bindings } = scopeSql({
+    const { sql, params } = scopeSql({
       organizationId: CHILD_ORG_B,
       organizationIds: [PARENT_ORG, CHILD_ORG_A],
     });
-    expect(sql).toContain('organization_id IN (?, ?)');
-    expect(bindings).toEqual([PARENT_ORG, CHILD_ORG_A]);
+    expect(sql).toContain('organization_id');
+    expect(params).toEqual([PARENT_ORG, CHILD_ORG_A]);
   });
 
   it('falls back to personal scope with no org', () => {
-    const { sql, bindings } = scopeSql({});
-    expect(sql).toContain('kilo_user_id = ?');
-    expect(sql).toContain('organization_id = ?');
-    // personal-only pins kilo_user_id to caller and org to the empty-string sentinel
-    expect(bindings).toEqual([CTX_USER, '']);
+    const { sql, params } = scopeSql({});
+    expect(sql).toContain('kilo_user_id');
+    expect(sql).toContain('organization_id');
+    // personal-only pins kilo_user_id to caller and organization_id IS NULL
+    expect(sql).toMatch(/is null/i);
+    expect(params).toEqual([CTX_USER]);
+  });
+
+  it('includes org-attributed rows when personalScope is include-orgs', () => {
+    const { sql, params } = scopeSql({ personalScope: 'include-orgs' });
+    expect(sql).toContain('kilo_user_id');
+    expect(sql).not.toContain('organization_id');
+    expect(params).toEqual([CTX_USER]);
   });
 
   it('caps organizationIds at the boundary to bound auth fan-out', () => {
@@ -171,7 +194,7 @@ describe('usage analytics organization breakdown', () => {
     ).toBe(false);
   });
 
-  it('maps the organization dimension to the Snowflake organization column', () => {
-    expect(dimensionColumn('organization')).toBe('organization_id');
+  it('maps the organization dimension to organization_id', () => {
+    expect(dialect.sqlToQuery(dimensionColumn('organization')).sql).toContain('organization_id');
   });
 });