From b167eaec1dd97c0bc96fc1323cb83283cc0360cd Mon Sep 17 00:00:00 2001 From: Chris Estreich Date: Sun, 18 May 2025 11:48:40 -0700 Subject: [PATCH] More usage queries (#31) --- src/actions/__tests__/syncCurrentUser.test.ts | 3 +- src/actions/__tests__/syncOrg.test.ts | 3 +- src/actions/analytics.ts | 104 ++++++++++++---- src/actions/auditLogs.ts | 11 +- src/actions/defaultParameters.ts | 3 +- src/actions/organizationSettings.ts | 3 +- src/actions/providerWhitelist.ts | 2 +- src/actions/sync.ts | 12 +- .../(authenticated)/audit-logs/AuditLogs.tsx | 2 +- src/app/(authenticated)/usage/Developers.tsx | 113 +++++------------- src/app/(authenticated)/usage/Usage.tsx | 5 +- src/app/(authenticated)/usage/types.ts | 10 -- src/components/audit-logs/AuditLogCard.tsx | 2 +- src/components/audit-logs/AuditLogDetails.tsx | 2 +- src/components/audit-logs/AuditLogDrawer.tsx | 2 +- src/db/client.ts | 33 +++++ src/db/index.ts | 35 +----- 17 files changed, 178 insertions(+), 167 deletions(-) create mode 100644 src/db/client.ts diff --git a/src/actions/__tests__/syncCurrentUser.test.ts b/src/actions/__tests__/syncCurrentUser.test.ts index b341500175..07988ecd20 100644 --- a/src/actions/__tests__/syncCurrentUser.test.ts +++ b/src/actions/__tests__/syncCurrentUser.test.ts @@ -1,10 +1,9 @@ // pnpm test src/actions/__tests__/syncCurrentUser.test.ts -import { db } from '@/db'; -import { users, orgs } from '@/db/schema'; import { eq } from 'drizzle-orm'; import { logger } from '@/lib/server/logger'; +import { client as db, users, orgs } from '@/db'; import { syncCurrentUser } from '../sync'; diff --git a/src/actions/__tests__/syncOrg.test.ts b/src/actions/__tests__/syncOrg.test.ts index c7260a9726..f06d6075ea 100644 --- a/src/actions/__tests__/syncOrg.test.ts +++ b/src/actions/__tests__/syncOrg.test.ts @@ -3,9 +3,8 @@ import { eq } from 'drizzle-orm'; import { clerkClient } from '@clerk/nextjs/server'; -import { db } from '@/db'; -import { orgs } from '@/db/schema'; import { logger } from '@/lib/server/logger'; +import { client as db, orgs } from '@/db'; import { syncOrg } from '../sync'; diff --git a/src/actions/analytics.ts b/src/actions/analytics.ts index 0c1d39038e..52489f148b 100644 --- a/src/actions/analytics.ts +++ b/src/actions/analytics.ts @@ -9,6 +9,8 @@ import { import type { TimePeriod } from '@/types'; import { analytics } from '@/lib/server'; +import * as db from '@/db'; +import { inArray } from 'drizzle-orm'; /** * captureEvent @@ -38,7 +40,7 @@ export const captureEvent = async ({ }; /** - * Usage + * getUsage */ const usageSchema = z.object({ @@ -52,10 +54,6 @@ const usageSchema = z.object({ export type Usage = z.infer; -/** - * getUsage - */ - type UsageRecord = Partial>; export const getUsage = async ({ @@ -64,7 +62,7 @@ export const getUsage = async ({ }: { orgId?: string | null; timePeriod: TimePeriod; -}): Promise => { +}): Promise => { if (!orgId) { return {}; } @@ -78,23 +76,89 @@ export const getUsage = async ({ SUM(COALESCE(inputTokens, 0)) AS inputTokens, SUM(COALESCE(outputTokens, 0)) AS outputTokens, SUM(COALESCE(cost, 0)) AS cost - FROM - events - WHERE - orgId = {orgId: String} - AND timestamp >= toUnixTimestamp(now() - INTERVAL {timePeriod: Int32} DAY) - GROUP BY - type + FROM events + WHERE orgId = {orgId: String} AND timestamp >= toUnixTimestamp(now() - INTERVAL {timePeriod: Int32} DAY) + GROUP BY 1 `, format: 'JSONEachRow', query_params: { orgId, timePeriod }, }); - const data = await resultSet.json(); - const usages = z.array(usageSchema).parse(data); - - return usages.reduce( - (collect, usage) => ({ ...collect, [usage.type]: usage }), - {} as UsageRecord, - ); + return z + .array(usageSchema) + .parse(await resultSet.json()) + .reduce( + (collect, usage) => ({ ...collect, [usage.type]: usage }), + {} as UsageRecord, + ); +}; + +/** + * getDeveloperUsage + */ + +const developerUsageSchema = z.object({ + userId: z.string(), + tasksStarted: z.coerce.number(), + tasksCompleted: z.coerce.number(), + tokens: z.coerce.number(), + cost: z.coerce.number(), +}); + +export type DeveloperUsage = z.infer & { + user: db.User; +}; + +export const getDeveloperUsage = async ({ + orgId, + timePeriod, +}: { + orgId?: string | null; + timePeriod: TimePeriod; +}): Promise => { + if (!orgId) { + return []; + } + + const resultSet = await analytics.query({ + query: ` + SELECT + userId, + SUM(CASE WHEN type = '${TelemetryEventName.TASK_CREATED}' THEN 1 ELSE 0 END) AS tasksStarted, + SUM(CASE WHEN type = '${TelemetryEventName.TASK_COMPLETED}' THEN 1 ELSE 0 END) AS tasksCompleted, + SUM(CASE WHEN type = '${TelemetryEventName.LLM_COMPLETION}' THEN COALESCE(inputTokens, 0) + COALESCE(outputTokens, 0) ELSE 0 END) AS tokens, + SUM(CASE WHEN type = '${TelemetryEventName.LLM_COMPLETION}' THEN COALESCE(cost, 0) ELSE 0 END) AS cost + FROM events + WHERE orgId = {orgId: String} AND timestamp >= toUnixTimestamp(now() - INTERVAL {timePeriod: Int32} DAY) + GROUP BY 1 + `, + format: 'JSONEachRow', + query_params: { orgId, timePeriod }, + }); + + const developerUsages = z + .array(developerUsageSchema) + .parse(await resultSet.json()); + + const users = ( + await db.client + .select() + .from(db.users) + .where( + inArray( + db.users.id, + developerUsages.map(({ userId }) => userId), + ), + ) + ).reduce( + (acc, user) => ({ ...acc, [user.id]: user }), + {} as Record, + ); + + return developerUsages + .map((usage) => ({ + ...usage, + user: users[usage.userId], + })) + .filter((usage): usage is DeveloperUsage => !!usage.user); }; diff --git a/src/actions/auditLogs.ts b/src/actions/auditLogs.ts index 741e29e076..3a9e309ed0 100644 --- a/src/actions/auditLogs.ts +++ b/src/actions/auditLogs.ts @@ -3,8 +3,13 @@ import { eq, gte, and, desc } from 'drizzle-orm'; import { logger } from '@/lib/server'; -import { db, type DB_OR_TX } from '@/db'; -import { auditLogs, type AuditLog, type CreateAuditLog } from '@/db/schema'; +import { + type DatabaseOrTransaction, + type AuditLog, + type CreateAuditLog, + client as db, + auditLogs, +} from '@/db'; export const getAuditLogs = async ({ orgId, @@ -59,7 +64,7 @@ export const getAuditLogs = async ({ }; export async function insertAuditLog( - db: DB_OR_TX, + db: DatabaseOrTransaction, values: CreateAuditLog, ): Promise { await db.insert(auditLogs).values(values); diff --git a/src/actions/defaultParameters.ts b/src/actions/defaultParameters.ts index 63369c0ea5..348ed51724 100644 --- a/src/actions/defaultParameters.ts +++ b/src/actions/defaultParameters.ts @@ -4,8 +4,7 @@ import { sql } from 'drizzle-orm'; import { z } from 'zod'; import { type ApiResponse, ORGANIZATION_ALLOW_ALL } from '@/types'; -import { db } from '@/db'; -import { AuditLogTargetType, orgSettings } from '@/db/schema'; +import { client as db, AuditLogTargetType, orgSettings } from '@/db'; import { isAuthSuccess, handleError } from '@/lib/server'; import { validateAuth } from './auth'; diff --git a/src/actions/organizationSettings.ts b/src/actions/organizationSettings.ts index 0fab7f71c2..e840927b44 100644 --- a/src/actions/organizationSettings.ts +++ b/src/actions/organizationSettings.ts @@ -12,8 +12,7 @@ import { organizationAllowListSchema, organizationDefaultSettingsSchema, } from '@/types'; -import { db } from '@/db'; -import { AuditLogTargetType, orgSettings } from '@/db/schema'; +import { client as db, AuditLogTargetType, orgSettings } from '@/db'; import { handleError, isAuthSuccess } from '@/lib/server'; import { validateAuth } from './auth'; diff --git a/src/actions/providerWhitelist.ts b/src/actions/providerWhitelist.ts index 1aa2f4b968..64d0e731a9 100644 --- a/src/actions/providerWhitelist.ts +++ b/src/actions/providerWhitelist.ts @@ -4,7 +4,7 @@ import { z } from 'zod'; import type { ApiResponse } from '@/types'; import { handleError, isAuthSuccess } from '@/lib/server'; -import { AuditLogTargetType } from '@/db/schema'; +import { AuditLogTargetType } from '@/db'; import { validateAuth } from './auth'; import { createAuditLog } from './auditLogs'; diff --git a/src/actions/sync.ts b/src/actions/sync.ts index c0684807aa..2431fcb674 100644 --- a/src/actions/sync.ts +++ b/src/actions/sync.ts @@ -3,10 +3,14 @@ import { eq } from 'drizzle-orm'; import { clerkClient, currentUser } from '@clerk/nextjs/server'; -import { db } from '@/db'; -import { users, type CreateUser, orgs, type CreateOrg } from '@/db/schema'; - -import { logger } from '../lib/server/logger'; +import { + type CreateUser, + type CreateOrg, + client as db, + users, + orgs, +} from '@/db'; +import { logger } from '@/lib/server'; export async function syncCurrentUser({ userId, diff --git a/src/app/(authenticated)/audit-logs/AuditLogs.tsx b/src/app/(authenticated)/audit-logs/AuditLogs.tsx index f7878cdd4a..277c0dfaf4 100644 --- a/src/app/(authenticated)/audit-logs/AuditLogs.tsx +++ b/src/app/(authenticated)/audit-logs/AuditLogs.tsx @@ -5,7 +5,7 @@ import { useAuth } from '@clerk/nextjs'; import { useQuery } from '@tanstack/react-query'; import { type TimePeriod, timePeriods } from '@/types'; -import type { AuditLog } from '@/db/schema'; +import type { AuditLog } from '@/db'; import { getAuditLogs } from '@/actions/auditLogs'; import { Button, diff --git a/src/app/(authenticated)/usage/Developers.tsx b/src/app/(authenticated)/usage/Developers.tsx index e8bf5461d2..9dd383275b 100644 --- a/src/app/(authenticated)/usage/Developers.tsx +++ b/src/app/(authenticated)/usage/Developers.tsx @@ -1,81 +1,40 @@ +import { useQuery } from '@tanstack/react-query'; import type { ColumnDef } from '@tanstack/react-table'; +import { useAuth } from '@clerk/nextjs'; +import { type User } from '@/db'; +import { type DeveloperUsage, getDeveloperUsage } from '@/actions/analytics'; import { DataTable } from '@/components/layout/DataTable'; - -import type { Developer } from './types'; - -const mockDevelopers: Developer[] = [ - { - id: 'dev1', - name: 'John Doe', - email: 'john@example.com', - tasksStarted: 42, - tasksCompleted: 38, - tokensConsumed: 1200000, - cost: 24.5, - }, - { - id: 'dev2', - name: 'Jane Smith', - email: 'jane@example.com', - tasksStarted: 35, - tasksCompleted: 32, - tokensConsumed: 980000, - cost: 19.6, - }, - { - id: 'dev3', - name: 'Bob Johnson', - email: 'bob@example.com', - tasksStarted: 28, - tasksCompleted: 25, - tokensConsumed: 750000, - cost: 15.0, - }, - { - id: 'dev4', - name: 'Alice Williams', - email: 'alice@example.com', - tasksStarted: 31, - tasksCompleted: 29, - tokensConsumed: 820000, - cost: 16.4, - }, - { - id: 'dev5', - name: 'Charlie Brown', - email: 'charlie@example.com', - tasksStarted: 22, - tasksCompleted: 19, - tokensConsumed: 650000, - cost: 13.0, - }, -]; +import { formatCurrency, formatNumber } from '@/lib/formatters'; export const Developers = ({ - onClick, + onDeveloperSelected, }: { - onClick: (developer: Developer) => void; + onDeveloperSelected: (user: User) => void; }) => { - const developerColumns: ColumnDef[] = [ - { - accessorKey: 'name', - header: 'Developer', - cell: ({ row }) => { - const developer = row.original; + const { orgId } = useAuth(); - return ( - - ); - }, + const { data = [] } = useQuery({ + queryKey: ['developers'], + queryFn: () => getDeveloperUsage({ orgId, timePeriod: 30 }), + enabled: !!orgId, + }); + + const columns: ColumnDef[] = [ + { + accessorKey: 'user.name', + header: 'Developer', + cell: ({ row }) => ( + + ), }, { - accessorKey: 'email', + accessorKey: 'user.email', header: 'Email', }, { @@ -87,26 +46,16 @@ export const Developers = ({ header: 'Tasks Completed', }, { - accessorKey: 'tokensConsumed', + accessorKey: 'tokens', header: 'Tokens', - cell: ({ row }) => { - const tokens = row.getValue('tokensConsumed') as number; - return tokens >= 1000000 - ? `${(tokens / 1000000).toFixed(1)}M` - : tokens >= 1000 - ? `${(tokens / 1000).toFixed(1)}K` - : tokens; - }, + cell: ({ row }) => formatNumber(row.original.tokens), }, { accessorKey: 'cost', header: 'Cost (USD)', - cell: ({ row }) => { - const cost = row.getValue('cost') as number; - return `$${cost.toFixed(2)}`; - }, + cell: ({ row }) => formatCurrency(row.original.cost), }, ]; - return ; + return ; }; diff --git a/src/app/(authenticated)/usage/Usage.tsx b/src/app/(authenticated)/usage/Usage.tsx index e854425e41..1845d3d0df 100644 --- a/src/app/(authenticated)/usage/Usage.tsx +++ b/src/app/(authenticated)/usage/Usage.tsx @@ -2,9 +2,10 @@ import React, { useState } from 'react'; +import { type User } from '@/db'; import { UsageCard } from '@/components/usage/UsageCard'; -import type { Developer, Filter, Model, Task, ViewMode } from './types'; +import type { Filter, Model, Task, ViewMode } from './types'; import { ViewModeToggle } from './ViewModeToggle'; import { ActiveFilter } from './ActiveFilter'; import { Developers } from './Developers'; @@ -32,7 +33,7 @@ export const Usage = () => { /> ) : viewMode === 'developers' ? ( { + onDeveloperSelected={({ id, name }: User) => { setFilter({ type: 'developer', id, name }); setViewMode('tasks'); }} diff --git a/src/app/(authenticated)/usage/types.ts b/src/app/(authenticated)/usage/types.ts index e246e3c8d3..38d3602a88 100644 --- a/src/app/(authenticated)/usage/types.ts +++ b/src/app/(authenticated)/usage/types.ts @@ -2,16 +2,6 @@ export const viewModes = ['developers', 'models', 'tasks'] as const; export type ViewMode = (typeof viewModes)[number]; -export type Developer = { - id: string; - name: string; - email: string; - tasksStarted: number; - tasksCompleted: number; - tokensConsumed: number; - cost: number; -}; - export type Model = { id: string; name: string; diff --git a/src/components/audit-logs/AuditLogCard.tsx b/src/components/audit-logs/AuditLogCard.tsx index ba14f50920..f54900537c 100644 --- a/src/components/audit-logs/AuditLogCard.tsx +++ b/src/components/audit-logs/AuditLogCard.tsx @@ -7,7 +7,7 @@ import { useAuth } from '@clerk/nextjs'; import { useQuery } from '@tanstack/react-query'; import { ArrowRightIcon } from 'lucide-react'; -import type { AuditLog } from '@/db/schema'; +import type { AuditLog } from '@/db'; import { getAuditLogs } from '@/actions/auditLogs'; import { Card, diff --git a/src/components/audit-logs/AuditLogDetails.tsx b/src/components/audit-logs/AuditLogDetails.tsx index 450e28672a..1f10034c88 100644 --- a/src/components/audit-logs/AuditLogDetails.tsx +++ b/src/components/audit-logs/AuditLogDetails.tsx @@ -3,7 +3,7 @@ import { ArrowRight, Calendar, Clock, User } from 'lucide-react'; import Link from 'next/link'; -import type { AuditLog } from '@/db/schema'; +import type { AuditLog } from '@/db'; type AuditLogDetailsProps = { log: AuditLog; diff --git a/src/components/audit-logs/AuditLogDrawer.tsx b/src/components/audit-logs/AuditLogDrawer.tsx index 33fb7c15e4..4aec6353fe 100644 --- a/src/components/audit-logs/AuditLogDrawer.tsx +++ b/src/components/audit-logs/AuditLogDrawer.tsx @@ -1,4 +1,4 @@ -import type { AuditLog } from '@/db/schema'; +import type { AuditLog } from '@/db'; import { Drawer, DrawerContent, diff --git a/src/db/client.ts b/src/db/client.ts new file mode 100644 index 0000000000..4e1b44448a --- /dev/null +++ b/src/db/client.ts @@ -0,0 +1,33 @@ +import { drizzle } from 'drizzle-orm/node-postgres'; +import { Pool } from 'pg'; + +import { Env } from '@/lib/server'; + +const pool = new Pool({ + connectionString: process.env.DATABASE_URL, +}); + +let testDb: ReturnType | undefined = undefined; + +if (process.env.NODE_ENV === 'test') { + if ( + !Env.DATABASE_URL?.includes('test') || + !Env.DATABASE_URL?.includes('localhost') + ) { + throw new Error('DATABASE_URL is not a test database'); + } + + testDb = drizzle(pool); +} + +const client = process.env.NODE_ENV === 'test' ? testDb! : drizzle(pool); + +const disconnect = async () => { + await pool.end(); +}; + +type DatabaseOrTransaction = + | typeof client + | Parameters[0]>[0]; + +export { client, testDb, disconnect, type DatabaseOrTransaction }; diff --git a/src/db/index.ts b/src/db/index.ts index 9984eab722..ad941abdc6 100644 --- a/src/db/index.ts +++ b/src/db/index.ts @@ -1,33 +1,2 @@ -import { drizzle } from 'drizzle-orm/node-postgres'; -import { Pool } from 'pg'; - -import { Env } from '@/lib/server'; - -const pool = new Pool({ - connectionString: process.env.DATABASE_URL, -}); - -let testDb: ReturnType | undefined = undefined; - -if (process.env.NODE_ENV === 'test') { - if ( - !Env.DATABASE_URL?.includes('test') || - !Env.DATABASE_URL?.includes('localhost') - ) { - throw new Error('DATABASE_URL is not a test database'); - } - - testDb = drizzle(pool); -} - -const db = process.env.NODE_ENV === 'test' ? testDb! : drizzle(pool); - -const disconnect = async () => { - await pool.end(); -}; - -export type DB_OR_TX = - | typeof db - | Parameters[0]>[0]; - -export { db, testDb, disconnect }; +export * from './client'; +export * from './schema';