test(testcontainers): add DB blip harness for connection-resilience tests

Adds DbBlipController + a postgresBlipTest fixture that sever a test
Postgres connection via pg_terminate_backend, so a vertical can prove
its DB code survives a disconnect without a proxy or extra container.

Includes tests that exercise the harness: a severed read fails, an
in-flight statement is terminated mid-flight, a non-idempotent write
double-applies on retry while the idempotent form does not, and a
pg-driver-adapter client recovers a model read through a blip.

TRI-13551

Author: d-cs

internal-packages/testcontainers/package.json

@@ -15,12 +15,15 @@
   "dependencies": {
     "@clickhouse/client": "^1.11.1",
     "@trigger.dev/database": "workspace:*",
-    "ioredis": "~5.6.0"
+    "ioredis": "~5.6.0",
+    "pg": "8.15.6"
   },
   "devDependencies": {
     "@internal/run-ops-database": "workspace:*",
+    "@prisma/adapter-pg": "6.14.0",
     "@testcontainers/postgresql": "^11.14.0",
     "@testcontainers/redis": "^11.14.0",
+    "@types/pg": "8.11.14",
     "std-env": "^3.9.0",
     "testcontainers": "^11.14.0",
     "tinyexec": "^0.3.0"

internal-packages/testcontainers/src/dbBlip.test.ts

@@ -0,0 +1,142 @@
+import { expect } from "vitest";
+import { Pool } from "pg";
+import { PrismaPg } from "@prisma/adapter-pg";
+import { PrismaClient } from "@trigger.dev/database";
+import { postgresBlipTest } from "./index";
+
+// A minimal infra retry, standing in for the shared read-retry util so this
+// file can demonstrate the harness end-to-end on its own.
+async function withRetry<T>(fn: () => Promise<T>, maxAttempts = 8): Promise<T> {
+  let lastError: unknown;
+  for (let attempt = 0; attempt < maxAttempts; attempt++) {
+    try {
+      return await fn();
+    } catch (error) {
+      lastError = error;
+      await new Promise((r) => setTimeout(r, Math.min(50 * (attempt + 1), 250)));
+    }
+  }
+  throw lastError;
+}
+
+// Production runs the pg driver adapter, so the client under test is adapter-backed.
+async function adapterClient(connectionString: string) {
+  const pool = new Pool({ connectionString });
+  const client = new PrismaClient({ adapter: new PrismaPg(pool) });
+  return { client, dispose: async () => void (await client.$disconnect(), await pool.end()) };
+}
+
+async function createProbeTable(client: PrismaClient) {
+  await client.$executeRawUnsafe(
+    `CREATE TABLE IF NOT EXISTS blip_probe (id uuid PRIMARY KEY, tag text NOT NULL)`
+  );
+}
+
+async function countTag(client: PrismaClient, tag: string): Promise<number> {
+  const rows = await client.$queryRawUnsafe<{ n: number }[]>(
+    `SELECT count(*)::int AS n FROM blip_probe WHERE tag = $1`,
+    tag
+  );
+  return rows[0]?.n ?? 0;
+}
+
+postgresBlipTest(
+  "a pooled adapter client transparently survives an idle-connection drop",
+  { timeout: 60_000 },
+  async ({ postgresContainer, blip }) => {
+    const { client, dispose } = await adapterClient(postgresContainer.getConnectionUri());
+    try {
+      await client.user.count(); // warm the pool
+      const terminated = await blip.severIdle();
+      expect(terminated).toBeGreaterThan(0);
+      // The pool evicts the dead idle connection; the next read just works.
+      await new Promise((r) => setTimeout(r, 200));
+      const count = await client.user.count();
+      expect(typeof count).toBe("number");
+    } finally {
+      await dispose();
+    }
+  }
+);
+
+postgresBlipTest(
+  "severDuringNextStatement fails an in-flight statement",
+  { timeout: 60_000 },
+  async ({ postgresContainer, blip }) => {
+    const { client, dispose } = await adapterClient(postgresContainer.getConnectionUri());
+    try {
+      const slow = client.$queryRawUnsafe(`SELECT pg_sleep(3)`);
+      // PrismaPromise is lazy — form the assertion so the query actually starts.
+      const rejected = expect(slow).rejects.toThrow();
+      await blip.severDuringNextStatement({ queryContains: "pg_sleep" });
+      await rejected;
+    } finally {
+      await dispose();
+    }
+  }
+);
+
+postgresBlipTest(
+  "a read recovers after a mid-flight blip",
+  { timeout: 60_000 },
+  async ({ postgresContainer, blip }) => {
+    const { client, dispose } = await adapterClient(postgresContainer.getConnectionUri());
+    try {
+      const severed = client.$queryRawUnsafe(`SELECT pg_sleep(3)`).catch(() => undefined);
+      await blip.severDuringNextStatement({ queryContains: "pg_sleep" });
+      await severed;
+      const count = await withRetry(() => client.user.count());
+      expect(typeof count).toBe("number");
+    } finally {
+      await dispose();
+    }
+  }
+);
+
+postgresBlipTest(
+  "a non-idempotent write double-applies on retry after a post-commit blip; the idempotent form does not",
+  { timeout: 60_000 },
+  async ({ postgresContainer, blip }) => {
+    const { client, dispose } = await adapterClient(postgresContainer.getConnectionUri());
+    try {
+      await createProbeTable(client);
+
+      // Model the dangerous case: the write commits, then a later statement in
+      // the same op is severed mid-flight (ack lost), and the caller retries.
+      let nonIdempotentAttempts = 0;
+      const nonIdempotentWrite = async () => {
+        nonIdempotentAttempts++;
+        await client.$executeRawUnsafe(
+          `INSERT INTO blip_probe (id, tag) VALUES (gen_random_uuid(), 'non-idempotent')`
+        );
+        if (nonIdempotentAttempts === 1) {
+          await client.$queryRawUnsafe(`SELECT pg_sleep(3)`); // severed → throws after the commit
+        }
+      };
+      const nonIdempotentDone = withRetry(nonIdempotentWrite);
+      await blip.severDuringNextStatement({ queryContains: "pg_sleep" });
+      await nonIdempotentDone;
+      expect(await countTag(client, "non-idempotent")).toBe(2); // the hazard, proven
+
+      // The idempotent form: a fixed id + ON CONFLICT makes the replay a no-op.
+      let idempotentAttempts = 0;
+      const idempotentWrite = async () => {
+        idempotentAttempts++;
+        await client.$executeRawUnsafe(
+          `INSERT INTO blip_probe (id, tag)
+           VALUES ('00000000-0000-0000-0000-000000000001', 'idempotent')
+           ON CONFLICT (id) DO NOTHING`
+        );
+        if (idempotentAttempts === 1) {
+          await client.$queryRawUnsafe(`SELECT pg_sleep(3)`);
+        }
+      };
+      const idempotentDone = withRetry(idempotentWrite);
+      await blip.severDuringNextStatement({ queryContains: "pg_sleep" });
+      await idempotentDone;
+      expect(await countTag(client, "idempotent")).toBe(1); // exactly once despite retry
+    } finally {
+      await dispose();
+    }
+  }
+);

internal-packages/testcontainers/src/dbBlip.ts

@@ -0,0 +1,96 @@
+import { Client } from "pg";
+
+/**
+ * Simulates a connection blip against a test Postgres (via a separate admin
+ * connection that terminates backends), so a vertical can prove its DB code
+ * survives a disconnect. Reproduces the mid-statement / stale-connection
+ * signatures (P1017, "Connection terminated unexpectedly").
+ */
+export type DbBlipController = {
+  /** Terminate every client backend except this harness's own, so the next
+   * operation hits a dead connection. Returns the number terminated. */
+  severIdle(): Promise<number>;
+
+  /** Poll for an active statement (optionally matching `queryContains`), then
+   * terminate it mid-flight. Rejects if none appears within `timeoutMs`. */
+  severDuringNextStatement(opts?: {
+    queryContains?: string;
+    timeoutMs?: number;
+    pollMs?: number;
+  }): Promise<void>;
+};
+
+/** A {@link DbBlipController} plus the teardown for its admin connection. */
+export type DbBlipHandle = DbBlipController & { close(): Promise<void> };
+
+// Tag the admin connection so we can exclude it from the backends we kill.
+const ADMIN_APPLICATION_NAME = "trigger-db-blip-admin";
+
+const sleep = (ms: number) => new Promise<void>((resolve) => setTimeout(resolve, ms));
+
+/** Opens an isolated admin connection and returns a handle that can sever the
+ * other connections on that database. `close()` in teardown. */
+export async function createDbBlipController(connectionUri: string): Promise<DbBlipHandle> {
+  // Raw pg (not Prisma): the control connection must be one identifiable backend we can exclude from the sever, independent of the client under test.
+  const admin = new Client({
+    connectionString: connectionUri,
+    application_name: ADMIN_APPLICATION_NAME,
+  });
+  await admin.connect();
+  // Swallow async connection errors so a consumer that severs a DB the admin
+  // isn't excluded from (or drops it while open) can't crash the test worker.
+  admin.on("error", () => {});
+
+  async function severIdle(): Promise<number> {
+    const result = await admin.query(
+      `SELECT pg_terminate_backend(pid)
+         FROM pg_stat_activity
+        WHERE datname = current_database()
+          AND pid <> pg_backend_pid()
+          AND application_name IS DISTINCT FROM $1`,
+      [ADMIN_APPLICATION_NAME]
+    );
+    return result.rowCount ?? 0;
+  }
+
+  async function severDuringNextStatement(opts?: {
+    queryContains?: string;
+    timeoutMs?: number;
+    pollMs?: number;
+  }): Promise<void> {
+    const queryContains = opts?.queryContains ?? null;
+    const timeoutMs = opts?.timeoutMs ?? 5000;
+    const pollMs = opts?.pollMs ?? 25;
+    const deadline = Date.now() + timeoutMs;
+
+    while (Date.now() < deadline) {
+      const found = await admin.query<{ pid: number }>(
+        `SELECT pid
+           FROM pg_stat_activity
+          WHERE datname = current_database()
+            AND state = 'active'
+            AND pid <> pg_backend_pid()
+            AND application_name IS DISTINCT FROM $1
+            AND ($2::text IS NULL OR query ILIKE '%' || $2 || '%')
+          LIMIT 1`,
+        [ADMIN_APPLICATION_NAME, queryContains]
+      );
+
+      const pid = found.rows[0]?.pid;
+      if (pid !== undefined) {
+        await admin.query("SELECT pg_terminate_backend($1)", [pid]);
+        return;
+      }
+
+      await sleep(pollMs);
+    }
+
+    throw new Error(
+      `severDuringNextStatement: no active statement${
+        queryContains ? ` matching ${JSON.stringify(queryContains)}` : ""
+      } appeared within ${timeoutMs}ms`
+    );
+  }
+
+  return { severIdle, severDuringNextStatement, close: () => admin.end() };
+}

internal-packages/testcontainers/src/index.ts

@@ -13,6 +13,7 @@ import {
   runClickhouseMigrations,
   truncateClickhouseTables,
 } from "./clickhouse";
+import { createDbBlipController, type DbBlipController } from "./dbBlip";
 import { getTaskMetadata, logCleanup, logSetup } from "./logs";
 import { type MinIOConnectionConfig, type StartedMinIOContainer, MinIOContainer } from "./minio";
 import {
@@ -35,6 +36,7 @@ export {
 } from "./utils";
 export { OtelCollectorContainer, StartedOtelCollectorContainer } from "./otelCollector";
 export { laggingReplica, type LaggingModel } from "./laggingReplica";
+export { createDbBlipController, type DbBlipController, type DbBlipHandle } from "./dbBlip";
 export { logCleanup };
 export type { MinIOConnectionConfig };
 
@@ -353,6 +355,32 @@ export const postgresTest = withWarmup(
   }
 );
 
+export type PostgresBlipTestContext = PostgresTestContext & { blip: DbBlipController };
+
+const blipFromContainer = async (
+  { postgresContainer }: { postgresContainer: StartedPostgreSqlContainer } & TestContext,
+  use: Use<DbBlipController>
+) => {
+  const handle = await createDbBlipController(postgresContainer.getConnectionUri());
+  try {
+    await use(handle);
+  } finally {
+    await handle.close();
+  }
+};
+
+// postgresTest + a DbBlipController bound to the same per-test database.
+export const postgresBlipTest = withWarmup(
+  test.extend<PostgresBlipTestContext>({
+    postgresContainer: clonedPostgresContainer,
+    prisma: prismaFromContainer,
+    blip: blipFromContainer,
+  }),
+  async () => {
+    await getWorkerPostgresContainer();
+  }
+);
+
 type HeteroPostgresTestContext = {
   // PG14 (legacy / control-plane DB analog)
   postgresContainer14: StartedPostgreSqlContainer;