Revamp SFN->SFN context injection

DataDog · Sep 26, 2024 · 0d7867d · 0d7867d
1 parent d36a07e
commit 0d7867d
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 15 deletions.
diff --git a/src/commands/stepfunctions/__tests__/helpers.test.ts b/src/commands/stepfunctions/__tests__/helpers.test.ts
@@ -290,7 +290,7 @@ describe('stepfunctions command helpers tests', () => {
 
  const changed = injectContextForStepFunctions(step, context, 'Step Functions StartExecution')
  expect(changed).toBeTruthy()
- expect(step.Parameters?.Input).toEqual({'CONTEXT.$': 'States.JsonMerge($$, $, false)'})
+ expect(step.Parameters?.Input).toEqual({'CONTEXT.$': `$$['Execution', 'State', 'StateMachine']`})
  })
 
  test('Case 1: is true when "CONTEXT.$" and "CONTEXT" fields are not set', () => {
@@ -304,9 +304,10 @@ describe('stepfunctions command helpers tests', () => {
  End: true,
  }
  expect(injectContextForStepFunctions(step, context, 'Step Functions StartExecution')).toBeTruthy()
+ expect(step.Parameters?.Input).toEqual({'CONTEXT.$': `$$['Execution', 'State', 'StateMachine']`})
  })
 
- test('is true for undefined', () => {
+ test('is true when Input field is not set', () => {
  const step: StepType = {
  Type: 'Task',
  Resource: 'arn:aws:states:::states:startExecution.sync:2',
@@ -316,9 +317,10 @@ describe('stepfunctions command helpers tests', () => {
  End: true,
  }
  expect(injectContextForStepFunctions(step, context, 'Step Functions StartExecution')).toBeTruthy()
+ expect(step.Parameters?.Input).toEqual({'CONTEXT.$': `$$['Execution', 'State', 'StateMachine']`})
  })
 
- test('is false when Input is an object that contains a CONTEXT key using JSONPath expression', () => {
+ test('Case 3.2: is false when Input contains a custom "CONTEXT.$" field', () => {
  const step: StepType = {
  Type: 'Task',
  Resource: 'arn:aws:states:::states:startExecution.sync:2',
@@ -331,7 +333,7 @@ describe('stepfunctions command helpers tests', () => {
  expect(injectContextForStepFunctions(step, context, 'Step Functions StartExecution')).toBeFalsy()
  })
 
- test('Case 3: is false when Input is an object that contains a CONTEXT key that is not a JSON object', () => {
+ test('Case 2: is false when Input contains "CONTEXT" field', () => {
  const step: StepType = {
  Type: 'Task',
  Resource: 'arn:aws:states:::states:startExecution.sync:2',
@@ -343,6 +345,36 @@ describe('stepfunctions command helpers tests', () => {
  }
  expect(injectContextForStepFunctions(step, context, 'Step Functions StartExecution')).toBeFalsy()
  })
+
+ test('Case 3.1: is false when context injection is already set up using States.JsonMerge($$, $, false)', () => {
+ const step: StepType = {
+ Type: 'Task',
+ Resource: 'arn:aws:states:::states:startExecution.sync:2',
+ Parameters: {
+ StateMachineArn: 'arn:aws:states:us-east-1:425362996713:stateMachine:agocs_inner_state_machine',
+ Input: {
+ 'CONTEXT.$': 'States.JsonMerge($$, $, false)',
+ },
+ },
+ End: true,
+ }
+ expect(injectContextForStepFunctions(step, context, 'Step Functions StartExecution')).toBeFalsy()
+ })
+
+ test(`Case 3.1: is false when context injection is already set up using $$['Execution', 'State', 'StateMachine']`, () => {
+ const step: StepType = {
+ Type: 'Task',
+ Resource: 'arn:aws:states:::states:startExecution.sync:2',
+ Parameters: {
+ StateMachineArn: 'arn:aws:states:us-east-1:425362996713:stateMachine:agocs_inner_state_machine',
+ Input: {
+ 'CONTEXT.$': `$$['Execution', 'State', 'StateMachine']`,
+ },
+ },
+ End: true,
+ }
+ expect(injectContextForStepFunctions(step, context, 'Step Functions StartExecution')).toBeFalsy()
+ })
  })
 
  test('buildLogAccessPolicyName test', () => {

diff --git a/src/commands/stepfunctions/helpers.ts b/src/commands/stepfunctions/helpers.ts
@@ -275,7 +275,7 @@ traces, check out https://docs.datadoghq.com/serverless/step_functions/troublesh
  }
 
  if (!step.Parameters.Input) {
- step.Parameters.Input = {'CONTEXT.$': 'States.JsonMerge($$, $, false)'}
+ step.Parameters.Input = {'CONTEXT.$': `$$['Execution', 'State', 'StateMachine']`}
 
  return true
  }
@@ -291,32 +291,34 @@ merge these traces, check out https://docs.datadoghq.com/serverless/step_functio
 
  // 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)'
+ step.Parameters.Input['CONTEXT.$'] = `$$['Execution', 'State', 'StateMachine']`
 
  return true
  }
 
+ // Case 2: Has 'CONTEXT' field.
+ // This case should be rare so we don't support context injection for now
  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 \
+ context.stdout.write(`[Warn] Step ${stepName}'s has 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
- }
+ return false
  }
 
- // context injection is already set up
- if (step.Parameters.Input['CONTEXT.$'] === 'States.JsonMerge($$, $, false)') {
+ // Case 3.1 context injection is already set up
+ if (
+ step.Parameters.Input['CONTEXT.$'] === 'States.JsonMerge($$, $, false)' ||
+ step.Parameters.Input['CONTEXT.$'] === `$$['Execution', 'State', 'StateMachine']`
+ ) {
  context.stdout.write(` Step ${stepName}: Context injection is already set up. Skipping context injection.\n`)
 
  return false
  }
 
+ // Case 3.2 custom CONTEXT.$ field
  context.stdout
- .write(`[Warn] Step ${stepName}'s Parameters.Input field has a custom CONTEXT field. Step Functions Context \
+ .write(`[Warn] Step ${stepName}'s 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`)