refactor(webapp): read waitpoint connected runs via ORM instead of raw SQL

Replaces the raw connected-runs query, which JOINed the waitpoint connection table onto the large TaskRun table, with two indexed ORM reads joined in memory: read the connection rows, then resolve the runs by id. The id lookup can only plan as a primary-key lookup, so the query planner never scans TaskRun, and Prisma qualifies the schema per client so the previous manual schema handling is gone.

Author: d-cs

apps/webapp/app/presenters/v3/WaitpointPresenter.server.ts

@@ -1,9 +1,5 @@
 import { isWaitpointOutputTimeout, prettyPrintPacket } from "@trigger.dev/core/v3";
-import {
-  DATABASE_SCHEMA,
-  type PrismaClientOrTransaction,
-  type PrismaReplicaClient,
-} from "~/db.server";
+import { type PrismaClientOrTransaction, type PrismaReplicaClient } from "~/db.server";
 import { clickhouseFactory } from "~/services/clickhouse/clickhouseFactoryInstance.server";
 import { generateHttpCallbackUrl } from "~/services/httpCallback.server";
 import { logger } from "~/services/logger.server";
@@ -15,10 +11,12 @@ import { waitpointStatusToApiStatus } from "./WaitpointListPresenter.server";
 
 export type WaitpointDetail = NonNullable<Awaited<ReturnType<WaitpointPresenter["call"]>>>;
 
-// Single-sourced bound for connected run friendlyIds: applied at the FETCH in #connectedRunIdsOn,
-// not just at display time.
+// Single-sourced display bound for a waitpoint's connected run friendlyIds.
 export const CONNECTED_RUNS_DISPLAY_LIMIT = 5;
 
+// Over-read connection rows on the FK-free dedicated join so danglers don't cost a display slot.
+export const CONNECTED_RUNS_CONNECTION_SCAN_LIMIT = CONNECTED_RUNS_DISPLAY_LIMIT * 5;
+
 export class WaitpointPresenter extends BasePresenter {
   constructor(
     prisma?: PrismaClientOrTransaction,
@@ -84,8 +82,7 @@ export class WaitpointPresenter extends BasePresenter {
 
   // Connected-run friendlyIds gathered across BOTH stores. The run<->waitpoint join co-locates with
   // the RUN (written on the run's DB), so the waitpoint's own store misses a cross-DB connection; we
-  // read the join on each client and resolve the run's friendlyId on that same client, then union.
-  // We never relation-select `connectedRuns`: it is not a field on the dedicated subset `Waitpoint`.
+  // read the join on each client, resolve the run's friendlyId on that same client, and union.
   async #connectedRunFriendlyIds(waitpointId: string): Promise<string[]> {
     const replica = this._replica as unknown as PrismaReplicaClient;
     const rawClients: PrismaReplicaClient[] =
@@ -99,17 +96,8 @@ export class WaitpointPresenter extends BasePresenter {
 
     const friendlyIds = new Set<string>();
     for (const client of clients) {
-      const runIds = await this.#connectedRunIdsOn(client, waitpointId);
-      if (runIds.length === 0) {
-        continue;
-      }
-      const runs = await client.taskRun.findMany({
-        where: { id: { in: runIds } },
-        select: { friendlyId: true },
-        take: CONNECTED_RUNS_DISPLAY_LIMIT,
-      });
-      for (const run of runs) {
-        friendlyIds.add(run.friendlyId);
+      for (const friendlyId of await this.#connectedRunFriendlyIdsOn(client, waitpointId)) {
+        friendlyIds.add(friendlyId);
       }
       if (friendlyIds.size >= CONNECTED_RUNS_DISPLAY_LIMIT) {
         break;
@@ -118,44 +106,56 @@ export class WaitpointPresenter extends BasePresenter {
     return Array.from(friendlyIds).slice(0, CONNECTED_RUNS_DISPLAY_LIMIT);
   }
 
-  // Schema-aware read of the run ids linked to a waitpoint: the dedicated subset uses the explicit
-  // `WaitpointRunConnection` model (scalar `taskRunId`, no FK -- a row can dangle after its run is
-  // deleted), the control-plane full schema the implicit `_WaitpointRunConnections` M2M
-  // (A = TaskRun.id, B = Waitpoint.id). Both branches existence-filter AT THE QUERY via a JOIN to
-  // TaskRun, so a dangling connection row can never occupy a LIMIT slot ahead of a real one.
-  // Tables are schema-qualified with DATABASE_SCHEMA (trusted boot constant) so a non-`public`
-  // schema= deployment resolves the right tables instead of leaning on search_path. $queryRawUnsafe,
-  // not a `sqlDatabaseSchema` Prisma.Sql fragment: `client` may be the dedicated run-ops client (a
-  // different Prisma runtime) which would bind a foreign runtime's Sql fragment as a param instead
-  // of inlining it. waitpointId and the constant limit stay bound params ($1/$2).
-  async #connectedRunIdsOn(client: PrismaReplicaClient, waitpointId: string): Promise<string[]> {
-    const isDedicated = Boolean(
-      (client as unknown as { waitpointRunConnection?: unknown }).waitpointRunConnection
-    );
-
-    if (isDedicated) {
-      const rows = await client.$queryRawUnsafe<{ taskRunId: string }[]>(
-        `SELECT c."taskRunId" AS "taskRunId"
-        FROM ${DATABASE_SCHEMA}."WaitpointRunConnection" c
-        JOIN ${DATABASE_SCHEMA}."TaskRun" t ON t."id" = c."taskRunId"
-        WHERE c."waitpointId" = $1
-        LIMIT $2`,
-        waitpointId,
-        CONNECTED_RUNS_DISPLAY_LIMIT
-      );
-      return rows.map((row) => row.taskRunId);
+  // Connected-run friendlyIds for one store, via the ORM. Two indexed reads joined in memory instead
+  // of a SQL JOIN onto the (very large) TaskRun table: `id IN (...)` can only plan as a PK lookup, so
+  // the planner can never scan TaskRun.
+  //
+  // Dedicated subset: the explicit `WaitpointRunConnection` is scalar (`taskRunId`, no FK), so a
+  // connection can outlive a deleted run. We over-read connection ids, resolve runs by id (a missing
+  // id just drops out -- danglers cost no display slot), and cap at the display limit.
+  //
+  // Control-plane full schema: no queryable join delegate (implicit M2M), so we traverse the
+  // `connectedRuns` relation; it cascade-deletes with the run, so no dangler can exist.
+  async #connectedRunFriendlyIdsOn(
+    client: PrismaReplicaClient,
+    waitpointId: string
+  ): Promise<string[]> {
+    const dedicated = (
+      client as unknown as {
+        waitpointRunConnection?: {
+          findMany: (args: unknown) => Promise<{ taskRunId: string }[]>;
+        };
+      }
+    ).waitpointRunConnection;
+
+    if (dedicated) {
+      const connections = await dedicated.findMany({
+        where: { waitpointId },
+        select: { taskRunId: true },
+        take: CONNECTED_RUNS_CONNECTION_SCAN_LIMIT,
+      });
+      if (connections.length === 0) {
+        return [];
+      }
+      const runs = await client.taskRun.findMany({
+        where: { id: { in: connections.map((connection) => connection.taskRunId) } },
+        select: { friendlyId: true },
+        take: CONNECTED_RUNS_DISPLAY_LIMIT,
+      });
+      return runs.map((run) => run.friendlyId);
     }
 
-    const rows = await client.$queryRawUnsafe<{ A: string }[]>(
-      `SELECT c."A" AS "A"
-      FROM ${DATABASE_SCHEMA}."_WaitpointRunConnections" c
-      JOIN ${DATABASE_SCHEMA}."TaskRun" t ON t."id" = c."A"
-      WHERE c."B" = $1
-      LIMIT $2`,
-      waitpointId,
-      CONNECTED_RUNS_DISPLAY_LIMIT
-    );
-    return rows.map((row) => row.A);
+    const waitpoint = (await (
+      client.waitpoint.findFirst as (
+        args: unknown
+      ) => Promise<{ connectedRuns: { friendlyId: string }[] } | null>
+    )({
+      where: { id: waitpointId },
+      select: {
+        connectedRuns: { select: { friendlyId: true }, take: CONNECTED_RUNS_DISPLAY_LIMIT },
+      },
+    })) as { connectedRuns: { friendlyId: string }[] } | null;
+    return (waitpoint?.connectedRuns ?? []).map((run) => run.friendlyId);
   }
 
   public async call({

apps/webapp/test/waitpointPresenter.connectedRunsBounded.test.ts

@@ -1,12 +1,11 @@
-// RED->GREEN guard: WaitpointPresenter#connectedRunIdsOn must BOUND the fetch of a waitpoint's
+// RED->GREEN guard: WaitpointPresenter connected-run gather must BOUND the fetch of a waitpoint's
 // connected-run ids, not just the number displayed. A "displayed count <= 5" assertion would
-// false-green on today's buggy code (the existing take:5 already bounds the DISPLAY). Instead
-// this captures the real `taskRun.findMany` call args via a Proxy over a REAL testcontainer
-// Postgres client (no mocks) and asserts the IN-list itself is bounded.
+// false-green (the take:5 on the run resolve already bounds the DISPLAY). Instead this captures the
+// real `taskRun.findMany` call args via a Proxy over a REAL testcontainer Postgres client (no mocks)
+// and asserts the IN-list is capped at the scan limit (danglers over-read), never unbounded.
 //
-// Exercises the dedicated-schema (Prisma `waitpointRunConnection`) branch of #connectedRunIdsOn:
-// waitpoint + 8 connected runs seeded on the NEW dedicated run-ops client (RunOpsPrismaClient,
-// prisma17), which uses the delegate directly (no raw SQL fallback).
+// Exercises the dedicated-schema (Prisma `waitpointRunConnection`) branch: waitpoint + more than the
+// scan-limit connected runs seeded on the NEW dedicated run-ops client (RunOpsPrismaClient, prisma17).
 import { describe, expect, vi } from "vitest";
 
 const legacyReplicaHolder = vi.hoisted(() => ({ client: undefined as any }));
@@ -66,7 +65,7 @@ import { heteroRunOpsPostgresTest } from "@internal/testcontainers";
 import type { PrismaClient } from "@trigger.dev/database";
 import type { RunOpsPrismaClient } from "@internal/run-ops-database";
 import {
-  CONNECTED_RUNS_DISPLAY_LIMIT,
+  CONNECTED_RUNS_CONNECTION_SCAN_LIMIT,
   WaitpointPresenter,
 } from "~/presenters/v3/WaitpointPresenter.server";
 
@@ -184,11 +183,13 @@ function capturingTaskRunFindMany(real: RunOpsPrismaClient): {
   return { client, calls };
 }
 
-const CONNECTED_RUN_COUNT = 8; // > CONNECTED_RUNS_DISPLAY_LIMIT (5), proves the fetch isn't bounded
+// Seed MORE than the scan limit so the assertion bites: an unbounded gather IN-lists all of them,
+// a correctly bounded one caps at CONNECTED_RUNS_CONNECTION_SCAN_LIMIT.
+const CONNECTED_RUN_COUNT = CONNECTED_RUNS_CONNECTION_SCAN_LIMIT + 5;
 
-describe("WaitpointPresenter#connectedRunIdsOn bounds the connected-run-id FETCH", () => {
+describe("WaitpointPresenter bounds the connected-run-id FETCH", () => {
   heteroRunOpsPostgresTest(
-    "reading a waitpoint with 8 connected runs never IN-lists more than the display limit",
+    "a waitpoint with more connections than the scan limit caps the IN-list at the scan limit",
     async ({ prisma14, prisma17 }) => {
       const ctx = await seedParents(prisma14, "bounded");
       const waitpoint = await seedWaitpoint(prisma17, ctx, "waitpoint_bounded");
@@ -220,13 +221,15 @@ describe("WaitpointPresenter#connectedRunIdsOn bounds the connected-run-id FETCH
       });
 
       // The guard: assert the FETCH (the IN-list built from the connected-run-id gather) is
-      // bounded, not just the eventual displayed count. On today's buggy code, the Prisma
-      // `waitpointRunConnection.findMany` branch of #connectedRunIdsOn returns all 8 connected
-      // run ids with no take/LIMIT, so this IN-list is 8. After the fix it must be
-      // <= CONNECTED_RUNS_DISPLAY_LIMIT.
+      // bounded, not just the eventual displayed count. An unbounded gather would IN-list every
+      // connection row; the dedicated branch over-reads up to CONNECTED_RUNS_CONNECTION_SCAN_LIMIT
+      // (25) so a display slot is never lost to a dangler, so the IN-list must be capped at the
+      // scan limit -- above the display limit (5), but never unbounded.
       expect(calls.length).toBeGreaterThan(0);
       for (const call of calls) {
-        expect(call.where?.id?.in?.length ?? 0).toBeLessThanOrEqual(CONNECTED_RUNS_DISPLAY_LIMIT);
+        expect(call.where?.id?.in?.length ?? 0).toBeLessThanOrEqual(
+          CONNECTED_RUNS_CONNECTION_SCAN_LIMIT
+        );
       }
     }
   );

apps/webapp/test/waitpointPresenter.readthrough.test.ts

@@ -68,11 +68,17 @@ import { setupClickhouseReplication } from "./utils/replicationUtils";
 
 vi.setConfig({ testTimeout: 90_000 });
 
-// A read client whose waitpoint.findFirst is recorded; throws if used after being marked
-// forbidden, so we can prove a store was NEVER read. Every other access forwards to the real
-// client (so the inlined `environment` join + connectedRuns relation still resolve).
+// A read client whose waitpoint.findFirst calls are recorded, split by kind. The presenter issues
+// two shapes of waitpoint.findFirst: a HYDRATE (loads the waitpoint scalar row -- no `connectedRuns`
+// in its select) and a connectedRuns-RELATION read (control-plane branch of the connected-runs
+// gather -- select is exactly `{ connectedRuns }`). The `forbidden` guard throws ONLY on a HYDRATE,
+// so a store can be proven to never be HYDRATED while still allowing the connectedRuns cross-DB
+// union to legitimately read it (that union always touched both stores; it was invisible when the
+// old code did it via raw SQL). Every other access forwards to the real client.
 function recording(client: PrismaClient, opts: { forbidden?: boolean } = {}) {
   const calls: unknown[] = [];
+  const hydrateCalls: unknown[] = [];
+  const connectedRunsCalls: unknown[] = [];
   return {
     handle: new Proxy(client, {
       get(target, prop) {
@@ -82,8 +88,16 @@ function recording(client: PrismaClient, opts: { forbidden?: boolean } = {}) {
               if (wpProp === "findFirst") {
                 return (args: unknown) => {
                   calls.push(args);
-                  if (opts.forbidden) {
-                    throw new Error("this store must never be read");
+                  const select = (args as { select?: Record<string, unknown> })?.select ?? {};
+                  const keys = Object.keys(select);
+                  const isConnectedRunsRead = keys.length === 1 && keys[0] === "connectedRuns";
+                  if (isConnectedRunsRead) {
+                    connectedRunsCalls.push(args);
+                  } else {
+                    hydrateCalls.push(args);
+                    if (opts.forbidden) {
+                      throw new Error("this store must never be hydrated");
+                    }
                   }
                   return (wpTarget as any).findFirst(args);
                 };
@@ -96,6 +110,8 @@ function recording(client: PrismaClient, opts: { forbidden?: boolean } = {}) {
       },
     }) as unknown as PrismaReplicaClient,
     calls,
+    hydrateCalls,
+    connectedRunsCalls,
   };
 }
 
@@ -255,9 +271,12 @@ describe("WaitpointPresenter dual-DB read-through (hetero PG14 + PG17, no connec
       const result = await presenter.call(callArgs(ctx, seeded.friendlyId));
 
       expect(result?.id).toBe(seeded.friendlyId);
-      // New-first short-circuit: legacy never probed (the throwing handle proves it).
-      expect(newClient.calls.length).toBe(1);
-      expect(legacy.calls.length).toBe(0);
+      // New-first short-circuit: the HYDRATE answered from new and never fell through to a legacy
+      // hydrate (the throwing handle proves it). The connectedRuns cross-DB union still reads legacy
+      // legitimately -- that read is allowed and must NOT count as a hydrate.
+      expect(newClient.hydrateCalls.length).toBe(1);
+      expect(legacy.hydrateCalls.length).toBe(0);
+      expect(legacy.connectedRunsCalls.length).toBeGreaterThan(0);
     }
   );
 
@@ -286,8 +305,11 @@ describe("WaitpointPresenter dual-DB read-through (hetero PG14 + PG17, no connec
 
       expect(result?.id).toBe(seeded.friendlyId);
       expect(result?.tags).toEqual(["one"]);
-      // Exactly one read on the single client; the second handle is never touched.
-      expect(single.calls.length).toBe(1);
+      // Two reads on the single client, both on the one handle: the first findFirst hydrates the
+      // waitpoint, the second loads the `connectedRuns` relation (the implicit M2M has no queryable
+      // join delegate, so the ORM must traverse the relation with a second findFirst). The second
+      // handle is never touched -- passthrough is structural, the split branch never fires.
+      expect(single.calls.length).toBe(2);
       expect(second.calls.length).toBe(0);
     }
   );