test(testcontainers): add a DB connection-blip harness

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,177 @@
+import { describe, 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 });
+  // A severed idle connection makes the pg Pool emit 'error'; swallow it so an
+  // unhandled event can't crash the test worker before recovery is asserted.
+  pool.on("error", () => {});
+  const client = new PrismaClient({ adapter: new PrismaPg(pool) });
+  const dispose = async () => {
+    try {
+      await client.$disconnect();
+    } finally {
+      await pool.end();
+    }
+  };
+  return { client, dispose };
+}
+
+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;
+}
+
+describe("DbBlipController", () => {
+  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();
+      }
+    }
+  );
+
+  // Regression: queryContains must match as literal text, not as an ILIKE pattern.
+  // The active query contains "fooXbar"; under ILIKE the pattern "foo_bar" (with the
+  // wildcard `_`) would wrongly match and terminate it. The literal matcher must not,
+  // so the sever times out instead of killing the wrong statement.
+  postgresBlipTest(
+    "severDuringNextStatement matches queryContains literally, not as an ILIKE pattern",
+    { timeout: 60_000 },
+    async ({ postgresContainer, blip }) => {
+      const { client, dispose } = await adapterClient(postgresContainer.getConnectionUri());
+      const slow = client
+        .$queryRawUnsafe(`SELECT pg_sleep(3) /* marker fooXbar */`)
+        .catch(() => undefined);
+      try {
+        await expect(
+          blip.severDuringNextStatement({ queryContains: "foo_bar", timeoutMs: 1000, pollMs: 25 })
+        ).rejects.toThrow(/no active statement/i);
+      } finally {
+        await dispose();
+        await slow;
+      }
+    }
+  );
+});

internal-packages/testcontainers/src/dbBlip.ts

@@ -0,0 +1,107 @@
+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 idle 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 client statement (optionally matching `queryContains`
+   * literally), then terminate it mid-flight. Rejects if none appears within
+   * `timeoutMs`. Terminating by pid isn't atomic with statement completion, so
+   * target a statement with a real execution window (e.g. `pg_sleep`) — a query
+   * that finishes first leaves its connection idle and it is closed anyway. */
+  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> };
+
+// Reserved application_name for the harness's control connections. The severs
+// exclude every connection using it (by name, plus their own pid), so multiple
+// controllers on one database can't kill each other's admin. A client-under-test
+// must not use this name.
+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<{ terminated: boolean }>(
+      `SELECT pg_terminate_backend(pid) AS terminated
+         FROM pg_stat_activity
+        WHERE datname = current_database()
+          AND pid <> pg_backend_pid()
+          AND backend_type = 'client backend'
+          AND application_name IS DISTINCT FROM $1
+          AND state = 'idle'`,
+      [ADMIN_APPLICATION_NAME]
+    );
+    // Count only backends that were actually terminated (a backend that exits
+    // between selection and signalling returns false).
+    return result.rows.filter((row) => row.terminated === true).length;
+  }
+
+  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) {
+      // Select and terminate in one statement so the backend can't go idle
+      // between picking it and killing it; return only when it was terminated.
+      const terminated = await admin.query<{ ok: boolean }>(
+        `SELECT pg_terminate_backend(pid) AS ok
+           FROM pg_stat_activity
+          WHERE datname = current_database()
+            AND state = 'active'
+            AND pid <> pg_backend_pid()
+            AND backend_type = 'client backend'
+            AND application_name IS DISTINCT FROM $1
+            AND ($2::text IS NULL OR strpos(lower(query), lower($2)) > 0)
+          LIMIT 1`,
+        [ADMIN_APPLICATION_NAME, queryContains]
+      );
+
+      if (terminated.rows[0]?.ok === true) {
+        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;