feat: Add per user storage quota

3 files changed, 63 insertions, 0 deletions

diff --git a/packages/trpc/lib/storageQuota.ts b/packages/trpc/lib/storageQuota.ts
new file mode 100644
index 00000000..49b96af8
--- /dev/null
+++ b/packages/trpc/lib/storageQuota.ts
@@ -0,0 +1,57 @@
+import { eq, sum } from "drizzle-orm";
+
+import type { DB, KarakeepDBTransaction } from "@karakeep/db";
+import { assets, users } from "@karakeep/db/schema";
+import { QuotaApproved } from "@karakeep/shared/storageQuota";
+
+export class StorageQuotaError extends Error {
+  constructor(
+    public readonly currentUsage: number,
+    public readonly quota: number,
+    public readonly requestedSize: number,
+  ) {
+    super(
+      `Storage quota exceeded. Current usage: ${Math.round(currentUsage / 1024 / 1024)}MB, Quota: ${Math.round(quota / 1024 / 1024)}MB, Requested: ${Math.round(requestedSize / 1024 / 1024)}MB`,
+    );
+    this.name = "StorageQuotaError";
+  }
+}
+
+export async function checkStorageQuota(
+  db: DB | KarakeepDBTransaction,
+  userId: string,
+  requestedSize: number,
+): Promise<QuotaApproved> {
+  const user = await db.query.users.findFirst({
+    where: eq(users.id, userId),
+    columns: {
+      storageQuota: true,
+    },
+  });
+
+  if (user?.storageQuota === null || user?.storageQuota === undefined) {
+    // No quota limit - approve the request
+    return QuotaApproved._create(userId, requestedSize);
+  }
+
+  const currentUsage = await getCurrentStorageUsage(db, userId);
+
+  if (currentUsage + requestedSize > user.storageQuota) {
+    throw new StorageQuotaError(currentUsage, user.storageQuota, requestedSize);
+  }
+
+  // Quota check passed - return approval token
+  return QuotaApproved._create(userId, requestedSize);
+}
+
+export async function getCurrentStorageUsage(
+  db: DB | KarakeepDBTransaction,
+  userId: string,
+): Promise<number> {
+  const currentUsageResult = await db
+    .select({ totalSize: sum(assets.size) })
+    .from(assets)
+    .where(eq(assets.userId, userId));
+
+  return Number(currentUsageResult[0]?.totalSize ?? 0);
+}
diff --git a/packages/trpc/routers/admin.ts b/packages/trpc/routers/admin.ts
index f46e17f2..5e169857 100644
--- a/packages/trpc/routers/admin.ts
+++ b/packages/trpc/routers/admin.ts
@@ -351,6 +351,10 @@ export const adminAppRouter = router({
         updateData.bookmarkQuota = input.bookmarkQuota;
       }
 
+      if (input.storageQuota !== undefined) {
+        updateData.storageQuota = input.storageQuota;
+      }
+
       if (Object.keys(updateData).length === 0) {
         throw new TRPCError({
           code: "BAD_REQUEST",
diff --git a/packages/trpc/routers/users.ts b/packages/trpc/routers/users.ts
index ea5e6944..3c52b321 100644
--- a/packages/trpc/routers/users.ts
+++ b/packages/trpc/routers/users.ts
@@ -148,6 +148,7 @@ export const usersAppRouter = router({
             role: z.enum(["user", "admin"]).nullable(),
             localUser: z.boolean(),
             bookmarkQuota: z.number().nullable(),
+            storageQuota: z.number().nullable(),
           }),
         ),
       }),
@@ -161,6 +162,7 @@ export const usersAppRouter = router({
           role: users.role,
           password: users.password,
           bookmarkQuota: users.bookmarkQuota,
+          storageQuota: users.storageQuota,
         })
         .from(users);