fix(stats): run production migration safely

.github/workflows/deploy.yml

@@ -47,7 +47,8 @@ jobs:
           VITE_SENTRY_DSN: ${{ vars.WEB_SENTRY_DSN }}
           VITE_SENTRY_RELEASE: web@${{ github.sha }}
 
-      - run: bun sst shell --stage=${{ github.ref_name }} -- bun run --cwd packages/stats/core db:ensure-unique-users
+      - if: github.ref_name != 'production'
+        run: bun sst shell --stage=${{ github.ref_name }} -- bun run --cwd packages/stats/core db:ensure-unique-users
         env:
           CLOUDFLARE_API_TOKEN: ${{ secrets.CLOUDFLARE_API_TOKEN }}
           PLANETSCALE_SERVICE_TOKEN_NAME: ${{ secrets.PLANETSCALE_SERVICE_TOKEN_NAME }}
@@ -60,3 +61,67 @@ jobs:
           SENTRY_RELEASE: web@${{ github.sha }}
           VITE_SENTRY_DSN: ${{ vars.WEB_SENTRY_DSN }}
           VITE_SENTRY_RELEASE: web@${{ github.sha }}
+
+      - if: github.ref_name == 'production'
+        uses: planetscale/setup-pscale-action@v1
+
+      - if: github.ref_name == 'production'
+        run: |
+          set -euo pipefail
+
+          database="opencode-stats"
+          organization="anomalyco"
+          branch="unique-users-${GITHUB_SHA::12}"
+          password_id=""
+
+          cleanup() {
+            if [ -n "$password_id" ]; then
+              pscale password delete "$database" "$branch" "$password_id" --org "$organization" --force >/dev/null 2>&1 || true
+            fi
+            pscale branch delete "$database" "$branch" --org "$organization" --force >/dev/null 2>&1 || true
+          }
+
+          trap cleanup EXIT
+
+          if bun sst shell --stage=production -- bun run --cwd packages/stats/core db:check-unique-users; then
+            echo "unique_users columns already exist in production"
+            exit 0
+          fi
+
+          pscale branch delete "$database" "$branch" --org "$organization" --force >/dev/null 2>&1 || true
+          pscale branch create "$database" "$branch" --org "$organization" --from production --wait
+
+          response="$(pscale password create "$database" "$branch" "unique-users-${GITHUB_RUN_ID}" --org "$organization" --format json)"
+          password_id="$(echo "$response" | jq -r '.id')"
+
+          export PLANETSCALE_HOST="$(echo "$response" | jq -r '.access_host_url')"
+          export PLANETSCALE_USERNAME="$(echo "$response" | jq -r '.username')"
+          export PLANETSCALE_PASSWORD="$(echo "$response" | jq -r '.plain_text')"
+          export PLANETSCALE_DATABASE="$database"
+
+          echo "::add-mask::$PLANETSCALE_PASSWORD"
+          bun run --cwd packages/stats/core db:ensure-unique-users
+
+          deploy_response="$(pscale deploy-request create "$database" "$branch" --org "$organization" --deploy-to production --format json)"
+          deploy_number="$(echo "$deploy_response" | jq -r '.number')"
+
+          if [ -z "$deploy_number" ] || [ "$deploy_number" = "null" ]; then
+            echo "Could not read deploy request number"
+            exit 1
+          fi
+
+          pscale deploy-request review "$database" "$deploy_number" --org "$organization" --approve || true
+          pscale deploy-request deploy "$database" "$deploy_number" --org "$organization"
+        env:
+          CLOUDFLARE_API_TOKEN: ${{ secrets.CLOUDFLARE_API_TOKEN }}
+          PLANETSCALE_SERVICE_TOKEN_NAME: ${{ secrets.PLANETSCALE_SERVICE_TOKEN_NAME }}
+          PLANETSCALE_SERVICE_TOKEN_ID: ${{ secrets.PLANETSCALE_SERVICE_TOKEN_NAME }}
+          PLANETSCALE_SERVICE_TOKEN: ${{ secrets.PLANETSCALE_SERVICE_TOKEN }}
+          STRIPE_SECRET_KEY: ${{ github.ref_name == 'production' && secrets.STRIPE_SECRET_KEY_PROD || secrets.STRIPE_SECRET_KEY_DEV }}
+          HONEYCOMB_API_KEY: ${{ secrets.HONEYCOMB_API_KEY }}
+          SENTRY_AUTH_TOKEN: ${{ secrets.SENTRY_AUTH_TOKEN }}
+          SENTRY_ORG: ${{ vars.SENTRY_ORG }}
+          SENTRY_PROJECT: ${{ vars.WEB_SENTRY_PROJECT }}
+          SENTRY_RELEASE: web@${{ github.sha }}
+          VITE_SENTRY_DSN: ${{ vars.WEB_SENTRY_DSN }}
+          VITE_SENTRY_RELEASE: web@${{ github.sha }}

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}`)
 }