implement column in db to track using own key

Author: ahmedkhaleel2004

src/app/_actions/cache.ts

@@ -3,6 +3,7 @@
 import { db } from "~/server/db";
 import { eq, and } from "drizzle-orm";
 import { diagramCache } from "~/server/db/schema";
+import { sql } from "drizzle-orm";
 
 export async function getCachedDiagram(username: string, repo: string) {
   try {
@@ -43,24 +44,45 @@ export async function cacheDiagramAndExplanation(
   repo: string,
   diagram: string,
   explanation: string,
+  usedOwnKey = false,
 ) {
   try {
     await db
       .insert(diagramCache)
       .values({
         username,
         repo,
-        explanation,
         diagram,
+        explanation,
+        usedOwnKey,
       })
       .onConflictDoUpdate({
         target: [diagramCache.username, diagramCache.repo],
         set: {
           diagram,
+          explanation,
+          usedOwnKey,
           updatedAt: new Date(),
         },
       });
   } catch (error) {
     console.error("Error caching diagram:", error);
   }
 }
+
+export async function getDiagramStats() {
+  try {
+    const stats = await db
+      .select({
+        totalDiagrams: sql`COUNT(*)`,
+        ownKeyUsers: sql`COUNT(CASE WHEN ${diagramCache.usedOwnKey} = true THEN 1 END)`,
+        freeUsers: sql`COUNT(CASE WHEN ${diagramCache.usedOwnKey} = false THEN 1 END)`,
+      })
+      .from(diagramCache);
+
+    return stats[0];
+  } catch (error) {
+    console.error("Error getting diagram stats:", error);
+    return null;
+  }
+}

src/hooks/useDiagram.ts

@@ -220,12 +220,14 @@ export function useDiagram(username: string, repo: string) {
 
   useEffect(() => {
     if (state.status === "complete" && state.diagram) {
-      // Cache the completed diagram
+      // Cache the completed diagram with the usedOwnKey flag
+      const hasApiKey = !!localStorage.getItem("openrouter_key");
       void cacheDiagramAndExplanation(
         username,
         repo,
         state.diagram,
         state.explanation ?? "No explanation provided",
+        hasApiKey,
       );
       setDiagram(state.diagram);
       void getLastGeneratedDate(username, repo).then((date) =>

src/server/db/schema.ts

@@ -7,6 +7,7 @@ import {
   timestamp,
   varchar,
   primaryKey,
+  boolean,
 } from "drizzle-orm/pg-core";
 
 /**
@@ -32,6 +33,7 @@ export const diagramCache = createTable(
     updatedAt: timestamp("updated_at", { withTimezone: true }).$onUpdate(
       () => new Date(),
     ),
+    usedOwnKey: boolean("used_own_key").default(false),
   },
   (table) => ({
     pk: primaryKey({ columns: [table.username, table.repo] }),