Improve support for workspaces / folders with 1M+ links

apps/web/app/api/links/count/route.ts

@@ -1,16 +1,23 @@
 import { getFolderIdsToFilter } from "@/lib/analytics/get-folder-ids-to-filter";
+import { getStartEndDates } from "@/lib/analytics/utils/get-start-end-dates";
 import { getDomainOrThrow } from "@/lib/api/domains/get-domain-or-throw";
-import { DubApiError } from "@/lib/api/errors";
 import { getLinksCount } from "@/lib/api/links";
 import { withWorkspace } from "@/lib/auth";
 import { verifyFolderAccess } from "@/lib/folder/permissions";
-import { getLinksCountQuerySchema } from "@/lib/zod/schemas/links";
+import { getLinksCountTB } from "@/lib/tinybird/get-links-count";
+import { WorkspaceProps } from "@/lib/types";
+import { getLinksCountQuerySchemaExtended } from "@/lib/zod/schemas/links";
 import { NextResponse } from "next/server";
 
+// A mega workspace is a workspace with more than 1M links
+function isMegaWorkspace(workspace: Pick<WorkspaceProps, "totalLinks">) {
+  return workspace.totalLinks > 1_000_000;
+}
+
 // GET /api/links/count – get the number of links for a workspace
 export const GET = withWorkspace(
   async ({ headers, searchParams, workspace, session }) => {
-    const params = getLinksCountQuerySchema.parse(searchParams);
+    const params = getLinksCountQuerySchemaExtended.parse(searchParams);
     const {
       groupBy,
       domain,
@@ -20,26 +27,51 @@ export const GET = withWorkspace(
       tagIds,
       tagNames,
       tenantId,
+      start,
+      end,
+      interval,
+      timezone,
     } = params;
 
     if (domain) {
       await getDomainOrThrow({ domain, workspace: workspace });
     }
 
     if (folderId) {
-      const selectedFolder = await verifyFolderAccess({
+      await verifyFolderAccess({
         workspace,
         userId: session.user.id,
         folderId,
         requiredPermission: "folders.read",
       });
+    }
 
-      if (selectedFolder.type === "mega") {
-        throw new DubApiError({
-          code: "bad_request",
-          message: "Cannot get links count for mega folders.",
+    // For mega workspaces, we fetch the count via Tinybird instead of MySQL
+    if (isMegaWorkspace(workspace)) {
+      // We don't support groupBy for mega workspaces
+      if (groupBy) {
+        return NextResponse.json([], {
+          headers,
         });
       }
+
+      const { startDate, endDate } = getStartEndDates({
+        start,
+        end,
+        interval,
+      });
+
+      const { data } = await getLinksCountTB({
+        workspaceId: workspace.id,
+        folderId,
+        timezone,
+        start: startDate.toISOString().replace("T", " ").replace("Z", ""),
+        end: endDate.toISOString().replace("T", " ").replace("Z", ""),
+      });
+
+      return NextResponse.json(data[0].count, {
+        headers,
+      });
     }
 
     /* we only need to get the folder ids if we are:

apps/web/app/app.dub.co/(dashboard)/[slug]/links/page-client.tsx

@@ -4,7 +4,7 @@ import {
   useCheckFolderPermission,
   useFolderPermissions,
 } from "@/lib/swr/use-folder-permissions";
-import { useIsMegaFolder } from "@/lib/swr/use-is-mega-folder";
+import { useIsMegaWorkspace } from "@/lib/swr/use-is-mega-workspace";
 import useLinks from "@/lib/swr/use-links";
 import useWorkspace from "@/lib/swr/use-workspace";
 import { useWorkspaceStore } from "@/lib/swr/use-workspace-store";
@@ -72,6 +72,7 @@ function WorkspaceLinks() {
   const { isValidating } = useLinks();
   const searchParams = useSearchParams();
   const workspace = useWorkspace();
+  const { isMegaWorkspace } = useIsMegaWorkspace();
   const { LinkBuilder, CreateLinkButton } = useLinkBuilder();
   const { AddEditTagModal, setShowAddEditTagModal } = useAddEditTagModal();
 
@@ -86,7 +87,6 @@ function WorkspaceLinks() {
   } = useLinkFilters();
 
   const folderId = searchParams.get("folderId");
-  const { isMegaFolder } = useIsMegaFolder();
 
   const { isLoading } = useFolderPermissions();
   const canCreateLinks = useCheckFolderPermission(
@@ -140,7 +140,7 @@ function WorkspaceLinks() {
         <PageWidthWrapper className="flex flex-col gap-y-3">
           <div className="flex flex-wrap items-center justify-between gap-2">
             <div className="flex w-full grow gap-2 md:w-auto">
-              {!isMegaFolder && (
+              {!isMegaWorkspace && (
                 <div className="grow basis-0 md:grow-0">
                   <Filter.Select
                     filters={filters}
@@ -209,7 +209,7 @@ function WorkspaceLinks() {
                   loading={isValidating}
                   inputClassName="h-10"
                   placeholder={
-                    isMegaFolder
+                    isMegaWorkspace
                       ? "Search by short link"
                       : "Search by short link or URL"
                   }

apps/web/lib/api/links/get-links-for-workspace.ts

@@ -18,6 +18,8 @@ export async function getLinksForWorkspace({
   sortOrder,
   page,
   pageSize,
+  startingAfter,
+  endingBefore,
   userId,
   showArchived,
   withTags,
@@ -57,6 +59,8 @@ export async function getLinksForWorkspace({
     } catch (e) {}
   }
 
+  const cursorId = startingAfter || endingBefore;
+
   const links = await prisma.link.findMany({
     where: {
       ...(linkIds && { id: { in: linkIds } }),
@@ -149,8 +153,20 @@ export async function getLinksForWorkspace({
     orderBy: {
       [sortBy]: sortOrder,
     },
-    take: pageSize,
-    skip: (page - 1) * pageSize,
+    ...(cursorId
+      ? {
+          cursor: { id: cursorId },
+          skip: 1, // skip the cursor item itself
+          take: startingAfter ? pageSize : -pageSize,
+        }
+      : page
+        ? {
+            skip: (page - 1) * pageSize,
+            take: pageSize,
+          }
+        : {
+            take: pageSize, // default if neither page nor cursorId
+          }),
   });
 
   return links.map((link) => transformLink(link));

apps/web/lib/swr/use-is-mega-workspace.ts

@@ -0,0 +1,9 @@
+import useWorkspace from "./use-workspace";
+
+export function useIsMegaWorkspace() {
+  const { totalLinks } = useWorkspace();
+
+  return {
+    isMegaWorkspace: totalLinks && totalLinks > 1_000_000,
+  };
+}

apps/web/lib/swr/use-links-count.ts

@@ -4,7 +4,6 @@ import { useEffect, useState } from "react";
 import useSWR from "swr";
 import z from "../zod";
 import { getLinksCountQuerySchema } from "../zod/schemas/links";
-import { useIsMegaFolder } from "./use-is-mega-folder";
 import useWorkspace from "./use-workspace";
 
 const partialQuerySchema = getLinksCountQuerySchema.partial();
@@ -20,17 +19,16 @@ export default function useLinksCount<T = any>({
 } = {}) {
   const { id: workspaceId } = useWorkspace();
   const { getQueryString } = useRouterStuff();
-
   const [admin, setAdmin] = useState(false);
+
   useEffect(() => {
     if (window.location.host.startsWith("admin.")) {
       setAdmin(true);
     }
   }, []);
-  const { isMegaFolder } = useIsMegaFolder();
 
   const { data, error } = useSWR<any>(
-    workspaceId && !isMegaFolder && enabled
+    workspaceId && enabled
       ? `/api/links/count${getQueryString(
           {
             workspaceId,

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

@@ -5,7 +5,7 @@ import useSWR, { SWRConfiguration } from "swr";
 import { z } from "zod";
 import { ExpandedLinkProps, UserProps } from "../types";
 import { getLinksQuerySchemaExtended } from "../zod/schemas/links";
-import { useIsMegaFolder } from "./use-is-mega-folder";
+import { useIsMegaWorkspace } from "./use-is-mega-workspace";
 import useWorkspace from "./use-workspace";
 
 const partialQuerySchema = getLinksQuerySchemaExtended.partial();
@@ -16,17 +16,17 @@ export default function useLinks(
 ) {
   const { id: workspaceId } = useWorkspace();
   const { getQueryString } = useRouterStuff();
-  const { isMegaFolder } = useIsMegaFolder();
-
+  const { isMegaWorkspace } = useIsMegaWorkspace();
   const [admin, setAdmin] = useState(false);
+
   useEffect(() => {
     if (window.location.host.startsWith("admin.")) {
       setAdmin(true);
     }
   }, []);
 
   const {
-    data: links,
+    data: response,
     isValidating,
     error,
   } = useSWR<
@@ -41,8 +41,8 @@ export default function useLinks(
             includeUser: "true",
             includeDashboard: "true",
             ...opts,
-            // don't show archived on mega folders
-            ...(isMegaFolder
+            // don't show archived on mega workspaces
+            ...(isMegaWorkspace
               ? {
                   showArchived: "false",
                 }
@@ -59,6 +59,8 @@ export default function useLinks(
               "sortBy",
               "sortOrder",
               "showArchived",
+              "startingAfter",
+              "endingBefore",
             ],
           },
         )}`
@@ -75,7 +77,7 @@ export default function useLinks(
   );
 
   return {
-    links,
+    links: response,
     isValidating,
     error,
   };

apps/web/lib/tinybird/get-links-count.ts

@@ -0,0 +1,21 @@
+import z from "../zod";
+import { analyticsFilterTB } from "../zod/schemas/analytics";
+import { tb } from "./client";
+
+const parameters = analyticsFilterTB.pick({
+  workspaceId: true,
+  folderId: true,
+  start: true,
+  end: true,
+  timezone: true,
+});
+
+const data = z.object({
+  count: z.number(),
+});
+
+export const getLinksCountTB = tb.buildPipe({
+  pipe: "get_links_count",
+  parameters,
+  data,
+});

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

@@ -1,3 +1,4 @@
+import { DATE_RANGE_INTERVAL_PRESETS } from "@/lib/analytics/constants";
 import { ErrorCode } from "@/lib/api/errors";
 import z from "@/lib/zod";
 import {
@@ -187,6 +188,15 @@ export const getLinksCountQuerySchema = LinksQuerySchema.merge(
   }),
 );
 
+export const getLinksCountQuerySchemaExtended = getLinksCountQuerySchema.merge(
+  z.object({
+    start: parseDateSchema.optional(),
+    end: parseDateSchema.optional(),
+    interval: z.enum(DATE_RANGE_INTERVAL_PRESETS).optional(),
+    timezone: z.string().optional(),
+  }),
+);
+
 export const linksExportQuerySchema = getLinksQuerySchemaBase
   .omit({ page: true, pageSize: true })
   .merge(
@@ -769,6 +779,20 @@ export const getLinksQuerySchemaExtended = getLinksQuerySchemaBase.merge(
       .enum(["fuzzy", "exact"])
       .default("fuzzy")
       .describe("Search mode to filter by."),
+    startingAfter: z
+      .string()
+      .trim()
+      .optional()
+      .describe(
+        "A cursor to use in pagination. Returns items after the given link ID.",
+      ),
+    endingBefore: z
+      .string()
+      .trim()
+      .optional()
+      .describe(
+        "A cursor to use in pagination. Returns items before the given link ID.",
+      ),
   }),
 );