Skip to content

Commit

Permalink
Update documentation for xfn->Lambda context injection
Browse files Browse the repository at this point in the history
  • Loading branch information
lym953 committed Sep 18, 2024
1 parent e04a645 commit 3d3f2df
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 14 deletions.
18 changes: 9 additions & 9 deletions src/commands/stepfunctions/__tests__/helpers.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ const createMockContext = (): BaseContext => {
describe('stepfunctions command helpers tests', () => {
const context = createMockContext()
describe('injectContextForLambdaFunctions test', () => {
test('already has JsonMerge added to payload field', () => {
test('Case 4.2: already has JsonMerge added to payload field', () => {
const step: StepType = {
Type: 'Task',
Resource: 'arn:aws:states:::lambda:invoke',
Expand All @@ -44,7 +44,7 @@ describe('stepfunctions command helpers tests', () => {
expect(injectContextForLambdaFunctions(step, context, 'Lambda Invoke')).toBeFalsy()
})

test('no payload field', () => {
test('Case 1: no Payload or Payload.$', () => {
const step: StepType = {
Type: 'Task',
Resource: 'arn:aws:states:::lambda:invoke',
Expand All @@ -57,7 +57,7 @@ describe('stepfunctions command helpers tests', () => {
expect(step.Parameters?.['Payload.$']).toEqual(`$$['Execution', 'State', 'StateMachine']`)
})

test('payload is not a JSON object', () => {
test('Case 3: Payload is not a JSON object', () => {
const step: StepType = {
Type: 'Task',
Resource: 'arn:aws:states:::lambda:invoke',
Expand All @@ -70,7 +70,7 @@ describe('stepfunctions command helpers tests', () => {
expect(injectContextForLambdaFunctions(step, context, 'Lambda Invoke')).toBeFalsy()
})

test('already injected Execution and State into Payload', () => {
test('Case 2.1: already injected Execution, State and StateMachine into Payload', () => {
const step: StepType = {
Type: 'Task',
Resource: 'arn:aws:states:::lambda:invoke',
Expand All @@ -87,7 +87,7 @@ describe('stepfunctions command helpers tests', () => {
expect(injectContextForLambdaFunctions(step, context, 'Lambda Invoke')).toBeFalsy()
})

test('custom State field in Payload', () => {
test('Case 2.2: custom State field in Payload', () => {
const step: StepType = {
Type: 'Task',
Resource: 'arn:aws:states:::lambda:invoke',
Expand All @@ -102,7 +102,7 @@ describe('stepfunctions command helpers tests', () => {
expect(injectContextForLambdaFunctions(step, context, 'Lambda Invoke')).toBeFalsy()
})

test('no Execution, State, or StateMachine field in Payload', () => {
test('Case 2.3: no Execution, State, or StateMachine field in Payload', () => {
const step: StepType = {
Type: 'Task',
Resource: 'arn:aws:states:::lambda:invoke',
Expand All @@ -119,7 +119,7 @@ describe('stepfunctions command helpers tests', () => {
expect(payload['StateMachine.$']).toEqual('$$.StateMachine')
})

test('default payload field of $', () => {
test('Case 4.1: default payload field of $', () => {
const step: StepType = {
Type: 'Task',
Resource: 'arn:aws:states:::lambda:invoke',
Expand All @@ -133,7 +133,7 @@ describe('stepfunctions command helpers tests', () => {
expect(step.Parameters?.['Payload.$']).toEqual('States.JsonMerge($$, $, false)')
})

test('custom payload field not using JsonPath expression', () => {
test('Case 4.3: custom payload field not using JsonPath expression', () => {
const step: StepType = {
Type: 'Task',
Resource: 'arn:aws:states:::lambda:invoke',
Expand All @@ -146,7 +146,7 @@ describe('stepfunctions command helpers tests', () => {
expect(injectContextForLambdaFunctions(step, context, 'Lambda Invoke')).toBeFalsy()
})

test('custom payload field using JsonPath expression', () => {
test('Case 4.3: custom payload field using JsonPath expression', () => {
const step: StepType = {
Type: 'Task',
Resource: 'arn:aws:states:::lambda:invoke',
Expand Down
25 changes: 20 additions & 5 deletions src/commands/stepfunctions/helpers.ts
Original file line number Diff line number Diff line change
Expand Up @@ -214,6 +214,17 @@ export type ParametersType = {
}
}

// Truth table
// Case | Input | Expected
// -----|----------------------------------------------------------|---------
// 1 | No "Payload" or "Payload.$" | true
// 2.1 | "Payload" is object, already injected | false
// 2.2 | "Payload" object has Execution, State or StateMachine | false
// 2.3 | "Payload" object has no Execution, State or StateMachine | true
// 3 | "Payload" is not object | false
// 4.1 | "Payload.$": "$" (default payload) | true
// 4.2 | "Payload.$": "States.JsonMerge($$, $, false)" | false
// 4.3 | Custom "Payload.$" | false
export const injectContextForLambdaFunctions = (step: StepType, context: BaseContext, stepName: string): boolean => {
if (step.Resource?.startsWith('arn:aws:lambda')) {
context.stdout.write(
Expand All @@ -239,7 +250,7 @@ check out https://docs.datadoghq.com/serverless/step_functions/troubleshooting/\
return false
}

// payload field not set
// Case 1: payload field not set
if (!step.Parameters.hasOwnProperty('Payload.$') && !step.Parameters.hasOwnProperty('Payload')) {
step.Parameters[`Payload.$`] = `$$['Execution', 'State', 'StateMachine']`

Expand All @@ -248,20 +259,22 @@ check out https://docs.datadoghq.com/serverless/step_functions/troubleshooting/\

if (step.Parameters.hasOwnProperty('Payload')) {
if (typeof step.Parameters['Payload'] !== 'object') {
// payload is not a JSON object
// Case 3: payload is not a JSON object
context.stdout
.write(`[Warn] Step ${stepName}'s Payload field is not a JSON object. 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 false
} else {
// Case 2: payload is not a JSON object
const payload = step.Parameters.Payload
if (
payload['Execution.$'] === '$$.Execution' &&
payload['State.$'] === '$$.State' &&
payload['StateMachine.$'] === '$$.StateMachine'
) {
// Case 2.1: already injected into "Payload"
context.stdout.write(`Step ${stepName}: Context injection is already set up. Skipping context injection.\n`)

return false
Expand All @@ -273,13 +286,15 @@ merge these traces, check out https://docs.datadoghq.com/serverless/step_functio
payload.hasOwnProperty('StateMachine.$') ||
payload.hasOwnProperty('StateMachine')
) {
// Case 2.2: "Payload" object has Execution, State or StateMachine
context.stdout
.write(`[Warn] Step ${stepName} 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 false
} else {
// Case 2.3: "Payload" object has no Execution, State or StateMachine
payload['Execution.$'] = '$$.Execution'
payload['State.$'] = '$$.State'
payload['StateMachine.$'] = '$$.StateMachine'
Expand All @@ -289,21 +304,21 @@ merge these traces, check out https://docs.datadoghq.com/serverless/step_functio
}
}

// default payload
// Case 4.1: default payload
if (step.Parameters['Payload.$'] === '$') {
step.Parameters[`Payload.$`] = 'States.JsonMerge($$, $, false)'

return true
}

// context injection is already set up
// Case 4.2: context injection is already set up using "Payload.$"
if (step.Parameters['Payload.$'] === 'States.JsonMerge($$, $, false)') {
context.stdout.write(` Step ${stepName}: Context injection is already set up. Skipping context injection.\n`)

return false
}

// custom payload
// Case 4.3: custom "Payload.$"
context.stdout
.write(`[Warn] Step ${stepName} has a custom Payload field. Step Functions Context Object injection skipped. \
Your Step Functions trace will not be merged with downstream Lambda traces. To manually merge these traces, \
Expand Down

0 comments on commit 3d3f2df

Please sign in to comment.