Protect deployment database connections

Rhys Sullivan committed 1 month ago

Commit 4fb6ddd1 parent 3890d6f7 files · +286 −2Browse files

Unified · Split

7 files

.github/scripts/check-database-capacity.py +59 −0
.github/workflows/database-capacity.yml +27 −0
apps/cloud/docs/database-connections.md +43 −0
apps/cloud/scripts/database-connection.ts +56 −0
apps/cloud/scripts/ensure-workos-mirror-ready.ts +4 −1
apps/cloud/scripts/migrate.ts +4 −1
apps/cloud/src/db/deployment-connection.test.ts +93 −0

.github/scripts/check-database-capacity.py+59−0

@@ -0,0 +1,59 @@
	1	+"""Read aggregate connection capacity with libpq; never print credentials or SQL data."""
	2	+
	3	+import json
	4	+import os
	5	+import subprocess
	6	+import sys
	7	+from urllib.parse import unquote, urlparse
	8	+
	9	+
	10	+def main() -> int:
	11	+ try:
	12	+ url = urlparse(os.environ["DATABASE_URL"])
	13	+ if url.scheme not in ("postgres", "postgresql") or not url.hostname:
	14	+ raise ValueError("Invalid database URL")
	15	+ if url.hostname.endswith(".psdb.cloud") and url.port not in (None, 5432):
	16	+ raise ValueError("Capacity checks require the direct endpoint")
	17	+ env = {
	18	+ **os.environ,
	19	+ "PGHOST": url.hostname,
	20	+ "PGPORT": str(url.port or 5432),
	21	+ "PGUSER": unquote(url.username or ""),
	22	+ "PGPASSWORD": unquote(url.password or ""),
	23	+ "PGDATABASE": unquote(url.path.removeprefix("/")),
	24	+ "PGSSLMODE": "require",
	25	+ "PGCONNECT_TIMEOUT": "10",
	26	+ "PGAPPNAME": "database-capacity-check",
	27	+ "PGOPTIONS": "-c default_transaction_read_only=on -c statement_timeout=10000",
	28	+ }
	29	+ result = subprocess.run(
	30	+ ["psql", "-X", "-A", "-t", "-v", "ON_ERROR_STOP=1", "-c", """
	31	+ SELECT json_build_object(
	32	+ 'limit', current_setting('max_connections')::int,
	33	+ 'reserved', current_setting('superuser_reserved_connections')::int
	34	+ + current_setting('reserved_connections')::int,
	35	+ 'used', count(*)::int
	36	+ ) FROM pg_stat_activity WHERE backend_type = 'client backend'
	37	+ """],
	38	+ env=env,
	39	+ capture_output=True,
	40	+ text=True,
	41	+ timeout=25,
	42	+ check=True,
	43	+ )
	44	+ capacity = json.loads(result.stdout)
	45	+ if any(type(capacity[key]) is not int for key in ("limit", "reserved", "used")):
	46	+ raise ValueError("Invalid capacity response")
	47	+ free = capacity["limit"] - capacity["reserved"] - capacity["used"]
	48	+ print(json.dumps({**capacity, "ordinary_free": free, "minimum_free": 10}))
	49	+ if free < 10:
	50	+ print("::error::Database connection headroom is below 10 slots. Inspect direct clients and the PgBouncer budget.")
	51	+ return 1
	52	+ return 0
	53	+ except (KeyError, ValueError, OSError, subprocess.SubprocessError):
	54	+ print("::error::Database capacity check failed. Check direct endpoint access and provider health.")
	55	+ return 1
	56	+
	57	+
	58	+if __name__ == "__main__":
	59	+ sys.exit(main())

.github/workflows/database-capacity.yml+27−0

@@ -0,0 +1,27 @@
	1	+name: Database capacity
	2	+
	3	+on:
	4	+ schedule:
	5	+ - cron: "2-57/5 * * * *"
	6	+ workflow_dispatch:
	7	+
	8	+permissions:
	9	+ contents: read
	10	+
	11	+concurrency:
	12	+ group: database-capacity
	13	+ cancel-in-progress: false
	14	+
	15	+jobs:
	16	+ check:
	17	+ runs-on: ubuntu-24.04
	18	+ timeout-minutes: 2
	19	+ environment: production
	20	+ steps:
	21	+ - uses: actions/checkout@v4
	22	+ # psql is supplied by the Ubuntu runner image. A failed check uses the
	23	+ # repository's Actions failure notifications; no customer data is logged.
	24	+ - name: Check ordinary connection headroom
	25	+ run: python3 .github/scripts/check-database-capacity.py
	26	+ env:
	27	+ DATABASE_URL: ${{ secrets.DATABASE_URL }}

apps/cloud/docs/database-connections.md+43−0

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

apps/cloud/scripts/database-connection.ts+56−0

@@ -0,0 +1,56 @@
	1	+/* oxlint-disable executor/no-error-constructor, executor/no-try-catch-or-throw -- boundary: deployment CLI connection acquisition */
	2	+
	3	+import { setTimeout } from "node:timers/promises";
	4	+
	5	+const MAX_ATTEMPTS = 7;
	6	+const RETRY_DELAY_MS = 10_000;
	7	+
	8	+/**
	9	+ * Validate the deploy transport without logging credentials. PlanetScale schema
	10	+ * migrations use the direct endpoint because code migrations hold session locks.
	11	+ */
	12	+export const directDatabaseUrl = (value: string): string => {
	13	+ let url: URL;
	14	+ try {
	15	+ url = new URL(value);
	16	+ } catch {
	17	+ throw new Error("DATABASE_URL must be a valid PostgreSQL URL");
	18	+ }
	19	+ if (url.protocol !== "postgres:" && url.protocol !== "postgresql:") {
	20	+ throw new Error("DATABASE_URL must use the postgres or postgresql protocol");
	21	+ }
	22	+ if (url.hostname.endsWith(".psdb.cloud") && url.port !== "" && url.port !== "5432") {
	23	+ throw new Error("PlanetScale deploy scripts require the direct endpoint on port 5432");
	24	+ }
	25	+ return value;
	26	+};
	27	+
	28	+/**
	29	+ * Open the CLI's single connection before starting work. Retry only PostgreSQL
	30	+ * admission failures (53300), at most six times with ten seconds between tries.
	31	+ * The caller must set connect_timeout and close the client on every exit.
	32	+ * Migration and readiness mutations remain outside this retry boundary.
	33	+ */
	34	+export const waitForDatabaseConnection = async (
	35	+ sql: { readonly unsafe: (query: string) => PromiseLike<unknown> },
	36	+ options: {
	37	+ readonly log: (message: string) => void;
	38	+ readonly sleep?: (milliseconds: number) => Promise<void>;
	39	+ },
	40	+): Promise<void> => {
	41	+ const sleep = options.sleep ?? ((milliseconds: number) => setTimeout(milliseconds));
	42	+ for (let attempt = 1; attempt <= MAX_ATTEMPTS; attempt += 1) {
	43	+ try {
	44	+ await sql.unsafe("SELECT 1");
	45	+ return;
	46	+ } catch (cause) {
	47	+ const isCapacityError =
	48	+ typeof cause === "object" && cause !== null && "code" in cause && cause.code === "53300";
	49	+ if (!isCapacityError \|\| attempt === MAX_ATTEMPTS) throw cause;
	50	+ options.log(
	51	+ `Database connection capacity is full (53300). Retrying connection ${attempt}/${MAX_ATTEMPTS - 1} in 10s; no work has started.`,
	52	+ );
	53	+ await sleep(RETRY_DELAY_MS);
	54	+ }
	55	+ }
	56	+};

apps/cloud/scripts/ensure-workos-mirror-ready.ts+4−1

@@ -36,6 +36,7 @@ import { fileURLToPath } from "node:url"; · 35 unmodified lines
36	36
37	37	import { drizzle } from "drizzle-orm/postgres-js";
38	38	import postgres from "postgres";
	39	+import { directDatabaseUrl, waitForDatabaseConnection } from "./database-connection";
39	40
40	41	import {
41	42	MirrorReadinessState,
@@ -56,9 +57,10 @@ if (!connectionString) { · 14 unmodified lines
56	57	const usesLocalDatabase =
57	58	connectionString.includes("127.0.0.1") \|\| connectionString.includes("localhost");
58	59
59		−const sql = postgres(connectionString, {
	60	+const sql = postgres(directDatabaseUrl(connectionString), {
60	61	max: 1,
61	62	prepare: false,
	63	+ connect_timeout: 10,
62	64	...(usesLocalDatabase ? {} : { ssl: "require" as const }),
63	65	});
64	66	const db = drizzle(sql);
@@ -84,6 +86,7 @@ const runScript = (what: string, script: string) => { · 19 unmodified lines
84	86	};
85	87
86	88	try {
	89	+ await waitForDatabaseConnection(sql, { log });
87	90	let state = await readiness();
88	91	log(describeMirrorReadiness(state));
89	92

apps/cloud/scripts/migrate.ts+4−1

@@ -8,6 +8,7 @@ import { migrate as migrateDrizzle } from "drizzle-orm/postgres-js/migrator"; · 7 unmodified lines
8	8	import postgres from "postgres";
9	9
10	10	import { cloudCodeMigrations, runCodeMigrations } from "./code-migrations/index";
	11	+import { directDatabaseUrl, waitForDatabaseConnection } from "./database-connection";
11	12
12	13	const __dirname = dirname(fileURLToPath(import.meta.url));
13	14	const MIGRATIONS_FOLDER = resolve(__dirname, "../drizzle");
@@ -41,13 +42,15 @@ if (!connectionString) { · 27 unmodified lines
41	42	const usesLocalDatabase =
42	43	connectionString.includes("127.0.0.1") \|\| connectionString.includes("localhost");
43	44
44		−const sql = postgres(connectionString, {
	45	+const sql = postgres(directDatabaseUrl(connectionString), {
45	46	max: 1,
46	47	prepare: false,
	48	+ connect_timeout: 10,
47	49	...(usesLocalDatabase ? {} : { ssl: "require" as const }),
48	50	});
49	51
50	52	try {
	53	+ await waitForDatabaseConnection(sql, { log: console.log });
51	54	if (!codeOnly) {
52	55	if (dryRun) {
53	56	console.log("[schema-migrate] dry run: Drizzle SQL migrations are not applied");

apps/cloud/src/db/deployment-connection.test.ts+93−0

@@ -0,0 +1,93 @@
	1	+/* oxlint-disable executor/no-promise-reject -- boundary: simulate the Postgres.js driver's rejected promises */
	2	+
	3	+import { describe, expect, it } from "@effect/vitest";
	4	+import { Effect } from "effect";
	5	+
	6	+import { directDatabaseUrl, waitForDatabaseConnection } from "../../scripts/database-connection";
	7	+
	8	+describe("deployment database connection", () => {
	9	+ it.effect("waits for admission before allowing deployment work", () =>
	10	+ Effect.promise(async () => {
	11	+ let remainingFailures = 2;
	12	+ const waits: number[] = [];
	13	+ const logs: string[] = [];
	14	+ const queries: string[] = [];
	15	+ await waitForDatabaseConnection(
	16	+ {
	17	+ unsafe: (query) => {
	18	+ queries.push(query);
	19	+ return remainingFailures-- > 0
	20	+ ? Promise.reject({ code: "53300", detail: "private connection data" })
	21	+ : Promise.resolve([]);
	22	+ },
	23	+ },
	24	+ {
	25	+ log: (line) => logs.push(line),
	26	+ sleep: async (ms) => {
	27	+ waits.push(ms);
	28	+ },
	29	+ },
	30	+ );
	31	+ expect(queries).toEqual(["SELECT 1", "SELECT 1", "SELECT 1"]);
	32	+ expect(waits).toEqual([10_000, 10_000]);
	33	+ expect(logs).toHaveLength(2);
	34	+ expect(logs.join()).not.toContain("private connection data");
	35	+ }),
	36	+ );
	37	+
	38	+ it.effect("fails after the bounded admission budget", () =>
	39	+ Effect.promise(async () => {
	40	+ const failure = { code: "53300" };
	41	+ let attempts = 0;
	42	+ const waits: number[] = [];
	43	+ await expect(
	44	+ waitForDatabaseConnection(
	45	+ {
	46	+ unsafe: () => {
	47	+ attempts += 1;
	48	+ return Promise.reject(failure);
	49	+ },
	50	+ },
	51	+ {
	52	+ log: () => {},
	53	+ sleep: async (ms) => {
	54	+ waits.push(ms);
	55	+ },
	56	+ },
	57	+ ),
	58	+ ).rejects.toBe(failure);
	59	+ expect(attempts).toBe(7);
	60	+ expect(waits).toEqual(Array(6).fill(10_000));
	61	+ }),
	62	+ );
	63	+
	64	+ it.effect("fails immediately for authentication, transport and SQL errors", () =>
	65	+ Effect.promise(async () => {
	66	+ for (const code of ["28P01", "CONNECT_TIMEOUT", "CONNECTION_CLOSED", "42601", "40001"]) {
	67	+ const failure = { code };
	68	+ const waits: number[] = [];
	69	+ await expect(
	70	+ waitForDatabaseConnection(
	71	+ { unsafe: () => Promise.reject(failure) },
	72	+ {
	73	+ log: () => {},
	74	+ sleep: async (ms) => {
	75	+ waits.push(ms);
	76	+ },
	77	+ },
	78	+ ),
	79	+ ).rejects.toBe(failure);
	80	+ expect(waits).toEqual([]);
	81	+ }
	82	+ }),
	83	+ );
	84	+
	85	+ it("keeps PlanetScale deployment traffic on the direct endpoint", () => {
	86	+ const direct = "postgres://example:secret@region.pg.psdb.cloud:5432/database";
	87	+ expect(directDatabaseUrl(direct)).toBe(direct);
	88	+ expect(directDatabaseUrl("postgres://localhost:25432/postgres")).toContain(":25432");
	89	+ expect(() => directDatabaseUrl(direct.replace(":5432", ":6432"))).toThrow("direct endpoint");
	90	+ expect(() => directDatabaseUrl("invalid-secret")).toThrow("valid PostgreSQL URL");
	91	+ expect(() => directDatabaseUrl("https://localhost/database")).toThrow("protocol");
	92	+ });
	93	+});