Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Print warning if traces merging from sfn to sfn is skipped #1447

Merged
merged 2 commits into from
Sep 16, 2024
Merged

Print warning if traces merging from sfn to sfn is skipped #1447

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4a069a4
Print a warning message if traces merging is skipped for Lambda in St…
lym953 Sep 10, 2024
5ef1845
Print warning if traces merging from Step Function to Step Function i…
lym953 Sep 12, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 27 additions & 9 deletions src/commands/stepfunctions/__tests__/helpers.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import {DescribeStateMachineCommandOutput, LogLevel} from '@aws-sdk/client-sfn'
import {BaseContext} from 'clipanion'

import {
buildArn,
Expand All @@ -16,7 +17,18 @@ import {

import {describeStateMachineFixture} from './fixtures/aws-resources'

const createMockContext = (): BaseContext => {
return {
stdout: {
write: (input: string) => {
return true
},
},
} as BaseContext
}

describe('stepfunctions command helpers tests', () => {
const context = createMockContext()
describe('shouldUpdateStepForTracesMerging test', () => {
test('already has JsonMerge added to payload field', () => {
const step: StepType = {
Expand All @@ -28,7 +40,7 @@ describe('stepfunctions command helpers tests', () => {
},
End: true,
}
expect(shouldUpdateStepForTracesMerging(step)).toBeFalsy()
expect(shouldUpdateStepForTracesMerging(step, context, 'Lambda Invoke')).toBeFalsy()
})

test('no payload field', () => {
Expand All @@ -40,7 +52,7 @@ describe('stepfunctions command helpers tests', () => {
},
End: true,
}
expect(shouldUpdateStepForTracesMerging(step)).toBeTruthy()
expect(shouldUpdateStepForTracesMerging(step, context, 'Lambda Invoke')).toBeTruthy()
})

test('default payload field of $', () => {
Expand All @@ -53,7 +65,7 @@ describe('stepfunctions command helpers tests', () => {
},
End: true,
}
expect(shouldUpdateStepForTracesMerging(step)).toBeTruthy()
expect(shouldUpdateStepForTracesMerging(step, context, 'Lambda Invoke')).toBeTruthy()
})

test('none-lambda step should not be updated', () => {
Expand All @@ -65,7 +77,7 @@ describe('stepfunctions command helpers tests', () => {
},
End: true,
}
expect(shouldUpdateStepForTracesMerging(step)).toBeFalsy()
expect(shouldUpdateStepForTracesMerging(step, context, 'DynamoDB Update')).toBeFalsy()
})

test('legacy lambda api should not be updated', () => {
Expand All @@ -74,7 +86,7 @@ describe('stepfunctions command helpers tests', () => {
Resource: 'arn:aws:lambda:sa-east-1:601427271234:function:hello-function',
End: true,
}
expect(shouldUpdateStepForTracesMerging(step)).toBeFalsy()
expect(shouldUpdateStepForTracesMerging(step, context, 'Legacy Lambda')).toBeFalsy()
})
})

Expand Down Expand Up @@ -185,7 +197,9 @@ describe('stepfunctions command helpers tests', () => {
},
End: true,
}
expect(shouldUpdateStepForStepFunctionContextInjection(step)).toBeTruthy()
expect(
shouldUpdateStepForStepFunctionContextInjection(step, context, 'Step Functions StartExecution')
).toBeTruthy()
})

test('is true for undefined', () => {
Expand All @@ -197,7 +211,9 @@ describe('stepfunctions command helpers tests', () => {
},
End: true,
}
expect(shouldUpdateStepForStepFunctionContextInjection(step)).toBeTruthy()
expect(
shouldUpdateStepForStepFunctionContextInjection(step, context, 'Step Functions StartExecution')
).toBeTruthy()
})

test('is false when Input is an object that contains a CONTEXT key', () => {
Expand All @@ -210,7 +226,9 @@ describe('stepfunctions command helpers tests', () => {
},
End: true,
}
expect(shouldUpdateStepForStepFunctionContextInjection(step)).toBeFalsy()
expect(
shouldUpdateStepForStepFunctionContextInjection(step, context, 'Step Functions StartExecution')
).toBeFalsy()
})
})

Expand All @@ -226,7 +244,7 @@ describe('stepfunctions command helpers tests', () => {
End: true,
}

const changed = injectContextForStepFunctions(step)
const changed = injectContextForStepFunctions(step, context, 'Step Functions StartExecution')
expect(changed).toBeTruthy()
expect(step.Parameters?.Input).toEqual({'CONTEXT.$': 'States.JsonMerge($$, $, false)'})
})
Expand Down
51 changes: 45 additions & 6 deletions src/commands/stepfunctions/helpers.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -95,7 +95,7 @@ export const injectContextIntoTasks = async (
if (definitionObj.States.hasOwnProperty(stepName)) {
const step = definitionObj.States[stepName]
const lambdaUpdated = injectContextForLambdaFunctions(step, context, stepName)
const stepUpdated = injectContextForStepFunctions(step)
const stepUpdated = injectContextForStepFunctions(step, context, stepName)
definitionHasBeenUpdated = lambdaUpdated || stepUpdated
}
}
Expand Down Expand Up @@ -125,20 +125,35 @@ export const addTraceContextToStepFunctionParameters = ({Parameters}: StepType):
}
}

export const shouldUpdateStepForTracesMerging = (step: StepType): boolean => {
export const shouldUpdateStepForTracesMerging = (step: StepType, context: BaseContext, stepName: string): boolean => {
// is default lambda api
if (step.Resource === 'arn:aws:states:::lambda:invoke') {
if (!step.Parameters) {
context.stdout
.write(`[Warn] Step ${stepName} does not have a Parameters 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
}

// payload field not set
if (!step.Parameters.hasOwnProperty('Payload.$')) {
return true
}

// default payload
if (step.Parameters['Payload.$'] === '$') {
return true
}

// 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, \
check out https://docs.datadoghq.com/serverless/step_functions/troubleshooting/\n`)

return false
}

return false
Expand All @@ -152,21 +167,45 @@ export const shouldUpdateStepForTracesMerging = (step: StepType): boolean => {
// not object | false
// object without CONTEXT.$ | true
// object with CONTEXT.$ | false
export const shouldUpdateStepForStepFunctionContextInjection = (step: StepType): boolean => {
export const shouldUpdateStepForStepFunctionContextInjection = (
step: StepType,
context: BaseContext,
stepName: string
): boolean => {
// is default lambda api
if (step.Resource?.startsWith('arn:aws:states:::states:startExecution')) {
if (!step.Parameters) {
context.stdout
.write(`[Warn] Step ${stepName} does not have a Parameters 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
}

if (!step.Parameters.Input) {
return true
}

if (typeof step.Parameters.Input !== 'object') {
context.stdout
.write(`[Warn] Step ${stepName}'s Parameters.Input 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
}

if (!step.Parameters.Input['CONTEXT.$']) {
return true
}

context.stdout
.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`)

return false
}

return false
Expand Down Expand Up @@ -198,7 +237,7 @@ export type ParametersType = {
}

const injectContextForLambdaFunctions = (step: StepType, context: BaseContext, stepName: string): boolean => {
if (shouldUpdateStepForTracesMerging(step)) {
if (shouldUpdateStepForTracesMerging(step, context, stepName)) {
addTraceContextToLambdaParameters(step)

return true
Expand All @@ -213,8 +252,8 @@ const injectContextForLambdaFunctions = (step: StepType, context: BaseContext, s
return false
}

export const injectContextForStepFunctions = (step: StepType): boolean => {
if (shouldUpdateStepForStepFunctionContextInjection(step)) {
export const injectContextForStepFunctions = (step: StepType, context: BaseContext, stepName: string): boolean => {
if (shouldUpdateStepForStepFunctionContextInjection(step, context, stepName)) {
addTraceContextToStepFunctionParameters(step)

return true
Expand Down