feat: Add per user bookmark count quota

4 files changed, 308 insertions, 6 deletions

diff --git a/packages/trpc/routers/admin.ts b/packages/trpc/routers/admin.ts
index 91f4a34f..f46e17f2 100644
--- a/packages/trpc/routers/admin.ts
+++ b/packages/trpc/routers/admin.ts
@@ -18,8 +18,8 @@ import {
 } from "@karakeep/shared/queues";
 import { getSearchIdxClient } from "@karakeep/shared/search";
 import {
-  changeRoleSchema,
   resetPasswordSchema,
+  updateUserSchema,
   zAdminCreateUserSchema,
 } from "@karakeep/shared/types/admin";
 
@@ -331,18 +331,36 @@ export const adminAppRouter = router({
     .mutation(async ({ input, ctx }) => {
       return createUser(input, ctx, input.role);
     }),
-  changeRole: adminProcedure
-    .input(changeRoleSchema)
+  updateUser: adminProcedure
+    .input(updateUserSchema)
     .mutation(async ({ input, ctx }) => {
       if (ctx.user.id == input.userId) {
         throw new TRPCError({
           code: "BAD_REQUEST",
-          message: "Cannot change own role",
+          message: "Cannot update own user",
         });
       }
+
+      const updateData: Partial<typeof users.$inferInsert> = {};
+
+      if (input.role !== undefined) {
+        updateData.role = input.role;
+      }
+
+      if (input.bookmarkQuota !== undefined) {
+        updateData.bookmarkQuota = input.bookmarkQuota;
+      }
+
+      if (Object.keys(updateData).length === 0) {
+        throw new TRPCError({
+          code: "BAD_REQUEST",
+          message: "No fields to update",
+        });
+      }
+
       const result = await ctx.db
         .update(users)
-        .set({ role: input.role })
+        .set(updateData)
         .where(eq(users.id, input.userId));
 
       if (!result.changes) {
diff --git a/packages/trpc/routers/bookmarks.test.ts b/packages/trpc/routers/bookmarks.test.ts
index 575b4d9a..87f34b31 100644
--- a/packages/trpc/routers/bookmarks.test.ts
+++ b/packages/trpc/routers/bookmarks.test.ts
@@ -5,6 +5,7 @@ import {
   bookmarkLinks,
   bookmarks,
   rssFeedImportsTable,
+  users,
 } from "@karakeep/db/schema";
 import { BookmarkTypes } from "@karakeep/shared/types/bookmarks";
 
@@ -550,4 +551,262 @@ describe("Bookmark Routes", () => {
     const emptyResult = await api.getBrokenLinks();
     expect(emptyResult.bookmarks.length).toEqual(0);
   });
+
+  describe("Bookmark Quotas", () => {
+    test<CustomTestContext>("create bookmark with no quota (unlimited)", async ({
+      apiCallers,
+    }) => {
+      const api = apiCallers[0].bookmarks;
+
+      // User should be able to create bookmarks without any quota restrictions
+      const bookmark1 = await api.createBookmark({
+        url: "https://example1.com",
+        type: BookmarkTypes.LINK,
+      });
+      expect(bookmark1.alreadyExists).toEqual(false);
+
+      const bookmark2 = await api.createBookmark({
+        url: "https://example2.com",
+        type: BookmarkTypes.LINK,
+      });
+      expect(bookmark2.alreadyExists).toEqual(false);
+
+      const bookmark3 = await api.createBookmark({
+        text: "Test text bookmark",
+        type: BookmarkTypes.TEXT,
+      });
+      expect(bookmark3.alreadyExists).toEqual(false);
+    });
+
+    test<CustomTestContext>("create bookmark with quota limit", async ({
+      apiCallers,
+      db,
+    }) => {
+      const user = await apiCallers[0].users.whoami();
+      const api = apiCallers[0].bookmarks;
+
+      // Set quota to 2 bookmarks for this user
+      await db
+        .update(users)
+        .set({ bookmarkQuota: 2 })
+        .where(eq(users.id, user.id));
+
+      // First bookmark should succeed
+      const bookmark1 = await api.createBookmark({
+        url: "https://example1.com",
+        type: BookmarkTypes.LINK,
+      });
+      expect(bookmark1.alreadyExists).toEqual(false);
+
+      // Second bookmark should succeed
+      const bookmark2 = await api.createBookmark({
+        url: "https://example2.com",
+        type: BookmarkTypes.LINK,
+      });
+      expect(bookmark2.alreadyExists).toEqual(false);
+
+      // Third bookmark should fail due to quota
+      await expect(() =>
+        api.createBookmark({
+          url: "https://example3.com",
+          type: BookmarkTypes.LINK,
+        }),
+      ).rejects.toThrow(
+        /Bookmark quota exceeded. You can only have 2 bookmarks./,
+      );
+    });
+
+    test<CustomTestContext>("create bookmark with quota limit - different types", async ({
+      apiCallers,
+      db,
+    }) => {
+      const user = await apiCallers[0].users.whoami();
+      const api = apiCallers[0].bookmarks;
+
+      // Set quota to 2 bookmarks for this user
+      await db
+        .update(users)
+        .set({ bookmarkQuota: 2 })
+        .where(eq(users.id, user.id));
+
+      // Create one link bookmark
+      await api.createBookmark({
+        url: "https://example1.com",
+        type: BookmarkTypes.LINK,
+      });
+
+      // Create one text bookmark
+      await api.createBookmark({
+        text: "Test text content",
+        type: BookmarkTypes.TEXT,
+      });
+
+      // Third bookmark (any type) should fail
+      await expect(() =>
+        api.createBookmark({
+          text: "Another text bookmark",
+          type: BookmarkTypes.TEXT,
+        }),
+      ).rejects.toThrow(
+        /Bookmark quota exceeded. You can only have 2 bookmarks./,
+      );
+    });
+
+    test<CustomTestContext>("quota enforcement after deletion", async ({
+      apiCallers,
+      db,
+    }) => {
+      const user = await apiCallers[0].users.whoami();
+      const api = apiCallers[0].bookmarks;
+
+      // Set quota to 1 bookmark for this user
+      await db
+        .update(users)
+        .set({ bookmarkQuota: 1 })
+        .where(eq(users.id, user.id));
+
+      // Create first bookmark
+      const bookmark1 = await api.createBookmark({
+        url: "https://example1.com",
+        type: BookmarkTypes.LINK,
+      });
+
+      // Second bookmark should fail
+      await expect(() =>
+        api.createBookmark({
+          url: "https://example2.com",
+          type: BookmarkTypes.LINK,
+        }),
+      ).rejects.toThrow(
+        /Bookmark quota exceeded. You can only have 1 bookmarks./,
+      );
+
+      // Delete the first bookmark
+      await api.deleteBookmark({ bookmarkId: bookmark1.id });
+
+      // Now should be able to create a new bookmark
+      const bookmark2 = await api.createBookmark({
+        url: "https://example2.com",
+        type: BookmarkTypes.LINK,
+      });
+      expect(bookmark2.alreadyExists).toEqual(false);
+    });
+
+    test<CustomTestContext>("quota isolation between users", async ({
+      apiCallers,
+      db,
+    }) => {
+      const user1 = await apiCallers[0].users.whoami();
+
+      // Set quota to 1 for user1, unlimited for user2
+      await db
+        .update(users)
+        .set({ bookmarkQuota: 1 })
+        .where(eq(users.id, user1.id));
+
+      // User1 creates one bookmark (reaches quota)
+      await apiCallers[0].bookmarks.createBookmark({
+        url: "https://user1-example.com",
+        type: BookmarkTypes.LINK,
+      });
+
+      // User1 cannot create another bookmark
+      await expect(() =>
+        apiCallers[0].bookmarks.createBookmark({
+          url: "https://user1-example2.com",
+          type: BookmarkTypes.LINK,
+        }),
+      ).rejects.toThrow(
+        /Bookmark quota exceeded. You can only have 1 bookmarks./,
+      );
+
+      // User2 should be able to create multiple bookmarks (no quota)
+      await apiCallers[1].bookmarks.createBookmark({
+        url: "https://user2-example1.com",
+        type: BookmarkTypes.LINK,
+      });
+
+      await apiCallers[1].bookmarks.createBookmark({
+        url: "https://user2-example2.com",
+        type: BookmarkTypes.LINK,
+      });
+
+      await apiCallers[1].bookmarks.createBookmark({
+        text: "User2 text bookmark",
+        type: BookmarkTypes.TEXT,
+      });
+    });
+
+    test<CustomTestContext>("quota with zero limit", async ({
+      apiCallers,
+      db,
+    }) => {
+      const user = await apiCallers[0].users.whoami();
+      const api = apiCallers[0].bookmarks;
+
+      // Set quota to 0 bookmarks for this user
+      await db
+        .update(users)
+        .set({ bookmarkQuota: 0 })
+        .where(eq(users.id, user.id));
+
+      // Any bookmark creation should fail
+      await expect(() =>
+        api.createBookmark({
+          url: "https://example.com",
+          type: BookmarkTypes.LINK,
+        }),
+      ).rejects.toThrow(
+        /Bookmark quota exceeded. You can only have 0 bookmarks./,
+      );
+
+      await expect(() =>
+        api.createBookmark({
+          text: "Test text",
+          type: BookmarkTypes.TEXT,
+        }),
+      ).rejects.toThrow(
+        /Bookmark quota exceeded. You can only have 0 bookmarks./,
+      );
+    });
+
+    test<CustomTestContext>("quota does not affect duplicate link detection", async ({
+      apiCallers,
+      db,
+    }) => {
+      const user = await apiCallers[0].users.whoami();
+      const api = apiCallers[0].bookmarks;
+
+      // Set quota to 1 bookmark for this user
+      await db
+        .update(users)
+        .set({ bookmarkQuota: 1 })
+        .where(eq(users.id, user.id));
+
+      // Create first bookmark
+      const bookmark1 = await api.createBookmark({
+        url: "https://example.com",
+        type: BookmarkTypes.LINK,
+      });
+      expect(bookmark1.alreadyExists).toEqual(false);
+
+      // Try to create the same URL again - should return existing bookmark, not fail with quota
+      const bookmark2 = await api.createBookmark({
+        url: "https://example.com",
+        type: BookmarkTypes.LINK,
+      });
+      expect(bookmark2.alreadyExists).toEqual(true);
+      expect(bookmark2.id).toEqual(bookmark1.id);
+
+      // But creating a different URL should fail due to quota
+      await expect(() =>
+        api.createBookmark({
+          url: "https://different-example.com",
+          type: BookmarkTypes.LINK,
+        }),
+      ).rejects.toThrow(
+        /Bookmark quota exceeded. You can only have 1 bookmarks./,
+      );
+    });
+  });
 });
diff --git a/packages/trpc/routers/bookmarks.ts b/packages/trpc/routers/bookmarks.ts
index 2a02a0cd..f1fe10d7 100644
--- a/packages/trpc/routers/bookmarks.ts
+++ b/packages/trpc/routers/bookmarks.ts
@@ -1,5 +1,5 @@
 import { experimental_trpcMiddleware, TRPCError } from "@trpc/server";
-import { and, eq, gt, inArray, lt, or } from "drizzle-orm";
+import { and, count, eq, gt, inArray, lt, or } from "drizzle-orm";
 import invariant from "tiny-invariant";
 import { z } from "zod";
 
@@ -19,6 +19,7 @@ import {
   bookmarkTexts,
   customPrompts,
   tagsOnBookmarks,
+  users,
 } from "@karakeep/db/schema";
 import {
   deleteAsset,
@@ -267,6 +268,28 @@ export const bookmarksAppRouter = router({
           return { ...alreadyExists, alreadyExists: true };
         }
       }
+
+      // Check user quota
+      const user = await ctx.db.query.users.findFirst({
+        where: eq(users.id, ctx.user.id),
+        columns: {
+          bookmarkQuota: true,
+        },
+      });
+
+      if (user?.bookmarkQuota !== null && user?.bookmarkQuota !== undefined) {
+        const currentBookmarkCount = await ctx.db
+          .select({ count: count() })
+          .from(bookmarks)
+          .where(eq(bookmarks.userId, ctx.user.id));
+
+        if (currentBookmarkCount[0].count >= user.bookmarkQuota) {
+          throw new TRPCError({
+            code: "FORBIDDEN",
+            message: `Bookmark quota exceeded. You can only have ${user.bookmarkQuota} bookmarks.`,
+          });
+        }
+      }
       const bookmark = await ctx.db.transaction(async (tx) => {
         const bookmark = (
           await tx
diff --git a/packages/trpc/routers/users.ts b/packages/trpc/routers/users.ts
index bc1064e8..33aac2b7 100644
--- a/packages/trpc/routers/users.ts
+++ b/packages/trpc/routers/users.ts
@@ -144,6 +144,7 @@ export const usersAppRouter = router({
             email: z.string(),
             role: z.enum(["user", "admin"]).nullable(),
             localUser: z.boolean(),
+            bookmarkQuota: z.number().nullable(),
           }),
         ),
       }),
@@ -156,6 +157,7 @@ export const usersAppRouter = router({
           email: users.email,
           role: users.role,
           password: users.password,
+          bookmarkQuota: users.bookmarkQuota,
         })
         .from(users);