Retry deploy scripts when Postgres has no free connection slots

apps/cloud/scripts/backfill-workos-mirror.ts

@@ -42,6 +42,7 @@ import { WorkOS } from "@workos-inc/node";
 import { backfillWorkOsMirror } from "../src/auth/workos-mirror-backfill";
 import { makeWorkOsMirrorStore } from "../src/auth/workos-mirror-store";
 import { organizations } from "../src/db/schema";
+import { describeRefusedAttempt, waitForConnectionSlot } from "../src/db/too-many-connections";
 
 const dryRun = process.argv.includes("--dry-run");
 
@@ -71,32 +72,46 @@ const workos = new WorkOS(apiKey);
 const fromPromise = <A>(fn: () => Promise<A>) =>
   Effect.tryPromise({ try: fn, catch: (cause) => cause });
 
+// A full server refuses the connection with SQLSTATE 53300 when it is opened.
+// Open it first, waiting for a slot, rather than fail the deploy gate that
+// spawned this run (src/db/too-many-connections.ts); the mirror store's own
+// failures do not carry the driver code, so the wait cannot sit around the
+// backfill itself.
+const connected = waitForConnectionSlot(sql, {
+  onRefused: (_failure, attempt) => console.log(describeRefusedAttempt(attempt)),
+});
+
+const backfill = backfillWorkOsMirror(
+  {
+    listOrganizationIds: () =>
+      fromPromise(async () => {
+        // Never a deleted organization: its row is a tombstone (its
+        // memberships are purged, WorkOS no longer has it) and the mirror
+        // refuses a scan of it anyway.
+        const rows = await db
+          .select({ id: organizations.id })
+          .from(organizations)
+          .where(isNull(organizations.deletedAt))
+          .orderBy(asc(organizations.createdAt));
+        return rows.map((row) => row.id);
+      }),
+    listOrgMembers: (organizationId) =>
+      fromPromise(async () => {
+        const page = await workos.userManagement.listOrganizationMemberships({
+          organizationId,
+          statuses: ["active", "pending", "inactive"],
+        });
+        return page.listMetadata.after ? page.autoPagination() : page.data;
+      }),
+    getUser: (userId) => fromPromise(() => workos.userManagement.getUser(userId)),
+  },
+  makeWorkOsMirrorStore(db),
+  { dryRun, log: (line) => console.log(line) },
+);
+
 await Effect.runPromise(
-  backfillWorkOsMirror(
-    {
-      listOrganizationIds: () =>
-        fromPromise(async () => {
-          // Never a deleted organization: its row is a tombstone (its
-          // memberships are purged, WorkOS no longer has it) and the mirror
-          // refuses a scan of it anyway.
-          const rows = await db
-            .select({ id: organizations.id })
-            .from(organizations)
-            .where(isNull(organizations.deletedAt))
-            .orderBy(asc(organizations.createdAt));
-          return rows.map((row) => row.id);
-        }),
-      listOrgMembers: (organizationId) =>
-        fromPromise(async () => {
-          const page = await workos.userManagement.listOrganizationMemberships({
-            organizationId,
-            statuses: ["active", "pending", "inactive"],
-          });
-          return page.listMetadata.after ? page.autoPagination() : page.data;
-        }),
-      getUser: (userId) => fromPromise(() => workos.userManagement.getUser(userId)),
-    },
-    makeWorkOsMirrorStore(db),
-    { dryRun, log: (line) => console.log(line) },
-  ).pipe(Effect.ensuring(Effect.promise(() => sql.end({ timeout: 5 })))),
+  connected.pipe(
+    Effect.andThen(backfill),
+    Effect.ensuring(Effect.promise(() => sql.end({ timeout: 5 }))),
+  ),
 );

apps/cloud/scripts/drain-workos-events.ts

@@ -26,6 +26,7 @@ import { WorkOS } from "@workos-inc/node";
 import { makeUserStore } from "../src/auth/user-store";
 import { replayWorkOsEvents, type WorkOsEventsSyncReport } from "../src/auth/workos-events-replay";
 import { makeWorkOsMirrorStore } from "../src/auth/workos-mirror-store";
+import { describeRefusedAttempt, waitForConnectionSlot } from "../src/db/too-many-connections";
 
 const connectionString = process.env.DATABASE_URL;
 if (!connectionString) {
@@ -111,8 +112,21 @@ const drain = Effect.gen(function* () {
   return last;
 });
 
+// A full server refuses the connection with SQLSTATE 53300 when it is opened.
+// Open it first, waiting for a slot, rather than fail the deploy gate that
+// spawned this run (src/db/too-many-connections.ts); the mirror store's own
+// failures do not carry the driver code, so the wait cannot sit around the
+// drain itself.
+const connected = waitForConnectionSlot(sql, {
+  onRefused: (_failure, attempt) =>
+    console.log(`[drain-events] ${describeRefusedAttempt(attempt)}`),
+});
+
 const report = await Effect.runPromise(
-  drain.pipe(Effect.ensuring(Effect.promise(() => sql.end({ timeout: 5 })))),
+  connected.pipe(
+    Effect.andThen(drain),
+    Effect.ensuring(Effect.promise(() => sql.end({ timeout: 5 }))),
+  ),
 );
 
 if (report === null || report.stopped !== "drained") {

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

@@ -35,13 +35,18 @@ import { dirname, resolve } from "node:path";
 import { fileURLToPath } from "node:url";
 
 import { drizzle } from "drizzle-orm/postgres-js";
+import { Result } from "effect";
 import postgres from "postgres";
 
 import {
   MirrorReadinessState,
   describeMirrorReadiness,
   readMirrorReadiness,
 } from "../src/auth/mirror-readiness-store";
+import {
+  describeRefusedAttempt,
+  retryWhileTooManyConnections,
+} from "../src/db/too-many-connections";
 
 const __dirname = dirname(fileURLToPath(import.meta.url));
 const BACKFILL_SCRIPT = resolve(__dirname, "backfill-workos-mirror.ts");
@@ -65,7 +70,19 @@ const db = drizzle(sql);
 
 const log = (line: string) => console.log(`[mirror-ready] ${line}`);
 
-const readiness = () => readMirrorReadiness(db, new Date());
+// A full server refuses the connection with SQLSTATE 53300 before a statement
+// runs — on the first read, or on a later one after postgres.js has reopened a
+// dropped connection — so every read waits for a slot instead of failing the
+// deploy, as the migration step before this one does
+// (src/db/too-many-connections.ts). The backfill and drain scripts below each
+// open their own connection and wait for their own slot.
+const readiness = async () => {
+  const outcome = await retryWhileTooManyConnections(() => readMirrorReadiness(db, new Date()), {
+    onRefused: (_failure, attempt) => log(describeRefusedAttempt(attempt)),
+  });
+  if (Result.isFailure(outcome)) throw outcome.failure;
+  return outcome.success;
+};
 
 // The backfill and drain scripts own their own WorkOS + database wiring;
 // running them as subprocesses (with this process's env) keeps that wiring

apps/cloud/scripts/migrate.ts

@@ -5,8 +5,13 @@ import { fileURLToPath } from "node:url";
 
 import { drizzle } from "drizzle-orm/postgres-js";
 import { migrate as migrateDrizzle } from "drizzle-orm/postgres-js/migrator";
+import { Result } from "effect";
 import postgres from "postgres";
 
+import {
+  describeRefusedAttempt,
+  retryWhileTooManyConnections,
+} from "../src/db/too-many-connections";
 import { cloudCodeMigrations, runCodeMigrations } from "./code-migrations/index";
 
 const __dirname = dirname(fileURLToPath(import.meta.url));
@@ -47,13 +52,29 @@ const sql = postgres(connectionString, {
   ...(usesLocalDatabase ? {} : { ssl: "require" as const }),
 });
 
+// A full server refuses the connection with SQLSTATE 53300 before a statement
+// runs — on the first statement, or on a later one after postgres.js has
+// reopened a dropped connection. So each step over the direct connection
+// waits for a slot instead of failing the deploy (src/db/too-many-connections.ts).
+// A step is safe to repeat: Drizzle and the code-migration ledger each skip
+// what has already been applied.
+const withConnectionSlot = async <A>(run: () => Promise<A>): Promise<A> => {
+  const outcome = await retryWhileTooManyConnections(run, {
+    onRefused: (_failure, attempt) => console.warn(`[migrate] ${describeRefusedAttempt(attempt)}`),
+  });
+  if (Result.isFailure(outcome)) throw outcome.failure;
+  return outcome.success;
+};
+
 try {
   if (!codeOnly) {
     if (dryRun) {
       console.log("[schema-migrate] dry run: Drizzle SQL migrations are not applied");
     } else {
       console.log(`[schema-migrate] running Drizzle migrations from ${MIGRATIONS_FOLDER}`);
-      await migrateDrizzle(drizzle(sql), { migrationsFolder: MIGRATIONS_FOLDER });
+      await withConnectionSlot(() =>
+        migrateDrizzle(drizzle(sql), { migrationsFolder: MIGRATIONS_FOLDER }),
+      );
       console.log("[schema-migrate] complete");
     }
   }
@@ -63,7 +84,9 @@ try {
     if (migrations.length === 0) {
       console.log("[code-migrate] no code migrations configured");
     } else {
-      const applied = await runCodeMigrations(sql, migrations, { dryRun });
+      const applied = await withConnectionSlot(() =>
+        runCodeMigrations(sql, migrations, { dryRun }),
+      );
       console.log(
         dryRun
           ? `[code-migrate] dry run planned ${applied.length} migration(s)`

apps/cloud/src/db/too-many-connections.test.ts

@@ -0,0 +1,137 @@
+/* oxlint-disable executor/no-try-catch-or-throw, executor/no-error-constructor -- test doubles: simulate the driver's rejected promise and its Error-shaped cause chain */
+import { describe, expect, it } from "@effect/vitest";
+import { Effect, Fiber, Result, Schedule } from "effect";
+import { TestClock } from "effect/testing";
+
+import {
+  TOO_MANY_CONNECTIONS_RETRIES,
+  TOO_MANY_CONNECTIONS_SQLSTATE,
+  describeRefusedAttempt,
+  isTooManyConnectionsError,
+  retryTooManyConnections,
+  retryWhileTooManyConnections,
+} from "./too-many-connections";
+
+// The shape postgres.js + Drizzle produce in the deploy log: Drizzle's
+// "Failed query" error with the driver's PostgresError (SQLSTATE `code`) as
+// its cause.
+const refused = () =>
+  Object.assign(new Error('Failed query: CREATE SCHEMA IF NOT EXISTS "drizzle"'), {
+    cause: Object.assign(
+      new Error("remaining connection slots are reserved for roles with the SUPERUSER attribute"),
+      { code: TOO_MANY_CONNECTIONS_SQLSTATE },
+    ),
+  });
+
+const noDelay = Schedule.recurs(2);
+
+describe("isTooManyConnectionsError", () => {
+  it("matches SQLSTATE 53300 anywhere in the cause chain", () => {
+    expect(isTooManyConnectionsError(refused())).toBe(true);
+    expect(isTooManyConnectionsError({ code: "53300" })).toBe(true);
+  });
+
+  it("rejects other driver codes and non-errors", () => {
+    expect(isTooManyConnectionsError({ code: "CONNECT_TIMEOUT" })).toBe(false);
+    expect(isTooManyConnectionsError(new Error("boom"))).toBe(false);
+    expect(isTooManyConnectionsError(undefined)).toBe(false);
+    expect(isTooManyConnectionsError("53300")).toBe(false);
+  });
+});
+
+describe("retryWhileTooManyConnections", () => {
+  it("retries a refused connection and resolves with the first success", async () => {
+    let calls = 0;
+    const refusals: number[] = [];
+    const result = await retryWhileTooManyConnections(
+      async () => {
+        calls += 1;
+        if (calls < 3) throw refused();
+        return "applied";
+      },
+      { schedule: noDelay, onRefused: (_, attempt) => refusals.push(attempt) },
+    );
+    expect(result).toEqual(Result.succeed("applied"));
+    expect(calls).toBe(3);
+    expect(refusals).toEqual([1, 2]);
+  });
+
+  it("rethrows any other failure without retrying", async () => {
+    let calls = 0;
+    const failure = Object.assign(new Error("Failed query: alter table"), {
+      cause: { code: "42P01" },
+    });
+    const result = await retryWhileTooManyConnections(
+      async () => {
+        calls += 1;
+        throw failure;
+      },
+      { schedule: noDelay },
+    );
+    expect(Result.isFailure(result) && result.failure).toBe(failure);
+    expect(calls).toBe(1);
+  });
+
+  it("rethrows the last refusal once the schedule is spent", async () => {
+    let calls = 0;
+    const failures: unknown[] = [];
+    const result = await retryWhileTooManyConnections(
+      async () => {
+        calls += 1;
+        const failure = refused();
+        failures.push(failure);
+        throw failure;
+      },
+      { schedule: noDelay },
+    );
+    expect(calls).toBe(3);
+    expect(Result.isFailure(result) && result.failure).toBe(failures[2]);
+  });
+});
+
+describe("the production schedule", () => {
+  // Virtual time: the production cadence is thirty seconds between attempts,
+  // and a test that waited it out for real would take a quarter of an hour.
+  // `it.effect` runs under the TestClock, which the schedule's sleeps use.
+  it.effect("retries every thirty seconds and gives up after about fifteen minutes", () =>
+    Effect.gen(function* () {
+      let calls = 0;
+      const refusals: number[] = [];
+      const fiber = yield* Effect.forkChild(
+        Effect.result(
+          retryTooManyConnections(
+            Effect.suspend(() => {
+              calls += 1;
+              return Effect.fail(refused());
+            }),
+            { onRefused: (_, attempt) => refusals.push(attempt) },
+          ),
+        ),
+      );
+
+      yield* TestClock.adjust("29 seconds");
+      expect(calls).toBe(1);
+      yield* TestClock.adjust("1 second");
+      expect(calls).toBe(2);
+
+      // Attempt n runs at (n - 1) × 30 s: the thirtieth at 14:30, the last at 15:00.
+      yield* TestClock.adjust("14 minutes");
+      expect(calls).toBe(TOO_MANY_CONNECTIONS_RETRIES);
+      yield* TestClock.adjust("30 seconds");
+      const result = yield* Fiber.join(fiber);
+
+      expect(calls).toBe(TOO_MANY_CONNECTIONS_RETRIES + 1);
+      expect(refusals.at(-1)).toBe(TOO_MANY_CONNECTIONS_RETRIES + 1);
+      expect(Result.isFailure(result) && isTooManyConnectionsError(result.failure)).toBe(true);
+    }),
+  );
+});
+
+describe("describeRefusedAttempt", () => {
+  it("says when it will retry and when it is giving up", () => {
+    expect(describeRefusedAttempt(1)).toBe(
+      `Postgres refused the connection: no free connection slots (attempt 1 of ${TOO_MANY_CONNECTIONS_RETRIES + 1}); retrying in 30 seconds`,
+    );
+    expect(describeRefusedAttempt(TOO_MANY_CONNECTIONS_RETRIES + 1)).toMatch(/giving up$/);
+  });
+});