Skip xfn->Lambda context propagation if custom Payload is used

DataDog · Sep 12, 2024 · a9f365f · a9f365f
1 parent 5ef1845
commit a9f365f
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 1 deletion.
diff --git a/src/commands/stepfunctions/__tests__/helpers.test.ts b/src/commands/stepfunctions/__tests__/helpers.test.ts
@@ -68,6 +68,32 @@ describe('stepfunctions command helpers tests', () => {
       expect(shouldUpdateStepForTracesMerging(step, context, 'Lambda Invoke')).toBeTruthy()
     })
 
+    test('custom payload field not using JsonPath expression', () => {
+      const step: StepType = {
+        Type: 'Task',
+        Resource: 'arn:aws:states:::lambda:invoke',
+        Parameters: {
+          FunctionName: 'arn:aws:lambda:sa-east-1:425362991234:function:unit-test-lambda-function',
+          Payload: '{"action": "service/delete_customer"}',
+        },
+        End: true,
+      }
+      expect(shouldUpdateStepForTracesMerging(step, context, 'Lambda Invoke')).toBeFalsy()
+    })
+
+    test('custom payload field using JsonPath expression', () => {
+      const step: StepType = {
+        Type: 'Task',
+        Resource: 'arn:aws:states:::lambda:invoke',
+        Parameters: {
+          FunctionName: 'arn:aws:lambda:sa-east-1:425362991234:function:unit-test-lambda-function',
+          'Payload.$': '{"customer.$": "$.customer"}',
+        },
+        End: true,
+      }
+      expect(shouldUpdateStepForTracesMerging(step, context, 'Lambda Invoke')).toBeFalsy()
+    })
+
     test('none-lambda step should not be updated', () => {
       const step: StepType = {
         Type: 'Task',

diff --git a/src/commands/stepfunctions/helpers.ts b/src/commands/stepfunctions/helpers.ts
@@ -138,7 +138,7 @@ check out https://docs.datadoghq.com/serverless/step_functions/troubleshooting/\
     }
 
     // payload field not set
-    if (!step.Parameters.hasOwnProperty('Payload.$')) {
+    if (!step.Parameters.hasOwnProperty('Payload.$') && !step.Parameters.hasOwnProperty('Payload')) {
       return true
     }
 
@@ -228,6 +228,7 @@ export type StepType = {
 
 export type ParametersType = {
   'Payload.$'?: string
+  Payload?: string
   FunctionName?: string
   StateMachineArn?: string
   TableName?: string