refactor: Extract quota logic into its own class

11 files changed, 133 insertions, 102 deletions

diff --git a/apps/workers/workers/assetPreprocessingWorker.ts b/apps/workers/workers/assetPreprocessingWorker.ts
index 73cf8bb5..d059e21c 100644
--- a/apps/workers/workers/assetPreprocessingWorker.ts
+++ b/apps/workers/workers/assetPreprocessingWorker.ts
@@ -14,6 +14,7 @@ import {
   bookmarkAssets,
   bookmarks,
 } from "@karakeep/db/schema";
+import { QuotaService, StorageQuotaError } from "@karakeep/shared-server";
 import { newAssetId, readAsset, saveAsset } from "@karakeep/shared/assetdb";
 import serverConfig from "@karakeep/shared/config";
 import logger from "@karakeep/shared/logger";
@@ -22,10 +23,6 @@ import {
   OpenAIQueue,
   triggerSearchReindex,
 } from "@karakeep/shared/queues";
-import {
-  checkStorageQuota,
-  StorageQuotaError,
-} from "@karakeep/trpc/lib/storageQuota";
 
 export class AssetPreprocessingWorker {
   static build() {
@@ -136,7 +133,7 @@ export async function extractAndSavePDFScreenshot(
     }
 
     // Check storage quota before inserting
-    const quotaApproved = await checkStorageQuota(
+    const quotaApproved = await QuotaService.checkStorageQuota(
       db,
       bookmark.userId,
       screenshot.buffer.byteLength,
diff --git a/apps/workers/workers/crawlerWorker.ts b/apps/workers/workers/crawlerWorker.ts
index 36068f14..2aaab776 100644
--- a/apps/workers/workers/crawlerWorker.ts
+++ b/apps/workers/workers/crawlerWorker.ts
@@ -41,6 +41,7 @@ import {
   bookmarks,
   users,
 } from "@karakeep/db/schema";
+import { QuotaService } from "@karakeep/shared-server";
 import {
   ASSET_TYPES,
   getAssetSize,
@@ -65,7 +66,6 @@ import {
 } from "@karakeep/shared/queues";
 import { tryCatch } from "@karakeep/shared/tryCatch";
 import { BookmarkTypes } from "@karakeep/shared/types/bookmarks";
-import { checkStorageQuota } from "@karakeep/trpc/lib/storageQuota";
 
 import metascraperReddit from "../metascraper-plugins/metascraper-reddit";
 
@@ -536,7 +536,7 @@ async function storeScreenshot(
 
   // Check storage quota before saving the screenshot
   const { data: quotaApproved, error: quotaError } = await tryCatch(
-    checkStorageQuota(db, userId, screenshot.byteLength),
+    QuotaService.checkStorageQuota(db, userId, screenshot.byteLength),
   );
 
   if (quotaError) {
@@ -586,7 +586,7 @@ async function downloadAndStoreFile(
 
     // Check storage quota before saving the asset
     const { data: quotaApproved, error: quotaError } = await tryCatch(
-      checkStorageQuota(db, userId, buffer.byteLength),
+      QuotaService.checkStorageQuota(db, userId, buffer.byteLength),
     );
 
     if (quotaError) {
@@ -655,7 +655,7 @@ async function archiveWebpage(
   const fileSize = stats.size;
 
   const { data: quotaApproved, error: quotaError } = await tryCatch(
-    checkStorageQuota(db, userId, fileSize),
+    QuotaService.checkStorageQuota(db, userId, fileSize),
   );
 
   if (quotaError) {
@@ -813,7 +813,7 @@ async function storeHtmlContent(
   }
 
   const { data: quotaApproved, error: quotaError } = await tryCatch(
-    checkStorageQuota(db, userId, contentBuffer.byteLength),
+    QuotaService.checkStorageQuota(db, userId, contentBuffer.byteLength),
   );
   if (quotaError) {
     logger.warn(
diff --git a/apps/workers/workers/videoWorker.ts b/apps/workers/workers/videoWorker.ts
index 69f88a21..68be0126 100644
--- a/apps/workers/workers/videoWorker.ts
+++ b/apps/workers/workers/videoWorker.ts
@@ -7,6 +7,7 @@ import { workerStatsCounter } from "metrics";
 

 import { db } from "@karakeep/db";

 import { AssetTypes } from "@karakeep/db/schema";

+import { QuotaService, StorageQuotaError } from "@karakeep/shared-server";

 import {

   ASSET_TYPES,

   newAssetId,

@@ -20,10 +21,6 @@ import {
   ZVideoRequest,

   zvideoRequestSchema,

 } from "@karakeep/shared/queues";

-import {

-  checkStorageQuota,

-  StorageQuotaError,

-} from "@karakeep/trpc/lib/storageQuota";

 

 import { getBookmarkDetails, updateAsset } from "../workerUtils";

 

@@ -148,7 +145,11 @@ async function runWorker(job: DequeuedJob<ZVideoRequest>) {
   const fileSize = stats.size;

 

   try {

-    const quotaApproved = await checkStorageQuota(db, userId, fileSize);

+    const quotaApproved = await QuotaService.checkStorageQuota(

+      db,

+      userId,

+      fileSize,

+    );

 

     await saveAssetFromFile({

       userId,

diff --git a/packages/api/utils/upload.ts b/packages/api/utils/upload.ts
index 6f18790b..a843e29c 100644
--- a/packages/api/utils/upload.ts
+++ b/packages/api/utils/upload.ts
@@ -5,6 +5,7 @@ import { Readable } from "stream";
 import { pipeline } from "stream/promises";
 
 import { assets, AssetTypes } from "@karakeep/db/schema";
+import { QuotaService, StorageQuotaError } from "@karakeep/shared-server";
 import {
   newAssetId,
   saveAssetFromFile,
@@ -12,10 +13,6 @@ import {
 } from "@karakeep/shared/assetdb";
 import serverConfig from "@karakeep/shared/config";
 import { AuthedContext } from "@karakeep/trpc";
-import {
-  checkStorageQuota,
-  StorageQuotaError,
-} from "@karakeep/trpc/lib/storageQuota";
 
 const MAX_UPLOAD_SIZE_BYTES = serverConfig.maxAssetSizeMb * 1024 * 1024;
 
@@ -73,7 +70,11 @@ export async function uploadAsset(
 
   let quotaApproved;
   try {
-    quotaApproved = await checkStorageQuota(db, user.id, data.size);
+    quotaApproved = await QuotaService.checkStorageQuota(
+      db,
+      user.id,
+      data.size,
+    );
   } catch (error) {
     if (error instanceof StorageQuotaError) {
       return { error: error.message, status: 403 };
diff --git a/packages/shared-server/package.json b/packages/shared-server/package.json
index 8ac98e21..6ba6b6d9 100644
--- a/packages/shared-server/package.json
+++ b/packages/shared-server/package.json
@@ -5,6 +5,7 @@
   "private": true,
   "type": "module",
   "dependencies": {
+    "@karakeep/db": "workspace:^0.1.0",
     "@karakeep/plugins-search-meilisearch": "workspace:^0.1.0",
     "@karakeep/shared": "workspace:^0.1.0"
   },
diff --git a/packages/shared-server/src/index.ts b/packages/shared-server/src/index.ts
index a17576ad..ff3c6abc 100644
--- a/packages/shared-server/src/index.ts
+++ b/packages/shared-server/src/index.ts
@@ -1 +1,2 @@
 export { loadAllPlugins } from "./plugins";
+export { QuotaService, StorageQuotaError } from "./services/quotaService";
diff --git a/packages/shared-server/src/services/quotaService.ts b/packages/shared-server/src/services/quotaService.ts
new file mode 100644
index 00000000..a09b76bf
--- /dev/null
+++ b/packages/shared-server/src/services/quotaService.ts
@@ -0,0 +1,93 @@
+import { count, eq, sum } from "drizzle-orm";
+
+import type { DB, KarakeepDBTransaction } from "@karakeep/db";
+import { assets, bookmarks, 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";
+  }
+}
+
+// TODO: Change the API of this class to either return a boolean
+// or throw an exception on lack of quota because now, it's inconsistent.
+export class QuotaService {
+  // TODO: Use quota approval tokens for bookmark creation when
+  // bookmark creation logic is in the model.
+  static async canCreateBookmark(db: DB, userId: string) {
+    const user = await db.query.users.findFirst({
+      where: eq(users.id, userId),
+      columns: {
+        bookmarkQuota: true,
+      },
+    });
+
+    if (user?.bookmarkQuota !== null && user?.bookmarkQuota !== undefined) {
+      const currentBookmarkCount = await db
+        .select({ count: count() })
+        .from(bookmarks)
+        .where(eq(bookmarks.userId, userId));
+
+      if (currentBookmarkCount[0].count >= user.bookmarkQuota) {
+        return {
+          result: false,
+          error: `Bookmark quota exceeded. You can only have ${user.bookmarkQuota} bookmarks.`,
+        } as const;
+      }
+    }
+    return {
+      result: true,
+    } as const;
+  }
+
+  static async 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 this.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);
+  }
+
+  static async 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/lib/storageQuota.ts b/packages/trpc/lib/storageQuota.ts
deleted file mode 100644
index 49b96af8..00000000
--- a/packages/trpc/lib/storageQuota.ts
+++ /dev/null
@@ -1,57 +0,0 @@
-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/package.json b/packages/trpc/package.json
index 31cb3d9a..c4e16675 100644
--- a/packages/trpc/package.json
+++ b/packages/trpc/package.json
@@ -14,8 +14,8 @@
   },
   "dependencies": {
     "@karakeep/db": "workspace:*",
-    "@karakeep/plugins-search-meilisearch": "workspace:*",
     "@karakeep/shared": "workspace:*",
+    "@karakeep/shared-server": "workspace:*",
     "@trpc/server": "^11.4.3",
     "bcryptjs": "^2.4.3",
     "deep-equal": "^2.2.3",
diff --git a/packages/trpc/routers/bookmarks.ts b/packages/trpc/routers/bookmarks.ts
index 298f0961..db9d33fc 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, count, eq, gt, inArray, lt, or } from "drizzle-orm";
+import { and, eq, gt, inArray, lt, or } from "drizzle-orm";
 import { EnqueueOptions } from "liteque";
 import invariant from "tiny-invariant";
 import { z } from "zod";
@@ -20,8 +20,8 @@ import {
   bookmarkTexts,
   customPrompts,
   tagsOnBookmarks,
-  users,
 } from "@karakeep/db/schema";
+import { QuotaService } from "@karakeep/shared-server";
 import {
   deleteAsset,
   SUPPORTED_BOOKMARK_ASSET_TYPES,
@@ -273,26 +273,17 @@ export const bookmarksAppRouter = router({
       }
 
       // 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 quotaResult = await QuotaService.canCreateBookmark(
+        ctx.db,
+        ctx.user.id,
+      );
+      if (!quotaResult.result) {
+        throw new TRPCError({
+          code: "FORBIDDEN",
+          message: quotaResult.error,
+        });
       }
+
       const bookmark = await ctx.db.transaction(async (tx) => {
         const bookmark = (
           await tx
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index e81712c6..50eb8062 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -1280,6 +1280,9 @@ importers:
 
   packages/shared-server:
     dependencies:
+      '@karakeep/db':
+        specifier: workspace:^0.1.0
+        version: link:../db
       '@karakeep/plugins-search-meilisearch':
         specifier: workspace:^0.1.0
         version: link:../plugins-search-meilisearch
@@ -1299,12 +1302,12 @@ importers:
       '@karakeep/db':
         specifier: workspace:*
         version: link:../db
-      '@karakeep/plugins-search-meilisearch':
-        specifier: workspace:*
-        version: link:../plugins-search-meilisearch
       '@karakeep/shared':
         specifier: workspace:*
         version: link:../shared
+      '@karakeep/shared-server':
+        specifier: workspace:*
+        version: link:../shared-server
       '@trpc/server':
         specifier: ^11.4.3
         version: 11.4.3(typescript@5.8.3)