From c16173ea0fdbf6cc47b13756c0a77e8399669055 Mon Sep 17 00:00:00 2001
From: MohamedBassem <me@mbassem.com>
Date: Sat, 12 Oct 2024 16:47:22 +0000
Subject: feature: Introduce a mechanism to cleanup dangling assets

---
 .../components/dashboard/admin/AdminActions.tsx    |  22 +++++
 .../web/components/dashboard/admin/ServerStats.tsx |   6 ++
 apps/web/components/dashboard/admin/UserList.tsx   |  18 +++-
 apps/workers/index.ts                              |  12 ++-
 apps/workers/tidyAssetsWorker.ts                   | 107 +++++++++++++++++++++
 5 files changed, 161 insertions(+), 4 deletions(-)
 create mode 100644 apps/workers/tidyAssetsWorker.ts

(limited to 'apps')

diff --git a/apps/web/components/dashboard/admin/AdminActions.tsx b/apps/web/components/dashboard/admin/AdminActions.tsx
index dfdf65eb..e5d6ed69 100644
--- a/apps/web/components/dashboard/admin/AdminActions.tsx
+++ b/apps/web/components/dashboard/admin/AdminActions.tsx
@@ -52,6 +52,21 @@ export default function AdminActions() {
     },
   });
 
+  const { mutateAsync: tidyAssets, isPending: isTidyAssetsPending } =
+    api.admin.tidyAssets.useMutation({
+      onSuccess: () => {
+        toast({
+          description: "Tidy assets request has been enqueued!",
+        });
+      },
+      onError: (e) => {
+        toast({
+          variant: "destructive",
+          description: e.message,
+        });
+      },
+    });
+
   return (
     <div>
       <div className="mb-2 mt-8 text-xl font-medium">Actions</div>
@@ -97,6 +112,13 @@ export default function AdminActions() {
         >
           Reindex All Bookmarks
         </ActionButton>
+        <ActionButton
+          variant="destructive"
+          loading={isTidyAssetsPending}
+          onClick={() => tidyAssets()}
+        >
+          Compact Assets
+        </ActionButton>
       </div>
     </div>
   );
diff --git a/apps/web/components/dashboard/admin/ServerStats.tsx b/apps/web/components/dashboard/admin/ServerStats.tsx
index e95dc437..f45d86c5 100644
--- a/apps/web/components/dashboard/admin/ServerStats.tsx
+++ b/apps/web/components/dashboard/admin/ServerStats.tsx
@@ -122,6 +122,12 @@ export default function ServerStats() {
               <TableCell>{serverStats.inferenceStats.pending}</TableCell>
               <TableCell>{serverStats.inferenceStats.failed}</TableCell>
             </TableRow>
+            <TableRow>
+              <TableCell>Tidy Assets Jobs</TableCell>
+              <TableCell>{serverStats.tidyAssetsStats.queued}</TableCell>
+              <TableCell>-</TableCell>
+              <TableCell>-</TableCell>
+            </TableRow>
           </TableBody>
         </Table>
       </div>
diff --git a/apps/web/components/dashboard/admin/UserList.tsx b/apps/web/components/dashboard/admin/UserList.tsx
index 024325a3..65bf4068 100644
--- a/apps/web/components/dashboard/admin/UserList.tsx
+++ b/apps/web/components/dashboard/admin/UserList.tsx
@@ -15,10 +15,18 @@ import { api } from "@/lib/trpc";
 import { Trash } from "lucide-react";
 import { useSession } from "next-auth/react";
 
+function toHumanReadableSize(size: number) {
+  const sizes = ["Bytes", "KB", "MB", "GB", "TB"];
+  if (size === 0) return "0 Bytes";
+  const i = Math.floor(Math.log(size) / Math.log(1024));
+  return (size / Math.pow(1024, i)).toFixed(2) + " " + sizes[i];
+}
+
 export default function UsersSection() {
   const { data: session } = useSession();
   const invalidateUserList = api.useUtils().users.list.invalidate;
   const { data: users } = api.users.list.useQuery();
+  const { data: userStats } = api.admin.userStats.useQuery();
   const { mutate: deleteUser, isPending: isDeletionPending } =
     api.users.delete.useMutation({
       onSuccess: () => {
@@ -35,7 +43,7 @@ export default function UsersSection() {
       },
     });
 
-  if (!users) {
+  if (!users || !userStats) {
     return <LoadingSpinner />;
   }
 
@@ -47,6 +55,8 @@ export default function UsersSection() {
         <TableHeader className="bg-gray-200">
           <TableHead>Name</TableHead>
           <TableHead>Email</TableHead>
+          <TableHead>Num Bookmarks</TableHead>
+          <TableHead>Asset Sizes</TableHead>
           <TableHead>Role</TableHead>
           <TableHead>Action</TableHead>
         </TableHeader>
@@ -55,6 +65,12 @@ export default function UsersSection() {
             <TableRow key={u.id}>
               <TableCell className="py-1">{u.name}</TableCell>
               <TableCell className="py-1">{u.email}</TableCell>
+              <TableCell className="py-1">
+                {userStats[u.id].numBookmarks}
+              </TableCell>
+              <TableCell className="py-1">
+                {toHumanReadableSize(userStats[u.id].assetSizes)}
+              </TableCell>
               <TableCell className="py-1 capitalize">{u.role}</TableCell>
               <TableCell className="py-1">
                 <ActionButton
diff --git a/apps/workers/index.ts b/apps/workers/index.ts
index e576776a..f9a05e59 100644
--- a/apps/workers/index.ts
+++ b/apps/workers/index.ts
@@ -1,5 +1,7 @@
 import "dotenv/config";
 
+import { TidyAssetsWorker } from "tidyAssetsWorker";
+
 import serverConfig from "@hoarder/shared/config";
 import logger from "@hoarder/shared/logger";
 import { runQueueDBMigrations } from "@hoarder/shared/queues";
@@ -13,21 +15,25 @@ async function main() {
   logger.info(`Workers version: ${serverConfig.serverVersion ?? "not set"}`);
   runQueueDBMigrations();
 
-  const [crawler, openai, search] = [
+  const [crawler, openai, search, tidyAssets] = [
     await CrawlerWorker.build(),
     OpenAiWorker.build(),
     SearchIndexingWorker.build(),
+    TidyAssetsWorker.build(),
   ];
 
   await Promise.any([
-    Promise.all([crawler.run(), openai.run(), search.run()]),
+    Promise.all([crawler.run(), openai.run(), search.run(), tidyAssets.run()]),
     shutdownPromise,
   ]);
-  logger.info("Shutting down crawler, openai and search workers ...");
+  logger.info(
+    "Shutting down crawler, openai, tidyAssets and search workers ...",
+  );
 
   crawler.stop();
   openai.stop();
   search.stop();
+  tidyAssets.stop();
 }
 
 main();
diff --git a/apps/workers/tidyAssetsWorker.ts b/apps/workers/tidyAssetsWorker.ts
new file mode 100644
index 00000000..bc14aab9
--- /dev/null
+++ b/apps/workers/tidyAssetsWorker.ts
@@ -0,0 +1,107 @@
+import { eq } from "drizzle-orm";
+
+import { db } from "@hoarder/db";
+import { assets } from "@hoarder/db/schema";
+import { DequeuedJob, Runner } from "@hoarder/queue";
+import { deleteAsset, getAllAssets } from "@hoarder/shared/assetdb";
+import logger from "@hoarder/shared/logger";
+import {
+  TidyAssetsQueue,
+  ZTidyAssetsRequest,
+  zTidyAssetsRequestSchema,
+} from "@hoarder/shared/queues";
+
+export class TidyAssetsWorker {
+  static build() {
+    logger.info("Starting tidy assets worker ...");
+    const worker = new Runner<ZTidyAssetsRequest>(
+      TidyAssetsQueue,
+      {
+        run: runTidyAssets,
+        onComplete: (job) => {
+          const jobId = job?.id ?? "unknown";
+          logger.info(`[tidyAssets][${jobId}] Completed successfully`);
+          return Promise.resolve();
+        },
+        onError: (job) => {
+          const jobId = job?.id ?? "unknown";
+          logger.error(
+            `[tidyAssets][${jobId}] tidy assets job failed: ${job.error}\n${job.error.stack}`,
+          );
+          return Promise.resolve();
+        },
+      },
+      {
+        concurrency: 1,
+        pollIntervalMs: 1000,
+        timeoutSecs: 30,
+      },
+    );
+
+    return worker;
+  }
+}
+
+async function handleAsset(
+  asset: {
+    assetId: string;
+    userId: string;
+    size: number;
+    contentType: string;
+    fileName?: string | null;
+  },
+  request: ZTidyAssetsRequest,
+  jobId: string,
+) {
+  const dbRow = await db.query.assets.findFirst({
+    where: eq(assets.id, asset.assetId),
+  });
+  if (!dbRow) {
+    if (request.cleanDanglingAssets) {
+      await deleteAsset({ userId: asset.userId, assetId: asset.assetId });
+      logger.info(
+        `[tidyAssets][${jobId}] Asset ${asset.assetId} not found in the database. Deleting it.`,
+      );
+    } else {
+      logger.warn(
+        `[tidyAssets][${jobId}] Asset ${asset.assetId} not found in the database. Not deleting it because cleanDanglingAssets is false.`,
+      );
+    }
+    return;
+  }
+
+  if (request.syncAssetMetadata) {
+    await db
+      .update(assets)
+      .set({
+        contentType: asset.contentType,
+        fileName: asset.fileName,
+        size: asset.size,
+      })
+      .where(eq(assets.id, asset.assetId));
+    logger.info(
+      `[tidyAssets][${jobId}] Updated metadata for asset ${asset.assetId}`,
+    );
+  }
+}
+
+async function runTidyAssets(job: DequeuedJob<ZTidyAssetsRequest>) {
+  const jobId = job.id ?? "unknown";
+
+  const request = zTidyAssetsRequestSchema.safeParse(job.data);
+  if (!request.success) {
+    throw new Error(
+      `[tidyAssets][${jobId}] Got malformed job request: ${request.error.toString()}`,
+    );
+  }
+
+  for await (const asset of getAllAssets()) {
+    try {
+      handleAsset(asset, request.data, jobId);
+    } catch (e) {
+      logger.error(
+        `[tidyAssets][${jobId}] Failed to tidy asset ${asset.assetId}: ${e}`,
+      );
+    }
+  }
+}
-- 
cgit v1.2.3-70-g09d2