clean up referral tables

packages/console/app/src/routes/stripe/webhook.ts

@@ -177,8 +177,6 @@ export async function POST(input: APIEvent) {
           await Referral.createFromLiteSubscription({
             workspaceID,
             userID,
-            customerID,
-            subscriptionID,
             inviteCode,
           }).catch((error) => {
             console.error("Referral sync failed", error)

packages/console/core/migrations/20260516082200_long_spirit/migration.sql

@@ -0,0 +1,20 @@
+DROP TABLE `referral_code`;--> statement-breakpoint
+DROP INDEX `referral_reward_referral_source` ON `referral_reward`;--> statement-breakpoint
+DROP INDEX `referral_stripe_subscription_id` ON `referral`;--> statement-breakpoint
+DROP INDEX `referral_inviter_workspace_id` ON `referral`;--> statement-breakpoint
+DROP INDEX `referral_code_id` ON `referral`;--> statement-breakpoint
+ALTER TABLE `referral_reward` DROP PRIMARY KEY;--> statement-breakpoint
+ALTER TABLE `referral` DROP PRIMARY KEY;--> statement-breakpoint
+ALTER TABLE `referral_reward` MODIFY COLUMN `workspace_id` varchar(30);--> statement-breakpoint
+ALTER TABLE `workspace` ADD `referral_code` varchar(16);--> statement-breakpoint
+ALTER TABLE `referral_reward` ADD PRIMARY KEY (`id`);--> statement-breakpoint
+ALTER TABLE `referral` ADD PRIMARY KEY (`id`);--> statement-breakpoint
+CREATE INDEX `referral_workspace_id` ON `referral` (`workspace_id`);--> statement-breakpoint
+CREATE UNIQUE INDEX `workspace_referral_code` ON `workspace` (`referral_code`);--> statement-breakpoint
+ALTER TABLE `referral_reward` DROP COLUMN `source`;--> statement-breakpoint
+ALTER TABLE `referral_reward` DROP COLUMN `applied_by_user_id`;--> statement-breakpoint
+ALTER TABLE `referral` DROP COLUMN `inviter_workspace_id`;--> statement-breakpoint
+ALTER TABLE `referral` DROP COLUMN `invitee_user_id`;--> statement-breakpoint
+ALTER TABLE `referral` DROP COLUMN `referral_code_id`;--> statement-breakpoint
+ALTER TABLE `referral` DROP COLUMN `stripe_customer_id`;--> statement-breakpoint
+ALTER TABLE `referral` DROP COLUMN `stripe_subscription_id`;

packages/console/core/src/identifier.ts

@@ -13,7 +13,6 @@ export namespace Identifier {
     payment: "pay",
     provider: "prv",
     referral: "ref",
-    referralCode: "rfc",
     referralReward: "rrw",
     subscription: "sub",
     usage: "usg",

packages/console/core/src/referral.ts

@@ -4,8 +4,9 @@ import { and, desc, eq, isNull, sql, Database } from "./drizzle"
 import { Actor } from "./actor"
 import { Identifier } from "./identifier"
 import { LiteTable } from "./schema/billing.sql"
-import { ReferralCodeTable, ReferralRewardTable, ReferralTable } from "./schema/referral.sql"
+import { ReferralRewardTable, ReferralTable } from "./schema/referral.sql"
 import { UserTable } from "./schema/user.sql"
+import { WorkspaceTable } from "./schema/workspace.sql"
 import { LiteData } from "./lite"
 import { centsToMicroCents, microCentsToCents } from "./util/price"
 import { getMonthlyBounds, getWeekBounds } from "./util/date"
@@ -13,7 +14,7 @@ import { fn } from "./util/fn"
 
 export namespace Referral {
   export const REWARD_AMOUNT = centsToMicroCents(500)
-  const CODE_LENGTH = 10
+  const CODE_LENGTH = 16
 
   function normalizeCode(code?: string) {
     return code?.toUpperCase().replace(/[^A-Z0-9]/g, "").slice(0, CODE_LENGTH)
@@ -26,32 +27,27 @@ export namespace Referral {
   export async function ensureCode(workspaceID = Actor.workspace()) {
     return Database.transaction(async (tx) => {
       const existing = await tx
-        .select({ id: ReferralCodeTable.id, code: ReferralCodeTable.code })
-        .from(ReferralCodeTable)
-        .where(and(eq(ReferralCodeTable.workspaceID, workspaceID), isNull(ReferralCodeTable.timeDeleted)))
+        .select({ code: WorkspaceTable.referralCode })
+        .from(WorkspaceTable)
+        .where(and(eq(WorkspaceTable.id, workspaceID), isNull(WorkspaceTable.timeDeleted)))
         .then((rows) => rows[0])
-      if (existing) return existing
+      if (!existing) throw new Error("Workspace not found")
+      if (existing.code) return { code: existing.code }
 
       for (const _ of Array.from({ length: 5 })) {
         await tx
-          .insert(ReferralCodeTable)
-          .values({
-            workspaceID,
-            id: Identifier.create("referralCode"),
-            code: generateCode(),
-          })
-          .onDuplicateKeyUpdate({
-            set: {
-              code: sql`${ReferralCodeTable.code}`,
-            },
-          })
+          .update(WorkspaceTable)
+          .set({ referralCode: generateCode() })
+          .where(
+            and(eq(WorkspaceTable.id, workspaceID), isNull(WorkspaceTable.referralCode), isNull(WorkspaceTable.timeDeleted)),
+          )
 
         const created = await tx
-          .select({ id: ReferralCodeTable.id, code: ReferralCodeTable.code })
-          .from(ReferralCodeTable)
-          .where(and(eq(ReferralCodeTable.workspaceID, workspaceID), isNull(ReferralCodeTable.timeDeleted)))
+          .select({ code: WorkspaceTable.referralCode })
+          .from(WorkspaceTable)
+          .where(and(eq(WorkspaceTable.id, workspaceID), isNull(WorkspaceTable.timeDeleted)))
           .then((rows) => rows[0])
-        if (created) return created
+        if (created?.code) return { code: created.code }
       }
 
       throw new Error("Failed to generate referral code")
@@ -65,19 +61,27 @@ export namespace Referral {
       const rewards = await tx
         .select({
           id: ReferralRewardTable.id,
-          source: ReferralRewardTable.source,
+          workspaceID: ReferralRewardTable.workspaceID,
+          referralWorkspaceID: ReferralTable.workspaceID,
           amount: ReferralRewardTable.amount,
           timeCreated: ReferralRewardTable.timeCreated,
           timeApplied: ReferralRewardTable.timeApplied,
         })
         .from(ReferralRewardTable)
-        .where(and(eq(ReferralRewardTable.workspaceID, workspaceID), isNull(ReferralRewardTable.timeDeleted)))
+        .innerJoin(ReferralTable, eq(ReferralTable.id, ReferralRewardTable.referralID))
+        .where(
+          and(
+            eq(ReferralRewardTable.workspaceID, workspaceID),
+            isNull(ReferralRewardTable.timeDeleted),
+            isNull(ReferralTable.timeDeleted),
+          ),
+        )
         .orderBy(desc(ReferralRewardTable.timeCreated))
 
       const invites = await tx
         .select({ id: ReferralTable.id })
         .from(ReferralTable)
-        .where(and(eq(ReferralTable.inviterWorkspaceID, workspaceID), isNull(ReferralTable.timeDeleted)))
+        .where(and(eq(ReferralTable.workspaceID, workspaceID), isNull(ReferralTable.timeDeleted)))
 
       const lite = await tx
         .select({ id: LiteTable.id })
@@ -89,8 +93,11 @@ export namespace Referral {
     })
 
     const rewards = rows.rewards.map((reward) => ({
-      ...reward,
+      id: reward.id,
+      source: reward.workspaceID === reward.referralWorkspaceID ? ("inviter" as const) : ("invitee" as const),
       amount: microCentsToCents(reward.amount),
+      timeCreated: reward.timeCreated,
+      timeApplied: reward.timeApplied,
     }))
 
     return {
@@ -108,7 +115,6 @@ export namespace Referral {
 
   export const applyReward = fn(z.object({ rewardID: z.string() }), async (input) => {
     const workspaceID = Actor.workspace()
-    const userID = Actor.userID()
 
     return Database.transaction(async (tx) => {
       const reward = await tx
@@ -135,7 +141,6 @@ export namespace Referral {
       const update = await tx
         .update(ReferralRewardTable)
         .set({
-          appliedByUserID: userID,
           timeApplied: sql`now()`,
         })
         .where(
@@ -182,8 +187,6 @@ export namespace Referral {
   export async function createFromLiteSubscription(input: {
     workspaceID: string
     userID: string
-    customerID: string
-    subscriptionID: string
     inviteCode?: string
   }) {
     const inviteCode = normalizeCode(input.inviteCode)
@@ -191,9 +194,9 @@ export namespace Referral {
 
     return Database.transaction(async (tx) => {
       const code = await tx
-        .select({ id: ReferralCodeTable.id, workspaceID: ReferralCodeTable.workspaceID })
-        .from(ReferralCodeTable)
-        .where(and(eq(ReferralCodeTable.code, inviteCode), isNull(ReferralCodeTable.timeDeleted)))
+        .select({ workspaceID: WorkspaceTable.id })
+        .from(WorkspaceTable)
+        .where(and(eq(WorkspaceTable.referralCode, inviteCode), isNull(WorkspaceTable.timeDeleted)))
         .then((rows) => rows[0])
       if (!code) return { status: "invalid-code" as const }
 
@@ -207,94 +210,70 @@ export namespace Referral {
       if (!invitee?.accountID) return { status: "missing-account" as const }
 
       const existingReferral = await tx
-        .select({
-          id: ReferralTable.id,
-          workspaceID: ReferralTable.workspaceID,
-          inviterWorkspaceID: ReferralTable.inviterWorkspaceID,
-          stripeSubscriptionID: ReferralTable.stripeSubscriptionID,
-        })
+        .select({ id: ReferralTable.id })
         .from(ReferralTable)
         .where(and(eq(ReferralTable.inviteeAccountID, invitee.accountID), isNull(ReferralTable.timeDeleted)))
         .then((rows) => rows[0])
-      if (existingReferral && existingReferral.stripeSubscriptionID !== input.subscriptionID)
-        return { status: "already-redeemed" as const }
+      if (existingReferral) return { status: "already-redeemed" as const }
 
-      if (!existingReferral) {
-        const selfReferral = await tx
-          .select({ id: UserTable.id })
-          .from(UserTable)
-          .where(
-            and(
-              eq(UserTable.workspaceID, code.workspaceID),
-              eq(UserTable.accountID, invitee.accountID),
-              isNull(UserTable.timeDeleted),
-            ),
-          )
-          .then((rows) => rows[0])
-        if (selfReferral) return { status: "self-referral" as const }
+      const selfReferral = await tx
+        .select({ id: UserTable.id })
+        .from(UserTable)
+        .where(
+          and(
+            eq(UserTable.workspaceID, code.workspaceID),
+            eq(UserTable.accountID, invitee.accountID),
+            isNull(UserTable.timeDeleted),
+          ),
+        )
+        .then((rows) => rows[0])
+      if (selfReferral) return { status: "self-referral" as const }
 
-        const existingGo = await tx
-          .select({ workspaceID: LiteTable.workspaceID })
-          .from(LiteTable)
-          .innerJoin(UserTable, and(eq(UserTable.workspaceID, LiteTable.workspaceID), eq(UserTable.id, LiteTable.userID)))
-          .where(and(eq(UserTable.accountID, invitee.accountID), isNull(UserTable.timeDeleted), isNull(LiteTable.timeDeleted)))
-        if (existingGo.some((row) => row.workspaceID !== input.workspaceID)) return { status: "already-subscribed" as const }
+      const existingGo = await tx
+        .select({ workspaceID: LiteTable.workspaceID })
+        .from(LiteTable)
+        .innerJoin(UserTable, and(eq(UserTable.workspaceID, LiteTable.workspaceID), eq(UserTable.id, LiteTable.userID)))
+        .where(and(eq(UserTable.accountID, invitee.accountID), isNull(UserTable.timeDeleted), isNull(LiteTable.timeDeleted)))
+      if (existingGo.some((row) => row.workspaceID !== input.workspaceID)) return { status: "already-subscribed" as const }
 
-        await tx
-          .insert(ReferralTable)
-          .values({
-            workspaceID: input.workspaceID,
-            id: Identifier.create("referral"),
-            inviterWorkspaceID: code.workspaceID,
-            inviteeAccountID: invitee.accountID,
-            inviteeUserID: input.userID,
-            referralCodeID: code.id,
-            stripeCustomerID: input.customerID,
-            stripeSubscriptionID: input.subscriptionID,
-          })
-          .onDuplicateKeyUpdate({
-            set: {
-              stripeSubscriptionID: sql`${ReferralTable.stripeSubscriptionID}`,
-            },
-          })
-      }
+      const referralID = Identifier.create("referral")
+      await tx
+        .insert(ReferralTable)
+        .values({
+          workspaceID: code.workspaceID,
+          id: referralID,
+          inviteeAccountID: invitee.accountID,
+        })
+        .onDuplicateKeyUpdate({
+          set: {
+            inviteeAccountID: sql`${ReferralTable.inviteeAccountID}`,
+          },
+        })
 
-      const referral =
-        existingReferral ??
-        (await tx
-          .select({
-            id: ReferralTable.id,
-            workspaceID: ReferralTable.workspaceID,
-            inviterWorkspaceID: ReferralTable.inviterWorkspaceID,
-          })
-          .from(ReferralTable)
-          .where(and(eq(ReferralTable.stripeSubscriptionID, input.subscriptionID), isNull(ReferralTable.timeDeleted)))
-          .then((rows) => rows[0]))
+      const referral = await tx
+        .select({ id: ReferralTable.id, workspaceID: ReferralTable.workspaceID })
+        .from(ReferralTable)
+        .where(and(eq(ReferralTable.inviteeAccountID, invitee.accountID), isNull(ReferralTable.timeDeleted)))
+        .then((rows) => rows[0])
       if (!referral) return { status: "duplicate" as const }
+      if (referral.id !== referralID) return { status: "already-redeemed" as const }
 
       await tx
         .insert(ReferralRewardTable)
         .values([
-          {
-            workspaceID: referral.inviterWorkspaceID,
-            id: Identifier.create("referralReward"),
-            referralID: referral.id,
-            source: "inviter",
-            amount: REWARD_AMOUNT,
-          },
           {
             workspaceID: referral.workspaceID,
             id: Identifier.create("referralReward"),
             referralID: referral.id,
-            source: "invitee",
+            amount: REWARD_AMOUNT,
+          },
+          {
+            workspaceID: input.workspaceID,
+            id: Identifier.create("referralReward"),
+            referralID: referral.id,
             amount: REWARD_AMOUNT,
           },
         ])
-        .onDuplicateKeyUpdate({
-          set: {
-            amount: sql`${ReferralRewardTable.amount}`,
-          },
-        })
 
       return { status: "created" as const }
     })

packages/console/core/src/schema/referral.sql.ts

@@ -1,58 +1,29 @@
-import { bigint, index, mysqlEnum, mysqlTable, uniqueIndex, varchar } from "drizzle-orm/mysql-core"
-import { timestamps, ulid, utc, workspaceColumns } from "../drizzle/types"
-import { workspaceIndexes } from "./workspace.sql"
-
-export const ReferralRewardSource = ["inviter", "invitee"] as const
-
-export const ReferralCodeTable = mysqlTable(
-  "referral_code",
-  {
-    ...workspaceColumns,
-    ...timestamps,
-    code: varchar("code", { length: 10 }).notNull(),
-  },
-  (table) => [
-    ...workspaceIndexes(table),
-    uniqueIndex("referral_code_workspace_id").on(table.workspaceID),
-    uniqueIndex("referral_code_code").on(table.code),
-  ],
-)
+import { bigint, index, mysqlTable, uniqueIndex } from "drizzle-orm/mysql-core"
+import { timestamps, ulid, utc } from "../drizzle/types"
 
 export const ReferralTable = mysqlTable(
   "referral",
   {
-    ...workspaceColumns,
+    id: ulid("id").notNull().primaryKey(),
+    workspaceID: ulid("workspace_id").notNull(),
     ...timestamps,
-    inviterWorkspaceID: ulid("inviter_workspace_id").notNull(),
     inviteeAccountID: ulid("invitee_account_id").notNull(),
-    inviteeUserID: ulid("invitee_user_id").notNull(),
-    referralCodeID: ulid("referral_code_id").notNull(),
-    stripeCustomerID: varchar("stripe_customer_id", { length: 255 }).notNull(),
-    stripeSubscriptionID: varchar("stripe_subscription_id", { length: 255 }).notNull(),
   },
   (table) => [
-    ...workspaceIndexes(table),
     uniqueIndex("referral_invitee_account_id").on(table.inviteeAccountID),
-    uniqueIndex("referral_stripe_subscription_id").on(table.stripeSubscriptionID),
-    index("referral_inviter_workspace_id").on(table.inviterWorkspaceID),
-    index("referral_code_id").on(table.referralCodeID),
+    index("referral_workspace_id").on(table.workspaceID),
   ],
 )
 
 export const ReferralRewardTable = mysqlTable(
   "referral_reward",
   {
-    ...workspaceColumns,
+    id: ulid("id").notNull().primaryKey(),
+    workspaceID: ulid("workspace_id"),
     ...timestamps,
     referralID: ulid("referral_id").notNull(),
-    source: mysqlEnum("source", ReferralRewardSource).notNull(),
     amount: bigint("amount", { mode: "number" }).notNull(),
-    appliedByUserID: ulid("applied_by_user_id"),
     timeApplied: utc("time_applied"),
   },
-  (table) => [
-    ...workspaceIndexes(table),
-    uniqueIndex("referral_reward_referral_source").on(table.referralID, table.source),
-    index("referral_reward_workspace_time").on(table.workspaceID, table.timeCreated),
-  ],
+  (table) => [index("referral_reward_workspace_time").on(table.workspaceID, table.timeCreated)],
 )

packages/console/core/src/schema/workspace.sql.ts

@@ -6,10 +6,14 @@ export const WorkspaceTable = mysqlTable(
   {
     id: ulid("id").notNull().primaryKey(),
     slug: varchar("slug", { length: 255 }),
+    referralCode: varchar("referral_code", { length: 16 }),
     name: varchar("name", { length: 255 }).notNull(),
     ...timestamps,
   },
-  (table) => [uniqueIndex("slug").on(table.slug)],
+  (table) => [
+    uniqueIndex("slug").on(table.slug),
+    uniqueIndex("workspace_referral_code").on(table.referralCode),
+  ],
 )
 
 export function workspaceIndexes(table: any) {