HealthLog/tests/integration/backup-restore-statement-timeout.test.ts

/**
 * #1031 — a restore is not bounded by the per-request statement timeout.
 *
 * Every pooled connection starts with `statement_timeout` and
 * `idle_in_transaction_session_timeout` from `DATABASE_STATEMENT_TIMEOUT_MS`
 * (`src/lib/db.ts`, 60 s by default). A restore deletes the account's
 * readings in one statement, and on an account of 1.9 million readings on a
 * slow host that statement ran past 60 s, was cancelled, and the restore
 * rolled back in its clearing step with a message that named no cause.
 *
 * This file lowers the session limit to five seconds and makes the clearing
 * delete take six (a statement trigger that sleeps), which is the same
 * situation at a size a test can hold. Five, not one: the suite's own
 * truncation between tests runs under the same limit. The first test checks the limit
 * really cancels that delete outside the restore, so the second cannot pass
 * because nothing was slow. The restore then has to finish and give back
 * exactly the rows the backup holds.
 *
 * The last test makes the database refuse the clearing step with a disk-full
 * error and checks the job says so, and that nothing changed.
 *
 * Mutation check: remove the `set_config('statement_timeout', …)` call at the
 * top of the restore transaction and the restore test goes red with
 * `cause: "timeout"`.
 */
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";

vi.hoisted(() => {
  // Read once, when `@/lib/db` builds its pool; hoisted above every import.
  process.env.DATABASE_STATEMENT_TIMEOUT_MS = "5000";
});

import { cookieJar, headerJar } from "./mock-next-headers";
import { getPrismaClient, truncateAllTables } from "./setup";
import { streamFullBackupJson } from "@/lib/export/full-backup-stream";
import { storeBackupBlob } from "@/lib/export/store-backup-blob";
import { RESTORE_FAILURE_CAUSE_MESSAGES } from "@/lib/export/restore-failure-cause";
import { readBackupRestoreJob } from "@/lib/jobs/backup-restore";

vi.mock("next/headers", async () => {
  const { cookieJar, headerJar } = await import("./mock-next-headers");
  return {
    headers: vi.fn(async () => ({
      get: (name: string) => headerJar.get(name.toLowerCase()) ?? null,
    })),
    cookies: vi.fn(async () => ({
      get: (name: string) => {
        const value = cookieJar.get(name);
        return value ? { name, value } : undefined;
      },
      set: (name: string, value: string) => {
        cookieJar.set(name, value);
      },
      delete: (name: string) => {
        cookieJar.delete(name);
      },
    })),
  };
});
vi.mock("@/lib/db-compat", () => ({
  ensureDbCompatibility: vi.fn().mockResolvedValue(undefined),
}));

const COUNT = 600;

async function dropTrigger() {
  const prisma = getPrismaClient();
  await prisma.$executeRawUnsafe(
    `DROP TRIGGER IF EXISTS restore_test_clearing ON measurements`,
  );
  await prisma.$executeRawUnsafe(
    `DROP FUNCTION IF EXISTS restore_test_clearing()`,
  );
}

/** Run `body` on every DELETE statement against `measurements`. */
async function onMeasurementDelete(body: string) {
  const prisma = getPrismaClient();
  await prisma.$executeRawUnsafe(
    `CREATE FUNCTION restore_test_clearing() RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN ${body} RETURN NULL; END $$`,
  );
  await prisma.$executeRawUnsafe(
    `CREATE TRIGGER restore_test_clearing BEFORE DELETE ON measurements FOR EACH STATEMENT EXECUTE FUNCTION restore_test_clearing()`,
  );
}

beforeEach(async () => {
  await dropTrigger();
  await truncateAllTables(getPrismaClient());
  cookieJar.clear();
  headerJar.clear();
});

afterEach(async () => {
  await dropTrigger();
});

async function seedAccountWithBackup() {
  const prisma = getPrismaClient();
  const user = await prisma.user.create({
    data: {
      username: "slow-host",
      email: "slow-host@example.test",
      role: "ADMIN",
    },
  });
  const session = await prisma.session.create({
    data: { userId: user.id, expiresAt: new Date(Date.now() + 600_000) },
  });
  cookieJar.set("healthlog_session", session.id);
  const start = Date.UTC(2024, 0, 1);
  await prisma.measurement.createMany({
    data: Array.from({ length: COUNT }, (_, i) => ({
      id: `m-${String(i).padStart(4, "0")}`,
      userId: user.id,
      type: "PULSE" as const,
      value: 50 + (i % 40),
      unit: "bpm",
      source: "APPLE_HEALTH" as const,
      measuredAt: new Date(start + i * 60_000),
      externalId: `hk-${i}`,
      // Most of them deleted, as after the nightly consolidation.
      deletedAt: i % 8 === 0 ? null : new Date(start),
    })),
  });
  const { id } = await storeBackupBlob(
    prisma,
    { userId: user.id, type: "WEEKLY_AUTO" },
    (write) =>
      streamFullBackupJson(prisma, user.id, write, {
        purpose: "disaster-recovery",
      }),
  );
  const before = await measurementsOf(user.id);
  // Something for the restore to undo, without a DELETE.
  await prisma.measurement.updateMany({
    where: { userId: user.id },
    data: { value: 1 },
  });
  return { user, backupId: id, before };
}

async function measurementsOf(userId: string) {
  return getPrismaClient().measurement.findMany({
    where: { userId },
    orderBy: { id: "asc" },
  });
}

async function restore(id: string) {
  const { POST } = await import("./restore-job-driver");
  return POST(
    new Request(`http://localhost/api/admin/backups/${id}/restore`, {
      method: "POST",
      headers: { "content-type": "application/json" },
      body: JSON.stringify({ confirm: "RESTORE" }),
    }) as never,
    { params: Promise.resolve({ id }) },
  );
}

async function lastJob(userId: string) {
  const row = await getPrismaClient().backupRestoreJob.findFirstOrThrow({
    where: { userId },
    orderBy: { createdAt: "desc" },
    select: { id: true },
  });
  return readBackupRestoreJob(row.id);
}

describe("restore under a short statement timeout (#1031)", () => {
  it("the session limit cancels the slow clearing delete outside a restore", async () => {
    const { user } = await seedAccountWithBackup();
    await onMeasurementDelete("PERFORM pg_sleep(6);");
    const prisma = getPrismaClient();
    const attempt = prisma.$transaction(async (tx) => {
      await tx.measurement.deleteMany({ where: { userId: user.id } });
      throw new Error("not cancelled");
    });
    await expect(attempt).rejects.toMatchObject({
      meta: { driverAdapterError: { cause: { originalCode: "57014" } } },
    });
  });

  it("a restore whose clearing delete outlasts the session limit still completes", async () => {
    const { user, backupId, before } = await seedAccountWithBackup();
    await onMeasurementDelete("PERFORM pg_sleep(6);");

    const res = await restore(backupId);
    const job = await lastJob(user.id);
    expect(job?.failure ?? null).toBeNull();
    expect(res.status).toBe(200);
    expect(job?.status).toBe("succeeded");
    expect(await measurementsOf(user.id)).toEqual(before);
  });

  it("a refused clearing step names its cause and changes nothing", async () => {
    const { user, backupId } = await seedAccountWithBackup();
    const changed = await measurementsOf(user.id);
    await onMeasurementDelete(
      "RAISE EXCEPTION 'restore test' USING ERRCODE = 'disk_full';",
    );

    const res = await restore(backupId);
    expect(res.status).toBe(500);
    const job = await lastJob(user.id);
    expect(job?.status).toBe("failed");
    expect(job?.failure).toEqual({
      code: "transaction_failed",
      cause: "storage",
      message: RESTORE_FAILURE_CAUSE_MESSAGES.storage,
    });
    expect(await measurementsOf(user.id)).toEqual(changed);

    const audit = await getPrismaClient().auditLog.findFirstOrThrow({
      where: { action: "admin.backups.restore.failed" },
      orderBy: { createdAt: "desc" },
    });
    expect(JSON.parse(audit.details ?? "{}")).toMatchObject({
      reason: "transaction_failed",
      cause: "storage",
      code: "53100",
    });
  });
});