1 files changed, 57 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);
+}