From 278b19836d38817c295a1d85fb5cd8dde036e7f9 Mon Sep 17 00:00:00 2001 From: Ben Senescu <44480372+bensenescu@users.noreply.github.com> Date: Mon, 20 Jul 2026 18:16:17 -0400 Subject: [PATCH] feat(db): retry transient connection errors in the Postgres client (#408) --- src/db/pg/client.ts | 15 +++-- src/db/pg/retry.test.ts | 122 +++++++++++++++++++++++++++++++++++ src/db/pg/retry.ts | 138 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 271 insertions(+), 4 deletions(-) create mode 100644 src/db/pg/retry.test.ts create mode 100644 src/db/pg/retry.ts diff --git a/src/db/pg/client.ts b/src/db/pg/client.ts index d3d3223..b155e65 100644 --- a/src/db/pg/client.ts +++ b/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(fn: () => Promise): Promise { if (pgClientStore.getStore()) { return fn(); } - const sql = postgres(getPostgresConnectionString(), { - max: 1, - fetch_types: false, - }); + const sql = withQueryRetries( + postgres(getPostgresConnectionString(), { + 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); } diff --git a/src/db/pg/retry.test.ts b/src/db/pg/retry.test.ts new file mode 100644 index 0000000..5b00c7f --- /dev/null +++ b/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; + +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(thenable: PromiseLike): Promise { + const started = new Promise((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); + }); +}); diff --git a/src/db/pg/retry.ts b/src/db/pg/retry.ts new file mode 100644 index 0000000..be4dae4 --- /dev/null +++ b/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; +type PendingUnsafe = ReturnType; + +/** + * 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 ( + execute: (pending: PendingUnsafe) => Promise | T, + ): Promise => { + 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; + }, + }); +}