cloud: enable workspace slots for partners

apps/web/src/components/WorkspaceMenu.tsx

@@ -1,6 +1,6 @@
+import { useRouter } from "next/navigation";
 import { Button, Menu, Transition } from "@headlessui/react";
 import { t } from "@lingui/core/macro";
-import { useRouter } from "next/navigation";
 import { env } from "next-runtime-env";
 import { Fragment, useState } from "react";
 import { HiCheck, HiMagnifyingGlass } from "react-icons/hi2";
@@ -9,6 +9,7 @@ import { twMerge } from "tailwind-merge";
 import { useKeyboardShortcut } from "~/providers/keyboard-shortcuts";
 import { useModal } from "~/providers/modal";
 import { useWorkspace } from "~/providers/workspace";
+import { api } from "~/utils/api";
 import CommandPallette from "./CommandPallette";
 import { Tooltip } from "./Tooltip";
 
@@ -20,6 +21,8 @@ export default function WorkspaceMenu({
   const { workspace, isLoading, availableWorkspaces, switchWorkspace } =
     useWorkspace();
   const { openModal } = useModal();
+  const { data: hasPartnerSlot } =
+    api.workspace.hasAvailablePartnerSlot.useQuery();
   const router = useRouter();
   const [isOpen, setIsOpen] = useState(false);
 
@@ -150,11 +153,19 @@ export default function WorkspaceMenu({
             <div className="border-t-[1px] border-light-600 p-1 dark:border-dark-500">
               <Menu.Item>
                 <button
-                  onClick={() =>
+                  onClick={() => {
-                    env("NEXT_PUBLIC_KAN_ENV") === "cloud"
+                    if (env("NEXT_PUBLIC_KAN_ENV") !== "cloud") {
-                      ? router.push(`/onboarding/select-plan?returnUrl=${encodeURIComponent(window.location.pathname)}`)
+                      openModal("NEW_WORKSPACE");
-                      : openModal("NEW_WORKSPACE")
+                    } else if (hasPartnerSlot) {
-                  }
+                      router.push(
+                        `/onboarding/workspace?partner=1&returnUrl=${encodeURIComponent(window.location.pathname)}`,
+                      );
+                    } else {
+                      router.push(
+                        `/onboarding/select-plan?returnUrl=${encodeURIComponent(window.location.pathname)}`,
+                      );
+                    }
+                  }}
                   className="flex w-full items-center justify-between rounded-[5px] px-3 py-2 text-left text-xs text-neutral-900 hover:bg-light-200 dark:text-dark-1000 dark:hover:bg-dark-400"
                 >
                   {t`Create workspace`}

apps/web/src/pages/api/partner/_utils.ts

@@ -4,12 +4,13 @@ export interface TierConfig {
   plan: WorkspacePlan;
   seats: number | null;
   unlimitedSeats: boolean;
+  workspaceSlots: number;
 }
 
 const TIER_MAP: Record<number, TierConfig> = {
-  1: { plan: "team", seats: 5, unlimitedSeats: false },
+  1: { plan: "team", seats: 5, unlimitedSeats: false, workspaceSlots: 1 },
-  2: { plan: "pro", seats: 15, unlimitedSeats: false },
+  2: { plan: "pro", seats: 15, unlimitedSeats: false, workspaceSlots: 2 },
-  3: { plan: "pro", seats: null, unlimitedSeats: true },
+  3: { plan: "pro", seats: null, unlimitedSeats: true, workspaceSlots: 4 },
 };
 
 export function tierConfig(tier: number): TierConfig {

apps/web/src/pages/api/partner/callback.ts

@@ -4,7 +4,6 @@ import { createNextApiContext } from "@kan/api/trpc";
 import { withApiLogging } from "@kan/api/utils/apiLogging";
 import { withRateLimit } from "@kan/api/utils/rateLimit";
 import * as subscriptionRepo from "@kan/db/repository/subscription.repo";
-import * as workspaceRepo from "@kan/db/repository/workspace.repo";
 import { createLogger } from "@kan/logger";
 
 import { tierConfig } from "./_utils";
@@ -110,11 +109,15 @@ export default withRateLimit(
     const { db, user } = await createNextApiContext(req);
 
     const cfg = tierConfig(license.tier);
-    const isActive = license.status === "active";
+    const status = license.status === "active" ? "active" : "inactive";
-    const status = isActive ? "active" : "inactive";
 
-    if (!user) {
+    // Ensure subscription slots exist — webhook may have already created them
-      await subscriptionRepo.upsertByPartnerLicenseKey(
+    const existing = await subscriptionRepo.getAllByPartnerLicenseKey(
+      db,
+      license.license_key,
+    );
+    if (existing.length === 0) {
+      await subscriptionRepo.createPartnerLicenseSlots(
         db,
         license.license_key,
         {
@@ -124,45 +127,18 @@ export default withRateLimit(
           seats: cfg.seats,
           unlimitedSeats: cfg.unlimitedSeats,
         },
+        cfg.workspaceSlots,
       );
+    }
+
+    if (!user) {
       return res.redirect(
         `/partner/activate?license_key=${encodeURIComponent(license.license_key)}`,
       );
     }
 
-    const memberships = await workspaceRepo.getAllByUserId(db, user.id);
+    return res.redirect(
-    const workspace = memberships?.[0]?.workspace;
+      `/api/partner/link?license_key=${encodeURIComponent(license.license_key)}`,
-
+    );
-    if (!workspace) {
-      await subscriptionRepo.upsertByPartnerLicenseKey(
-        db,
-        license.license_key,
-        {
-          plan: cfg.plan,
-          status,
-          partnerTier: license.tier,
-          seats: cfg.seats,
-          unlimitedSeats: cfg.unlimitedSeats,
-        },
-      );
-      return res.redirect(
-        `/onboarding/workspace?license_key=${encodeURIComponent(license.license_key)}`,
-      );
-    }
-
-    await subscriptionRepo.upsertByPartnerLicenseKey(db, license.license_key, {
-      plan: cfg.plan,
-      status,
-      partnerTier: license.tier,
-      seats: cfg.seats,
-      unlimitedSeats: cfg.unlimitedSeats,
-      referenceId: workspace.publicId,
-    });
-
-    if (isActive) {
-      await workspaceRepo.update(db, workspace.publicId, { plan: cfg.plan });
-    }
-
-    return res.redirect(`/?partner_activated=1`);
   }),
 );

apps/web/src/pages/api/partner/link.ts

@@ -27,38 +27,72 @@ export default withRateLimit(
       );
     }
 
-    const sub = await subscriptionRepo.getByPartnerLicenseKey(db, license_key);
+    const allSlots = await subscriptionRepo.getAllByPartnerLicenseKey(
+      db,
+      license_key,
+    );
 
-    if (!sub) {
+    if (!allSlots.length) {
       return res.redirect("/boards?partner_error=invalid_license");
     }
 
-    if (sub.status !== "active") {
+    const activeSlots = allSlots.filter((s) =>
+      ["active", "trialing"].includes(s.status),
+    );
+
+    if (!activeSlots.length) {
       return res.redirect("/boards?partner_error=license_inactive");
     }
 
-    const memberships = await workspaceRepo.getAllByUserId(db, user.id);
+    const unlinkedSlot = activeSlots.find((s) => !s.referenceId);
-    const workspace = memberships?.[0]?.workspace;
 
-    if (!workspace) {
+    if (!unlinkedSlot) {
+      return res.redirect("/boards?partner_activated=1");
+    }
+
+    const linkedIds = new Set(
+      activeSlots.filter((s) => s.referenceId).map((s) => s.referenceId!),
+    );
+
+    const memberships = await workspaceRepo.getAllByUserId(db, user.id);
+    const availableWorkspace = memberships
+      .map((m) => m.workspace)
+      .find((w) => w && !w.deletedAt && !linkedIds.has(w.publicId));
+
+    if (!availableWorkspace) {
       return res.redirect(
         `/onboarding/workspace?license_key=${encodeURIComponent(license_key)}`,
       );
     }
 
-    await subscriptionRepo.upsertByPartnerLicenseKey(db, license_key, {
+    await subscriptionRepo.updateById(db, unlinkedSlot.id, {
-      plan: sub.plan,
+      referenceId: availableWorkspace.publicId,
-      status: sub.status,
-      partnerTier: sub.partnerTier ?? 1,
-      seats: sub.seats ?? null,
-      unlimitedSeats: sub.unlimitedSeats,
-      referenceId: workspace.publicId,
     });
 
-    await workspaceRepo.update(db, workspace.publicId, {
+    await workspaceRepo.update(db, availableWorkspace.publicId, {
-      plan: sub.plan as "free" | "team" | "pro" | "enterprise",
+      plan: unlinkedSlot.plan as "free" | "team" | "pro" | "enterprise",
     });
 
+    const remainingUnlinked = activeSlots.filter(
+      (s) => !s.referenceId && s.id !== unlinkedSlot.id,
+    );
+
+    if (remainingUnlinked.length > 0) {
+      const updatedLinkedIds = new Set([
+        ...linkedIds,
+        availableWorkspace.publicId,
+      ]);
+      const hasMoreAvailableWorkspace = memberships
+        .map((m) => m.workspace)
+        .some((w) => w && !w.deletedAt && !updatedLinkedIds.has(w.publicId));
+
+      if (hasMoreAvailableWorkspace) {
+        return res.redirect(
+          `/api/partner/link?license_key=${encodeURIComponent(license_key)}`,
+        );
+      }
+    }
+
     return res.redirect("/boards?partner_activated=1");
   }),
 );

apps/web/src/pages/api/partner/webhook.ts

@@ -43,6 +43,12 @@ function verifySignature(
   }
 }
 
+function hasReferenceId<T extends { referenceId: string | null | undefined }>(
+  s: T,
+): s is T & { referenceId: string } {
+  return !!s.referenceId;
+}
+
 interface WebhookPayload {
   event:
     | "purchase"
@@ -94,107 +100,169 @@ export default withApiLogging(
     const { db } = await createNextApiContext(req);
 
     switch (event) {
-      case "purchase": {
+      case "purchase":
-        const cfg = tierConfig(tier);
-        await subscriptionRepo.upsertByPartnerLicenseKey(db, license_key, {
-          plan: cfg.plan,
-          status: license_status,
-          partnerTier: tier,
-          seats: cfg.seats,
-          unlimitedSeats: cfg.unlimitedSeats,
-        });
-        break;
-      }
-
       case "activate": {
         const cfg = tierConfig(tier);
-        await subscriptionRepo.upsertByPartnerLicenseKey(db, license_key, {
+        const existing = await subscriptionRepo.getAllByPartnerLicenseKey(
-          plan: cfg.plan,
+          db,
-          status: "active",
+          license_key,
-          partnerTier: tier,
+        );
-          seats: cfg.seats,
+        const status = event === "activate" ? "active" : license_status;
-          unlimitedSeats: cfg.unlimitedSeats,
+
-        });
+        if (existing.length === 0) {
+          await subscriptionRepo.createPartnerLicenseSlots(
+            db,
+            license_key,
+            {
+              plan: cfg.plan,
+              status,
+              partnerTier: tier,
+              seats: cfg.seats,
+              unlimitedSeats: cfg.unlimitedSeats,
+            },
+            cfg.workspaceSlots,
+          );
+        } else {
+          await subscriptionRepo.updateAllByPartnerLicenseKey(db, license_key, {
+            plan: cfg.plan,
+            status,
+            partnerTier: tier,
+            seats: cfg.seats,
+            unlimitedSeats: cfg.unlimitedSeats,
+          });
+        }
         break;
       }
 
       case "deactivate": {
-        const sub = await subscriptionRepo.getByPartnerLicenseKey(
+        const allSlots = await subscriptionRepo.getAllByPartnerLicenseKey(
           db,
           license_key,
         );
-        if (sub) {
+
-          const [, allSubs] = await Promise.all([
+        await Promise.all(
-            subscriptionRepo.updateById(db, sub.id, {
+          allSlots.filter(hasReferenceId).map(async (slot) => {
-              plan: "free",
+            const siblingSubs = await subscriptionRepo.getByReferenceId(
-              status: "canceled",
+              db,
-            }),
+              slot.referenceId,
-            sub.referenceId
-              ? subscriptionRepo.getByReferenceId(db, sub.referenceId)
-              : Promise.resolve([]),
-          ]);
-          if (sub.referenceId) {
-            const hasActiveSub = getActiveSubscriptions(allSubs).some(
-              (s) => s.id !== sub.id,
             );
-            if (!hasActiveSub) {
+            const hasOtherActiveSub = getActiveSubscriptions(siblingSubs).some(
-              await cancelWorkspaceAccess(db, sub.referenceId);
+              (s) => s.id !== slot.id,
+            );
+            if (!hasOtherActiveSub) {
+              await cancelWorkspaceAccess(db, slot.referenceId);
             }
-          }
+          }),
-        }
+        );
+
+        await subscriptionRepo.updateAllByPartnerLicenseKey(db, license_key, {
+          plan: "free",
+          status: "canceled",
+          unlimitedSeats: false,
+          seats: null,
+        });
+
         break;
       }
 
       case "upgrade":
       case "downgrade": {
         const lookupKey = prev_license_key ?? license_key;
-        const sub = await subscriptionRepo.getByPartnerLicenseKey(
+        const existing = await subscriptionRepo.getAllByPartnerLicenseKey(
           db,
           lookupKey,
         );
-        if (sub) {
-          const cfg = tierConfig(tier);
-          await subscriptionRepo.upsertByPartnerLicenseKey(db, license_key, {
-            plan: cfg.plan,
-            status: "active",
-            partnerTier: tier,
-            seats: cfg.seats,
-            unlimitedSeats: cfg.unlimitedSeats,
-            referenceId: sub.referenceId ?? undefined,
-          });
-          if (prev_license_key) {
-            await subscriptionRepo.updateById(db, sub.id, {
-              status: "inactive",
-            });
-          }
-          if (sub.referenceId) {
-            await workspaceRepo.update(db, sub.referenceId, { plan: cfg.plan });
-          }
-        }
-        break;
-      }
 
-      case "migrate": {
+        if (existing.length === 0) break;
-        if (prev_license_key) {
+
-          const sub = await subscriptionRepo.getByPartnerLicenseKey(
+        const cfg = tierConfig(tier);
+        const newCount = cfg.workspaceSlots;
+
+        // Prefer keeping linked slots; among linked, keep in insertion order (LIFO removal)
+        const preferKeep = [
+          ...existing.filter((s) => s.referenceId),
+          ...existing.filter((s) => !s.referenceId),
+        ];
+
+        const slotsToKeep = preferKeep.slice(0, newCount);
+        const slotsToRemove = preferKeep.slice(newCount);
+
+        await Promise.all([
+          ...slotsToKeep.map((slot) =>
+            subscriptionRepo.updateById(db, slot.id, {
+              plan: cfg.plan,
+              seats: cfg.seats,
+              unlimitedSeats: cfg.unlimitedSeats,
+              partnerTier: tier,
+              status: "active",
+              ...(prev_license_key ? { partnerLicenseKey: license_key } : {}),
+            }),
+          ),
+          ...slotsToKeep
+            .filter(hasReferenceId)
+            .map((s) =>
+              workspaceRepo.update(db, s.referenceId, { plan: cfg.plan }),
+            ),
+        ]);
+
+        if (newCount > existing.length) {
+          await subscriptionRepo.createPartnerLicenseSlots(
             db,
-            prev_license_key,
+            license_key,
-          );
+            {
-          if (sub) {
-            const cfg = tierConfig(tier);
-            await subscriptionRepo.upsertByPartnerLicenseKey(db, license_key, {
               plan: cfg.plan,
               status: "active",
               partnerTier: tier,
               seats: cfg.seats,
               unlimitedSeats: cfg.unlimitedSeats,
-              referenceId: sub.referenceId ?? undefined,
+            },
-            });
+            newCount - existing.length,
-            await subscriptionRepo.updateById(db, sub.id, {
+          );
-              status: "inactive",
-            });
-          }
         }
+
+        if (slotsToRemove.length > 0) {
+          await Promise.all([
+            ...slotsToRemove
+              .filter(hasReferenceId)
+              .map((s) => cancelWorkspaceAccess(db, s.referenceId)),
+            ...slotsToRemove.map((s) =>
+              subscriptionRepo.updateById(db, s.id, {
+                plan: "free",
+                status: "inactive",
+                unlimitedSeats: false,
+                seats: null,
+              }),
+            ),
+          ]);
+        }
+
+        break;
+      }
+
+      case "migrate": {
+        if (!prev_license_key) break;
+
+        const existing = await subscriptionRepo.getAllByPartnerLicenseKey(
+          db,
+          prev_license_key,
+        );
+        if (existing.length === 0) break;
+
+        const cfg = tierConfig(tier);
+
+        await Promise.all(
+          existing.map((slot) =>
+            subscriptionRepo.updateById(db, slot.id, {
+              partnerLicenseKey: license_key,
+              plan: cfg.plan,
+              status: "active",
+              partnerTier: tier,
+              seats: cfg.seats,
+              unlimitedSeats: cfg.unlimitedSeats,
+            }),
+          ),
+        );
+
         break;
       }
 

apps/web/src/views/onboarding/workspace-details/index.tsx

@@ -40,7 +40,8 @@ export default function WorkspaceNameView() {
   const billing = searchParams.get("billing") ?? "annual";
   const returnUrl = searchParams.get("returnUrl") ?? "/boards";
   const licenseKeyParam = searchParams.get("license_key");
-  const isLicenseFlow = !!licenseKeyParam;
+  const isLicenseFlow =
+    !!licenseKeyParam || searchParams.get("partner") === "1";
   const { showPopup } = usePopup();
 
   useEffect(() => {
@@ -103,6 +104,7 @@ export default function WorkspaceNameView() {
       if (!workspace.publicId) return;
       localStorage.setItem("workspacePublicId", workspace.publicId);
       void utils.workspace.all.invalidate();
+      void utils.workspace.hasAvailablePartnerSlot.invalidate();
       const storedLicenseKey = localStorage.getItem("partnerLicenseKey");
       if (storedLicenseKey) {
         localStorage.removeItem("partnerLicenseKey");

packages/api/src/routers/workspace.ts

@@ -2,19 +2,21 @@ import { TRPCError } from "@trpc/server";
 import { env } from "next-runtime-env";
 import { z } from "zod";
 
+import type { WorkspacePlan } from "@kan/db/schema";
+import * as subscriptionRepo from "@kan/db/repository/subscription.repo";
 import * as workspaceRepo from "@kan/db/repository/workspace.repo";
 import * as workspaceSlugRepo from "@kan/db/repository/workspaceSlug.repo";
 import { generateAvatarUrl, generateUID } from "@kan/shared/utils";
 
-import { createTRPCRouter, protectedProcedure, publicProcedure } from "../trpc";
 import {
-  workspaceListItemSchema,
-  workspaceDetailSchema,
-  workspaceWithBoardsSchema,
   workspaceCreateResponseSchema,
-  workspaceUpdateResponseSchema,
   workspaceDeleteResponseSchema,
+  workspaceDetailSchema,
+  workspaceListItemSchema,
+  workspaceUpdateResponseSchema,
+  workspaceWithBoardsSchema,
 } from "../schemas";
+import { createTRPCRouter, protectedProcedure, publicProcedure } from "../trpc";
 import { assertPermission } from "../utils/permissions";
 
 export const workspaceRouter = createTRPCRouter({
@@ -270,12 +272,47 @@ export const workspaceRouter = createTRPCRouter({
           code: "INTERNAL_SERVER_ERROR",
         });
 
+      let unlinkedSlot: Awaited<
+        ReturnType<typeof subscriptionRepo.getFirstUnlinkedSlotByLicenseKey>
+      >;
+
+      if (env("NEXT_PUBLIC_KAN_ENV") === "cloud") {
+        const memberships = await workspaceRepo.getAllByUserId(ctx.db, userId);
+        const otherWorkspaceIds = memberships
+          .map((m) => m.workspace?.publicId)
+          .filter((id): id is string => !!id && id !== workspacePublicId);
+
+        const partnerSub = otherWorkspaceIds.length
+          ? await subscriptionRepo.getFirstActivePartnerSubByWorkspaceIds(
+              ctx.db,
+              otherWorkspaceIds,
+            )
+          : undefined;
+        unlinkedSlot = partnerSub?.partnerLicenseKey
+          ? await subscriptionRepo.getFirstUnlinkedSlotByLicenseKey(
+              ctx.db,
+              partnerSub.partnerLicenseKey,
+            )
+          : undefined;
+
+        if (unlinkedSlot) {
+          await Promise.all([
+            subscriptionRepo.updateById(ctx.db, unlinkedSlot.id, {
+              referenceId: workspacePublicId,
+            }),
+            workspaceRepo.update(ctx.db, workspacePublicId, {
+              plan: unlinkedSlot.plan as WorkspacePlan,
+            }),
+          ]);
+        }
+      }
+
       return {
         publicId: result.publicId,
         name: result.name!,
         slug: result.slug!,
         description: result.description ?? null,
-        plan: result.plan!,
+        plan: (unlinkedSlot?.plan ?? result.plan!) as WorkspacePlan,
         cardPrefix: result.cardPrefix!,
       };
     }),
@@ -412,10 +449,7 @@ export const workspaceRouter = createTRPCRouter({
         });
       await assertPermission(ctx.db, userId, workspace.id, "workspace:delete");
 
-      await workspaceRepo.hardDelete(
+      await workspaceRepo.hardDelete(ctx.db, input.workspacePublicId);
-        ctx.db,
-        input.workspacePublicId,
-      );
 
       return { success: true };
     }),
@@ -560,4 +594,38 @@ export const workspaceRouter = createTRPCRouter({
 
       return result;
     }),
+  hasAvailablePartnerSlot: protectedProcedure
+    .input(z.void())
+    .output(z.boolean())
+    .query(async ({ ctx }) => {
+      const userId = ctx.user?.id;
+
+      if (!userId)
+        throw new TRPCError({
+          message: `User not authenticated`,
+          code: "UNAUTHORIZED",
+        });
+
+      const memberships = await workspaceRepo.getAllByUserId(ctx.db, userId);
+      const workspaceIds = memberships
+        .map((m) => m.workspace?.publicId)
+        .filter((id): id is string => !!id);
+
+      if (!workspaceIds.length) return false;
+
+      const partnerSub =
+        await subscriptionRepo.getFirstActivePartnerSubByWorkspaceIds(
+          ctx.db,
+          workspaceIds,
+        );
+      if (!partnerSub?.partnerLicenseKey) return false;
+
+      const unlinkedSlot =
+        await subscriptionRepo.getFirstUnlinkedSlotByLicenseKey(
+          ctx.db,
+          partnerSub.partnerLicenseKey,
+        );
+
+      return !!unlinkedSlot;
+    }),
 });

packages/db/migrations/20260529123231_DropPartnerLicenseKeyUniqueConstraint.sql

@@ -0,0 +1 @@
+DROP INDEX IF EXISTS "subscription_partner_license_key_idx";

packages/db/migrations/meta/_journal.json

@@ -239,6 +239,13 @@
       "when": 1778617946519,
       "tag": "20260512203226_AddPartnerLicenseToSubscription",
       "breakpoints": true
+    },
+    {
+      "idx": 34,
+      "version": "7",
+      "when": 1780057951781,
+      "tag": "20260529123231_DropPartnerLicenseKeyUniqueConstraint",
+      "breakpoints": true
     }
   ]
 }

packages/db/src/repository/subscription.repo.ts

@@ -1,4 +1,4 @@
-import { eq } from "drizzle-orm";
+import { and, asc, eq, inArray, isNotNull, isNull } from "drizzle-orm";
 
 import type { dbClient } from "@kan/db/client";
 import { subscription } from "@kan/db/schema";
@@ -15,6 +15,9 @@ export const updateById = async (
     periodEnd?: Date | null;
     cancelAtPeriodEnd?: boolean | null;
     stripeSubscriptionId?: string | null;
+    referenceId?: string | null;
+    partnerLicenseKey?: string;
+    partnerTier?: number;
   },
 ) => {
   const [result] = await db
@@ -29,6 +32,7 @@ export const updateById = async (
       plan: subscription.plan,
       status: subscription.status,
       unlimitedSeats: subscription.unlimitedSeats,
+      referenceId: subscription.referenceId,
     });
 
   return result;
@@ -63,6 +67,24 @@ export const updateByStripeSubscriptionId = async (
   return result;
 };
 
+export const updateAllByPartnerLicenseKey = async (
+  db: dbClient,
+  partnerLicenseKey: string,
+  updates: {
+    plan?: string;
+    status?: string;
+    partnerTier?: number;
+    seats?: number | null;
+    unlimitedSeats?: boolean;
+  },
+) => {
+  return await db
+    .update(subscription)
+    .set({ ...updates, updatedAt: new Date() })
+    .where(eq(subscription.partnerLicenseKey, partnerLicenseKey))
+    .returning({ id: subscription.id });
+};
+
 export const getByStripeSubscriptionId = async (
   db: dbClient,
   stripeSubscriptionId: string,
@@ -96,13 +118,49 @@ export const getByPartnerLicenseKey = async (
   db: dbClient,
   partnerLicenseKey: string,
 ) => {
-  const result = await db.query.subscription.findFirst({
+  return await db.query.subscription.findFirst({
     where: eq(subscription.partnerLicenseKey, partnerLicenseKey),
   });
-  return result;
 };
 
-export const upsertByPartnerLicenseKey = async (
+export const getAllByPartnerLicenseKey = async (
+  db: dbClient,
+  partnerLicenseKey: string,
+) => {
+  return await db.query.subscription.findMany({
+    where: eq(subscription.partnerLicenseKey, partnerLicenseKey),
+    orderBy: [asc(subscription.id)],
+  });
+};
+
+export const getFirstUnlinkedSlotByLicenseKey = async (
+  db: dbClient,
+  partnerLicenseKey: string,
+) => {
+  return await db.query.subscription.findFirst({
+    where: and(
+      eq(subscription.partnerLicenseKey, partnerLicenseKey),
+      isNull(subscription.referenceId),
+      inArray(subscription.status, ["active", "trialing"]),
+    ),
+    orderBy: [asc(subscription.id)],
+  });
+};
+
+export const getFirstActivePartnerSubByWorkspaceIds = async (
+  db: dbClient,
+  workspacePublicIds: string[],
+) => {
+  return await db.query.subscription.findFirst({
+    where: and(
+      inArray(subscription.referenceId, workspacePublicIds),
+      isNotNull(subscription.partnerLicenseKey),
+      inArray(subscription.status, ["active", "trialing"]),
+    ),
+  });
+};
+
+export const createPartnerLicenseSlots = async (
   db: dbClient,
   partnerLicenseKey: string,
   data: {
@@ -111,20 +169,13 @@ export const upsertByPartnerLicenseKey = async (
     partnerTier: number;
     seats: number | null;
     unlimitedSeats: boolean;
-    referenceId?: string;
   },
+  count: number,
 ) => {
-  const [result] = await db
+  const rows = Array.from({ length: count }, () => ({
-    .insert(subscription)
+    partnerLicenseKey,
-    .values({
+    ...data,
-      partnerLicenseKey,
+    referenceId: null,
-      ...data,
+  }));
-      referenceId: data.referenceId ?? null,
+  return await db.insert(subscription).values(rows).returning();
-    })
-    .onConflictDoUpdate({
-      target: subscription.partnerLicenseKey,
-      set: { ...data, updatedAt: new Date() },
-    })
-    .returning();
-  return result;
 };

packages/db/src/schema/subscriptions.ts

@@ -5,42 +5,33 @@ import {
   integer,
   pgTable,
   timestamp,
-  uniqueIndex,
   varchar,
 } from "drizzle-orm/pg-core";
 
 import { workspaces } from "./workspaces";
 
-export const subscription = pgTable(
+export const subscription = pgTable("subscription", {
-  "subscription",
+  id: bigserial("id", { mode: "number" }).primaryKey(),
-  {
+  plan: varchar("plan", { length: 255 }).notNull(),
-    id: bigserial("id", { mode: "number" }).primaryKey(),
+  referenceId: varchar("referenceId", { length: 12 }).references(
-    plan: varchar("plan", { length: 255 }).notNull(),
+    () => workspaces.publicId,
-    referenceId: varchar("referenceId", { length: 12 }).references(
+    { onDelete: "set null" },
-      () => workspaces.publicId,
+  ),
-      { onDelete: "set null" },
+  stripeCustomerId: varchar("stripeCustomerId", { length: 255 }),
-    ),
+  stripeSubscriptionId: varchar("stripeSubscriptionId", { length: 255 }),
-    stripeCustomerId: varchar("stripeCustomerId", { length: 255 }),
+  status: varchar("status", { length: 255 }).notNull(),
-    stripeSubscriptionId: varchar("stripeSubscriptionId", { length: 255 }),
+  periodStart: timestamp("periodStart"),
-    status: varchar("status", { length: 255 }).notNull(),
+  periodEnd: timestamp("periodEnd"),
-    periodStart: timestamp("periodStart"),
+  cancelAtPeriodEnd: boolean("cancelAtPeriodEnd"),
-    periodEnd: timestamp("periodEnd"),
+  seats: integer("seats"),
-    cancelAtPeriodEnd: boolean("cancelAtPeriodEnd"),
+  unlimitedSeats: boolean("unlimitedSeats").default(false).notNull(),
-    seats: integer("seats"),
+  trialStart: timestamp("trialStart"),
-    unlimitedSeats: boolean("unlimitedSeats").default(false).notNull(),
+  trialEnd: timestamp("trialEnd"),
-    trialStart: timestamp("trialStart"),
+  partnerLicenseKey: varchar("partnerLicenseKey", { length: 255 }),
-    trialEnd: timestamp("trialEnd"),
+  partnerTier: integer("partnerTier"),
-    partnerLicenseKey: varchar("partnerLicenseKey", { length: 255 }),
+  createdAt: timestamp("createdAt").notNull().defaultNow(),
-    partnerTier: integer("partnerTier"),
+  updatedAt: timestamp("updatedAt").notNull().defaultNow(),
-    createdAt: timestamp("createdAt").notNull().defaultNow(),
+}).enableRLS();
-    updatedAt: timestamp("updatedAt").notNull().defaultNow(),
-  },
-  (table) => [
-    uniqueIndex("subscription_partner_license_key_idx").on(
-      table.partnerLicenseKey,
-    ),
-  ],
-).enableRLS();
 
 export const subscriptionsRelations = relations(subscription, ({ one }) => ({
   workspace: one(workspaces, {