fix(migrations): Preserve PGlite progress on failure

Serialize PGlite migration runs outside database transactions so data migration checkpoints and failure state survive an interrupted run.

Author: dcramer

packages/junior-testing/src/pglite.ts

@@ -21,11 +21,16 @@ export interface LocalPgliteFixture<TDatabase> {
   ): Promise<T[]>;
   transaction<T>(callback: () => Promise<T>): Promise<T>;
   withLock<T>(lockName: string, callback: () => Promise<T>): Promise<T>;
+  withMigrationLock<T>(
+    lockName: string,
+    callback: () => Promise<T>,
+  ): Promise<T>;
   close(): Promise<void>;
 }
 
 class LocalPgliteExecutor<TDatabase> implements LocalPgliteFixture<TDatabase> {
   private activeTransaction: Transaction | undefined;
+  private readonly migrationLocks = new Map<string, Promise<void>>();
 
   constructor(
     readonly client: PGlite,
@@ -85,6 +90,30 @@ class LocalPgliteExecutor<TDatabase> implements LocalPgliteFixture<TDatabase> {
     });
   }
 
+  async withMigrationLock<T>(
+    lockName: string,
+    callback: () => Promise<T>,
+  ): Promise<T> {
+    if (!lockName) {
+      throw new Error("Migration lock name is required");
+    }
+    const previous = this.migrationLocks.get(lockName) ?? Promise.resolve();
+    let release!: () => void;
+    const current = new Promise<void>((resolve) => {
+      release = resolve;
+    });
+    this.migrationLocks.set(lockName, current);
+    await previous;
+    try {
+      return await callback();
+    } finally {
+      release();
+      if (this.migrationLocks.get(lockName) === current) {
+        this.migrationLocks.delete(lockName);
+      }
+    }
+  }
+
   close(): Promise<void> {
     return this.client.close();
   }

packages/junior/tests/fixtures/sql.ts

@@ -47,7 +47,7 @@ export async function createLocalJuniorSqlFixture(): Promise<LocalJuniorSqlFixtu
     transaction: (callback) => fixture.transaction(callback),
     withLock: (lockName, callback) => fixture.withLock(lockName, callback),
     withMigrationLock: (migrationTable, callback) =>
-      fixture.withLock(`junior:migrate:${migrationTable}`, callback),
+      fixture.withMigrationLock(`junior:migrate:${migrationTable}`, callback),
   };
 
   return {

packages/junior/tests/unit/sql/pglite.test.ts

@@ -0,0 +1,37 @@
+import { afterEach, describe, expect, it } from "vitest";
+import {
+  createLocalPgliteFixture,
+  type LocalPgliteFixture,
+} from "@sentry/junior-testing/pglite";
+
+let fixture: LocalPgliteFixture<unknown> | undefined;
+
+afterEach(async () => {
+  await fixture?.close();
+  fixture = undefined;
+});
+
+describe("PGlite migration lock", () => {
+  it("preserves writes when a migration callback fails", async () => {
+    fixture = await createLocalPgliteFixture<unknown>({});
+    await fixture.execute(`
+CREATE TABLE migration_progress (
+  cursor INTEGER NOT NULL
+)
+`);
+
+    await expect(
+      fixture.withMigrationLock("test-journal", async () => {
+        await fixture!.execute(
+          "INSERT INTO migration_progress (cursor) VALUES ($1)",
+          [1],
+        );
+        throw new Error("migration interrupted");
+      }),
+    ).rejects.toThrow("migration interrupted");
+
+    await expect(
+      fixture.query("SELECT cursor FROM migration_progress"),
+    ).resolves.toEqual([{ cursor: 1 }]);
+  });
+});