Support sfn->sfn context injection case 1&3

DataDog · Sep 23, 2024 · 76112fd · 76112fd
1 parent 575e459
commit 76112fd
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 2 deletions.
diff --git a/src/commands/stepfunctions/__tests__/helpers.test.ts b/src/commands/stepfunctions/__tests__/helpers.test.ts
@@ -293,7 +293,7 @@ describe('stepfunctions command helpers tests', () => {
       expect(step.Parameters?.Input).toEqual({'CONTEXT.$': 'States.JsonMerge($$, $, false)'})
     })
 
-    test('is true for an empty object', () => {
+    test('Case 1: is true when "CONTEXT.$" and "CONTEXT" fields are not set', () => {
       const step: StepType = {
         Type: 'Task',
         Resource: 'arn:aws:states:::states:startExecution.sync:2',
@@ -331,7 +331,7 @@ describe('stepfunctions command helpers tests', () => {
       expect(injectContextForStepFunctions(step, context, 'Step Functions StartExecution')).toBeFalsy()
     })
 
-    test('is false when Input is an object that contains a CONTEXT key not using JSONPath expression', () => {
+    test('Case 3: is false when Input is an object that contains a CONTEXT key that is not a JSON object', () => {
       const step: StepType = {
         Type: 'Task',
         Resource: 'arn:aws:states:::states:startExecution.sync:2',

diff --git a/src/commands/stepfunctions/helpers.ts b/src/commands/stepfunctions/helpers.ts
@@ -289,12 +289,25 @@ merge these traces, check out https://docs.datadoghq.com/serverless/step_functio
     return false
   }
 
+  // Case 1: 'CONTEXT.$' and 'CONTEXT' fields are not set
   if (!step.Parameters.Input['CONTEXT.$'] && !step.Parameters.Input['CONTEXT']) {
     step.Parameters.Input['CONTEXT.$'] = 'States.JsonMerge($$, $, false)'
 
     return true
   }
 
+  if (step.Parameters.Input.hasOwnProperty('CONTEXT')) {
+    if (typeof step.Parameters.Input.CONTEXT !== 'object') {
+      // Case 3: 'CONTEXT' field is not a JSON object
+      context.stdout
+        .write(`[Warn] Step ${stepName}'s Parameters.Input.CONTEXT field is not a JSON object. 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
+    }
+  }
+
   // context injection is already set up
   if (step.Parameters.Input['CONTEXT.$'] === 'States.JsonMerge($$, $, false)') {
     context.stdout.write(` Step ${stepName}: Context injection is already set up. Skipping context injection.\n`)