From c64242702fe9ff993dd6d862105d4faef1aed23c Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Thu, 9 Jul 2026 15:14:55 +0100 Subject: [PATCH 1/2] perf(webapp,run-store): group per-item run reads into batched queries Several run-ops read paths issued one database query per item where a single batched query returns the same data: batch results hydrated each member run separately, the run retrieve endpoint resolved each locked worker version separately, the dedicated-schema relation hydrator ran per row, and a few services read runs in per-id loops. These now group into batched reads. Adds a grouped findRunsByIds to the run store (one residency-partitioned read for a set of run ids) and routes the presenters and services through it. The dedicated-schema relation hydrator batches across all rows per relation. The waitpoint connected-runs read also regains its display limit at the fetch, so it no longer loads every connected run into memory before slicing to five. --- .../v3/ApiBatchResultsPresenter.server.ts | 75 ++-- .../v3/ApiRetrieveRunPresenter.server.ts | 16 +- .../v3/WaitpointPresenter.server.ts | 13 +- .../controlPlaneResolver.server.test.ts | 168 ++++++++ .../controlPlaneResolver.server.ts | 90 ++++ .../services/cancelDevSessionRuns.server.ts | 211 ++++++--- ...atchResultsPresenter.dedicatedSeam.test.ts | 337 +++++++++++++++ ...iBatchResultsPresenter.splitNPlus1.test.ts | 251 +++++++++++ ...veRunPresenter.groupedLockedWorker.test.ts | 266 ++++++++++++ .../cancelDevSessionRunsGroupedReads.test.ts | 183 ++++++++ ...ointPresenter.connectedRunsBounded.test.ts | 232 ++++++++++ ...tpointPresenter.splitConnectedRuns.test.ts | 207 +++++++++ .../tests/legacyRunNewTokenResumeFlow.test.ts | 302 +++++++++++++ .../PostgresRunStore.findRunsByIds.test.ts | 170 ++++++++ .../PostgresRunStore.groupedRelations.test.ts | 228 ++++++++++ .../run-store/src/PostgresRunStore.ts | 401 +++++++++++------- .../src/runOpsStore.groupedReads.test.ts | 388 +++++++++++++++++ .../run-store/src/runOpsStore.ts | 84 +++- .../run-store/src/scheduleSeam.test.ts | 306 +++++++++++++ internal-packages/run-store/src/types.ts | 15 + 20 files changed, 3651 insertions(+), 292 deletions(-) create mode 100644 apps/webapp/app/v3/runOpsMigration/controlPlaneResolver.server.test.ts create mode 100644 apps/webapp/test/apiBatchResultsPresenter.dedicatedSeam.test.ts create mode 100644 apps/webapp/test/apiBatchResultsPresenter.splitNPlus1.test.ts create mode 100644 apps/webapp/test/apiRetrieveRunPresenter.groupedLockedWorker.test.ts create mode 100644 apps/webapp/test/cancelDevSessionRunsGroupedReads.test.ts create mode 100644 apps/webapp/test/waitpointPresenter.connectedRunsBounded.test.ts create mode 100644 apps/webapp/test/waitpointPresenter.splitConnectedRuns.test.ts create mode 100644 internal-packages/run-engine/src/engine/tests/legacyRunNewTokenResumeFlow.test.ts create mode 100644 internal-packages/run-store/src/PostgresRunStore.findRunsByIds.test.ts create mode 100644 internal-packages/run-store/src/PostgresRunStore.groupedRelations.test.ts create mode 100644 internal-packages/run-store/src/runOpsStore.groupedReads.test.ts create mode 100644 internal-packages/run-store/src/scheduleSeam.test.ts diff --git a/apps/webapp/app/presenters/v3/ApiBatchResultsPresenter.server.ts b/apps/webapp/app/presenters/v3/ApiBatchResultsPresenter.server.ts index 0e04e909eb..c9179d5912 100644 --- a/apps/webapp/app/presenters/v3/ApiBatchResultsPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/ApiBatchResultsPresenter.server.ts @@ -1,4 +1,5 @@ import type { BatchTaskRunExecutionResult } from "@trigger.dev/core/v3"; +import { ownerEngine } from "@trigger.dev/core/v3/isomorphic"; import { $replica, type PrismaClientOrTransaction, @@ -8,7 +9,6 @@ import { import type { TaskRunWithAttempts } from "~/models/taskRun.server"; import { executionResultForTaskRun } from "~/models/taskRun.server"; import type { AuthenticatedEnvironment } from "~/services/apiAuth.server"; -import { readThroughRun } from "~/v3/runOpsMigration/readThrough.server"; import { runStore as defaultRunStore } from "~/v3/runStore.server"; import { BasePresenter } from "./basePresenter.server"; @@ -43,11 +43,14 @@ const memberRunSelect = { } as const; /** - * Split on: the batch row + its item rows resolve new-run-ops first, then the LEGACY RUN-OPS - * READ REPLICA ONLY (never the legacy primary — there is no such handle); each member run is - * hydrated independently via readThroughRun keyed on the member runId, so a batch whose members - * span migrated + abandoned runs returns the complete reachable set (the batch-spanning-the-line - * read; the dangling-reference termination gate is a separate, adjacent unit). + * Split on: the batch row + its members resolve new-run-ops first, then the LEGACY RUN-OPS READ + * REPLICA ONLY (never the legacy primary — there is no such handle). Members hydrate via ONE + * grouped read against `newClient` for the whole id set, then ONE grouped read against + * `legacyReplica` for just the misses that could still be legacy-resident — the same + * residency-partitioned shape as the old per-member read-through, but batched instead of fanned + * out one query per member. A batch whose members span migrated + abandoned runs returns the + * complete reachable set (the batch-spanning-the-line read; the dangling-reference termination + * gate is a separate, adjacent unit). * * Split off (single-DB / self-host): one passthrough read for the batch row + a single store * id-set hydrate for the members — no legacy read, no known-migrated probe, no second connection. @@ -133,7 +136,7 @@ export class ApiBatchResultsPresenter extends BasePresenter { // Split: resolve the batch row new-first then off the legacy READ REPLICA only (a batch id may // be cuid or run-ops id, and a cuid-shaped id can still have been backfilled onto NEW, so id-shape // residency is not authoritative for the row — the new-first-then-legacy probe is), then - // hydrate every member run independently via the per-run read-through primitive. + // hydrate every member run in ONE grouped new-then-legacy read. async #callSplit( friendlyId: string, env: AuthenticatedEnvironment @@ -175,41 +178,35 @@ export class ApiBatchResultsPresenter extends BasePresenter { }; } - const readMemberRun = (client: PrismaClientOrTransaction, taskRunId: string) => - client.taskRun.findFirst({ - where: { id: taskRunId }, - select: memberRunSelect, - }) as Promise; - - // Per-member fan-out: each member may live on a different DB, so a single nested include cannot - // cross the seam. Promise.all preserves batchRun.items order, unchanged from today. - const memberResults = await Promise.all( - batchRun.items.map(async (item) => { - const result = await readThroughRun({ - runId: item.taskRunId, - environmentId: env.id, - readNew: (client) => readMemberRun(client, item.taskRunId), - readLegacy: (replica) => readMemberRun(replica, item.taskRunId), - deps: { - splitEnabled: true, - // Pass the SAME resolved handles the batch row used, so the batch row and its members - // never resolve against different DBs. (Letting these fall through to readThroughRun's - // own module-level defaults would diverge from the batch read's `?? this._replica`.) - newClient, - legacyReplica, - isPastRetention: this.readThrough?.isPastRetention, - }, - }); + const taskRunIds = batchRun.items.map((item) => item.taskRunId); - // not-found / past-retention members are omitted (matches today's drop-undefined behavior); - // the dangling-reference termination gate (separate unit) governs whether that's permitted. - if (result.source === "not-found" || result.source === "past-retention") { - return undefined; - } + const newRows = (await newClient.taskRun.findMany({ + where: { id: { in: taskRunIds } }, + select: memberRunSelect, + })) as TaskRunWithAttempts[]; + const runsById = new Map(newRows.map((run) => [run.id, run])); - return executionResultForTaskRun(result.value); - }) + // A run-ops id can only live on NEW, so only misses that AREN'T run-ops-shaped are candidates + // for the legacy probe — mirrors readThroughRun's per-id "NEW residency skips legacy" rule. + const legacyCandidateIds = taskRunIds.filter( + (id) => !runsById.has(id) && ownerEngine(id) !== "NEW" ); + if (legacyCandidateIds.length > 0) { + const legacyRows = (await legacyReplica.taskRun.findMany({ + where: { id: { in: legacyCandidateIds } }, + select: memberRunSelect, + })) as TaskRunWithAttempts[]; + for (const run of legacyRows) { + runsById.set(run.id, run); + } + } + + // not-found members are omitted (matches today's drop-undefined behavior); the + // dangling-reference termination gate (separate unit) governs whether that's permitted. + const memberResults = batchRun.items.map((item) => { + const run = runsById.get(item.taskRunId); + return run ? executionResultForTaskRun(run) : undefined; + }); return { id: batchRun.friendlyId, diff --git a/apps/webapp/app/presenters/v3/ApiRetrieveRunPresenter.server.ts b/apps/webapp/app/presenters/v3/ApiRetrieveRunPresenter.server.ts index a5c71a5555..c6a14a4492 100644 --- a/apps/webapp/app/presenters/v3/ApiRetrieveRunPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/ApiRetrieveRunPresenter.server.ts @@ -164,17 +164,13 @@ export class ApiRetrieveRunPresenter { collect(child.lockedToVersionId); } + const lockedWorkersByVersionId = + await controlPlaneResolver.resolveRunLockedWorkersByVersionIds([...distinctVersionIds]); const versionById = new Map( - await Promise.all( - [...distinctVersionIds].map( - async (id) => - [ - id, - (await controlPlaneResolver.resolveRunLockedWorker({ lockedToVersionId: id })) - ?.lockedToVersion?.version ?? null, - ] as const - ) - ) + [...distinctVersionIds].map((id) => [ + id, + lockedWorkersByVersionId.get(id)?.lockedToVersion?.version ?? null, + ]) ); const resolveVersion = (id: string | null): { version: string } | null => { diff --git a/apps/webapp/app/presenters/v3/WaitpointPresenter.server.ts b/apps/webapp/app/presenters/v3/WaitpointPresenter.server.ts index 74f0c06fdd..c7c4e67605 100644 --- a/apps/webapp/app/presenters/v3/WaitpointPresenter.server.ts +++ b/apps/webapp/app/presenters/v3/WaitpointPresenter.server.ts @@ -11,6 +11,10 @@ import { waitpointStatusToApiStatus } from "./WaitpointListPresenter.server"; export type WaitpointDetail = NonNullable>>; +// Single-sourced bound for connected run friendlyIds: applied at the FETCH in #connectedRunIdsOn, +// not just at display time. +export const CONNECTED_RUNS_DISPLAY_LIMIT = 5; + export class WaitpointPresenter extends BasePresenter { constructor( prisma?: PrismaClientOrTransaction, @@ -98,16 +102,16 @@ export class WaitpointPresenter extends BasePresenter { const runs = await client.taskRun.findMany({ where: { id: { in: runIds } }, select: { friendlyId: true }, - take: 5, + take: CONNECTED_RUNS_DISPLAY_LIMIT, }); for (const run of runs) { friendlyIds.add(run.friendlyId); } - if (friendlyIds.size >= 5) { + if (friendlyIds.size >= CONNECTED_RUNS_DISPLAY_LIMIT) { break; } } - return Array.from(friendlyIds).slice(0, 5); + 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 @@ -125,11 +129,14 @@ export class WaitpointPresenter extends BasePresenter { const links = await joinDelegate.findMany({ where: { waitpointId }, select: { taskRunId: true }, + take: CONNECTED_RUNS_DISPLAY_LIMIT, }); return links.map((link) => link.taskRunId); } + // CONNECTED_RUNS_DISPLAY_LIMIT is a compile-time constant int, not an interpolated value. const rows = await client.$queryRaw<{ A: string }[]>` SELECT "A" FROM "_WaitpointRunConnections" WHERE "B" = ${waitpointId} + LIMIT ${CONNECTED_RUNS_DISPLAY_LIMIT} `; return rows.map((row) => row.A); } diff --git a/apps/webapp/app/v3/runOpsMigration/controlPlaneResolver.server.test.ts b/apps/webapp/app/v3/runOpsMigration/controlPlaneResolver.server.test.ts new file mode 100644 index 0000000000..cd08c8a077 --- /dev/null +++ b/apps/webapp/app/v3/runOpsMigration/controlPlaneResolver.server.test.ts @@ -0,0 +1,168 @@ +// Real control-plane (single Postgres) proof that resolving the locked-worker version for many +// runs' distinct `lockedToVersionId`s is a GROUPED query, not one `backgroundWorker.findFirst` +// per id. The DB is never mocked: the call-counting proxy below delegates every call to the +// real Prisma client (the real query still runs against the real container) - it only tallies +// how many times each model.method pair was invoked. +import { postgresTest } from "@internal/testcontainers"; +import { describe, expect } from "vitest"; +import type { PrismaClient, PrismaReplicaClient } from "@trigger.dev/database"; +import { ControlPlaneCache } from "./controlPlaneCache.server"; +import { ControlPlaneResolver } from "./controlPlaneResolver.server"; + +// Wraps a real Prisma client so every `client..(...)` call is tallied by +// `"."` before being forwarded, unmodified, to the real delegate. No behavior is +// faked or short-circuited - this is instrumentation, not a mock. +function createCallCountingProxy( + client: T +): { client: T; counts: Map } { + const counts = new Map(); + const bump = (key: string) => counts.set(key, (counts.get(key) ?? 0) + 1); + const wrappedModels = new Map(); + + const wrapModel = (model: string, delegate: object) => + new Proxy(delegate, { + get(target, prop, receiver) { + const value = Reflect.get(target, prop, receiver); + if (typeof prop === "string" && typeof value === "function") { + return (...args: unknown[]) => { + bump(`${model}.${prop}`); + return (value as (...a: unknown[]) => unknown).apply(target, args); + }; + } + return value; + }, + }); + + const proxy = new Proxy(client, { + get(target, prop, receiver) { + const value = Reflect.get(target, prop, receiver); + if ( + typeof prop === "string" && + value && + typeof value === "object" && + typeof (value as { findMany?: unknown }).findMany === "function" + ) { + if (!wrappedModels.has(prop)) { + wrappedModels.set(prop, wrapModel(prop, value as object)); + } + return wrappedModels.get(prop); + } + return value; + }, + }); + + return { client: proxy as T, counts }; +} + +let n = 0; + +async function seedEnv(prisma: PrismaClient) { + const s = n++; + const organization = await prisma.organization.create({ + data: { title: `Org ${s}`, slug: `org-${s}` }, + }); + const project = await prisma.project.create({ + data: { + name: `P ${s}`, + slug: `p-${s}`, + externalRef: `proj_${s}`, + organizationId: organization.id, + }, + }); + const environment = await prisma.runtimeEnvironment.create({ + data: { + type: "PRODUCTION", + slug: `env-${s}`, + projectId: project.id, + organizationId: organization.id, + apiKey: `tr_${s}`, + pkApiKey: `pk_${s}`, + shortcode: `sc_${s}`, + }, + }); + return { organization, project, environment }; +} + +async function seedBackgroundWorker( + prisma: PrismaClient, + ctx: { projectId: string; runtimeEnvironmentId: string }, + version: string +) { + const s = n++; + return prisma.backgroundWorker.create({ + data: { + friendlyId: `worker_${s}`, + version, + contentHash: `hash_${s}`, + projectId: ctx.projectId, + runtimeEnvironmentId: ctx.runtimeEnvironmentId, + metadata: {}, + }, + }); +} + +describe("ControlPlaneResolver.resolveRunLockedWorkersByVersionIds", () => { + postgresTest( + "issues ONE grouped findMany and ZERO per-id findFirst for N distinct ids", + async ({ prisma }) => { + const { project, environment } = await seedEnv(prisma); + const workers = await Promise.all( + [0, 1, 2].map((i) => + seedBackgroundWorker( + prisma, + { projectId: project.id, runtimeEnvironmentId: environment.id }, + `2024010${i}.0` + ) + ) + ); + + const { client: countedReplica, counts } = createCallCountingProxy(prisma); + + const resolver = new ControlPlaneResolver({ + controlPlanePrimary: prisma, + controlPlaneReplica: countedReplica as unknown as PrismaReplicaClient, + cache: new ControlPlaneCache(), + splitEnabled: () => true, + }); + + const ids = workers.map((w) => w.id); + const result = await resolver.resolveRunLockedWorkersByVersionIds(ids); + + expect(counts.get("backgroundWorker.findMany") ?? 0).toBe(1); + expect(counts.get("backgroundWorker.findFirst") ?? 0).toBe(0); + + for (const worker of workers) { + expect(result.get(worker.id)?.lockedToVersion?.version).toBe(worker.version); + } + } + ); + + postgresTest("cache-hit ids issue no query at all", async ({ prisma }) => { + const { project, environment } = await seedEnv(prisma); + const worker = await seedBackgroundWorker( + prisma, + { projectId: project.id, runtimeEnvironmentId: environment.id }, + "20240101.0" + ); + + const cache = new ControlPlaneCache(); + const { client: countedReplica, counts } = createCallCountingProxy(prisma); + + const resolver = new ControlPlaneResolver({ + controlPlanePrimary: prisma, + controlPlaneReplica: countedReplica as unknown as PrismaReplicaClient, + cache, + splitEnabled: () => true, + }); + + // Warm the cache. + await resolver.resolveRunLockedWorkersByVersionIds([worker.id]); + expect(counts.get("backgroundWorker.findMany") ?? 0).toBe(1); + + // Second call for the same id is served entirely from cache - no query at all. + const result = await resolver.resolveRunLockedWorkersByVersionIds([worker.id]); + expect(counts.get("backgroundWorker.findMany") ?? 0).toBe(1); + expect(counts.get("backgroundWorker.findFirst") ?? 0).toBe(0); + expect(result.get(worker.id)?.lockedToVersion?.version).toBe(worker.version); + }); +}); diff --git a/apps/webapp/app/v3/runOpsMigration/controlPlaneResolver.server.ts b/apps/webapp/app/v3/runOpsMigration/controlPlaneResolver.server.ts index fb903bc8ea..d6a706c966 100644 --- a/apps/webapp/app/v3/runOpsMigration/controlPlaneResolver.server.ts +++ b/apps/webapp/app/v3/runOpsMigration/controlPlaneResolver.server.ts @@ -70,6 +70,15 @@ function lockedWorkerKey(lockedById?: string | null, lockedToVersionId?: string return `${lockedById ?? "_"}:${lockedToVersionId ?? "_"}`; } +type LockedToVersionRow = { + id: string; + version: string; + sdkVersion: string; + runtime: string | null; + runtimeVersion: string | null; + supportsLazyAttempts: boolean; +}; + export class ControlPlaneResolver { private readonly controlPlanePrimary: PrismaClient; private readonly controlPlaneReplica: PrismaReplicaClient; @@ -243,6 +252,87 @@ export class ControlPlaneResolver { }; } + /** + * Grouped counterpart to `resolveRunLockedWorker({ lockedToVersionId })`: resolves the + * `lockedToVersion` for many distinct BackgroundWorker ids with ONE `findMany` for the + * cache misses (BackgroundWorker is control-plane only, no run-ops residency split), instead + * of one `findFirst` per id. Cache keys/shape match `resolveRunLockedWorker`'s + * `lockedById=undefined` slot, so entries populated here are reused by that method and + * vice versa. + */ + async resolveRunLockedWorkersByVersionIds( + ids: string[] + ): Promise> { + const result = new Map(); + const uniqueIds = [...new Set(ids)]; + if (uniqueIds.length === 0) { + return result; + } + + if (!this.splitEnabled()) { + const rows = await this.#queryLockedToVersionRows(this.controlPlanePrimary, uniqueIds); + for (const id of uniqueIds) { + result.set(id, this.#toResolvedRunLockedWorker(rows.get(id))); + } + return result; + } + + const misses: string[] = []; + for (const id of uniqueIds) { + const cached = this.cache.getLockedWorker(lockedWorkerKey(undefined, id)); + if (cached !== undefined) { + result.set(id, cached); + } else { + misses.push(id); + } + } + + if (misses.length > 0) { + const rows = await this.#queryLockedToVersionRows(this.controlPlaneReplica, misses); + for (const id of misses) { + const value = this.#toResolvedRunLockedWorker(rows.get(id)); + this.cache.setLockedWorker(lockedWorkerKey(undefined, id), value); + result.set(id, value); + } + } + + return result; + } + + async #queryLockedToVersionRows( + client: CpClient, + ids: string[] + ): Promise> { + const rows = await client.backgroundWorker.findMany({ + where: { id: { in: ids } }, + select: { + id: true, + version: true, + sdkVersion: true, + runtime: true, + runtimeVersion: true, + supportsLazyAttempts: true, + }, + }); + return new Map(rows.map((row) => [row.id, row])); + } + + #toResolvedRunLockedWorker(row: LockedToVersionRow | undefined): ResolvedRunLockedWorker | null { + if (!row) { + return null; + } + return { + lockedBy: null, + lockedToVersion: { + version: row.version, + sdkVersion: row.sdkVersion, + runtime: row.runtime, + runtimeVersion: row.runtimeVersion, + supportsLazyAttempts: row.supportsLazyAttempts, + }, + }; + } + async resolveWorkerVersion(args: { environmentId: string; backgroundWorkerId?: string; diff --git a/apps/webapp/app/v3/services/cancelDevSessionRuns.server.ts b/apps/webapp/app/v3/services/cancelDevSessionRuns.server.ts index 3575a75052..3ebc5df76a 100644 --- a/apps/webapp/app/v3/services/cancelDevSessionRuns.server.ts +++ b/apps/webapp/app/v3/services/cancelDevSessionRuns.server.ts @@ -1,10 +1,17 @@ import { type RunStore } from "@internal/run-store"; +import { ownerEngine } from "@trigger.dev/core/v3/isomorphic"; import { z } from "zod"; -import { type PrismaClientOrTransaction } from "~/db.server"; +import { + runOpsLegacyReplica as defaultLegacyReplica, + runOpsNewReplica as defaultNewClient, + type PrismaClientOrTransaction, + type PrismaReplicaClient, +} from "~/db.server"; import { findLatestSession } from "~/models/runtimeEnvironment.server"; import { logger } from "~/services/logger.server"; import { commonWorker } from "../commonWorker.server"; -import { type ReadThroughDeps, readThroughRun } from "../runOpsMigration/readThrough.server"; +import { type ReadThroughDeps } from "../runOpsMigration/readThrough.server"; +import { isSplitEnabled } from "../runOpsMigration/splitMode.server"; import { BaseService } from "./baseService.server"; import { type CancelableTaskRun, CancelTaskRunService } from "./cancelTaskRun.server"; @@ -17,9 +24,19 @@ export const CancelDevSessionRunsServiceOptions = z.object({ export type CancelDevSessionRunsServiceOptions = z.infer; +const RUN_SELECT = { + id: true, + engine: true, + status: true, + friendlyId: true, + taskEventStore: true, + createdAt: true, + completedAt: true, +} as const; + export class CancelDevSessionRunsService extends BaseService { // Injectable read-through deps for the run-ops TaskRun read. Undefined in production: - // readThroughRun then uses its ~/db.server singleton handles and the boot split flag, + // the grouped read then uses its ~/db.server singleton handles and the boot split flag, // so single-DB is unchanged. Tests inject the hetero new/legacy handles + splitEnabled. readonly #readThroughDeps?: ReadThroughDeps; @@ -69,80 +86,116 @@ export class CancelDevSessionRunsService extends BaseService { const cancelTaskRunService = new CancelTaskRunService(); - // readThroughRun resolves residency from the run id alone; an env scope is only - // available when a cancelled session was resolved. - const environmentId = cancelledSession?.environmentId ?? ""; + // Read every run up front, grouped by id shape and (when the split is on) by + // residency, instead of the old per-run readThroughRun probe. Cancellation stays a + // per-run mutation loop below, unchanged. + const runsById = await this.#readRunsGrouped(options.runIds); for (const runId of options.runIds) { - await this.#cancelInProgressRun( - runId, - cancelTaskRunService, - options.cancelledAt, - options.reason, - environmentId - ); + logger.debug("Cancelling in progress run", { runId }); + + const taskRun = runsById.get(runId); + + if (!taskRun) { + continue; + } + + try { + await cancelTaskRunService.call(taskRun, { + reason: options.reason, + cancelAttempts: true, + cancelledAt: options.cancelledAt, + }); + } catch (e) { + logger.error("Failed to cancel in progress run", { + runId, + error: e, + }); + } } } - async #cancelInProgressRun( - runId: string, - service: CancelTaskRunService, - cancelledAt: Date, - reason: string, - environmentId: string - ) { - logger.debug("Cancelling in progress run", { runId }); - - // Read-through: new store first, legacy read replica for an old - // in-retention run; single plain read in single-DB passthrough. - const where = runId.startsWith("run_") ? { friendlyId: runId } : { id: runId }; - - const result = await readThroughRun({ - runId, - environmentId, - readNew: (client) => - client.taskRun.findFirst({ - where, - select: { - id: true, - engine: true, - status: true, - friendlyId: true, - taskEventStore: true, - createdAt: true, - completedAt: true, - }, - }), - readLegacy: (replica) => - replica.taskRun.findFirst({ - where, - select: { - id: true, - engine: true, - status: true, - friendlyId: true, - taskEventStore: true, - createdAt: true, - completedAt: true, - }, - }), - deps: this.#readThroughDeps, - }); + // Grouped replacement for readThroughRun-per-id: a constant number of queries + // regardless of N. Same fan-out order as before (NEW-classified ids read only the + // new store; LEGACY-classified ids probe new first, then the legacy replica for + // misses; passthrough reads the single store once), just batched via findMany. + async #readRunsGrouped(runIds: string[]): Promise> { + const result = new Map(); + + if (runIds.length === 0) { + return result; + } + + const deps = this.#readThroughDeps; + const newClient = deps?.newClient ?? defaultNewClient; + const legacyReplica = deps?.legacyReplica ?? defaultLegacyReplica; + const splitEnabled = deps?.splitEnabled ?? (await isSplitEnabled()); + + const friendlyIds: string[] = []; + const internalIds: string[] = []; + + for (const runId of runIds) { + (runId.startsWith("run_") ? friendlyIds : internalIds).push(runId); + } - if (result.source === "not-found" || result.source === "past-retention") { - return; + const applyRows = ( + rows: CancelableTaskRun[], + requestedIds: string[], + key: "id" | "friendlyId" + ) => { + const byKey = new Map(rows.map((row) => [row[key], row])); + for (const id of requestedIds) { + const row = byKey.get(id); + if (row) { + result.set(id, row); + } + } + }; + + if (!splitEnabled) { + applyRows(await readGroup(newClient, internalIds, "id"), internalIds, "id"); + applyRows(await readGroup(newClient, friendlyIds, "friendlyId"), friendlyIds, "friendlyId"); + return result; } - const taskRun = result.value; + // classifyResidency is total (never throws): NEW ids read only the new store, + // LEGACY ids fan out new-then-legacy-replica. + const newIds: string[] = []; + const newFriendlyIds: string[] = []; + const legacyIds: string[] = []; + const legacyFriendlyIds: string[] = []; - try { - await service.call(taskRun, { reason, cancelAttempts: true, cancelledAt }); - } catch (e) { - logger.error("Failed to cancel in progress run", { - runId, - error: e, - }); + for (const id of internalIds) { + (ownerEngine(id) === "NEW" ? newIds : legacyIds).push(id); } + for (const id of friendlyIds) { + (ownerEngine(id) === "NEW" ? newFriendlyIds : legacyFriendlyIds).push(id); + } + + applyRows(await readGroup(newClient, newIds, "id"), newIds, "id"); + applyRows( + await readGroup(newClient, newFriendlyIds, "friendlyId"), + newFriendlyIds, + "friendlyId" + ); + + applyRows(await readGroup(newClient, legacyIds, "id"), legacyIds, "id"); + const legacyIdMisses = legacyIds.filter((id) => !result.has(id)); + applyRows(await readGroup(legacyReplica, legacyIdMisses, "id"), legacyIdMisses, "id"); + + applyRows( + await readGroup(newClient, legacyFriendlyIds, "friendlyId"), + legacyFriendlyIds, + "friendlyId" + ); + const legacyFriendlyMisses = legacyFriendlyIds.filter((id) => !result.has(id)); + applyRows( + await readGroup(legacyReplica, legacyFriendlyMisses, "friendlyId"), + legacyFriendlyMisses, + "friendlyId" + ); + + return result; } static async enqueue(options: CancelDevSessionRunsServiceOptions, runAt?: Date) { @@ -156,3 +209,27 @@ export class CancelDevSessionRunsService extends BaseService { }); } } + +// One `findFirst` for a lone id (byte-identical to the old per-run read), or one +// `findMany` for 2+ ids in the same bucket — the N+1 fix. +async function readGroup( + client: PrismaReplicaClient, + ids: string[], + field: "id" | "friendlyId" +): Promise { + if (ids.length === 0) { + return []; + } + + if (ids.length === 1) { + const row = + field === "id" + ? await client.taskRun.findFirst({ where: { id: ids[0] }, select: RUN_SELECT }) + : await client.taskRun.findFirst({ where: { friendlyId: ids[0] }, select: RUN_SELECT }); + return row ? [row] : []; + } + + return field === "id" + ? client.taskRun.findMany({ where: { id: { in: ids } }, select: RUN_SELECT }) + : client.taskRun.findMany({ where: { friendlyId: { in: ids } }, select: RUN_SELECT }); +} diff --git a/apps/webapp/test/apiBatchResultsPresenter.dedicatedSeam.test.ts b/apps/webapp/test/apiBatchResultsPresenter.dedicatedSeam.test.ts new file mode 100644 index 0000000000..eb322c48a1 --- /dev/null +++ b/apps/webapp/test/apiBatchResultsPresenter.dedicatedSeam.test.ts @@ -0,0 +1,337 @@ +// Flow-level seam proof for ApiBatchResultsPresenter, ABOVE the run-store-level coverage in +// internal-packages/run-store (batchItemMisroute, batchCompletionResidency). Proves the batch +// RESULTS READ assembles correctly when one batch's members are genuinely split across the real +// dedicated run-ops subset schema (prisma17 / RunOpsPrismaClient) and the full control-plane +// schema (prisma14) — not a mirrored full schema on both sides. No mocks. +import { heteroRunOpsPostgresTest } from "@internal/testcontainers"; +import type { RunOpsPrismaClient } from "@internal/run-ops-database"; +import type { PrismaClient } from "@trigger.dev/database"; +import { generateRunOpsId } from "@trigger.dev/core/v3/isomorphic"; +import { describe, expect, vi } from "vitest"; +import type { PrismaReplicaClient } from "~/db.server"; +import type { AuthenticatedEnvironment } from "~/services/apiAuth.server"; +import { ApiBatchResultsPresenter } from "~/presenters/v3/ApiBatchResultsPresenter.server"; + +vi.setConfig({ testTimeout: 60_000 }); + +// A prisma handle that throws on any access — proves the passthrough constructor args are never +// touched on the split path. +const throwingPrisma = new Proxy( + {}, + { + get(_t, prop) { + throw new Error( + `passthrough handle must not be touched on the split path (got .${String(prop)})` + ); + }, + } +) as unknown as PrismaReplicaClient; + +// 25-char cuid-shaped id, no v1 version marker at index 25 -> classifies LEGACY. +function generateLegacyCuid(): string { + const alphabet = "0123456789abcdefghijklmnopqrstuvwxyz"; + const suffix = Array.from({ length: 24 }, () => alphabet[Math.floor(Math.random() * 36)]).join( + "" + ); + return `c${suffix}`; +} + +let seedCounter = 0; + +// TaskRun on the full control-plane schema has real FKs into RuntimeEnvironment/Project. +async function seedLegacyEnv(prisma14: PrismaClient, slug: string) { + const n = seedCounter++; + const organization = await prisma14.organization.create({ + data: { title: `Org ${slug}`, slug: `org-${slug}-${n}` }, + }); + const project = await prisma14.project.create({ + data: { + name: `Proj ${slug}`, + slug: `proj-${slug}-${n}`, + organizationId: organization.id, + externalRef: `ext-${slug}-${n}`, + }, + }); + const environment = await prisma14.runtimeEnvironment.create({ + data: { + slug: `env-${slug}-${n}`, + type: "PRODUCTION", + projectId: project.id, + organizationId: organization.id, + apiKey: `api-${slug}-${n}`, + pkApiKey: `pk-${slug}-${n}`, + shortcode: `sc-${slug}-${n}`, + }, + }); + return { organization, project, environment }; +} + +type SeedCtx = Awaited>; + +type MemberSeed = { + id: string; + friendlyId: string; + status: "COMPLETED_SUCCESSFULLY" | "COMPLETED_WITH_ERRORS"; + output?: string; + error?: unknown; +}; + +// Drop the TaskRunAttempt worker/queue FKs so attempts can be seeded without standing up +// BackgroundWorker/TaskQueue parents — incidental to this read path. +async function relaxLegacyAttemptFk(prisma14: PrismaClient) { + for (const sql of [ + `ALTER TABLE "TaskRunAttempt" DROP CONSTRAINT IF EXISTS "TaskRunAttempt_backgroundWorkerId_fkey"`, + `ALTER TABLE "TaskRunAttempt" DROP CONSTRAINT IF EXISTS "TaskRunAttempt_backgroundWorkerTaskId_fkey"`, + `ALTER TABLE "TaskRunAttempt" DROP CONSTRAINT IF EXISTS "TaskRunAttempt_queueId_fkey"`, + ]) { + await prisma14.$executeRawUnsafe(sql); + } +} + +async function seedLegacyMember(prisma14: PrismaClient, ctx: SeedCtx, m: MemberSeed) { + const run = await prisma14.taskRun.create({ + data: { + id: m.id, + friendlyId: m.friendlyId, + taskIdentifier: "my-task", + status: m.status, + payload: JSON.stringify({}), + payloadType: "application/json", + traceId: m.id, + spanId: m.id, + queue: "main", + runtimeEnvironmentId: ctx.environment.id, + projectId: ctx.project.id, + organizationId: ctx.organization.id, + environmentType: "PRODUCTION", + engine: "V2", + }, + }); + + await prisma14.taskRunAttempt.create({ + data: { + friendlyId: `attempt_${m.id}`, + number: 1, + taskRunId: run.id, + backgroundWorkerId: "bw", + backgroundWorkerTaskId: "bwt", + runtimeEnvironmentId: ctx.environment.id, + queueId: "q", + status: m.status === "COMPLETED_SUCCESSFULLY" ? "COMPLETED" : "FAILED", + output: m.output, + outputType: "application/json", + error: m.error as any, + }, + }); + + return run; +} + +// TaskRun/TaskRunAttempt on the DEDICATED (run-ops subset) schema: runtimeEnvironmentId, +// organizationId, projectId, backgroundWorkerId/backgroundWorkerTaskId/queueId are all +// scalar-only there (no real relation) — no parent rows to seed, unlike the legacy side. +async function seedNewMember( + prisma17: RunOpsPrismaClient, + ctx: { envId: string; orgId: string; projectId: string }, + m: MemberSeed +) { + await prisma17.taskRun.create({ + data: { + id: m.id, + engine: "V2", + status: m.status, + friendlyId: m.friendlyId, + runtimeEnvironmentId: ctx.envId, + environmentType: "PRODUCTION", + organizationId: ctx.orgId, + projectId: ctx.projectId, + taskIdentifier: "my-task", + payload: JSON.stringify({}), + payloadType: "application/json", + traceContext: {}, + traceId: m.id, + spanId: m.id, + queue: "main", + isTest: false, + taskEventStore: "taskEvent", + depth: 0, + }, + }); + + await prisma17.taskRunAttempt.create({ + data: { + friendlyId: `attempt_${m.id}`, + number: 1, + taskRunId: m.id, + backgroundWorkerId: "bw", + backgroundWorkerTaskId: "bwt", + runtimeEnvironmentId: ctx.envId, + queueId: "q", + status: m.status === "COMPLETED_SUCCESSFULLY" ? "COMPLETED" : "FAILED", + output: m.output, + outputType: "application/json", + error: m.error as any, + }, + }); +} + +// BatchTaskRunItem.taskRunId is a real FK to TaskRun on the dedicated schema too. A batch seeded +// on the dedicated DB whose items reference a LEGACY-resident (or wholly missing) member has no +// local row for that member — drop the FK so the cross-seam item row can exist, matching the +// physical reality of a split batch. +async function relaxNewBatchItemFk(prisma17: RunOpsPrismaClient) { + await prisma17.$executeRawUnsafe( + `ALTER TABLE "BatchTaskRunItem" DROP CONSTRAINT IF EXISTS "BatchTaskRunItem_taskRunId_fkey"` + ); +} + +async function seedBatchOnNew( + prisma17: RunOpsPrismaClient, + envId: string, + friendlyId: string, + memberIds: string[] +) { + const batch = await prisma17.batchTaskRun.create({ + data: { + friendlyId, + runtimeEnvironmentId: envId, + runCount: memberIds.length, + runIds: [], + batchVersion: "runengine:v2", + }, + }); + // Items in a deterministic order so the result `items` order is assertable. + for (const taskRunId of memberIds) { + await prisma17.batchTaskRunItem.create({ + data: { batchTaskRunId: batch.id, taskRunId, status: "COMPLETED" }, + }); + } + return batch; +} + +const env = (ctx: SeedCtx) => + ({ + id: ctx.environment.id, + type: ctx.environment.type, + slug: ctx.environment.slug, + organizationId: ctx.organization.id, + organization: { slug: ctx.organization.slug, title: ctx.organization.title }, + projectId: ctx.project.id, + project: { name: ctx.project.name }, + }) as unknown as AuthenticatedEnvironment; + +describe("ApiBatchResultsPresenter split mode — real run-ops dedicated schema seam", () => { + // The core assembly proof: one batch, members genuinely split across the two physically + // distinct, differently-shaped DBs (dedicated subset vs full control-plane). + heteroRunOpsPostgresTest( + "a batch with members split across the dedicated NEW schema and the legacy schema returns the complete union", + async ({ prisma14, prisma17 }: { prisma14: PrismaClient; prisma17: RunOpsPrismaClient }) => { + const ctx = await seedLegacyEnv(prisma14, "split-new-legacy"); + await relaxLegacyAttemptFk(prisma14); + await relaxNewBatchItemFk(prisma17); + + const newMemberId = generateRunOpsId(); + expect(newMemberId.length).toBe(26); + const legacyMemberId = generateLegacyCuid(); + expect(legacyMemberId.length).toBe(25); + + // NEW-resident member: lives ONLY on the dedicated run-ops schema (prisma17). + await seedNewMember( + prisma17, + { envId: ctx.environment.id, orgId: ctx.organization.id, projectId: ctx.project.id }, + { + id: newMemberId, + friendlyId: "run_new_member", + status: "COMPLETED_SUCCESSFULLY", + output: JSON.stringify({ from: "new" }), + } + ); + + // LEGACY-resident member: lives ONLY on the full control-plane schema (prisma14). + await seedLegacyMember(prisma14, ctx, { + id: legacyMemberId, + friendlyId: "run_legacy_member", + status: "COMPLETED_WITH_ERRORS", + error: { type: "BUILT_IN_ERROR", name: "Err", message: "boom", stackTrace: "" }, + }); + + // The batch row + items live on the NEW dedicated DB; items reference both members. + const batchFriendlyId = "batch_split_seam"; + await seedBatchOnNew(prisma17, ctx.environment.id, batchFriendlyId, [ + newMemberId, + legacyMemberId, + ]); + + const presenter = new ApiBatchResultsPresenter(throwingPrisma, throwingPrisma, { + splitEnabled: true, + newClient: prisma17 as unknown as PrismaReplicaClient, + legacyReplica: prisma14 as unknown as PrismaReplicaClient, + }); + + const result = await presenter.call(batchFriendlyId, env(ctx)); + + expect(result).toBeDefined(); + expect(result!.id).toBe(batchFriendlyId); + expect(result!.items).toHaveLength(2); + + // Order follows item order: the NEW-resident member first, the legacy-resident second. + const [first, second] = result!.items; + expect(first).toEqual({ + ok: true, + id: "run_new_member", + taskIdentifier: "my-task", + output: JSON.stringify({ from: "new" }), + outputType: "application/json", + }); + expect(second).toMatchObject({ + ok: false, + id: "run_legacy_member", + taskIdentifier: "my-task", + }); + } + ); + + // Dangling member: an item id that resolves on NEITHER the dedicated NEW schema nor the legacy + // schema must be dropped from the result, not thrown — the presenter degrades to the reachable + // set. Give the missing id a legacy shape so it also exercises the "not NEW-shaped -> legacy + // candidate" branch, proving the legacy probe finding nothing doesn't throw either. + heteroRunOpsPostgresTest( + "a member id that resolves on neither the dedicated NEW schema nor the legacy schema is dropped, not thrown", + async ({ prisma14, prisma17 }: { prisma14: PrismaClient; prisma17: RunOpsPrismaClient }) => { + const ctx = await seedLegacyEnv(prisma14, "dangling"); + await relaxNewBatchItemFk(prisma17); + + const presentId = generateRunOpsId(); + const missingId = generateLegacyCuid(); // referenced by an item but seeded on NEITHER DB + + await seedNewMember( + prisma17, + { envId: ctx.environment.id, orgId: ctx.organization.id, projectId: ctx.project.id }, + { + id: presentId, + friendlyId: "run_present", + status: "COMPLETED_SUCCESSFULLY", + output: JSON.stringify({ present: true }), + } + ); + + const batchFriendlyId = "batch_dangling_seam"; + await seedBatchOnNew(prisma17, ctx.environment.id, batchFriendlyId, [presentId, missingId]); + + const presenter = new ApiBatchResultsPresenter(throwingPrisma, throwingPrisma, { + splitEnabled: true, + newClient: prisma17 as unknown as PrismaReplicaClient, + legacyReplica: prisma14 as unknown as PrismaReplicaClient, + }); + + const result = await presenter.call(batchFriendlyId, env(ctx)); + + expect(result).toBeDefined(); + expect(result!.id).toBe(batchFriendlyId); + // The dangling member is silently dropped; the reachable member still returns — the call + // must not throw despite the item referencing a run absent from both physical DBs. + expect(result!.items).toHaveLength(1); + expect(result!.items[0]).toMatchObject({ ok: true, id: "run_present" }); + } + ); +}); diff --git a/apps/webapp/test/apiBatchResultsPresenter.splitNPlus1.test.ts b/apps/webapp/test/apiBatchResultsPresenter.splitNPlus1.test.ts new file mode 100644 index 0000000000..614e112992 --- /dev/null +++ b/apps/webapp/test/apiBatchResultsPresenter.splitNPlus1.test.ts @@ -0,0 +1,251 @@ +// RED→GREEN: kill the #callSplit per-member N+1 in ApiBatchResultsPresenter. +// +// Today, #callSplit hydrates every batch member independently via `readThroughRun` inside +// `Promise.all(batchRun.items.map(...))` — that is one (up to two, new-then-legacy) `taskRun` +// query PER member. `#callPassthrough` already does the grouped one-query form via +// `this.runStore.findRuns({ where: { id: { in: taskRunIds } } })`. +// +// The fix replaces the per-member fan-out with ONE grouped call to the RunStore's +// `findRunsByIds` method, mirroring `#callPassthrough`. This test proves the query-count +// reduction with a call-counting Proxy over a REAL testcontainer Postgres client: every call is +// delegated unchanged to the real client (the DB still runs the query) — this is instrumentation, +// not a mock. +import { PostgresRunStore } from "@internal/run-store"; +import { postgresTest } from "@internal/testcontainers"; +import type { PrismaClient } from "@trigger.dev/database"; +import { describe, expect } from "vitest"; +import type { PrismaReplicaClient } from "~/db.server"; +import type { AuthenticatedEnvironment } from "~/services/apiAuth.server"; +import { ApiBatchResultsPresenter } from "~/presenters/v3/ApiBatchResultsPresenter.server"; + +// 26-char run-ops v1 body: 24-char base32hex core + region char + version "1" at index 25. +// ownerEngine classifies residency by the VERSION CHAR AT INDEX 25, not by length — a naive id +// generator would misclassify this as LEGACY unless the last two chars are a valid region+version. +function newRunId(c: string) { + return c.repeat(24) + "01"; +} + +type CallCounts = { findMany: number; findFirst: number }; + +// Wrap a REAL client's `taskRun` delegate to tally findMany/findFirst calls, delegating every +// call unchanged to the real client (the DB still runs the query — pure instrumentation). +function countingClient(real: PrismaClient): { client: PrismaClient; counts: CallCounts } { + const counts: CallCounts = { findMany: 0, findFirst: 0 }; + const countingTaskRun = new Proxy((real as any).taskRun, { + get(target, prop) { + if (prop === "findMany" || prop === "findFirst") { + counts[prop as "findMany" | "findFirst"]++; + } + return (target as any)[prop]; + }, + }); + const client = new Proxy(real, { + get(target, prop) { + if (prop === "taskRun") { + return countingTaskRun; + } + return (target as any)[prop]; + }, + }) as PrismaClient; + return { client, counts }; +} + +let seedCounter = 0; + +async function seedEnv(prisma: PrismaClient, slug: string) { + const n = seedCounter++; + const organization = await prisma.organization.create({ + data: { title: `Org ${slug}`, slug: `org-${slug}-${n}` }, + }); + const project = await prisma.project.create({ + data: { + name: `Proj ${slug}`, + slug: `proj-${slug}-${n}`, + organizationId: organization.id, + externalRef: `ext-${slug}-${n}`, + }, + }); + const environment = await prisma.runtimeEnvironment.create({ + data: { + slug: `env-${slug}-${n}`, + type: "PRODUCTION", + projectId: project.id, + organizationId: organization.id, + apiKey: `api-${slug}-${n}`, + pkApiKey: `pk-${slug}-${n}`, + shortcode: `sc-${slug}-${n}`, + }, + }); + return { organization, project, environment }; +} + +type SeedCtx = Awaited>; + +type MemberSeed = { + id: string; + friendlyId: string; + status: "COMPLETED_SUCCESSFULLY" | "COMPLETED_WITH_ERRORS"; + output?: string; + error?: unknown; +}; + +async function seedMember(prisma: PrismaClient, ctx: SeedCtx, m: MemberSeed) { + const run = await prisma.taskRun.create({ + data: { + id: m.id, + friendlyId: m.friendlyId, + taskIdentifier: "my-task", + status: m.status, + payload: JSON.stringify({}), + payloadType: "application/json", + traceId: m.id, + spanId: m.id, + queue: "main", + runtimeEnvironmentId: ctx.environment.id, + projectId: ctx.project.id, + organizationId: ctx.organization.id, + environmentType: "PRODUCTION", + engine: "V2", + }, + }); + + await prisma.taskRunAttempt.create({ + data: { + friendlyId: `attempt_${m.id}`, + number: 1, + taskRunId: run.id, + backgroundWorkerId: "bw", + backgroundWorkerTaskId: "bwt", + runtimeEnvironmentId: ctx.environment.id, + queueId: "q", + status: m.status === "COMPLETED_SUCCESSFULLY" ? "COMPLETED" : "FAILED", + output: m.output, + outputType: "application/json", + error: m.error as any, + }, + }); + + return run; +} + +// Drop the TaskRunAttempt worker/queue FKs so attempts can be seeded (their output/error is what's +// under test) without standing up BackgroundWorker/TaskQueue parents — incidental to this read path. +async function relaxFks(prisma: PrismaClient) { + for (const sql of [ + `ALTER TABLE "TaskRunAttempt" DROP CONSTRAINT IF EXISTS "TaskRunAttempt_backgroundWorkerId_fkey"`, + `ALTER TABLE "TaskRunAttempt" DROP CONSTRAINT IF EXISTS "TaskRunAttempt_backgroundWorkerTaskId_fkey"`, + `ALTER TABLE "TaskRunAttempt" DROP CONSTRAINT IF EXISTS "TaskRunAttempt_queueId_fkey"`, + ]) { + await prisma.$executeRawUnsafe(sql); + } +} + +async function seedBatch( + prisma: PrismaClient, + ctx: SeedCtx, + friendlyId: string, + memberIds: string[] +) { + const batch = await prisma.batchTaskRun.create({ + data: { + friendlyId, + runtimeEnvironmentId: ctx.environment.id, + runCount: memberIds.length, + runIds: [], + batchVersion: "runengine:v2", + }, + }); + // Items in a deterministic order so the result `items` order is assertable. + for (const taskRunId of memberIds) { + await prisma.batchTaskRunItem.create({ + data: { batchTaskRunId: batch.id, taskRunId, status: "COMPLETED" }, + }); + } + return batch; +} + +const env = (ctx: SeedCtx) => + ({ + id: ctx.environment.id, + type: ctx.environment.type, + slug: ctx.environment.slug, + organizationId: ctx.organization.id, + organization: { slug: ctx.organization.slug, title: ctx.organization.title }, + projectId: ctx.project.id, + project: { name: ctx.project.name }, + }) as unknown as AuthenticatedEnvironment; + +describe("ApiBatchResultsPresenter split mode — member hydration is grouped, not per-member", () => { + postgresTest( + "a batch of N members costs ONE findMany (never N findFirst) and returns the same result", + async ({ prisma }) => { + const ctx = await seedEnv(prisma, "n1"); + await relaxFks(prisma); + + const { client: countingPrisma, counts } = countingClient(prisma); + + const memberIds = [newRunId("a"), newRunId("b"), newRunId("c")]; + await seedMember(prisma, ctx, { + id: memberIds[0], + friendlyId: "run_a", + status: "COMPLETED_SUCCESSFULLY", + output: JSON.stringify({ from: "a" }), + }); + await seedMember(prisma, ctx, { + id: memberIds[1], + friendlyId: "run_b", + status: "COMPLETED_WITH_ERRORS", + error: { type: "BUILT_IN_ERROR", name: "Err", message: "boom", stackTrace: "" }, + }); + await seedMember(prisma, ctx, { + id: memberIds[2], + friendlyId: "run_c", + status: "COMPLETED_SUCCESSFULLY", + output: JSON.stringify({ from: "c" }), + }); + + await seedBatch(prisma, ctx, "batch_n1", memberIds); + + const runStore = new PostgresRunStore({ + prisma: countingPrisma, + readOnlyPrisma: countingPrisma, + }); + + const presenter = new ApiBatchResultsPresenter( + countingPrisma, + countingPrisma, + { + splitEnabled: true, + newClient: countingPrisma as unknown as PrismaReplicaClient, + legacyReplica: countingPrisma as unknown as PrismaReplicaClient, + }, + runStore + ); + + const result = await presenter.call("batch_n1", env(ctx)); + + expect(result).toBeDefined(); + expect(result!.id).toBe("batch_n1"); + expect(result!.items).toHaveLength(3); + expect(result!.items[0]).toEqual({ + ok: true, + id: "run_a", + taskIdentifier: "my-task", + output: JSON.stringify({ from: "a" }), + outputType: "application/json", + }); + expect(result!.items[1]).toMatchObject({ ok: false, id: "run_b" }); + expect(result!.items[2]).toEqual({ + ok: true, + id: "run_c", + taskIdentifier: "my-task", + output: JSON.stringify({ from: "c" }), + outputType: "application/json", + }); + + // The grouped-read proof: ONE findMany for the whole member set, never a findFirst per member. + expect(counts.findMany).toBe(1); + expect(counts.findFirst).toBe(0); + } + ); +}); diff --git a/apps/webapp/test/apiRetrieveRunPresenter.groupedLockedWorker.test.ts b/apps/webapp/test/apiRetrieveRunPresenter.groupedLockedWorker.test.ts new file mode 100644 index 0000000000..1b4a200751 --- /dev/null +++ b/apps/webapp/test/apiRetrieveRunPresenter.groupedLockedWorker.test.ts @@ -0,0 +1,266 @@ +import { containerTest } from "@internal/testcontainers"; +import { generateRunOpsId } from "@trigger.dev/core/v3/isomorphic"; +import type { PrismaClient } from "@trigger.dev/database"; +import { beforeEach, describe, expect, vi } from "vitest"; + +// `findRun` reads the module-level `prisma`/`$replica` (control-plane handles) both directly +// and, transitively, via the `runStore`/`controlPlaneResolver` singletons that also import them. +// A lazy Proxy - not a plain value - is captured by those singletons at their own (one-time) +// construction, so pointing it at a real container client after the fact still routes every +// call there. Not a DB mock: every call forwards to a real Postgres container. +const dbHolder = vi.hoisted(() => ({ client: undefined as PrismaClient | undefined })); + +vi.mock("~/db.server", () => { + const proxy = new Proxy( + {}, + { + get(_t, prop) { + if (!dbHolder.client) throw new Error("dbHolder.client not set for this test"); + return (dbHolder.client as any)[prop]; + }, + } + ); + return { prisma: proxy, $replica: proxy }; +}); + +vi.mock("~/v3/objectStore.server", () => ({ + generatePresignedUrl: vi.fn(async () => ({ success: false, error: "not-used" })), +})); + +import { ApiRetrieveRunPresenter } from "~/presenters/v3/ApiRetrieveRunPresenter.server"; +import type { AuthenticatedEnvironment } from "~/services/apiAuth.server"; + +vi.setConfig({ testTimeout: 60_000 }); + +// Wraps a real Prisma client so every `client..(...)` call is tallied by +// `"."` before being forwarded, unmodified, to the real delegate. +function createCallCountingProxy(client: PrismaClient): { + client: PrismaClient; + counts: Map; +} { + const counts = new Map(); + const wrappedModels = new Map(); + + const wrapModel = (model: string, delegate: object) => + new Proxy(delegate, { + get(target, prop, receiver) { + const value = Reflect.get(target, prop, receiver); + if (typeof prop === "string" && typeof value === "function") { + return (...args: unknown[]) => { + counts.set(`${model}.${prop}`, (counts.get(`${model}.${prop}`) ?? 0) + 1); + return (value as (...a: unknown[]) => unknown).apply(target, args); + }; + } + return value; + }, + }); + + const proxy = new Proxy(client, { + get(target, prop, receiver) { + const value = Reflect.get(target, prop, receiver); + if ( + typeof prop === "string" && + value && + typeof value === "object" && + typeof (value as { findMany?: unknown }).findMany === "function" + ) { + if (!wrappedModels.has(prop)) { + wrappedModels.set(prop, wrapModel(prop, value as object)); + } + return wrappedModels.get(prop); + } + return value; + }, + }); + + return { client: proxy as PrismaClient, counts }; +} + +async function seedOrgProjectEnv(prisma: PrismaClient, suffix: string) { + const organization = await prisma.organization.create({ + data: { title: `test-${suffix}`, slug: `test-${suffix}` }, + }); + const project = await prisma.project.create({ + data: { + name: `test-${suffix}`, + slug: `test-${suffix}`, + organizationId: organization.id, + externalRef: `test-${suffix}`, + }, + }); + const runtimeEnvironment = await prisma.runtimeEnvironment.create({ + data: { + slug: `test-${suffix}`, + type: "DEVELOPMENT", + projectId: project.id, + organizationId: organization.id, + apiKey: `tr_dev_${suffix}`, + pkApiKey: `pk_dev_${suffix}`, + shortcode: `short-${suffix}`, + }, + }); + return { organization, project, runtimeEnvironment }; +} + +function authEnv( + organization: { id: string }, + project: { id: string; externalRef: string }, + runtimeEnvironment: { id: string; slug: string } +): AuthenticatedEnvironment { + return { + id: runtimeEnvironment.id, + slug: runtimeEnvironment.slug, + organizationId: organization.id, + organization: { id: organization.id }, + project: { id: project.id, externalRef: project.externalRef }, + } as unknown as AuthenticatedEnvironment; +} + +async function seedBackgroundWorker( + prisma: PrismaClient, + ctx: { projectId: string; runtimeEnvironmentId: string }, + version: string +) { + return prisma.backgroundWorker.create({ + data: { + friendlyId: `worker_${generateRunOpsId()}`, + version, + contentHash: `hash_${generateRunOpsId()}`, + projectId: ctx.projectId, + runtimeEnvironmentId: ctx.runtimeEnvironmentId, + metadata: {}, + }, + }); +} + +interface SeedRunOpts { + id: string; + friendlyId: string; + runtimeEnvironmentId: string; + projectId: string; + organizationId: string; + lockedToVersionId?: string; + parentTaskRunId?: string; + rootTaskRunId?: string; +} + +async function seedRun(prisma: PrismaClient, opts: SeedRunOpts) { + return prisma.taskRun.create({ + data: { + id: opts.id, + friendlyId: opts.friendlyId, + taskIdentifier: "my-task", + payload: JSON.stringify({ hello: "world" }), + payloadType: "application/json", + traceId: `trace_${opts.id}`, + spanId: `span_${opts.id}`, + queue: "task/my-task", + runtimeEnvironmentId: opts.runtimeEnvironmentId, + projectId: opts.projectId, + organizationId: opts.organizationId, + environmentType: "DEVELOPMENT", + engine: "V2", + lockedToVersionId: opts.lockedToVersionId, + parentTaskRunId: opts.parentTaskRunId, + rootTaskRunId: opts.rootTaskRunId, + }, + }); +} + +beforeEach(() => { + dbHolder.client = undefined; +}); + +describe("ApiRetrieveRunPresenter.findRun locked-worker version resolution", () => { + containerTest( + "resolves run+parent+root+children lockedToVersion with ONE grouped query, not one per id", + async ({ prisma }) => { + const proxied = createCallCountingProxy(prisma); + dbHolder.client = proxied.client; + + const { organization, project, runtimeEnvironment } = await seedOrgProjectEnv( + prisma, + "grouped" + ); + const workerCtx = { projectId: project.id, runtimeEnvironmentId: runtimeEnvironment.id }; + + // Two distinct versions - `workerA` is deliberately reused across run/root/one child to + // prove the pre-existing dedup-by-Set still collapses to the same distinct-id count. + const workerA = await seedBackgroundWorker(prisma, workerCtx, "2024.1.0"); + const workerB = await seedBackgroundWorker(prisma, workerCtx, "2024.2.0"); + const workerC = await seedBackgroundWorker(prisma, workerCtx, "2024.3.0"); + + const rootId = generateRunOpsId(); + const parentId = generateRunOpsId(); + const runId = generateRunOpsId(); + const childId1 = generateRunOpsId(); + const childId2 = generateRunOpsId(); + + await seedRun(prisma, { + id: rootId, + friendlyId: `run_${rootId}`, + runtimeEnvironmentId: runtimeEnvironment.id, + projectId: project.id, + organizationId: organization.id, + lockedToVersionId: workerA.id, + }); + await seedRun(prisma, { + id: parentId, + friendlyId: `run_${parentId}`, + runtimeEnvironmentId: runtimeEnvironment.id, + projectId: project.id, + organizationId: organization.id, + rootTaskRunId: rootId, + lockedToVersionId: workerB.id, + }); + await seedRun(prisma, { + id: runId, + friendlyId: `run_${runId}`, + runtimeEnvironmentId: runtimeEnvironment.id, + projectId: project.id, + organizationId: organization.id, + parentTaskRunId: parentId, + rootTaskRunId: rootId, + lockedToVersionId: workerA.id, + }); + await seedRun(prisma, { + id: childId1, + friendlyId: `run_${childId1}`, + runtimeEnvironmentId: runtimeEnvironment.id, + projectId: project.id, + organizationId: organization.id, + parentTaskRunId: runId, + rootTaskRunId: rootId, + lockedToVersionId: workerC.id, + }); + await seedRun(prisma, { + id: childId2, + friendlyId: `run_${childId2}`, + runtimeEnvironmentId: runtimeEnvironment.id, + projectId: project.id, + organizationId: organization.id, + parentTaskRunId: runId, + rootTaskRunId: rootId, + lockedToVersionId: workerA.id, + }); + + const env = authEnv(organization, project, runtimeEnvironment); + const found = await ApiRetrieveRunPresenter.findRun(`run_${runId}`, env); + + expect(found).not.toBeNull(); + expect(found!.lockedToVersion?.version).toBe(workerA.version); + expect(found!.parentTaskRun?.lockedToVersion?.version).toBe(workerB.version); + expect(found!.rootTaskRun?.lockedToVersion?.version).toBe(workerA.version); + const versionByChildFriendlyId = new Map( + found!.childRuns.map((c) => [c.friendlyId, c.lockedToVersion?.version]) + ); + expect(versionByChildFriendlyId.get(`run_${childId1}`)).toBe(workerC.version); + expect(versionByChildFriendlyId.get(`run_${childId2}`)).toBe(workerA.version); + + // 3 distinct lockedToVersionIds (workerA, workerB, workerC) across 5 rows -> ONE grouped + // findMany, ZERO per-id findFirst. + expect(proxied.counts.get("backgroundWorker.findMany") ?? 0).toBe(1); + expect(proxied.counts.get("backgroundWorker.findFirst") ?? 0).toBe(0); + } + ); +}); diff --git a/apps/webapp/test/cancelDevSessionRunsGroupedReads.test.ts b/apps/webapp/test/cancelDevSessionRunsGroupedReads.test.ts new file mode 100644 index 0000000000..617e68b47c --- /dev/null +++ b/apps/webapp/test/cancelDevSessionRunsGroupedReads.test.ts @@ -0,0 +1,183 @@ +// Proves the dev-session-cancel read is grouped into a constant number of queries +// instead of one findFirst per run id (N+1). Real Postgres via testcontainers — the +// DB is never mocked. A call-counting proxy sits over the REAL Prisma client: it +// tallies findFirst/findMany per delegate and forwards every call unchanged, so the +// real query still runs against the real database. This is instrumentation, not a +// mock. +import { postgresTest } from "@internal/testcontainers"; +import type { PrismaClient } from "@trigger.dev/database"; +import { describe, expect, vi } from "vitest"; +import type { PrismaReplicaClient } from "~/db.server"; +import { CancelDevSessionRunsService } from "~/v3/services/cancelDevSessionRuns.server"; + +vi.setConfig({ testTimeout: 60_000 }); + +async function seedOrgProjectEnv(prisma: PrismaClient, suffix: string) { + const organization = await prisma.organization.create({ + data: { title: `test-${suffix}`, slug: `test-${suffix}` }, + }); + const project = await prisma.project.create({ + data: { + name: `test-${suffix}`, + slug: `test-${suffix}`, + organizationId: organization.id, + externalRef: `test-${suffix}`, + }, + }); + const runtimeEnvironment = await prisma.runtimeEnvironment.create({ + data: { + slug: `test-${suffix}`, + type: "DEVELOPMENT", + projectId: project.id, + organizationId: organization.id, + apiKey: `test-${suffix}`, + pkApiKey: `test-${suffix}`, + shortcode: `test-${suffix}`, + }, + }); + return { organization, project, runtimeEnvironment }; +} + +async function seedRun( + prisma: PrismaClient, + ids: { id: string; friendlyId: string }, + env: { runtimeEnvironmentId: string; projectId: string; organizationId: string } +) { + return prisma.taskRun.create({ + data: { + id: ids.id, + friendlyId: ids.friendlyId, + taskIdentifier: "my-task", + payload: JSON.stringify({ foo: "bar" }), + payloadType: "application/json", + traceId: "1234", + spanId: "1234", + queue: "test", + runtimeEnvironmentId: env.runtimeEnvironmentId, + projectId: env.projectId, + organizationId: env.organizationId, + environmentType: "DEVELOPMENT", + // V1 so the (best-effort, error-swallowed) cancel does not require the V2 engine; + // the unit under test is the READ resolution, not the cancel side effect. + engine: "V1", + status: "EXECUTING", + }, + }); +} + +// Call-counting instrumentation over the REAL Prisma client. Every findFirst/findMany +// invocation is tallied per delegate and then forwarded to the real underlying method +// (`target[key].apply(target, args)`), so the real query still executes against the +// real container. Every other property passes through untouched. +function countingClient(client: PrismaClient) { + const counts = { findFirst: 0, findMany: 0 }; + const realTaskRun = client.taskRun; + const taskRunProxy = new Proxy(realTaskRun, { + get(target, prop, receiver) { + if (prop === "findFirst" || prop === "findMany") { + const key = prop as "findFirst" | "findMany"; + return (...args: unknown[]) => { + counts[key]++; + return (target[key] as (...a: unknown[]) => unknown).apply(target, args); + }; + } + return Reflect.get(target, prop, receiver); + }, + }); + const clientProxy = new Proxy(client, { + get(target, prop, receiver) { + if (prop === "taskRun") { + return taskRunProxy; + } + return Reflect.get(target, prop, receiver); + }, + }); + return { handle: clientProxy as unknown as PrismaReplicaClient, counts }; +} + +describe("CancelDevSessionRunsService grouped reads", () => { + postgresTest( + "reads N runs in a constant number of grouped queries instead of one findFirst per run id", + async ({ prisma }) => { + const { project, organization, runtimeEnvironment } = await seedOrgProjectEnv( + prisma, + "grouped" + ); + + const seeded = await Promise.all( + ["a", "b", "c"].map((suffix) => + seedRun( + prisma, + { id: `run-internal-${suffix}`, friendlyId: `run_run-internal-${suffix}` }, + { + runtimeEnvironmentId: runtimeEnvironment.id, + projectId: project.id, + organizationId: organization.id, + } + ) + ) + ); + + const counting = countingClient(prisma); + + const service = new CancelDevSessionRunsService({ + prisma, + replica: prisma, + readThroughDeps: { + splitEnabled: false, + newClient: counting.handle, + }, + }); + + await service.call({ + runIds: seeded.map((r) => r.id), + cancelledAt: new Date(), + reason: "test", + }); + + // The N+1 fix: one grouped findMany for the whole id batch, zero per-run findFirst calls. + expect(counting.counts.findMany).toBe(1); + expect(counting.counts.findFirst).toBe(0); + } + ); + + postgresTest( + "a single run id still resolves correctly (grouped read collapses to the same single read)", + async ({ prisma }) => { + const { project, organization, runtimeEnvironment } = await seedOrgProjectEnv( + prisma, + "single" + ); + + const run = await seedRun( + prisma, + { id: "run-internal-solo", friendlyId: "run_run-internal-solo" }, + { + runtimeEnvironmentId: runtimeEnvironment.id, + projectId: project.id, + organizationId: organization.id, + } + ); + + const counting = countingClient(prisma); + + const service = new CancelDevSessionRunsService({ + prisma, + replica: prisma, + readThroughDeps: { + splitEnabled: false, + newClient: counting.handle, + }, + }); + + await service.call({ + runIds: [run.id], + cancelledAt: new Date(), + reason: "test", + }); + + expect(counting.counts.findFirst).toBe(1); + expect(counting.counts.findMany).toBe(0); + } + ); +}); diff --git a/apps/webapp/test/waitpointPresenter.connectedRunsBounded.test.ts b/apps/webapp/test/waitpointPresenter.connectedRunsBounded.test.ts new file mode 100644 index 0000000000..6a0e24b28b --- /dev/null +++ b/apps/webapp/test/waitpointPresenter.connectedRunsBounded.test.ts @@ -0,0 +1,232 @@ +// RED->GREEN guard: WaitpointPresenter#connectedRunIdsOn 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. +// +// 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). +import { describe, expect, vi } from "vitest"; + +const legacyReplicaHolder = vi.hoisted(() => ({ client: undefined as any })); +const newClientHolder = vi.hoisted(() => ({ client: undefined as any })); + +vi.mock("~/db.server", async () => { + const { Prisma } = await import("@trigger.dev/database"); + const lazyProxy = (holder: { client: any }, label: string) => + new Proxy( + {}, + { + get(_t, prop) { + if (!holder.client) { + throw new Error(`${label} not set for this test`); + } + return holder.client[prop]; + }, + } + ); + const replicaProxy = lazyProxy(legacyReplicaHolder, "legacyReplicaHolder.client"); + return { + prisma: replicaProxy, + $replica: replicaProxy, + runOpsNewPrisma: lazyProxy(newClientHolder, "newClientHolder.client"), + runOpsNewReplica: lazyProxy(newClientHolder, "newClientHolder.client"), + runOpsLegacyPrisma: replicaProxy, + runOpsLegacyReplica: replicaProxy, + sqlDatabaseSchema: Prisma.sql([`public`]), + }; +}); + +vi.mock("~/services/clickhouse/clickhouseFactoryInstance.server", () => ({ + clickhouseFactory: { + getClickhouseForOrganization: async () => ({}), + }, +})); + +// Echo the runId set back as runs so the presenter's final CH hydrate never runs for real -- the +// thing under test is the connected-run-id GATHER (the taskRun.findMany args), not this step. +vi.mock("~/presenters/v3/NextRunListPresenter.server", () => ({ + NextRunListPresenter: class { + // eslint-disable-next-line @typescript-eslint/no-unused-vars + constructor(..._args: unknown[]) {} + async call(_organizationId: string, _environmentId: string, opts: { runId?: string[] }) { + return { + runs: (opts.runId ?? []).map((friendlyId) => ({ + friendlyId, + taskIdentifier: "echoed", + })), + }; + } + }, +})); + +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, + WaitpointPresenter, +} from "~/presenters/v3/WaitpointPresenter.server"; + +vi.setConfig({ testTimeout: 90_000 }); + +type SeedContext = { + organizationId: string; + projectId: string; + environmentId: string; +}; + +async function seedParents(prisma: PrismaClient, slug: string): Promise { + const organization = await prisma.organization.create({ + data: { title: `org-${slug}`, slug: `org-${slug}` }, + }); + const project = await prisma.project.create({ + data: { + name: `proj-${slug}`, + slug: `proj-${slug}`, + organizationId: organization.id, + externalRef: `proj-${slug}`, + }, + }); + const runtimeEnvironment = await prisma.runtimeEnvironment.create({ + data: { + slug: `env-${slug}`, + type: "DEVELOPMENT", + projectId: project.id, + organizationId: organization.id, + apiKey: `tr_dev_${slug}`, + pkApiKey: `pk_dev_${slug}`, + shortcode: `sc-${slug}`, + }, + }); + return { + organizationId: organization.id, + projectId: project.id, + environmentId: runtimeEnvironment.id, + }; +} + +async function seedWaitpoint( + prisma: PrismaClient | RunOpsPrismaClient, + ctx: SeedContext, + friendlyId: string +) { + return (prisma as PrismaClient).waitpoint.create({ + data: { + friendlyId, + type: "MANUAL", + status: "COMPLETED", + idempotencyKey: `idem-${friendlyId}`, + userProvidedIdempotencyKey: false, + outputType: "application/json", + outputIsError: false, + completedAt: new Date(), + tags: [], + projectId: ctx.projectId, + environmentId: ctx.environmentId, + }, + }); +} + +async function seedRun( + prisma: PrismaClient | RunOpsPrismaClient, + ctx: SeedContext, + friendlyId: string +) { + return (prisma as PrismaClient).taskRun.create({ + data: { + friendlyId, + taskIdentifier: "my-task", + status: "PENDING", + payload: JSON.stringify({ foo: friendlyId }), + payloadType: "application/json", + traceId: friendlyId, + spanId: friendlyId, + queue: "test", + runtimeEnvironmentId: ctx.environmentId, + projectId: ctx.projectId, + organizationId: ctx.organizationId, + environmentType: "DEVELOPMENT", + engine: "V2", + }, + }); +} + +// Wrap the REAL dedicated run-ops client's `taskRun.findMany` to capture the args of every call, +// delegating unchanged to the real client (the DB still runs the query -- pure instrumentation, +// never a mock). +function capturingTaskRunFindMany(real: RunOpsPrismaClient): { + client: RunOpsPrismaClient; + calls: { where?: { id?: { in?: string[] } } }[]; +} { + const calls: { where?: { id?: { in?: string[] } } }[] = []; + const wrappedTaskRun = new Proxy((real as any).taskRun, { + get(target, prop) { + if (prop === "findMany") { + return (...args: any[]) => { + calls.push(args[0]); + return (target as any)[prop](...args); + }; + } + return (target as any)[prop]; + }, + }); + const client = new Proxy(real as object, { + get(target, prop) { + if (prop === "taskRun") { + return wrappedTaskRun; + } + return (target as any)[prop]; + }, + }) as RunOpsPrismaClient; + return { client, calls }; +} + +const CONNECTED_RUN_COUNT = 8; // > CONNECTED_RUNS_DISPLAY_LIMIT (5), proves the fetch isn't bounded + +describe("WaitpointPresenter#connectedRunIdsOn bounds the connected-run-id FETCH", () => { + heteroRunOpsPostgresTest( + "reading a waitpoint with 8 connected runs never IN-lists more than the display limit", + async ({ prisma14, prisma17 }) => { + const ctx = await seedParents(prisma14, "bounded"); + const waitpoint = await seedWaitpoint(prisma17, ctx, "waitpoint_bounded"); + + const runs = await Promise.all( + Array.from({ length: CONNECTED_RUN_COUNT }, (_, i) => seedRun(prisma17, ctx, `run_b${i}`)) + ); + for (const run of runs) { + await prisma17.waitpointRunConnection.create({ + data: { taskRunId: run.id, waitpointId: waitpoint.id }, + }); + } + + legacyReplicaHolder.client = prisma14; + newClientHolder.client = prisma17; + + const { client: countingPrisma17, calls } = capturingTaskRunFindMany(prisma17); + + const presenter = new WaitpointPresenter(undefined, undefined, { + splitEnabled: true, + newClient: countingPrisma17 as unknown as PrismaClient, + legacyReplica: prisma14, + }); + + await presenter.call({ + friendlyId: waitpoint.friendlyId, + environmentId: ctx.environmentId, + projectId: ctx.projectId, + }); + + // 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. + expect(calls.length).toBeGreaterThan(0); + for (const call of calls) { + expect(call.where?.id?.in?.length ?? 0).toBeLessThanOrEqual(CONNECTED_RUNS_DISPLAY_LIMIT); + } + } + ); +}); diff --git a/apps/webapp/test/waitpointPresenter.splitConnectedRuns.test.ts b/apps/webapp/test/waitpointPresenter.splitConnectedRuns.test.ts new file mode 100644 index 0000000000..2511a25af1 --- /dev/null +++ b/apps/webapp/test/waitpointPresenter.splitConnectedRuns.test.ts @@ -0,0 +1,207 @@ +// FLOW-level gap: a waitpoint's connected runs SPLIT across BOTH physical DBs at once (some legacy, +// some new-resident) rather than the existing single-direction cross-DB cases in +// waitpointPresenter.dedicatedConnectedRuns.readthrough.test.ts. Asserts #connectedRunFriendlyIds +// unions friendlyIds gathered from BOTH stores in one read AND stays bounded to +// CONNECTED_RUNS_DISPLAY_LIMIT even when the combined total exceeds it. Real two-physical-DB +// topology (heteroRunOpsPostgresTest); never mocked beyond the same db.server/clickhouse/ +// NextRunListPresenter echo seams the sibling presenter tests use. +import { describe, expect, vi } from "vitest"; + +const legacyReplicaHolder = vi.hoisted(() => ({ client: undefined as any })); +const newClientHolder = vi.hoisted(() => ({ client: undefined as any })); + +vi.mock("~/db.server", async () => { + const { Prisma } = await import("@trigger.dev/database"); + const lazyProxy = (holder: { client: any }, label: string) => + new Proxy( + {}, + { + get(_t, prop) { + if (!holder.client) { + throw new Error(`${label} not set for this test`); + } + return holder.client[prop]; + }, + } + ); + const replicaProxy = lazyProxy(legacyReplicaHolder, "legacyReplicaHolder.client"); + return { + prisma: replicaProxy, + $replica: replicaProxy, + runOpsNewPrisma: lazyProxy(newClientHolder, "newClientHolder.client"), + runOpsNewReplica: lazyProxy(newClientHolder, "newClientHolder.client"), + runOpsLegacyPrisma: replicaProxy, + runOpsLegacyReplica: replicaProxy, + sqlDatabaseSchema: Prisma.sql([`public`]), + }; +}); + +vi.mock("~/services/clickhouse/clickhouseFactoryInstance.server", () => ({ + clickhouseFactory: { + getClickhouseForOrganization: async () => ({}), + }, +})); + +// Echo the runId set back as runs so `result.connectedRuns` == the friendlyIds the presenter +// gathered cross-DB (isolates the gather from the CH hydrate, as the sibling tests do). +vi.mock("~/presenters/v3/NextRunListPresenter.server", () => ({ + NextRunListPresenter: class { + // eslint-disable-next-line @typescript-eslint/no-unused-vars + constructor(..._args: unknown[]) {} + async call(_organizationId: string, _environmentId: string, opts: { runId?: string[] }) { + return { + runs: (opts.runId ?? []).map((friendlyId) => ({ + friendlyId, + taskIdentifier: "echoed", + })), + }; + } + }, +})); + +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, + WaitpointPresenter, +} from "~/presenters/v3/WaitpointPresenter.server"; + +vi.setConfig({ testTimeout: 90_000 }); + +type SeedContext = { + organizationId: string; + projectId: string; + environmentId: string; +}; + +async function seedParents(prisma: PrismaClient, slug: string): Promise { + const organization = await prisma.organization.create({ + data: { title: `org-${slug}`, slug: `org-${slug}` }, + }); + const project = await prisma.project.create({ + data: { + name: `proj-${slug}`, + slug: `proj-${slug}`, + organizationId: organization.id, + externalRef: `proj-${slug}`, + }, + }); + const runtimeEnvironment = await prisma.runtimeEnvironment.create({ + data: { + slug: `env-${slug}`, + type: "DEVELOPMENT", + projectId: project.id, + organizationId: organization.id, + apiKey: `tr_dev_${slug}`, + pkApiKey: `pk_dev_${slug}`, + shortcode: `sc-${slug}`, + }, + }); + return { + organizationId: organization.id, + projectId: project.id, + environmentId: runtimeEnvironment.id, + }; +} + +async function seedRun( + prisma: PrismaClient | RunOpsPrismaClient, + ctx: SeedContext, + friendlyId: string +) { + return (prisma as PrismaClient).taskRun.create({ + data: { + friendlyId, + taskIdentifier: "my-task", + status: "PENDING", + payload: JSON.stringify({ foo: friendlyId }), + payloadType: "application/json", + traceId: friendlyId, + spanId: friendlyId, + queue: "test", + runtimeEnvironmentId: ctx.environmentId, + projectId: ctx.projectId, + organizationId: ctx.organizationId, + environmentType: "DEVELOPMENT", + engine: "V2", + }, + }); +} + +describe("WaitpointPresenter — connected runs SPLIT across both physical DBs", () => { + heteroRunOpsPostgresTest( + "unions connected-run friendlyIds from both stores and stays bounded to the display limit", + async ({ prisma14, prisma17 }) => { + const ctx = await seedParents(prisma14, "split"); + + // Waitpoint resident on LEGACY. + const waitpoint = await prisma14.waitpoint.create({ + data: { + friendlyId: "waitpoint_split", + type: "MANUAL", + status: "COMPLETED", + idempotencyKey: "idem-waitpoint_split", + userProvidedIdempotencyKey: false, + outputType: "application/json", + outputIsError: false, + completedAt: new Date(), + tags: [], + projectId: ctx.projectId, + environmentId: ctx.environmentId, + }, + }); + + // 2 connected runs resident + joined on NEW (below the limit on its own). + const NEW_RUN_FRIENDLY_IDS = ["run_split_new0", "run_split_new1"]; + for (const friendlyId of NEW_RUN_FRIENDLY_IDS) { + const run = await seedRun(prisma17, ctx, friendlyId); + await prisma17.waitpointRunConnection.create({ + data: { taskRunId: run.id, waitpointId: waitpoint.id }, + }); + } + + // 6 connected runs resident + joined on LEGACY via the implicit M2M — more than enough, on its + // own, to push the combined total past CONNECTED_RUNS_DISPLAY_LIMIT (5). + const LEGACY_RUN_FRIENDLY_IDS = Array.from({ length: 6 }, (_, i) => `run_split_leg${i}`); + for (const friendlyId of LEGACY_RUN_FRIENDLY_IDS) { + const run = await seedRun(prisma14, ctx, friendlyId); + await prisma14.waitpoint.update({ + where: { id: waitpoint.id }, + data: { connectedRuns: { connect: [{ id: run.id }] } }, + }); + } + + legacyReplicaHolder.client = prisma14; + newClientHolder.client = prisma17; + + const presenter = new WaitpointPresenter(undefined, undefined, { + splitEnabled: true, + newClient: prisma17 as unknown as PrismaClient, + legacyReplica: prisma14, + }); + + const result = await presenter.call({ + friendlyId: waitpoint.friendlyId, + environmentId: ctx.environmentId, + projectId: ctx.projectId, + }); + + const returnedIds = result?.connectedRuns.map((r) => r.friendlyId) ?? []; + + // Bounded: the combined total across both DBs (8) must never surface more than the display + // limit, proving the fetch is capped globally, not just per-DB. + expect(returnedIds.length).toBeLessThanOrEqual(CONNECTED_RUNS_DISPLAY_LIMIT); + expect(returnedIds.length).toBeGreaterThan(0); + + // BOTH stores contributed: at least one NEW-resident and one LEGACY-resident connected run + // friendlyId made it into the result. A single-store gather (the bug this guards against) + // would surface only one side. + const fromNew = returnedIds.filter((id) => NEW_RUN_FRIENDLY_IDS.includes(id)); + const fromLegacy = returnedIds.filter((id) => LEGACY_RUN_FRIENDLY_IDS.includes(id)); + expect(fromNew.length).toBeGreaterThan(0); + expect(fromLegacy.length).toBeGreaterThan(0); + expect(fromNew.length + fromLegacy.length).toBe(returnedIds.length); + } + ); +}); diff --git a/internal-packages/run-engine/src/engine/tests/legacyRunNewTokenResumeFlow.test.ts b/internal-packages/run-engine/src/engine/tests/legacyRunNewTokenResumeFlow.test.ts new file mode 100644 index 0000000000..20f68f13d2 --- /dev/null +++ b/internal-packages/run-engine/src/engine/tests/legacyRunNewTokenResumeFlow.test.ts @@ -0,0 +1,302 @@ +// Completion -> resume FLOW test for the direction crossDbTokenBlock.test.ts's writer fix enables but +// never itself exercises end-to-end: a LEGACY (cuid) run blocks on a standalone MANUAL token that is +// NEW-resident (run-ops id) — e.g. the token was minted co-located with a run-ops sibling. This test +// drives the real engine path (blockRunWithWaitpoint -> completeWaitpoint -> continueRunIfUnblocked) +// across the two PHYSICAL DBs to prove the blocking edge is discovered and the run actually resumes, +// not just that the edge write succeeds (that unit-level assertion already lives in +// runOpsStore.crossDbTokenBlock.test.ts). Mirrors completeWaitpointReadResidency.test.ts's own +// cross-DB case, but with the run/token residencies swapped. Real two-physical-DB topology; never +// mocked. + +import { + heteroRunOpsPostgresTest, + network, + redisContainer, + redisOptions, +} from "@internal/testcontainers"; +import { trace } from "@internal/tracing"; +import { PostgresRunStore, RoutingRunStore, type CreateRunInput } from "@internal/run-store"; +import { generateRunOpsId } from "@trigger.dev/core/v3/isomorphic"; +import type { PrismaClient } from "@trigger.dev/database"; +import type { RunOpsPrismaClient } from "@internal/run-ops-database"; +import { expect, vi } from "vitest"; +import { RunEngine } from "../index.js"; + +const twoDbEngineTest = heteroRunOpsPostgresTest.extend<{ + redisContainer: any; + redisOptions: any; +}>({ + network, + redisContainer, + redisOptions, +}); + +// 25-char cuid body (no v1 version marker at index 25) -> classifies LEGACY -> #legacy (prisma14). +const CUID_25 = "c".repeat(25); + +function baseEngineOptions(redisOptions: any, prisma: any) { + return { + prisma, + worker: { redis: redisOptions, workers: 1, tasksPerWorker: 10, pollIntervalMs: 100 }, + queue: { + redis: redisOptions, + masterQueueConsumersDisabled: true, + processWorkerQueueDebounceMs: 50, + }, + runLock: { redis: redisOptions }, + machines: { + defaultMachine: "small-1x" as const, + machines: { + "small-1x": { name: "small-1x" as const, cpu: 0.5, memory: 0.5, centsPerMs: 0.0001 }, + }, + baseCostInCents: 0.0001, + }, + tracer: trace.getTracer("test", "0.0.0"), + }; +} + +async function seedControlPlaneEnv(prisma: PrismaClient, suffix: string) { + const organization = await prisma.organization.create({ + data: { title: `Org ${suffix}`, slug: `org-${suffix}` }, + }); + const project = await prisma.project.create({ + data: { + name: `Project ${suffix}`, + slug: `project-${suffix}`, + externalRef: `proj_${suffix}`, + organizationId: organization.id, + }, + }); + const environment = await prisma.runtimeEnvironment.create({ + data: { + type: "PRODUCTION", + slug: `prod-${suffix}`, + projectId: project.id, + organizationId: organization.id, + apiKey: `tr_prod_${suffix}`, + pkApiKey: `pk_prod_${suffix}`, + shortcode: `short_${suffix}`, + maximumConcurrencyLimit: 10, + }, + }); + return { organization, project, environment }; +} + +function buildCreateRunInput(p: { + runId: string; + friendlyId: string; + organizationId: string; + projectId: string; + runtimeEnvironmentId: string; +}): CreateRunInput { + return { + data: { + id: p.runId, + engine: "V2", + status: "EXECUTING", + friendlyId: p.friendlyId, + runtimeEnvironmentId: p.runtimeEnvironmentId, + environmentType: "PRODUCTION", + organizationId: p.organizationId, + projectId: p.projectId, + taskIdentifier: "parent-task", + payload: "{}", + payloadType: "application/json", + context: {}, + traceContext: {}, + traceId: `trace_${p.runId}`, + spanId: `span_${p.runId}`, + runTags: [], + queue: "task/parent-task", + isTest: false, + taskEventStore: "taskEvent", + depth: 0, + createdAt: new Date("2024-01-01T00:00:00.000Z"), + }, + snapshot: { + engine: "V2", + executionStatus: "RUN_CREATED", + description: "Run was created", + runStatus: "PENDING", + environmentId: p.runtimeEnvironmentId, + environmentType: "PRODUCTION", + projectId: p.projectId, + organizationId: p.organizationId, + }, + }; +} + +// Seed an EXECUTING LEGACY (cuid) parent run on #legacy (prisma14), plus a standalone MANUAL token +// that is NEW-resident (run-ops id, minted co-located with some other run-ops sibling) on #new +// (prisma17) ONLY — the token is never created on #legacy. +async function seedExecutingLegacyParentAndNewToken( + prisma14: PrismaClient, + prisma17: RunOpsPrismaClient, + router: RoutingRunStore, + parentRunId: string, + waitpointId: string, + suffix: string +) { + const env = await seedControlPlaneEnv(prisma14, suffix); + + await router.createRun( + buildCreateRunInput({ + runId: parentRunId, + friendlyId: `run_${suffix}_parent`, + organizationId: env.organization.id, + projectId: env.project.id, + runtimeEnvironmentId: env.environment.id, + }) + ); + + const created = await router.findLatestExecutionSnapshot(parentRunId); + await router.createExecutionSnapshot( + { + run: { id: parentRunId, status: "EXECUTING", attemptNumber: 1 }, + snapshot: { executionStatus: "EXECUTING", description: "parent executing" }, + previousSnapshotId: created!.id, + environmentId: env.environment.id, + environmentType: "PRODUCTION", + projectId: env.project.id, + organizationId: env.organization.id, + }, + prisma14 + ); + + // Standalone MANUAL token lives on #new ONLY (run-ops id) — NOT on #legacy. + await prisma17.waitpoint.create({ + data: { + id: waitpointId, + friendlyId: `wp_${suffix}`, + type: "MANUAL", + status: "PENDING", + idempotencyKey: `idem_${waitpointId}`, + userProvidedIdempotencyKey: false, + projectId: env.project.id, + environmentId: env.environment.id, + }, + }); + + return env; +} + +function makeRouter(prisma14: PrismaClient, prisma17: RunOpsPrismaClient) { + const newStore = new PostgresRunStore({ + prisma: prisma17 as never, + readOnlyPrisma: prisma17 as never, + schemaVariant: "dedicated", + }); + const legacyStore = new PostgresRunStore({ + prisma: prisma14, + readOnlyPrisma: prisma14, + schemaVariant: "legacy", + }); + return new RoutingRunStore({ new: newStore, legacy: legacyStore }); +} + +describe("RunEngine completion->resume flow — a LEGACY run blocked on a NEW-resident token resumes across the seam", () => { + twoDbEngineTest( + "completeWaitpoint on the NEW token finds the #legacy edge and unblocks the LEGACY run", + async ({ prisma14, prisma17, redisOptions }) => { + const router = makeRouter(prisma14 as unknown as PrismaClient, prisma17); + + // The harness builds the schema with `prisma db push`, which re-creates the FKs that the + // run-ops split migrations drop in prod (see runOpsStore.crossDbTokenBlock.test.ts). Mirror + // those drops on this clone so the cross-DB edge write exercises real prod state. + await (prisma14 as unknown as PrismaClient).$executeRawUnsafe( + `ALTER TABLE "TaskRunWaitpoint" DROP CONSTRAINT IF EXISTS "TaskRunWaitpoint_waitpointId_fkey"` + ); + await (prisma14 as unknown as PrismaClient).$executeRawUnsafe( + `ALTER TABLE "_WaitpointRunConnections" DROP CONSTRAINT IF EXISTS "_WaitpointRunConnections_B_fkey"` + ); + // continueRunIfUnblocked's resume snapshot connects the now-COMPLETED (NEW-resident) token via + // _completedWaitpoints (migration 20260705230000 drops this FK in prod; see + // runOpsStore.crossDbCompletedWaitpoint.test.ts). + await (prisma14 as unknown as PrismaClient).$executeRawUnsafe( + `ALTER TABLE "_completedWaitpoints" DROP CONSTRAINT IF EXISTS "_completedWaitpoints_B_fkey"` + ); + + const engine = new RunEngine({ + store: router, + ...baseEngineOptions(redisOptions, prisma14), + }); + + try { + const parentRunId = `run_${CUID_25}`; // LEGACY run -> #legacy + const tokenId = `waitpoint_${generateRunOpsId()}`; // NEW-resident standalone token -> #new + const env = await seedExecutingLegacyParentAndNewToken( + prisma14 as unknown as PrismaClient, + prisma17, + router, + parentRunId, + tokenId, + "xdbresume" + ); + + // Block the LEGACY run on the NEW token. The edge must land on #legacy (the run's own DB, + // FK-free now that the migration-dropped constraints are mirrored above). + await engine.blockRunWithWaitpoint({ + runId: parentRunId, + waitpoints: tokenId, + projectId: env.project.id, + organizationId: env.organization.id, + tx: prisma14 as unknown as PrismaClient, + }); + + expect( + await (prisma14 as unknown as PrismaClient).taskRunWaitpoint.count({ + where: { taskRunId: parentRunId }, + }) + ).toBe(1); + expect(await prisma17.taskRunWaitpoint.count({ where: { taskRunId: parentRunId } })).toBe( + 0 + ); + + const blocked = await engine.getRunExecutionData({ runId: parentRunId }); + expect(blocked?.snapshot.executionStatus).toBe("EXECUTING_WITH_WAITPOINTS"); + + const enqueueSpy = vi.spyOn((engine as any).worker, "enqueue"); + + // Complete the NEW token via the engine path. forWaitpointCompletion resolves #new (run-ops + // id, no pins); completeWaitpoint marks it COMPLETED there, then fans the waitpointId edge + // read across BOTH DBs (findManyTaskRunWaitpoints) to discover the #legacy-resident edge. + const completed = await engine.completeWaitpoint({ + id: tokenId, + output: { value: '{"resumed":"legacy-run-new-token"}', isError: false }, + }); + expect(completed.status).toBe("COMPLETED"); + + // Token COMPLETED on #new only. + expect((await prisma17.waitpoint.findFirst({ where: { id: tokenId } }))?.status).toBe( + "COMPLETED" + ); + expect( + await (prisma14 as unknown as PrismaClient).waitpoint.findFirst({ + where: { id: tokenId }, + }) + ).toBeNull(); + + // The fan-out found the #legacy edge and enqueued the unblock for the LEGACY run. + const continueEnqueued = enqueueSpy.mock.calls.some( + ([arg]) => + (arg as any)?.job === "continueRunIfUnblocked" && + (arg as any)?.payload?.runId === parentRunId + ); + expect(continueEnqueued).toBe(true); + + // Driving the unblock body re-resolves the NEW token's COMPLETED status across the seam + // (continueRunIfUnblocked's blockingWaitpoints read rehydrates `waitpoint` cross-DB) and + // actually resumes the LEGACY run — the higher-level outcome, not just the edge write. + const result = await (engine as any).waitpointSystem.continueRunIfUnblocked({ + runId: parentRunId, + }); + expect(result.status).toBe("unblocked"); + + const after = await engine.getRunExecutionData({ runId: parentRunId }); + expect(after?.snapshot.executionStatus).not.toBe("EXECUTING_WITH_WAITPOINTS"); + } finally { + await engine.quit(); + } + } + ); +}); diff --git a/internal-packages/run-store/src/PostgresRunStore.findRunsByIds.test.ts b/internal-packages/run-store/src/PostgresRunStore.findRunsByIds.test.ts new file mode 100644 index 0000000000..09dbab106d --- /dev/null +++ b/internal-packages/run-store/src/PostgresRunStore.findRunsByIds.test.ts @@ -0,0 +1,170 @@ +// RED→GREEN: `findRunsByIds` is the grouped replacement for `Promise.all(ids.map(id => +// findRun(id)))` — it must issue ONE `findMany` for the whole id batch, never a `findFirst` +// per id. `postgresTest` is a REAL PrismaClient over the standard (non-dedicated) schema. The +// counting proxy wraps its `taskRun` delegate to tally `findFirst`/`findMany` calls while +// delegating every call to the real client — the DB still runs the query; this is +// instrumentation, not a mock. + +import { postgresTest } from "@internal/testcontainers"; +import type { PrismaClient } from "@trigger.dev/database"; +import { describe, expect } from "vitest"; +import { PostgresRunStore } from "./PostgresRunStore.js"; +import type { CreateRunInput } from "./types.js"; + +type CallCounts = { findMany: number; findFirst: number }; + +// Wraps the named delegates on a REAL client to tally `findMany`/`findFirst` calls PER delegate, +// delegating every call unchanged to the real client (the DB still runs the query). +function countingClient( + real: PrismaClient, + delegateNames: string[] +): { client: PrismaClient; counts: Record } { + const counts: Record = {}; + for (const name of delegateNames) { + counts[name] = { findMany: 0, findFirst: 0 }; + } + const wrapped = new Map( + delegateNames.map((name) => [ + name, + new Proxy((real as any)[name], { + get(target, prop) { + if (prop === "findMany" || prop === "findFirst") { + counts[name][prop as "findMany" | "findFirst"]++; + } + return (target as any)[prop]; + }, + }), + ]) + ); + const client = new Proxy(real, { + get(target, prop) { + if (typeof prop === "string" && wrapped.has(prop)) { + return wrapped.get(prop); + } + return (target as any)[prop]; + }, + }) as PrismaClient; + return { client, counts }; +} + +async function seedEnvironment(prisma: PrismaClient) { + const organization = await prisma.organization.create({ + data: { + title: "Test Organization", + slug: "test-organization-findrunsbyids", + }, + }); + + const project = await prisma.project.create({ + data: { + name: "Test Project", + slug: "test-project-findrunsbyids", + externalRef: "proj_findrunsbyids", + organizationId: organization.id, + }, + }); + + const environment = await prisma.runtimeEnvironment.create({ + data: { + type: "DEVELOPMENT", + slug: "dev", + projectId: project.id, + organizationId: organization.id, + apiKey: "tr_dev_apikey_findrunsbyids", + pkApiKey: "pk_dev_apikey_findrunsbyids", + shortcode: "short_code_findrunsbyids", + }, + }); + + return { organization, project, environment }; +} + +function buildCreateRunInput(params: { + runId: string; + friendlyId: string; + organizationId: string; + projectId: string; + runtimeEnvironmentId: string; +}): CreateRunInput { + return { + data: { + id: params.runId, + engine: "V2", + status: "PENDING", + friendlyId: params.friendlyId, + runtimeEnvironmentId: params.runtimeEnvironmentId, + environmentType: "DEVELOPMENT", + organizationId: params.organizationId, + projectId: params.projectId, + taskIdentifier: "my-task", + payload: "{}", + payloadType: "application/json", + traceContext: {}, + traceId: `trace_${params.runId}`, + spanId: `span_${params.runId}`, + queue: "task/my-task", + isTest: false, + taskEventStore: "taskEvent", + depth: 0, + }, + snapshot: { + engine: "V2", + executionStatus: "RUN_CREATED", + description: "Run was created", + runStatus: "PENDING", + environmentId: params.runtimeEnvironmentId, + environmentType: "DEVELOPMENT", + projectId: params.projectId, + organizationId: params.organizationId, + }, + }; +} + +describe("PostgresRunStore.findRunsByIds", () => { + postgresTest( + "resolves N ids with ONE grouped findMany, never a findFirst per id", + async ({ prisma }) => { + const { organization, project, environment } = await seedEnvironment(prisma); + + const seedStore = new PostgresRunStore({ prisma, readOnlyPrisma: prisma }); + + const runId1 = "run_findbyids_1"; + const runId2 = "run_findbyids_2"; + await seedStore.createRun( + buildCreateRunInput({ + runId: runId1, + friendlyId: "run_findbyids_1_friendly", + organizationId: organization.id, + projectId: project.id, + runtimeEnvironmentId: environment.id, + }) + ); + await seedStore.createRun( + buildCreateRunInput({ + runId: runId2, + friendlyId: "run_findbyids_2_friendly", + organizationId: organization.id, + projectId: project.id, + runtimeEnvironmentId: environment.id, + }) + ); + + const counting = countingClient(prisma, ["taskRun"]); + const readStore = new PostgresRunStore({ prisma, readOnlyPrisma: counting.client }); + + const byId = await readStore.findRunsByIds([runId1, runId2], { + select: { friendlyId: true, status: true }, + }); + + expect(byId.size).toBe(2); + expect(byId.get(runId1)?.friendlyId).toBe("run_findbyids_1_friendly"); + expect(byId.get(runId1)?.status).toBe("PENDING"); + expect(byId.get(runId2)?.friendlyId).toBe("run_findbyids_2_friendly"); + expect(byId.get(runId2)?.status).toBe("PENDING"); + + // GROUPED: one findMany for the WHOLE batch, never one findFirst per id. + expect(counting.counts.taskRun.findMany).toBe(1); + expect(counting.counts.taskRun.findFirst).toBe(0); + } + ); +}); diff --git a/internal-packages/run-store/src/PostgresRunStore.groupedRelations.test.ts b/internal-packages/run-store/src/PostgresRunStore.groupedRelations.test.ts new file mode 100644 index 0000000000..948b904748 --- /dev/null +++ b/internal-packages/run-store/src/PostgresRunStore.groupedRelations.test.ts @@ -0,0 +1,228 @@ +// RED→GREEN: the dedicated-schema relation hydrators (`hydrateBlockingTaskRuns` / +// `#hydrateDedicatedRelations`) resolve a relation PER PARENT ROW, so for N parents each requesting +// the same relation, the store issues N round trips (a `findFirst`/`findMany` per row) instead of one +// grouped query for the whole batch. `heteroRunOpsPostgresTest.prisma17` is a REAL RunOpsPrismaClient +// over the dedicated subset schema. The counting proxy wraps its `taskRun` delegate to tally +// `findFirst`/`findMany` calls while delegating every call to the real client — the DB still runs the +// query; this is instrumentation, not a mock. + +import { heteroRunOpsPostgresTest } from "@internal/testcontainers"; +import type { RunOpsPrismaClient } from "@internal/run-ops-database"; +import { describe, expect } from "vitest"; +import { PostgresRunStore } from "./PostgresRunStore.js"; +import type { CreateRunInput } from "./types.js"; + +type CallCounts = { findMany: number; findFirst: number }; + +// Wraps the named delegates on a REAL client to tally `findMany`/`findFirst` calls PER delegate, +// delegating every call unchanged to the real client (the DB still runs the query). +function countingClient( + real: RunOpsPrismaClient, + delegateNames: string[] +): { client: RunOpsPrismaClient; counts: Record } { + const counts: Record = {}; + for (const name of delegateNames) { + counts[name] = { findMany: 0, findFirst: 0 }; + } + const wrapped = new Map( + delegateNames.map((name) => [ + name, + new Proxy((real as any)[name], { + get(target, prop) { + if (prop === "findMany" || prop === "findFirst") { + counts[name][prop as "findMany" | "findFirst"]++; + } + return (target as any)[prop]; + }, + }), + ]) + ); + const client = new Proxy(real, { + get(target, prop) { + if (typeof prop === "string" && wrapped.has(prop)) { + return wrapped.get(prop); + } + return (target as any)[prop]; + }, + }) as RunOpsPrismaClient; + return { client, counts }; +} + +function seedEnvironmentDedicated(suffix: string) { + return { + organization: { id: `org_${suffix}` }, + project: { id: `proj_${suffix}` }, + environment: { id: `env_${suffix}` }, + }; +} + +function buildCreateRunInput(params: { + runId: string; + friendlyId: string; + organizationId: string; + projectId: string; + runtimeEnvironmentId: string; +}): CreateRunInput { + return { + data: { + id: params.runId, + engine: "V2", + status: "PENDING", + friendlyId: params.friendlyId, + runtimeEnvironmentId: params.runtimeEnvironmentId, + environmentType: "DEVELOPMENT", + organizationId: params.organizationId, + projectId: params.projectId, + taskIdentifier: "my-task", + payload: '{"hello":"world"}', + payloadType: "application/json", + traceContext: { trace: "ctx" }, + traceId: `trace_${params.runId}`, + spanId: `span_${params.runId}`, + runTags: [], + queue: "task/my-task", + isTest: false, + taskEventStore: "taskEvent", + depth: 0, + createdAt: new Date("2024-01-01T00:00:00.000Z"), + }, + snapshot: { + engine: "V2", + executionStatus: "RUN_CREATED", + description: "Run was created", + runStatus: "PENDING", + environmentId: params.runtimeEnvironmentId, + environmentType: "DEVELOPMENT", + projectId: params.projectId, + organizationId: params.organizationId, + }, + }; +} + +function makeStore(prisma: RunOpsPrismaClient) { + return new PostgresRunStore({ + prisma: prisma as never, + readOnlyPrisma: prisma as never, + schemaVariant: "dedicated", + }); +} + +describe("PostgresRunStore dedicated relation hydrators — grouped batch reads", () => { + heteroRunOpsPostgresTest( + "hydrateBlockingTaskRuns resolves N edges' `taskRun` with ONE grouped findMany, not N findFirst", + async ({ prisma17 }) => { + const seedStore = makeStore(prisma17); + const env = seedEnvironmentDedicated("hbt"); + const waitpointId = "wp_hbt_shared"; + await prisma17.waitpoint.create({ + data: { + id: waitpointId, + friendlyId: "wp_hbt_friendly", + type: "MANUAL", + status: "PENDING", + idempotencyKey: `idem_${waitpointId}`, + userProvidedIdempotencyKey: false, + projectId: env.project.id, + environmentId: env.environment.id, + }, + }); + + const runIds = ["run_hbt_1", "run_hbt_2", "run_hbt_3"]; + for (const runId of runIds) { + await seedStore.createRun( + buildCreateRunInput({ + runId, + friendlyId: `${runId}_friendly`, + organizationId: env.organization.id, + projectId: env.project.id, + runtimeEnvironmentId: env.environment.id, + }) + ); + await seedStore.blockRunWithWaitpointEdges({ + runId, + waitpointIds: [waitpointId], + projectId: env.project.id, + }); + } + + const counting = countingClient(prisma17, ["taskRun"]); + const readStore = makeStore(counting.client); + + const waitpoint = (await readStore.findWaitpoint({ + where: { id: waitpointId }, + include: { + blockingTaskRuns: { select: { taskRun: { select: { id: true, friendlyId: true } } } }, + }, + })) as { blockingTaskRuns: { taskRun: { id: string; friendlyId: string } | null }[] } | null; + + const blocking = waitpoint?.blockingTaskRuns ?? []; + expect(blocking).toHaveLength(3); + expect(blocking.map((b) => b.taskRun?.id).sort()).toEqual([...runIds].sort()); + + // GROUPED: one findMany for the WHOLE batch, never one findFirst per edge. + expect(counting.counts.taskRun.findMany).toBe(1); + expect(counting.counts.taskRun.findFirst).toBe(0); + } + ); + + heteroRunOpsPostgresTest( + "findManyWaitpoints hydrates N parents' `connectedRuns` with ONE grouped join query + ONE grouped target query", + async ({ prisma17 }) => { + const seedStore = makeStore(prisma17); + const env = seedEnvironmentDedicated("cr"); + + const waitpointIds = ["wp_cr_1", "wp_cr_2", "wp_cr_3"]; + const runIds = ["run_cr_1", "run_cr_2", "run_cr_3"]; + for (const [i, waitpointId] of waitpointIds.entries()) { + await prisma17.waitpoint.create({ + data: { + id: waitpointId, + friendlyId: `${waitpointId}_friendly`, + type: "MANUAL", + status: "PENDING", + idempotencyKey: `idem_${waitpointId}`, + userProvidedIdempotencyKey: false, + projectId: env.project.id, + environmentId: env.environment.id, + }, + }); + await seedStore.createRun( + buildCreateRunInput({ + runId: runIds[i], + friendlyId: `${runIds[i]}_friendly`, + organizationId: env.organization.id, + projectId: env.project.id, + runtimeEnvironmentId: env.environment.id, + }) + ); + // One connection per waitpoint, to its own distinct run. + await seedStore.blockRunWithWaitpointEdges({ + runId: runIds[i], + waitpointIds: [waitpointId], + projectId: env.project.id, + }); + } + + const counting = countingClient(prisma17, ["waitpointRunConnection", "taskRun"]); + const readStore = makeStore(counting.client); + + const rows = (await readStore.findManyWaitpoints({ + where: { id: { in: waitpointIds } }, + include: { connectedRuns: { select: { id: true, friendlyId: true } } }, + })) as { id: string; connectedRuns: { id: string; friendlyId: string }[] }[]; + + expect(rows).toHaveLength(3); + const byWaitpointId = new Map(rows.map((r) => [r.id, r.connectedRuns])); + for (const [i, waitpointId] of waitpointIds.entries()) { + const connected = byWaitpointId.get(waitpointId) ?? []; + expect(connected).toHaveLength(1); + expect(connected[0].id).toBe(runIds[i]); + } + + // GROUPED: one join query + one target query for the WHOLE batch, never one pair per parent. + expect(counting.counts.waitpointRunConnection.findMany).toBe(1); + expect(counting.counts.taskRun.findMany).toBe(1); + expect(counting.counts.taskRun.findFirst).toBe(0); + } + ); +}); diff --git a/internal-packages/run-store/src/PostgresRunStore.ts b/internal-packages/run-store/src/PostgresRunStore.ts index 5d8ea6c8df..c863874220 100644 --- a/internal-packages/run-store/src/PostgresRunStore.ts +++ b/internal-packages/run-store/src/PostgresRunStore.ts @@ -99,13 +99,15 @@ export type PostgresRunStoreOptions = { // A caller sub-select for a relation: `{ select?, include? }` or `true` for a bare `key: true`. type SubProjection = { select?: any; include?: any } | true | undefined; -// Hydrates one dedicated-schema relation for a single parent row, honoring the caller's sub-projection. +// Hydrates one dedicated-schema relation for a WHOLE batch of parent rows in one grouped pass: +// one query for the join/target rows spanning every parent id, never one per parent. Returns a +// Map keyed by parent `id` to the already-defaulted (null / []) hydrated value. type DedicatedRelationHydrator = ( client: RunOpsCapableClient, - parent: Record, + parents: Record[], projection: { select?: any; include?: any } | undefined, store: PostgresRunStore -) => Promise; +) => Promise>; // The dedicated-schema relation keys (with hydrators) for a single Prisma model. type DedicatedRelationSpec = Record; @@ -175,153 +177,192 @@ function stripDedicatedRelations( return { stripped: args, requested }; } -// --- per-model dedicated-schema relation hydrators --- - -// Waitpoint where completedByTaskRunId = run.id (the @unique scalar back-pointer); at most one. -const hydrateAssociatedWaitpoint: DedicatedRelationHydrator = async ( - client, - parent, - projection -) => { - const wp = (await client.waitpoint.findFirst({ - where: { completedByTaskRunId: parent.id as string }, - })) as Record | null; - return applyProjection(wp, projection); -}; - -// Display connections for a run: WaitpointRunConnection → Waitpoint rows. -const hydrateConnectedWaitpoints: DedicatedRelationHydrator = async ( - client, - parent, - projection -) => { - const join = client.waitpointRunConnection; - if (!join) { - return []; +// --- per-model dedicated-schema relation hydrators (batched across the WHOLE parent array) --- + +// Generic to-many relation reached via an explicit join model: one grouped query for the join rows +// spanning every parent id, then one grouped query for the distinct target rows, then an in-memory +// (DB-free) assembly per parent. `joinParentField`/`joinTargetField` name the join row's two FK +// columns; `targetDelegate` is the model the join points at. +async function batchHydrateJoinRelation( + join: RunOpsDelegate<"findMany"> | undefined, + targetDelegate: RunOpsDelegate<"findMany">, + parentIds: string[], + joinParentField: string, + joinTargetField: string, + projection: { select?: any; include?: any } | undefined +): Promise> { + const byParent = new Map(parentIds.map((id) => [id, []])); + if (!join || parentIds.length === 0) { + return byParent; } const links = (await join.findMany({ - where: { taskRunId: parent.id as string }, - select: { waitpointId: true }, - })) as { waitpointId: string }[]; + where: { [joinParentField]: { in: parentIds } }, + select: { [joinParentField]: true, [joinTargetField]: true }, + })) as Record[]; if (links.length === 0) { - return []; + return byParent; } - const rows = (await client.waitpoint.findMany({ - where: { id: { in: links.map((l) => l.waitpointId) } }, + const targetIds = [...new Set(links.map((l) => l[joinTargetField]))]; + const rows = (await targetDelegate.findMany({ + where: { id: { in: targetIds } }, })) as Record[]; - return rows.map((r) => applyProjection(r, projection)); -}; + const byTargetId = new Map(rows.map((r) => [r.id as string, r])); + for (const link of links) { + const target = byTargetId.get(link[joinTargetField]); + const bucket = byParent.get(link[joinParentField]); + if (target && bucket) { + bucket.push(applyProjection(target, projection)); + } + } + return byParent; +} -// Completed waitpoints for a snapshot: CompletedWaitpoint join → Waitpoint rows. -const hydrateCompletedWaitpoints: DedicatedRelationHydrator = async ( +// Waitpoint where completedByTaskRunId = run.id (the @unique scalar back-pointer); at most one. +const hydrateAssociatedWaitpoint: DedicatedRelationHydrator = async ( client, - parent, + parents, projection ) => { - const join = client.completedWaitpoint; - if (!join) { - return []; - } - const links = (await join.findMany({ - where: { snapshotId: parent.id as string }, - select: { waitpointId: true }, - })) as { waitpointId: string }[]; - if (links.length === 0) { - return []; + const parentIds = parents.map((p) => p.id as string); + const byParent = new Map(parentIds.map((id) => [id, null])); + if (parentIds.length === 0) { + return byParent; } const rows = (await client.waitpoint.findMany({ - where: { id: { in: links.map((l) => l.waitpointId) } }, + where: { completedByTaskRunId: { in: parentIds } }, })) as Record[]; - return rows.map((r) => applyProjection(r, projection)); + for (const row of rows) { + const runId = row.completedByTaskRunId as string | undefined; + if (runId && byParent.has(runId)) { + byParent.set(runId, applyProjection(row, projection)); + } + } + return byParent; }; +// Display connections for a run: WaitpointRunConnection → Waitpoint rows. +const hydrateConnectedWaitpoints: DedicatedRelationHydrator = async (client, parents, projection) => + batchHydrateJoinRelation( + client.waitpointRunConnection, + client.waitpoint, + parents.map((p) => p.id as string), + "taskRunId", + "waitpointId", + projection + ); + +// Completed waitpoints for a snapshot: CompletedWaitpoint join → Waitpoint rows. +const hydrateCompletedWaitpoints: DedicatedRelationHydrator = async (client, parents, projection) => + batchHydrateJoinRelation( + client.completedWaitpoint, + client.waitpoint, + parents.map((p) => p.id as string), + "snapshotId", + "waitpointId", + projection + ); + // Runs a waitpoint is blocking: TaskRunWaitpoint rows keyed by waitpointId. A nested `taskRun` // select (the run-engine's getWaitpoint shape) is resolved from the scalar TaskRunWaitpoint.taskRunId. -const hydrateBlockingTaskRuns: DedicatedRelationHydrator = async (client, parent, projection) => { +const hydrateBlockingTaskRuns: DedicatedRelationHydrator = async (client, parents, projection) => { + const parentIds = parents.map((p) => p.id as string); + const byParent = new Map(parentIds.map((id) => [id, []])); + if (parentIds.length === 0) { + return byParent; + } const edges = (await client.taskRunWaitpoint.findMany({ - where: { waitpointId: parent.id as string }, + where: { waitpointId: { in: parentIds } }, })) as Record[]; const nestedTaskRun = projection?.select?.taskRun; - if (!nestedTaskRun) { - return edges; - } - const runProjection = projectionOf(nestedTaskRun as SubProjection); - return Promise.all( - edges.map(async (edge) => { - const run = (await client.taskRun.findFirst({ - where: { id: edge.taskRunId as string }, - })) as Record | null; - return { ...edge, taskRun: applyProjection(run, runProjection) }; - }) - ); + const runProjection = nestedTaskRun ? projectionOf(nestedTaskRun as SubProjection) : undefined; + let byRunId = new Map>(); + if (nestedTaskRun) { + const runIds = [...new Set(edges.map((e) => e.taskRunId as string))]; + const runs = ( + runIds.length > 0 ? await client.taskRun.findMany({ where: { id: { in: runIds } } }) : [] + ) as Record[]; + byRunId = new Map(runs.map((r) => [r.id as string, r])); + } + for (const edge of edges) { + const bucket = byParent.get(edge.waitpointId as string); + if (!bucket) continue; + bucket.push( + nestedTaskRun + ? { + ...edge, + taskRun: applyProjection(byRunId.get(edge.taskRunId as string) ?? null, runProjection), + } + : edge + ); + } + return byParent; }; // Display connections for a waitpoint: WaitpointRunConnection → TaskRun rows. -const hydrateConnectedRuns: DedicatedRelationHydrator = async (client, parent, projection) => { - const join = client.waitpointRunConnection; - if (!join) { - return []; - } - const links = (await join.findMany({ - where: { waitpointId: parent.id as string }, - select: { taskRunId: true }, - })) as { taskRunId: string }[]; - if (links.length === 0) { - return []; - } - const rows = (await client.taskRun.findMany({ - where: { id: { in: links.map((l) => l.taskRunId) } }, - })) as Record[]; - return rows.map((r) => applyProjection(r, projection)); -}; +const hydrateConnectedRuns: DedicatedRelationHydrator = async (client, parents, projection) => + batchHydrateJoinRelation( + client.waitpointRunConnection, + client.taskRun, + parents.map((p) => p.id as string), + "waitpointId", + "taskRunId", + projection + ); // Snapshots that completed a waitpoint: CompletedWaitpoint join → TaskRunExecutionSnapshot rows. const hydrateCompletedExecutionSnapshots: DedicatedRelationHydrator = async ( client, - parent, + parents, projection -) => { - const join = client.completedWaitpoint; - if (!join) { - return []; - } - const links = (await join.findMany({ - where: { waitpointId: parent.id as string }, - select: { snapshotId: true }, - })) as { snapshotId: string }[]; - if (links.length === 0) { - return []; - } - const rows = (await client.taskRunExecutionSnapshot.findMany({ - where: { id: { in: links.map((l) => l.snapshotId) } }, - })) as Record[]; - return rows.map((r) => applyProjection(r, projection)); -}; - -// The waitpoint a block edge points at, resolved from the edge's scalar `waitpointId`. The edge's -// own client only finds a co-resident token; the router re-resolves cross-DB. -const hydrateEdgeWaitpoint: DedicatedRelationHydrator = async (client, parent, projection) => { - const waitpointId = parent.waitpointId as string | undefined; - if (!waitpointId) { - return null; - } - const wp = (await client.waitpoint.findFirst({ - where: { id: waitpointId }, - })) as Record | null; - return applyProjection(wp, projection); -}; +) => + batchHydrateJoinRelation( + client.completedWaitpoint, + client.taskRunExecutionSnapshot, + parents.map((p) => p.id as string), + "waitpointId", + "snapshotId", + projection + ); -// The run a block edge belongs to, resolved from the edge's scalar `taskRunId`. -const hydrateEdgeTaskRun: DedicatedRelationHydrator = async (client, parent, projection) => { - const taskRunId = parent.taskRunId as string | undefined; - if (!taskRunId) { - return null; +// The waitpoint each block edge points at, resolved from its scalar `waitpointId`. The edge's own +// client only finds a co-resident token; the router re-resolves cross-DB. +const hydrateEdgeWaitpoint: DedicatedRelationHydrator = async (client, parents, projection) => + batchHydrateEdgeTarget(client.waitpoint, parents, "waitpointId", projection); + +// The run each block edge belongs to, resolved from its scalar `taskRunId`. +const hydrateEdgeTaskRun: DedicatedRelationHydrator = async (client, parents, projection) => + batchHydrateEdgeTarget(client.taskRun, parents, "taskRunId", projection); + +// Generic to-one relation reached via a scalar FK ON the parent itself (not a join model): one +// grouped target query for every distinct FK value across the batch. +async function batchHydrateEdgeTarget( + targetDelegate: RunOpsDelegate<"findMany">, + parents: Record[], + fkField: string, + projection: { select?: any; include?: any } | undefined +): Promise> { + const byParent = new Map(); + const targetIds: string[] = []; + for (const p of parents) { + byParent.set(p.id as string, null); + const fk = p[fkField] as string | undefined; + if (fk) targetIds.push(fk); + } + if (targetIds.length === 0) { + return byParent; + } + const rows = (await targetDelegate.findMany({ + where: { id: { in: [...new Set(targetIds)] } }, + })) as Record[]; + const byTargetId = new Map(rows.map((r) => [r.id as string, r])); + for (const p of parents) { + const fk = p[fkField] as string | undefined; + if (fk) { + byParent.set(p.id as string, applyProjection(byTargetId.get(fk) ?? null, projection)); + } } - const run = (await client.taskRun.findFirst({ - where: { id: taskRunId }, - })) as Record | null; - return applyProjection(run, projection); -}; + return byParent; +} const TASK_RUN_DEDICATED: DedicatedRelationSpec = { associatedWaitpoint: hydrateAssociatedWaitpoint, @@ -1444,15 +1485,65 @@ export class PostgresRunStore implements RunStore { cursor, ...stripped, })) as Record[]; + await this.#hydrateDedicatedRelations( + prisma as RunOpsCapableClient, + rows, + requested, + TASK_RUN_DEDICATED + ); + return rows; + } + + // Grouped replacement for `Promise.all(ids.map(id => findRun(id)))`: a thin wrapper over + // `findRuns` (so it inherits dedicated-relation hydration + read-your-writes client routing), + // bounding the whole id batch into one round trip instead of one per id. + findRunsByIds( + ids: string[], + args: { select: S }, + client?: ReadClient + ): Promise>>; + findRunsByIds( + ids: string[], + args: { include: I }, + client?: ReadClient + ): Promise>>; + findRunsByIds(ids: string[], client?: ReadClient): Promise>; + async findRunsByIds( + ids: string[], + argsOrClient?: { select?: Prisma.TaskRunSelect; include?: Prisma.TaskRunInclude } | ReadClient, + _client?: ReadClient + ): Promise> { + if (ids.length === 0) { + return new Map(); + } + const hasSelectOrInclude = + argsOrClient != null && + typeof argsOrClient === "object" && + ("select" in argsOrClient || "include" in argsOrClient); + const args = hasSelectOrInclude + ? (argsOrClient as { select?: Prisma.TaskRunSelect; include?: Prisma.TaskRunInclude }) + : undefined; + // Slot recovery mirrors `findRuns`'s overloads: when `argsOrClient` isn't a + // `{ select | include }` object it may itself BE the client (2-arg call) or be undefined + // with the client in the 3rd slot (an explicit `(ids, undefined, client)` call). + const client = + args === undefined ? ((argsOrClient as ReadClient | undefined) ?? _client) : _client; + // Force `id` into the projection so the map can key off it, even when the caller's select + // omits it — `findRuns` would otherwise strip it back out as an added-for-merge-only field. + const projected = args?.select + ? { select: { ...args.select, id: true } } + : args?.include + ? { include: args.include } + : {}; + const rows = (await this.findRuns( + { where: { id: { in: ids } }, ...projected } as Parameters[0], + client + )) as Record[]; + const byId = new Map(); for (const row of rows) { - await this.#hydrateDedicatedRelations( - prisma as RunOpsCapableClient, - row, - requested, - TASK_RUN_DEDICATED - ); + byId.set(row.id as string, row); } - return rows; + return byId; } // --- run-ops persistence --- @@ -1563,14 +1654,12 @@ export class PostgresRunStore implements RunStore { cursor, ...stripped, })) as Record[]; - for (const row of rows) { - await this.#hydrateDedicatedRelations( - prisma as RunOpsCapableClient, - row, - requested, - SNAPSHOT_DEDICATED - ); - } + await this.#hydrateDedicatedRelations( + prisma as RunOpsCapableClient, + rows, + requested, + SNAPSHOT_DEDICATED + ); return rows as Prisma.TaskRunExecutionSnapshotGetPayload[]; } @@ -1949,14 +2038,12 @@ export class PostgresRunStore implements RunStore { cursor, ...stripped, })) as Record[]; - for (const row of rows) { - await this.#hydrateDedicatedRelations( - prisma as RunOpsCapableClient, - row, - requested, - WAITPOINT_DEDICATED - ); - } + await this.#hydrateDedicatedRelations( + prisma as RunOpsCapableClient, + rows, + requested, + WAITPOINT_DEDICATED + ); return rows as Prisma.WaitpointGetPayload[]; } @@ -2019,14 +2106,12 @@ export class PostgresRunStore implements RunStore { cursor, ...stripped, })) as Record[]; - for (const row of rows) { - await this.#hydrateDedicatedRelations( - prisma as RunOpsCapableClient, - row, - requested, - TASK_RUN_WAITPOINT_DEDICATED - ); - } + await this.#hydrateDedicatedRelations( + prisma as RunOpsCapableClient, + rows, + requested, + TASK_RUN_WAITPOINT_DEDICATED + ); return rows as Prisma.TaskRunWaitpointGetPayload[]; } @@ -2267,14 +2352,15 @@ export class PostgresRunStore implements RunStore { if (!row) { return row; } - await this.#hydrateDedicatedRelations(client, row, requested, spec); + await this.#hydrateDedicatedRelations(client, [row], requested, spec); return row; } - // Hydrate each requested dedicated-schema relation key onto `row` in place, honoring the caller's sub-select. + // Hydrate each requested dedicated-schema relation key onto EVERY row in `rows` in ONE grouped + // pass per key (never one query per row), honoring the caller's sub-select. async #hydrateDedicatedRelations( client: RunOpsCapableClient, - row: Record, + rows: Record[], requested: Record, spec: DedicatedRelationSpec ): Promise { @@ -2284,7 +2370,10 @@ export class PostgresRunStore implements RunStore { continue; } const subArgs = requested[key]; - row[key] = await hydrator(client, row, projectionOf(subArgs), this); + const byParentId = await hydrator(client, rows, projectionOf(subArgs), this); + for (const row of rows) { + row[key] = byParentId.get(row.id as string); + } } } diff --git a/internal-packages/run-store/src/runOpsStore.groupedReads.test.ts b/internal-packages/run-store/src/runOpsStore.groupedReads.test.ts new file mode 100644 index 0000000000..fb1fbd3a6a --- /dev/null +++ b/internal-packages/run-store/src/runOpsStore.groupedReads.test.ts @@ -0,0 +1,388 @@ +// RED→GREEN for the grouped-read primitive `RoutingRunStore.findRunsByIds`. +// +// Today, callers that need N runs by id do `Promise.all(ids.map(id => readThroughRun(id)))`, which +// fans out to `findRun` PER ID — each one either a direct `taskRun.findFirst` on the owning store, or +// (for an unclassifiable id) a NEW-then-LEGACY probe. For N ids that is O(N) round trips. +// +// `findRunsByIds` is the grouped replacement: it reuses the router's existing bounded id-set +// machinery (`#findRunsByIdSet` via `findRuns`), which queries NEW for the WHOLE id set in one +// `findMany`, then probes LEGACY in one more `findMany` only for the ids NEW missed. So for any N +// (spread across both stores), the call count is CONSTANT (one grouped `findMany` per store queried), +// never N `findFirst` calls. +// +// `heteroRunOpsPostgresTest` gives the REAL split topology: prisma17 = dedicated subset schema +// (#new), prisma14 = full legacy schema on a SEPARATE physical PG container (#legacy). The counting +// proxy wraps each REAL client's `taskRun` delegate to tally `findFirst`/`findMany` calls while +// delegating every call to the real client — the DB still runs the query; this is instrumentation, +// not a mock. + +import { heteroRunOpsPostgresTest } from "@internal/testcontainers"; +import type { PrismaClient } from "@trigger.dev/database"; +import type { RunOpsPrismaClient } from "@internal/run-ops-database"; +import { describe, expect } from "vitest"; +import { PostgresRunStore } from "./PostgresRunStore.js"; +import { RoutingRunStore } from "./runOpsStore.js"; + +type AnyClient = PrismaClient | RunOpsPrismaClient; + +// ownerEngine classifies a 26-char body with version "1" at index 25 (and a valid base32hex/region +// char elsewhere) as NEW; anything else (including any 25-char id, regardless of content) as LEGACY. +// The distinguishing prefix varies by `n` so multiple ids in one test stay unique. +const CUID_25 = (n: number) => `c${String(n).padStart(2, "0")}`.padEnd(25, "c"); +const NEW_ID_26 = (n: number) => `${String(n).padStart(2, "0")}${"k".repeat(22)}01`; + +type CallCounts = { findMany: number; findFirst: number }; + +// Wrap a REAL client's `taskRun` delegate to tally `findFirst`/`findMany` invocations, delegating +// every call unchanged to the real client (the DB still runs the query — pure instrumentation). +function countingClient(real: C): { client: C; counts: CallCounts } { + const counts: CallCounts = { findMany: 0, findFirst: 0 }; + const countingTaskRun = new Proxy((real as any).taskRun, { + get(target, prop) { + if (prop === "findMany" || prop === "findFirst") { + counts[prop as "findMany" | "findFirst"]++; + } + return (target as any)[prop]; + }, + }); + const client = new Proxy(real, { + get(target, prop) { + if (prop === "taskRun") { + return countingTaskRun; + } + return (target as any)[prop]; + }, + }) as C; + return { client, counts }; +} + +async function seedEnvironmentLegacy(prisma: PrismaClient, suffix: string) { + const organization = await prisma.organization.create({ + data: { title: `Org ${suffix}`, slug: `org-${suffix}` }, + }); + const project = await prisma.project.create({ + data: { + name: `Project ${suffix}`, + slug: `project-${suffix}`, + externalRef: `proj_${suffix}`, + organizationId: organization.id, + }, + }); + const environment = await prisma.runtimeEnvironment.create({ + data: { + type: "DEVELOPMENT", + slug: "dev", + projectId: project.id, + organizationId: organization.id, + apiKey: `tr_dev_${suffix}`, + pkApiKey: `pk_dev_${suffix}`, + shortcode: `short_${suffix}`, + }, + }); + return { organization, project, environment }; +} + +function seedEnvironmentDedicated(suffix: string) { + return { + organization: { id: `org_${suffix}` }, + project: { id: `proj_${suffix}` }, + environment: { id: `env_${suffix}` }, + }; +} + +function taskRunData(opts: { + id: string; + friendlyId: string; + organizationId: string; + projectId: string; + runtimeEnvironmentId: string; +}) { + return { + id: opts.id, + engine: "V2" as const, + status: "PENDING" as const, + friendlyId: opts.friendlyId, + runtimeEnvironmentId: opts.runtimeEnvironmentId, + environmentType: "DEVELOPMENT" as const, + organizationId: opts.organizationId, + projectId: opts.projectId, + taskIdentifier: "my-task", + payload: "{}", + payloadType: "application/json", + traceContext: {}, + traceId: `trace_${opts.id}`, + spanId: `span_${opts.id}`, + queue: "task/my-task", + isTest: false, + taskEventStore: "taskEvent", + depth: 0, + }; +} + +describe("RoutingRunStore.findRunsByIds — grouped, residency-partitioned read of a run id set", () => { + heteroRunOpsPostgresTest( + "returns all N rows with a CONSTANT number of grouped findMany calls (never per-id findFirst)", + async ({ prisma14, prisma17 }) => { + const legacyCounting = countingClient(prisma14); + const newCounting = countingClient(prisma17); + + const legacyStore = new PostgresRunStore({ + prisma: legacyCounting.client, + readOnlyPrisma: legacyCounting.client, + schemaVariant: "legacy", + }); + const newStore = new PostgresRunStore({ + prisma: newCounting.client as never, + readOnlyPrisma: newCounting.client as never, + schemaVariant: "dedicated", + }); + const router = new RoutingRunStore({ new: newStore, legacy: legacyStore }); + + // Seed 2 LEGACY-resident (cuid) runs + 2 NEW-resident (run-ops id) runs — ids spread across + // both stores. + const legEnv = await seedEnvironmentLegacy(prisma14, "grp_leg"); + const newEnv = seedEnvironmentDedicated("grp_new"); + + const legacyIds = [`run_${CUID_25(1)}`, `run_${CUID_25(2)}`]; + const newIds = [`run_${NEW_ID_26(1)}`, `run_${NEW_ID_26(2)}`]; + + for (const [i, id] of legacyIds.entries()) { + await prisma14.taskRun.create({ + data: taskRunData({ + id, + friendlyId: `run_grp_leg_${i}`, + organizationId: legEnv.organization.id, + projectId: legEnv.project.id, + runtimeEnvironmentId: legEnv.environment.id, + }), + }); + } + for (const [i, id] of newIds.entries()) { + await prisma17.taskRun.create({ + data: taskRunData({ + id, + friendlyId: `run_grp_new_${i}`, + organizationId: newEnv.organization.id, + projectId: newEnv.project.id, + runtimeEnvironmentId: newEnv.environment.id, + }), + }); + } + + const allIds = [...legacyIds, ...newIds]; + + // Reset counters after seeding (seeding uses `prisma.taskRun.create`, not find*). + legacyCounting.counts.findMany = 0; + legacyCounting.counts.findFirst = 0; + newCounting.counts.findMany = 0; + newCounting.counts.findFirst = 0; + + const result = await router.findRunsByIds(allIds, { + select: { id: true, friendlyId: true }, + }); + + // All N rows returned, keyed by id. + expect(result.size).toBe(4); + for (const id of allIds) { + expect(result.has(id)).toBe(true); + } + expect(result.get(legacyIds[0])?.friendlyId).toBe("run_grp_leg_0"); + expect(result.get(newIds[0])?.friendlyId).toBe("run_grp_new_0"); + + // GROUPED: exactly one findMany call per store queried (NEW always queried for the whole + // set; LEGACY queried once more for the misses) — never N per-id findFirst calls. + expect(newCounting.counts.findMany).toBe(1); + expect(legacyCounting.counts.findMany).toBe(1); + expect(newCounting.counts.findFirst).toBe(0); + expect(legacyCounting.counts.findFirst).toBe(0); + } + ); +}); + +describe("RoutingRunStore.findManyTaskRunWaitpoints — grouped taskRun hydration", () => { + heteroRunOpsPostgresTest( + "hydrates N edges' `taskRun` with a CONSTANT number of grouped findMany calls (never per-edge findFirst)", + async ({ prisma14, prisma17 }) => { + const legacyCounting = countingClient(prisma14); + const newCounting = countingClient(prisma17); + + const legacyStore = new PostgresRunStore({ + prisma: legacyCounting.client, + readOnlyPrisma: legacyCounting.client, + schemaVariant: "legacy", + }); + const newStore = new PostgresRunStore({ + prisma: newCounting.client as never, + readOnlyPrisma: newCounting.client as never, + schemaVariant: "dedicated", + }); + const router = new RoutingRunStore({ new: newStore, legacy: legacyStore }); + + const legEnv = await seedEnvironmentLegacy(prisma14, "edge_leg"); + const newEnv = seedEnvironmentDedicated("edge_new"); + + const legacyRunIds = [`run_${CUID_25(11)}`, `run_${CUID_25(12)}`]; + const newRunIds = [`run_${NEW_ID_26(11)}`, `run_${NEW_ID_26(12)}`]; + for (const [i, id] of legacyRunIds.entries()) { + await prisma14.taskRun.create({ + data: taskRunData({ + id, + friendlyId: `run_edge_leg_${i}`, + organizationId: legEnv.organization.id, + projectId: legEnv.project.id, + runtimeEnvironmentId: legEnv.environment.id, + }), + }); + } + for (const [i, id] of newRunIds.entries()) { + await prisma17.taskRun.create({ + data: taskRunData({ + id, + friendlyId: `run_edge_new_${i}`, + organizationId: newEnv.organization.id, + projectId: newEnv.project.id, + runtimeEnvironmentId: newEnv.environment.id, + }), + }); + } + + // Block edges are FK-free on the dedicated (#new) subset schema, so they can point at a run on + // EITHER DB — create one edge per run, all on #new. + const allRunIds = [...legacyRunIds, ...newRunIds]; + const edgeIds: string[] = []; + for (const runId of allRunIds) { + const edge = await prisma17.taskRunWaitpoint.create({ + data: { + taskRunId: runId, + waitpointId: `wp_${runId}`, + projectId: newEnv.project.id, + }, + }); + edgeIds.push(edge.id); + } + + legacyCounting.counts.findMany = 0; + legacyCounting.counts.findFirst = 0; + newCounting.counts.findMany = 0; + newCounting.counts.findFirst = 0; + + const rows = await router.findManyTaskRunWaitpoints({ + where: { id: { in: edgeIds } }, + select: { + id: true, + taskRunId: true, + taskRun: { select: { id: true, friendlyId: true } }, + }, + }); + + expect(rows).toHaveLength(4); + const byRunId = new Map(rows.map((r) => [r.taskRunId, r.taskRun])); + expect((byRunId.get(legacyRunIds[0]) as any)?.friendlyId).toBe("run_edge_leg_0"); + expect((byRunId.get(newRunIds[0]) as any)?.friendlyId).toBe("run_edge_new_0"); + + // GROUPED: one findMany per store queried for the WHOLE edge set, never one findFirst per edge. + expect(newCounting.counts.findMany).toBe(1); + expect(legacyCounting.counts.findMany).toBe(1); + expect(newCounting.counts.findFirst).toBe(0); + expect(legacyCounting.counts.findFirst).toBe(0); + } + ); +}); + +describe("RoutingRunStore.findWaitpoint connectedRuns — grouped, order-preserving hydration", () => { + heteroRunOpsPostgresTest( + "hydrates N connectedRuns with a CONSTANT number of grouped findMany calls, in the join's own order", + async ({ prisma14, prisma17 }) => { + const legacyCounting = countingClient(prisma14); + const newCounting = countingClient(prisma17); + + const legacyStore = new PostgresRunStore({ + prisma: legacyCounting.client, + readOnlyPrisma: legacyCounting.client, + schemaVariant: "legacy", + }); + const newStore = new PostgresRunStore({ + prisma: newCounting.client as never, + readOnlyPrisma: newCounting.client as never, + schemaVariant: "dedicated", + }); + const router = new RoutingRunStore({ new: newStore, legacy: legacyStore }); + + const legEnv = await seedEnvironmentLegacy(prisma14, "conn_leg"); + const newEnv = seedEnvironmentDedicated("conn_new"); + const waitpointId = `waitpoint_${CUID_25(20)}`; + await prisma14.waitpoint.create({ + data: { + id: waitpointId, + friendlyId: "wp_conn", + type: "MANUAL", + status: "PENDING", + idempotencyKey: `idem_${waitpointId}`, + userProvidedIdempotencyKey: false, + projectId: legEnv.project.id, + environmentId: legEnv.environment.id, + }, + }); + + const legacyRunIds = [`run_${CUID_25(21)}`, `run_${CUID_25(22)}`]; + const newRunIds = [`run_${NEW_ID_26(21)}`, `run_${NEW_ID_26(22)}`]; + for (const [i, id] of legacyRunIds.entries()) { + await prisma14.taskRun.create({ + data: taskRunData({ + id, + friendlyId: `run_conn_leg_${i}`, + organizationId: legEnv.organization.id, + projectId: legEnv.project.id, + runtimeEnvironmentId: legEnv.environment.id, + }), + }); + await router.blockRunWithWaitpointEdges({ + runId: id, + waitpointIds: [waitpointId], + projectId: legEnv.project.id, + }); + } + for (const [i, id] of newRunIds.entries()) { + await prisma17.taskRun.create({ + data: taskRunData({ + id, + friendlyId: `run_conn_new_${i}`, + organizationId: newEnv.organization.id, + projectId: newEnv.project.id, + runtimeEnvironmentId: newEnv.environment.id, + }), + }); + await router.blockRunWithWaitpointEdges({ + runId: id, + waitpointIds: [waitpointId], + projectId: newEnv.project.id, + }); + } + + // The order the connection join itself returns, independent of any assumption about Postgres's + // internal row order — the fix must preserve exactly THIS order, not resort by residency/leg. + const expectedOrder = await router.findWaitpointConnectedRunIds(waitpointId); + expect(expectedOrder).toHaveLength(4); + + legacyCounting.counts.findMany = 0; + legacyCounting.counts.findFirst = 0; + newCounting.counts.findMany = 0; + newCounting.counts.findFirst = 0; + + const waitpoint = (await router.findWaitpoint({ + where: { id: waitpointId }, + include: { connectedRuns: { select: { id: true, friendlyId: true } } }, + })) as { connectedRuns: { id: string; friendlyId: string }[] } | null; + + const connected = waitpoint?.connectedRuns ?? []; + expect(connected).toHaveLength(4); + expect(connected.map((r) => r.id)).toEqual(expectedOrder); + + // GROUPED: one findMany per store for the WHOLE connected-run set, never one findFirst per run. + expect(newCounting.counts.findMany).toBe(1); + expect(legacyCounting.counts.findMany).toBe(1); + expect(newCounting.counts.findFirst).toBe(0); + expect(legacyCounting.counts.findFirst).toBe(0); + } + ); +}); diff --git a/internal-packages/run-store/src/runOpsStore.ts b/internal-packages/run-store/src/runOpsStore.ts index 27cab57d1a..3341bad9b6 100644 --- a/internal-packages/run-store/src/runOpsStore.ts +++ b/internal-packages/run-store/src/runOpsStore.ts @@ -376,6 +376,55 @@ export class RoutingRunStore implements RunStore { return finalizeRows([...byId.values()], args, addedFields); } + // Canonical grouped replacement for `Promise.all(ids.map(id => readThroughRun(id)))`: reuses + // `findRuns`'s bounded id-set path (`#findRunsByIdSet`), so NEW is queried once for the whole set + // and LEGACY once more only for the misses — never one round trip per id. Returns an id-keyed Map; + // missing/duplicate ids are simply absent/collapsed. + findRunsByIds( + ids: string[], + args: { select: S }, + client?: ReadClient + ): Promise>>; + findRunsByIds( + ids: string[], + args: { include: I }, + client?: ReadClient + ): Promise>>; + findRunsByIds(ids: string[], client?: ReadClient): Promise>; + async findRunsByIds( + ids: string[], + argsOrClient?: + | { select?: Record; include?: Record } + | ReadClient, + _client?: ReadClient + ): Promise> { + if (ids.length === 0) { + return new Map(); + } + const args = selectOrIncludeArgs(argsOrClient); + // Mirrors `readYourWrites`'s slot recovery: when `argsOrClient` isn't a `{select|include}` + // object it may itself BE the client (2-arg call) or be undefined with the client in the + // 3rd slot (an explicit `(ids, undefined, client)` call, e.g. from a relation hydrator). + const client = + args === undefined ? ((argsOrClient as ReadClient | undefined) ?? _client) : _client; + // Force `id` into the projection so the map can key off it, even when the caller's select + // omits it — `findRuns` would otherwise strip it back out as an added-for-merge-only field. + const projected = args?.select + ? { select: { ...args.select, id: true } } + : args?.include + ? { include: args.include } + : {}; + const rows = (await this.findRuns( + { where: { id: { in: ids } }, ...projected } as FindRunsArgs, + client + )) as Record[]; + const byId = new Map(); + for (const row of rows) { + byId.set(row.id as string, row); + } + return byId; + } + // --------------------------------------------------------------------------- // TaskRun-core: update-family — route by run id in params // --------------------------------------------------------------------------- @@ -1123,23 +1172,20 @@ export class RoutingRunStore implements RunStore { } // connectedRuns: the WaitpointRunConnection join co-locates with the run, so read the connected run - // ids from EACH store, then resolve the TaskRun rows across BOTH DBs (findRun routes by id). + // ids from EACH store, then resolve the TaskRun rows in ONE grouped, residency-partitioned read + // (`findRunsByIds`) and reorder to match `runIds` (the join's own order). async #reresolveConnectedRunsCrossDb( waitpointId: string, projection: SubProjection, client?: ReadClient ): Promise { const runIds = await this.findWaitpointConnectedRunIds(waitpointId, client); - const findRun = (this.findRun as (...rest: unknown[]) => Promise).bind(this); const args = projectionAsArgs(projection); - const runs: unknown[] = []; - for (const runId of runIds) { - const run = await findRun({ id: runId }, args, client); - if (run != null) { - runs.push(run); - } - } - return runs; + const findRunsByIds = ( + this.findRunsByIds as (...rest: unknown[]) => Promise> + ).bind(this); + const byId = runIds.length > 0 ? await findRunsByIds(runIds, args, client) : new Map(); + return runIds.map((id) => byId.get(id)).filter((run) => run != null); } // completedExecutionSnapshots: the CompletedWaitpoint join co-locates with the snapshot/run, so read @@ -1300,21 +1346,25 @@ export class RoutingRunStore implements RunStore { } } - // Resolve each edge's `taskRun` from its scalar `taskRunId` across BOTH stores (findRun routes by - // id and falls back NEW→LEGACY). A missing run is left null (display-only callers tolerate it; the - // blocked-run resume path keys off `waitpoint`). + // Resolve every edge's `taskRun` from its scalar `taskRunId` in ONE grouped, residency-partitioned + // read (`findRunsByIds`) rather than one `findRun` per edge. A missing run is left null + // (display-only callers tolerate it; the blocked-run resume path keys off `waitpoint`). async #hydrateEdgeTaskRunsCrossDb( edges: Record[], projection: SubProjection, client?: ReadClient ): Promise { - // Bind to `this`: findRun reaches the private #routeOrNew/#findRunUnrouted members, so an unbound - // reference loses `this` and throws on the first private access. - const findRun = (this.findRun as (...rest: unknown[]) => Promise).bind(this); + const ids = uniqueStrings(edges.map((e) => e.taskRunId)); const args = projectionAsArgs(projection); + // Bind to `this`: findRunsByIds reaches private members, so an unbound reference throws. + const findRunsByIds = ( + this.findRunsByIds as (...rest: unknown[]) => Promise> + ).bind(this); + const byId = + ids.length > 0 ? await findRunsByIds(ids, args, client) : new Map(); for (const edge of edges) { const id = edge.taskRunId as string | undefined; - const run = id ? await findRun({ id }, args, client) : null; + const run = id ? byId.get(id) : undefined; edge.taskRun = applyEdgeProjection((run as Record) ?? null, projection); } } diff --git a/internal-packages/run-store/src/scheduleSeam.test.ts b/internal-packages/run-store/src/scheduleSeam.test.ts new file mode 100644 index 0000000000..43c4ce63bf --- /dev/null +++ b/internal-packages/run-store/src/scheduleSeam.test.ts @@ -0,0 +1,306 @@ +// SCHEDULES over the run-ops cutover SEAM. +// +// schedule-engine (internal-packages/schedule-engine) has NO residency logic of its own: a cron +// schedule fires and calls back into `onTriggerScheduledTask` (apps/webapp/app/v3/scheduleEngine.server.ts), +// which mints the run via the ORDINARY TriggerTaskService -> RunEngine.trigger -> RoutingRunStore.createRun +// path — the same mint path already exercised by internal-packages/run-engine/src/engine/tests/triggerCreateRouting.test.ts +// and the `case 1` / `case 1b` findRuns-by-id-set tests in runOpsStore.mixedResidency.test.ts. There is no +// schedule-specific residency branch to test at the mint step; schedules simply inherit whatever the trigger +// path already does. Case A below adds exactly ONE small confirming test (using the previously-untested +// `scheduleId`/`scheduleInstanceId` scalar fields specifically) rather than re-deriving that generic coverage. +// +// The one piece of schedule-adjacent logic that genuinely spans the seam and had NO hetero coverage is +// `rescheduleRun` (RoutingRunStore.rescheduleRun, runOpsStore.ts:617) — the write delegate behind the +// delayed-run "reschedule" API (apps/webapp/app/v3/services/rescheduleTaskRun.server.ts) and +// `DelayedRunSystem.rescheduleDelayedRun`. It is a pure `#routeForWrite(runId)` mechanical delegate with +// no dedicated mixed-residency test anywhere in the existing matrix (runOpsStore.mixedResidency.test.ts +// covers batch/waitpoint/find methods but not rescheduleRun). Case B below closes that gap. +// +// Real two-physical-DB fixture, NO MOCKS: heteroRunOpsPostgresTest (prisma14 = full control-plane/LEGACY +// schema, prisma17 = RunOpsPrismaClient dedicated NEW-subset schema). + +import { heteroRunOpsPostgresTest } from "@internal/testcontainers"; +import type { PrismaClient } from "@trigger.dev/database"; +import type { RunOpsPrismaClient } from "@internal/run-ops-database"; +import { describe, expect } from "vitest"; +import { PostgresRunStore } from "./PostgresRunStore.js"; +import { RoutingRunStore } from "./runOpsStore.js"; +import type { CreateRunInput, RunStoreSchemaVariant } from "./types.js"; + +type AnyClient = PrismaClient | RunOpsPrismaClient; + +// ownerEngine classifies by internal-id LENGTH after stripping a leading `_` +// (runOpsResidency.ts): 25 chars → cuid → LEGACY, a v1 body (version "1" at index 25 of a 26-char +// body) → run-ops id → NEW. Mirrors the helpers in runOpsStore.mixedResidency.test.ts so ids classify +// the same way here (the documented ID pitfall: a naive generator misclassifies NEW ids as LEGACY). +function cuidLegacy(seed: string): string { + return (seed + "c".repeat(25)).slice(0, 25); // 25 chars → LEGACY (#legacy / prisma14) +} +function runOpsNew(seed: string): string { + return (seed.replace(/[^0-9a-v]/g, "0") + "k".repeat(24)).slice(0, 24) + "01"; // 26 chars, version "1" at index 25 → NEW (#new / prisma17) +} + +async function seedEnvironment( + prisma: AnyClient, + schemaVariant: RunStoreSchemaVariant, + suffix: string +) { + if (schemaVariant === "dedicated") { + return { + organization: { id: `org_${suffix}` }, + project: { id: `proj_${suffix}` }, + environment: { id: `env_${suffix}` }, + }; + } + const organization = await (prisma as PrismaClient).organization.create({ + data: { title: `Org ${suffix}`, slug: `org-${suffix}` }, + }); + const project = await (prisma as PrismaClient).project.create({ + data: { + name: `Project ${suffix}`, + slug: `project-${suffix}`, + externalRef: `proj_${suffix}`, + organizationId: organization.id, + }, + }); + const environment = await (prisma as PrismaClient).runtimeEnvironment.create({ + data: { + type: "DEVELOPMENT", + slug: "dev", + projectId: project.id, + organizationId: organization.id, + apiKey: `tr_dev_${suffix}`, + pkApiKey: `pk_dev_${suffix}`, + shortcode: `short_${suffix}`, + }, + }); + return { organization, project, environment }; +} + +async function seedSharedEnv(prisma14: PrismaClient, suffix: string) { + const legacy = await seedEnvironment(prisma14, "legacy", suffix); + return { + organizationId: legacy.organization.id, + projectId: legacy.project.id, + runtimeEnvironmentId: legacy.environment.id, + environmentId: legacy.environment.id, + }; +} + +function buildCreateRunInput(params: { + runId: string; + friendlyId: string; + organizationId: string; + projectId: string; + runtimeEnvironmentId: string; + status?: "PENDING" | "DELAYED"; + scheduleId?: string; + scheduleInstanceId?: string; +}): CreateRunInput { + return { + data: { + id: params.runId, + engine: "V2", + status: params.status ?? "PENDING", + friendlyId: params.friendlyId, + runtimeEnvironmentId: params.runtimeEnvironmentId, + environmentType: "DEVELOPMENT", + organizationId: params.organizationId, + projectId: params.projectId, + taskIdentifier: "my-scheduled-task", + payload: '{"hello":"world"}', + payloadType: "application/json", + traceContext: { trace: "ctx" }, + traceId: `trace_${params.runId}`, + spanId: `span_${params.runId}`, + runTags: [], + queue: "task/my-scheduled-task", + isTest: false, + taskEventStore: "taskEvent", + depth: 0, + createdAt: new Date("2024-01-01T00:00:00.000Z"), + ...(params.scheduleId && { scheduleId: params.scheduleId }), + ...(params.scheduleInstanceId && { scheduleInstanceId: params.scheduleInstanceId }), + }, + snapshot: { + engine: "V2", + executionStatus: "RUN_CREATED", + description: "Run was created", + runStatus: params.status ?? "PENDING", + environmentId: params.runtimeEnvironmentId, + environmentType: "DEVELOPMENT", + projectId: params.projectId, + organizationId: params.organizationId, + }, + }; +} + +function makeDedicatedStore(prisma17: RunOpsPrismaClient) { + return new PostgresRunStore({ + prisma: prisma17 as never, + readOnlyPrisma: prisma17 as never, + schemaVariant: "dedicated", + }); +} + +function makeLegacyStore(prisma14: PrismaClient) { + return new PostgresRunStore({ + prisma: prisma14, + readOnlyPrisma: prisma14, + schemaVariant: "legacy", + }); +} + +// The REAL production split topology: #new = dedicated subset on prisma17, #legacy = full schema on +// prisma14. +function makeSplitRouter(prisma14: PrismaClient, prisma17: RunOpsPrismaClient) { + const legacyStore = makeLegacyStore(prisma14); + const newStore = makeDedicatedStore(prisma17); + return { + router: new RoutingRunStore({ new: newStore, legacy: legacyStore }), + legacyStore, + newStore, + }; +} + +describe("Schedules over the cutover seam", () => { + // ── Case A: a schedule-minted run (scheduleId/scheduleInstanceId set) lands on the correct + // physical store under mixed residency, and is found by plain id lookup (the ONLY lookup that + // exists — there is no scheduleInstanceId-keyed TaskRun query anywhere in the codebase; the field + // is write-only metadata stamped at creation). This confirms schedules have no distinct mint-routing + // seam beyond the already-covered generic createRun routing. ── + heteroRunOpsPostgresTest( + "case A: a schedule-minted run routes to the owning store and resolves by id, for both residencies", + async ({ prisma14, prisma17 }) => { + const { router } = makeSplitRouter(prisma14, prisma17); + const env = await seedSharedEnv(prisma14, "schedA"); + + const legacyRun = cuidLegacy("schAl"); // pre-cutover mint shape → #legacy + const newRun = runOpsNew("schAn"); // run-ops mint shape → #new + + await router.createRun( + buildCreateRunInput({ + runId: legacyRun, + friendlyId: "run_schedA_legacy", + scheduleId: "sched_A", + scheduleInstanceId: "schedinst_A_legacy", + ...env, + }) + ); + await router.createRun( + buildCreateRunInput({ + runId: newRun, + friendlyId: "run_schedA_new", + scheduleId: "sched_A", + scheduleInstanceId: "schedinst_A_new", + ...env, + }) + ); + + // Physical residency: each landed on its OWN DB only, with the scheduleInstanceId intact. + const onLegacy = await prisma14.taskRun.findUnique({ where: { id: legacyRun } }); + expect(onLegacy?.scheduleInstanceId).toBe("schedinst_A_legacy"); + expect(await prisma17.taskRun.findUnique({ where: { id: legacyRun } })).toBeNull(); + + const onNew = await prisma17.taskRun.findUnique({ where: { id: newRun } }); + expect(onNew?.scheduleInstanceId).toBe("schedinst_A_new"); + expect(await prisma14.taskRun.findUnique({ where: { id: newRun } })).toBeNull(); + + // The only lookup a schedule-minted run ever needs (by id) resolves through the router on + // BOTH residencies. + const foundLegacy = (await router.findRun( + { id: legacyRun }, + { select: { id: true, scheduleInstanceId: true } } + )) as Record | null; + expect(foundLegacy?.id).toBe(legacyRun); + expect(foundLegacy?.scheduleInstanceId).toBe("schedinst_A_legacy"); + + const foundNew = (await router.findRun( + { id: newRun }, + { select: { id: true, scheduleInstanceId: true } } + )) as Record | null; + expect(foundNew?.id).toBe(newRun); + expect(foundNew?.scheduleInstanceId).toBe("schedinst_A_new"); + } + ); + + // ── Case B: rescheduleRun (the write delegate behind the delayed-run "reschedule" API and + // DelayedRunSystem.rescheduleDelayedRun) routes to the OWNING store for a mixed-residency + // population, and does NOT touch the other physical DB. No existing hetero test covers this + // write path. ── + heteroRunOpsPostgresTest( + "case B: rescheduleRun routes the write to the owning store only, for a mixed-residency population", + async ({ prisma14, prisma17 }) => { + const { router } = makeSplitRouter(prisma14, prisma17); + const env = await seedSharedEnv(prisma14, "schedB"); + + const legacyRun = cuidLegacy("schBl"); + const newRun = runOpsNew("schBn"); + + await router.createRun( + buildCreateRunInput({ + runId: legacyRun, + friendlyId: "run_schedB_legacy", + status: "DELAYED", + ...env, + }) + ); + await router.createRun( + buildCreateRunInput({ + runId: newRun, + friendlyId: "run_schedB_new", + status: "DELAYED", + ...env, + }) + ); + + const legacyDelayUntil = new Date("2027-05-01T00:00:00.000Z"); + const newDelayUntil = new Date("2027-06-01T00:00:00.000Z"); + + const updatedLegacy = await router.rescheduleRun(legacyRun, { + delayUntil: legacyDelayUntil, + snapshot: { + environmentId: env.environmentId, + environmentType: "DEVELOPMENT", + projectId: env.projectId, + organizationId: env.organizationId, + }, + }); + const updatedNew = await router.rescheduleRun(newRun, { + delayUntil: newDelayUntil, + snapshot: { + environmentId: env.environmentId, + environmentType: "DEVELOPMENT", + projectId: env.projectId, + organizationId: env.organizationId, + }, + }); + + expect(updatedLegacy.id).toBe(legacyRun); + expect(updatedLegacy.delayUntil).toEqual(legacyDelayUntil); + expect(updatedNew.id).toBe(newRun); + expect(updatedNew.delayUntil).toEqual(newDelayUntil); + + // The write landed on the OWNING physical DB only. + const legacyRow = await prisma14.taskRun.findUnique({ where: { id: legacyRun } }); + expect(legacyRow?.delayUntil).toEqual(legacyDelayUntil); + const legacySnapshots = await prisma14.taskRunExecutionSnapshot.findMany({ + where: { runId: legacyRun, executionStatus: "DELAYED" }, + }); + expect(legacySnapshots).toHaveLength(1); + + const newRow = await prisma17.taskRun.findUnique({ where: { id: newRun } }); + expect(newRow?.delayUntil).toEqual(newDelayUntil); + const newSnapshots = await prisma17.taskRunExecutionSnapshot.findMany({ + where: { runId: newRun, executionStatus: "DELAYED" }, + }); + expect(newSnapshots).toHaveLength(1); + + // Cross-DB isolation: the OTHER physical DB's row for each run is untouched by the other + // reschedule call — neither run exists on the non-owning DB at all, so there is nothing there + // to have been (mis)updated. + expect(await prisma17.taskRun.findUnique({ where: { id: legacyRun } })).toBeNull(); + expect(await prisma14.taskRun.findUnique({ where: { id: newRun } })).toBeNull(); + } + ); +}); diff --git a/internal-packages/run-store/src/types.ts b/internal-packages/run-store/src/types.ts index 1d39498f39..26ccb52aea 100644 --- a/internal-packages/run-store/src/types.ts +++ b/internal-packages/run-store/src/types.ts @@ -573,6 +573,21 @@ export interface RunStore { client?: ReadClient ): Promise; + // Grouped replacement for `Promise.all(ids.map(id => findRun(id)))`: one round trip for the + // whole id batch instead of one per id. Returns an id-keyed Map; missing/duplicate ids are + // simply absent/collapsed. + findRunsByIds( + ids: string[], + args: { select: S }, + client?: ReadClient + ): Promise>>; + findRunsByIds( + ids: string[], + args: { include: I }, + client?: ReadClient + ): Promise>>; + findRunsByIds(ids: string[], client?: ReadClient): Promise>; + // --- run-ops persistence --- // Snapshots, waitpoints, implicit M:N joins, dependents, attempts and checkpoints. The // generic model wrappers are thin generics over the Prisma `*Args` types so include/select From e7fb6bdd527da5ecbff9664b22d016944f5cfc84 Mon Sep 17 00:00:00 2001 From: Dan Sutton Date: Thu, 9 Jul 2026 15:50:18 +0100 Subject: [PATCH 2/2] fix(run-store): drop force-injected id from findRunsByIds results when unrequested findRunsByIds forces `id` into the select so it can key the result map, but did not remove it afterwards, so returned rows carried an `id` the declared payload type did not include. Delete the injected id from each value when the caller's select did not request it; the map stays keyed by id. --- .../run-store/src/PostgresRunStore.findRunsByIds.test.ts | 5 +++++ internal-packages/run-store/src/PostgresRunStore.ts | 9 ++++++++- .../run-store/src/runOpsStore.groupedReads.test.ts | 6 +++++- internal-packages/run-store/src/runOpsStore.ts | 9 ++++++++- 4 files changed, 26 insertions(+), 3 deletions(-) diff --git a/internal-packages/run-store/src/PostgresRunStore.findRunsByIds.test.ts b/internal-packages/run-store/src/PostgresRunStore.findRunsByIds.test.ts index 09dbab106d..10e9aa8086 100644 --- a/internal-packages/run-store/src/PostgresRunStore.findRunsByIds.test.ts +++ b/internal-packages/run-store/src/PostgresRunStore.findRunsByIds.test.ts @@ -162,6 +162,11 @@ describe("PostgresRunStore.findRunsByIds", () => { expect(byId.get(runId2)?.friendlyId).toBe("run_findbyids_2_friendly"); expect(byId.get(runId2)?.status).toBe("PENDING"); + // The select omitted `id`; the id we force-inject for keying must not leak into the value + // (it would otherwise violate the declared payload type and expose an unrequested id). + expect("id" in (byId.get(runId1) as object)).toBe(false); + expect("id" in (byId.get(runId2) as object)).toBe(false); + // GROUPED: one findMany for the WHOLE batch, never one findFirst per id. expect(counting.counts.taskRun.findMany).toBe(1); expect(counting.counts.taskRun.findFirst).toBe(0); diff --git a/internal-packages/run-store/src/PostgresRunStore.ts b/internal-packages/run-store/src/PostgresRunStore.ts index c863874220..974034ff27 100644 --- a/internal-packages/run-store/src/PostgresRunStore.ts +++ b/internal-packages/run-store/src/PostgresRunStore.ts @@ -1540,8 +1540,15 @@ export class PostgresRunStore implements RunStore { client )) as Record[]; const byId = new Map(); + // Strip the id we force-injected for map keying when the caller's select did not ask for it, + // so returned values match the declared payload type and never leak an unrequested id. + const stripInjectedId = args?.select != null && !("id" in args.select); for (const row of rows) { - byId.set(row.id as string, row); + const key = row.id as string; + if (stripInjectedId) { + delete row.id; + } + byId.set(key, row); } return byId; } diff --git a/internal-packages/run-store/src/runOpsStore.groupedReads.test.ts b/internal-packages/run-store/src/runOpsStore.groupedReads.test.ts index fb1fbd3a6a..4b2230bec0 100644 --- a/internal-packages/run-store/src/runOpsStore.groupedReads.test.ts +++ b/internal-packages/run-store/src/runOpsStore.groupedReads.test.ts @@ -178,7 +178,7 @@ describe("RoutingRunStore.findRunsByIds — grouped, residency-partitioned read newCounting.counts.findFirst = 0; const result = await router.findRunsByIds(allIds, { - select: { id: true, friendlyId: true }, + select: { friendlyId: true }, }); // All N rows returned, keyed by id. @@ -189,6 +189,10 @@ describe("RoutingRunStore.findRunsByIds — grouped, residency-partitioned read expect(result.get(legacyIds[0])?.friendlyId).toBe("run_grp_leg_0"); expect(result.get(newIds[0])?.friendlyId).toBe("run_grp_new_0"); + // The select omitted `id`; the id we force-inject for keying must not leak into the value. + expect("id" in (result.get(legacyIds[0]) as object)).toBe(false); + expect("id" in (result.get(newIds[0]) as object)).toBe(false); + // GROUPED: exactly one findMany call per store queried (NEW always queried for the whole // set; LEGACY queried once more for the misses) — never N per-id findFirst calls. expect(newCounting.counts.findMany).toBe(1); diff --git a/internal-packages/run-store/src/runOpsStore.ts b/internal-packages/run-store/src/runOpsStore.ts index 3341bad9b6..1884742356 100644 --- a/internal-packages/run-store/src/runOpsStore.ts +++ b/internal-packages/run-store/src/runOpsStore.ts @@ -419,8 +419,15 @@ export class RoutingRunStore implements RunStore { client )) as Record[]; const byId = new Map(); + // Strip the id we force-injected for map keying when the caller's select did not ask for it, + // so returned values match the declared payload type and never leak an unrequested id. + const stripInjectedId = args?.select != null && !("id" in (args.select as object)); for (const row of rows) { - byId.set(row.id as string, row); + const key = row.id as string; + if (stripInjectedId) { + delete row.id; + } + byId.set(key, row); } return byId; }