support transactionLimits

CHANGELOG.md

@@ -1,5 +1,16 @@
 # Changelog
 
+## Unreleased
+
+- Support the `transactionLimits` option on nested `ctx.runQuery` /
+  `ctx.runMutation` calls (Convex 1.41). The nested call is enforced against
+  its own limits, capped at the global transaction limits so they can only be
+  lowered, never raised.
+- Bandwidth tracking now mirrors the database's nested-transaction layers: a
+  nested call's usage folds into its parent only when it commits, so the writes
+  of a rolled-back nested `ctx.runMutation` no longer count against the
+  transaction's limits.
+
 ## 0.0.53
 
 - Support scheduled functions correctly with or without fake timers,

convex/_generated/api.d.ts

@@ -21,6 +21,7 @@ import type * as helpers from "../helpers.js";
 import type * as http from "../http.js";
 import type * as messages from "../messages.js";
 import type * as mutations from "../mutations.js";
+import type * as nestedLimits from "../nestedLimits.js";
 import type * as pagination from "../pagination.js";
 import type * as queries from "../queries.js";
 import type * as returnsValidation from "../returnsValidation.js";
@@ -49,6 +50,7 @@ declare const fullApi: ApiFromModules<{
   http: typeof http;
   messages: typeof messages;
   mutations: typeof mutations;
+  nestedLimits: typeof nestedLimits;
   pagination: typeof pagination;
   queries: typeof queries;
   returnsValidation: typeof returnsValidation;

convex/nestedLimits.test.ts

@@ -0,0 +1,140 @@
+import { expect, test } from "vitest";
+import { convexTest } from "../index";
+import { api } from "./_generated/api";
+import schema from "./schema";
+
+const MiB = 1 << 20;
+
+async function seed(t: ReturnType<typeof convexTest>, count: number) {
+  await t.run(async (ctx) => {
+    for (let i = 0; i < count; i++) {
+      await ctx.db.insert("messages", { author: "sarah", body: `msg${i}` });
+    }
+  });
+}
+
+test("nested runQuery enforces its own documentsRead limit", async () => {
+  const t = convexTest({ schema });
+  await seed(t, 10);
+  // Global limits are disabled, but the nested call explicitly asks for a
+  // tight limit, so it should be enforced.
+  await expect(
+    t.query(api.nestedLimits.parentReadWithLimits, {
+      transactionLimits: { documentsRead: 5 },
+    }),
+  ).rejects.toThrow(/Scanned too many documents/);
+});
+
+test("nested runQuery under its limit succeeds", async () => {
+  const t = convexTest({ schema });
+  await seed(t, 3);
+  const count = await t.query(api.nestedLimits.parentReadWithLimits, {
+    transactionLimits: { documentsRead: 5 },
+  });
+  expect(count).toBe(3);
+});
+
+test("nested runMutation enforces its own documentsWritten limit", async () => {
+  const t = convexTest({ schema });
+  await expect(
+    t.mutation(api.nestedLimits.parentInsertWithLimits, {
+      count: 10,
+      transactionLimits: { documentsWritten: 3 },
+    }),
+  ).rejects.toThrow(/Wrote too many documents/);
+});
+
+test("nested runMutation under its limit succeeds", async () => {
+  const t = convexTest({ schema });
+  await t.mutation(api.nestedLimits.parentInsertWithLimits, {
+    count: 2,
+    transactionLimits: { documentsWritten: 5 },
+  });
+  const count = await t.query(api.nestedLimits.readAll, {});
+  expect(count).toBe(2);
+});
+
+test("nested limits cannot raise above the global limit", async () => {
+  const t = convexTest({
+    schema,
+    transactionLimits: { documentsRead: 5 },
+  });
+  await seed(t, 10);
+  // The nested call asks for documentsRead: 100000, but it is capped at the
+  // global limit of 5, so reading 10 docs still throws.
+  await expect(
+    t.query(api.nestedLimits.parentReadWithLimits, {
+      transactionLimits: { documentsRead: 100000 },
+    }),
+  ).rejects.toThrow(/Scanned too many documents/);
+});
+
+test("global limit still applies when the nested limit is looser", async () => {
+  const t = convexTest({
+    schema,
+    transactionLimits: { documentsRead: 5 },
+  });
+  await seed(t, 10);
+  // Nested limit (8) is looser, but the global limit (5) is the binding
+  // constraint and is enforced too.
+  await expect(
+    t.query(api.nestedLimits.parentReadWithLimits, {
+      transactionLimits: { documentsRead: 8 },
+    }),
+  ).rejects.toThrow(/Scanned too many documents/);
+});
+
+test("nested limit does not leak to the parent transaction", async () => {
+  const t = convexTest({
+    schema,
+    // Generous global limit so only the nested scope is restrictive.
+    transactionLimits: { documentsRead: 100, bytesRead: 16 * MiB },
+  });
+  await seed(t, 10);
+  // The nested query reads 2 docs under a tight limit of 3 (ok). After it
+  // returns, the parent reads all 10 docs itself. If the nested limit leaked,
+  // this would throw; under the global limit of 100 it succeeds.
+  const count = await t.query(api.nestedLimits.nestThenReadAll, {
+    nestedReadCount: 2,
+    transactionLimits: { documentsRead: 3 },
+  });
+  expect(count).toBe(10);
+});
+
+test("rolled-back nested writes are not counted against the transaction", async () => {
+  const t = convexTest({ schema });
+  // The nested mutation writes 5 docs then throws, so its writes roll back.
+  // The parent swallows the error and reads back the transaction metrics: the
+  // rolled-back writes should not be counted.
+  const documentsWritten = await t.mutation(
+    api.nestedLimits.insertRollbackThenReportMetrics,
+    { count: 5 },
+  );
+  expect(documentsWritten).toBe(0);
+  // And nothing was actually persisted.
+  const count = await t.query(api.nestedLimits.readAll, {});
+  expect(count).toBe(0);
+});
+
+test("rolled-back nested writes do not consume the global write limit", async () => {
+  const t = convexTest({
+    schema,
+    transactionLimits: { documentsWritten: 5 },
+  });
+  // Nested mutation writes 4 docs then throws (rolled back). Afterwards the
+  // parent can still write 4 docs of its own: if the rolled-back writes had
+  // been counted, 4 + 4 would exceed the limit of 5.
+  await t.mutation(api.nestedLimits.parentInsertAfterRollback, {
+    rolledBackCount: 4,
+    keptCount: 4,
+  });
+  const count = await t.query(api.nestedLimits.readAll, {});
+  expect(count).toBe(4);
+});
+
+test("nested transactionLimits is optional", async () => {
+  const t = convexTest({ schema });
+  await seed(t, 10);
+  const count = await t.query(api.nestedLimits.parentReadNoLimits, {});
+  expect(count).toBe(10);
+});

convex/nestedLimits.ts

@@ -0,0 +1,139 @@
+import { v } from "convex/values";
+import { api } from "./_generated/api";
+import { mutation, query } from "./_generated/server";
+
+const transactionLimitsValidator = v.object({
+  bytesRead: v.optional(v.number()),
+  bytesWritten: v.optional(v.number()),
+  databaseQueries: v.optional(v.number()),
+  documentsRead: v.optional(v.number()),
+  documentsWritten: v.optional(v.number()),
+  functionsScheduled: v.optional(v.number()),
+  scheduledFunctionArgsBytes: v.optional(v.number()),
+});
+
+// Child query that reads every message.
+export const readAll = query(async (ctx) => {
+  const docs = await ctx.db.query("messages").collect();
+  return docs.length;
+});
+
+// Child query that reads exactly `n` messages.
+export const readN = query({
+  args: { n: v.number() },
+  handler: async (ctx, { n }) => {
+    const docs = await ctx.db.query("messages").take(n);
+    return docs.length;
+  },
+});
+
+// Child mutation that inserts `count` messages.
+export const insertMany = mutation({
+  args: { count: v.number() },
+  handler: async (ctx, { count }) => {
+    for (let i = 0; i < count; i++) {
+      await ctx.db.insert("messages", { author: "child", body: `msg${i}` });
+    }
+  },
+});
+
+// Child mutation that inserts `count` messages and then throws, so all of its
+// writes are rolled back.
+export const insertManyThenThrow = mutation({
+  args: { count: v.number() },
+  handler: async (ctx, { count }) => {
+    for (let i = 0; i < count; i++) {
+      await ctx.db.insert("messages", { author: "child", body: `msg${i}` });
+    }
+    throw new Error("nested boom");
+  },
+});
+
+// Parent mutation that calls a nested mutation which rolls back, swallows the
+// error, then reports the transaction metrics. Used to verify the rolled-back
+// nested writes are not counted against the transaction.
+export const insertRollbackThenReportMetrics = mutation({
+  args: { count: v.number() },
+  handler: async (ctx, { count }) => {
+    try {
+      await ctx.runMutation(api.nestedLimits.insertManyThenThrow, { count });
+    } catch {
+      // Swallow: the nested mutation's writes have been rolled back.
+    }
+    const metrics = await ctx.meta.getTransactionMetrics();
+    return metrics.documentsWritten.used;
+  },
+});
+
+// Parent mutation that runs a nested mutation which rolls back, then inserts
+// `keptCount` messages of its own. Used to verify rolled-back nested writes
+// don't consume the global write limit.
+export const parentInsertAfterRollback = mutation({
+  args: { rolledBackCount: v.number(), keptCount: v.number() },
+  handler: async (ctx, { rolledBackCount, keptCount }) => {
+    try {
+      await ctx.runMutation(api.nestedLimits.insertManyThenThrow, {
+        count: rolledBackCount,
+      });
+    } catch {
+      // Swallow: the nested mutation's writes have been rolled back.
+    }
+    for (let i = 0; i < keptCount; i++) {
+      await ctx.db.insert("messages", { author: "parent", body: `kept${i}` });
+    }
+  },
+});
+
+// Parent query that calls a nested query with custom transaction limits.
+export const parentReadWithLimits = query({
+  args: { transactionLimits: transactionLimitsValidator },
+  handler: async (ctx, { transactionLimits }): Promise<number> => {
+    return await ctx.runQuery(
+      api.nestedLimits.readAll,
+      {},
+      { transactionLimits },
+    );
+  },
+});
+
+// Parent query that calls a nested query without any custom limits.
+export const parentReadNoLimits = query(async (ctx): Promise<number> => {
+  return await ctx.runQuery(api.nestedLimits.readAll, {});
+});
+
+// Parent mutation that calls a nested mutation with custom transaction limits.
+export const parentInsertWithLimits = mutation({
+  args: {
+    count: v.number(),
+    transactionLimits: transactionLimitsValidator,
+  },
+  handler: async (ctx, { count, transactionLimits }) => {
+    await ctx.runMutation(
+      api.nestedLimits.insertMany,
+      { count },
+      { transactionLimits },
+    );
+  },
+});
+
+// Parent query that runs a tightly-limited nested query, then reads more on
+// its own. Used to verify the nested scope's limit does not leak back to the
+// parent transaction once the nested call returns.
+export const nestThenReadAll = query({
+  args: {
+    nestedReadCount: v.number(),
+    transactionLimits: transactionLimitsValidator,
+  },
+  handler: async (
+    ctx,
+    { nestedReadCount, transactionLimits },
+  ): Promise<number> => {
+    await ctx.runQuery(
+      api.nestedLimits.readN,
+      { n: nestedReadCount },
+      { transactionLimits },
+    );
+    const docs = await ctx.db.query("messages").collect();
+    return docs.length;
+  },
+});