Skip sfn->sfn context injection when CONTEXT.$ field exists (case 3) (#…

…538)
DataDog · Oct 4, 2024 · 07b5dd7 · 07b5dd7
1 parent 5f1ec79
commit 07b5dd7
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 1 deletion.
diff --git a/src/step-functions-helper.spec.ts b/src/step-functions-helper.spec.ts
@@ -361,7 +361,19 @@ describe("test updateDefinitionForStepFunctionInvocationStep", () => {
     expect(updateDefinitionForStepFunctionInvocationStep(stepName, step, serverless, stateMachineName)).toBeFalsy();
   });
 
-  it("Input field has CONTEXT.$ already", async () => {
+  it("Case 3.1: Context injection already set up using States.JsonMerge($$, $, false)", async () => {
+    const parameters = { FunctionName: "bla", Input: { "CONTEXT.$": "States.JsonMerge($$, $, false)" } };
+    const step = { Parameters: parameters };
+    expect(updateDefinitionForStepFunctionInvocationStep(stepName, step, serverless, stateMachineName)).toBeFalsy();
+  });
+
+  it("Case 3.1: Context injection already set up using $$['Execution', 'State', 'StateMachine']", async () => {
+    const parameters = { FunctionName: "bla", Input: { "CONTEXT.$": "$$['Execution', 'State', 'StateMachine']" } };
+    const step = { Parameters: parameters };
+    expect(updateDefinitionForStepFunctionInvocationStep(stepName, step, serverless, stateMachineName)).toBeFalsy();
+  });
+
+  it("Case 3.2: Input field has a custom CONTEXT.$ field", async () => {
     const parameters = { FunctionName: "bla", Input: { "CONTEXT.$": "something else" } };
     const step = { Parameters: parameters };
     expect(updateDefinitionForStepFunctionInvocationStep(stepName, step, serverless, stateMachineName)).toBeFalsy();

diff --git a/src/step-functions-helper.ts b/src/step-functions-helper.ts
@@ -278,6 +278,24 @@ merge these traces, check out https://docs.datadoghq.com/serverless/step_functio
     return false;
   }
 
+  // Case 3.1 context injection is already set up
+  if (
+    parameters.Input["CONTEXT.$"] === "States.JsonMerge($$, $, false)" ||
+    parameters.Input["CONTEXT.$"] === `$$['Execution', 'State', 'StateMachine']`
+  ) {
+    serverless.cli.log(
+      `Step ${stepName} of state machine ${stateMachineName}: Context injection is already set up. Skipping context injection.\n`,
+    );
+
+    return false;
+  }
+
+  // Case 3.2 custom CONTEXT.$ field
+  serverless.cli
+    .log(`[Warn] Step ${stepName} of state machine ${stateMachineName}: Parameters.Input field has a custom CONTEXT.$ field. Step \
+Functions Context Object injection skipped. Your Step Functions trace will not be merged with downstream Step Function traces. To \
+manually merge these traces, check out https://docs.datadoghq.com/serverless/step_functions/troubleshooting/\n`);
+
   return false;
 }