add extra params to customFunction

Author: ianmacartney

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/server/customFunctions.test.ts

@@ -359,6 +359,22 @@ 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 +404,7 @@ const testApi: ApiFromModules<{
     create: typeof create;
     outerAdds: typeof outerAdds;
     outerRemoves: typeof outerRemoves;
+    extraArgQuery: typeof extraArgQuery;
   };
 }>["fns"] = anyApi["customFunctions.test"] as any;
 
@@ -569,3 +586,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 };
@@ -147,17 +149,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 +229,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 +246,8 @@ export function customMutation<
     ModCtx,
     ModMadeArgs,
     GenericMutationCtx<DataModel>,
-    Visibility
+    Visibility,
+    ExtraArgs
   >;
 }
 
@@ -298,44 +311,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 +375,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 +454,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 +464,14 @@ export type CustomBuilder<
       ArgsArrayForOptionalValidator<ArgsValidator> = DefaultArgsForOptionalValidator<ArgsValidator>,
   >(
     func:
-      | {
+      | ({
           args?: ArgsValidator;
           returns?: ReturnsValidator;
           handler: (
             ctx: Overwrite<InputCtx, ModCtx>,
             ...args: ArgsForHandlerType<OneOrZeroArgs, ModMadeArgs>
           ) => ReturnValue;
-        }
+        } & ExtraArgs)
       | {
           (
             ctx: Overwrite<InputCtx, ModCtx>,
@@ -474,6 +499,7 @@ export type CustomCtx<Builder> =
     infer ModCtx,
     any,
     infer InputCtx,
+    any,
     any
   >
     ? Overwrite<InputCtx, ModCtx>