feat: read usage from cache priority (#963)

olzhik11 · web-flow · commit f2e2cb634f37 · 2025-10-17T18:47:09.000+01:00
* feat: read usage from cache priority

* feat: remove log

* feat: conditionally fetch workspace usage
diff --git a/frontend/app/project/[projectId]/layout.tsx b/frontend/app/project/[projectId]/layout.tsx
@@ -1,6 +1,5 @@
 import "@/app/globals.css";
 
-import { eq } from "drizzle-orm";
 import { cookies } from "next/headers";
 import { redirect } from "next/navigation";
 import { getServerSession } from "next-auth";
@@ -13,79 +12,11 @@ import ProjectUsageBanner from "@/components/project/usage-banner";
 import { SidebarInset, SidebarProvider } from "@/components/ui/sidebar";
 import { ProjectContextProvider } from "@/contexts/project-context";
 import { UserContextProvider } from "@/contexts/user-context";
+import { getProjectDetails } from "@/lib/actions/project";
 import { getProjectsByWorkspace } from "@/lib/actions/projects";
-import { getWorkspaceInfo, getWorkspaceUsage } from "@/lib/actions/workspace";
+import { getWorkspaceInfo } from "@/lib/actions/workspace";
 import { authOptions } from "@/lib/auth";
-import { db } from "@/lib/db/drizzle";
-import { projects, subscriptionTiers, workspaces } from "@/lib/db/migrations/schema";
 import { Feature, isFeatureEnabled } from "@/lib/features/features";
-import { GetProjectResponse } from "@/lib/workspaces/types";
-
-async function getProjectDetails(projectId: string): Promise<GetProjectResponse> {
-  const projectResult = await db
-    .select({
-      id: projects.id,
-      name: projects.name,
-      workspaceId: projects.workspaceId,
-    })
-    .from(projects)
-    .where(eq(projects.id, projectId))
-    .limit(1);
-
-  if (projectResult.length === 0) {
-    throw new Error("Project not found");
-  }
-  const project = projectResult[0];
-
-  const workspaceResult = await db
-    .select({
-      id: workspaces.id,
-      tierId: workspaces.tierId,
-    })
-    .from(workspaces)
-    .where(eq(workspaces.id, project.workspaceId))
-    .limit(1);
-
-  if (workspaceResult.length === 0) {
-    throw new Error("Workspace not found for project");
-  }
-  const workspace = workspaceResult[0];
-  const usageResult = await getWorkspaceUsage(project.workspaceId);
-
-  const tierResult = await db
-    .select({
-      name: subscriptionTiers.name,
-      stepsLimit: subscriptionTiers.steps,
-      bytesLimit: subscriptionTiers.bytesIngested,
-    })
-    .from(subscriptionTiers)
-    .where(eq(subscriptionTiers.id, workspace.tierId))
-    .limit(1);
-
-  if (tierResult.length === 0) {
-    throw new Error("Subscription tier not found for workspace");
-  }
-  const tier = tierResult[0];
-
-  // Convert bytes to GB (1 GB = 1024^3 bytes)
-  const bytesToGB = (bytes: number): number => bytes / (1024 * 1024 * 1024);
-
-  const gbUsedThisMonth = bytesToGB(
-    Number(
-      usageResult.spansBytesIngested + usageResult.browserSessionEventsBytesIngested + usageResult.eventsBytesIngested
-    )
-  );
-  const gbLimit = bytesToGB(Number(tier.bytesLimit));
-
-  return {
-    id: project.id,
-    name: project.name,
-    workspaceId: project.workspaceId,
-    gbUsedThisMonth,
-    gbLimit,
-    isFreeTier: tier.name.toLowerCase().trim() === "free",
-  };
-}
 
 export default async function ProjectIdLayout(props: { children: ReactNode; params: Promise<{ projectId: string }> }) {
   const params = await props.params;
diff --git a/frontend/lib/actions/project/index.ts b/frontend/lib/actions/project/index.ts
@@ -1,10 +1,11 @@
 import { eq } from "drizzle-orm";
 import { z } from "zod/v4";
 
+import { getWorkspaceUsage } from "@/lib/actions/workspace";
 import { cache, PROJECT_API_KEY_CACHE_KEY, PROJECT_CACHE_KEY } from "@/lib/cache";
 import { clickhouseClient } from "@/lib/clickhouse/client";
 import { db } from "@/lib/db/drizzle";
-import { projectApiKeys, projects } from "@/lib/db/migrations/schema";
+import { projectApiKeys, projects, subscriptionTiers, workspaces } from "@/lib/db/migrations/schema";
 
 export const DeleteProjectSchema = z.object({
   projectId: z.uuid(),
@@ -162,3 +163,87 @@ async function deleteProjectWorkspaceInfoFromCache(projectId: string) {
   const cacheKey = `${PROJECT_CACHE_KEY}:${projectId}`;
   await cache.remove(cacheKey);
 }
+
+export const getProjectDetails = async (
+  projectId: string
+): Promise<{
+  id: string;
+  name: string;
+  workspaceId: string;
+  gbUsedThisMonth: number;
+  gbLimit: number;
+  isFreeTier: boolean;
+}> => {
+  const projectResult = await db
+    .select({
+      id: projects.id,
+      name: projects.name,
+      workspaceId: projects.workspaceId,
+    })
+    .from(projects)
+    .where(eq(projects.id, projectId))
+    .limit(1);
+
+  if (projectResult.length === 0) {
+    throw new Error("Project not found");
+  }
+
+  const project = projectResult[0];
+
+  const workspaceResult = await db
+    .select({
+      id: workspaces.id,
+      tierId: workspaces.tierId,
+    })
+    .from(workspaces)
+    .where(eq(workspaces.id, project.workspaceId))
+    .limit(1);
+
+  if (workspaceResult.length === 0) {
+    throw new Error("Workspace not found for project");
+  }
+  const workspace = workspaceResult[0];
+
+  const tierResult = await db
+    .select({
+      name: subscriptionTiers.name,
+      stepsLimit: subscriptionTiers.steps,
+      bytesLimit: subscriptionTiers.bytesIngested,
+    })
+    .from(subscriptionTiers)
+    .where(eq(subscriptionTiers.id, workspace.tierId))
+    .limit(1);
+
+  if (tierResult.length === 0) {
+    throw new Error("Subscription tier not found for workspace");
+  }
+  const tier = tierResult[0];
+  const isFreeTier = tier.name.toLowerCase().trim() === "free";
+
+  const bytesToGB = (bytes: number): number => bytes / (1024 * 1024 * 1024);
+  const gbLimit = bytesToGB(Number(tier.bytesLimit));
+
+  if (!isFreeTier) {
+    return {
+      id: project.id,
+      name: project.name,
+      workspaceId: project.workspaceId,
+      // not used in ui
+      gbUsedThisMonth: 0,
+      gbLimit,
+      isFreeTier,
+    };
+  }
+
+  const usageResult = await getWorkspaceUsage(project.workspaceId);
+  const gbUsedThisMonth = bytesToGB(usageResult.totalBytesIngested);
+
+  return {
+    id: project.id,
+    name: project.name,
+    workspaceId: project.workspaceId,
+    gbUsedThisMonth,
+    gbLimit,
+    isFreeTier,
+  };
+};
diff --git a/frontend/lib/actions/workspace/index.ts b/frontend/lib/actions/workspace/index.ts
@@ -5,6 +5,7 @@ import { z } from "zod/v4";
 import { deleteProject } from "@/lib/actions/project";
 import { checkUserWorkspaceRole } from "@/lib/actions/workspace/utils";
 import { completeMonthsElapsed } from "@/lib/actions/workspaces/utils";
+import { cache, WORKSPACE_BYTES_USAGE_CACHE_KEY } from "@/lib/cache";
 import { clickhouseClient } from "@/lib/clickhouse/client";
 import { db } from "@/lib/db/drizzle";
 import { membersOfWorkspaces, projects, subscriptionTiers, users, workspaces } from "@/lib/db/migrations/schema";
@@ -129,13 +130,6 @@ export const getWorkspaceInfo = async (workspaceId: string): Promise<Workspace>
 };
 
 export const getWorkspaceUsage = async (workspaceId: string): Promise<WorkspaceUsage> => {
-  const projectIds = await db.query.projects.findMany({
-    where: eq(projects.workspaceId, workspaceId),
-    columns: {
-      id: true,
-    },
-  });
-
   const resetTime = await db.query.workspaces.findFirst({
     where: eq(workspaces.id, workspaceId),
     columns: {
@@ -147,18 +141,38 @@ export const getWorkspaceUsage = async (workspaceId: string): Promise<WorkspaceU
     throw new Error("Workspace not found");
   }
 
+  const resetTimeDate = new Date(resetTime.resetTime);
+  const latestResetTime = addMonths(resetTimeDate, completeMonthsElapsed(resetTimeDate, new Date()));
+
+  const cacheKey = `${WORKSPACE_BYTES_USAGE_CACHE_KEY}:${workspaceId}`;
+  try {
+    const cachedUsage = await cache.get<number>(cacheKey);
+    if (cachedUsage !== null) {
+      return {
+        totalBytesIngested: Number(cachedUsage),
+        resetTime: latestResetTime,
+      };
+    }
+  } catch (error) {
+    // If cache fails, continue to ClickHouse query
+    console.error("Error reading from cache:", error);
+  }
+
+  // Cache miss - query ClickHouse for the actual breakdown
+  const projectIds = await db.query.projects.findMany({
+    where: eq(projects.workspaceId, workspaceId),
+    columns: {
+      id: true,
+    },
+  });
+
   if (projectIds.length === 0) {
     return {
-      spansBytesIngested: 0,
-      browserSessionEventsBytesIngested: 0,
-      eventsBytesIngested: 0,
-      resetTime: new Date(resetTime.resetTime),
+      totalBytesIngested: 0,
+      resetTime: latestResetTime,
     };
   }
 
-  const resetTimeDate = new Date(resetTime.resetTime);
-
-  const latestResetTime = addMonths(resetTimeDate, completeMonthsElapsed(resetTimeDate, new Date()));
   const query = `WITH spans_bytes_ingested AS (
       SELECT
         SUM(spans.size_bytes) as spans_bytes_ingested
@@ -205,10 +219,13 @@ export const getWorkspaceUsage = async (workspaceId: string): Promise<WorkspaceU
     throw new Error("Error getting workspace usage");
   }
 
+  const totalBytesIngested =
+    Number(result[0].spans_bytes_ingested) +
+    Number(result[0].browser_session_events_bytes_ingested) +
+    Number(result[0].events_bytes_ingested);
+
   return {
-    spansBytesIngested: Number(result[0].spans_bytes_ingested),
-    browserSessionEventsBytesIngested: Number(result[0].browser_session_events_bytes_ingested),
-    eventsBytesIngested: Number(result[0].events_bytes_ingested),
+    totalBytesIngested,
     resetTime: latestResetTime,
   };
 };
diff --git a/frontend/lib/cache.ts b/frontend/lib/cache.ts
@@ -118,5 +118,6 @@ export const PROJECT_API_KEY_CACHE_KEY = "project_api_key";
 export const PROJECT_EVALUATORS_BY_PATH_CACHE_KEY = "project_evaluators_by_path";
 export const PROJECT_CACHE_KEY = "project";
 export const WORKSPACE_LIMITS_CACHE_KEY = "workspace_limits";
+export const WORKSPACE_BYTES_USAGE_CACHE_KEY = "workspace_bytes_usage";
 export const TRACE_CHATS_CACHE_KEY = "trace_chats";
 export const TRACE_SUMMARIES_CACHE_KEY = "trace_summaries";
diff --git a/frontend/lib/usage/workspace-stats.ts b/frontend/lib/usage/workspace-stats.ts
@@ -38,9 +38,7 @@ export async function getWorkspaceStats(workspaceId: string): Promise<WorkspaceS
 
   const usage = await getWorkspaceUsage(workspaceId);
 
-  const gbUsedThisMonth = bytesToGB(
-    Number(usage.spansBytesIngested + usage.browserSessionEventsBytesIngested + usage.eventsBytesIngested)
-  );
+  const gbUsedThisMonth = bytesToGB(usage.totalBytesIngested);
   const gbLimit = bytesToGB(Number(limits.bytesLimit));
 
   // Calculate GB overages
diff --git a/frontend/lib/workspaces/types.ts b/frontend/lib/workspaces/types.ts
@@ -41,25 +41,13 @@ export interface WorkspaceInvitation {
   createdAt: string;
 }
 
-export type GetProjectResponse = {
-  id: string;
-  name: string;
-  workspaceId: string;
-  // New GB-based usage fields
-  gbUsedThisMonth: number;
-  gbLimit: number;
-  isFreeTier: boolean;
-};
-
 export interface ProjectStats {
   datasetsCount: number;
   spansCount: number;
   evaluationsCount: number;
 }
 
 export interface WorkspaceUsage {
-  spansBytesIngested: number;
-  browserSessionEventsBytesIngested: number;
-  eventsBytesIngested: number;
+  totalBytesIngested: number;
   resetTime: Date;
 }
