Enforce cloud website limits by subscription tier

src/app/api/websites/route.ts

@@ -1,16 +1,15 @@
 import { z } from 'zod';
 import { ENTITY_TYPE } from '@/lib/constants';
 import { uuid } from '@/lib/crypto';
-import { fetchAccount } from '@/lib/load';
+import { fetchAccount, fetchTeam } from '@/lib/load';
 import { getQueryFilters, parseRequest } from '@/lib/request';
 import { json, unauthorized } from '@/lib/response';
 import { pagingParams, searchParams, sortingParams } from '@/lib/schema';
+import { getCloudWebsiteLimit } from '@/lib/subscription';
 import { canCreateTeamWebsite, canCreateWebsite } from '@/permissions';
-import { createShare, createWebsite, getWebsiteCount } from '@/queries/prisma';
+import { createShare, createWebsite, getTeamWebsiteCount, getWebsiteCount } from '@/queries/prisma';
 import { getAllUserWebsitesIncludingTeamAccess, getUserWebsites } from '@/queries/prisma/website';
 
-const CLOUD_WEBSITE_LIMIT = 3;
-
 export async function GET(request: Request) {
   const schema = z.object({
     ...pagingParams,
@@ -53,13 +52,16 @@ export async function POST(request: Request) {
 
   const { id, name, domain, shareId, teamId } = body;
 
-  if (process.env.CLOUD_MODE && !teamId) {
-    const account = await fetchAccount(auth.user.id);
+  if (process.env.CLOUD_MODE) {
+    const account = teamId ? await fetchTeam(teamId) : await fetchAccount(auth.user.id);
+    const websiteLimit = getCloudWebsiteLimit(account);
 
-    if (!account?.hasSubscription) {
-      const count = await getWebsiteCount(auth.user.id);
+    if (websiteLimit !== null) {
+      const count = teamId
+        ? await getTeamWebsiteCount(teamId)
+        : await getWebsiteCount(auth.user.id);
 
-      if (count >= CLOUD_WEBSITE_LIMIT) {
+      if (count >= websiteLimit) {
         return unauthorized({ message: 'Website limit reached.' });
       }
     }

src/lib/subscription.test.ts

@@ -0,0 +1,24 @@
+import { describe, expect, test } from 'vitest';
+import {
+  CLOUD_FREE_WEBSITE_LIMIT,
+  CLOUD_PRO_WEBSITE_LIMIT,
+  getCloudWebsiteLimit,
+} from './subscription';
+
+describe('getCloudWebsiteLimit', () => {
+  test('limits accounts without a subscription to the free website limit', () => {
+    expect(getCloudWebsiteLimit(null)).toBe(CLOUD_FREE_WEBSITE_LIMIT);
+    expect(getCloudWebsiteLimit({ hasSubscription: false })).toBe(CLOUD_FREE_WEBSITE_LIMIT);
+  });
+
+  test('limits Pro accounts to the Pro website limit', () => {
+    expect(getCloudWebsiteLimit({ hasSubscription: true, isPro: true })).toBe(
+      CLOUD_PRO_WEBSITE_LIMIT,
+    );
+  });
+
+  test('does not limit Business or no-billing accounts', () => {
+    expect(getCloudWebsiteLimit({ hasSubscription: true, isBusiness: true })).toBeNull();
+    expect(getCloudWebsiteLimit({ hasSubscription: true, isNoBilling: true })).toBeNull();
+  });
+});

src/lib/subscription.ts

@@ -0,0 +1,25 @@
+export interface SubscriptionAccount {
+  isPro?: boolean | null;
+  isBusiness?: boolean | null;
+  isNoBilling?: boolean | null;
+  hasSubscription?: boolean | null;
+}
+
+export const CLOUD_FREE_WEBSITE_LIMIT = 3;
+export const CLOUD_PRO_WEBSITE_LIMIT = 20;
+
+export function getCloudWebsiteLimit(account?: SubscriptionAccount | null): number | null {
+  if (!account?.hasSubscription) {
+    return CLOUD_FREE_WEBSITE_LIMIT;
+  }
+
+  if (account.isNoBilling || account.isBusiness) {
+    return null;
+  }
+
+  if (account.isPro) {
+    return CLOUD_PRO_WEBSITE_LIMIT;
+  }
+
+  return null;
+}

src/queries/prisma/website.ts

@@ -46,7 +46,10 @@ export async function getWebsites(criteria: Prisma.WebsiteFindManyArgs, filters:
   return attachShareIdToWebsites(websites);
 }
 
-export async function getAllUserWebsitesIncludingTeamAccess(userId: string, filters?: QueryFilters) {
+export async function getAllUserWebsitesIncludingTeamAccess(
+  userId: string,
+  filters?: QueryFilters,
+) {
   return getWebsites(
     {
       where: {
@@ -263,6 +266,15 @@ export async function getWebsiteCount(userId: string) {
   });
 }
 
+export async function getTeamWebsiteCount(teamId: string) {
+  return prisma.client.website.count({
+    where: {
+      teamId,
+      deletedAt: null,
+    },
+  });
+}
+
 export async function attachShareIdToWebsite(website: Website) {
   const share = await prisma.client.share.findFirst({
     where: {