diff --git a/src/lib/jobs/integration-backfill-drain.ts b/src/lib/jobs/integration-backfill-drain.ts index e2cacf471..76f91c4d6 100644 --- a/src/lib/jobs/integration-backfill-drain.ts +++ b/src/lib/jobs/integration-backfill-drain.ts @@ -115,7 +115,10 @@ async function runAdmitted( return; } case "lab-biomarker-backfill": { - const { markers, linked } = await runLabBiomarkerBackfillForUser(userId); + const { markers, linked } = await runLabBiomarkerBackfillForUser( + userId, + shouldStop, + ); totals.markers += markers; totals.linked += linked; workerLog( diff --git a/src/lib/jobs/lab-biomarker-backfill.ts b/src/lib/jobs/lab-biomarker-backfill.ts index 351a161d1..61ffd162f 100644 --- a/src/lib/jobs/lab-biomarker-backfill.ts +++ b/src/lib/jobs/lab-biomarker-backfill.ts @@ -103,9 +103,16 @@ export function groupRowsByAnalyte(rows: LegacyRow[]): AnalyteGroup[] { * Backfill one user: group un-linked live readings, create/reuse a Biomarker * per group, link the rows. Returns the count of markers touched + rows * linked. + * + * `shouldStop` is the admission job's budget (`jobBudget`). It is asked + * before each analyte group; once it says stop, the pass throws without + * touching the rest, like the Google Health backfill does. Each group commits + * on its own and a linked reading drops out of the scan above, so the retry + * pg-boss schedules for the failed job starts where this one stopped. */ export async function runLabBiomarkerBackfillForUser( userId: string, + shouldStop: () => boolean = () => false, ): Promise<{ markers: number; linked: number }> { const rows = await prisma.labResult.findMany({ where: { userId, deletedAt: null, biomarkerId: null }, @@ -129,6 +136,18 @@ export async function runLabBiomarkerBackfillForUser( let markers = 0; let linked = 0; for (const group of groups) { + if (shouldStop()) { + if (linked > 0) invalidateUserHealthScore(userId); + annotate({ + action: { + name: "labs.biomarker.backfill.stopped", + details: { markers, linked, groupsLeft: groups.length - markers }, + }, + }); + throw new Error( + `lab-biomarker backfill for user ${userId} stopped on its job budget after ${markers} of ${groups.length} groups; the retry resumes`, + ); + } // Re-use a hand-defined marker of the same name if one exists; otherwise // mint it. `(userId, name)` is unique, so a concurrent re-run converges. const biomarker = await prisma.biomarker.upsert({ diff --git a/tests/integration/lab-biomarker-backfill-budget.test.ts b/tests/integration/lab-biomarker-backfill-budget.test.ts new file mode 100644 index 000000000..b24e45b05 --- /dev/null +++ b/tests/integration/lab-biomarker-backfill-budget.test.ts @@ -0,0 +1,76 @@ +/** + * The lab-biomarker backfill stops on its job budget between analyte groups, + * against a real Postgres. + * + * A stop must leave every group it finished linked and every other reading + * untouched, and the retry must pick up exactly the rest: a linked reading + * drops out of the backfill's scan, so nothing is linked twice and nothing is + * skipped. + */ +import { beforeEach, describe, expect, it } from "vitest"; + +import { getPrismaClient, truncateAllTables } from "./setup"; +import { runLabBiomarkerBackfillForUser } from "@/lib/jobs/lab-biomarker-backfill"; + +const USER_ID = "user-lab-backfill-budget"; + +beforeEach(async () => { + const prisma = getPrismaClient(); + await truncateAllTables(prisma); + await prisma.user.create({ + data: { + id: USER_ID, + username: "lab-budget", + email: "lab-budget@example.test", + }, + }); + const takenAt = new Date("2026-06-01T08:00:00.000Z"); + await prisma.labResult.createMany({ + data: ["LDL", "LDL", "HbA1c", "TSH"].map((analyte, i) => ({ + userId: USER_ID, + analyte, + unit: "mg/dL", + value: 100 + i, + takenAt: new Date(takenAt.getTime() + i * 86_400_000), + })), + }); +}); + +async function unlinked() { + return getPrismaClient().labResult.count({ + where: { userId: USER_ID, biomarkerId: null }, + }); +} + +describe("lab-biomarker backfill under a job budget", () => { + it("stops after the group in hand and the retry links the rest", async () => { + let asked = 0; + // Allow exactly one group, then say stop. + const shouldStop = () => asked++ >= 1; + + await expect( + runLabBiomarkerBackfillForUser(USER_ID, shouldStop), + ).rejects.toThrow(/stopped on its job budget/); + + const prisma = getPrismaClient(); + expect(await prisma.biomarker.count({ where: { userId: USER_ID } })).toBe( + 1, + ); + const linkedAfterStop = 4 - (await unlinked()); + // One group: either the two LDL rows or one single-row analyte. + expect([1, 2]).toContain(linkedAfterStop); + + const retry = await runLabBiomarkerBackfillForUser(USER_ID); + expect(retry.linked).toBe(4 - linkedAfterStop); + expect(await unlinked()).toBe(0); + expect(await prisma.biomarker.count({ where: { userId: USER_ID } })).toBe( + 3, + ); + }); + + it("runs to the end when the budget never stops it", async () => { + const result = await runLabBiomarkerBackfillForUser(USER_ID, () => false); + expect(result).toEqual({ markers: 3, linked: 4 }); + expect(await unlinked()).toBe(0); + }); +});