fix: optimize tagging db queries (#2287)

* fix: optimize tagging db queries

* review

* parallel queries

* refactoring

2 files changed, 206 insertions, 86 deletions

diff --git a/packages/trpc/routers/bookmarks.test.ts b/packages/trpc/routers/bookmarks.test.ts
index c272e015..7e0e65e0 100644
--- a/packages/trpc/routers/bookmarks.test.ts
+++ b/packages/trpc/routers/bookmarks.test.ts
@@ -331,6 +331,130 @@ describe("Bookmark Routes", () => {
     ).rejects.toThrow(/You must provide either a tagId or a tagName/);
   });
 
+  test<CustomTestContext>("update tags - comprehensive edge cases", async ({
+    apiCallers,
+  }) => {
+    const api = apiCallers[0].bookmarks;
+
+    // Create two bookmarks
+    const bookmark1 = await api.createBookmark({
+      url: "https://bookmark1.com",
+      type: BookmarkTypes.LINK,
+    });
+    const bookmark2 = await api.createBookmark({
+      url: "https://bookmark2.com",
+      type: BookmarkTypes.LINK,
+    });
+
+    // Test 1: Attach tags by name to bookmark1 (creates new tags)
+    await api.updateTags({
+      bookmarkId: bookmark1.id,
+      attach: [{ tagName: "existing-tag" }, { tagName: "shared-tag" }],
+      detach: [],
+    });
+
+    let b1 = await api.getBookmark({ bookmarkId: bookmark1.id });
+    expect(b1.tags.map((t) => t.name).sort()).toEqual([
+      "existing-tag",
+      "shared-tag",
+    ]);
+
+    const existingTagId = b1.tags.find((t) => t.name === "existing-tag")!.id;
+    const sharedTagId = b1.tags.find((t) => t.name === "shared-tag")!.id;
+
+    // Test 2: Attach existing tag by ID to bookmark2 (tag already exists in DB from bookmark1)
+    await api.updateTags({
+      bookmarkId: bookmark2.id,
+      attach: [{ tagId: existingTagId }],
+      detach: [],
+    });
+
+    let b2 = await api.getBookmark({ bookmarkId: bookmark2.id });
+    expect(b2.tags.map((t) => t.name)).toEqual(["existing-tag"]);
+
+    // Test 3: Attach existing tag by NAME to bookmark2 (tag already exists in DB)
+    await api.updateTags({
+      bookmarkId: bookmark2.id,
+      attach: [{ tagName: "shared-tag" }],
+      detach: [],
+    });
+
+    b2 = await api.getBookmark({ bookmarkId: bookmark2.id });
+    expect(b2.tags.map((t) => t.name).sort()).toEqual([
+      "existing-tag",
+      "shared-tag",
+    ]);
+
+    // Test 4: Re-attaching the same tag (idempotency) - should be no-op
+    await api.updateTags({
+      bookmarkId: bookmark2.id,
+      attach: [{ tagId: existingTagId }],
+      detach: [],
+    });
+
+    b2 = await api.getBookmark({ bookmarkId: bookmark2.id });
+    expect(b2.tags.map((t) => t.name).sort()).toEqual([
+      "existing-tag",
+      "shared-tag",
+    ]);
+
+    // Test 5: Detach non-existent tag by name (should be no-op)
+    await api.updateTags({
+      bookmarkId: bookmark2.id,
+      attach: [],
+      detach: [{ tagName: "non-existent-tag" }],
+    });
+
+    b2 = await api.getBookmark({ bookmarkId: bookmark2.id });
+    expect(b2.tags.map((t) => t.name).sort()).toEqual([
+      "existing-tag",
+      "shared-tag",
+    ]);
+
+    // Test 6: Mixed attach/detach with pre-existing tags
+    await api.updateTags({
+      bookmarkId: bookmark2.id,
+      attach: [{ tagName: "new-tag" }, { tagId: sharedTagId }], // sharedTagId already attached
+      detach: [{ tagName: "existing-tag" }],
+    });
+
+    b2 = await api.getBookmark({ bookmarkId: bookmark2.id });
+    expect(b2.tags.map((t) => t.name).sort()).toEqual([
+      "new-tag",
+      "shared-tag",
+    ]);
+
+    // Test 7: Detach by ID and re-attach by name in same operation
+    await api.updateTags({
+      bookmarkId: bookmark2.id,
+      attach: [{ tagName: "new-tag" }], // Already exists, should be idempotent
+      detach: [{ tagId: sharedTagId }],
+    });
+
+    b2 = await api.getBookmark({ bookmarkId: bookmark2.id });
+    expect(b2.tags.map((t) => t.name).sort()).toEqual(["new-tag"]);
+
+    // Verify bookmark1 still has its original tags (operations on bookmark2 didn't affect it)
+    b1 = await api.getBookmark({ bookmarkId: bookmark1.id });
+    expect(b1.tags.map((t) => t.name).sort()).toEqual([
+      "existing-tag",
+      "shared-tag",
+    ]);
+
+    // Test 8: Attach same tag multiple times in one operation (deduplication)
+    await api.updateTags({
+      bookmarkId: bookmark1.id,
+      attach: [{ tagName: "duplicate-test" }, { tagName: "duplicate-test" }],
+      detach: [],
+    });
+
+    b1 = await api.getBookmark({ bookmarkId: bookmark1.id });
+    const duplicateTagCount = b1.tags.filter(
+      (t) => t.name === "duplicate-test",
+    ).length;
+    expect(duplicateTagCount).toEqual(1); // Should only be attached once
+  });
+
   test<CustomTestContext>("update bookmark text", async ({ apiCallers }) => {
     const api = apiCallers[0].bookmarks;
     const createdBookmark = await api.createBookmark({
diff --git a/packages/trpc/routers/bookmarks.ts b/packages/trpc/routers/bookmarks.ts
index 509aa1d2..4f18e11b 100644
--- a/packages/trpc/routers/bookmarks.ts
+++ b/packages/trpc/routers/bookmarks.ts
@@ -1,6 +1,5 @@
 import { experimental_trpcMiddleware, TRPCError } from "@trpc/server";
 import { and, eq, gt, inArray, lt, or } from "drizzle-orm";
-import invariant from "tiny-invariant";
 import { z } from "zod";
 
 import type { ZBookmarkContent } from "@karakeep/shared/types/bookmarks";
@@ -714,36 +713,79 @@ export const bookmarksAppRouter = router({
     )
     .use(ensureBookmarkOwnership)
     .mutation(async ({ input, ctx }) => {
-      const res = await ctx.db.transaction(async (tx) => {
-        // Detaches
-        const idsToRemove: string[] = [];
-        if (input.detach.length > 0) {
-          const namesToRemove: string[] = [];
-          input.detach.forEach((detachInfo) => {
-            if (detachInfo.tagId) {
-              idsToRemove.push(detachInfo.tagId);
-            }
-            if (detachInfo.tagName) {
-              namesToRemove.push(detachInfo.tagName);
-            }
-          });
+      // Helper function to fetch tag IDs from a list of tag identifiers
+      const fetchTagIds = async (
+        tagIdentifiers: { tagId?: string; tagName?: string }[],
+      ): Promise<string[]> => {
+        const tagIds = tagIdentifiers.flatMap((t) =>
+          t.tagId ? [t.tagId] : [],
+        );
+        const tagNames = tagIdentifiers.flatMap((t) =>
+          t.tagName ? [t.tagName] : [],
+        );
 
-          if (namesToRemove.length > 0) {
-            (
-              await tx.query.bookmarkTags.findMany({
-                where: and(
-                  eq(bookmarkTags.userId, ctx.user.id),
-                  inArray(bookmarkTags.name, namesToRemove),
-                ),
-                columns: {
-                  id: true,
-                },
-              })
-            ).forEach((tag) => {
-              idsToRemove.push(tag.id);
-            });
-          }
+        // Fetch tag IDs in parallel
+        const [byIds, byNames] = await Promise.all([
+          tagIds.length > 0
+            ? ctx.db
+                .select({ id: bookmarkTags.id })
+                .from(bookmarkTags)
+                .where(
+                  and(
+                    eq(bookmarkTags.userId, ctx.user.id),
+                    inArray(bookmarkTags.id, tagIds),
+                  ),
+                )
+            : Promise.resolve([]),
+          tagNames.length > 0
+            ? ctx.db
+                .select({ id: bookmarkTags.id })
+                .from(bookmarkTags)
+                .where(
+                  and(
+                    eq(bookmarkTags.userId, ctx.user.id),
+                    inArray(bookmarkTags.name, tagNames),
+                  ),
+                )
+            : Promise.resolve([]),
+        ]);
+
+        // Union results and deduplicate tag IDs
+        const results = [...byIds, ...byNames];
+        return [...new Set(results.map((t) => t.id))];
+      };
+
+      // Normalize tag names and create new tags outside transaction to reduce transaction duration
+      const normalizedAttachTags = input.attach.map((tag) => ({
+        tagId: tag.tagId,
+        tagName: tag.tagName ? normalizeTagName(tag.tagName) : undefined,
+      }));
+
+      {
+        // Create new tags
+        const toAddTagNames = normalizedAttachTags
+          .flatMap((i) => (i.tagName ? [i.tagName] : []))
+          .filter((n) => n.length > 0); // drop empty results
+
+        if (toAddTagNames.length > 0) {
+          await ctx.db
+            .insert(bookmarkTags)
+            .values(
+              toAddTagNames.map((name) => ({ name, userId: ctx.user.id })),
+            )
+            .onConflictDoNothing();
+        }
+      }
+
+      // Fetch tag IDs for attachment/detachment now that we know that they all exist
+      const [allIdsToAttach, idsToRemove] = await Promise.all([
+        fetchTagIds(normalizedAttachTags),
+        fetchTagIds(input.detach),
+      ]);
 
+      const res = await ctx.db.transaction(async (tx) => {
+        // Detaches
+        if (idsToRemove.length > 0) {
           await tx
             .delete(tagsOnBookmarks)
             .where(
@@ -754,67 +796,21 @@ export const bookmarksAppRouter = router({
             );
         }
 
-        if (input.attach.length == 0) {
-          return {
-            bookmarkId: input.bookmarkId,
-            attached: [],
-            detached: idsToRemove,
-          };
-        }
-
-        const toAddTagNames = input.attach
-          .flatMap((i) => (i.tagName ? [i.tagName] : []))
-          .map(normalizeTagName) // strip leading #
-          .filter((n) => n.length > 0); // drop empty results
-
-        const toAddTagIds = input.attach.flatMap((i) =>
-          i.tagId ? [i.tagId] : [],
-        );
-
-        // New Tags
-        if (toAddTagNames.length > 0) {
+        // Attach tags
+        if (allIdsToAttach.length > 0) {
           await tx
-            .insert(bookmarkTags)
+            .insert(tagsOnBookmarks)
             .values(
-              toAddTagNames.map((name) => ({ name, userId: ctx.user.id })),
+              allIdsToAttach.map((i) => ({
+                tagId: i,
+                bookmarkId: input.bookmarkId,
+                attachedBy: "human" as const,
+              })),
             )
-            .onConflictDoNothing()
-            .returning();
+            .onConflictDoNothing();
         }
 
-        // If there is nothing to add, the "or" statement will become useless and
-        // the query below will simply select all the existing tags for this user and assign them to the bookmark
-        invariant(toAddTagNames.length > 0 || toAddTagIds.length > 0);
-        const allIds = (
-          await tx.query.bookmarkTags.findMany({
-            where: and(
-              eq(bookmarkTags.userId, ctx.user.id),
-              or(
-                toAddTagIds.length > 0
-                  ? inArray(bookmarkTags.id, toAddTagIds)
-                  : undefined,
-                toAddTagNames.length > 0
-                  ? inArray(bookmarkTags.name, toAddTagNames)
-                  : undefined,
-              ),
-            ),
-            columns: {
-              id: true,
-            },
-          })
-        ).map((t) => t.id);
-
-        await tx
-          .insert(tagsOnBookmarks)
-          .values(
-            allIds.map((i) => ({
-              tagId: i,
-              bookmarkId: input.bookmarkId,
-              attachedBy: "human" as const,
-              userId: ctx.user.id,
-            })),
-          )
-          .onConflictDoNothing();
+        // Update bookmark modified timestamp
         await tx
           .update(bookmarks)
           .set({ modifiedAt: new Date() })
@@ -827,7 +823,7 @@ export const bookmarksAppRouter = router({
 
         return {
           bookmarkId: input.bookmarkId,
-          attached: allIds,
+          attached: allIdsToAttach,
           detached: idsToRemove,
         };
       });