More usage queries (#31)

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';
 

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';
 

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<typeof usageSchema>;
 
-/**
- * getUsage
- */
-
 type UsageRecord = Partial<Record<TelemetryEventName, Usage>>;
 
 export const getUsage = async ({
@@ -64,7 +62,7 @@ export const getUsage = async ({
 }: {
   orgId?: string | null;
   timePeriod: TimePeriod;
-}): Promise<UsageRecord | undefined> => {
+}): Promise<UsageRecord> => {
   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<typeof developerUsageSchema> & {
+  user: db.User;
+};
+
+export const getDeveloperUsage = async ({
+  orgId,
+  timePeriod,
+}: {
+  orgId?: string | null;
+  timePeriod: TimePeriod;
+}): Promise<DeveloperUsage[]> => {
+  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<string, db.User>,
+  );
+
+  return developerUsages
+    .map((usage) => ({
+      ...usage,
+      user: users[usage.userId],
+    }))
+    .filter((usage): usage is DeveloperUsage => !!usage.user);
 };

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<void> {
   await db.insert(auditLogs).values(values);

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';

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';

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';

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,

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,

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<Developer>[] = [
-    {
-      accessorKey: 'name',
-      header: 'Developer',
-      cell: ({ row }) => {
-        const developer = row.original;
+  const { orgId } = useAuth();
 
-        return (
-          <button
-            onClick={() => onClick(developer)}
-            className="text-left font-medium text-primary hover:underline"
-          >
-            {developer.name}
-          </button>
-        );
-      },
+  const { data = [] } = useQuery({
+    queryKey: ['developers'],
+    queryFn: () => getDeveloperUsage({ orgId, timePeriod: 30 }),
+    enabled: !!orgId,
+  });
+
+  const columns: ColumnDef<DeveloperUsage>[] = [
+    {
+      accessorKey: 'user.name',
+      header: 'Developer',
+      cell: ({ row }) => (
+        <button
+          onClick={() => onDeveloperSelected(row.original.user)}
+          className="text-left font-medium text-primary hover:underline"
+        >
+          {row.original.user.name}
+        </button>
+      ),
     },
     {
-      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 <DataTable columns={developerColumns} data={mockDevelopers} />;
+  return <DataTable columns={columns} data={data} />;
 };

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' ? (
           <Developers
-            onClick={({ id, name }: Developer) => {
+            onDeveloperSelected={({ id, name }: User) => {
               setFilter({ type: 'developer', id, name });
               setViewMode('tasks');
             }}

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;

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,

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;

src/components/audit-logs/AuditLogDrawer.tsx

@@ -1,4 +1,4 @@
-import type { AuditLog } from '@/db/schema';
+import type { AuditLog } from '@/db';
 import {
   Drawer,
   DrawerContent,

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<typeof drizzle> | 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<Parameters<typeof client.transaction>[0]>[0];
+
+export { client, testDb, disconnect, type DatabaseOrTransaction };

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<typeof drizzle> | 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<Parameters<typeof db.transaction>[0]>[0];
-
-export { db, testDb, disconnect };
+export * from './client';
+export * from './schema';