Support multiple QR code logos for each domain

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

@@ -9,12 +9,14 @@ import { DubApiError } from "@/lib/api/errors";
 import { parseRequestBody } from "@/lib/api/utils";
 import { withWorkspace } from "@/lib/auth";
 import { prisma } from "@/lib/prisma";
+import { storage } from "@/lib/storage";
 import { recordLink } from "@/lib/tinybird";
 import { redis } from "@/lib/upstash";
 import {
   DomainSchema,
   updateDomainBodySchema,
 } from "@/lib/zod/schemas/domains";
+import { nanoid } from "@dub/utils";
 import { waitUntil } from "@vercel/functions";
 import { NextResponse } from "next/server";
 
@@ -37,7 +39,11 @@ export const GET = withWorkspace(
 // PUT /api/domains/[domain] – edit a workspace's domain
 export const PATCH = withWorkspace(
   async ({ req, workspace, params }) => {
-    const { slug: domain, registeredDomain } = await getDomainOrThrow({
+    const {
+      slug: domain,
+      registeredDomain,
+      id: domainId,
+    } = await getDomainOrThrow({
       workspace,
       domain: params.domain,
       dubDomainChecks: true,
@@ -49,6 +55,7 @@ export const PATCH = withWorkspace(
       expiredUrl,
       notFoundUrl,
       archived,
+      logo,
     } = updateDomainBodySchema.parse(await parseRequestBody(req));
 
     if (workspace.plan === "free" && expiredUrl) {
@@ -85,6 +92,11 @@ export const PATCH = withWorkspace(
       }
     }
 
+    const logoUploaded =
+      logo && workspace.plan !== "free"
+        ? await storage.upload(`logos/${domainId}_${nanoid(7)}`, logo)
+        : null;
+
     const domainRecord = await prisma.domain.update({
       where: {
         slug: domain,
@@ -96,6 +108,7 @@ export const PATCH = withWorkspace(
         ...(workspace.plan != "free" && {
           expiredUrl,
           notFoundUrl,
+          ...(logoUploaded && { logo: logoUploaded.url }),
         }),
       },
       include: {

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

@@ -4,12 +4,13 @@ import { createLink, transformLink } from "@/lib/api/links";
 import { createId, parseRequestBody } from "@/lib/api/utils";
 import { withWorkspace } from "@/lib/auth";
 import { prisma } from "@/lib/prisma";
+import { storage } from "@/lib/storage";
 import {
   DomainSchema,
   createDomainBodySchema,
   getDomainsQuerySchemaExtended,
 } from "@/lib/zod/schemas/domains";
-import { DEFAULT_LINK_PROPS } from "@dub/utils";
+import { DEFAULT_LINK_PROPS, nanoid } from "@dub/utils";
 import { NextResponse } from "next/server";
 
 // GET /api/domains – get all domains for a workspace
@@ -77,7 +78,7 @@ export const GET = withWorkspace(
 export const POST = withWorkspace(
   async ({ req, workspace, session }) => {
     const body = await parseRequestBody(req);
-    const { slug, placeholder, expiredUrl, notFoundUrl } =
+    const { slug, placeholder, expiredUrl, notFoundUrl, logo } =
       createDomainBodySchema.parse(body);
 
     const totalDomains = await prisma.domain.count({
@@ -123,17 +124,25 @@ export const POST = withWorkspace(
       return new Response(vercelResponse.error.message, { status: 422 });
     }
 
+    const domainId = createId({ prefix: "dom_" });
+
+    const logoUploaded =
+      logo && workspace.plan !== "free"
+        ? await storage.upload(`logos/${domainId}_${nanoid(7)}`, logo)
+        : null;
+
     const [domainRecord, _] = await Promise.all([
       prisma.domain.create({
         data: {
-          id: createId({ prefix: "dom_" }),
+          id: domainId,
           slug: slug,
           projectId: workspace.id,
           primary: totalDomains === 0,
           ...(placeholder && { placeholder }),
           ...(workspace.plan !== "free" && {
             expiredUrl,
             notFoundUrl,
+            ...(logoUploaded && { logo: logoUploaded.url }),
           }),
         },
       }),

apps/web/app/api/qr/route.tsx

@@ -1,9 +1,10 @@
 import { handleAndReturnErrorResponse } from "@/lib/api/errors";
 import { ratelimitOrThrow } from "@/lib/api/utils";
 import { getShortLinkViaEdge, getWorkspaceViaEdge } from "@/lib/planetscale";
+import { getDomainViaEdge } from "@/lib/planetscale/get-domain-via-edge";
 import { QRCodeSVG } from "@/lib/qr/utils";
 import { getQRCodeQuerySchema } from "@/lib/zod/schemas/qr";
-import { DUB_QR_LOGO, getSearchParams } from "@dub/utils";
+import { DUB_QR_LOGO, getSearchParams, isDubDomain } from "@dub/utils";
 import { ImageResponse } from "next/og";
 import { NextRequest } from "next/server";
 
@@ -16,33 +17,14 @@ const CORS_HEADERS = {
 
 export async function GET(req: NextRequest) {
   try {
-    const params = getSearchParams(req.url);
-
-    let { url, logo, size, level, fgColor, bgColor, margin, hideLogo } =
-      getQRCodeQuerySchema.parse(params);
+    const paramsParsed = getQRCodeQuerySchema.parse(getSearchParams(req.url));
 
     await ratelimitOrThrow(req, "qr");
 
-    const shortLink = await getShortLinkViaEdge(url.split("?")[0]);
-    if (shortLink) {
-      const workspace = await getWorkspaceViaEdge(shortLink.projectId);
-      // Free workspaces should always use the default logo.
-      if (!workspace || workspace.plan === "free") {
-        logo = DUB_QR_LOGO;
-        /*
-          If:
-          - no logo is passed
-          - the workspace has a logo
-          - the hideLogo flag is not set
-          then we should use the workspace logo.
-        */
-      } else if (!logo && workspace.logo && !hideLogo) {
-        logo = workspace.logo;
-      }
-      // if the link is not on Dub, use the default logo.
-    } else {
-      logo = DUB_QR_LOGO;
-    }
+    const { logo, url, size, level, fgColor, bgColor, margin, hideLogo } =
+      paramsParsed;
+
+    const qrCodeLogo = await getQRCodeLogo({ url, logo, hideLogo });
 
     return new ImageResponse(
       QRCodeSVG({
@@ -52,10 +34,10 @@ export async function GET(req: NextRequest) {
         fgColor,
         bgColor,
         margin,
-        ...(logo
+        ...(qrCodeLogo
           ? {
               imageSettings: {
-                src: logo,
+                src: qrCodeLogo,
                 height: size / 4,
                 width: size / 4,
                 excavate: true,
@@ -75,6 +57,48 @@ export async function GET(req: NextRequest) {
   }
 }
 
+const getQRCodeLogo = async ({
+  url,
+  logo,
+  hideLogo,
+}: {
+  url: string;
+  logo: string | undefined;
+  hideLogo: boolean;
+}) => {
+  const shortLink = await getShortLinkViaEdge(url.split("?")[0]);
+
+  // Not a Dub link
+  if (!shortLink) {
+    return DUB_QR_LOGO;
+  }
+
+  // Dub owned domain
+  if (isDubDomain(shortLink.domain)) {
+    return DUB_QR_LOGO;
+  }
+
+  // hideLogo is set or logo is passed
+  if (hideLogo || logo) {
+    const workspace = await getWorkspaceViaEdge(shortLink.projectId);
+
+    if (workspace?.plan === "free") {
+      return DUB_QR_LOGO;
+    }
+
+    if (hideLogo) {
+      return null;
+    }
+
+    return logo;
+  }
+
+  // if no logo is passed, use domain logo
+  const domain = await getDomainViaEdge(shortLink.domain);
+
+  return domain?.logo || DUB_QR_LOGO;
+};
+
 export function OPTIONS() {
   return new Response(null, {
     status: 204,

apps/web/app/api/stripe/webhook/customer-subscription-deleted.ts

@@ -111,6 +111,16 @@ export async function customerSubscriptionDeleted(event: Stripe.Event) {
       },
     }),
 
+    // remove logo from all domains for the workspace
+    prisma.domain.updateMany({
+      where: {
+        projectId: workspace.id,
+      },
+      data: {
+        logo: null,
+      },
+    }),
+
     // remove root domain link for all domains from MySQL
     prisma.link.updateMany({
       where: {

apps/web/lib/planetscale/get-domain-via-edge.ts

@@ -3,9 +3,10 @@ import { EdgeDomainProps } from "./types";
 
 export const getDomainViaEdge = async (domain: string) => {
   const { rows } =
-    (await conn.execute("SELECT * FROM Domain WHERE slug = ?", [domain])) || {};
+    (await conn.execute<EdgeDomainProps>(
+      "SELECT * FROM Domain WHERE slug = ?",
+      [domain],
+    )) || {};
 
-  return rows && Array.isArray(rows) && rows.length > 0
-    ? (rows[0] as EdgeDomainProps)
-    : null;
+  return rows && Array.isArray(rows) && rows.length > 0 ? rows[0] : null;
 };

apps/web/lib/planetscale/types.ts

@@ -23,6 +23,7 @@ export interface EdgeLinkProps {
 export interface EdgeDomainProps {
   id: string;
   slug: string;
+  logo: string | null;
   verified: number;
   placeholder: string;
   expiredUrl: string | null;

apps/web/lib/swr/use-domain.ts

@@ -0,0 +1,18 @@
+import { fetcher } from "@dub/utils";
+import useSWR from "swr";
+import { DomainProps } from "../types";
+import useWorkspace from "./use-workspace";
+
+export default function useDomain(slug: string) {
+  const { id: workspaceId } = useWorkspace();
+
+  const { data: domain, error } = useSWR<DomainProps>(
+    workspaceId && slug && `/api/domains/${slug}?workspaceId=${workspaceId}`,
+    fetcher,
+  );
+
+  return {
+    ...domain,
+    loading: !domain && !error,
+  };
+}

apps/web/lib/types.ts

@@ -166,6 +166,7 @@ export interface DomainProps {
   projectId: string;
   link?: LinkProps;
   registeredDomain?: RegisteredDomainProps;
+  logo?: string;
 }
 
 export interface RegisteredDomainProps {

apps/web/lib/zod/schemas/domains.ts

@@ -41,6 +41,7 @@ export const DomainSchema = z.object({
       "The URL to redirect to when a link under this domain doesn't exist.",
     )
     .openapi({ example: "https://acme.com/not-found" }),
+  logo: z.string().nullable().describe("The logo of the domain."),
   createdAt: z.date().describe("The date the domain was created."),
   updatedAt: z.date().describe("The date the domain was last updated."),
   registeredDomain: z
@@ -118,6 +119,7 @@ export const createDomainBodySchema = z.object({
       "Provide context to your teammates in the link creation modal by showing them an example of a link to be shortened.",
     )
     .openapi({ example: "https://dub.co/help/article/what-is-dub" }),
+  logo: z.string().url().nullish().describe("The logo of the domain."),
 });
 
 export const updateDomainBodySchema = createDomainBodySchema.partial();

apps/web/prisma/schema/domain.prisma

@@ -8,6 +8,7 @@ model Domain {
   primary          Boolean           @default(false)
   archived         Boolean           @default(false)
   lastChecked      DateTime          @default(now())
+  logo             String?
   links            Link[]
   project          Project?          @relation(fields: [projectId], references: [id], onDelete: Cascade)
   projectId        String?

apps/web/scripts/backfill-domain-logo.ts

@@ -0,0 +1,38 @@
+import { prisma } from "@/lib/prisma";
+import "dotenv-flow/config";
+
+async function main() {
+  const workspaces = await prisma.project.findMany({
+    where: {
+      plan: {
+        not: "free",
+      },
+      logo: {
+        not: null,
+      },
+    },
+    take: 1000,
+  });
+
+  if (!workspaces.length) {
+    console.log("No workspaces found.");
+    return;
+  }
+
+  const updated = await Promise.all(
+    workspaces.map((workspace) =>
+      prisma.domain.updateMany({
+        where: {
+          projectId: workspace.id,
+        },
+        data: {
+          logo: workspace.logo,
+        },
+      }),
+    ),
+  );
+
+  console.log(`Updated ${updated.length} domains.`);
+}
+
+main();