Protect deployment database connections

.github/scripts/check-database-capacity.py

@@ -0,0 +1,59 @@
+"""Read aggregate connection capacity with libpq; never print credentials or SQL data."""
+
+import json
+import os
+import subprocess
+import sys
+from urllib.parse import unquote, urlparse
+
+
+def main() -> int:
+    try:
+        url = urlparse(os.environ["DATABASE_URL"])
+        if url.scheme not in ("postgres", "postgresql") or not url.hostname:
+            raise ValueError("Invalid database URL")
+        if url.hostname.endswith(".psdb.cloud") and url.port not in (None, 5432):
+            raise ValueError("Capacity checks require the direct endpoint")
+        env = {
+            **os.environ,
+            "PGHOST": url.hostname,
+            "PGPORT": str(url.port or 5432),
+            "PGUSER": unquote(url.username or ""),
+            "PGPASSWORD": unquote(url.password or ""),
+            "PGDATABASE": unquote(url.path.removeprefix("/")),
+            "PGSSLMODE": "require",
+            "PGCONNECT_TIMEOUT": "10",
+            "PGAPPNAME": "database-capacity-check",
+            "PGOPTIONS": "-c default_transaction_read_only=on -c statement_timeout=10000",
+        }
+        result = subprocess.run(
+            ["psql", "-X", "-A", "-t", "-v", "ON_ERROR_STOP=1", "-c", """
+                SELECT json_build_object(
+                    'limit', current_setting('max_connections')::int,
+                    'reserved', current_setting('superuser_reserved_connections')::int
+                        + current_setting('reserved_connections')::int,
+                    'used', count(*)::int
+                ) FROM pg_stat_activity WHERE backend_type = 'client backend'
+            """],
+            env=env,
+            capture_output=True,
+            text=True,
+            timeout=25,
+            check=True,
+        )
+        capacity = json.loads(result.stdout)
+        if any(type(capacity[key]) is not int for key in ("limit", "reserved", "used")):
+            raise ValueError("Invalid capacity response")
+        free = capacity["limit"] - capacity["reserved"] - capacity["used"]
+        print(json.dumps({**capacity, "ordinary_free": free, "minimum_free": 10}))
+        if free < 10:
+            print("::error::Database connection headroom is below 10 slots. Inspect direct clients and the PgBouncer budget.")
+            return 1
+        return 0
+    except (KeyError, ValueError, OSError, subprocess.SubprocessError):
+        print("::error::Database capacity check failed. Check direct endpoint access and provider health.")
+        return 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

.github/workflows/database-capacity.yml

@@ -0,0 +1,27 @@
+name: Database capacity
+
+on:
+  schedule:
+    - cron: "2-57/5 * * * *"
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+concurrency:
+  group: database-capacity
+  cancel-in-progress: false
+
+jobs:
+  check:
+    runs-on: ubuntu-24.04
+    timeout-minutes: 2
+    environment: production
+    steps:
+      - uses: actions/checkout@v4
+      # psql is supplied by the Ubuntu runner image. A failed check uses the
+      # repository's Actions failure notifications; no customer data is logged.
+      - name: Check ordinary connection headroom
+        run: python3 .github/scripts/check-database-capacity.py
+        env:
+          DATABASE_URL: ${{ secrets.DATABASE_URL }}

apps/cloud/docs/database-connections.md

@@ -0,0 +1,43 @@
+# Production database connections
+
+Application traffic uses Hyperdrive, then PlanetScale's local transaction-mode
+PgBouncer on port 6432. Deployment scripts use the direct endpoint on port 5432.
+Code migrations hold session advisory locks, so they must bypass transaction pooling.
+
+The connection budget is:
+
+| Setting                                   | Value           |
+| ----------------------------------------- | --------------- |
+| PostgreSQL max_connections                | 50              |
+| PostgreSQL superuser_reserved_connections | 3               |
+| Local PgBouncer processes                 | 1               |
+| PgBouncer default_pool_size               | 12              |
+| PgBouncer max_db_connections              | 12              |
+| PgBouncer max_client_conn                 | 400             |
+| PgBouncer max_prepared_statements         | 200             |
+| Hyperdrive origin connection limit        | 12 (soft limit) |
+
+Hyperdrive's origin limit is advisory. PgBouncer's database limit enforces the
+backend budget across users of one database. The cap is per PgBouncer process:
+adding processes, databases, direct clients, or other poolers requires a new
+aggregate budget. Keep capacity for provider sessions, deploys and administration.
+Prepared statements require protocol-level support to remain enabled in PgBouncer.
+
+The migration and membership-readiness scripts retry only the initial `SELECT 1`
+when PostgreSQL returns SQLSTATE `53300`. They make at most seven attempts, with
+ten seconds between attempts and a ten-second connection timeout. They never
+retry migration bodies or readiness mutations. Other errors fail immediately.
+
+The Database capacity workflow checks direct access and aggregate connection
+headroom every five minutes. It fails when fewer than ten ordinary slots remain.
+Counts include the monitor and conservatively count privileged client sessions
+against ordinary capacity. GitHub schedule delays and notification preferences
+apply; this is not a real-time paging service. Check PlanetScale CPU and PgBouncer
+waiting clients alongside Cloudflare query errors and latency during load spikes.
+
+For a routing change, first account for overlapping old and new pools. Verify
+the active PostgreSQL limit and applied pool settings before changing Hyperdrive.
+Afterward, check an authenticated application page, direct database access,
+backend counts, and provider errors. Roll back by restoring the prior Hyperdrive
+origin port only while there is capacity for both pools. Do not kill idle sessions
+as routine maintenance: clients can reconnect and consume the slots again.

apps/cloud/scripts/database-connection.ts

@@ -0,0 +1,56 @@
+/* oxlint-disable executor/no-error-constructor, executor/no-try-catch-or-throw -- boundary: deployment CLI connection acquisition */
+
+import { setTimeout } from "node:timers/promises";
+
+const MAX_ATTEMPTS = 7;
+const RETRY_DELAY_MS = 10_000;
+
+/**
+ * Validate the deploy transport without logging credentials. PlanetScale schema
+ * migrations use the direct endpoint because code migrations hold session locks.
+ */
+export const directDatabaseUrl = (value: string): string => {
+  let url: URL;
+  try {
+    url = new URL(value);
+  } catch {
+    throw new Error("DATABASE_URL must be a valid PostgreSQL URL");
+  }
+  if (url.protocol !== "postgres:" && url.protocol !== "postgresql:") {
+    throw new Error("DATABASE_URL must use the postgres or postgresql protocol");
+  }
+  if (url.hostname.endsWith(".psdb.cloud") && url.port !== "" && url.port !== "5432") {
+    throw new Error("PlanetScale deploy scripts require the direct endpoint on port 5432");
+  }
+  return value;
+};
+
+/**
+ * Open the CLI's single connection before starting work. Retry only PostgreSQL
+ * admission failures (53300), at most six times with ten seconds between tries.
+ * The caller must set connect_timeout and close the client on every exit.
+ * Migration and readiness mutations remain outside this retry boundary.
+ */
+export const waitForDatabaseConnection = async (
+  sql: { readonly unsafe: (query: string) => PromiseLike<unknown> },
+  options: {
+    readonly log: (message: string) => void;
+    readonly sleep?: (milliseconds: number) => Promise<void>;
+  },
+): Promise<void> => {
+  const sleep = options.sleep ?? ((milliseconds: number) => setTimeout(milliseconds));
+  for (let attempt = 1; attempt <= MAX_ATTEMPTS; attempt += 1) {
+    try {
+      await sql.unsafe("SELECT 1");
+      return;
+    } catch (cause) {
+      const isCapacityError =
+        typeof cause === "object" && cause !== null && "code" in cause && cause.code === "53300";
+      if (!isCapacityError || attempt === MAX_ATTEMPTS) throw cause;
+      options.log(
+        `Database connection capacity is full (53300). Retrying connection ${attempt}/${MAX_ATTEMPTS - 1} in 10s; no work has started.`,
+      );
+      await sleep(RETRY_DELAY_MS);
+    }
+  }
+};

apps/cloud/scripts/ensure-workos-mirror-ready.ts

@@ -36,6 +36,7 @@ import { fileURLToPath } from "node:url";
 
 import { drizzle } from "drizzle-orm/postgres-js";
 import postgres from "postgres";
+import { directDatabaseUrl, waitForDatabaseConnection } from "./database-connection";
 
 import {
   MirrorReadinessState,
@@ -56,9 +57,10 @@ if (!connectionString) {
 const usesLocalDatabase =
   connectionString.includes("127.0.0.1") || connectionString.includes("localhost");
 
-const sql = postgres(connectionString, {
+const sql = postgres(directDatabaseUrl(connectionString), {
   max: 1,
   prepare: false,
+  connect_timeout: 10,
   ...(usesLocalDatabase ? {} : { ssl: "require" as const }),
 });
 const db = drizzle(sql);
@@ -84,6 +86,7 @@ const runScript = (what: string, script: string) => {
 };
 
 try {
+  await waitForDatabaseConnection(sql, { log });
   let state = await readiness();
   log(describeMirrorReadiness(state));
 

apps/cloud/scripts/migrate.ts

@@ -8,6 +8,7 @@ import { migrate as migrateDrizzle } from "drizzle-orm/postgres-js/migrator";
 import postgres from "postgres";
 
 import { cloudCodeMigrations, runCodeMigrations } from "./code-migrations/index";
+import { directDatabaseUrl, waitForDatabaseConnection } from "./database-connection";
 
 const __dirname = dirname(fileURLToPath(import.meta.url));
 const MIGRATIONS_FOLDER = resolve(__dirname, "../drizzle");
@@ -41,13 +42,15 @@ if (!connectionString) {
 const usesLocalDatabase =
   connectionString.includes("127.0.0.1") || connectionString.includes("localhost");
 
-const sql = postgres(connectionString, {
+const sql = postgres(directDatabaseUrl(connectionString), {
   max: 1,
   prepare: false,
+  connect_timeout: 10,
   ...(usesLocalDatabase ? {} : { ssl: "require" as const }),
 });
 
 try {
+  await waitForDatabaseConnection(sql, { log: console.log });
   if (!codeOnly) {
     if (dryRun) {
       console.log("[schema-migrate] dry run: Drizzle SQL migrations are not applied");

apps/cloud/src/db/deployment-connection.test.ts

@@ -0,0 +1,93 @@
+/* oxlint-disable executor/no-promise-reject -- boundary: simulate the Postgres.js driver's rejected promises */
+
+import { describe, expect, it } from "@effect/vitest";
+import { Effect } from "effect";
+
+import { directDatabaseUrl, waitForDatabaseConnection } from "../../scripts/database-connection";
+
+describe("deployment database connection", () => {
+  it.effect("waits for admission before allowing deployment work", () =>
+    Effect.promise(async () => {
+      let remainingFailures = 2;
+      const waits: number[] = [];
+      const logs: string[] = [];
+      const queries: string[] = [];
+      await waitForDatabaseConnection(
+        {
+          unsafe: (query) => {
+            queries.push(query);
+            return remainingFailures-- > 0
+              ? Promise.reject({ code: "53300", detail: "private connection data" })
+              : Promise.resolve([]);
+          },
+        },
+        {
+          log: (line) => logs.push(line),
+          sleep: async (ms) => {
+            waits.push(ms);
+          },
+        },
+      );
+      expect(queries).toEqual(["SELECT 1", "SELECT 1", "SELECT 1"]);
+      expect(waits).toEqual([10_000, 10_000]);
+      expect(logs).toHaveLength(2);
+      expect(logs.join()).not.toContain("private connection data");
+    }),
+  );
+
+  it.effect("fails after the bounded admission budget", () =>
+    Effect.promise(async () => {
+      const failure = { code: "53300" };
+      let attempts = 0;
+      const waits: number[] = [];
+      await expect(
+        waitForDatabaseConnection(
+          {
+            unsafe: () => {
+              attempts += 1;
+              return Promise.reject(failure);
+            },
+          },
+          {
+            log: () => {},
+            sleep: async (ms) => {
+              waits.push(ms);
+            },
+          },
+        ),
+      ).rejects.toBe(failure);
+      expect(attempts).toBe(7);
+      expect(waits).toEqual(Array(6).fill(10_000));
+    }),
+  );
+
+  it.effect("fails immediately for authentication, transport and SQL errors", () =>
+    Effect.promise(async () => {
+      for (const code of ["28P01", "CONNECT_TIMEOUT", "CONNECTION_CLOSED", "42601", "40001"]) {
+        const failure = { code };
+        const waits: number[] = [];
+        await expect(
+          waitForDatabaseConnection(
+            { unsafe: () => Promise.reject(failure) },
+            {
+              log: () => {},
+              sleep: async (ms) => {
+                waits.push(ms);
+              },
+            },
+          ),
+        ).rejects.toBe(failure);
+        expect(waits).toEqual([]);
+      }
+    }),
+  );
+
+  it("keeps PlanetScale deployment traffic on the direct endpoint", () => {
+    const direct = "postgres://example:secret@region.pg.psdb.cloud:5432/database";
+    expect(directDatabaseUrl(direct)).toBe(direct);
+    expect(directDatabaseUrl("postgres://localhost:25432/postgres")).toContain(":25432");
+    expect(() => directDatabaseUrl(direct.replace(":5432", ":6432"))).toThrow("direct endpoint");
+    expect(() => directDatabaseUrl("invalid-secret")).toThrow("valid PostgreSQL URL");
+    expect(() => directDatabaseUrl("https://localhost/database")).toThrow("protocol");
+  });
+});