fix(stats): run production migration safely

packages/stats/core/package.json

@@ -17,6 +17,7 @@
   },
   "scripts": {
     "db:generate": "drizzle-kit generate --config=drizzle.config.ts",
+    "db:check-unique-users": "bun src/ensure-unique-users.ts --check",
     "db:ensure-unique-users": "bun src/ensure-unique-users.ts",
     "db:migrate": "bun src/migrate.ts",
     "db:push": "drizzle-kit push --config=drizzle.config.ts",

packages/stats/core/src/ensure-unique-users.ts

@@ -2,22 +2,63 @@ import { Client } from "@planetscale/database"
 import { Resource } from "sst/resource"
 
 const tables = ["geo_stat", "model_stat", "provider_stat"] as const
+const checkOnly = process.argv.includes("--check")
 
-const client = new Client({ url: Resource.StatsDatabase.url })
+const client = new Client({ url: databaseUrl() })
 
-await tables.reduce((promise, table) => promise.then(() => ensureUniqueUsersColumn(table)), Promise.resolve())
+const missing = await tables.reduce<Promise<(typeof tables)[number][]>>(
+  async (promise, table) => {
+    const result = await promise
+    if (await hasUniqueUsersColumn(table)) {
+      console.log(`unique_users column already exists on ${table}`)
+      return result
+    }
+    return [...result, table]
+  },
+  Promise.resolve([]),
+)
 
-async function ensureUniqueUsersColumn(table: (typeof tables)[number]) {
+if (missing.length === 0) {
+  console.log("unique_users columns complete")
+  process.exit(0)
+}
+
+if (checkOnly) {
+  console.log(`unique_users columns missing on ${missing.join(", ")}`)
+  process.exit(1)
+}
+
+await missing.reduce(
+  (promise, table) => promise.then(() => addUniqueUsersColumn(table)),
+  Promise.resolve(),
+)
+
+function databaseUrl() {
+  if (
+    process.env.PLANETSCALE_HOST &&
+    process.env.PLANETSCALE_USERNAME &&
+    process.env.PLANETSCALE_PASSWORD &&
+    process.env.PLANETSCALE_DATABASE
+  )
+    return `mysql://${encodeURIComponent(process.env.PLANETSCALE_USERNAME)}:${encodeURIComponent(
+      process.env.PLANETSCALE_PASSWORD,
+    )}@${process.env.PLANETSCALE_HOST}/${process.env.PLANETSCALE_DATABASE}?ssl=${encodeURIComponent(
+      JSON.stringify({ rejectUnauthorized: true }),
+    )}`
+
+  return process.env.DATABASE_URL ?? Resource.StatsDatabase.url
+}
+
+async function hasUniqueUsersColumn(table: (typeof tables)[number]) {
   const result = await client.execute<{ column_name: string }>(
     "SELECT column_name FROM information_schema.columns WHERE table_schema = database() AND table_name = ? AND column_name = 'unique_users'",
     [table],
   )
 
-  if (result.rows.length > 0) {
-    console.log(`unique_users column already exists on ${table}`)
-    return
-  }
+  return result.rows.length > 0
+}
 
+async function addUniqueUsersColumn(table: (typeof tables)[number]) {
   await client.execute(`ALTER TABLE \`${table}\` ADD \`unique_users\` bigint NOT NULL DEFAULT 0`)
   console.log(`added unique_users column to ${table}`)
 }