feat: Inject context for sfn->Lambda when Payload is an object (case …

…2) (#530)
DataDog · Sep 25, 2024 · bcd0646 · bcd0646
1 parent 11be473
commit bcd0646
Show file tree

Hide file tree

Showing 2 changed files with 87 additions and 4 deletions.
diff --git 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
@@ -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,13 +163,50 @@ 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 \
   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.$"] === "$") {