wip: black

packages/console/core/migrations/0042_flat_nightmare.sql

@@ -0,0 +1,7 @@
+ALTER TABLE `billing` ADD `subscription_id` varchar(28);--> statement-breakpoint
+ALTER TABLE `usage` ADD `data` json;--> statement-breakpoint
+ALTER TABLE `user` ADD `time_subscribed` timestamp(3);--> statement-breakpoint
+ALTER TABLE `user` ADD `sub_recent_usage` bigint;--> statement-breakpoint
+ALTER TABLE `user` ADD `sub_monthly_usage` bigint;--> statement-breakpoint
+ALTER TABLE `user` ADD `sub_time_recent_usage_updated` timestamp(3);--> statement-breakpoint
+ALTER TABLE `user` ADD `sub_time_monthly_usage_updated` timestamp(3);

packages/console/core/migrations/meta/_journal.json

@@ -295,6 +295,13 @@
       "when": 1767732559197,
       "tag": "0041_odd_misty_knight",
       "breakpoints": true
+    },
+    {
+      "idx": 42,
+      "version": "5",
+      "when": 1767744077346,
+      "tag": "0042_flat_nightmare",
+      "breakpoints": true
     }
   ]
 }

packages/console/core/script/onboard-zen-black.ts

@@ -0,0 +1,152 @@
+import { Billing } from "../src/billing.js"
+import { Database, eq, and, sql } from "../src/drizzle/index.js"
+import { AuthTable } from "../src/schema/auth.sql.js"
+import { UserTable } from "../src/schema/user.sql.js"
+import { BillingTable, PaymentTable } from "../src/schema/billing.sql.js"
+import { Identifier } from "../src/identifier.js"
+import { centsToMicroCents } from "../src/util/price.js"
+
+const workspaceID = process.argv[2]
+const email = process.argv[3]
+
+if (!workspaceID || !email) {
+  console.error("Usage: bun onboard-zen-black.ts <workspaceID> <email>")
+  process.exit(1)
+}
+
+// Look up the Stripe customer by email
+const customers = await Billing.stripe().customers.list({ email, limit: 1 })
+const customer = customers.data[0]
+if (!customer) {
+  console.error(`Error: No Stripe customer found for email ${email}`)
+  process.exit(1)
+}
+const customerID = customer.id
+
+// Get the subscription id
+const subscriptions = await Billing.stripe().subscriptions.list({ customer: customerID, limit: 1 })
+const subscription = subscriptions.data[0]
+if (!subscription) {
+  console.error(`Error: Customer ${customerID} does not have a subscription`)
+  process.exit(1)
+}
+const subscriptionID = subscription.id
+
+// Validate the subscription is $200
+const amountInCents = subscription.items.data[0]?.price.unit_amount ?? 0
+if (amountInCents !== 20000) {
+  console.error(`Error: Subscription amount is $${amountInCents / 100}, expected $200`)
+  process.exit(1)
+}
+
+// Check if subscription is already tied to another workspace
+const existingSubscription = await Database.use((tx) =>
+  tx
+    .select({ workspaceID: BillingTable.workspaceID })
+    .from(BillingTable)
+    .where(eq(BillingTable.subscriptionID, subscriptionID))
+    .then((rows) => rows[0]),
+)
+if (existingSubscription) {
+  console.error(
+    `Error: Subscription ${subscriptionID} is already tied to workspace ${existingSubscription.workspaceID}`,
+  )
+  process.exit(1)
+}
+
+// Look up the workspace billing and check if it already has a customer id or subscription
+const billing = await Database.use((tx) =>
+  tx
+    .select({ customerID: BillingTable.customerID, subscriptionID: BillingTable.subscriptionID })
+    .from(BillingTable)
+    .where(eq(BillingTable.workspaceID, workspaceID))
+    .then((rows) => rows[0]),
+)
+if (billing?.subscriptionID) {
+  console.error(`Error: Workspace ${workspaceID} already has a subscription: ${billing.subscriptionID}`)
+  process.exit(1)
+}
+if (billing?.customerID) {
+  console.warn(
+    `Warning: Workspace ${workspaceID} already has a customer id: ${billing.customerID}, replacing with ${customerID}`,
+  )
+}
+
+// Get the latest invoice and payment from the subscription
+const invoices = await Billing.stripe().invoices.list({
+  subscription: subscriptionID,
+  limit: 1,
+  expand: ["data.payments"],
+})
+const invoice = invoices.data[0]
+const invoiceID = invoice?.id
+const paymentID = invoice?.payments?.data[0]?.payment.payment_intent as string | undefined
+
+// Look up the user by email via AuthTable
+const auth = await Database.use((tx) =>
+  tx
+    .select({ accountID: AuthTable.accountID })
+    .from(AuthTable)
+    .where(and(eq(AuthTable.provider, "email"), eq(AuthTable.subject, email)))
+    .then((rows) => rows[0]),
+)
+if (!auth) {
+  console.error(`Error: No user found with email ${email}`)
+  process.exit(1)
+}
+
+// Look up the user in the workspace
+const user = await Database.use((tx) =>
+  tx
+    .select({ id: UserTable.id })
+    .from(UserTable)
+    .where(and(eq(UserTable.workspaceID, workspaceID), eq(UserTable.accountID, auth.accountID)))
+    .then((rows) => rows[0]),
+)
+if (!user) {
+  console.error(`Error: User with email ${email} is not a member of workspace ${workspaceID}`)
+  process.exit(1)
+}
+
+// Set workspaceID in Stripe customer metadata
+await Billing.stripe().customers.update(customerID, {
+  metadata: {
+    workspaceID,
+  },
+})
+
+await Database.transaction(async (tx) => {
+  // Set customer id and subscription id on workspace billing
+  await tx
+    .update(BillingTable)
+    .set({
+      customerID,
+      subscriptionID,
+    })
+    .where(eq(BillingTable.workspaceID, workspaceID))
+
+  // Set current time as timeSubscribed on user
+  await tx
+    .update(UserTable)
+    .set({
+      timeSubscribed: sql`now()`,
+    })
+    .where(eq(UserTable.id, user.id))
+
+  // Create a row in payments table
+  await tx.insert(PaymentTable).values({
+    workspaceID,
+    id: Identifier.create("payment"),
+    amount: centsToMicroCents(amountInCents),
+    customerID,
+    invoiceID,
+    paymentID,
+  })
+})
+
+console.log(`Successfully onboarded workspace ${workspaceID}`)
+console.log(`  Customer ID: ${customerID}`)
+console.log(`  Subscription ID: ${subscriptionID}`)
+console.log(`  User ID: ${user.id}`)
+console.log(`  Invoice ID: ${invoiceID ?? "(none)"}`)
+console.log(`  Payment ID: ${paymentID ?? "(none)"}`)

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

@@ -1,4 +1,4 @@
-import { bigint, boolean, index, int, mysqlTable, uniqueIndex, varchar } from "drizzle-orm/mysql-core"
+import { bigint, boolean, index, int, json, mysqlTable, uniqueIndex, varchar } from "drizzle-orm/mysql-core"
 import { timestamps, ulid, utc, workspaceColumns } from "../drizzle/types"
 import { workspaceIndexes } from "./workspace.sql"
 
@@ -21,6 +21,7 @@ export const BillingTable = mysqlTable(
     reloadError: varchar("reload_error", { length: 255 }),
     timeReloadError: utc("time_reload_error"),
     timeReloadLockedTill: utc("time_reload_locked_till"),
+    subscriptionID: varchar("subscription_id", { length: 28 }),
   },
   (table) => [...workspaceIndexes(table), uniqueIndex("global_customer_id").on(table.customerID)],
 )
@@ -54,6 +55,9 @@ export const UsageTable = mysqlTable(
     cacheWrite1hTokens: int("cache_write_1h_tokens"),
     cost: bigint("cost", { mode: "number" }).notNull(),
     keyID: ulid("key_id"),
+    enrichment: json("data").$type<{
+      plan: "sub"
+    }>(),
   },
   (table) => [...workspaceIndexes(table), index("usage_time_created").on(table.workspaceID, table.timeCreated)],
 )

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

@@ -18,6 +18,12 @@ export const UserTable = mysqlTable(
     monthlyLimit: int("monthly_limit"),
     monthlyUsage: bigint("monthly_usage", { mode: "number" }),
     timeMonthlyUsageUpdated: utc("time_monthly_usage_updated"),
+    // subscription
+    timeSubscribed: utc("time_subscribed"),
+    subRecentUsage: bigint("sub_recent_usage", { mode: "number" }),
+    subMonthlyUsage: bigint("sub_monthly_usage", { mode: "number" }),
+    timeSubRecentUsageUpdated: utc("sub_time_recent_usage_updated"),
+    timeSubMonthlyUsageUpdated: utc("sub_time_monthly_usage_updated"),
   },
   (table) => [
     ...workspaceIndexes(table),