feat(orchestration,workflows): pipe oncomplete and workflow preparation (#4697)

* chore: pipe onComplete and workflow preparation step

* changeset

* fix: tests

---------

Co-authored-by: Adrien de Peretti <adrien.deperetti@gmail.com>
Co-authored-by: Oli Juhl <59018053+olivermrbl@users.noreply.github.com>

packages/workflows/src/helper/__tests__/pipe.spec.ts

@@ -45,4 +45,44 @@ describe("Pipe", function () {
     expect(result).toBeDefined()
     expect(result).toEqual(output)
   })
+
+  it("should execute onComplete function if available but the output result shouldn't change", async function () {
+    const payload = { input: "input" }
+    const output = { test: "test" }
+    const invoke = {
+      input: payload,
+    }
+
+    const onComplete = jest.fn(async ({ data }) => {
+      data.__changed = true
+
+      return
+    })
+
+    const handler = jest.fn().mockImplementation(async () => output)
+    const input = {
+      inputAlias: "payload",
+      invoke: [
+        {
+          from: "payload",
+          alias: "input",
+        },
+      ],
+      onComplete,
+    }
+
+    const result = await pipe(input, handler)({ invoke, payload } as any)
+
+    expect(handler).toHaveBeenCalled()
+    expect(handler).toHaveBeenCalledWith(
+      expect.objectContaining({
+        data: {
+          input: payload,
+        },
+      })
+    )
+
+    expect(onComplete).toHaveBeenCalled()
+    expect(result).toEqual(output)
+  })
 })

packages/workflows/src/helper/__tests__/workflow-export.spec.ts

@@ -0,0 +1,64 @@
+import { exportWorkflow } from "../workflow-export"
+
+jest.mock("@medusajs/orchestration", () => {
+  return {
+    TransactionHandlerType: {
+      INVOKE: "invoke",
+      COMPENSATE: "compensate",
+    },
+    TransactionState: {
+      FAILED: "failed",
+      REVERTED: "reverted",
+    },
+    LocalWorkflow: jest.fn(() => {
+      return {
+        run: jest.fn(() => {
+          return {
+            getErrors: jest.fn(),
+            getState: jest.fn(() => "done"),
+            getContext: jest.fn(() => {
+              return {
+                invoke: { result_step: "invoke_test" },
+              }
+            }),
+          }
+        }),
+      }
+    }),
+  }
+})
+
+describe("Export Workflow", function () {
+  it("should prepare the input data before initializing the transaction", async function () {
+    let transformedInput
+    const prepare = jest.fn().mockImplementation(async (data) => {
+      data.__transformed = true
+      transformedInput = data
+
+      return data
+    })
+
+    const work = exportWorkflow("id" as any, "result_step", prepare)
+
+    const wfHandler = work()
+
+    const input = {
+      test: "payload",
+    }
+
+    const { result } = await wfHandler.run({
+      input,
+    })
+
+    expect(input).toEqual({
+      test: "payload",
+    })
+
+    expect(transformedInput).toEqual({
+      test: "payload",
+      __transformed: true,
+    })
+
+    expect(result).toEqual("invoke_test")
+  })
+})

packages/workflows/src/helper/pipe.ts

@@ -1,9 +1,10 @@
-import { Context, MedusaContainer, SharedContext } from "@medusajs/types"
 import {
   TransactionMetadata,
   WorkflowStepHandler,
 } from "@medusajs/orchestration"
+import { Context, MedusaContainer, SharedContext } from "@medusajs/types"
 
+import { DistributedTransaction } from "@medusajs/orchestration"
 import { InputAlias } from "../definitions"
 
 export type WorkflowStepMiddlewareReturn = {
@@ -20,6 +21,7 @@ interface PipelineInput {
   inputAlias?: InputAlias | string
   invoke?: WorkflowStepMiddlewareInput | WorkflowStepMiddlewareInput[]
   compensate?: WorkflowStepMiddlewareInput | WorkflowStepMiddlewareInput[]
+  onComplete?: (args: WorkflowOnCompleteArguments) => {}
 }
 
 export type WorkflowArguments<T = any> = {
@@ -30,6 +32,15 @@ export type WorkflowArguments<T = any> = {
   context: Context | SharedContext
 }
 
+export type WorkflowOnCompleteArguments<T = any> = {
+  container: MedusaContainer
+  payload: unknown
+  data: T
+  metadata: TransactionMetadata
+  transaction: DistributedTransaction
+  context: Context | SharedContext
+}
+
 export type PipelineHandler<T extends any = undefined> = (
   args: WorkflowArguments
 ) => Promise<
@@ -48,6 +59,7 @@ export function pipe<T>(
     invoke,
     compensate,
     metadata,
+    transaction,
     context,
   }) => {
     let data = {}
@@ -61,8 +73,9 @@ export function pipe<T>(
       Object.assign(original.invoke, { [input.inputAlias]: payload })
     }
 
-    for (const key in input) {
-      if (!input[key] || key === "inputAlias") {
+    const dataKeys = ["invoke", "compensate"]
+    for (const key of dataKeys) {
+      if (!input[key]) {
         continue
       }
 
@@ -111,6 +124,18 @@ export function pipe<T>(
       finalResult = result
     }
 
+    if (typeof input.onComplete === "function") {
+      const dataCopy = JSON.parse(JSON.stringify(data))
+      await input.onComplete({
+        container,
+        payload,
+        data: dataCopy,
+        metadata,
+        transaction,
+        context: context as Context,
+      })
+    }
+
     return finalResult
   }
 }

packages/workflows/src/helper/workflow-export.ts

@@ -27,7 +27,8 @@ export type WorkflowResult<TResult = unknown> = {
 
 export const exportWorkflow = <TData = unknown, TResult = unknown>(
   workflowId: Workflows,
-  defaultResult?: string
+  defaultResult?: string,
+  dataPreparation?: (data: TData) => Promise<unknown>
 ) => {
   return function <TDataOverride = undefined, TResultOverride = undefined>(
     container?: LoadedModule[] | MedusaContainer
@@ -60,6 +61,22 @@ export const exportWorkflow = <TData = unknown, TResult = unknown>(
       resultFrom ??= defaultResult
       throwOnError ??= true
 
+      if (typeof dataPreparation === "function") {
+        try {
+          const copyInput = JSON.parse(JSON.stringify(input))
+          input = await dataPreparation(copyInput as TData)
+        } catch (err) {
+          if (throwOnError) {
+            throw new Error(
+              `Data preparation failed: ${err.message}${EOL}${err.stack}`
+            )
+          }
+          return {
+            errors: [err],
+          }
+        }
+      }
+
       const transaction = await originalRun(
         context?.transactionId ?? ulid(),
         input,