fix(migrations): Address migration recovery gaps

dcramer · dcramer · commit 5a93cff28dd8 · 2026-07-23T09:06:17.000-07:00
Allow edited TypeScript migrations to recover from stale running rows once the migration lock is reacquired. Pass the host Redis capability into plugin data migrations and cover both paths with tests.
diff --git a/packages/junior-migrations/src/runner.ts b/packages/junior-migrations/src/runner.ts
@@ -234,7 +234,12 @@ async function runTypeScriptMigration(args: {
   const qualified = qualifiedTable(args.table);
   const sourceChanged =
     args.row !== undefined && args.row.hash !== args.migration.hash;
-  if (args.row && sourceChanged && args.row.status !== "failed") {
+  if (
+    args.row &&
+    sourceChanged &&
+    args.row.status !== "failed" &&
+    args.row.status !== "running"
+  ) {
     throw new Error(`Migration ${args.migration.tag} changed after it started`);
   }
   if (args.row) {
@@ -354,7 +359,10 @@ export async function runMigrationJournal(
         if (
           row &&
           row.hash !== migration.hash &&
-          !(migration.kind === "typescript" && row.status === "failed")
+          !(
+            migration.kind === "typescript" &&
+            (row.status === "failed" || row.status === "running")
+          )
         ) {
           throw new Error(
             `Migration ${migration.tag} changed after it started`,
diff --git a/packages/junior-migrations/tests/runner.test.ts b/packages/junior-migrations/tests/runner.test.ts
@@ -304,60 +304,68 @@ describe("runMigrationJournal", () => {
     ]);
   });
 
-  it("restarts a failed TypeScript migration after its source changes", async () => {
-    const folder = await mixedFolder();
-    const executor = new FakeExecutor();
-    let attempts = 0;
-    let resumedProgress: unknown;
-    const migration: MigrationV1 = {
-      apiVersion: 1,
-      async up(context) {
-        attempts += 1;
-        if (attempts === 1) {
-          await context.progress.save({ cursor: 1 });
-          throw new Error("needs a source fix");
-        }
-        resumedProgress = await context.progress.load();
-      },
-    };
-    const options = {
-      executor,
-      migrationsFolder: folder,
-      migrationsTable: "__drizzle_test",
-      loadTypeScript: async () => ({ default: migration }),
-      createContext: ({
-        progress,
-      }: {
-        progress: MigrationContextV1["progress"];
-      }) => ({
-        database: executor,
-        log: () => {},
-        progress,
-        state: fakeMigrationState(),
-      }),
-    };
+  it.each(["failed", "running"] as const)(
+    "restarts a %s TypeScript migration after its source changes",
+    async (retryStatus) => {
+      const folder = await mixedFolder();
+      const executor = new FakeExecutor();
+      let attempts = 0;
+      let resumedProgress: unknown;
+      const migration: MigrationV1 = {
+        apiVersion: 1,
+        async up(context) {
+          attempts += 1;
+          if (attempts === 1) {
+            await context.progress.save({ cursor: 1 });
+            throw new Error("needs a source fix");
+          }
+          resumedProgress = await context.progress.load();
+        },
+      };
+      const options = {
+        executor,
+        migrationsFolder: folder,
+        migrationsTable: "__drizzle_test",
+        loadTypeScript: async () => ({ default: migration }),
+        createContext: ({
+          progress,
+        }: {
+          progress: MigrationContextV1["progress"];
+        }) => ({
+          database: executor,
+          log: () => {},
+          progress,
+          state: fakeMigrationState(),
+        }),
+      };
 
-    await expect(runMigrationJournal(options)).rejects.toThrow(
-      "needs a source fix",
-    );
-    const failedHash = executor.rows.get(2_001)?.hash;
-    await writeFile(
-      join(folder, "0001_data.ts"),
-      "export default { apiVersion: 1, async up() { return { fixed: true }; } };\n",
-    );
+      await expect(runMigrationJournal(options)).rejects.toThrow(
+        "needs a source fix",
+      );
+      const interruptedRow = executor.rows.get(2_001);
+      expect(interruptedRow?.status).toBe("failed");
+      if (interruptedRow) {
+        interruptedRow.status = retryStatus;
+      }
+      const failedHash = interruptedRow?.hash;
+      await writeFile(
+        join(folder, "0001_data.ts"),
+        "export default { apiVersion: 1, async up() { return { fixed: true }; } };\n",
+      );
 
-    await expect(runMigrationJournal(options)).resolves.toEqual({
-      existing: 1,
-      migrated: 2,
-      scanned: 3,
-      skipped: 0,
-    });
-    expect(resumedProgress).toBeUndefined();
-    expect(executor.rows.get(2_001)).toMatchObject({
-      status: "completed",
-    });
-    expect(executor.rows.get(2_001)?.hash).not.toBe(failedHash);
-  });
+      await expect(runMigrationJournal(options)).resolves.toEqual({
+        existing: 1,
+        migrated: 2,
+        scanned: 3,
+        skipped: 0,
+      });
+      expect(resumedProgress).toBeUndefined();
+      expect(executor.rows.get(2_001)).toMatchObject({
+        status: "completed",
+      });
+      expect(executor.rows.get(2_001)?.hash).not.toBe(failedHash);
+    },
+  );
 
   it("rejects source changes after a TypeScript migration completes", async () => {
     const folder = await mixedFolder();
diff --git a/packages/junior/src/chat/plugins/migrations.ts b/packages/junior/src/chat/plugins/migrations.ts
@@ -7,6 +7,7 @@ import {
   type TypeScriptMigrationLoader,
 } from "@sentry/junior-migrations";
 import type { StateAdapter } from "chat";
+import type { RedisStateAdapter } from "@chat-adapter/state-redis";
 import { createPluginMigrationStateV1 } from "@/chat/plugins/migration-state";
 import type { JuniorSqlMigrationExecutor } from "@/db/db";
 
@@ -26,7 +27,10 @@ type PluginMigrationResult = {
 type PluginMigrationOptions =
   | { mode: "schema-bootstrap" }
   | {
-      getStateAdapter: () => Promise<StateAdapter>;
+      getStateContext: () => Promise<{
+        redisStateAdapter?: RedisStateAdapter;
+        stateAdapter: StateAdapter;
+      }>;
       loadTypeScript: TypeScriptMigrationLoader;
       log?: MigrationContextV1["log"];
       mode: "all";
@@ -165,15 +169,29 @@ export async function migratePluginSchemas(
     const pluginResult = runAll
       ? await runMigrationJournal({
           ...baseOptions,
-          createContext: async ({ progress }): Promise<MigrationContextV1> => ({
-            database: executor,
-            log: options.log ?? (() => {}),
-            progress,
-            state: createPluginMigrationStateV1(
-              root.pluginName,
-              await options.getStateAdapter(),
-            ),
-          }),
+          createContext: async ({ progress }): Promise<MigrationContextV1> => {
+            const { redisStateAdapter, stateAdapter } =
+              await options.getStateContext();
+            return {
+              database: executor,
+              log: options.log ?? (() => {}),
+              progress,
+              ...(redisStateAdapter
+                ? {
+                    redis: {
+                      sendCommand: async <T>(args: readonly string[]) =>
+                        await redisStateAdapter
+                          .getClient()
+                          .sendCommand<T>([...args]),
+                    },
+                  }
+                : {}),
+              state: createPluginMigrationStateV1(
+                root.pluginName,
+                stateAdapter,
+              ),
+            };
+          },
           loadTypeScript: options.loadTypeScript,
           mode: "all",
         })
diff --git a/packages/junior/src/cli/upgrade/migrations/plugin-journal.ts b/packages/junior/src/cli/upgrade/migrations/plugin-journal.ts
@@ -24,8 +24,7 @@ export async function migratePluginJournals(
       executor,
       pluginCatalogRuntime.getMigrationRoots(),
       {
-        getStateAdapter: async () =>
-          (await context.getStateContext()).stateAdapter,
+        getStateContext: context.getStateContext,
         loadTypeScript: async (path) =>
           await migrationLoader.import<Record<string, unknown>>(path),
         log: context.io.info,
diff --git a/packages/junior/tests/component/migrations/mixed-runner-database.test.ts b/packages/junior/tests/component/migrations/mixed-runner-database.test.ts
@@ -6,6 +6,9 @@ import {
   type MigrationContextV1,
   type MigrationStateV1,
 } from "@sentry/junior-migrations";
+import type { RedisStateAdapter } from "@chat-adapter/state-redis";
+import type { StateAdapter } from "chat";
+import { migratePluginSchemas } from "@/chat/plugins/migrations";
 import { createJiti } from "jiti";
 import { afterEach, describe, expect, it, vi } from "vitest";
 import {
@@ -92,6 +95,54 @@ afterEach(async () => {
 });
 
 describe("mixed migration runner database contract", () => {
+  it("provides the host Redis capability to plugin data migrations", async () => {
+    const fixture = await createLocalJuniorSqlFixture();
+    const folder = await createMigrationFolder({
+      source: `
+export default {
+  apiVersion: 1,
+  async up(context) {
+    return {
+      response: await context.redis.sendCommand(["PING"]),
+    };
+  },
+};
+`,
+      tag: "0000_plugin_redis",
+      when: 2_026_071_600_000,
+    });
+    const sendCommand = vi.fn(async () => "PONG");
+
+    try {
+      await expect(
+        migratePluginSchemas(
+          fixture.sql,
+          [{ dir: folder, pluginName: "redis-test" }],
+          {
+            getStateContext: async () => ({
+              redisStateAdapter: {
+                getClient: () => ({ sendCommand }),
+              } as unknown as RedisStateAdapter,
+              stateAdapter: {} as StateAdapter,
+            }),
+            loadTypeScript: async (migrationPath) =>
+              await migrationLoader.import<Record<string, unknown>>(
+                migrationPath,
+              ),
+            mode: "all",
+          },
+        ),
+      ).resolves.toEqual({
+        existing: 0,
+        migrated: 1,
+        scanned: 1,
+      });
+      expect(sendCommand).toHaveBeenCalledWith(["PING"]);
+    } finally {
+      await fixture.close();
+    }
+  });
+
   it("persists failed progress and resumes a TypeScript migration", async () => {
     const fixture = await createLocalJuniorSqlFixture();
     const folder = await createMigrationFolder({
