cloud: enable workspace slots for partners

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(
-        ctx.db,
-        input.workspacePublicId,
-      );
+      await workspaceRepo.hardDelete(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
-    .insert(subscription)
-    .values({
-      partnerLicenseKey,
-      ...data,
-      referenceId: data.referenceId ?? null,
-    })
-    .onConflictDoUpdate({
-      target: subscription.partnerLicenseKey,
-      set: { ...data, updatedAt: new Date() },
-    })
-    .returning();
-  return result;
+  const rows = Array.from({ length: count }, () => ({
+    partnerLicenseKey,
+    ...data,
+    referenceId: null,
+  }));
+  return await db.insert(subscription).values(rows).returning();
 };

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(
-  "subscription",
-  {
-    id: bigserial("id", { mode: "number" }).primaryKey(),
-    plan: varchar("plan", { length: 255 }).notNull(),
-    referenceId: varchar("referenceId", { length: 12 }).references(
-      () => workspaces.publicId,
-      { onDelete: "set null" },
-    ),
-    stripeCustomerId: varchar("stripeCustomerId", { length: 255 }),
-    stripeSubscriptionId: varchar("stripeSubscriptionId", { length: 255 }),
-    status: varchar("status", { length: 255 }).notNull(),
-    periodStart: timestamp("periodStart"),
-    periodEnd: timestamp("periodEnd"),
-    cancelAtPeriodEnd: boolean("cancelAtPeriodEnd"),
-    seats: integer("seats"),
-    unlimitedSeats: boolean("unlimitedSeats").default(false).notNull(),
-    trialStart: timestamp("trialStart"),
-    trialEnd: timestamp("trialEnd"),
-    partnerLicenseKey: varchar("partnerLicenseKey", { length: 255 }),
-    partnerTier: integer("partnerTier"),
-    createdAt: timestamp("createdAt").notNull().defaultNow(),
-    updatedAt: timestamp("updatedAt").notNull().defaultNow(),
-  },
-  (table) => [
-    uniqueIndex("subscription_partner_license_key_idx").on(
-      table.partnerLicenseKey,
-    ),
-  ],
-).enableRLS();
+export const subscription = pgTable("subscription", {
+  id: bigserial("id", { mode: "number" }).primaryKey(),
+  plan: varchar("plan", { length: 255 }).notNull(),
+  referenceId: varchar("referenceId", { length: 12 }).references(
+    () => workspaces.publicId,
+    { onDelete: "set null" },
+  ),
+  stripeCustomerId: varchar("stripeCustomerId", { length: 255 }),
+  stripeSubscriptionId: varchar("stripeSubscriptionId", { length: 255 }),
+  status: varchar("status", { length: 255 }).notNull(),
+  periodStart: timestamp("periodStart"),
+  periodEnd: timestamp("periodEnd"),
+  cancelAtPeriodEnd: boolean("cancelAtPeriodEnd"),
+  seats: integer("seats"),
+  unlimitedSeats: boolean("unlimitedSeats").default(false).notNull(),
+  trialStart: timestamp("trialStart"),
+  trialEnd: timestamp("trialEnd"),
+  partnerLicenseKey: varchar("partnerLicenseKey", { length: 255 }),
+  partnerTier: integer("partnerTier"),
+  createdAt: timestamp("createdAt").notNull().defaultNow(),
+  updatedAt: timestamp("updatedAt").notNull().defaultNow(),
+}).enableRLS();
 
 export const subscriptionsRelations = relations(subscription, ({ one }) => ({
   workspace: one(workspaces, {