feat(workflows-sdk): run, registerStepSuccess and registerStepFailure shortcut (#6164)

**What**

Currently, when exporting a workflow using the workflowExport util, it is mandatory to first call the resulted function passing the container before being able to call run, registerStepSuccess or failure on it. Now, it is possible to either continue that way, or to directly call the run, registerStepSuccess or failure on the exported workflow and at that moment it is possible to pass a container if needed

e.g
```ts
const workflow = exportWorkflow("id" as any, "result_step", prepare)
const wfRunner = workflow(container)
wfRunner.run(...) // Here the container is not expected
```
or
```ts
const workflow = exportWorkflow("id" as any, "result_step", prepare)
workflow.run(...) // Here we can now pass an optional container 
```

.changeset/red-emus-carry.md

@@ -0,0 +1,5 @@
+---
+"@medusajs/workflows-sdk": patch
+---
+
+feat(workflows-sdk): run, registerStepSuccess and registerStepFailure shortcut

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

@@ -1,4 +1,5 @@
 import { exportWorkflow } from "../workflow-export"
+import { createMedusaContainer } from "@medusajs/utils"
 
 jest.mock("@medusajs/orchestration", () => {
   return {
@@ -62,7 +63,8 @@ describe("Export Workflow", function () {
 
     const work = exportWorkflow("id" as any, "result_step", prepare)
 
-    const wfHandler = work()
+    const container = createMedusaContainer()
+    const wfHandler = work(container)
 
     const input = {
       test: "payload",
@@ -83,4 +85,40 @@ describe("Export Workflow", function () {
 
     expect(result).toEqual("invoke_test")
   })
+
+  describe("Using the exported workflow run", 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 input = {
+        test: "payload",
+      }
+
+      const container = createMedusaContainer()
+
+      const { result } = await work.run({
+        input,
+        container,
+      })
+
+      expect(input).toEqual({
+        test: "payload",
+      })
+
+      expect(transformedInput).toEqual({
+        test: "payload",
+        __transformed: true,
+      })
+
+      expect(result).toEqual("invoke_test")
+    })
+  })
 })

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

@@ -81,23 +81,77 @@ export type ExportedWorkflow<
   >
 }
 
-export const exportWorkflow = <TData = unknown, TResult = unknown>(
+export type MainExportedWorkflow<TData = unknown, TResult = unknown> = {
-  workflowId: string,
+  // Main function on the exported workflow
-  defaultResult?: string | Symbol,
+  <TDataOverride = undefined, TResultOverride = undefined>(
-  dataPreparation?: (data: TData) => Promise<unknown>,
-  options?: {
-    wrappedInput?: boolean
-  }
-) => {
-  function exportedWorkflow<
-    TDataOverride = undefined,
-    TResultOverride = undefined
-  >(
     container?: LoadedModule[] | MedusaContainer
   ): Omit<
     LocalWorkflow,
     "run" | "registerStepSuccess" | "registerStepFailure"
   > &
+    ExportedWorkflow<TData, TResult, TDataOverride, TResultOverride>
+
+  /**
+   * You can also directly call run, registerStepSuccess and registerStepFailure on the exported workflow
+   */
+
+  run<TDataOverride = undefined, TResultOverride = undefined>(
+    args?: FlowRunOptions<
+      TDataOverride extends undefined ? TData : TDataOverride
+    > & {
+      container?: LoadedModule[] | MedusaContainer
+    }
+  ): Promise<
+    WorkflowResult<
+      TResultOverride extends undefined ? TResult : TResultOverride
+    >
+  >
+
+  registerStepSuccess<TDataOverride = undefined, TResultOverride = undefined>(
+    args?: FlowRegisterStepSuccessOptions<
+      TDataOverride extends undefined ? TData : TDataOverride
+    > & {
+      container?: LoadedModule[] | MedusaContainer
+    }
+  ): Promise<
+    WorkflowResult<
+      TResultOverride extends undefined ? TResult : TResultOverride
+    >
+  >
+
+  registerStepFailure<TDataOverride = undefined, TResultOverride = undefined>(
+    args?: FlowRegisterStepFailureOptions<
+      TDataOverride extends undefined ? TData : TDataOverride
+    > & {
+      container?: LoadedModule[] | MedusaContainer
+    }
+  ): Promise<
+    WorkflowResult<
+      TResultOverride extends undefined ? TResult : TResultOverride
+    >
+  >
+}
+
+function createContextualWorkflowRunner<
+  TData = unknown,
+  TResult = unknown,
+  TDataOverride = undefined,
+  TResultOverride = undefined
+>({
+  workflowId,
+  defaultResult,
+  dataPreparation,
+  options,
+  container,
+}: {
+  workflowId: string
+  defaultResult?: string | Symbol
+  dataPreparation?: (data: TData) => Promise<unknown>
+  options?: {
+    wrappedInput?: boolean
+  }
+  container?: LoadedModule[] | MedusaContainer
+}): Omit<LocalWorkflow, "run" | "registerStepSuccess" | "registerStepFailure"> &
   ExportedWorkflow<TData, TResult, TDataOverride, TResultOverride> {
   if (!container) {
     container = MedusaModule.getLoadedModules().map(
@@ -238,6 +292,153 @@ export const exportWorkflow = <TData = unknown, TResult = unknown>(
     ExportedWorkflow<TData, TResult, TDataOverride, TResultOverride>
 }
 
+export const exportWorkflow = <TData = unknown, TResult = unknown>(
+  workflowId: string,
+  defaultResult?: string | Symbol,
+  dataPreparation?: (data: TData) => Promise<unknown>,
+  options?: {
+    wrappedInput?: boolean
+  }
+): MainExportedWorkflow<TData, TResult> => {
+  function exportedWorkflow<
+    TDataOverride = undefined,
+    TResultOverride = undefined
+  >(
+    container?: LoadedModule[] | MedusaContainer
+  ): Omit<
+    LocalWorkflow,
+    "run" | "registerStepSuccess" | "registerStepFailure"
+  > &
+    ExportedWorkflow<TData, TResult, TDataOverride, TResultOverride> {
+    return createContextualWorkflowRunner<
+      TData,
+      TResult,
+      TDataOverride,
+      TResultOverride
+    >({
+      workflowId,
+      defaultResult,
+      dataPreparation,
+      options,
+      container,
+    })
+  }
+
+  const buildRunnerFn = <
+    TAction extends "run" | "registerStepSuccess" | "registerStepFailure",
+    TDataOverride,
+    TResultOverride
+  >(
+    action: "run" | "registerStepSuccess" | "registerStepFailure",
+    container?: LoadedModule[] | MedusaContainer
+  ) => {
+    const contextualRunner = createContextualWorkflowRunner<
+      TData,
+      TResult,
+      TDataOverride,
+      TResultOverride
+    >({
+      workflowId,
+      defaultResult,
+      dataPreparation,
+      options,
+      container,
+    })
+
+    return contextualRunner[action] as ExportedWorkflow<
+      TData,
+      TResult,
+      TDataOverride,
+      TResultOverride
+    >[TAction]
+  }
+
+  exportedWorkflow.run = async <
+    TDataOverride = undefined,
+    TResultOverride = undefined
+  >(
+    args?: FlowRunOptions<
+      TDataOverride extends undefined ? TData : TDataOverride
+    > & {
+      container?: LoadedModule[] | MedusaContainer
+    }
+  ): Promise<
+    WorkflowResult<
+      TResultOverride extends undefined ? TResult : TResultOverride
+    >
+  > => {
+    const container = args?.container
+    delete args?.container
+    const inputArgs = { ...args } as FlowRunOptions<
+      TDataOverride extends undefined ? TData : TDataOverride
+    >
+
+    return await buildRunnerFn<"run", TDataOverride, TResultOverride>(
+      "run",
+      container
+    )(inputArgs)
+  }
+
+  exportedWorkflow.registerStepSuccess = async <
+    TDataOverride = undefined,
+    TResultOverride = undefined
+  >(
+    args?: FlowRegisterStepSuccessOptions<
+      TDataOverride extends undefined ? TData : TDataOverride
+    > & {
+      container?: LoadedModule[] | MedusaContainer
+    }
+  ): Promise<
+    WorkflowResult<
+      TResultOverride extends undefined ? TResult : TResultOverride
+    >
+  > => {
+    const container = args?.container
+    delete args?.container
+    const inputArgs = { ...args } as FlowRegisterStepSuccessOptions<
+      TDataOverride extends undefined ? TData : TDataOverride
+    >
+
+    return await buildRunnerFn<
+      "registerStepSuccess",
+      TDataOverride,
+      TResultOverride
+    >(
+      "registerStepSuccess",
+      container
+    )(inputArgs)
+  }
+
+  exportedWorkflow.registerStepFailure = async <
+    TDataOverride = undefined,
+    TResultOverride = undefined
+  >(
+    args?: FlowRegisterStepFailureOptions<
+      TDataOverride extends undefined ? TData : TDataOverride
+    > & {
+      container?: LoadedModule[] | MedusaContainer
+    }
+  ): Promise<
+    WorkflowResult<
+      TResultOverride extends undefined ? TResult : TResultOverride
+    >
+  > => {
+    const container = args?.container
+    delete args?.container
+    const inputArgs = { ...args } as FlowRegisterStepFailureOptions<
+      TDataOverride extends undefined ? TData : TDataOverride
+    >
+
+    return await buildRunnerFn<
+      "registerStepFailure",
+      TDataOverride,
+      TResultOverride
+    >(
+      "registerStepFailure",
+      container
+    )(inputArgs)
+  }
+
   MedusaWorkflow.registerWorkflow(workflowId, exportedWorkflow)
   return exportedWorkflow
 }