add finally callback

Author: ianmacartney

packages/convex-helpers/README.md

@@ -47,6 +47,8 @@ define custom behavior, allowing you to:
 - Consume arguments from the client that are not passed to the action, such
   as taking in an authentication parameter like an API key or session ID.
   These arguments must be sent up by the client along with each request.
+- Execute finalization logic after function execution using the `onSuccess`
+  callback, which has access to the function's result.
 
 See the associated [Stack Post](https://stack.convex.dev/custom-functions)
 
@@ -60,7 +62,15 @@ const myQueryBuilder = customQuery(query, {
   input: async (ctx, args) => {
     const apiUser = await getApiUser(args.apiToken);
     const db = wrapDatabaseReader({ apiUser }, ctx.db, rlsRules);
-    return { ctx: { db, apiUser }, args: {} };
+    return {
+      ctx: { db, apiUser },
+      args: {},
+      onSuccess: (result) => {
+        // Optional callback that runs after the function executes
+        // Has access to resources created during input processing
+        console.log(`Query for ${apiUser.name} completed:`, result);
+      },
+    };
   },
 });
 

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

@@ -36,6 +36,7 @@ import {
   describe,
   expect,
   test,
+  vi,
 } from "vitest";
 import { modules } from "./setup.test.js";
 
@@ -621,3 +622,116 @@ describe("extra args", () => {
     });
   });
 });
+
+describe("finally callback", () => {
+  test("finally callback is called with result and context", async () => {
+    const t = convexTest(schema, modules);
+    const finallyMock = vi.fn();
+
+    const withFinally = customQuery(query, {
+      args: {},
+      input: async () => ({
+        ctx: { foo: "bar" },
+        args: {},
+        onSuccess: (params) => {
+          finallyMock(params);
+        },
+      }),
+    });
+
+    const successFn = withFinally({
+      args: {},
+      handler: async (ctx) => {
+        return { success: true, foo: ctx.foo };
+      },
+    });
+
+    await t.run(async (ctx) => {
+      const result = await (successFn as any)._handler(ctx, {});
+      expect(result).toEqual({ success: true, foo: "bar" });
+
+      expect(finallyMock).toHaveBeenCalledWith({ success: true, foo: "bar" });
+    });
+
+    finallyMock.mockClear();
+
+    const errorFn = withFinally({
+      args: {},
+      handler: async () => {
+        throw new Error("Test error");
+      },
+    });
+
+    await t.run(async (ctx) => {
+      try {
+        await (errorFn as any)._handler(ctx, {});
+        expect.fail("Should have thrown an error");
+      } catch (e: unknown) {
+        const error = e as Error;
+        expect(error.message).toContain("Test error");
+      }
+
+      expect(finallyMock).not.toHaveBeenCalled();
+    });
+  });
+
+  test("finally callback with mutation", async () => {
+    const t = convexTest(schema, modules);
+    const finallyMock = vi.fn();
+
+    const withFinally = customMutation(mutation, {
+      args: {},
+      input: async () => ({
+        ctx: { foo: "bar" },
+        args: {},
+        onSuccess: (params) => {
+          finallyMock(params);
+        },
+      }),
+    });
+
+    const mutationFn = withFinally({
+      args: {},
+      handler: async (ctx) => {
+        return { updated: true, foo: ctx.foo };
+      },
+    });
+
+    await t.run(async (ctx) => {
+      const result = await (mutationFn as any)._handler(ctx, {});
+      expect(result).toEqual({ updated: true, foo: "bar" });
+
+      expect(finallyMock).toHaveBeenCalledWith({ updated: true, foo: "bar" });
+    });
+  });
+
+  test("finally callback with action", async () => {
+    const t = convexTest(schema, modules);
+    const finallyMock = vi.fn();
+
+    const withFinally = customAction(action, {
+      args: {},
+      input: async () => ({
+        ctx: { foo: "bar" },
+        args: {},
+        onSuccess: (params) => {
+          finallyMock(params);
+        },
+      }),
+    });
+
+    const actionFn = withFinally({
+      args: {},
+      handler: async (ctx) => {
+        return { executed: true, foo: ctx.foo };
+      },
+    });
+
+    await t.run(async (ctx) => {
+      const result = await (actionFn as any)._handler(ctx, {});
+      expect(result).toEqual({ executed: true, foo: "bar" });
+
+      expect(finallyMock).toHaveBeenCalledWith({ executed: true, foo: "bar" });
+    });
+  });
+});

packages/convex-helpers/server/customFunctions.ts

@@ -15,6 +15,7 @@ import type {
   ObjectType,
   PropertyValidators,
   Validator,
+  Value,
 } from "convex/values";
 import { asObjectValidator, v } from "convex/values";
 import type {
@@ -70,6 +71,10 @@ import { omit, pick } from "../index.js";
  * modified function. All returned ctx and args will show up in the type
  * signature for the modified function. To remove something from `ctx`, you
  * can return it as `undefined`.
+
+ *  The `input` function can also return an `onSuccess` callback that will be
+ * called after the function executes successfully. The `onSuccess` callback
+ * has access to resources created during input processing via closure.
  */
 export type Customization<
   Ctx extends Record<string, any>,
@@ -84,8 +89,16 @@ export type Customization<
     args: ObjectType<CustomArgsValidator>,
     extra: ExtraArgs,
   ) =>
-    | Promise<{ ctx: CustomCtx; args: CustomMadeArgs }>
-    | { ctx: CustomCtx; args: CustomMadeArgs };
+    | Promise<{
+        ctx: CustomCtx;
+        args: CustomMadeArgs;
+        onSuccess?: (result: Value) => void | Promise<void>;
+      }>
+    | {
+        ctx: CustomCtx;
+        args: CustomMadeArgs;
+        onSuccess?: (result: Value) => void | Promise<void>;
+      };
 };
 
 /**
@@ -185,7 +198,15 @@ export const NoOp = {
  *     const user = await getUserOrNull(ctx);
  *     const session = await db.get(sessionId);
  *     const db = wrapDatabaseReader({ user }, ctx.db, rlsRules);
- *     return { ctx: { db, user, session }, args: {} };
+ *     return {
+ *       ctx: { db, user, session },
+ *       args: {},
+ *       onSuccess: (result) => {
+ *         // Optional callback that runs after the function executes
+ *         // Has access to resources created during input processing
+ *         console.log(`Query for ${user.name} returned:`, result);
+ *       }
+ *     };
  *   },
  * });
  *
@@ -265,7 +286,15 @@ export function customQuery<
  *     const user = await getUserOrNull(ctx);
  *     const session = await db.get(sessionId);
  *     const db = wrapDatabaseReader({ user }, ctx.db, rlsRules);
- *     return { ctx: { db, user, session }, args: {} };
+ *     return {
+ *       ctx: { db, user, session },
+ *       args: {},
+ *       onSuccess: (result) => {
+ *         // Optional callback that runs after the function executes
+ *         // Has access to resources created during input processing
+ *         console.log(`User ${user.name} returned:`, result);
+ *       }
+ *     };
  *   },
  * });
  *
@@ -347,7 +376,17 @@ export function customMutation<
  *       throw new Error("Invalid secret key");
  *     }
  *     const user = await ctx.runQuery(internal.users.getUser, {});
- *     return { ctx: { user }, args: {} };
+ *     // Create resources that can be used in the onSuccess callback
+ *     const logger = createLogger();
+ *     return {
+ *       ctx: { user },
+ *       args: {},
+ *       onSuccess: (result) => {
+ *         // Optional callback that runs after the function executes
+ *         // Has access to resources created during input processing
+ *         logger.info(`Action for user ${user.name} returned:`, result);
+ *       }
+ *     };
  *   },
  * });
  *
@@ -444,7 +483,12 @@ function customFnBuilder(
             extra,
           );
           const args = omit(allArgs, Object.keys(inputArgs));
-          return handler({ ...ctx, ...added.ctx }, { ...args, ...added.args });
+          const finalCtx = { ...ctx, ...added.ctx };
+          const result = await handler(finalCtx, { ...args, ...added.args });
+          if (added.onSuccess) {
+            await added.onSuccess(result ?? null);
+          }
+          return result;
         },
       });
     }
@@ -458,7 +502,12 @@ function customFnBuilder(
       returns: fn.returns,
       handler: async (ctx: any, args: any) => {
         const added = await customInput(ctx, args, extra);
-        return handler({ ...ctx, ...added.ctx }, { ...args, ...added.args });
+        const finalCtx = { ...ctx, ...added.ctx };
+        const result = await handler(finalCtx, { ...args, ...added.args });
+        if (added.onSuccess) {
+          await added.onSuccess(result ?? null);
+        }
+        return result;
       },
     });
   };