diff --git a/apps/web/src/lib/drizzle.test.ts b/apps/web/src/lib/drizzle.test.ts index d2e1409737..82ed9d423d 100644 --- a/apps/web/src/lib/drizzle.test.ts +++ b/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({ diff --git a/apps/web/src/lib/drizzle.ts b/apps/web/src/lib/drizzle.ts index 0136e743db..cc3f0abfb7 100644 --- a/apps/web/src/lib/drizzle.ts +++ b/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,6 +148,25 @@ 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') { @@ -128,6 +174,9 @@ if (process.env.NODE_ENV !== 'test') { 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 { 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 5ec6559e3f..1bb4a273a2 100644 --- a/apps/web/src/routers/usage-analytics-router.test.ts +++ b/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) { 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'); }); }); diff --git a/apps/web/src/routers/usage-analytics-router.ts b/apps/web/src/routers/usage-analytics-router.ts index 1df2a013c5..f590b64af8 100644 --- a/apps/web/src/routers/usage-analytics-router.ts +++ b/apps/web/src/routers/usage-analytics-router.ts @@ -1,16 +1,15 @@ import { TRPCError } from '@trpc/server'; import * as z from 'zod'; -import { and, asc, eq, inArray, isNull, or } from 'drizzle-orm'; +import { and, asc, eq, gte, inArray, isNull, lt, notInArray, or, sql, type SQL } from 'drizzle-orm'; import { baseProcedure, createTRPCRouter, type TRPCContext } from '@/lib/trpc/init'; -import { readDb } from '@/lib/drizzle'; -import { getEnvVariable } from '@/lib/dotenvx'; -import { - executeSnowflakeStatement, - resolveSnowflakeConfig, - type SnowflakeBinding, -} from '@/lib/snowflake'; +import { readDb, usageReadDb } from '@/lib/drizzle'; +import { timedUsageQuery } from '@/lib/usage-query'; import { + feature, kilocode_users, + microdollar_usage, + microdollar_usage_metadata, + mode, organization_memberships, organizations, user_auth_provider, @@ -91,118 +90,34 @@ function resolveTier(granularity: Granularity, startDate: string): TableMeta { if (ageDays < 8) { return { tier: 'hourly', effectiveGranularity: 'hour' }; } - // Auto-downgrade: hourly data is only available for the past 7 days. + // Auto-downgrade: hourly buckets are only used 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 with positional `?` bindings. - * Callers push conditions in any order; `sql()` joins them with AND. + * Accumulates SQL WHERE clauses. Callers push conditions in any order; + * `toSQL()` joins them with AND. */ export class WhereBuilder { - readonly clauses: string[] = []; - readonly bindings: SnowflakeBinding[] = []; + readonly conditions: SQL[] = []; - private push(clause: string, ...bindings: SnowflakeBinding[]) { - bindings.forEach(b => this.bindings.push(b)); - this.clauses.push(clause); + add(condition: SQL): void { + this.conditions.push(condition); } - 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'; + toSQL(): SQL | undefined { + return this.conditions.length > 0 ? and(...this.conditions) : undefined; } } @@ -258,29 +173,9 @@ async function ensureScopeAccess(ctx: TRPCContext, filters: UsageAnalyticsFilter // WHERE clause helpers // --------------------------------------------------------------------------- -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) - ); - } +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)); } export function buildScopeConditions( @@ -291,65 +186,83 @@ 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.addIn('organization_id', filters.organizationIds); + where.add(inArray(microdollar_usage.organization_id, filters.organizationIds)); if (filters.userIds && filters.userIds.length > 0) { - where.addIn('kilo_user_id', filters.userIds); + where.add(inArray(microdollar_usage.kilo_user_id, filters.userIds)); } if (filters.excludedUserIds && filters.excludedUserIds.length > 0) { - where.addNotIn('kilo_user_id', filters.excludedUserIds); + where.add(notInArray(microdollar_usage.kilo_user_id, filters.excludedUserIds)); } return; } if (filters.organizationId) { - where.addEq('organization_id', filters.organizationId); + where.add(eq(microdollar_usage.organization_id, filters.organizationId)); if (filters.viewAs === 'self') { - where.addEq('kilo_user_id', ctxUserId); + where.add(eq(microdollar_usage.kilo_user_id, ctxUserId)); } else { if (filters.userIds && filters.userIds.length > 0) { - where.addIn('kilo_user_id', filters.userIds); + where.add(inArray(microdollar_usage.kilo_user_id, filters.userIds)); } if (filters.excludedUserIds && filters.excludedUserIds.length > 0) { - where.addNotIn('kilo_user_id', filters.excludedUserIds); + where.add(notInArray(microdollar_usage.kilo_user_id, filters.excludedUserIds)); } } } else { - where.addEq('kilo_user_id', ctxUserId); + where.add(eq(microdollar_usage.kilo_user_id, ctxUserId)); if (filters.personalScope === 'personal-only') { - // DBT coalesces personal Snowflake usage rollups to an empty-string sentinel - // so incremental merges can match on organization_id. - where.addEq('organization_id', ''); + // Personal usage is stored with a NULL organization_id. + where.add(isNull(microdollar_usage.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: string, values: string[] | undefined) => { - if (values && values.length > 0) where.addIn(column, values); + const addInIfNonEmpty = (column: SQL, values: string[] | undefined) => { + if (values && values.length > 0) where.add(inValues(column, values)); }; - const addNotInIfNonEmpty = (column: string, values: string[] | undefined) => { - if (values && values.length > 0) where.addNotIn(column, values); + const addNotInIfNonEmpty = (column: SQL, values: string[] | undefined) => { + if (values && values.length > 0) where.add(notInValues(column, values)); }; - 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); + 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); } function buildWhereClause( - tier: GranularityTier, filters: UsageAnalyticsFilters, ctxUserId: string, includeDimensions: boolean ): WhereBuilder { const where = new WhereBuilder(); - buildDateConditions(where, tier, filters); + buildDateConditions(where, filters); buildScopeConditions(where, filters, ctxUserId); if (includeDimensions) { buildDimensionConditions(where, filters); @@ -361,59 +274,62 @@ function buildWhereClause( // Metric SQL expression // --------------------------------------------------------------------------- -export function costColumnFor(costSource: CostSource): string { +export function costColumnFor(costSource: CostSource): SQL { switch (costSource) { case 'cost': - return 'total_cost_microdollars'; + return sql`${microdollar_usage.cost}`; case 'market': - return 'total_market_cost_microdollars'; + return sql`COALESCE(${microdollar_usage_metadata.market_cost}, 0)`; } } -export function costSumExprSql(costSource: CostSource): string { - return `COALESCE(SUM(${costColumnFor(costSource)}), 0)`; +export function costSumExprSql(costSource: CostSource): SQL { + return sql`COALESCE(SUM(${costColumnFor(costSource)}), 0)`; } -function metricExprSql(metric: Metric, tier: GranularityTier, costSource: CostSource): string { +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_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 { const costSumExpr = costSumExprSql(costSource); switch (metric) { case 'cost': return costSumExpr; case 'requests': - return 'COALESCE(SUM(request_count), 0)'; + return sql`COALESCE(${requestCountExpr}, 0)`; case 'inputTokens': - return 'COALESCE(SUM(total_input_tokens), 0)'; + return inputTokensExpr; case 'outputTokens': - return 'COALESCE(SUM(total_output_tokens), 0)'; + return outputTokensExpr; case 'tokens': - return 'COALESCE(SUM(total_tokens), 0)'; + return totalTokensExpr; case 'errorRate': - return 'CASE WHEN COALESCE(SUM(request_count), 0) = 0 THEN 0 ELSE COALESCE(SUM(error_count), 0)::FLOAT / SUM(request_count)::FLOAT END'; + return sql`CASE WHEN ${requestCountExpr} = 0 THEN 0 ELSE (${errorCountExpr})::FLOAT / (${requestCountExpr})::FLOAT END`; case 'avgLatencyMs': - 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`; - } + 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`; case 'costPerRequest': - return `CASE WHEN COALESCE(SUM(request_count), 0) = 0 THEN 0 ELSE ${costSumExpr}::FLOAT / SUM(request_count)::FLOAT END`; + return sql`CASE WHEN ${requestCountExpr} = 0 THEN 0 ELSE (${costSumExpr})::FLOAT / (${requestCountExpr})::FLOAT END`; case 'tokensPerRequest': - return 'CASE WHEN COALESCE(SUM(request_count), 0) = 0 THEN 0 ELSE COALESCE(SUM(total_tokens), 0)::FLOAT / SUM(request_count)::FLOAT END'; + return sql`CASE WHEN ${requestCountExpr} = 0 THEN 0 ELSE (${totalTokensExpr})::FLOAT / (${requestCountExpr})::FLOAT END`; case 'cacheHitRatio': - 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'; + 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`; case 'outputInputRatio': - 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)'; + 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`; } - // 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)'; } // --------------------------------------------------------------------------- @@ -427,124 +343,49 @@ function generationTimeCountExprSql(tier: GranularityTier): string { * 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 (from daily table) + * Month → 'YYYY-MM-DD' of the first of the month */ -function bucketExprSql(effectiveGranularity: Granularity, tier: GranularityTier): string { - const timeCol = getTimeColumn(tier); - +function bucketExprSql(effectiveGranularity: Granularity): SQL { + const createdAtUtc = sql`${microdollar_usage.created_at} AT TIME ZONE 'UTC'`; if (effectiveGranularity === 'hour') { - return `TO_VARCHAR(${timeCol}, 'YYYY-MM-DD HH24:MI:SS')`; + return sql`TO_CHAR(DATE_TRUNC('hour', ${createdAtUtc}), 'YYYY-MM-DD HH24:MI:SS')`; } if (effectiveGranularity === 'week') { - return `TO_VARCHAR(DATE_TRUNC('WEEK', ${timeCol}), 'YYYY-MM-DD')`; + return sql`TO_CHAR(DATE_TRUNC('week', ${createdAtUtc}), 'YYYY-MM-DD')`; } if (effectiveGranularity === 'month') { - // Daily table, group by month - return `TO_VARCHAR(DATE_TRUNC('MONTH', ${timeCol}), 'YYYY-MM-DD')`; + return sql`TO_CHAR(DATE_TRUNC('month', ${createdAtUtc}), 'YYYY-MM-DD')`; } // 'day' - return `TO_VARCHAR(${timeCol}, 'YYYY-MM-DD')`; + return sql`TO_CHAR(DATE_TRUNC('day', ${createdAtUtc}), 'YYYY-MM-DD')`; } // --------------------------------------------------------------------------- // Dimension column name // --------------------------------------------------------------------------- -export function dimensionColumn(dimension: BreakdownDimension): string { +export function dimensionColumn(dimension: BreakdownDimension): SQL { switch (dimension) { case 'feature': - return 'feature'; + return featureName; case 'model': - return 'model'; + return modelName; case 'mode': - return 'mode'; + return modeName; case 'user': - return 'kilo_user_id'; + return sql`${microdollar_usage.kilo_user_id}`; case 'provider': - return 'provider'; + return providerName; case 'project': - return 'project_id'; + return projectName; case 'organization': - return 'organization_id'; + return sql`COALESCE(${microdollar_usage.organization_id}::text, '')`; } } -// --------------------------------------------------------------------------- -// 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 (error) { - settled = true; - console.error( - JSON.stringify({ - type: 'usage_query_error', - route: params.route, - queryLabel: params.queryLabel, - scope: params.scope, - period: params.period, - message: error instanceof Error ? error.message : String(error), - }) - ); - 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, - }) - ); - } -} +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); // --------------------------------------------------------------------------- // getSummary @@ -556,7 +397,7 @@ function ratioSafe(numerator: number, denominator: number): number { } /** - * Convert an aggregate value (often returned as a string by Snowflake) to a + * Convert an aggregate value (often returned as a string by Postgres) to a * JS number. Values above `MAX_SAFE_INTEGER` are logged as a warning but still * returned so the UI does not crash. */ @@ -643,6 +484,14 @@ 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}`; +} + // --------------------------------------------------------------------------- // Router definition // --------------------------------------------------------------------------- @@ -654,100 +503,61 @@ export const usageAnalyticsRouter = createTRPCRouter({ .query(async ({ input, ctx }): Promise => { await ensureScopeAccess(ctx, input); - const config = resolveSnowflakeConfig(); const meta = resolveTier(input.granularity, input.startDate); - if (!config) { - return { - costMicrodollars: 0, - requestCount: 0, - inputTokens: 0, - outputTokens: 0, - cacheWriteTokens: 0, - cacheHitTokens: 0, - errorCount: 0, - cancelledCount: 0, - freeRequestCount: 0, - byokRequestCount: 0, - totalLatencyMs: 0, - totalGenerationTimeMs: 0, - latencyCount: 0, - generationTimeCount: 0, - totalTokens: 0, - distinctUsers: 0, - errorRate: 0, - avgLatencyMs: 0, - avgGenerationTimeMs: 0, - costPerRequest: 0, - tokensPerRequest: 0, - cacheHitRatio: 0, - outputInputRatio: 0, - effectiveGranularity: meta.effectiveGranularity, - }; - } - 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(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( + const where = buildWhereClause(input, ctx.user.id, true); + + const rows = await timedUsageQuery( { + db: usageReadDb, route: 'usageAnalytics.getSummary', queryLabel: `summary_${meta.tier}`, - scope: isOrgScope(input) ? 'org' : 'user', - period: `${input.startDate}/${input.endDate}`, + scope: queryScope(input), + period: queryPeriod(input), }, - signal => - executeSnowflakeStatement({ - config, - statement, - bindings: where.bindings, - timeoutSeconds: Math.ceil( - defaultTimeoutForScope(isOrgScope(input) ? 'org' : 'user') / 1000 - ), - signal, - }) + 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()) ); - 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]); + 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); return { costMicrodollars, @@ -783,65 +593,44 @@ export const usageAnalyticsRouter = createTRPCRouter({ .query(async ({ input, ctx }) => { await ensureScopeAccess(ctx, input); - const config = resolveSnowflakeConfig(); const meta = resolveTier(input.granularity, input.startDate); - if (!config) { - return { timeseries: [], effectiveGranularity: meta.effectiveGranularity }; - } - 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 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 rows = await timedSnowflakeQuery( + const rows = await timedUsageQuery( { + db: usageReadDb, route: 'usageAnalytics.getTimeseries', queryLabel: `timeseries_${meta.tier}${input.splitBy ? `_split_${input.splitBy}` : ''}`, - scope: isOrgScope(input) ? 'org' : 'user', - period: `${input.startDate}/${input.endDate}`, + scope: queryScope(input), + period: queryPeriod(input), }, - signal => - executeSnowflakeStatement({ - config, - statement, - bindings: where.bindings, - timeoutSeconds: Math.ceil( - defaultTimeoutForScope(isOrgScope(input) ? 'org' : 'user') / 1000 - ), - signal, - }) + 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); + } ); return { timeseries: rows.map(row => ({ - datetime: row[0] ?? '', - value: toSafeNumber(row[1]), - label: input.splitBy ? (row[2] ?? undefined) : undefined, + datetime: row.datetime ?? '', + value: toSafeNumber(row.value), + label: input.splitBy ? (row.label ?? undefined) : undefined, })), effectiveGranularity: meta.effectiveGranularity, }; @@ -853,52 +642,36 @@ export const usageAnalyticsRouter = createTRPCRouter({ .query(async ({ input, ctx }) => { await ensureScopeAccess(ctx, input); - const config = resolveSnowflakeConfig(); const meta = resolveTier(input.granularity, input.startDate); - if (!config) { - return { breakdown: [], totalValue: 0, effectiveGranularity: meta.effectiveGranularity }; - } - const table = getTableName(meta.tier); const dimCol = dimensionColumn(input.dimension); - 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( + const metricExpr = metricExprSql(input.metric, input.costSource); + const where = buildWhereClause(input, ctx.user.id, true); + + const rows = await timedUsageQuery( { + db: usageReadDb, route: 'usageAnalytics.getBreakdown', queryLabel: `breakdown_${meta.tier}_by_${input.dimension}`, - scope: isOrgScope(input) ? 'org' : 'user', - period: `${input.startDate}/${input.endDate}`, + scope: queryScope(input), + period: queryPeriod(input), }, - signal => - executeSnowflakeStatement({ - config, - statement, - bindings: where.bindings, - timeoutSeconds: Math.ceil( - defaultTimeoutForScope(isOrgScope(input) ? 'org' : 'user') / 1000 - ), - signal, - }) + 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) ); - const values = rows.map(row => ({ key: row[0] ?? '', value: toSafeNumber(row[1]) })); + const values = rows.map(row => ({ key: row.key ?? '', value: toSafeNumber(row.value) })); // 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); @@ -921,101 +694,86 @@ export const usageAnalyticsRouter = createTRPCRouter({ .query(async ({ input, ctx }) => { await ensureScopeAccess(ctx, input); - const config = resolveSnowflakeConfig(); const meta = resolveTier(input.granularity, input.startDate); - if (!config) { - return { rows: [], effectiveGranularity: meta.effectiveGranularity }; - } - 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 bucketExpr = bucketExprSql(meta.effectiveGranularity); + const where = buildWhereClause(input, ctx.user.id, true); 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') ? '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( + 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( { + db: usageReadDb, route: 'usageAnalytics.getTable', queryLabel: `table_${meta.tier}_groupby_${requestedDims.join('+') || 'none'}`, - scope: isOrgScope(input) ? 'org' : 'user', - period: `${input.startDate}/${input.endDate}`, + scope: queryScope(input), + period: queryPeriod(input), }, - signal => - executeSnowflakeStatement({ - config, - statement, - bindings: where.bindings, - timeoutSeconds: Math.ceil( - defaultTimeoutForScope(isOrgScope(input) ? 'org' : 'user') / 1000 - ), - signal, - }) + 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) ); - 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 = row[dimIndexMap[d]]; + const raw = { + feature: row.dimFeature, + model: row.dimModel, + mode: row.dimMode, + user: row.dimUser, + provider: row.dimProvider, + project: row.dimProject, + }[d]; dimensions[d] = typeof raw === 'string' ? raw : ''; } return { - datetime: row[0] ?? '', + datetime: row.datetime ?? '', dimensions, - 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]), + 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), }; }), effectiveGranularity: meta.effectiveGranularity,