Merge pull request dubinc#2292 from dubinc/improve-domain-patch

Improve domain update

Author: steven-tey

apps/web/app/api/cron/domains/update/route.ts

@@ -4,28 +4,24 @@ import { linkCache } from "@/lib/api/links/cache";
 import { verifyQstashSignature } from "@/lib/cron/verify-qstash";
 import { recordLink } from "@/lib/tinybird";
 import { prisma } from "@dub/prisma";
+import { Link } from "@dub/prisma/client";
+import { linkConstructorSimple } from "@dub/utils";
 import { z } from "zod";
 
 export const dynamic = "force-dynamic";
 
 const schema = z.object({
   newDomain: z.string(),
   oldDomain: z.string(),
-  workspaceId: z.string(),
-  page: z.number(),
 });
 
-const pageSize = 100;
-
 // POST /api/cron/domains/update
 export async function POST(req: Request) {
   try {
     const rawBody = await req.text();
     await verifyQstashSignature({ req, rawBody });
 
-    const { newDomain, oldDomain, workspaceId, page } = schema.parse(
-      JSON.parse(rawBody),
-    );
+    const { newDomain, oldDomain } = schema.parse(JSON.parse(rawBody));
 
     const newDomainRecord = await prisma.domain.findUnique({
       where: {
@@ -37,46 +33,85 @@ export async function POST(req: Request) {
       return new Response(`Domain ${newDomain} not found. Skipping update...`);
     }
 
-    const links = await prisma.link.findMany({
+    const linksToUpdate = await prisma.link.findMany({
       where: {
+        domain: oldDomain,
+      },
+      take: 100,
+    });
+
+    if (linksToUpdate.length === 0) {
+      return new Response("No more links to update. Exiting...");
+    }
+
+    const linkIdsToUpdate = linksToUpdate.map((link) => link.id);
+
+    await prisma.link.updateMany({
+      where: {
+        id: {
+          in: linkIdsToUpdate,
+        },
+      },
+      data: {
         domain: newDomain,
       },
+    });
+
+    const updatedLinks = await prisma.link.findMany({
+      where: {
+        id: {
+          in: linkIdsToUpdate,
+        },
+      },
       include: {
         tags: {
           select: {
             tag: true,
           },
         },
       },
-      skip: (page - 1) * pageSize,
-      take: pageSize,
     });
 
-    if (links.length === 0) {
-      return new Response("No more links to update. Exiting...");
-    }
-
-    await Promise.all([
-      // rename redis keys
-      linkCache.rename({
-        links,
-        oldDomain,
-      }),
-
-      // update links in Tinybird
-      recordLink(links),
+    await Promise.allSettled([
+      // update the `shortLink` field for each of the short links
+      updateShortLinks(updatedLinks),
+      // record new link values in Tinybird (dub_links_metadata)
+      recordLink(updatedLinks),
+      // expire the redis cache for the old links
+      linkCache.expireMany(linksToUpdate),
     ]);
 
     await queueDomainUpdate({
-      workspaceId,
-      oldDomain,
       newDomain,
-      page: page + 1,
-      delay: 2,
+      oldDomain,
+      delay: 1,
     });
 
     return new Response("Domain's links updated.");
   } catch (error) {
     return handleAndReturnErrorResponse(error);
   }
 }
+
+// Update the shortLink column for a list of links
+const updateShortLinks = async (
+  links: Pick<Link, "id" | "domain" | "key">[],
+) => {
+  if (!links || links.length === 0) {
+    return new Response("No links found.");
+  }
+
+  for (const link of links) {
+    await prisma.link.update({
+      where: {
+        id: link.id,
+      },
+      data: {
+        shortLink: linkConstructorSimple({
+          domain: link.domain,
+          key: link.key,
+        }),
+      },
+    });
+  }
+};

apps/web/app/api/domains/[domain]/route.ts

@@ -151,10 +151,8 @@ export const PATCH = withWorkspace(
 
             // trigger the queue to rename the redis keys and update the links in Tinybird
             queueDomainUpdate({
-              workspaceId: workspace.id,
               oldDomain: domain,
               newDomain: newDomain,
-              page: 1,
             }),
           ]);
         }

apps/web/app/api/domains/route.ts

@@ -12,7 +12,7 @@ import {
 } from "@/lib/zod/schemas/domains";
 import { prisma } from "@dub/prisma";
 import { combineWords, DEFAULT_LINK_PROPS, nanoid } from "@dub/utils";
-import { Prisma } from "@prisma/client";
+import { Link, Prisma } from "@prisma/client";
 import { NextResponse } from "next/server";
 
 // GET /api/domains – get all domains for a workspace
@@ -33,40 +33,46 @@ export const GET = withWorkspace(
       },
       include: {
         registeredDomain: true,
-        ...(includeLink && {
-          links: {
-            where: {
-              key: {
-                in: ["_root", "akoJCU0="],
-              },
+      },
+      take: pageSize,
+      skip: (page - 1) * pageSize,
+    });
+
+    const links = includeLink
+      ? await prisma.link.findMany({
+          where: {
+            key: {
+              in: ["_root", "akoJCU0="],
+            },
+            domain: {
+              in: domains.map((domain) => domain.slug),
             },
-            include: {
-              tags: {
-                select: {
-                  tag: {
-                    select: {
-                      id: true,
-                      name: true,
-                      color: true,
-                    },
-                  },
-                },
+          },
+          include: {
+            tags: {
+              select: {
+                tag: true,
               },
             },
           },
-        }),
+        })
+      : [];
+
+    const linkMap = links.reduce(
+      (acc, link) => {
+        acc[link.domain] = link;
+        return acc;
       },
-      take: pageSize,
-      skip: (page - 1) * pageSize,
-    });
+      {} as Record<string, Link>,
+    );
 
     const response = domains.map((domain) => ({
       ...transformDomain(domain),
       ...(includeLink &&
-        domain.links.length > 0 && {
+        linkMap[domain.slug] && {
           link: transformLink({
-            ...domain.links[0],
-            tags: domain.links[0]["tags"].map((tag) => tag),
+            ...linkMap[domain.slug],
+            tags: linkMap[domain.slug]["tags"].map((tag) => tag),
           }),
         }),
     }));

apps/web/lib/api/domains/queue.ts

@@ -2,26 +2,20 @@ import { qstash } from "@/lib/cron";
 import { APP_DOMAIN_WITH_NGROK } from "@dub/utils";
 
 export async function queueDomainUpdate({
-  workspaceId,
   oldDomain,
   newDomain,
-  page,
   delay,
 }: {
-  workspaceId: string;
   oldDomain: string;
   newDomain: string;
-  page: number;
   delay?: number;
 }) {
   return await qstash.publishJSON({
     url: `${APP_DOMAIN_WITH_NGROK}/api/cron/domains/update`,
     ...(delay && { delay }),
     body: {
-      workspaceId,
       oldDomain,
       newDomain,
-      page,
     },
   });
 }

apps/web/lib/api/links/cache.ts

@@ -75,29 +75,6 @@ class LinkCache {
     return await pipeline.exec();
   }
 
-  async rename({
-    links,
-    oldDomain,
-  }: {
-    links: Pick<LinkProps, "domain" | "key">[];
-    oldDomain: string;
-  }) {
-    if (links.length === 0) {
-      return;
-    }
-
-    const pipeline = redis.pipeline();
-
-    links.forEach(({ domain, key }) => {
-      const oldCacheKey = this._createKey({ domain: oldDomain, key });
-      const newCacheKey = this._createKey({ domain, key });
-
-      pipeline.renamenx(oldCacheKey, newCacheKey);
-    });
-
-    return await pipeline.exec();
-  }
-
   async expireMany(links: Pick<LinkProps, "domain" | "key">[]) {
     if (links.length === 0) {
       return;

apps/web/scripts/fix-broken-root-domains.ts

@@ -1,3 +1,5 @@
+// @ts-nocheck
+
 import { createLink } from "@/lib/api/links";
 import { prisma } from "@dub/prisma";
 import { DEFAULT_LINK_PROPS } from "@dub/utils";

packages/prisma/schema/domain.prisma

@@ -11,7 +11,6 @@ model Domain {
   logo                    String?
   appleAppSiteAssociation Json?
   assetLinks              Json?
-  links                   Link[]
   project                 Project?          @relation(fields: [projectId], references: [id], onDelete: Cascade)
   projectId               String?
   registeredDomain        RegisteredDomain?

packages/prisma/schema/link.prisma

@@ -43,12 +43,11 @@ model Link {
   userId String?
 
   // Project that the link belongs to
-  project       Project? @relation(fields: [projectId], references: [id], onUpdate: Cascade, onDelete: Cascade)
-  projectId     String?
-  // Relational reference to the project domain
-  projectDomain Domain?  @relation(fields: [domain], references: [slug], onUpdate: Cascade, onDelete: Cascade)
+  project   Project? @relation(fields: [projectId], references: [id], onUpdate: Cascade, onDelete: Cascade)
+  projectId String?
+
   // Program that the link belongs to
-  programId     String?
+  programId String?
 
   folderId String?
   folder   Folder? @relation(fields: [folderId], references: [id], onUpdate: Cascade, onDelete: SetNull)
@@ -84,9 +83,9 @@ model Link {
   commissions Commission[]
 
   @@unique([domain, key])
+  @@unique([projectId, externalId])
   @@index([projectId, url(length: 500)])
   @@index([projectId, folderId, archived, createdAt(sort: Desc)])
-  @@unique([projectId, externalId])
   @@index([programId, partnerId])
   @@index(folderId) // used in /api/folders
   @@index(userId) // for relation to User table