fix: optimize memory usage of tag listing

2 files changed, 137 insertions, 51 deletions

diff --git a/packages/trpc/models/tags.ts b/packages/trpc/models/tags.ts
index a91dcbdf..dadb20f7 100644
--- a/packages/trpc/models/tags.ts
+++ b/packages/trpc/models/tags.ts
@@ -1,5 +1,5 @@
 import { TRPCError } from "@trpc/server";
-import { and, eq, inArray, notExists } from "drizzle-orm";
+import { and, count, eq, inArray, notExists } from "drizzle-orm";
 import { z } from "zod";
 
 import type { ZAttachedByEnum } from "@karakeep/shared/types/tags";
@@ -70,40 +70,45 @@ export class Tag implements PrivacyAware {
     }
   }
 
-  static async getAll(ctx: AuthedContext): Promise<Tag[]> {
-    const tags = await ctx.db.query.bookmarkTags.findMany({
-      where: eq(bookmarkTags.userId, ctx.user.id),
-    });
-
-    return tags.map((t) => new Tag(ctx, t));
-  }
-
   static async getAllWithStats(ctx: AuthedContext) {
-    const tags = await ctx.db.query.bookmarkTags.findMany({
-      where: eq(bookmarkTags.userId, ctx.user.id),
-      with: {
-        tagsOnBookmarks: {
-          columns: {
-            attachedBy: true,
-          },
-        },
-      },
-    });
+    const tags = await ctx.db
+      .select({
+        id: bookmarkTags.id,
+        name: bookmarkTags.name,
+        attachedBy: tagsOnBookmarks.attachedBy,
+        count: count(),
+      })
+      .from(bookmarkTags)
+      .leftJoin(tagsOnBookmarks, eq(bookmarkTags.id, tagsOnBookmarks.tagId))
+      .where(and(eq(bookmarkTags.userId, ctx.user.id)))
+      .groupBy(bookmarkTags.id, tagsOnBookmarks.attachedBy);
 
-    return tags.map(({ tagsOnBookmarks, ...rest }) => ({
-      ...rest,
-      numBookmarks: tagsOnBookmarks.length,
-      numBookmarksByAttachedType: tagsOnBookmarks.reduce<
-        Record<ZAttachedByEnum, number>
-      >(
-        (acc, curr) => {
-          if (curr.attachedBy) {
-            acc[curr.attachedBy]++;
-          }
-          return acc;
-        },
-        { ai: 0, human: 0 },
-      ),
+    if (tags.length === 0) {
+      return [];
+    }
+
+    const tagsById = tags.reduce<
+      Record<
+        string,
+        {
+          id: string;
+          name: string;
+          attachedBy: "ai" | "human" | null;
+          count: number;
+        }[]
+      >
+    >((acc, curr) => {
+      if (!acc[curr.id]) {
+        acc[curr.id] = [];
+      }
+      acc[curr.id].push(curr);
+      return acc;
+    }, {});
+
+    return Object.entries(tagsById).map(([k, t]) => ({
+      id: k,
+      name: t[0].name,
+      ...Tag._aggregateStats(t),
     }));
   }
 
@@ -310,12 +315,34 @@ export class Tag implements PrivacyAware {
     }
   }
 
+  static _aggregateStats(
+    res: { attachedBy: "ai" | "human" | null; count: number }[],
+  ) {
+    const numBookmarksByAttachedType = res.reduce<
+      Record<ZAttachedByEnum, number>
+    >(
+      (acc, curr) => {
+        if (curr.attachedBy) {
+          acc[curr.attachedBy] += curr.count;
+        }
+        return acc;
+      },
+      { ai: 0, human: 0 },
+    );
+    return {
+      numBookmarks:
+        numBookmarksByAttachedType.ai + numBookmarksByAttachedType.human,
+      numBookmarksByAttachedType,
+    };
+  }
+
   async getStats(): Promise<z.infer<typeof zGetTagResponseSchema>> {
     const res = await this.ctx.db
       .select({
         id: bookmarkTags.id,
         name: bookmarkTags.name,
         attachedBy: tagsOnBookmarks.attachedBy,
+        count: count(),
       })
       .from(bookmarkTags)
       .leftJoin(tagsOnBookmarks, eq(bookmarkTags.id, tagsOnBookmarks.tagId))
@@ -324,32 +351,17 @@ export class Tag implements PrivacyAware {
           eq(bookmarkTags.id, this.tag.id),
           eq(bookmarkTags.userId, this.ctx.user.id),
         ),
-      );
+      )
+      .groupBy(tagsOnBookmarks.attachedBy);
 
     if (res.length === 0) {
       throw new TRPCError({ code: "NOT_FOUND" });
     }
 
-    const numBookmarksByAttachedType = res.reduce<
-      Record<ZAttachedByEnum, number>
-    >(
-      (acc, curr) => {
-        if (curr.attachedBy) {
-          acc[curr.attachedBy]++;
-        }
-        return acc;
-      },
-      { ai: 0, human: 0 },
-    );
-
     return {
       id: res[0].id,
       name: res[0].name,
-      numBookmarks: Object.values(numBookmarksByAttachedType).reduce(
-        (s, a) => s + a,
-        0,
-      ),
-      numBookmarksByAttachedType,
+      ...Tag._aggregateStats(res),
     };
   }
 
diff --git a/packages/trpc/routers/tags.test.ts b/packages/trpc/routers/tags.test.ts
index a4d690ee..4004cc2c 100644
--- a/packages/trpc/routers/tags.test.ts
+++ b/packages/trpc/routers/tags.test.ts
@@ -20,6 +20,42 @@ describe("Tags Routes", () => {
     expect(res.numBookmarks).toBeGreaterThanOrEqual(0);
   });
 
+  test<CustomTestContext>("get tag returns bookmark stats", async ({
+    apiCallers,
+  }) => {
+    const tagsApi = apiCallers[0].tags;
+    const bookmarksApi = apiCallers[0].bookmarks;
+
+    const firstBookmark = await bookmarksApi.createBookmark({
+      url: "https://example.com/first",
+      type: BookmarkTypes.LINK,
+    });
+    const secondBookmark = await bookmarksApi.createBookmark({
+      url: "https://example.com/second",
+      type: BookmarkTypes.LINK,
+    });
+
+    const firstAttachment = await bookmarksApi.updateTags({
+      bookmarkId: firstBookmark.id,
+      attach: [{ tagName: "stats-tag" }],
+      detach: [],
+    });
+
+    const tagId = firstAttachment.attached[0];
+
+    await bookmarksApi.updateTags({
+      bookmarkId: secondBookmark.id,
+      attach: [{ tagId }],
+      detach: [],
+    });
+
+    const stats = await tagsApi.get({ tagId });
+
+    expect(stats.numBookmarks).toBe(2);
+    expect(stats.numBookmarksByAttachedType.human).toBe(2);
+    expect(stats.numBookmarksByAttachedType.ai).toBe(0);
+  });
+
   test<CustomTestContext>("get tag - not found", async ({ apiCallers }) => {
     const api = apiCallers[0].tags;
     await expect(() => api.get({ tagId: "nonExistentId" })).rejects.toThrow(
@@ -135,6 +171,44 @@ describe("Tags Routes", () => {
     expect(res.tags.some((tag) => tag.name === "tag2")).toBeTruthy();
   });
 
+  test<CustomTestContext>("list tags includes bookmark stats", async ({
+    apiCallers,
+  }) => {
+    const tagsApi = apiCallers[0].tags;
+    const bookmarksApi = apiCallers[0].bookmarks;
+
+    const firstBookmark = await bookmarksApi.createBookmark({
+      url: "https://example.com/list-first",
+      type: BookmarkTypes.LINK,
+    });
+    const secondBookmark = await bookmarksApi.createBookmark({
+      url: "https://example.com/list-second",
+      type: BookmarkTypes.LINK,
+    });
+
+    const firstAttachment = await bookmarksApi.updateTags({
+      bookmarkId: firstBookmark.id,
+      attach: [{ tagName: "list-stats-tag" }],
+      detach: [],
+    });
+
+    const tagId = firstAttachment.attached[0];
+
+    await bookmarksApi.updateTags({
+      bookmarkId: secondBookmark.id,
+      attach: [{ tagId }],
+      detach: [],
+    });
+
+    const list = await tagsApi.list();
+    const tagStats = list.tags.find((tag) => tag.id === tagId);
+
+    expect(tagStats).toBeDefined();
+    expect(tagStats!.numBookmarks).toBe(2);
+    expect(tagStats!.numBookmarksByAttachedType.human).toBe(2);
+    expect(tagStats!.numBookmarksByAttachedType.ai).toBe(0);
+  });
+
   test<CustomTestContext>("list tags - privacy", async ({ apiCallers }) => {
     const apiUser1 = apiCallers[0].tags;
     await apiUser1.create({ name: "user1Tag" });