diff --git a/apps/web/src/lib/drizzle.test.ts b/apps/web/src/lib/drizzle.test.ts index 82ed9d423d..d2e1409737 100644 --- a/apps/web/src/lib/drizzle.test.ts +++ b/apps/web/src/lib/drizzle.test.ts @@ -1,10 +1,4 @@ -import { - pool, - db, - selectReplicaUrl, - selectUsageReplicaUrl, - shouldExitOnPoolError, -} from '@/lib/drizzle'; +import { pool, db, selectReplicaUrl, shouldExitOnPoolError } from '@/lib/drizzle'; describe('drizzle', () => { describe('pool', () => { @@ -112,39 +106,6 @@ 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({ diff --git a/apps/web/src/lib/drizzle.ts b/apps/web/src/lib/drizzle.ts index cc3f0abfb7..0136e743db 100644 --- a/apps/web/src/lib/drizzle.ts +++ b/apps/web/src/lib/drizzle.ts @@ -68,48 +68,21 @@ 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 use POSTGRES_REPLICA_EU_URL. POSTGRES_REPLICA_EU_URL_2 is - * reserved for usage-analytics scans so those queries stay off this pool. + * - EU deployments randomly select one of two EU replicas to split read traffic + * across ~2,200 concurrent Vercel instances (~50/50 statistical distribution) * - 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: 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(), + euReplicaUrls: [ + getEnvVariable('POSTGRES_REPLICA_EU_URL'), + getEnvVariable('POSTGRES_REPLICA_EU_URL_2'), + ].filter(Boolean) as string[], }); } @@ -148,25 +121,6 @@ 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') { @@ -174,9 +128,6 @@ if (process.env.NODE_ENV !== 'test') { if (usesSeparateReplica) { attachDatabasePool(replicaPool); } - if (usesDedicatedUsageReplica) { - attachDatabasePool(usageReplicaPool); - } } /** @@ -223,15 +174,6 @@ 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. @@ -254,13 +196,6 @@ 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. @@ -290,9 +225,6 @@ export async function closeAllDrizzleConnections(): Promise { if (usesSeparateReplica) { await replicaPool.end(); } - if (usesDedicatedUsageReplica) { - await usageReplicaPool.end(); - } } export type DrizzleTransaction = Parameters[0]>[0]; diff --git a/apps/web/src/routers/usage-analytics-router.test.ts b/apps/web/src/routers/usage-analytics-router.test.ts index 1bb4a273a2..33e98d1f92 100644 --- a/apps/web/src/routers/usage-analytics-router.test.ts +++ b/apps/web/src/routers/usage-analytics-router.test.ts @@ -1,6 +1,19 @@ jest.mock('@/lib/redis', () => ({ redisClient: {} })); +jest.mock('@/lib/snowflake', () => ({ + resolveSnowflakeConfig: jest.fn(), + executeSnowflakeStatement: jest.fn(), +})); + +import { + executeSnowflakeStatement, + resolveSnowflakeConfig, + type SnowflakeConfig, +} from '@/lib/snowflake'; +import { defineTestUser } from '@/tests/helpers/user.helper'; + +const mockResolveSnowflakeConfig = jest.mocked(resolveSnowflakeConfig); +const mockExecuteSnowflakeStatement = jest.mocked(executeSnowflakeStatement); -import { PgDialect } from 'drizzle-orm/pg-core'; import { BreakdownInputSchema, CostSourceSchema, @@ -13,6 +26,7 @@ import { costColumnFor, costSumExprSql, dimensionColumn, + usageAnalyticsRouter, } from './usage-analytics-router'; const baseFilters = { @@ -25,39 +39,42 @@ const CTX_USER = 'user-1'; 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 SNOWFLAKE_CONFIG = { + accountHost: 'account.snowflakecomputing.com', + jwtAccountIdentifier: 'ACCOUNT', + username: 'user', + role: 'role', + warehouse: 'warehouse', + database: 'database', + schema: 'schema', + privateKeyPem: 'key', + publicKeyFingerprint: 'SHA256:fingerprint', +} satisfies SnowflakeConfig; -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 caller() { + return usageAnalyticsRouter.createCaller({ user: defineTestUser({ id: CTX_USER }) }); } function scopeSql(rawFilters: Record) { const filters = UsageAnalyticsFiltersSchema.parse({ ...baseFilters, ...rawFilters }); const where = new WhereBuilder(); buildScopeConditions(where, filters, CTX_USER); - return compile(where); + return { sql: where.sql(), bindings: where.bindings.map(b => b.value) }; } describe('usage analytics cost source', () => { it('defaults to billable cost for existing clients', () => { expect(UsageAnalyticsFiltersSchema.parse(baseFilters).costSource).toBe('cost'); - 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'); + expect(costColumnFor('cost')).toBe('total_cost_microdollars'); + expect(costSumExprSql('cost')).toBe('COALESCE(SUM(total_cost_microdollars), 0)'); }); - it('uses the estimated market cost when selected', () => { + it('uses the estimated market cost rollup when selected', () => { expect( UsageAnalyticsFiltersSchema.parse({ ...baseFilters, costSource: 'market' }).costSource ).toBe('market'); - 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'); + expect(costColumnFor('market')).toBe('total_market_cost_microdollars'); + expect(costSumExprSql('market')).toBe('COALESCE(SUM(total_market_cost_microdollars), 0)'); }); it('rejects arbitrary cost source values', () => { @@ -69,63 +86,53 @@ describe('usage analytics cost source', () => { describe('usage analytics scope conditions', () => { it('pins a single org to the caller in self view', () => { - 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]); + 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]); }); it('does not pin to the caller in org-wide view', () => { - const { sql, params } = scopeSql({ organizationId: PARENT_ORG, viewAs: 'org-wide' }); - expect(sql).toContain('organization_id'); + const { sql, bindings } = scopeSql({ organizationId: PARENT_ORG, viewAs: 'org-wide' }); + expect(sql).toContain('organization_id = ?'); expect(sql).not.toContain('kilo_user_id'); - expect(params).toEqual([PARENT_ORG]); + expect(bindings).toEqual([PARENT_ORG]); }); it('aggregates org-wide across all orgs when organizationIds is set', () => { - const { sql, params } = scopeSql({ + const { sql, bindings } = scopeSql({ organizationIds: [PARENT_ORG, CHILD_ORG_A, CHILD_ORG_B], }); - expect(sql).toContain('organization_id'); - expect(sql).toMatch(/in/i); + expect(sql).toContain('organization_id IN (?, ?, ?)'); expect(sql).not.toContain('kilo_user_id'); - expect(params).toEqual([PARENT_ORG, CHILD_ORG_A, CHILD_ORG_B]); + expect(bindings).toEqual([PARENT_ORG, CHILD_ORG_A, CHILD_ORG_B]); }); it('honors explicit user filters in the all-orgs aggregate', () => { - const { sql, params } = scopeSql({ + const { sql, bindings } = scopeSql({ organizationIds: [PARENT_ORG, CHILD_ORG_A], userIds: [CTX_USER], }); - expect(sql).toContain('organization_id'); - expect(sql).toContain('kilo_user_id'); - expect(params).toEqual([PARENT_ORG, CHILD_ORG_A, 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]); }); it('takes precedence over a single organizationId', () => { - const { sql, params } = scopeSql({ + const { sql, bindings } = scopeSql({ organizationId: CHILD_ORG_B, organizationIds: [PARENT_ORG, CHILD_ORG_A], }); - expect(sql).toContain('organization_id'); - expect(params).toEqual([PARENT_ORG, CHILD_ORG_A]); + expect(sql).toContain('organization_id IN (?, ?)'); + expect(bindings).toEqual([PARENT_ORG, CHILD_ORG_A]); }); it('falls back to personal scope with no org', () => { - 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]); + 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, '']); }); it('caps organizationIds at the boundary to bound auth fan-out', () => { @@ -194,7 +201,171 @@ describe('usage analytics organization breakdown', () => { ).toBe(false); }); - it('maps the organization dimension to organization_id', () => { - expect(dialect.sqlToQuery(dimensionColumn('organization')).sql).toContain('organization_id'); + it('maps the organization dimension to the Snowflake organization column', () => { + expect(dimensionColumn('organization')).toBe('organization_id'); + }); +}); + +describe('usage analytics procedures', () => { + beforeEach(() => { + jest.clearAllMocks(); + mockResolveSnowflakeConfig.mockReturnValue(SNOWFLAKE_CONFIG); + mockExecuteSnowflakeStatement.mockResolvedValue([]); + }); + + it('keeps include-orgs personal usage scoped to the caller without the personal sentinel', async () => { + mockExecuteSnowflakeStatement.mockResolvedValue([ + ['100', '2', '3', '4', '5', '6', '0', '0', '0', '7', '8', '9', '2', '1', '7', '1'], + ]); + + await expect( + caller().getSummary({ ...baseFilters, personalScope: 'include-orgs', features: ['chat'] }) + ).resolves.toMatchObject({ + costMicrodollars: 100, + byokRequestCount: 7, + effectiveGranularity: 'day', + }); + + const request = mockExecuteSnowflakeStatement.mock.calls[0][0]; + expect(request.statement).toContain('COALESCE(SUM(user_byok_request_count), 0)'); + expect(request.statement).toContain('MICRODOLLAR_USAGE_DAILY'); + expect(request.bindings).toEqual([ + { type: 'TEXT', value: '2026-06-04' }, + { type: 'TEXT', value: '2026-06-05' }, + { type: 'TEXT', value: CTX_USER }, + { type: 'TEXT', value: 'chat' }, + ]); + }); + + it('uses the user BYOK rollup in the hourly summary tier', async () => { + const endDate = new Date().toISOString(); + const startDate = new Date(Date.now() - 60 * 60 * 1000).toISOString(); + + await caller().getSummary({ ...baseFilters, startDate, endDate, granularity: 'hour' }); + + expect(mockExecuteSnowflakeStatement.mock.calls[0][0].statement).toContain( + 'COALESCE(SUM(user_byok_request_count), 0)' + ); + expect(mockExecuteSnowflakeStatement.mock.calls[0][0].statement).toContain( + 'MICRODOLLAR_USAGE_HOURLY' + ); + }); + + it('maps timeseries, breakdown, and table Snowflake rows', async () => { + mockExecuteSnowflakeStatement + .mockResolvedValueOnce([['2026-06-04', '12', 'model-a']]) + .mockResolvedValueOnce([ + ['model-a', '3'], + ['model-b', '1'], + ]) + .mockResolvedValueOnce([ + ['2026-06-04', 'chat', '', '', '', '', '', '10', '2', '3', '4', '5', '6', '7'], + ]); + + await expect( + caller().getTimeseries({ ...baseFilters, metric: 'requests', splitBy: 'model' }) + ).resolves.toEqual({ + timeseries: [{ datetime: '2026-06-04', value: 12, label: 'model-a' }], + effectiveGranularity: 'day', + }); + await expect( + caller().getBreakdown({ ...baseFilters, dimension: 'model', metric: 'requests' }) + ).resolves.toEqual({ + breakdown: [ + { key: 'model-a', label: 'model-a', value: 3, percentage: 75 }, + { key: 'model-b', label: 'model-b', value: 1, percentage: 25 }, + ], + totalValue: 4, + effectiveGranularity: 'day', + }); + await expect(caller().getTable({ ...baseFilters, groupBy: ['feature'] })).resolves.toEqual({ + rows: [ + { + datetime: '2026-06-04', + dimensions: { feature: 'chat' }, + costMicrodollars: 10, + requestCount: 2, + inputTokens: 3, + outputTokens: 4, + cacheWriteTokens: 5, + cacheHitTokens: 6, + errorCount: 7, + }, + ], + effectiveGranularity: 'day', + }); + + expect(mockExecuteSnowflakeStatement.mock.calls[0][0].statement).toContain('GROUP BY 1, 3'); + expect(mockExecuteSnowflakeStatement.mock.calls[1][0].statement).toContain('GROUP BY 1'); + expect(mockExecuteSnowflakeStatement.mock.calls[2][0].statement).toContain('dim_feature'); + }); + + it('allows successful empty results for every Snowflake procedure', async () => { + await expect(caller().getSummary(baseFilters)).resolves.toMatchObject({ requestCount: 0 }); + await expect(caller().getTimeseries({ ...baseFilters, metric: 'cost' })).resolves.toMatchObject( + { + timeseries: [], + } + ); + await expect( + caller().getBreakdown({ ...baseFilters, dimension: 'model', metric: 'cost' }) + ).resolves.toMatchObject({ breakdown: [], totalValue: 0 }); + await expect(caller().getTable({ ...baseFilters, groupBy: [] })).resolves.toMatchObject({ + rows: [], + }); + }); + + it.each([ + () => caller().getSummary(baseFilters), + () => caller().getTimeseries({ ...baseFilters, metric: 'cost' }), + () => caller().getBreakdown({ ...baseFilters, dimension: 'model', metric: 'cost' }), + () => caller().getTable({ ...baseFilters, groupBy: [] }), + ])('rejects missing Snowflake configuration with a sanitized error', async invoke => { + mockResolveSnowflakeConfig.mockReturnValue(null); + + await expect(invoke()).rejects.toMatchObject({ + code: 'INTERNAL_SERVER_ERROR', + message: 'Usage data temporarily unavailable', + }); + expect(mockExecuteSnowflakeStatement).not.toHaveBeenCalled(); + }); + + it.each([ + () => caller().getSummary(baseFilters), + () => caller().getTimeseries({ ...baseFilters, metric: 'cost' }), + () => caller().getBreakdown({ ...baseFilters, dimension: 'model', metric: 'cost' }), + () => caller().getTable({ ...baseFilters, groupBy: [] }), + ])('sanitizes Snowflake query failures in responses and logs', async invoke => { + mockExecuteSnowflakeStatement.mockRejectedValue(new Error('upstream response body')); + const errorLog = jest.spyOn(console, 'error').mockImplementation(() => {}); + + try { + await expect(invoke()).rejects.toMatchObject({ + code: 'INTERNAL_SERVER_ERROR', + message: 'Usage data temporarily unavailable', + }); + expect(mockExecuteSnowflakeStatement).toHaveBeenCalledTimes(1); + expect(errorLog).toHaveBeenCalledTimes(1); + expect(errorLog).toHaveBeenCalledWith(expect.stringContaining('"reason":"query_failed"')); + expect(errorLog).not.toHaveBeenCalledWith(expect.stringContaining('upstream response body')); + } finally { + errorLog.mockRestore(); + } + }); + + it.each([ + () => caller().getSummary({ ...baseFilters, userIds: ['other-user'] }), + () => caller().getTimeseries({ ...baseFilters, metric: 'cost', userIds: ['other-user'] }), + () => + caller().getBreakdown({ + ...baseFilters, + dimension: 'model', + metric: 'cost', + userIds: ['other-user'], + }), + () => caller().getTable({ ...baseFilters, groupBy: [], userIds: ['other-user'] }), + ])('rejects another personal user before querying Snowflake', async invoke => { + await expect(invoke()).rejects.toMatchObject({ code: 'FORBIDDEN' }); + expect(mockExecuteSnowflakeStatement).not.toHaveBeenCalled(); }); }); diff --git a/apps/web/src/routers/usage-analytics-router.ts b/apps/web/src/routers/usage-analytics-router.ts index a183e26250..8cc2e6fa33 100644 --- a/apps/web/src/routers/usage-analytics-router.ts +++ b/apps/web/src/routers/usage-analytics-router.ts @@ -1,15 +1,16 @@ import { TRPCError } from '@trpc/server'; import * as z from 'zod'; -import { and, asc, eq, gte, inArray, isNull, lt, notInArray, or, sql, type SQL } from 'drizzle-orm'; +import { and, asc, eq, inArray, isNull, or } from 'drizzle-orm'; import { baseProcedure, createTRPCRouter, type TRPCContext } from '@/lib/trpc/init'; -import { readDb, usageReadDb } from '@/lib/drizzle'; -import { timedUsageQuery } from '@/lib/usage-query'; +import { readDb } from '@/lib/drizzle'; +import { getEnvVariable } from '@/lib/dotenvx'; +import { + executeSnowflakeStatement, + resolveSnowflakeConfig, + type SnowflakeBinding, +} from '@/lib/snowflake'; import { - feature, kilocode_users, - microdollar_usage, - microdollar_usage_metadata, - mode, organization_memberships, organizations, user_auth_provider, @@ -90,34 +91,118 @@ function resolveTier(granularity: Granularity, startDate: string): TableMeta { if (ageDays < 8) { return { tier: 'hourly', effectiveGranularity: 'hour' }; } - // Auto-downgrade: hourly buckets are only used for the past 7 days. + // Auto-downgrade: hourly data is only available for the past 7 days. return { tier: 'daily', effectiveGranularity: 'day' }; } if (granularity === 'day' || granularity === 'week') { + // MICRODOLLAR_USAGE_DAILY holds full history — no age-based downgrade needed. return { tier: 'daily', effectiveGranularity: granularity }; } return { tier: 'monthly', effectiveGranularity: 'month' }; } +/** + * Returns the Snowflake table name for a given tier. + * Both daily and monthly tiers use MICRODOLLAR_USAGE_DAILY; monthly queries + * add a DATE_TRUNC('MONTH', usage_day) bucket expression on top. + */ +function getTableName(tier: GranularityTier): string { + return tier === 'hourly' ? 'MICRODOLLAR_USAGE_HOURLY' : 'MICRODOLLAR_USAGE_DAILY'; +} + +/** The column that holds the time value for a given tier. */ +function getTimeColumn(tier: GranularityTier): string { + return tier === 'hourly' ? 'usage_hour' : 'usage_date'; +} + +// --------------------------------------------------------------------------- +// Date helpers +// --------------------------------------------------------------------------- + +function ceilIsoToUtcDayExclusive(iso: string): string { + const d = new Date(iso); + const dayStartMs = Date.UTC(d.getUTCFullYear(), d.getUTCMonth(), d.getUTCDate()); + if (d.getTime() === dayStartMs) { + return iso.slice(0, 10); + } + return new Date(dayStartMs + 86_400_000).toISOString().slice(0, 10); +} + +function floorIsoToUtcMonth(iso: string): string { + return `${iso.slice(0, 7)}-01`; +} + +function ceilIsoToUtcMonthExclusive(iso: string): string { + const d = new Date(iso); + const firstOfMonthMs = Date.UTC(d.getUTCFullYear(), d.getUTCMonth(), 1); + if (d.getTime() === firstOfMonthMs) { + return iso.slice(0, 10); + } + const next = new Date(Date.UTC(d.getUTCFullYear(), d.getUTCMonth() + 1, 1)); + return next.toISOString().slice(0, 10); +} + // --------------------------------------------------------------------------- // SQL WHERE clause builder // --------------------------------------------------------------------------- /** - * Accumulates SQL WHERE clauses. Callers push conditions in any order; - * `toSQL()` joins them with AND. + * Accumulates SQL WHERE clauses with positional `?` bindings. + * Callers push conditions in any order; `sql()` joins them with AND. */ export class WhereBuilder { - readonly conditions: SQL[] = []; + readonly clauses: string[] = []; + readonly bindings: SnowflakeBinding[] = []; - add(condition: SQL): void { - this.conditions.push(condition); + private push(clause: string, ...bindings: SnowflakeBinding[]) { + bindings.forEach(b => this.bindings.push(b)); + this.clauses.push(clause); } - toSQL(): SQL | undefined { - return this.conditions.length > 0 ? and(...this.conditions) : undefined; + addTimestampRange(column: string, gte: string, lt: string): void { + this.push( + `${column} >= ? AND ${column} < ?`, + { type: 'TEXT', value: gte }, + { type: 'TEXT', value: lt } + ); + } + + addDateRange(column: string, gte: string, lt: string): void { + this.push( + `${column} >= ? AND ${column} < ?`, + { type: 'TEXT', value: gte }, + { type: 'TEXT', value: lt } + ); + } + + addEq(column: string, value: string): void { + this.push(`${column} = ?`, { type: 'TEXT', value }); + } + + addIsNull(column: string): void { + this.clauses.push(`${column} IS NULL`); + } + + addIn(column: string, values: string[]): void { + const placeholders = values.map(() => '?').join(', '); + this.push( + `${column} IN (${placeholders})`, + ...values.map(v => ({ type: 'TEXT' as const, value: v })) + ); + } + + addNotIn(column: string, values: string[]): void { + const placeholders = values.map(() => '?').join(', '); + this.push( + `${column} NOT IN (${placeholders})`, + ...values.map(v => ({ type: 'TEXT' as const, value: v })) + ); + } + + sql(): string { + return this.clauses.length > 0 ? this.clauses.join('\n AND ') : '1=1'; } } @@ -173,9 +258,29 @@ async function ensureScopeAccess(ctx: TRPCContext, filters: UsageAnalyticsFilter // WHERE clause helpers // --------------------------------------------------------------------------- -function buildDateConditions(where: WhereBuilder, filters: UsageAnalyticsFilters): void { - where.add(gte(microdollar_usage.created_at, filters.startDate)); - where.add(lt(microdollar_usage.created_at, filters.endDate)); +function buildDateConditions( + where: WhereBuilder, + tier: GranularityTier, + filters: UsageAnalyticsFilters +): void { + const timeCol = getTimeColumn(tier); + + if (tier === 'hourly') { + where.addTimestampRange(timeCol, filters.startDate, filters.endDate); + } else if (tier === 'daily') { + where.addDateRange( + timeCol, + filters.startDate.slice(0, 10), + ceilIsoToUtcDayExclusive(filters.endDate) + ); + } else { + // monthly — daily table, filter by day boundaries aligned to month + where.addDateRange( + timeCol, + floorIsoToUtcMonth(filters.startDate), + ceilIsoToUtcMonthExclusive(filters.endDate) + ); + } } export function buildScopeConditions( @@ -186,83 +291,65 @@ export function buildScopeConditions( if (filters.organizationIds && filters.organizationIds.length > 0) { // Aggregate across the parent org and its children. Always org-wide, so // honor any explicit user include/exclude filters but never pin to self. - where.add(inArray(microdollar_usage.organization_id, filters.organizationIds)); + where.addIn('organization_id', filters.organizationIds); if (filters.userIds && filters.userIds.length > 0) { - where.add(inArray(microdollar_usage.kilo_user_id, filters.userIds)); + where.addIn('kilo_user_id', filters.userIds); } if (filters.excludedUserIds && filters.excludedUserIds.length > 0) { - where.add(notInArray(microdollar_usage.kilo_user_id, filters.excludedUserIds)); + where.addNotIn('kilo_user_id', filters.excludedUserIds); } return; } if (filters.organizationId) { - where.add(eq(microdollar_usage.organization_id, filters.organizationId)); + where.addEq('organization_id', filters.organizationId); if (filters.viewAs === 'self') { - where.add(eq(microdollar_usage.kilo_user_id, ctxUserId)); + where.addEq('kilo_user_id', ctxUserId); } else { if (filters.userIds && filters.userIds.length > 0) { - where.add(inArray(microdollar_usage.kilo_user_id, filters.userIds)); + where.addIn('kilo_user_id', filters.userIds); } if (filters.excludedUserIds && filters.excludedUserIds.length > 0) { - where.add(notInArray(microdollar_usage.kilo_user_id, filters.excludedUserIds)); + where.addNotIn('kilo_user_id', filters.excludedUserIds); } } } else { - where.add(eq(microdollar_usage.kilo_user_id, ctxUserId)); + where.addEq('kilo_user_id', ctxUserId); if (filters.personalScope === 'personal-only') { - // Personal usage is stored with a NULL organization_id. - where.add(isNull(microdollar_usage.organization_id)); + // DBT coalesces personal Snowflake usage rollups to an empty-string sentinel + // so incremental merges can match on organization_id. + where.addEq('organization_id', ''); } } } -const featureName: SQL = sql`COALESCE(${feature.feature}, '')`; -const modeName: SQL = sql`COALESCE(${mode.mode}, '')`; -const modelName: SQL = sql`COALESCE(${microdollar_usage.model}, '')`; -const providerName: SQL = sql`COALESCE(${microdollar_usage.provider}, '')`; -const projectName: SQL = sql`COALESCE(${microdollar_usage.project_id}, '')`; - -function inValues(column: SQL, values: string[]): SQL { - return sql`${column} IN (${sql.join( - values.map(value => sql`${value}`), - sql`, ` - )})`; -} - -function notInValues(column: SQL, values: string[]): SQL { - return sql`${column} NOT IN (${sql.join( - values.map(value => sql`${value}`), - sql`, ` - )})`; -} - function buildDimensionConditions(where: WhereBuilder, filters: UsageAnalyticsFilters): void { - const addInIfNonEmpty = (column: SQL, values: string[] | undefined) => { - if (values && values.length > 0) where.add(inValues(column, values)); + const addInIfNonEmpty = (column: string, values: string[] | undefined) => { + if (values && values.length > 0) where.addIn(column, values); }; - const addNotInIfNonEmpty = (column: SQL, values: string[] | undefined) => { - if (values && values.length > 0) where.add(notInValues(column, values)); + const addNotInIfNonEmpty = (column: string, values: string[] | undefined) => { + if (values && values.length > 0) where.addNotIn(column, values); }; - addInIfNonEmpty(featureName, filters.features); - addInIfNonEmpty(modelName, filters.models); - addInIfNonEmpty(modeName, filters.modes); - addInIfNonEmpty(providerName, filters.providers); - addInIfNonEmpty(projectName, filters.projects); - addNotInIfNonEmpty(featureName, filters.excludedFeatures); - addNotInIfNonEmpty(modelName, filters.excludedModels); - addNotInIfNonEmpty(modeName, filters.excludedModes); - addNotInIfNonEmpty(providerName, filters.excludedProviders); - addNotInIfNonEmpty(projectName, filters.excludedProjects); + addInIfNonEmpty('feature', filters.features); + addInIfNonEmpty('model', filters.models); + addInIfNonEmpty('mode', filters.modes); + addInIfNonEmpty('provider', filters.providers); + addInIfNonEmpty('project_id', filters.projects); + addNotInIfNonEmpty('feature', filters.excludedFeatures); + addNotInIfNonEmpty('model', filters.excludedModels); + addNotInIfNonEmpty('mode', filters.excludedModes); + addNotInIfNonEmpty('provider', filters.excludedProviders); + addNotInIfNonEmpty('project_id', filters.excludedProjects); } function buildWhereClause( + tier: GranularityTier, filters: UsageAnalyticsFilters, ctxUserId: string, includeDimensions: boolean ): WhereBuilder { const where = new WhereBuilder(); - buildDateConditions(where, filters); + buildDateConditions(where, tier, filters); buildScopeConditions(where, filters, ctxUserId); if (includeDimensions) { buildDimensionConditions(where, filters); @@ -274,64 +361,61 @@ function buildWhereClause( // Metric SQL expression // --------------------------------------------------------------------------- -export function costColumnFor(costSource: CostSource): SQL { +export function costColumnFor(costSource: CostSource): string { switch (costSource) { case 'cost': - return sql`${microdollar_usage.cost}`; + return 'total_cost_microdollars'; case 'market': - return sql`COALESCE(${microdollar_usage_metadata.market_cost}, 0)`; + return 'total_market_cost_microdollars'; } } -export function costSumExprSql(costSource: CostSource): SQL { - return sql`COALESCE(SUM(${costColumnFor(costSource)}), 0)`; +export function costSumExprSql(costSource: CostSource): string { + return `COALESCE(SUM(${costColumnFor(costSource)}), 0)`; } -const requestCountExpr = sql`COUNT(*)`; -const inputTokensExpr = sql`COALESCE(SUM(${microdollar_usage.input_tokens}), 0)`; -const outputTokensExpr = sql`COALESCE(SUM(${microdollar_usage.output_tokens}), 0)`; -const cacheWriteTokensExpr = sql`COALESCE(SUM(${microdollar_usage.cache_write_tokens}), 0)`; -const cacheHitTokensExpr = sql`COALESCE(SUM(${microdollar_usage.cache_hit_tokens}), 0)`; -const totalTokensExpr = sql`COALESCE(SUM(${microdollar_usage.input_tokens} + ${microdollar_usage.output_tokens} + ${microdollar_usage.cache_write_tokens} + ${microdollar_usage.cache_hit_tokens}), 0)`; -const errorCountExpr = sql`COUNT(*) FILTER (WHERE ${microdollar_usage.has_error})`; -const cancelledCountExpr = sql`COUNT(*) FILTER (WHERE ${microdollar_usage_metadata.cancelled})`; -const freeRequestCountExpr = sql`COUNT(*) FILTER (WHERE ${microdollar_usage_metadata.is_free})`; -const byokRequestCountExpr = sql`COUNT(*) FILTER (WHERE ${microdollar_usage_metadata.is_user_byok})`; -const totalLatencyMsExpr = sql`COALESCE(SUM(${microdollar_usage_metadata.latency}), 0)`; -const latencyCountExpr = sql`COUNT(${microdollar_usage_metadata.latency})`; -const totalGenerationTimeMsExpr = sql`COALESCE(SUM(${microdollar_usage_metadata.generation_time}), 0)`; -const generationTimeCountExpr = sql`COUNT(${microdollar_usage_metadata.generation_time})`; - -function metricExprSql(metric: Metric, costSource: CostSource): SQL { +function metricExprSql(metric: Metric, tier: GranularityTier, costSource: CostSource): string { const costSumExpr = costSumExprSql(costSource); switch (metric) { case 'cost': return costSumExpr; case 'requests': - return sql`COALESCE(${requestCountExpr}, 0)`; + return 'COALESCE(SUM(request_count), 0)'; case 'inputTokens': - return inputTokensExpr; + return 'COALESCE(SUM(total_input_tokens), 0)'; case 'outputTokens': - return outputTokensExpr; + return 'COALESCE(SUM(total_output_tokens), 0)'; case 'tokens': - return totalTokensExpr; + return 'COALESCE(SUM(total_tokens), 0)'; case 'errorRate': - return sql`CASE WHEN ${requestCountExpr} = 0 THEN 0 ELSE (${errorCountExpr})::FLOAT / (${requestCountExpr})::FLOAT END`; + return 'CASE WHEN COALESCE(SUM(request_count), 0) = 0 THEN 0 ELSE COALESCE(SUM(error_count), 0)::FLOAT / SUM(request_count)::FLOAT END'; case 'avgLatencyMs': - return sql`CASE WHEN ${latencyCountExpr} = 0 THEN 0 ELSE (${totalLatencyMsExpr})::FLOAT / (${latencyCountExpr})::FLOAT END`; - case 'avgGenerationTimeMs': - return sql`CASE WHEN ${generationTimeCountExpr} = 0 THEN 0 ELSE (${totalGenerationTimeMsExpr})::FLOAT / (${generationTimeCountExpr})::FLOAT END`; + return 'CASE WHEN COALESCE(SUM(latency_count), 0) = 0 THEN 0 ELSE COALESCE(SUM(total_latency_ms), 0)::FLOAT / SUM(latency_count)::FLOAT END'; + case 'avgGenerationTimeMs': { + const countExpr = generationTimeCountExprSql(tier); + return `CASE WHEN COALESCE(SUM(${countExpr}), 0) = 0 THEN 0 ELSE COALESCE(SUM(total_generation_time_ms), 0)::FLOAT / SUM(${countExpr})::FLOAT END`; + } case 'costPerRequest': - return sql`CASE WHEN ${requestCountExpr} = 0 THEN 0 ELSE (${costSumExpr})::FLOAT / (${requestCountExpr})::FLOAT END`; + return `CASE WHEN COALESCE(SUM(request_count), 0) = 0 THEN 0 ELSE ${costSumExpr}::FLOAT / SUM(request_count)::FLOAT END`; case 'tokensPerRequest': - return sql`CASE WHEN ${requestCountExpr} = 0 THEN 0 ELSE (${totalTokensExpr})::FLOAT / (${requestCountExpr})::FLOAT END`; + return 'CASE WHEN COALESCE(SUM(request_count), 0) = 0 THEN 0 ELSE COALESCE(SUM(total_tokens), 0)::FLOAT / SUM(request_count)::FLOAT END'; case 'cacheHitRatio': - return sql`CASE WHEN COALESCE(SUM(${microdollar_usage.input_tokens} + ${microdollar_usage.cache_hit_tokens}), 0) = 0 THEN 0 ELSE COALESCE(SUM(${microdollar_usage.cache_hit_tokens}), 0)::FLOAT / SUM(${microdollar_usage.input_tokens} + ${microdollar_usage.cache_hit_tokens})::FLOAT END`; + return 'CASE WHEN COALESCE(SUM(total_input_tokens + total_cache_hit_tokens), 0) = 0 THEN 0 ELSE COALESCE(SUM(total_cache_hit_tokens), 0)::FLOAT / SUM(total_input_tokens + total_cache_hit_tokens)::FLOAT END'; case 'outputInputRatio': - return sql`CASE WHEN COALESCE(SUM(${microdollar_usage.input_tokens}), 0) = 0 THEN 0 ELSE COALESCE(SUM(${microdollar_usage.output_tokens}), 0)::FLOAT / SUM(${microdollar_usage.input_tokens})::FLOAT END`; + return 'CASE WHEN COALESCE(SUM(total_input_tokens), 0) = 0 THEN 0 ELSE COALESCE(SUM(total_output_tokens), 0)::FLOAT / SUM(total_input_tokens)::FLOAT END'; } } +function generationTimeCountExprSql(tier: GranularityTier): string { + if (tier === 'hourly') { + return 'IFF(total_generation_time_ms IS NOT NULL, 1, 0)'; + } + // Daily rollups do not currently carry a generation-time observation count. + // Derive one only for the window backed by hourly rollups so older daily + // history does not reuse latency_count as an incorrect denominator. + return 'IFF(total_generation_time_ms IS NOT NULL AND usage_date >= DATEADD(day, -7, CURRENT_DATE), 1, 0)'; +} + // --------------------------------------------------------------------------- // Bucket expression for timeseries / table grouping // --------------------------------------------------------------------------- @@ -343,49 +427,124 @@ function metricExprSql(metric: Metric, costSource: CostSource): SQL { * Hourly → 'YYYY-MM-DD HH24:MI:SS' (matches what Postgres timestamp::text returns) * Day → 'YYYY-MM-DD' * Week → 'YYYY-MM-DD' of the Monday-aligned week start - * Month → 'YYYY-MM-DD' of the first of the month + * Month → 'YYYY-MM-DD' of the first of the month (from daily table) */ -function bucketExprSql(effectiveGranularity: Granularity): SQL { - const createdAtUtc = sql`${microdollar_usage.created_at} AT TIME ZONE 'UTC'`; +function bucketExprSql(effectiveGranularity: Granularity, tier: GranularityTier): string { + const timeCol = getTimeColumn(tier); + if (effectiveGranularity === 'hour') { - return sql`TO_CHAR(DATE_TRUNC('hour', ${createdAtUtc}), 'YYYY-MM-DD HH24:MI:SS')`; + return `TO_VARCHAR(${timeCol}, 'YYYY-MM-DD HH24:MI:SS')`; } if (effectiveGranularity === 'week') { - return sql`TO_CHAR(DATE_TRUNC('week', ${createdAtUtc}), 'YYYY-MM-DD')`; + return `TO_VARCHAR(DATE_TRUNC('WEEK', ${timeCol}), 'YYYY-MM-DD')`; } if (effectiveGranularity === 'month') { - return sql`TO_CHAR(DATE_TRUNC('month', ${createdAtUtc}), 'YYYY-MM-DD')`; + // Daily table, group by month + return `TO_VARCHAR(DATE_TRUNC('MONTH', ${timeCol}), 'YYYY-MM-DD')`; } // 'day' - return sql`TO_CHAR(DATE_TRUNC('day', ${createdAtUtc}), 'YYYY-MM-DD')`; + return `TO_VARCHAR(${timeCol}, 'YYYY-MM-DD')`; } // --------------------------------------------------------------------------- // Dimension column name // --------------------------------------------------------------------------- -export function dimensionColumn(dimension: BreakdownDimension): SQL { +export function dimensionColumn(dimension: BreakdownDimension): string { switch (dimension) { case 'feature': - return featureName; + return 'feature'; case 'model': - return modelName; + return 'model'; case 'mode': - return modeName; + return 'mode'; case 'user': - return sql`${microdollar_usage.kilo_user_id}`; + return 'kilo_user_id'; case 'provider': - return providerName; + return 'provider'; case 'project': - return projectName; + return 'project_id'; case 'organization': - return sql`COALESCE(${microdollar_usage.organization_id}::text, '')`; + return 'organization_id'; } } -const usageMetadataJoin = eq(microdollar_usage.id, microdollar_usage_metadata.id); -const usageFeatureJoin = eq(microdollar_usage_metadata.feature_id, feature.feature_id); -const usageModeJoin = eq(microdollar_usage_metadata.mode_id, mode.mode_id); +// --------------------------------------------------------------------------- +// Timed query wrapper +// --------------------------------------------------------------------------- + +function parseTimeoutEnv(envKey: string, fallback: number): number { + const raw = getEnvVariable(envKey); + if (!raw) return fallback; + const parsed = Number(raw); + if (!Number.isFinite(parsed) || parsed < 0) return fallback; + return parsed; +} + +function defaultTimeoutForScope(scope: 'user' | 'org' | 'admin'): number { + if (scope === 'admin') return parseTimeoutEnv('USAGE_QUERY_TIMEOUT_ADMIN_MS', 20_000); + if (scope === 'org') return parseTimeoutEnv('USAGE_QUERY_TIMEOUT_ORG_MS', 10_000); + return parseTimeoutEnv('USAGE_QUERY_TIMEOUT_USER_MS', 5_000); +} + +async function timedSnowflakeQuery( + params: { + route: string; + queryLabel: string; + scope: 'user' | 'org' | 'admin'; + period: string | null; + timeoutMs?: number; + }, + queryFn: (signal: AbortSignal) => Promise +): Promise { + const timeoutMs = params.timeoutMs ?? defaultTimeoutForScope(params.scope); + const start = performance.now(); + let rowCount = 0; + + const controller = new AbortController(); + let settled = false; + const timer = setTimeout(() => { + if (!settled) controller.abort(); + }, timeoutMs); + + try { + const result = await queryFn(controller.signal); + settled = true; + rowCount = Array.isArray(result) ? result.length : 1; + return result; + } catch { + settled = true; + console.error( + JSON.stringify({ + type: 'usage_query_error', + route: params.route, + queryLabel: params.queryLabel, + scope: params.scope, + period: params.period, + reason: controller.signal.aborted ? 'timeout' : 'query_failed', + }) + ); + throw new TRPCError({ + code: 'INTERNAL_SERVER_ERROR', + message: 'Usage data temporarily unavailable', + }); + } finally { + clearTimeout(timer); + const durationMs = Math.round((performance.now() - start) * 100) / 100; + console.log( + JSON.stringify({ + type: 'usage_query', + route: params.route, + queryLabel: params.queryLabel, + scope: params.scope, + period: params.period, + durationMs, + rowCount, + timeoutMs, + }) + ); + } +} // --------------------------------------------------------------------------- // getSummary @@ -397,7 +556,7 @@ function ratioSafe(numerator: number, denominator: number): number { } /** - * Convert an aggregate value (often returned as a string by Postgres) to a + * Convert an aggregate value (often returned as a string by Snowflake) to a * JS number. Values above `MAX_SAFE_INTEGER` are logged as a warning but still * returned so the UI does not crash. */ @@ -484,12 +643,15 @@ function legacyOAuthProviderKey(provider: AuthProviderId, providerAccountId: str return `${provider}:${providerAccountId}`; } -function queryScope(input: UsageAnalyticsFilters): 'org' | 'user' { - return isOrgScope(input) ? 'org' : 'user'; -} - -function queryPeriod(input: UsageAnalyticsFilters): string { - return `${input.startDate}/${input.endDate}`; +function requireSnowflakeConfig() { + const config = resolveSnowflakeConfig(); + if (!config) { + throw new TRPCError({ + code: 'INTERNAL_SERVER_ERROR', + message: 'Usage data temporarily unavailable', + }); + } + return config; } // --------------------------------------------------------------------------- @@ -503,61 +665,72 @@ export const usageAnalyticsRouter = createTRPCRouter({ .query(async ({ input, ctx }): Promise => { await ensureScopeAccess(ctx, input); + const config = requireSnowflakeConfig(); const meta = resolveTier(input.granularity, input.startDate); - const where = buildWhereClause(input, ctx.user.id, true); - - const rows = await timedUsageQuery( + const table = getTableName(meta.tier); + const where = buildWhereClause(meta.tier, input, ctx.user.id, true); + const generationTimeCountExpr = generationTimeCountExprSql(meta.tier); + const costSumExpr = costSumExprSql(input.costSource); + + const statement = ` + SELECT + ${costSumExpr}, + COALESCE(SUM(request_count), 0), + COALESCE(SUM(total_input_tokens), 0), + COALESCE(SUM(total_output_tokens), 0), + COALESCE(SUM(total_cache_write_tokens), 0), + COALESCE(SUM(total_cache_hit_tokens), 0), + COALESCE(SUM(error_count), 0), + COALESCE(SUM(cancelled_count), 0), + COALESCE(SUM(free_request_count), 0), + COALESCE(SUM(user_byok_request_count), 0), + COALESCE(SUM(total_latency_ms), 0), + COALESCE(SUM(total_generation_time_ms), 0), + COALESCE(SUM(latency_count), 0), + COALESCE(SUM(${generationTimeCountExpr}), 0), + COALESCE(SUM(total_tokens), 0), + COUNT(DISTINCT kilo_user_id) + FROM ${table} + WHERE ${where.sql()} + `; + + const rows = await timedSnowflakeQuery( { - db: usageReadDb, route: 'usageAnalytics.getSummary', queryLabel: `summary_${meta.tier}`, - scope: queryScope(input), - period: queryPeriod(input), + scope: isOrgScope(input) ? 'org' : 'user', + period: `${input.startDate}/${input.endDate}`, }, - tx => - tx - .select({ - costMicrodollars: costSumExprSql(input.costSource), - requestCount: requestCountExpr, - inputTokens: inputTokensExpr, - outputTokens: outputTokensExpr, - cacheWriteTokens: cacheWriteTokensExpr, - cacheHitTokens: cacheHitTokensExpr, - errorCount: errorCountExpr, - cancelledCount: cancelledCountExpr, - freeRequestCount: freeRequestCountExpr, - byokRequestCount: byokRequestCountExpr, - totalLatencyMs: totalLatencyMsExpr, - totalGenerationTimeMs: totalGenerationTimeMsExpr, - latencyCount: latencyCountExpr, - generationTimeCount: generationTimeCountExpr, - totalTokens: totalTokensExpr, - distinctUsers: sql`COUNT(DISTINCT ${microdollar_usage.kilo_user_id})`, - }) - .from(microdollar_usage) - .leftJoin(microdollar_usage_metadata, usageMetadataJoin) - .leftJoin(feature, usageFeatureJoin) - .leftJoin(mode, usageModeJoin) - .where(where.toSQL()) + signal => + executeSnowflakeStatement({ + config, + statement, + bindings: where.bindings, + timeoutSeconds: Math.ceil( + defaultTimeoutForScope(isOrgScope(input) ? 'org' : 'user') / 1000 + ), + signal, + }) ); - const row = rows[0]; - const costMicrodollars = toSafeNumber(row?.costMicrodollars); - const requestCount = toSafeNumber(row?.requestCount); - const inputTokens = toSafeNumber(row?.inputTokens); - const outputTokens = toSafeNumber(row?.outputTokens); - const cacheWriteTokens = toSafeNumber(row?.cacheWriteTokens); - const cacheHitTokens = toSafeNumber(row?.cacheHitTokens); - const errorCount = toSafeNumber(row?.errorCount); - const cancelledCount = toSafeNumber(row?.cancelledCount); - const freeRequestCount = toSafeNumber(row?.freeRequestCount); - const byokRequestCount = toSafeNumber(row?.byokRequestCount); - const totalLatencyMs = toSafeNumber(row?.totalLatencyMs); - const totalGenerationTimeMs = toSafeNumber(row?.totalGenerationTimeMs); - const latencyCount = toSafeNumber(row?.latencyCount); - const generationTimeCount = toSafeNumber(row?.generationTimeCount); - const totalTokens = toSafeNumber(row?.totalTokens); - const distinctUsers = toSafeNumber(row?.distinctUsers); + const row = rows[0] ?? []; + + const costMicrodollars = toSafeNumber(row[0]); + const requestCount = toSafeNumber(row[1]); + const inputTokens = toSafeNumber(row[2]); + const outputTokens = toSafeNumber(row[3]); + const cacheWriteTokens = toSafeNumber(row[4]); + const cacheHitTokens = toSafeNumber(row[5]); + const errorCount = toSafeNumber(row[6]); + const cancelledCount = toSafeNumber(row[7]); + const freeRequestCount = toSafeNumber(row[8]); + const byokRequestCount = toSafeNumber(row[9]); + const totalLatencyMs = toSafeNumber(row[10]); + const totalGenerationTimeMs = toSafeNumber(row[11]); + const latencyCount = toSafeNumber(row[12]); + const generationTimeCount = toSafeNumber(row[13]); + const totalTokens = toSafeNumber(row[14]); + const distinctUsers = toSafeNumber(row[15]); return { costMicrodollars, @@ -593,44 +766,62 @@ export const usageAnalyticsRouter = createTRPCRouter({ .query(async ({ input, ctx }) => { await ensureScopeAccess(ctx, input); + const config = requireSnowflakeConfig(); const meta = resolveTier(input.granularity, input.startDate); - const bucketExpr = bucketExprSql(meta.effectiveGranularity); - const metricExpr = metricExprSql(input.metric, input.costSource); - const where = buildWhereClause(input, ctx.user.id, true); - const splitCol = input.splitBy ? dimensionColumn(input.splitBy) : undefined; + const table = getTableName(meta.tier); + const bucketExpr = bucketExprSql(meta.effectiveGranularity, meta.tier); + const metricExpr = metricExprSql(input.metric, meta.tier, input.costSource); + const where = buildWhereClause(meta.tier, input, ctx.user.id, true); + + let statement: string; + if (input.splitBy) { + const splitCol = dimensionColumn(input.splitBy); + statement = ` + SELECT + ${bucketExpr} AS bucket, + ${metricExpr} AS value, + ${splitCol} AS label + FROM ${table} + WHERE ${where.sql()} + GROUP BY 1, 3 + ORDER BY 1 + `; + } else { + statement = ` + SELECT + ${bucketExpr} AS bucket, + ${metricExpr} AS value + FROM ${table} + WHERE ${where.sql()} + GROUP BY 1 + ORDER BY 1 + `; + } - const rows = await timedUsageQuery( + const rows = await timedSnowflakeQuery( { - db: usageReadDb, route: 'usageAnalytics.getTimeseries', queryLabel: `timeseries_${meta.tier}${input.splitBy ? `_split_${input.splitBy}` : ''}`, - scope: queryScope(input), - period: queryPeriod(input), + scope: isOrgScope(input) ? 'org' : 'user', + period: `${input.startDate}/${input.endDate}`, }, - tx => { - const query = tx - .select({ - datetime: bucketExpr, - value: metricExpr, - label: splitCol ?? sql`CAST(NULL AS TEXT)`, - }) - .from(microdollar_usage) - .leftJoin(microdollar_usage_metadata, usageMetadataJoin) - .leftJoin(feature, usageFeatureJoin) - .leftJoin(mode, usageModeJoin) - .where(where.toSQL()); - - return splitCol - ? query.groupBy(bucketExpr, splitCol).orderBy(bucketExpr) - : query.groupBy(bucketExpr).orderBy(bucketExpr); - } + signal => + executeSnowflakeStatement({ + config, + statement, + bindings: where.bindings, + timeoutSeconds: Math.ceil( + defaultTimeoutForScope(isOrgScope(input) ? 'org' : 'user') / 1000 + ), + signal, + }) ); return { timeseries: rows.map(row => ({ - datetime: row.datetime ?? '', - value: toSafeNumber(row.value), - label: input.splitBy ? (row.label ?? undefined) : undefined, + datetime: row[0] ?? '', + value: toSafeNumber(row[1]), + label: input.splitBy ? (row[2] ?? undefined) : undefined, })), effectiveGranularity: meta.effectiveGranularity, }; @@ -642,36 +833,49 @@ export const usageAnalyticsRouter = createTRPCRouter({ .query(async ({ input, ctx }) => { await ensureScopeAccess(ctx, input); + const config = requireSnowflakeConfig(); const meta = resolveTier(input.granularity, input.startDate); + const table = getTableName(meta.tier); const dimCol = dimensionColumn(input.dimension); - const metricExpr = metricExprSql(input.metric, input.costSource); - const where = buildWhereClause(input, ctx.user.id, true); - - const rows = await timedUsageQuery( + const metricExpr = metricExprSql(input.metric, meta.tier, input.costSource); + const where = buildWhereClause(meta.tier, input, ctx.user.id, true); + + const statement = ` + SELECT + ${dimCol} AS key, + ${metricExpr} AS value + FROM ${table} + WHERE ${where.sql()} + GROUP BY 1 + ORDER BY 2 DESC + LIMIT ${Number(input.limit)} + `; + + // SAFETY: LIMIT value is interpolated directly into SQL but is + // validated by Zod above: organization breakdowns allow at most + // MAX_SCOPE_ORGANIZATION_IDS rows, while other dimensions allow 100. + // Snowflake's SQL API v2 does not support parameter binding for LIMIT. + + const rows = await timedSnowflakeQuery( { - db: usageReadDb, route: 'usageAnalytics.getBreakdown', queryLabel: `breakdown_${meta.tier}_by_${input.dimension}`, - scope: queryScope(input), - period: queryPeriod(input), + scope: isOrgScope(input) ? 'org' : 'user', + period: `${input.startDate}/${input.endDate}`, }, - tx => - tx - .select({ - key: dimCol, - value: metricExpr, - }) - .from(microdollar_usage) - .leftJoin(microdollar_usage_metadata, usageMetadataJoin) - .leftJoin(feature, usageFeatureJoin) - .leftJoin(mode, usageModeJoin) - .where(where.toSQL()) - .groupBy(dimCol) - .orderBy(sql`${metricExpr} DESC`) - .limit(input.limit) + signal => + executeSnowflakeStatement({ + config, + statement, + bindings: where.bindings, + timeoutSeconds: Math.ceil( + defaultTimeoutForScope(isOrgScope(input) ? 'org' : 'user') / 1000 + ), + signal, + }) ); - const values = rows.map(row => ({ key: row.key ?? '', value: toSafeNumber(row.value) })); + const values = rows.map(row => ({ key: row[0] ?? '', value: toSafeNumber(row[1]) })); // Percentages are relative to the *returned* rows (limited by input.limit). // They will not reflect the true share when the result set is capped. const totalValue = values.reduce((s, r) => s + r.value, 0); @@ -694,86 +898,98 @@ export const usageAnalyticsRouter = createTRPCRouter({ .query(async ({ input, ctx }) => { await ensureScopeAccess(ctx, input); + const config = requireSnowflakeConfig(); const meta = resolveTier(input.granularity, input.startDate); - const bucketExpr = bucketExprSql(meta.effectiveGranularity); - const where = buildWhereClause(input, ctx.user.id, true); + const table = getTableName(meta.tier); + const bucketExpr = bucketExprSql(meta.effectiveGranularity, meta.tier); + const where = buildWhereClause(meta.tier, input, ctx.user.id, true); + const costSumExpr = costSumExprSql(input.costSource); + const requestedDims = input.groupBy; // For dimensions not in groupBy, emit an empty string constant so the // row shape stays stable regardless of which dimensions were requested. - const featExpr = requestedDims.includes('feature') ? featureName : sql`''`; - const modelExpr = requestedDims.includes('model') ? modelName : sql`''`; - const modeExpr = requestedDims.includes('mode') ? modeName : sql`''`; - const userExpr = requestedDims.includes('user') - ? sql`${microdollar_usage.kilo_user_id}` - : sql`''`; - const providerExpr = requestedDims.includes('provider') ? providerName : sql`''`; - const projectExpr = requestedDims.includes('project') ? projectName : sql`''`; - - // GROUP BY columns: bucket + each requested dimension column - const dimGroupBy = requestedDims.map(d => dimensionColumn(d)); - const groupByClause = [bucketExpr, ...dimGroupBy]; - - const rows = await timedUsageQuery( + const featExpr = requestedDims.includes('feature') ? 'feature' : "''"; + const modelExpr = requestedDims.includes('model') ? 'model' : "''"; + const modeExpr = requestedDims.includes('mode') ? 'mode' : "''"; + const userExpr = requestedDims.includes('user') ? 'kilo_user_id' : "''"; + const providerExpr = requestedDims.includes('provider') ? 'provider' : "''"; + const projectExpr = requestedDims.includes('project') ? 'project_id' : "''"; + + // GROUP BY columns: bucket (pos 1) + each requested dimension column + // SAFETY: dimensionColumn() returns only hardcoded string literals from + // a typed enum chain — never user input. + const dimGroupByCols = requestedDims.map(d => dimensionColumn(d)).join(', '); + const groupByClause = dimGroupByCols ? `1, ${dimGroupByCols}` : '1'; + + const statement = ` + SELECT + ${bucketExpr} AS datetime, + ${featExpr} AS dim_feature, + ${modelExpr} AS dim_model, + ${modeExpr} AS dim_mode, + ${userExpr} AS dim_user, + ${providerExpr} AS dim_provider, + ${projectExpr} AS dim_project, + ${costSumExpr}, + COALESCE(SUM(request_count), 0), + COALESCE(SUM(total_input_tokens), 0), + COALESCE(SUM(total_output_tokens), 0), + COALESCE(SUM(total_cache_write_tokens), 0), + COALESCE(SUM(total_cache_hit_tokens), 0), + COALESCE(SUM(error_count), 0) + FROM ${table} + WHERE ${where.sql()} + GROUP BY ${groupByClause} + ORDER BY 1 DESC + LIMIT ${Number(input.limit)} + `; + + const rows = await timedSnowflakeQuery( { - db: usageReadDb, route: 'usageAnalytics.getTable', queryLabel: `table_${meta.tier}_groupby_${requestedDims.join('+') || 'none'}`, - scope: queryScope(input), - period: queryPeriod(input), + scope: isOrgScope(input) ? 'org' : 'user', + period: `${input.startDate}/${input.endDate}`, }, - tx => - tx - .select({ - datetime: bucketExpr, - dimFeature: featExpr, - dimModel: modelExpr, - dimMode: modeExpr, - dimUser: userExpr, - dimProvider: providerExpr, - dimProject: projectExpr, - costMicrodollars: costSumExprSql(input.costSource), - requestCount: requestCountExpr, - inputTokens: inputTokensExpr, - outputTokens: outputTokensExpr, - cacheWriteTokens: cacheWriteTokensExpr, - cacheHitTokens: cacheHitTokensExpr, - errorCount: errorCountExpr, - }) - .from(microdollar_usage) - .leftJoin(microdollar_usage_metadata, usageMetadataJoin) - .leftJoin(feature, usageFeatureJoin) - .leftJoin(mode, usageModeJoin) - .where(where.toSQL()) - .groupBy(...groupByClause) - .orderBy(sql`${bucketExpr} DESC`) - .limit(input.limit) + signal => + executeSnowflakeStatement({ + config, + statement, + bindings: where.bindings, + timeoutSeconds: Math.ceil( + defaultTimeoutForScope(isOrgScope(input) ? 'org' : 'user') / 1000 + ), + signal, + }) ); + const dimIndexMap: Record = { + feature: 1, + model: 2, + mode: 3, + user: 4, + provider: 5, + project: 6, + }; + return { rows: rows.map(row => { const dimensions: Record = {}; for (const d of requestedDims) { - const raw = { - feature: row.dimFeature, - model: row.dimModel, - mode: row.dimMode, - user: row.dimUser, - provider: row.dimProvider, - project: row.dimProject, - }[d]; + const raw = row[dimIndexMap[d]]; dimensions[d] = typeof raw === 'string' ? raw : ''; } return { - datetime: row.datetime ?? '', + datetime: row[0] ?? '', dimensions, - costMicrodollars: toSafeNumber(row.costMicrodollars), - requestCount: toSafeNumber(row.requestCount), - inputTokens: toSafeNumber(row.inputTokens), - outputTokens: toSafeNumber(row.outputTokens), - cacheWriteTokens: toSafeNumber(row.cacheWriteTokens), - cacheHitTokens: toSafeNumber(row.cacheHitTokens), - errorCount: toSafeNumber(row.errorCount), + costMicrodollars: toSafeNumber(row[7]), + requestCount: toSafeNumber(row[8]), + inputTokens: toSafeNumber(row[9]), + outputTokens: toSafeNumber(row[10]), + cacheWriteTokens: toSafeNumber(row[11]), + cacheHitTokens: toSafeNumber(row[12]), + errorCount: toSafeNumber(row[13]), }; }), effectiveGranularity: meta.effectiveGranularity,