diff --git a/src/step-functions-helper.spec.ts b/src/step-functions-helper.spec.ts index e911090..7efc1b3 100644 --- a/src/step-functions-helper.spec.ts +++ b/src/step-functions-helper.spec.ts @@ -116,7 +116,7 @@ describe("test updateDefinitionString", () => { ); }); - it("test lambda step when Payload is not an object", async () => { + it("Case 3: test lambda step when Payload is not an object", async () => { const definitionString = { "Fn::Sub": [ '{"Comment":"fake comment","StartAt":"InvokeLambda","States":{"InvokeLambda":{"Type":"Task","Parameters":{"FunctionName":"fake-function-name","Payload":"Just a string!"},"Resource":"arn:aws:states:::lambda:invoke","End":true}}}', @@ -129,7 +129,39 @@ describe("test updateDefinitionString", () => { expect(definitionAfterUpdate.States?.InvokeLambda?.Parameters?.["Payload"]).toBe("Just a string!"); }); - it("test lambda step with custom Payload", async () => { + it("Case 2.1: test lambda step when Execution, State and StateMachine are already injected into Payload", async () => { + const definitionString = { + "Fn::Sub": [ + '{"Comment":"fake comment","StartAt":"InvokeLambda","States":{"InvokeLambda":{"Type":"Task","Parameters":{"FunctionName":"fake-function-name","Payload":{"Execution.$":"$$.Execution","State.$":"$$.State","StateMachine.$":"$$.StateMachine"}},"Resource":"arn:aws:states:::lambda:invoke","End":true}}}', + {}, + ], + }; + updateDefinitionString(definitionString, serverless, stateMachineName); + + const definitionAfterUpdate: StateMachineDefinition = JSON.parse(definitionString["Fn::Sub"][0] as string); + expect(definitionAfterUpdate.States?.InvokeLambda?.Parameters?.["Payload"]).toStrictEqual({ + "Execution.$": "$$.Execution", + "State.$": "$$.State", + "StateMachine.$": "$$.StateMachine", + }); + }); + + it("Case 2.2: test lambda step when some of Execution, State or StateMachine field but conject injection is not set up completely", async () => { + const definitionString = { + "Fn::Sub": [ + '{"Comment":"fake comment","StartAt":"InvokeLambda","States":{"InvokeLambda":{"Type":"Task","Parameters":{"FunctionName":"fake-function-name","Payload":{"Execution":"$$.Execution"}},"Resource":"arn:aws:states:::lambda:invoke","End":true}}}', + {}, + ], + }; + updateDefinitionString(definitionString, serverless, stateMachineName); + + const definitionAfterUpdate: StateMachineDefinition = JSON.parse(definitionString["Fn::Sub"][0] as string); + expect(definitionAfterUpdate.States?.InvokeLambda?.Parameters?.["Payload"]).toStrictEqual({ + Execution: "$$.Execution", + }); + }); + + it("Case 2.3: test lambda step when none of Execution, State, or StateMachine is in Payload", async () => { const definitionString = { "Fn::Sub": [ '{"Comment":"fake comment","StartAt":"InvokeLambda","States":{"InvokeLambda":{"Type":"Task","Parameters":{"FunctionName":"fake-function-name","Payload":{"CustomerId":42}},"Resource":"arn:aws:states:::lambda:invoke","End":true}}}', @@ -139,7 +171,12 @@ describe("test updateDefinitionString", () => { updateDefinitionString(definitionString, serverless, stateMachineName); const definitionAfterUpdate: StateMachineDefinition = JSON.parse(definitionString["Fn::Sub"][0] as string); - expect(definitionAfterUpdate.States?.InvokeLambda?.Parameters?.["Payload"]).toStrictEqual({ CustomerId: 42 }); + expect(definitionAfterUpdate.States?.InvokeLambda?.Parameters?.["Payload"]).toStrictEqual({ + CustomerId: 42, + "Execution.$": "$$.Execution", + "State.$": "$$.State", + "StateMachine.$": "$$.StateMachine", + }); }); it("test lambda step already has customized payload set do nothing", async () => { diff --git a/src/step-functions-helper.ts b/src/step-functions-helper.ts index fe8294f..918040b 100644 --- a/src/step-functions-helper.ts +++ b/src/step-functions-helper.ts @@ -42,12 +42,21 @@ export interface StateMachineDefinition { States: { [key: string]: StateMachineStep }; } +export type PayloadObject = { + "Execution.$"?: any; + Execution?: any; + "State.$"?: any; + State?: any; + "StateMachine.$"?: any; + StateMachine?: any; +}; + export interface StateMachineStep { Resource?: string; Parameters?: { FunctionName?: string; "Payload.$"?: string; - Payload?: string; + Payload?: string | PayloadObject; Input?: { "CONTEXT.$"?: string; }; @@ -154,6 +163,7 @@ https://docs.datadoghq.com/serverless/step_functions/troubleshooting/`, if (step.Parameters.hasOwnProperty("Payload")) { const payload = step.Parameters.Payload; if (typeof payload !== "object") { + // Case 3: payload is not a JSON object serverless.cli.log( `[Warn] Payload field is not a JSON object. Merging traces failed for step: ${stepName} of state machine: ${stateMachineName}. \ Your Step Functions trace will not be merged with downstream Lambda traces. To manually merge these traces, check out \ @@ -161,6 +171,42 @@ https://docs.datadoghq.com/serverless/step_functions/troubleshooting/`, ); return; } + + // Case 2: payload is a JSON object + if ( + payload["Execution.$"] === "$$.Execution" && + payload["State.$"] === "$$.State" && + payload["StateMachine.$"] === "$$.StateMachine" + ) { + // Case 2.1: already injected into "Payload" + serverless.cli.log( + `Context injection is already set up. Skipping merging traces for step: ${stepName} of state machine: \${stateMachineName}.\n`, + ); + + return; + } else if ( + payload.hasOwnProperty("Execution.$") || + payload.hasOwnProperty("Execution") || + payload.hasOwnProperty("State.$") || + payload.hasOwnProperty("State") || + payload.hasOwnProperty("StateMachine.$") || + payload.hasOwnProperty("StateMachine") + ) { + // Case 2.2: "Payload" object has Execution, State or StateMachine field but conject injection is not set up completely + serverless.cli + .log(`[Warn] Step ${stepName} of state machine: ${stateMachineName} may be using custom Execution, State or StateMachine field. \ +Step Functions Context Object injection skipped. Your Step Functions trace will not be merged with downstream Lambda traces. To manually \ +merge these traces, check out https://docs.datadoghq.com/serverless/step_functions/troubleshooting/\n`); + + return; + } else { + // Case 2.3: "Payload" object has no Execution, State or StateMachine field + payload["Execution.$"] = "$$.Execution"; + payload["State.$"] = "$$.State"; + payload["StateMachine.$"] = "$$.StateMachine"; + + return; + } } if (step.Parameters["Payload.$"] === "$") {