Merge pull request #23 from get-convex/ian/run-sync

Ian/run sync

Author: ianmacartney

README.md

@@ -442,4 +442,57 @@ await migrations.getStatus(ctx, { migrations: ["myNewMutation"] });
 await migrations.cancel(ctx, "myNewMutation");
 ```
 
+## Running migrations synchronously
+
+If you want to run a migration synchronously from a test or action, you can use
+`runToCompletion`. Note that if the action crashes or is canceled, it will not
+continue migrating in the background.
+
+From an action:
+
+```ts
+import { components, internal } from "./_generated/api";
+import { internalAction } from "./_generated/server";
+import { runToCompletion } from "@convex-dev/migrations";
+
+export const myAction = internalAction({
+  args: {},
+  handler: async (ctx) => {
+    //...
+    const toRun = internal.example.setDefaultValue;
+    await runToCompletion(ctx, components.migrations, toRun);
+  },
+});
+```
+
+In a test:
+
+```ts
+import { test } from "vitest";
+import { convexTest } from "convex-test";
+import component from "@convex-dev/migrations/test";
+import { runToCompletion } from "@convex-dev/migrations";
+import { components, internal } from "./_generated/api";
+import schema from "./schema";
+
+test("test setDefaultValue migration", async () => {
+  const t = convexTest(schema);
+  // Register the component in the test instance
+  component.register(t);
+
+  await t.run(async (ctx) => {
+    // Add sample data to migrate
+    await ctx.db.insert("myTable", { optionalField: undefined });
+
+    // Run the migration to completion
+    const migrationToTest = internal.example.setDefaultValue;
+    await runToCompletion(ctx, components.migrations, migrationToTest);
+
+    // Assert that the migration was successful by checking the data
+    const docs = await ctx.db.query("myTable").collect();
+    expect(docs.every((doc) => doc.optionalField !== undefined)).toBe(true);
+  });
+});
+```
+
 <!-- END: Include on https://convex.dev/components -->

example/convex/example.test.ts

@@ -1,7 +1,8 @@
 import { afterEach, beforeEach, describe, expect, test, vi } from "vitest";
 import { initConvexTest } from "./setup.test";
-import { internal } from "./_generated/api";
-import { migrations } from "./example";
+import { components, internal } from "./_generated/api";
+import { runToCompletion } from "@convex-dev/migrations";
+import { createFunctionHandle, getFunctionName } from "convex/server";
 
 describe("example", () => {
   beforeEach(async () => {
@@ -12,7 +13,7 @@ describe("example", () => {
     vi.useRealTimers();
   });
 
-  test("setDefaultValue", async () => {
+  test("test setDefaultValue migration", async () => {
     const t = initConvexTest();
     await t.mutation(internal.example.seed, { count: 10 });
     await t.run(async (ctx) => {
@@ -21,12 +22,51 @@ describe("example", () => {
       expect(docs.some((doc) => doc.optionalField === undefined)).toBe(true);
     });
     await t.run(async (ctx) => {
-      await migrations.runOne(ctx, internal.example.setDefaultValue, {
+      await runToCompletion(
+        ctx,
+        components.migrations,
+        internal.example.setDefaultValue,
+        { batchSize: 2 },
+      );
+    });
+    await t.run(async (ctx) => {
+      const after = await ctx.db.query("myTable").collect();
+      expect(after).toHaveLength(10);
+      expect(after.every((doc) => doc.optionalField !== undefined)).toBe(true);
+    });
+  });
+
+  test("test failingMigration", async () => {
+    const t = initConvexTest();
+    await t.mutation(internal.example.seed, { count: 10 });
+    await expect(
+      t.run(async (ctx) => {
+        await runToCompletion(
+          ctx,
+          components.migrations,
+          internal.example.failingMigration,
+        );
+      }),
+    ).rejects.toThrow("This migration fails after the first");
+  });
+
+  test("test migrating with function handle", async () => {
+    const t = initConvexTest();
+    await t.mutation(internal.example.seed, { count: 10 });
+    await t.run(async (ctx) => {
+      const docs = await ctx.db.query("myTable").collect();
+      expect(docs).toHaveLength(10);
+      expect(docs.some((doc) => doc.optionalField === undefined)).toBe(true);
+    });
+    await t.run(async (ctx) => {
+      const fnHandle = await createFunctionHandle(
+        internal.example.setDefaultValue,
+      );
+      await runToCompletion(ctx, components.migrations, fnHandle, {
+        name: getFunctionName(internal.example.setDefaultValue),
         batchSize: 2,
-        dryRun: false,
       });
     });
-    await t.finishAllScheduledFunctions(vi.runAllTimers);
     await t.run(async (ctx) => {
       const after = await ctx.db.query("myTable").collect();
       expect(after).toHaveLength(10);

package-lock.json

@@ -71,7 +71,7 @@
     "@vitejs/plugin-react": "5.0.4",
     "chokidar-cli": "3.0.0",
     "convex": "1.29.0",
-    "convex-test": "0.0.38",
+    "convex-test": "^0.0.40-alpha.0",
     "cpy-cli": "6.0.0",
     "eslint": "9.39.1",
     "eslint-plugin-react": "7.37.5",

package.json

@@ -6,6 +6,7 @@ import {
   type GenericDataModel,
   type GenericMutationCtx,
   type GenericQueryCtx,
+  getFunctionAddress,
   getFunctionName,
   internalMutationGeneric,
   makeFunctionReference,
@@ -27,6 +28,7 @@ export type { MigrationArgs, MigrationResult, MigrationStatus };
 import { ConvexError, type GenericId } from "convex/values";
 import type { ComponentApi } from "../component/_generated/component.js";
 import { logStatusAndInstructions } from "./log.js";
+import type { MigrationFunctionHandle } from "../component/lib.js";
 
 // Note: this value is hard-coded in the docstring below. Please keep in sync.
 export const DEFAULT_BATCH_SIZE = 100;
@@ -609,6 +611,92 @@ export type MigrationFunctionReference = FunctionReference<
   MigrationArgs
 >;
 
+/**
+ * Start a migration and run it synchronously until it's done.
+ * If this function crashes, the migration will not continue.
+ *
+ * ```ts
+ * // In an action
+ * const toRun = internal.migrations.myMigration;
+ * await runToCompletion(ctx, components.migrations, toRun);
+ * ```
+ *
+ * If it's already in progress, it will no-op.
+ * If you run a migration that had previously failed which was part of a series,
+ * it will not resume the series.
+ * To resume a series, call the series again: {@link Migrations.runSerially}.
+ *
+ * Note: It's up to you to determine if it's safe to run a migration while
+ * others are in progress. It won't run multiple instances of the same migration
+ * but it currently allows running multiple migrations on the same table.
+ *
+ * @param ctx Context from an action.
+ * @param component The migrations component, usually `components.migrations`.
+ * @param fnRef The migration function to run. Like `internal.migrations.foo`.
+ * @param opts Options to start the migration.
+ *   It's helpful to see what it would do without committing the transaction.
+ */
+export async function runToCompletion(
+  ctx: ActionCtx,
+  component: ComponentApi,
+  fnRef: MigrationFunctionReference | MigrationFunctionHandle,
+  opts?: {
+    /**
+     * The name of the migration function, generated with getFunctionName.
+     */
+    name?: string;
+    /**
+     * The cursor to start from.
+     * null: start from the beginning.
+     * undefined: start, or resume from where it failed. No-ops if already done.
+     */
+    cursor?: string | null;
+    /**
+     * The number of documents to process in a batch.
+     * Overrides the migrations's configured batch size.
+     */
+    batchSize?: number;
+    /**
+     * If true, it will run a batch and then throw an error.
+     * It's helpful to see what it would do without committing the transaction.
+     */
+    dryRun?: boolean;
+  },
+): Promise<MigrationStatus> {
+  let cursor = opts?.cursor;
+  const {
+    name = getFunctionName(fnRef),
+    batchSize,
+    dryRun = false,
+  } = opts ?? {};
+  const address = getFunctionAddress(fnRef);
+  const fnHandle =
+    address.functionHandle ?? (await createFunctionHandle(fnRef));
+  while (true) {
+    const status = await ctx.runMutation(component.lib.migrate, {
+      name,
+      fnHandle,
+      cursor,
+      batchSize,
+      dryRun,
+      oneBatchOnly: true,
+    });
+    if (status.isDone) {
+      return status;
+    }
+    if (status.error) {
+      throw new Error(status.error);
+    }
+    if (!status.cursor || status.cursor === cursor) {
+      throw new Error(
+        "Invariant violation: Migration did not make progress." +
+          `\nStatus: ${JSON.stringify(status)}`,
+      );
+    }
+    cursor = status.cursor;
+  }
+}
+
 /* Type utils follow */
 
 type QueryCtx = Pick<GenericQueryCtx<GenericDataModel>, "runQuery">;

src/client/index.ts

@@ -95,6 +95,7 @@ export type ComponentApi<Name extends string | undefined = string | undefined> =
           fnHandle: string;
           name: string;
           next?: Array<{ fnHandle: string; name: string }>;
+          oneBatchOnly?: boolean;
         },
         {
           batchSize?: number;

src/component/_generated/component.ts

@@ -27,6 +27,7 @@ const runMigrationArgs = {
   cursor: v.optional(v.union(v.string(), v.null())),
 
   batchSize: v.optional(v.number()),
+  oneBatchOnly: v.optional(v.boolean()),
   next: v.optional(
     v.array(
       v.object({
@@ -43,7 +44,7 @@ export const migrate = mutation({
   returns: migrationStatus,
   handler: async (ctx, args) => {
     // Step 1: Get or create the state.
-    const { fnHandle, batchSize, next: next_, dryRun, ...initialState } = args;
+    const { fnHandle, batchSize, next: next_, dryRun, name } = args;
     if (batchSize !== undefined && batchSize <= 0) {
       throw new Error("Batch size must be greater than 0");
     }
@@ -58,11 +59,11 @@ export const migrate = mutation({
     const state =
       (await ctx.db
         .query("migrations")
-        .withIndex("name", (q) => q.eq("name", args.name))
+        .withIndex("name", (q) => q.eq("name", name))
         .unique()) ??
       (await ctx.db.get(
         await ctx.db.insert("migrations", {
-          ...initialState,
+          name,
           cursor: args.cursor ?? null,
           isDone: false,
           processed: 0,
@@ -122,7 +123,9 @@ export const migrate = mutation({
       }
 
       // Step 3: Schedule the next batch or next migration.
-      if (!state.isDone) {
+      if (args.oneBatchOnly) {
+        state.workerId = undefined;
+      } else if (!state.isDone) {
         // Recursively schedule the next batch.
         state.workerId = await ctx.scheduler.runAfter(0, api.lib.migrate, {
           ...args,
@@ -159,7 +162,7 @@ export const migrate = mutation({
           }
         } else {
           console.info(
-            `Migration ${args.name} is done.` +
+            `Migration ${name} is done.` +
               (i < next.length ? ` Next: ${next[i]!.name}` : ""),
           );
         }
@@ -171,7 +174,7 @@ export const migrate = mutation({
         updateState(e.data.result);
       } else {
         state.error = e instanceof Error ? e.message : String(e);
-        console.error(`Migration ${args.name} failed: ${state.error}`);
+        console.error(`Migration ${name} failed: ${state.error}`);
       }
       if (dryRun) {
         const status = await getMigrationState(ctx, state);
@@ -296,7 +299,9 @@ async function cancelMigration(ctx: MutationCtx, migration: Doc<"migrations">) {
     return state;
   }
   if (state.state === "inProgress") {
-    await ctx.scheduler.cancel(migration.workerId!);
+    if (!migration.workerId) {
+      await ctx.scheduler.cancel(migration.workerId!);
+    }
     console.log(`Canceled migration ${migration.name}`);
     return { ...state, state: "canceled" as const };
   }