Custom function allows extra params

convex/testingFunctions.ts

@@ -21,12 +21,15 @@ export const testingQuery = customQuery(query, {
 
 export const testingMutation = customMutation(mutation, {
   args: {},
-  input: async (_ctx, _args) => {
+  input: async (_ctx, _args, { devOnly }: { devOnly: boolean }) => {
     if (process.env.IS_TEST === undefined) {
       throw new Error(
         "Calling a test only function in an unexpected environment",
       );
     }
+    if (devOnly && process.env.IS_PROD) {
+      throw new Error("This function is only available in development");
+    }
     return { ctx: {}, args: {} };
   },
 });
@@ -43,13 +46,16 @@ export const testingAction = customAction(action, {
   },
 });
 
-export const clearAll = testingMutation(async ({ db, scheduler, storage }) => {
-  for (const table of Object.keys(schema.tables)) {
-    const docs = await db.query(table as any).collect();
-    await Promise.all(docs.map((doc) => db.delete(doc._id)));
-  }
-  const scheduled = await db.system.query("_scheduled_functions").collect();
-  await Promise.all(scheduled.map((s) => scheduler.cancel(s._id)));
-  const storedFiles = await db.system.query("_storage").collect();
-  await Promise.all(storedFiles.map((s) => storage.delete(s._id)));
+export const clearAll = testingMutation({
+  devOnly: true,
+  handler: async ({ db, scheduler, storage }) => {
+    for (const table of Object.keys(schema.tables)) {
+      const docs = await db.query(table as any).collect();
+      await Promise.all(docs.map((doc) => db.delete(doc._id)));
+    }
+    const scheduled = await db.system.query("_scheduled_functions").collect();
+    await Promise.all(scheduled.map((s) => scheduler.cancel(s._id)));
+    const storedFiles = await db.system.query("_storage").collect();
+    await Promise.all(storedFiles.map((s) => storage.delete(s._id)));
+  },
 });

packages/convex-helpers/README.md

@@ -52,7 +52,7 @@ See the associated [Stack Post](https://stack.convex.dev/custom-functions)
 
 For example:
 
-```js
+```ts
 import { customQuery } from "convex-helpers/server/customFunctions.js";
 
 const myQueryBuilder = customQuery(query, {
@@ -75,6 +75,34 @@ export const getSomeData = myQueryBuilder({
 });
 ```
 
+### Taking in extra arguments
+
+You can take in extra arguments to a custom function by specifying the type of a third `input` arg.
+
+```ts
+const myQueryBuilder = customQuery(query, {
+  args: {},
+  input: async (ctx, args, { role }: { role: "admin" | "user" }) => {
+    const user = await getUser(ctx);
+    if (role === "admin" && user.role !== "admin") {
+      throw new Error("You are not an admin");
+    }
+    if (role === "user" && !user) {
+      throw new Error("You must be logged in to access this query");
+    }
+    return { ctx: { user }, args: {} };
+  },
+});
+
+const myAdminQuery = myQueryBuilder({
+  role: "admin",
+  args: {},
+  handler: async (ctx, args) => {
+    // ...
+  },
+});
+```
+
 ## Relationship helpers
 
 Traverse database relationships without all the query boilerplate.

packages/convex-helpers/server/customFunctions.test.ts

@@ -147,6 +147,30 @@ customAction(
   customCtx((ctx) => ({})),
 ) satisfies typeof action;
 
+customQuery({} as any, {
+  args: {},
+  input: async () => ({
+    ctx: {},
+    args: {},
+  }),
+}) satisfies typeof query;
+
+customMutation(mutation, {
+  args: {},
+  input: async () => ({
+    ctx: {},
+    args: {},
+  }),
+}) satisfies typeof mutation;
+
+customAction(action, {
+  args: {},
+  input: async () => ({
+    ctx: {},
+    args: {},
+  }),
+}) satisfies typeof action;
+
 /**
  * Testing custom function modifications.
  */
@@ -359,6 +383,25 @@ export const outerRemoves = outerRemover({
   },
 });
 
+/**
+ * Adding extra args to `input`
+ */
+const extraArgQueryBuilder = customQuery(query, {
+  args: { a: v.string() },
+  input: async (_ctx, args, { extraArg }: { extraArg: string }) => ({
+    ctx: { extraArg },
+    args,
+  }),
+});
+export const extraArgQuery = extraArgQueryBuilder({
+  args: {},
+  extraArg: "foo",
+  handler: async (ctx, args) => {
+    return { ctxA: ctx.extraArg };
+  },
+});
+queryMatches(extraArgQuery, {}, { ctxA: "foo" });
+
 /**
  * Test helpers
  */
@@ -388,6 +431,7 @@ const testApi: ApiFromModules<{
     create: typeof create;
     outerAdds: typeof outerAdds;
     outerRemoves: typeof outerRemoves;
+    extraArgQuery: typeof extraArgQuery;
   };
 }>["fns"] = anyApi["customFunctions.test"] as any;
 
@@ -569,3 +613,12 @@ describe("nested custom functions", () => {
     ).rejects.toThrow("Validator error: Expected `string`");
   });
 });
+
+describe("extra args", () => {
+  test("add extra args", async () => {
+    const t = convexTest(schema, modules);
+    expect(await t.query(testApi.extraArgQuery, { a: "foo" })).toMatchObject({
+      ctxA: "foo",
+    });
+  });
+});

packages/convex-helpers/server/customFunctions.ts

@@ -51,11 +51,13 @@ export type Mod<
   ModArgsValidator extends PropertyValidators,
   ModCtx extends Record<string, any>,
   ModMadeArgs extends Record<string, any>,
+  ExtraArgs extends Record<string, any> = {},
 > = {
   args: ModArgsValidator;
   input: (
     ctx: Ctx,
     args: ObjectType<ModArgsValidator>,
+    extra: ExtraArgs,
   ) =>
     | Promise<{ ctx: ModCtx; args: ModMadeArgs }>
     | { ctx: ModCtx; args: ModMadeArgs };
@@ -70,12 +72,13 @@ export type Mod<
 export function customCtx<
   InCtx extends Record<string, any>,
   OutCtx extends Record<string, any>,
+  ExtraArgs extends Record<string, any> = {},
 >(
-  mod: (original: InCtx) => Promise<OutCtx> | OutCtx,
-): Mod<InCtx, {}, OutCtx, {}> {
+  mod: (original: InCtx, extra: ExtraArgs) => Promise<OutCtx> | OutCtx,
+): Mod<InCtx, {}, OutCtx, {}, ExtraArgs> {
   return {
     args: {},
-    input: async (ctx) => ({ ctx: await mod(ctx), args: {} }),
+    input: async (ctx, _, extra) => ({ ctx: await mod(ctx, extra), args: {} }),
   };
 }
 
@@ -147,17 +150,25 @@ export function customQuery<
   ModMadeArgs extends Record<string, any>,
   Visibility extends FunctionVisibility,
   DataModel extends GenericDataModel,
+  ExtraArgs extends Record<string, any> = {},
 >(
   query: QueryBuilder<DataModel, Visibility>,
-  mod: Mod<GenericQueryCtx<DataModel>, ModArgsValidator, ModCtx, ModMadeArgs>,
+  mod: Mod<
+    GenericQueryCtx<DataModel>,
+    ModArgsValidator,
+    ModCtx,
+    ModMadeArgs,
+    ExtraArgs
+  >,
 ) {
   return customFnBuilder(query, mod) as CustomBuilder<
     "query",
     ModArgsValidator,
     ModCtx,
     ModMadeArgs,
     GenericQueryCtx<DataModel>,
-    Visibility
+    Visibility,
+    ExtraArgs
   >;
 }
 
@@ -219,13 +230,15 @@ export function customMutation<
   ModMadeArgs extends Record<string, any>,
   Visibility extends FunctionVisibility,
   DataModel extends GenericDataModel,
+  ExtraArgs extends Record<string, any> = {},
 >(
   mutation: MutationBuilder<DataModel, Visibility>,
   mod: Mod<
     GenericMutationCtx<DataModel>,
     ModArgsValidator,
     ModCtx,
-    ModMadeArgs
+    ModMadeArgs,
+    ExtraArgs
   >,
 ) {
   return customFnBuilder(mutation, mod) as CustomBuilder<
@@ -234,7 +247,8 @@ export function customMutation<
     ModCtx,
     ModMadeArgs,
     GenericMutationCtx<DataModel>,
-    Visibility
+    Visibility,
+    ExtraArgs
   >;
 }
 
@@ -298,44 +312,55 @@ export function customAction<
   ModMadeArgs extends Record<string, any>,
   Visibility extends FunctionVisibility,
   DataModel extends GenericDataModel,
+  ExtraArgs extends Record<string, any> = {},
 >(
   action: ActionBuilder<DataModel, Visibility>,
-  mod: Mod<GenericActionCtx<DataModel>, ModArgsValidator, ModCtx, ModMadeArgs>,
+  mod: Mod<
+    GenericActionCtx<DataModel>,
+    ModArgsValidator,
+    ModCtx,
+    ModMadeArgs,
+    ExtraArgs
+  >,
 ): CustomBuilder<
   "action",
   ModArgsValidator,
   ModCtx,
   ModMadeArgs,
   GenericActionCtx<DataModel>,
-  Visibility
+  Visibility,
+  ExtraArgs
 > {
   return customFnBuilder(action, mod) as CustomBuilder<
     "action",
     ModArgsValidator,
     ModCtx,
     ModMadeArgs,
     GenericActionCtx<DataModel>,
-    Visibility
+    Visibility,
+    ExtraArgs
   >;
 }
 
 function customFnBuilder(
   builder: (args: any) => any,
-  mod: Mod<any, any, any, any>,
+  mod: Mod<any, any, any, any, any>,
 ) {
   // Looking forward to when input / args / ... are optional
   const inputMod = mod.input ?? NoOp.input;
   const inputArgs = mod.args ?? NoOp.args;
   return function customBuilder(fn: any): any {
-    const handler = fn.handler ?? fn;
-    if ("args" in fn) {
+    // N.B.: This is fine if it's a function
+    const { args, handler = fn, returns, ...extra } = fn;
+    if (args) {
       return builder({
-        args: addArgs(fn.args, inputArgs),
-        returns: fn.returns,
+        args: addArgs(args, inputArgs),
+        returns,
         handler: async (ctx: any, allArgs: any) => {
           const added = await inputMod(
             ctx,
             pick(allArgs, Object.keys(inputArgs)) as any,
+            extra,
           );
           const args = omit(allArgs, Object.keys(inputArgs));
           return handler({ ...ctx, ...added.ctx }, { ...args, ...added.args });
@@ -351,7 +376,7 @@ function customFnBuilder(
     return builder({
       returns: fn.returns,
       handler: async (ctx: any, args: any) => {
-        const added = await inputMod(ctx, args);
+        const added = await inputMod(ctx, args, extra);
         return handler({ ...ctx, ...added.ctx }, { ...args, ...added.args });
       },
     });
@@ -430,6 +455,7 @@ export type CustomBuilder<
   ModMadeArgs extends Record<string, any>,
   InputCtx,
   Visibility extends FunctionVisibility,
+  ExtraArgs extends Record<string, any>,
 > = {
   <
     ArgsValidator extends PropertyValidators | void | Validator<any, any, any>,
@@ -439,14 +465,18 @@ export type CustomBuilder<
       ArgsArrayForOptionalValidator<ArgsValidator> = DefaultArgsForOptionalValidator<ArgsValidator>,
   >(
     func:
-      | {
+      | ({
           args?: ArgsValidator;
           returns?: ReturnsValidator;
           handler: (
             ctx: Overwrite<InputCtx, ModCtx>,
             ...args: ArgsForHandlerType<OneOrZeroArgs, ModMadeArgs>
           ) => ReturnValue;
-        }
+        } & {
+          [key in keyof ExtraArgs as key extends "args" | "returns" | "handler"
+            ? never
+            : key]: ExtraArgs[key];
+        })
       | {
           (
             ctx: Overwrite<InputCtx, ModCtx>,
@@ -474,9 +504,10 @@ export type CustomCtx<Builder> =
     infer ModCtx,
     any,
     infer InputCtx,
+    any,
     any
   >
     ? Overwrite<InputCtx, ModCtx>
     : never;
 
-type Overwrite<T, U> = Omit<T, keyof U> & U;
+type Overwrite<T, U> = keyof U extends never ? T : Omit<T, keyof U> & U;