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/orchestration/src/__tests__/transaction/transaction-orchestrator.ts

@@ -1,4 +1,5 @@
 import {
+  DistributedTransaction,
   TransactionHandlerType,
   TransactionOrchestrator,
   TransactionPayload,
@@ -858,4 +859,32 @@ describe("Transaction Orchestrator", () => {
     expect(mocks.oneCompensate).toBeCalledTimes(1)
     expect(mocks.twoCompensate).toBeCalledTimes(1)
   })
+
+  it("Should receive the current transaction as reference in the handler", async () => {
+    let transactionInHandler
+
+    async function handler(
+      actionId: string,
+      functionHandlerType: TransactionHandlerType,
+      payload: TransactionPayload,
+      transaction?: DistributedTransaction
+    ) {
+      transactionInHandler = transaction
+    }
+
+    const strategy = new TransactionOrchestrator("transaction-name", {
+      next: {
+        action: "firstMethod",
+      },
+    })
+
+    const transaction = await strategy.beginTransaction(
+      "transaction_id_123",
+      handler
+    )
+
+    await strategy.resume(transaction)
+
+    expect(transaction).toBe(transactionInHandler)
+  })
 })

packages/orchestration/src/transaction/distributed-transaction.ts

@@ -1,5 +1,6 @@
 import { isDefined } from "@medusajs/utils"
 import { TransactionFlow } from "./transaction-orchestrator"
+import { TransactionStepHandler } from "./transaction-step"
 import { TransactionHandlerType, TransactionState } from "./types"
 
 /**
@@ -79,11 +80,7 @@ export class DistributedTransaction {
 
   constructor(
     private flow: TransactionFlow,
-    public handler: (
-      actionId: string,
-      handlerType: TransactionHandlerType,
-      payload: TransactionPayload
-    ) => Promise<unknown>,
+    public handler: TransactionStepHandler,
     public payload?: any,
     errors?: TransactionStepError[],
     context?: TransactionContext

packages/orchestration/src/transaction/transaction-orchestrator.ts

@@ -3,14 +3,13 @@ import {
   TransactionCheckpoint,
   TransactionPayload,
 } from "./distributed-transaction"
+import { TransactionStep, TransactionStepHandler } from "./transaction-step"
 import {
   TransactionHandlerType,
-  TransactionModel,
   TransactionState,
   TransactionStepStatus,
   TransactionStepsDefinition,
 } from "./types"
-import { TransactionStep, TransactionStepHandler } from "./transaction-step"
 
 import { EventEmitter } from "events"
 
@@ -366,7 +365,7 @@ export class TransactionOrchestrator extends EventEmitter {
       if (!step.definition.async) {
         execution.push(
           transaction
-            .handler(step.definition.action + "", type, payload)
+            .handler(step.definition.action + "", type, payload, transaction)
             .then(async (response) => {
               await TransactionOrchestrator.setStepSuccess(
                 transaction,
@@ -387,7 +386,7 @@ export class TransactionOrchestrator extends EventEmitter {
         execution.push(
           transaction.saveCheckpoint().then(async () =>
             transaction
-              .handler(step.definition.action + "", type, payload)
+              .handler(step.definition.action + "", type, payload, transaction)
               .catch(async (error) => {
                 await TransactionOrchestrator.setStepFailure(
                   transaction,

packages/orchestration/src/transaction/transaction-step.ts

@@ -1,15 +1,19 @@
-import { TransactionPayload } from "./distributed-transaction"
 import {
-  TransactionStepsDefinition,
-  TransactionStepStatus,
-  TransactionState,
+  DistributedTransaction,
+  TransactionPayload,
+} from "./distributed-transaction"
+import {
   TransactionHandlerType,
+  TransactionState,
+  TransactionStepStatus,
+  TransactionStepsDefinition,
 } from "./types"
 
 export type TransactionStepHandler = (
   actionId: string,
   handlerType: TransactionHandlerType,
-  payload: TransactionPayload
+  payload: TransactionPayload,
+  transaction?: DistributedTransaction
 ) => Promise<unknown>
 
 /**

packages/orchestration/src/workflow/workflow-manager.ts

@@ -1,5 +1,6 @@
 import { Context, MedusaContainer } from "@medusajs/types"
 import {
+  DistributedTransaction,
   OrchestratorBuilder,
   TransactionHandlerType,
   TransactionMetadata,
@@ -35,6 +36,7 @@ export type WorkflowStepHandler = (args: {
   invoke: { [actions: string]: unknown }
   compensate: { [actions: string]: unknown }
   metadata: TransactionMetadata
+  transaction: DistributedTransaction
   context?: Context
 }) => unknown
 
@@ -136,7 +138,8 @@ export class WorkflowManager {
       return async (
         actionId: string,
         handlerType: TransactionHandlerType,
-        payload?: any
+        payload?: any,
+        transaction?: DistributedTransaction
       ) => {
         const command = handlers.get(actionId)
 
@@ -157,6 +160,7 @@ export class WorkflowManager {
           invoke,
           compensate,
           metadata,
+          transaction: transaction as DistributedTransaction,
           context,
         })
       }