Add transactional success completion for queries

README.md

@@ -131,10 +131,11 @@ export const emailSent = internalMutation({
 ```
 
 By default, the `onComplete` handler runs in a different transaction than the
-enqueued job. For mutations, `completeTransactionally: true` makes successful
-work and its selected callback commit together, as described below. You can also
-do that work at the end of the enqueued function, before returning, which avoids
-another function call and preserves the result's type directly.
+enqueued job. For mutations and queries, `completeTransactionally: true` makes
+successful work and its selected callback complete in one transaction, as
+described below. For mutations, you can also do callback work at the end of the
+enqueued function, before returning, which avoids another function call and
+preserves the result's type directly.
 
 You can also use this equivalent helper to define an `onComplete` mutation. Note
 the `DataModel` type parameter, if you want ctx.db to be type safe.
@@ -170,11 +171,12 @@ with `cancel` can still report `"success"` or a terminal `"failed"`. Direct
 scheduler cancelation (for example, from the dashboard) is treated as a failure
 and can trigger retries instead.
 
-### Transactional success for mutations
+### Transactional success for mutations and queries
 
 Use `completeTransactionally: true` with `enqueueMutation` or
 `enqueueMutationBatch` to commit each mutation, its selected success callback,
-and workpool completion bookkeeping together:
+and workpool completion bookkeeping together. `enqueueQuery` and
+`enqueueQueryBatch` support the same option:
 
 ```ts
 await pool.enqueueMutation(ctx, internal.tasks.process, args, {
@@ -184,20 +186,39 @@ await pool.enqueueMutation(ctx, internal.tasks.process, args, {
 });
 ```
 
-The callback can read the mutation's writes. If the success callback or
-completion bookkeeping throws, the scheduled mutation fails and all those writes
-roll back. Workpool's periodic recovery then marks the work as failed and
-invokes the callback for `"failed"`, unless excluded. Until recovery runs, the
-job remains running and occupies a concurrency slot. There is no separate
-fallback execution of the success callback.
+Queries always read an independent snapshot. With this option, the success
+callback and completion bookkeeping commit together, but the query's reads do
+not create dependencies that cause that transaction to retry on concurrent
+writes. The query result may be stale when the callback commits; read any data
+that must still be current within the callback. Without this option, queries run
+through the action batch worker.
+
+Each opted-in query is scheduled in its own mutation wrapper, including queries
+enqueued in a batch. For example:
+
+```ts
+await pool.enqueueQuery(ctx, internal.tasks.findPending, args, {
+  onComplete: internal.tasks.handlePending,
+  onCompleteExcludeKinds: ["canceled"],
+  completeTransactionally: true,
+});
+```
+
+For mutations, the callback can read the mutation's writes. For either function
+type, if the success callback or completion bookkeeping throws, the scheduled
+wrapper fails and all its writes roll back. Workpool's periodic recovery then
+marks the work as failed and invokes the callback for `"failed"`, unless
+excluded. Until recovery runs, the job remains running and occupies a
+concurrency slot. There is no separate fallback execution of the success
+callback.
 
 Ordinary work errors still use the existing failure-completion path. Failure and
 cancellation callbacks keep their usual transaction behavior. Excluding every
-kind, or omitting the callback, still allows transactional bookkeeping, avoiding
-the separately scheduled completion mutation. Jobs in a batch remain independent
-transactions. Actions and queries do not support this option.
+kind, or omitting the callback, still allows transactional bookkeeping, without
+a separate completion invocation. Jobs in a batch remain independent
+transactions. Actions do not support this option.
 
-This option requires Convex 1.41 or newer. The work and callback share one
+This option requires Convex 1.42 or newer. The work and callback share one
 transaction budget. Workpool uses nested transaction limits to reserve space for
 completion bookkeeping; work that fits by itself may exceed the combined budget.
 Limit overruns or execution failures can still fail the transaction, in which
@@ -341,9 +362,9 @@ See example usage in [example.ts](./example/convex/example.ts).
 
 Check out the [docstrings](./src/client/index.ts), but notable options include:
 
-- `maxParallelism`: How many actions/mutations can run at once within this pool.
-  Avoid exceeding 100 on Pro, 20 on the free plan, across all workpools and
-  workflows.
+- `maxParallelism`: How many actions, mutations, and queries can run at once
+  within this pool. Avoid exceeding 100 on Pro, 20 on the free plan, across all
+  workpools and workflows.
 - `retryActionsByDefault`: Whether to retry actions that fail by default.
 - `defaultRetryBehavior`: The default retry behavior for enqueued actions.
 
@@ -370,8 +391,9 @@ options include:
 - `onCompleteExcludeKinds`: Optional list of result kinds to skip: `"success"`,
   `"failed"`, and `"canceled"`. Defaults to excluding none, so every outcome
   invokes the callback.
-- `completeTransactionally`: Mutation-only option to commit successful work, its
-  selected callback, and completion bookkeeping together. Defaults to false.
+- `completeTransactionally`: Mutation and query option to complete successful
+  work, its selected callback, and completion bookkeeping together. Defaults to
+  false.
 - `context`: Any data you want to pass to the `onComplete` mutation.
 - `runAt` and `runAfter`: Similar to `ctx.scheduler.run*`, allows you to
   schedule the work to run later. By default it's immediate.

src/client/index.test.ts

@@ -273,6 +273,84 @@ describe("completion callbacks through the client and scheduler", () => {
     },
   );
 
+  test.each([false, true])(
+    "transactional success works through query enqueue (batch: %s)",
+    async (batch) => {
+      const args = [
+        { key: "atomic-success" },
+        { key: "atomic-failure", fail: true },
+      ];
+      const options = {
+        completeTransactionally: true,
+        onComplete: refs.complete,
+        context: { key: "atomic-callbacks" },
+      };
+      const ids = await t.mutation(async (ctx) =>
+        batch
+          ? pool.enqueueQueryBatch(ctx, refs.query, args, options)
+          : Promise.all(
+              args.map((args) =>
+                pool.enqueueQuery(ctx, refs.query, args, options),
+              ),
+            ),
+      );
+      await drain();
+      expect((await events("atomic-success")).map((e) => e.kind)).toEqual([]);
+      expect(await events("atomic-failure")).toEqual([]);
+      const results = await events("atomic-callbacks");
+      expect(results).toHaveLength(2);
+      expect(results.find((e) => e.workId === ids[0])?.result).toEqual({
+        kind: "success",
+        returnValue: "query result",
+      });
+      expect(results.find((e) => e.workId === ids[1])?.result).toEqual({
+        kind: "failed",
+        error: "work failed",
+      });
+      expect(await t.query((ctx) => pool.statusBatch(ctx, ids))).toEqual([
+        { state: "finished" },
+        { state: "finished" },
+      ]);
+    },
+  );
+
+  test.each([false, true])(
+    "transactional query enqueue rolls back completion when the callback fails (batch: %s)",
+    async (batch) => {
+      const args = [{ key: "rollback-first" }, { key: "rollback-second" }];
+      const options = {
+        completeTransactionally: true,
+        onComplete: refs.complete,
+        context: { key: "rollback-callback", throw: true },
+      };
+      try {
+        const ids = await t.mutation(async (ctx) =>
+          batch
+            ? pool.enqueueQueryBatch(ctx, refs.query, args, options)
+            : Promise.all(
+                args.map((args) =>
+                  pool.enqueueQuery(ctx, refs.query, args, options),
+                ),
+              ),
+        );
+        await vi.waitFor(() => expect(callback).toHaveBeenCalledTimes(2), {
+          timeout: 10_000,
+        });
+        await t.finishInProgressScheduledFunctions();
+        for (const { key } of args) expect(await events(key)).toEqual([]);
+        expect(await events("rollback-callback")).toEqual([]);
+        expect(
+          await t.query((ctx) => pool.statusBatch(ctx, ids)),
+        ).toMatchObject([{ state: "running" }, { state: "running" }]);
+      } finally {
+        // Inspect rollback before periodic recovery. worker.test.ts covers
+        // recovery with convex-test's missing scheduled-error text supplied.
+        await t.finishInProgressScheduledFunctions();
+        vi.clearAllTimers();
+      }
+    },
+  );
+
   test("canceling pending transactional work preserves the cancellation callback", async () => {
     const id = await t.mutation((ctx) =>
       pool.enqueueMutation(
@@ -966,7 +1044,7 @@ test("exclude filters only accept terminal result kinds", () => {
   }).toMatchTypeOf<EnqueueOptions>();
 });
 
-test("transactional completion is only exposed for mutations", () => {
+test("transactional completion is exposed for mutations and queries", () => {
   const options = {
     onComplete,
     onCompleteExcludeKinds: ["failed", "canceled"] as const,
@@ -976,18 +1054,11 @@ test("transactional completion is only exposed for mutations", () => {
   const mutation = makeFunctionReference<"mutation", { value: number }, number>(
     "work:mutation",
   );
-  const query = makeFunctionReference<"query", { value: number }, number>(
-    "work:query",
-  );
   expectTypeOf((pool: Workpool, ctx: MutationCtx) => {
     // @ts-expect-error Actions cannot opt into transactional completion.
     void pool.enqueueAction(ctx, action, { value: 1 }, options);
     // @ts-expect-error Batched actions cannot opt into transactional completion.
     void pool.enqueueActionBatch(ctx, action, [{ value: 1 }], options);
-    // @ts-expect-error Queries cannot opt into transactional completion.
-    void pool.enqueueQuery(ctx, query, { value: 1 }, options);
-    // @ts-expect-error Batched queries cannot opt into transactional completion.
-    void pool.enqueueQueryBatch(ctx, query, [{ value: 1 }], options);
     return pool.enqueueMutation(ctx, mutation, { value: 1 }, options);
   }).returns.toEqualTypeOf<Promise<WorkId>>();
   expectTypeOf((pool: Workpool, ctx: MutationCtx) =>
@@ -1012,16 +1083,19 @@ test("standalone enqueue gates transactional completion on fnType", () => {
     { value: number },
     number
   >("work:mixed");
+  const txFn = makeFunctionReference<
+    "mutation" | "query",
+    { value: number },
+    number
+  >("work:tx");
   const arg = { value: 1 };
   const batch = [arg];
   expectTypeOf((c: WorkpoolComponent, ctx: MutationCtx) => {
     // @ts-expect-error Actions cannot opt into transactional completion.
     void enqueue(c, ctx, "action", action, arg, options);
-    // @ts-expect-error Queries cannot opt into transactional completion.
-    void enqueue(c, ctx, "query", query, arg, options);
     // @ts-expect-error Batched actions cannot opt into transactional completion.
     void enqueueBatch(c, ctx, "action", action, batch, options);
-    // @ts-expect-error Batched queries cannot opt into transactional completion.
+    void enqueue(c, ctx, "query", query, arg, options);
     void enqueueBatch(c, ctx, "query", query, batch, options);
     // A union that includes a mutation must still be rejected as a whole: a
     // distributive conditional would collapse this to boolean and let it through.
@@ -1030,13 +1104,52 @@ test("standalone enqueue gates transactional completion on fnType", () => {
     void enqueue(c, ctx, mixed, mixedFn, arg, options);
     // @ts-expect-error A batched action-capable union cannot opt in.
     void enqueueBatch(c, ctx, mixed, mixedFn, batch, options);
+    // Every member of this union supports the option, so it is accepted.
+    const tx = "mutation" as "mutation" | "query";
+    void enqueue(c, ctx, tx, txFn, arg, options);
+    void enqueueBatch(c, ctx, tx, txFn, batch, options);
     return enqueue(c, ctx, "mutation", mutation, arg, options);
   }).returns.toEqualTypeOf<Promise<WorkId>>();
   expectTypeOf((c: WorkpoolComponent, ctx: MutationCtx) =>
     enqueueBatch(c, ctx, "mutation", mutation, batch, options),
   ).returns.toEqualTypeOf<Promise<WorkId[]>>();
 });
 
+test("transactional query options retain result and context types", () => {
+  const query = makeFunctionReference<"query", { value: number }, number>(
+    "work:query",
+  );
+  const options = {
+    onComplete,
+    onCompleteExcludeKinds: ["failed", "canceled"] as const,
+    context: { label: "query" },
+    completeTransactionally: true,
+  } satisfies TransactionalEnqueueOptions<{ label: string }, number>;
+  expectTypeOf((pool: Workpool, ctx: MutationCtx) => {
+    void pool.enqueueQuery(
+      ctx,
+      query,
+      { value: 1 },
+      {
+        ...options,
+        // @ts-expect-error The callback requires a string label.
+        context: { label: 123 },
+      },
+    );
+    const stringQuery = makeFunctionReference<
+      "query",
+      { value: number },
+      string
+    >("work:stringQuery");
+    // @ts-expect-error The callback expects a numeric query result.
+    void pool.enqueueQuery(ctx, stringQuery, { value: 1 }, options);
+    return pool.enqueueQuery(ctx, query, { value: 1 }, options);
+  }).returns.toEqualTypeOf<Promise<WorkId>>();
+  expectTypeOf((pool: Workpool, ctx: MutationCtx) =>
+    pool.enqueueQueryBatch(ctx, query, [{ value: 1 }], options),
+  ).returns.toEqualTypeOf<Promise<WorkId[]>>();
+});
+
 test("generated enqueue types match the component validators", () => {
   expectTypeOf<
     FunctionArgs<WorkpoolComponent["lib"]["enqueue"]>

src/client/index.ts

@@ -220,8 +220,9 @@ export class Workpool {
   /**
    * Enqueues a query to be run.
    * Usually not what you want, but it can be useful during workflows.
-   * The query is run in a mutation and the result is returned to the caller,
-   * so it can conflict if other mutations are writing the value.
+   * The query always runs in its own snapshot. With completeTransactionally,
+   * its success callback and completion bookkeeping commit together, without
+   * adding the query's reads to that transaction's conflict detection.
    *
    * @param ctx - The mutation or action context that can call ctx.runMutation.
    * @param fn - The query to run, like `internal.example.myQuery`.
@@ -238,9 +239,7 @@ export class Workpool {
       NoInfer<ReturnType>
     >,
     fnArgs: Args,
-    options?: EnqueueOptions<Context, ReturnType> & {
-      completeTransactionally?: never;
-    },
+    options?: TransactionalEnqueueOptions<Context, ReturnType>,
   ): Promise<WorkId> {
     return enqueue(this.component, ctx, "query", fn, fnArgs, {
       ...this.options,
@@ -272,9 +271,7 @@ export class Workpool {
       NoInfer<ReturnType>
     >,
     argsArray: Array<Args>,
-    options?: EnqueueOptions<Context, ReturnType> & {
-      completeTransactionally?: never;
-    },
+    options?: TransactionalEnqueueOptions<Context, ReturnType>,
   ): Promise<WorkId[]> {
     return enqueueBatch(this.component, ctx, "query", fn, argsArray, {
       ...this.options,
@@ -429,7 +426,7 @@ export type RetryOption = {
 };
 
 export type WorkpoolOptions = {
-  /** How many actions/mutations can be running at once within this pool.
+  /** How many actions, mutations, and queries can run at once within this pool.
    * Suggested max: 100 on Pro, 20 on free plan.
    * If set to 0, no new work will be started.
    */
@@ -542,7 +539,8 @@ export type TransactionalEnqueueOptions<
    * errors roll back the work; recovery later reports the job as failed.
    * The work and callback share transaction limits, with headroom reserved for
    * bookkeeping. Failure and cancellation callbacks keep their usual behavior.
-   * Only supported for mutations.
+   * Supported by mutation and query enqueue methods. Queries use an independent
+   * snapshot without adding read dependencies to the completion transaction.
    */
   completeTransactionally?: boolean;
 };
@@ -659,10 +657,12 @@ export async function enqueueBatch<
     retryBehavior?: RetryBehavior;
     maxParallelism?: number;
     logLevel?: LogLevel;
-    /** Mutations only; see {@link TransactionalEnqueueOptions}. */
+    /** Mutations and queries only; see {@link TransactionalEnqueueOptions}. */
     // Wrapped in a tuple so a union FnType is checked as a whole rather than
     // distributed, which would admit action-capable calls.
-    completeTransactionally?: [FnType] extends ["mutation"] ? boolean : never;
+    completeTransactionally?: [FnType] extends ["mutation" | "query"]
+      ? boolean
+      : never;
   },
 ): Promise<WorkId[]> {
   const { config, ...defaults } = await enqueueArgs(fn, options);
@@ -726,10 +726,12 @@ export async function enqueue<
     retryBehavior?: RetryBehavior;
     maxParallelism?: number;
     logLevel?: LogLevel;
-    /** Mutations only; see {@link TransactionalEnqueueOptions}. */
+    /** Mutations and queries only; see {@link TransactionalEnqueueOptions}. */
     // Wrapped in a tuple so a union FnType is checked as a whole rather than
     // distributed, which would admit action-capable calls.
-    completeTransactionally?: [FnType] extends ["mutation"] ? boolean : never;
+    completeTransactionally?: [FnType] extends ["mutation" | "query"]
+      ? boolean
+      : never;
   },
 ): Promise<WorkId> {
   const id = await ctx.runMutation(component.lib.enqueue, {