chore: simplify referral codes (#28491)

packages/console/core/src/referral.ts

@@ -3,7 +3,7 @@ import { and, asc, eq, isNull, sql, Database } from "./drizzle"
 import { Actor } from "./actor"
 import { Identifier } from "./identifier"
 import { LiteTable } from "./schema/billing.sql"
-import { ReferralRewardTable, ReferralTable } from "./schema/referral.sql"
+import { ReferralCodeTable, ReferralRewardTable, ReferralTable } from "./schema/referral.sql"
 import { AuthTable } from "./schema/auth.sql"
 import { UserTable } from "./schema/user.sql"
 import { WorkspaceTable } from "./schema/workspace.sql"
@@ -26,38 +26,29 @@ export namespace Referral {
   }
 
   function generateCode() {
-    return ulid().slice(-CODE_LENGTH)
+    return ulid().slice(-CODE_LENGTH).toUpperCase()
   }
 
   async function ensureCode(workspaceID = Actor.workspace()) {
-    return Database.transaction(async (tx) => {
-      const existing = await tx
-        .select({ code: WorkspaceTable.referralCode })
-        .from(WorkspaceTable)
-        .where(and(eq(WorkspaceTable.id, workspaceID), isNull(WorkspaceTable.timeDeleted)))
+    return Database.use(async (db) => {
+      const existing = await db
+        .select({ code: ReferralCodeTable.code })
+        .from(ReferralCodeTable)
+        .where(eq(ReferralCodeTable.workspaceID, workspaceID))
         .then((rows) => rows[0])
-      if (!existing) throw new Error("Workspace not found")
-      if (existing.code) return { code: existing.code }
+      if (existing) return { code: existing.code }
 
-      for (const _ of Array.from({ length: 5 })) {
-        await tx
-          .update(WorkspaceTable)
-          .set({ referralCode: generateCode() })
-          .where(
-            and(
-              eq(WorkspaceTable.id, workspaceID),
-              isNull(WorkspaceTable.referralCode),
-              isNull(WorkspaceTable.timeDeleted),
-            ),
-          )
+      await db.insert(ReferralCodeTable).ignore().values({
+        workspaceID,
+        code: generateCode(),
+      })
 
-        const created = await tx
-          .select({ code: WorkspaceTable.referralCode })
-          .from(WorkspaceTable)
-          .where(and(eq(WorkspaceTable.id, workspaceID), isNull(WorkspaceTable.timeDeleted)))
-          .then((rows) => rows[0])
-        if (created?.code) return { code: created.code }
-      }
+      const created = await db
+        .select({ code: ReferralCodeTable.code })
+        .from(ReferralCodeTable)
+        .where(eq(ReferralCodeTable.workspaceID, workspaceID))
+        .then((rows) => rows[0])
+      if (created) return { code: created.code }
 
       throw new Error("Failed to generate referral code")
     })
@@ -300,9 +291,15 @@ export namespace Referral {
 
     return Database.transaction(async (tx) => {
       const code = await tx
-        .select({ workspaceID: WorkspaceTable.id })
-        .from(WorkspaceTable)
-        .where(and(eq(WorkspaceTable.referralCode, referralCode), isNull(WorkspaceTable.timeDeleted)))
+        .select({ workspaceID: ReferralCodeTable.workspaceID })
+        .from(ReferralCodeTable)
+        .innerJoin(WorkspaceTable, eq(WorkspaceTable.id, ReferralCodeTable.workspaceID))
+        .where(
+          and(
+            eq(ReferralCodeTable.code, referralCode),
+            isNull(WorkspaceTable.timeDeleted),
+          ),
+        )
         .then((rows) => rows[0])
       if (!code) throw new Error("Referral code invalid")
 

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

@@ -1,7 +1,17 @@
-import { bigint, mysqlTable, primaryKey, uniqueIndex } from "drizzle-orm/mysql-core"
+import { bigint, mysqlTable, primaryKey, uniqueIndex, varchar } from "drizzle-orm/mysql-core"
 import { timestamps, ulid, utc, workspaceColumns } from "../drizzle/types"
 import { workspaceIndexes } from "./workspace.sql"
 
+export const ReferralCodeTable = mysqlTable(
+  "referral_code",
+  {
+    workspaceID: ulid("workspace_id").notNull(),
+    code: varchar("code", { length: 10 }).notNull(),
+    ...timestamps,
+  },
+  (table) => [primaryKey({ columns: [table.workspaceID] }), uniqueIndex("code").on(table.code)],
+)
+
 export const ReferralTable = mysqlTable(
   "referral",
   {
@@ -9,7 +19,7 @@ export const ReferralTable = mysqlTable(
     ...timestamps,
     inviteeAccountID: ulid("invitee_account_id").notNull(),
   },
-  (table) => [...workspaceIndexes(table), uniqueIndex("referral_invitee_account_id").on(table.inviteeAccountID)],
+  (table) => [...workspaceIndexes(table), uniqueIndex("invitee_account_id").on(table.inviteeAccountID)],
 )
 
 export const ReferralRewardTable = mysqlTable(

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

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