diff --git a/packages/convex-helpers/server/customFunctions.test.ts b/packages/convex-helpers/server/customFunctions.test.ts index 8b3156ce..3f4a71ad 100644 --- a/packages/convex-helpers/server/customFunctions.test.ts +++ b/packages/convex-helpers/server/customFunctions.test.ts @@ -27,7 +27,7 @@ import { type Auth, } from "convex/server"; import { v } from "convex/values"; -import { afterEach, beforeEach, describe, expect, test } from "vitest"; +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; import { modules } from "./setup.test.js"; const schema = defineSchema({ @@ -560,3 +560,119 @@ describe("nested custom functions", () => { ).rejects.toThrow("Validator error: Expected `string`"); }); }); + +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: {} }), + finally: (ctx, params) => { + finallyMock(ctx, 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( + expect.objectContaining({ foo: "bar" }), + { result: { 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).toHaveBeenCalledWith( + expect.objectContaining({ foo: "bar" }), + { error: expect.objectContaining({ message: expect.stringContaining("Test error") }) } + ); + }); + }); + + 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: {} }), + finally: (ctx, params) => { + finallyMock(ctx, 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( + expect.objectContaining({ foo: "bar" }), + { result: { 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: {} }), + finally: (ctx, params) => { + finallyMock(ctx, 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( + expect.objectContaining({ foo: "bar" }), + { result: { executed: true, foo: "bar" } } + ); + }); + }); +}); diff --git a/packages/convex-helpers/server/customFunctions.ts b/packages/convex-helpers/server/customFunctions.ts index 3e24bd26..0253d990 100644 --- a/packages/convex-helpers/server/customFunctions.ts +++ b/packages/convex-helpers/server/customFunctions.ts @@ -60,6 +60,10 @@ export type Mod< ) => | Promise<{ ctx: ModCtx; args: ModMadeArgs }> | { ctx: ModCtx; args: ModMadeArgs }; + finally?: (ctx: Ctx & ModCtx, params: { + result?: unknown; + error?: unknown; + }) => void | Promise; }; /** @@ -88,6 +92,7 @@ export const NoOp = { input() { return { args: {}, ctx: {} }; }, + }; /** @@ -339,7 +344,20 @@ function customFnBuilder( pick(allArgs, Object.keys(inputArgs)) as any, ); const args = omit(allArgs, Object.keys(inputArgs)); - return handler({ ...ctx, ...added.ctx }, { ...args, ...added.args }); + const finalCtx = { ...ctx, ...added.ctx }; + let result; + try { + result = await handler(finalCtx, { ...args, ...added.args }); + if (mod.finally) { + await mod.finally(finalCtx, { result }); + } + return result; + } catch (e) { + if (mod.finally) { + await mod.finally(finalCtx, { error: e }); + } + throw e; + } }, }); } @@ -353,7 +371,20 @@ function customFnBuilder( returns: fn.returns, handler: async (ctx: any, args: any) => { const added = await inputMod(ctx, args); - return handler({ ...ctx, ...added.ctx }, { ...args, ...added.args }); + const finalCtx = { ...ctx, ...added.ctx }; + let result; + try { + result = await handler(finalCtx, { ...args, ...added.args }); + if (mod.finally) { + await mod.finally(finalCtx, { result }); + } + return result; + } catch (e) { + if (mod.finally) { + await mod.finally(finalCtx, { error: e }); + } + throw e; + } }, }); };