Skip to content

Update Step Function Context Parsing #325

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

Merged
merged 6 commits into from
Dec 12, 2022
Merged

Update Step Function Context Parsing #325

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
91b6d9b
Update snapshots
DarcyRaynerDD Jan 11, 2021
1e6a582
Update Step Function Context Parsing
kimi-p Dec 9, 2022
4f31750
update context.input format
kimi-p Dec 9, 2022
86f72ff
lint
kimi-p Dec 9, 2022
6cd1236
merge
kimi-p Dec 12, 2022
34abf63
Revert "merge"
kimi-p Dec 12, 2022
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
2 changes: 1 addition & 1 deletion src/index.spec.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -331,7 +331,7 @@ describe("datadog", () => {

expect(mockedIncrementInvocations).toBeCalledTimes(1);
expect(mockedIncrementInvocations).toBeCalledWith(expect.anything(), mockContext);
expect(logger.debug).toHaveBeenCalledTimes(9);
expect(logger.debug).toHaveBeenCalledTimes(11);
expect(logger.debug).toHaveBeenLastCalledWith('{"status":"debug","message":"datadog:Unpatching HTTP libraries"}');
});
});
102 changes: 59 additions & 43 deletions src/trace/context.spec.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -668,30 +668,41 @@ describe("readTraceFromLambdaContext", () => {

describe("readStepFunctionContextFromEvent", () => {
const stepFunctionEvent = {
dd: {
Execution: {
Name: "fb7b1e15-e4a2-4cb2-963f-8f1fa4aec492",
StartTime: "2019-09-30T20:28:24.236Z",
},
State: {
Name: "step-one",
RetryCount: 2,
},
StateMachine: {
Id: "arn:aws:states:us-east-1:601427279990:stateMachine:HelloStepOneStepFunctionsStateMachine-z4T0mJveJ7pJ",
Name: "my-state-machine",
Execution: {
Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf",
Input: {
MyInput: "MyValue",
},
Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be",
RoleArn: "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03",
StartTime: "2022-12-08T21:08:17.924Z",
},
State: {
Name: "step-one",
EnteredTime: "2022-12-08T21:08:19.224Z",
RetryCount: 2,
},
StateMachine: {
Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential",
Name: "my-state-machine",
},
} as const;
it("reads a trace from an execution id", () => {

it("reads a step function context from event with Execution.Input", () => {
const result = readStepFunctionContextFromEvent(stepFunctionEvent);
expect(result).toEqual({
"step_function.execution_id": "fb7b1e15-e4a2-4cb2-963f-8f1fa4aec492",
"step_function.retry_count": 2,
"step_function.state_machine_arn":
"arn:aws:states:us-east-1:601427279990:stateMachine:HelloStepOneStepFunctionsStateMachine-z4T0mJveJ7pJ",
"step_function.execution_id":
"arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf",
"step_function.execution_input": { MyInput: "MyValue" },
"step_function.execution_name": "85a9933e-9e11-83dc-6a61-b92367b6c3be",
"step_function.execution_role_arn":
"arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03",
"step_function.execution_start_time": "2022-12-08T21:08:17.924Z",
"step_function.state_entered_time": "2022-12-08T21:08:19.224Z",
"step_function.state_machine_arn": "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential",
"step_function.state_machine_name": "my-state-machine",
"step_function.step_name": "step-one",
"step_function.state_name": "step-one",
"step_function.state_retry_count": 2,
});
});
it("returns undefined when event isn't an object", () => {
Expand All @@ -711,7 +722,7 @@ describe("readStepFunctionContextFromEvent", () => {
it("returns undefined when Execution is missing Name field", () => {
const result = readStepFunctionContextFromEvent({
dd: {
...stepFunctionEvent.dd,
...stepFunctionEvent,
Execution: {},
},
});
Expand All @@ -720,7 +731,7 @@ describe("readStepFunctionContextFromEvent", () => {
it("returns undefined when Name isn't a string", () => {
const result = readStepFunctionContextFromEvent({
dd: {
...stepFunctionEvent.dd,
...stepFunctionEvent,
Execution: {
Name: 12345,
},
Expand All @@ -731,7 +742,7 @@ describe("readStepFunctionContextFromEvent", () => {
it("returns undefined when State isn't defined", () => {
const result = readStepFunctionContextFromEvent({
dd: {
...stepFunctionEvent.dd,
...stepFunctionEvent,
State: undefined,
},
});
Expand All @@ -740,9 +751,9 @@ describe("readStepFunctionContextFromEvent", () => {
it("returns undefined when try retry count isn't a number", () => {
const result = readStepFunctionContextFromEvent({
dd: {
...stepFunctionEvent.dd,
...stepFunctionEvent,
State: {
...stepFunctionEvent.dd.State,
...stepFunctionEvent.State,
RetryCount: "1",
},
},
Expand All @@ -752,9 +763,9 @@ describe("readStepFunctionContextFromEvent", () => {
it("returns undefined when try step name isn't a string", () => {
const result = readStepFunctionContextFromEvent({
dd: {
...stepFunctionEvent.dd,
...stepFunctionEvent,
State: {
...stepFunctionEvent.dd.State,
...stepFunctionEvent.State,
Name: 1,
},
},
Expand All @@ -764,7 +775,7 @@ describe("readStepFunctionContextFromEvent", () => {
it("returns undefined when StateMachine is undefined", () => {
const result = readStepFunctionContextFromEvent({
dd: {
...stepFunctionEvent.dd,
...stepFunctionEvent,
StateMachine: undefined,
},
});
Expand All @@ -773,9 +784,9 @@ describe("readStepFunctionContextFromEvent", () => {
it("returns undefined when StateMachineId isn't a string", () => {
const result = readStepFunctionContextFromEvent({
dd: {
...stepFunctionEvent.dd,
...stepFunctionEvent,
StateMachine: {
...stepFunctionEvent.dd.StateMachine,
...stepFunctionEvent.StateMachine,
Id: 1,
},
},
Expand All @@ -785,9 +796,9 @@ describe("readStepFunctionContextFromEvent", () => {
it("returns undefined when StateMachineName isn't a string", () => {
const result = readStepFunctionContextFromEvent({
dd: {
...stepFunctionEvent.dd,
...stepFunctionEvent,
StateMachine: {
...stepFunctionEvent.dd.StateMachine,
...stepFunctionEvent.StateMachine,
Name: 1,
},
},
Expand Down Expand Up @@ -1062,21 +1073,26 @@ describe("extractTraceContext", () => {

it("adds step function metadata to xray", () => {
const stepFunctionEvent = {
dd: {
Execution: {
Name: "fb7b1e15-e4a2-4cb2-963f-8f1fa4aec492",
StartTime: "2019-09-30T20:28:24.236Z",
},
State: {
Name: "step-one",
RetryCount: 2,
},
StateMachine: {
Id: "arn:aws:states:us-east-1:601427279990:stateMachine:HelloStepOneStepFunctionsStateMachine-z4T0mJveJ7pJ",
Name: "my-state-machine",
Execution: {
Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf",
Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be",
RoleArn: "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03",
StartTime: "2022-12-08T21:08:17.924Z",
Input: {
MyInput: "MyValue",
},
},
State: {
Name: "step-one",
EnteredTime: "2022-12-08T21:08:19.224Z",
RetryCount: 2,
},
StateMachine: {
Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential",
Name: "my-state-machine",
},
} as const;

jest.spyOn(Date, "now").mockImplementation(() => 1487076708000);
process.env[xrayTraceEnvVar] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1";
process.env[awsXrayDaemonAddressEnvVar] = "localhost:127.0.0.1:2000";
Expand All @@ -1089,7 +1105,7 @@ describe("extractTraceContext", () => {
const sentMessage = sentSegment.toString();
expect(sentMessage).toMatchInlineSnapshot(`
"{\\"format\\": \\"json\\", \\"version\\": 1}
{\\"id\\":\\"11111\\",\\"trace_id\\":\\"1-5e272390-8c398be037738dc042009320\\",\\"parent_id\\":\\"94ae789b969f1cc5\\",\\"name\\":\\"datadog-metadata\\",\\"start_time\\":1487076708,\\"end_time\\":1487076708,\\"type\\":\\"subsegment\\",\\"metadata\\":{\\"datadog\\":{\\"root_span_metadata\\":{\\"step_function.execution_id\\":\\"fb7b1e15-e4a2-4cb2-963f-8f1fa4aec492\\",\\"step_function.retry_count\\":2,\\"step_function.state_machine_arn\\":\\"arn:aws:states:us-east-1:601427279990:stateMachine:HelloStepOneStepFunctionsStateMachine-z4T0mJveJ7pJ\\",\\"step_function.state_machine_name\\":\\"my-state-machine\\",\\"step_function.step_name\\":\\"step-one\\"}}}}"
{\\"id\\":\\"11111\\",\\"trace_id\\":\\"1-5e272390-8c398be037738dc042009320\\",\\"parent_id\\":\\"94ae789b969f1cc5\\",\\"name\\":\\"datadog-metadata\\",\\"start_time\\":1487076708,\\"end_time\\":1487076708,\\"type\\":\\"subsegment\\",\\"metadata\\":{\\"datadog\\":{\\"root_span_metadata\\":{\\"step_function.execution_name\\":\\"85a9933e-9e11-83dc-6a61-b92367b6c3be\\",\\"step_function.execution_id\\":\\"arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf\\",\\"step_function.execution_input\\":{\\"MyInput\\":\\"MyValue\\"},\\"step_function.execution_role_arn\\":\\"arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03\\",\\"step_function.execution_start_time\\":\\"2022-12-08T21:08:17.924Z\\",\\"step_function.state_entered_time\\":\\"2022-12-08T21:08:19.224Z\\",\\"step_function.state_machine_arn\\":\\"arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential\\",\\"step_function.state_machine_name\\":\\"my-state-machine\\",\\"step_function.state_name\\":\\"step-one\\",\\"step_function.state_retry_count\\":2}}}}"
`);
});
});
76 changes: 58 additions & 18 deletions src/trace/context.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ import {
isSQSEvent,
} from "../utils/event-type-guards";
import {
authorizingRequestIdHeader,
awsXrayDaemonAddressEnvVar,
parentIDHeader,
SampleMode,
Expand All @@ -26,8 +27,7 @@ import {
xrayTraceEnvVar,
} from "./constants";
import { TraceExtractor } from "./listener";
import { parseEventSourceSubType, eventSubTypes } from "./trigger";
import { authorizingRequestIdHeader } from "./constants";
import { eventSubTypes, parseEventSourceSubType } from "./trigger";

export interface XRayTraceHeader {
traceID: string;
Expand All @@ -43,11 +43,16 @@ export interface TraceContext {
}

export interface StepFunctionContext {
"step_function.retry_count": number;
"step_function.execution_name": string;
"step_function.execution_id": string;
"step_function.execution_input": object;
"step_function.execution_role_arn": string;
"step_function.execution_start_time": string;
"step_function.state_machine_name": string;
"step_function.state_machine_arn": string;
"step_function.step_name": string;
"step_function.state_entered_time": string;
"step_function.state_name": string;
"step_function.state_retry_count": number;
}

/**
Expand Down Expand Up @@ -485,50 +490,85 @@ export function readStepFunctionContextFromEvent(event: any): StepFunctionContex
if (typeof event !== "object") {
return;
}
const { dd } = event;
if (typeof dd !== "object") {
return;
}
const execution = dd.Execution;

const execution = event.Execution;
if (typeof execution !== "object") {
logDebug("event.Execution is not an object.");
return;
}
const executionID = execution.Name;
const executionID = execution.Id;
if (typeof executionID !== "string") {
logDebug("event.Execution.Id is not a string.");
return;
}
const executionInput = execution.Input;
const executionName = execution.Name;
if (typeof executionName !== "string") {
logDebug("event.Execution.Name is not a string.");
return;
}
const executionRoleArn = execution.RoleArn;
if (typeof executionRoleArn !== "string") {
logDebug("event.Execution.RoleArn is not a string.");
return;
}
const state = dd.State;
const executionStartTime = execution.StartTime;
if (typeof executionStartTime !== "string") {
logDebug("event.Execution.StartTime is not a string.");
return;
}

const state = event.State;
if (typeof state !== "object") {
logDebug("event.State is not an object.");
return;
}
const stateRetryCount = state.RetryCount;
if (typeof stateRetryCount !== "number") {
logDebug("event.State.RetryCount is not a string.");
return;
}
const retryCount = state.RetryCount;
if (typeof retryCount !== "number") {
const stateEnteredTime = state.EnteredTime;
if (typeof stateEnteredTime !== "string") {
logDebug("event.State.EnteredTime is not a string.");
return;
}
const stepName = state.Name;
if (typeof stepName !== "string") {
const stateName = state.Name;
if (typeof stateName !== "string") {
logDebug("event.State.Name is not a string.");
return;
}
const stateMachine = dd.StateMachine;

const stateMachine = event.StateMachine;
if (typeof stateMachine !== "object") {
logDebug("event.StateMachine is not an object.");
return;
}
const stateMachineArn = stateMachine.Id;
if (typeof stateMachineArn !== "string") {
logDebug("event.StateMachine.Id is not a string.");
return;
}
const stateMachineName = stateMachine.Name;
if (typeof stateMachineName !== "string") {
logDebug("event.StateMachine.Name is not a string.");
return;
}

return {
"step_function.execution_name": executionName,
"step_function.execution_id": executionID,
"step_function.retry_count": retryCount,
"step_function.execution_input": executionInput ?? {},
"step_function.execution_role_arn": executionRoleArn,
"step_function.execution_start_time": executionStartTime,
"step_function.state_entered_time": stateEnteredTime,
"step_function.state_machine_arn": stateMachineArn,
"step_function.state_machine_name": stateMachineName,
"step_function.step_name": stepName,
"step_function.state_name": stateName,
"step_function.state_retry_count": stateRetryCount,
};
}

export function convertToSampleMode(xraySampled: number): SampleMode {
return xraySampled === 1 ? SampleMode.USER_KEEP : SampleMode.USER_REJECT;
}
Expand Down