feat(db): retry transient connection errors in the Postgres client (#408)

src/db/pg/client.ts

@@ -6,6 +6,7 @@ import {
   getDatabaseProvider,
   getPostgresConnectionString,
 } from "@/db/provider";
+import { withQueryRetries } from "./retry";
 import * as schema from "./schema";
 
 // Postgres on Cloudflare Workers requires a PER-REQUEST client: the runtime
@@ -70,9 +71,15 @@ export async function withPgClient<T>(fn: () => Promise<T>): Promise<T> {
   if (pgClientStore.getStore()) {
     return fn();
   }
-  const sql = postgres(getPostgresConnectionString(), {
+  const sql = withQueryRetries(
-    max: 1,
+    postgres(getPostgresConnectionString(), {
-    fetch_types: false,
+      max: 1,
-  });
+      fetch_types: false,
+      // Bound connect stalls (seconds) so the per-query retry in
+      // withQueryRetries gets its turn within the request's lifetime instead
+      // of hanging on postgres.js's 30s default during a failover.
+      connect_timeout: 10,
+    }),
+  );
   return pgClientStore.run({ sql, db: createPgDb(sql) }, fn);
 }

src/db/pg/retry.test.ts

@@ -0,0 +1,122 @@
+/* oxlint-disable typescript/no-unsafe-type-assertion -- fakes stand in for the postgres.js client */
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import type postgres from "postgres";
+import { withQueryRetries } from "./retry";
+
+type Sql = ReturnType<typeof postgres>;
+
+function connectionError(code: string): Error {
+  return Object.assign(new Error(`${code}: boom`), { code });
+}
+
+/**
+ * Build a fake postgres.js client whose `unsafe` fails `failures` times with
+ * `error` before succeeding. Each unsafe() call returns a fresh thenable
+ * exposing the PendingQuery surface the wrapper covers.
+ */
+function fakeSql(failures: number, error: Error) {
+  let calls = 0;
+  const unsafe = vi.fn(() => {
+    calls++;
+    const outcome =
+      calls <= failures ? Promise.reject(error) : Promise.resolve([{ ok: 1 }]);
+    // Swallow the bare rejection; the wrapper observes it via then/values/etc.
+    void outcome.catch(() => {});
+    return {
+      // oxlint-disable-next-line unicorn/no-thenable -- mirrors postgres.js's awaitable PendingQuery
+      then: (f?: (v: unknown) => unknown, r?: (e: unknown) => unknown) =>
+        outcome.then(f, r),
+      values: () => outcome,
+      raw: () => outcome,
+      execute: () => outcome,
+    };
+  });
+  const sql = { unsafe, options: { parsers: {}, serializers: {} } };
+  return { sql: sql as unknown as Sql, unsafe };
+}
+
+/**
+ * The wrapper's pending query is lazy — the attempt only starts once a
+ * handler is attached. Attach one immediately (so retries are in flight
+ * before the fake timers advance) and silence unhandled-rejection noise.
+ */
+function start<T>(thenable: PromiseLike<T>): Promise<T> {
+  const started = new Promise<T>((resolve, reject) =>
+    thenable.then(resolve, reject),
+  );
+  started.catch(() => {});
+  return started;
+}
+
+beforeEach(() => {
+  vi.useFakeTimers();
+});
+afterEach(() => {
+  vi.useRealTimers();
+});
+
+// Generous enough to flush all retry delays (max ~4.5s including jitter).
+const flushRetries = () => vi.advanceTimersByTimeAsync(10_000);
+
+describe("withQueryRetries", () => {
+  it("retries a select after a mid-flight connection loss", async () => {
+    const { sql, unsafe } = fakeSql(1, connectionError("CONNECTION_CLOSED"));
+    const rows = start(withQueryRetries(sql).unsafe("select 1"));
+    await flushRetries();
+    await expect(rows).resolves.toEqual([{ ok: 1 }]);
+    expect(unsafe).toHaveBeenCalledTimes(2);
+  });
+
+  it("retries the .values() path drizzle uses", async () => {
+    const { sql, unsafe } = fakeSql(1, connectionError("ECONNRESET"));
+    const rows = start(withQueryRetries(sql).unsafe('select "id"').values());
+    await flushRetries();
+    await expect(rows).resolves.toEqual([{ ok: 1 }]);
+    expect(unsafe).toHaveBeenCalledTimes(2);
+  });
+
+  it("does not retry a write after a mid-flight connection loss", async () => {
+    const error = connectionError("CONNECTION_CLOSED");
+    const { sql, unsafe } = fakeSql(1, error);
+    const insert = start(
+      withQueryRetries(sql).unsafe('insert into "session" values ($1)'),
+    );
+    await flushRetries();
+    await expect(insert).rejects.toBe(error);
+    expect(unsafe).toHaveBeenCalledTimes(1);
+  });
+
+  it("retries a write on a connect-phase failure (query never sent)", async () => {
+    const { sql, unsafe } = fakeSql(2, connectionError("CONNECT_TIMEOUT"));
+    const insert = start(
+      withQueryRetries(sql).unsafe('insert into "session" values ($1)'),
+    );
+    await flushRetries();
+    await expect(insert).resolves.toEqual([{ ok: 1 }]);
+    expect(unsafe).toHaveBeenCalledTimes(3);
+  });
+
+  it("does not retry query errors (e.g. unique violations)", async () => {
+    const error = connectionError("23505");
+    const { sql, unsafe } = fakeSql(1, error);
+    const rows = start(withQueryRetries(sql).unsafe("select 1"));
+    await flushRetries();
+    await expect(rows).rejects.toBe(error);
+    expect(unsafe).toHaveBeenCalledTimes(1);
+  });
+
+  it("gives up after exhausting the retry budget", async () => {
+    const error = connectionError("ECONNREFUSED");
+    const { sql, unsafe } = fakeSql(99, error);
+    const rows = start(withQueryRetries(sql).unsafe("select 1"));
+    await flushRetries();
+    await expect(rows).rejects.toBe(error);
+    expect(unsafe).toHaveBeenCalledTimes(4); // initial + 3 retries
+  });
+
+  it("passes everything except unsafe through to the client", () => {
+    const { sql } = fakeSql(0, connectionError("unused"));
+    const wrapped = withQueryRetries(sql);
+    expect(wrapped.options).toBe(sql.options);
+  });
+});

src/db/pg/retry.ts

@@ -0,0 +1,138 @@
+/* oxlint-disable typescript/no-unsafe-type-assertion -- The lazy query stand-in mirrors postgres.js's PendingQuery surface, which has no constructible type. */
+import type postgres from "postgres";
+
+type Sql = ReturnType<typeof postgres>;
+type PendingUnsafe = ReturnType<Sql["unsafe"]>;
+
+/**
+ * Per-query retry for transient connection failures.
+ *
+ * PlanetScale HA operations (cluster resizes, image upgrades, unplanned
+ * failovers) terminate every connection when the primary moves, with the
+ * cutover completing in under ~5 seconds. Hyperdrive reconnects transparently
+ * for new queries, but a query in flight — or one issued before the new
+ * primary accepts connections — fails with a connection error. This wrapper
+ * retries those instead of failing the whole request.
+ *
+ * Safety: a write may have committed even though the connection died before
+ * the response arrived, so writes are only retried on errors that occur
+ * strictly before the query reaches the server (connect-phase failures).
+ * Selects are retried on any transient connection error. Transactions
+ * (`db.transaction` → `sql.begin`) are not retried: postgres.js hands the
+ * transaction callback a raw connection-scoped client that bypasses this
+ * wrapper, and replaying a partially applied transaction is not safe.
+ */
+
+// postgres.js client-side codes, socket errno codes, and Postgres server codes
+// that mean the connection died or the server is briefly unavailable — not
+// that the query is invalid.
+const TRANSIENT_ERROR_CODES = new Set([
+  // postgres.js (src/errors.js)
+  "CONNECTION_CLOSED",
+  "CONNECTION_ENDED",
+  "CONNECTION_DESTROYED",
+  "CONNECT_TIMEOUT",
+  // socket errno
+  "ECONNRESET",
+  "ECONNREFUSED",
+  "EPIPE",
+  "ETIMEDOUT",
+  // Postgres class 08 (connection exception)
+  "08000",
+  "08001",
+  "08003",
+  "08004",
+  "08006",
+  "08P01",
+  // shutdown / failover: admin_shutdown, crash_shutdown, cannot_connect_now
+  "57P01",
+  "57P02",
+  "57P03",
+]);
+
+// Codes that can only occur before the server received the query, so a retry
+// can never double-execute a write.
+const PRE_EXECUTION_CODES = new Set([
+  "CONNECT_TIMEOUT",
+  "ECONNREFUSED",
+  "08001", // sqlclient_unable_to_establish_sqlconnection
+  "08004", // sqlserver_rejected_establishment_of_sqlconnection
+  "57P03", // cannot_connect_now (server starting up / mid-failover)
+]);
+
+// Spaced to ride out a PlanetScale primary cutover (< ~5s total), with jitter
+// so concurrent requests don't reconnect in lockstep.
+const RETRY_DELAYS_MS = [250, 1000, 2500];
+
+function isRetryable(error: unknown, query: string): boolean {
+  const code = (error as { code?: unknown } | null)?.code;
+  if (typeof code !== "string" || !TRANSIENT_ERROR_CODES.has(code)) {
+    return false;
+  }
+  if (PRE_EXECUTION_CODES.has(code)) {
+    return true;
+  }
+  // The connection dropped with the query possibly in flight: the server may
+  // have already executed it, so only replay statements that are safe to run
+  // twice.
+  return /^\s*select\b/i.test(query);
+}
+
+const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms));
+
+/**
+ * Wrap a postgres.js client so `sql.unsafe` — the single entrypoint
+ * drizzle-orm/postgres-js issues queries through — retries transient
+ * connection errors. Everything else on the client passes through untouched.
+ */
+export function withQueryRetries(sql: Sql): Sql {
+  const retryingUnsafe = ((query: string, ...rest: unknown[]) => {
+    const run = () =>
+      (sql.unsafe as (...args: unknown[]) => PendingUnsafe)(query, ...rest);
+
+    const attempt = async <T>(
+      execute: (pending: PendingUnsafe) => Promise<T> | T,
+    ): Promise<T> => {
+      for (let i = 0; ; i++) {
+        try {
+          return await execute(run());
+        } catch (error) {
+          if (i >= RETRY_DELAYS_MS.length || !isRetryable(error, query)) {
+            throw error;
+          }
+          await sleep(RETRY_DELAYS_MS[i] + Math.random() * 250);
+        }
+      }
+    };
+
+    // Lazy stand-in for postgres.js's PendingQuery covering the surface
+    // drizzle-orm/postgres-js uses (await, .values(), .execute(), .raw()).
+    // Each accessor starts its own attempt so a retry re-issues a fresh query
+    // on a fresh connection.
+    const pending = {
+      // oxlint-disable-next-line unicorn/no-thenable -- deliberately thenable: it stands in for postgres.js's PendingQuery, which is awaited directly
+      then: (onFulfilled?: unknown, onRejected?: unknown) =>
+        attempt((q) => q).then(
+          onFulfilled as (value: unknown) => unknown,
+          onRejected as (reason: unknown) => unknown,
+        ),
+      catch: (onRejected?: unknown) =>
+        attempt((q) => q).catch(onRejected as (reason: unknown) => unknown),
+      finally: (onFinally?: unknown) =>
+        attempt((q) => q).finally(onFinally as () => void),
+      values: () => attempt((q) => q.values()),
+      raw: () => attempt((q) => q.raw()),
+      execute: () => attempt((q) => q.execute()),
+    };
+    return pending as unknown as PendingUnsafe;
+  }) as Sql["unsafe"];
+
+  return new Proxy(sql, {
+    get(target, prop, receiver) {
+      if (prop === "unsafe") {
+        return retryingUnsafe;
+      }
+      return Reflect.get(target, prop, receiver) as unknown;
+    },
+  });
+}