refactor: Move highlights object into models

2 files changed, 188 insertions, 131 deletions

diff --git a/packages/trpc/models/highlights.ts b/packages/trpc/models/highlights.ts
new file mode 100644
index 00000000..326f97f3
--- /dev/null
+++ b/packages/trpc/models/highlights.ts
@@ -0,0 +1,171 @@
+import { TRPCError } from "@trpc/server";
+import { and, desc, eq, lt, lte, or } from "drizzle-orm";
+import { z } from "zod";
+
+import { highlights } from "@karakeep/db/schema";
+import {
+  zHighlightSchema,
+  zNewHighlightSchema,
+  zUpdateHighlightSchema,
+} from "@karakeep/shared/types/highlights";
+import { zCursorV2 } from "@karakeep/shared/types/pagination";
+
+import { AuthedContext } from "..";
+import { PrivacyAware } from "./privacy";
+
+export class Highlight implements PrivacyAware {
+  constructor(
+    protected ctx: AuthedContext,
+    private highlight: typeof highlights.$inferSelect,
+  ) {}
+
+  static async fromId(ctx: AuthedContext, id: string): Promise<Highlight> {
+    const highlight = await ctx.db.query.highlights.findFirst({
+      where: eq(highlights.id, id),
+    });
+
+    if (!highlight) {
+      throw new TRPCError({
+        code: "NOT_FOUND",
+        message: "Highlight not found",
+      });
+    }
+
+    // If it exists but belongs to another user, throw forbidden error
+    if (highlight.userId !== ctx.user.id) {
+      throw new TRPCError({
+        code: "FORBIDDEN",
+        message: "User is not allowed to access resource",
+      });
+    }
+
+    return new Highlight(ctx, highlight);
+  }
+
+  static async create(
+    ctx: AuthedContext,
+    input: z.infer<typeof zNewHighlightSchema>,
+  ): Promise<Highlight> {
+    const [result] = await ctx.db
+      .insert(highlights)
+      .values({
+        bookmarkId: input.bookmarkId,
+        startOffset: input.startOffset,
+        endOffset: input.endOffset,
+        color: input.color,
+        text: input.text,
+        note: input.note,
+        userId: ctx.user.id,
+      })
+      .returning();
+
+    return new Highlight(ctx, result);
+  }
+
+  static async getForBookmark(
+    ctx: AuthedContext,
+    bookmarkId: string,
+  ): Promise<Highlight[]> {
+    const results = await ctx.db.query.highlights.findMany({
+      where: and(
+        eq(highlights.bookmarkId, bookmarkId),
+        eq(highlights.userId, ctx.user.id),
+      ),
+      orderBy: [desc(highlights.createdAt), desc(highlights.id)],
+    });
+
+    return results.map((h) => new Highlight(ctx, h));
+  }
+
+  static async getAll(
+    ctx: AuthedContext,
+    cursor?: z.infer<typeof zCursorV2> | null,
+    limit = 50,
+  ): Promise<{
+    highlights: Highlight[];
+    nextCursor: z.infer<typeof zCursorV2> | null;
+  }> {
+    const results = await ctx.db.query.highlights.findMany({
+      where: and(
+        eq(highlights.userId, ctx.user.id),
+        cursor
+          ? or(
+              lt(highlights.createdAt, cursor.createdAt),
+              and(
+                eq(highlights.createdAt, cursor.createdAt),
+                lte(highlights.id, cursor.id),
+              ),
+            )
+          : undefined,
+      ),
+      limit: limit + 1,
+      orderBy: [desc(highlights.createdAt), desc(highlights.id)],
+    });
+
+    let nextCursor: z.infer<typeof zCursorV2> | null = null;
+    if (results.length > limit) {
+      const nextItem = results.pop()!;
+      nextCursor = {
+        id: nextItem.id,
+        createdAt: nextItem.createdAt,
+      };
+    }
+
+    return {
+      highlights: results.map((h) => new Highlight(ctx, h)),
+      nextCursor,
+    };
+  }
+
+  ensureCanAccess(ctx: AuthedContext): void {
+    if (this.highlight.userId !== ctx.user.id) {
+      throw new TRPCError({
+        code: "FORBIDDEN",
+        message: "User is not allowed to access resource",
+      });
+    }
+  }
+
+  async delete(): Promise<z.infer<typeof zHighlightSchema>> {
+    const result = await this.ctx.db
+      .delete(highlights)
+      .where(
+        and(
+          eq(highlights.id, this.highlight.id),
+          eq(highlights.userId, this.ctx.user.id),
+        ),
+      )
+      .returning();
+
+    if (result.length === 0) {
+      throw new TRPCError({ code: "NOT_FOUND" });
+    }
+
+    return result[0];
+  }
+
+  async update(input: z.infer<typeof zUpdateHighlightSchema>): Promise<void> {
+    const result = await this.ctx.db
+      .update(highlights)
+      .set({
+        color: input.color,
+      })
+      .where(
+        and(
+          eq(highlights.id, this.highlight.id),
+          eq(highlights.userId, this.ctx.user.id),
+        ),
+      )
+      .returning();
+
+    if (result.length === 0) {
+      throw new TRPCError({ code: "NOT_FOUND" });
+    }
+
+    this.highlight = result[0];
+  }
+
+  asPublicHighlight(): z.infer<typeof zHighlightSchema> {
+    return this.highlight;
+  }
+}
diff --git a/packages/trpc/routers/highlights.ts b/packages/trpc/routers/highlights.ts
index 6124e35f..cb9b5e79 100644
--- a/packages/trpc/routers/highlights.ts
+++ b/packages/trpc/routers/highlights.ts
@@ -1,8 +1,5 @@
-import { experimental_trpcMiddleware, TRPCError } from "@trpc/server";
-import { and, desc, eq, lt, lte, or } from "drizzle-orm";
 import { z } from "zod";
 
-import { highlights } from "@karakeep/db/schema";
 import {
   DEFAULT_NUM_HIGHLIGHTS_PER_PAGE,
   zGetAllHighlightsResponseSchema,
@@ -10,174 +7,63 @@ import {
   zNewHighlightSchema,
   zUpdateHighlightSchema,
 } from "@karakeep/shared/types/highlights";
-import { zCursorV2 } from "@karakeep/shared/types/pagination";
 
-import { authedProcedure, Context, router } from "../index";
+import { authedProcedure, router } from "../index";
+import { Highlight } from "../models/highlights";
 import { ensureBookmarkOwnership } from "./bookmarks";
 
-const ensureHighlightOwnership = experimental_trpcMiddleware<{
-  ctx: Context;
-  input: { highlightId: string };
-}>().create(async (opts) => {
-  const highlight = await opts.ctx.db.query.highlights.findFirst({
-    where: eq(highlights.id, opts.input.highlightId),
-    columns: {
-      userId: true,
-    },
-  });
-  if (!opts.ctx.user) {
-    throw new TRPCError({
-      code: "UNAUTHORIZED",
-      message: "User is not authorized",
-    });
-  }
-  if (!highlight) {
-    throw new TRPCError({
-      code: "NOT_FOUND",
-      message: "Highlight not found",
-    });
-  }
-  if (highlight.userId != opts.ctx.user.id) {
-    throw new TRPCError({
-      code: "FORBIDDEN",
-      message: "User is not allowed to access resource",
-    });
-  }
-
-  return opts.next();
-});
-
 export const highlightsAppRouter = router({
   create: authedProcedure
     .input(zNewHighlightSchema)
     .output(zHighlightSchema)
     .use(ensureBookmarkOwnership)
     .mutation(async ({ input, ctx }) => {
-      const [result] = await ctx.db
-        .insert(highlights)
-        .values({
-          bookmarkId: input.bookmarkId,
-          startOffset: input.startOffset,
-          endOffset: input.endOffset,
-          color: input.color,
-          text: input.text,
-          note: input.note,
-          userId: ctx.user.id,
-        })
-        .returning();
-      return result;
+      const highlight = await Highlight.create(ctx, input);
+      return highlight.asPublicHighlight();
     }),
   getForBookmark: authedProcedure
     .input(z.object({ bookmarkId: z.string() }))
     .output(z.object({ highlights: z.array(zHighlightSchema) }))
     .use(ensureBookmarkOwnership)
     .query(async ({ input, ctx }) => {
-      const results = await ctx.db.query.highlights.findMany({
-        where: and(
-          eq(highlights.bookmarkId, input.bookmarkId),
-          eq(highlights.userId, ctx.user.id),
-        ),
-        orderBy: [desc(highlights.createdAt), desc(highlights.id)],
-      });
-      return { highlights: results };
+      const highlights = await Highlight.getForBookmark(ctx, input.bookmarkId);
+      return { highlights: highlights.map((h) => h.asPublicHighlight()) };
     }),
   get: authedProcedure
     .input(z.object({ highlightId: z.string() }))
     .output(zHighlightSchema)
-    .use(ensureHighlightOwnership)
     .query(async ({ input, ctx }) => {
-      const result = await ctx.db.query.highlights.findFirst({
-        where: and(
-          eq(highlights.id, input.highlightId),
-          eq(highlights.userId, ctx.user.id),
-        ),
-      });
-      if (!result) {
-        throw new TRPCError({ code: "NOT_FOUND" });
-      }
-      return result;
+      const highlight = await Highlight.fromId(ctx, input.highlightId);
+      return highlight.asPublicHighlight();
     }),
   getAll: authedProcedure
     .input(
       z.object({
-        cursor: zCursorV2.nullish(),
+        cursor: z.any().nullish(),
         limit: z.number().optional().default(DEFAULT_NUM_HIGHLIGHTS_PER_PAGE),
       }),
     )
     .output(zGetAllHighlightsResponseSchema)
     .query(async ({ input, ctx }) => {
-      const results = await ctx.db.query.highlights.findMany({
-        where: and(
-          eq(highlights.userId, ctx.user.id),
-          input.cursor
-            ? or(
-                lt(highlights.createdAt, input.cursor.createdAt),
-                and(
-                  eq(highlights.createdAt, input.cursor.createdAt),
-                  lte(highlights.id, input.cursor.id),
-                ),
-              )
-            : undefined,
-        ),
-        limit: input.limit + 1,
-        orderBy: [desc(highlights.createdAt), desc(highlights.id)],
-      });
-      let nextCursor: z.infer<typeof zCursorV2> | null = null;
-      if (results.length > input.limit) {
-        const nextItem = results.pop()!;
-        nextCursor = {
-          id: nextItem.id,
-          createdAt: nextItem.createdAt,
-        };
-      }
+      const result = await Highlight.getAll(ctx, input.cursor, input.limit);
       return {
-        highlights: results,
-        nextCursor,
+        highlights: result.highlights.map((h) => h.asPublicHighlight()),
+        nextCursor: result.nextCursor,
       };
     }),
   delete: authedProcedure
     .input(z.object({ highlightId: z.string() }))
     .output(zHighlightSchema)
-    .use(ensureHighlightOwnership)
     .mutation(async ({ input, ctx }) => {
-      const result = await ctx.db
-        .delete(highlights)
-        .where(
-          and(
-            eq(highlights.id, input.highlightId),
-            eq(highlights.userId, ctx.user.id),
-          ),
-        )
-        .returning();
-      if (result.length == 0) {
-        throw new TRPCError({
-          code: "NOT_FOUND",
-        });
-      }
-      return result[0];
+      const highlight = await Highlight.fromId(ctx, input.highlightId);
+      return await highlight.delete();
     }),
   update: authedProcedure
     .input(zUpdateHighlightSchema)
     .output(zHighlightSchema)
-    .use(ensureHighlightOwnership)
     .mutation(async ({ input, ctx }) => {
-      const result = await ctx.db
-        .update(highlights)
-        .set({
-          color: input.color,
-        })
-        .where(
-          and(
-            eq(highlights.id, input.highlightId),
-            eq(highlights.userId, ctx.user.id),
-          ),
-        )
-        .returning();
-      if (result.length == 0) {
-        throw new TRPCError({
-          code: "NOT_FOUND",
-        });
-      }
-      return result[0];
+      const highlight = await Highlight.fromId(ctx, input.highlightId);
+      await highlight.update(input);
+      return highlight.asPublicHighlight();
     }),
 });