From e0b4f3c427488fb132fa165137db109a6a941069 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marc-Andr=C3=A9=20Bombeck?= <mbombeck@gmail.com>
Date: Sun, 27 Sep 2026 11:57:37 +0200
Subject: [PATCH] fix(restore): size the restore budget by the rows written,
 and keep test housekeeping off the short limit

---
 src/lib/export/restore-backup.ts              | 14 ++++---
 .../backup-restore-statement-timeout.test.ts  | 39 ++++++++++++++++---
 2 files changed, 41 insertions(+), 12 deletions(-)

diff --git a/src/lib/export/restore-backup.ts b/src/lib/export/restore-backup.ts
index 336fb7c14..e2b317ce1 100644
--- a/src/lib/export/restore-backup.ts
+++ b/src/lib/export/restore-backup.ts
@@ -177,17 +177,19 @@ export interface RestoreResponse {
  * The work is deleting the readings the account holds now and writing the
  * ones the file carries. Both count: an account whose readings were folded
  * into hourly averages still holds the originals as deleted rows until the
- * nightly purge removes them, and the restore deletes those too (#1031). The
- * budget is two minutes plus a second per 2 000 of those rows together. A
- * seeded account of 1.89 million readings restored in 196 s on a fast host
- * and in 9 min on a database held to half a CPU and 40 MB/s of writes; the
- * budget gives it 33 min.
+ * nightly purge removes them, and the restore deletes and rewrites those too
+ * (#1031). Writing is the slow half, since every reading goes into thirteen
+ * indexes. The budget is two minutes, plus a second per 1 000 readings in the
+ * file and per 10 000 in the account. A seeded account of 1.89 million
+ * readings restored in 196 s on a fast host and in 19 min on a database held
+ * to half a CPU and 40 MB/s of writes, where its transaction ran 16 min of a
+ * 37 min budget.
  */
 function restoreTransactionTimeoutMs(
   fileMeasurements: number,
   currentMeasurements: number,
 ): number {
-  return 120_000 + Math.ceil((fileMeasurements + currentMeasurements) / 2);
+  return 120_000 + fileMeasurements + Math.ceil(currentMeasurements / 10);
 }
 
 function decodeEncryptedBytes(encoded: string): Uint8Array<ArrayBuffer> {
diff --git a/tests/integration/backup-restore-statement-timeout.test.ts b/tests/integration/backup-restore-statement-timeout.test.ts
index 2da896255..e58adf1f7 100644
--- a/tests/integration/backup-restore-statement-timeout.test.ts
+++ b/tests/integration/backup-restore-statement-timeout.test.ts
@@ -10,8 +10,7 @@
  *
  * 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
+ * situation at a size a test can hold. 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.
@@ -23,13 +22,24 @@
  * top of the restore transaction and the restore test goes red with
  * `cause: "timeout"`.
  */
-import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import {
+  afterAll,
+  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 { PrismaPg } from "@prisma/adapter-pg";
+
+import { PrismaClient } from "@/generated/prisma/client";
 import { cookieJar, headerJar } from "./mock-next-headers";
 import { getPrismaClient, truncateAllTables } from "./setup";
 import { streamFullBackupJson } from "@/lib/export/full-backup-stream";
@@ -63,8 +73,21 @@ vi.mock("@/lib/db-compat", () => ({
 
 const COUNT = 600;
 
+/**
+ * A second pool without the lowered limit, for the suite's own housekeeping:
+ * truncating every table between tests can take longer than five seconds on
+ * a busy host, and it is not what this file measures.
+ */
+let setupClient: PrismaClient | null = null;
+function housekeeping(): PrismaClient {
+  setupClient ??= new PrismaClient({
+    adapter: new PrismaPg({ connectionString: process.env.DATABASE_URL! }),
+  });
+  return setupClient;
+}
+
 async function dropTrigger() {
-  const prisma = getPrismaClient();
+  const prisma = housekeeping();
   await prisma.$executeRawUnsafe(
     `DROP TRIGGER IF EXISTS restore_test_clearing ON measurements`,
   );
@@ -75,7 +98,7 @@ async function dropTrigger() {
 
 /** Run `body` on every DELETE statement against `measurements`. */
 async function onMeasurementDelete(body: string) {
-  const prisma = getPrismaClient();
+  const prisma = housekeeping();
   await prisma.$executeRawUnsafe(
     `CREATE FUNCTION restore_test_clearing() RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN ${body} RETURN NULL; END $$`,
   );
@@ -86,7 +109,7 @@ async function onMeasurementDelete(body: string) {
 
 beforeEach(async () => {
   await dropTrigger();
-  await truncateAllTables(getPrismaClient());
+  await truncateAllTables(housekeeping());
   cookieJar.clear();
   headerJar.clear();
 });
@@ -95,6 +118,10 @@ afterEach(async () => {
   await dropTrigger();
 });
 
+afterAll(async () => {
+  await setupClient?.$disconnect();
+});
+
 async function seedAccountWithBackup() {
   const prisma = getPrismaClient();
   const user = await prisma.user.create({