Skip to content

Commit 9e1957e

Browse files
kimi-pkimi-p
authored
Update Step Function Context Parsing (#325)
1 parent a891948 commit 9e1957e

File tree

3 files changed

+118
-62
lines changed

3 files changed

+118
-62
lines changed

‎src/index.spec.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -331,7 +331,7 @@ describe("datadog", () => {
331331

332332
expect(mockedIncrementInvocations).toBeCalledTimes(1);
333333
expect(mockedIncrementInvocations).toBeCalledWith(expect.anything(), mockContext);
334-
expect(logger.debug).toHaveBeenCalledTimes(9);
334+
expect(logger.debug).toHaveBeenCalledTimes(11);
335335
expect(logger.debug).toHaveBeenLastCalledWith('{"status":"debug","message":"datadog:Unpatching HTTP libraries"}');
336336
});
337337
});

‎src/trace/context.spec.ts

Lines changed: 59 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -668,30 +668,41 @@ describe("readTraceFromLambdaContext", () => {
668668

669669
describe("readStepFunctionContextFromEvent", () => {
670670
const stepFunctionEvent = {
671-
dd: {
672-
Execution: {
673-
Name: "fb7b1e15-e4a2-4cb2-963f-8f1fa4aec492",
674-
StartTime: "2019-09-30T20:28:24.236Z",
675-
},
676-
State: {
677-
Name: "step-one",
678-
RetryCount: 2,
679-
},
680-
StateMachine: {
681-
Id: "arn:aws:states:us-east-1:601427279990:stateMachine:HelloStepOneStepFunctionsStateMachine-z4T0mJveJ7pJ",
682-
Name: "my-state-machine",
671+
Execution: {
672+
Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf",
673+
Input: {
674+
MyInput: "MyValue",
683675
},
676+
Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be",
677+
RoleArn: "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03",
678+
StartTime: "2022-12-08T21:08:17.924Z",
679+
},
680+
State: {
681+
Name: "step-one",
682+
EnteredTime: "2022-12-08T21:08:19.224Z",
683+
RetryCount: 2,
684+
},
685+
StateMachine: {
686+
Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential",
687+
Name: "my-state-machine",
684688
},
685689
} as const;
686-
it("reads a trace from an execution id", () => {
690+
691+
it("reads a step function context from event with Execution.Input", () => {
687692
const result = readStepFunctionContextFromEvent(stepFunctionEvent);
688693
expect(result).toEqual({
689-
"step_function.execution_id": "fb7b1e15-e4a2-4cb2-963f-8f1fa4aec492",
690-
"step_function.retry_count": 2,
691-
"step_function.state_machine_arn":
692-
"arn:aws:states:us-east-1:601427279990:stateMachine:HelloStepOneStepFunctionsStateMachine-z4T0mJveJ7pJ",
694+
"step_function.execution_id":
695+
"arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf",
696+
"step_function.execution_input": { MyInput: "MyValue" },
697+
"step_function.execution_name": "85a9933e-9e11-83dc-6a61-b92367b6c3be",
698+
"step_function.execution_role_arn":
699+
"arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03",
700+
"step_function.execution_start_time": "2022-12-08T21:08:17.924Z",
701+
"step_function.state_entered_time": "2022-12-08T21:08:19.224Z",
702+
"step_function.state_machine_arn": "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential",
693703
"step_function.state_machine_name": "my-state-machine",
694-
"step_function.step_name": "step-one",
704+
"step_function.state_name": "step-one",
705+
"step_function.state_retry_count": 2,
695706
});
696707
});
697708
it("returns undefined when event isn't an object", () => {
@@ -711,7 +722,7 @@ describe("readStepFunctionContextFromEvent", () => {
711722
it("returns undefined when Execution is missing Name field", () => {
712723
const result = readStepFunctionContextFromEvent({
713724
dd: {
714-
...stepFunctionEvent.dd,
725+
...stepFunctionEvent,
715726
Execution: {},
716727
},
717728
});
@@ -720,7 +731,7 @@ describe("readStepFunctionContextFromEvent", () => {
720731
it("returns undefined when Name isn't a string", () => {
721732
const result = readStepFunctionContextFromEvent({
722733
dd: {
723-
...stepFunctionEvent.dd,
734+
...stepFunctionEvent,
724735
Execution: {
725736
Name: 12345,
726737
},
@@ -731,7 +742,7 @@ describe("readStepFunctionContextFromEvent", () => {
731742
it("returns undefined when State isn't defined", () => {
732743
const result = readStepFunctionContextFromEvent({
733744
dd: {
734-
...stepFunctionEvent.dd,
745+
...stepFunctionEvent,
735746
State: undefined,
736747
},
737748
});
@@ -740,9 +751,9 @@ describe("readStepFunctionContextFromEvent", () => {
740751
it("returns undefined when try retry count isn't a number", () => {
741752
const result = readStepFunctionContextFromEvent({
742753
dd: {
743-
...stepFunctionEvent.dd,
754+
...stepFunctionEvent,
744755
State: {
745-
...stepFunctionEvent.dd.State,
756+
...stepFunctionEvent.State,
746757
RetryCount: "1",
747758
},
748759
},
@@ -752,9 +763,9 @@ describe("readStepFunctionContextFromEvent", () => {
752763
it("returns undefined when try step name isn't a string", () => {
753764
const result = readStepFunctionContextFromEvent({
754765
dd: {
755-
...stepFunctionEvent.dd,
766+
...stepFunctionEvent,
756767
State: {
757-
...stepFunctionEvent.dd.State,
768+
...stepFunctionEvent.State,
758769
Name: 1,
759770
},
760771
},
@@ -764,7 +775,7 @@ describe("readStepFunctionContextFromEvent", () => {
764775
it("returns undefined when StateMachine is undefined", () => {
765776
const result = readStepFunctionContextFromEvent({
766777
dd: {
767-
...stepFunctionEvent.dd,
778+
...stepFunctionEvent,
768779
StateMachine: undefined,
769780
},
770781
});
@@ -773,9 +784,9 @@ describe("readStepFunctionContextFromEvent", () => {
773784
it("returns undefined when StateMachineId isn't a string", () => {
774785
const result = readStepFunctionContextFromEvent({
775786
dd: {
776-
...stepFunctionEvent.dd,
787+
...stepFunctionEvent,
777788
StateMachine: {
778-
...stepFunctionEvent.dd.StateMachine,
789+
...stepFunctionEvent.StateMachine,
779790
Id: 1,
780791
},
781792
},
@@ -785,9 +796,9 @@ describe("readStepFunctionContextFromEvent", () => {
785796
it("returns undefined when StateMachineName isn't a string", () => {
786797
const result = readStepFunctionContextFromEvent({
787798
dd: {
788-
...stepFunctionEvent.dd,
799+
...stepFunctionEvent,
789800
StateMachine: {
790-
...stepFunctionEvent.dd.StateMachine,
801+
...stepFunctionEvent.StateMachine,
791802
Name: 1,
792803
},
793804
},
@@ -1062,21 +1073,26 @@ describe("extractTraceContext", () => {
10621073

10631074
it("adds step function metadata to xray", () => {
10641075
const stepFunctionEvent = {
1065-
dd: {
1066-
Execution: {
1067-
Name: "fb7b1e15-e4a2-4cb2-963f-8f1fa4aec492",
1068-
StartTime: "2019-09-30T20:28:24.236Z",
1069-
},
1070-
State: {
1071-
Name: "step-one",
1072-
RetryCount: 2,
1073-
},
1074-
StateMachine: {
1075-
Id: "arn:aws:states:us-east-1:601427279990:stateMachine:HelloStepOneStepFunctionsStateMachine-z4T0mJveJ7pJ",
1076-
Name: "my-state-machine",
1076+
Execution: {
1077+
Id: "arn:aws:states:sa-east-1:425362996713:express:logs-to-traces-sequential:85a9933e-9e11-83dc-6a61-b92367b6c3be:3f7ef5c7-c8b8-4c88-90a1-d54aa7e7e2bf",
1078+
Name: "85a9933e-9e11-83dc-6a61-b92367b6c3be",
1079+
RoleArn: "arn:aws:iam::425362996713:role/service-role/StepFunctions-logs-to-traces-sequential-role-ccd69c03",
1080+
StartTime: "2022-12-08T21:08:17.924Z",
1081+
Input: {
1082+
MyInput: "MyValue",
10771083
},
10781084
},
1085+
State: {
1086+
Name: "step-one",
1087+
EnteredTime: "2022-12-08T21:08:19.224Z",
1088+
RetryCount: 2,
1089+
},
1090+
StateMachine: {
1091+
Id: "arn:aws:states:sa-east-1:425362996713:stateMachine:logs-to-traces-sequential",
1092+
Name: "my-state-machine",
1093+
},
10791094
} as const;
1095+
10801096
jest.spyOn(Date, "now").mockImplementation(() => 1487076708000);
10811097
process.env[xrayTraceEnvVar] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1";
10821098
process.env[awsXrayDaemonAddressEnvVar] = "localhost:127.0.0.1:2000";
@@ -1089,7 +1105,7 @@ describe("extractTraceContext", () => {
10891105
const sentMessage = sentSegment.toString();
10901106
expect(sentMessage).toMatchInlineSnapshot(`
10911107
"{\\"format\\": \\"json\\", \\"version\\": 1}
1092-
{\\"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\\"}}}}"
1108+
{\\"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}}}}"
10931109
`);
10941110
});
10951111
});

‎src/trace/context.ts

Lines changed: 58 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import {
1212
isSQSEvent,
1313
} from "../utils/event-type-guards";
1414
import {
15+
authorizingRequestIdHeader,
1516
awsXrayDaemonAddressEnvVar,
1617
parentIDHeader,
1718
SampleMode,
@@ -26,8 +27,7 @@ import {
2627
xrayTraceEnvVar,
2728
} from "./constants";
2829
import { TraceExtractor } from "./listener";
29-
import { parseEventSourceSubType, eventSubTypes } from "./trigger";
30-
import { authorizingRequestIdHeader } from "./constants";
30+
import { eventSubTypes, parseEventSourceSubType } from "./trigger";
3131

3232
export interface XRayTraceHeader {
3333
traceID: string;
@@ -43,11 +43,16 @@ export interface TraceContext {
4343
}
4444

4545
export interface StepFunctionContext {
46-
"step_function.retry_count": number;
46+
"step_function.execution_name": string;
4747
"step_function.execution_id": string;
48+
"step_function.execution_input": object;
49+
"step_function.execution_role_arn": string;
50+
"step_function.execution_start_time": string;
4851
"step_function.state_machine_name": string;
4952
"step_function.state_machine_arn": string;
50-
"step_function.step_name": string;
53+
"step_function.state_entered_time": string;
54+
"step_function.state_name": string;
55+
"step_function.state_retry_count": number;
5156
}
5257

5358
/**
@@ -485,50 +490,85 @@ export function readStepFunctionContextFromEvent(event: any): StepFunctionContex
485490
if (typeof event !== "object") {
486491
return;
487492
}
488-
const { dd } = event;
489-
if (typeof dd !== "object") {
490-
return;
491-
}
492-
const execution = dd.Execution;
493+
494+
const execution = event.Execution;
493495
if (typeof execution !== "object") {
496+
logDebug("event.Execution is not an object.");
494497
return;
495498
}
496-
const executionID = execution.Name;
499+
const executionID = execution.Id;
497500
if (typeof executionID !== "string") {
501+
logDebug("event.Execution.Id is not a string.");
502+
return;
503+
}
504+
const executionInput = execution.Input;
505+
const executionName = execution.Name;
506+
if (typeof executionName !== "string") {
507+
logDebug("event.Execution.Name is not a string.");
508+
return;
509+
}
510+
const executionRoleArn = execution.RoleArn;
511+
if (typeof executionRoleArn !== "string") {
512+
logDebug("event.Execution.RoleArn is not a string.");
498513
return;
499514
}
500-
const state = dd.State;
515+
const executionStartTime = execution.StartTime;
516+
if (typeof executionStartTime !== "string") {
517+
logDebug("event.Execution.StartTime is not a string.");
518+
return;
519+
}
520+
521+
const state = event.State;
501522
if (typeof state !== "object") {
523+
logDebug("event.State is not an object.");
524+
return;
525+
}
526+
const stateRetryCount = state.RetryCount;
527+
if (typeof stateRetryCount !== "number") {
528+
logDebug("event.State.RetryCount is not a string.");
502529
return;
503530
}
504-
const retryCount = state.RetryCount;
505-
if (typeof retryCount !== "number") {
531+
const stateEnteredTime = state.EnteredTime;
532+
if (typeof stateEnteredTime !== "string") {
533+
logDebug("event.State.EnteredTime is not a string.");
506534
return;
507535
}
508-
const stepName = state.Name;
509-
if (typeof stepName !== "string") {
536+
const stateName = state.Name;
537+
if (typeof stateName !== "string") {
538+
logDebug("event.State.Name is not a string.");
510539
return;
511540
}
512-
const stateMachine = dd.StateMachine;
541+
542+
const stateMachine = event.StateMachine;
513543
if (typeof stateMachine !== "object") {
544+
logDebug("event.StateMachine is not an object.");
514545
return;
515546
}
516547
const stateMachineArn = stateMachine.Id;
517548
if (typeof stateMachineArn !== "string") {
549+
logDebug("event.StateMachine.Id is not a string.");
518550
return;
519551
}
520552
const stateMachineName = stateMachine.Name;
521553
if (typeof stateMachineName !== "string") {
554+
logDebug("event.StateMachine.Name is not a string.");
522555
return;
523556
}
557+
524558
return {
559+
"step_function.execution_name": executionName,
525560
"step_function.execution_id": executionID,
526-
"step_function.retry_count": retryCount,
561+
"step_function.execution_input": executionInput ?? {},
562+
"step_function.execution_role_arn": executionRoleArn,
563+
"step_function.execution_start_time": executionStartTime,
564+
"step_function.state_entered_time": stateEnteredTime,
527565
"step_function.state_machine_arn": stateMachineArn,
528566
"step_function.state_machine_name": stateMachineName,
529-
"step_function.step_name": stepName,
567+
"step_function.state_name": stateName,
568+
"step_function.state_retry_count": stateRetryCount,
530569
};
531570
}
571+
532572
export function convertToSampleMode(xraySampled: number): SampleMode {
533573
return xraySampled === 1 ? SampleMode.USER_KEEP : SampleMode.USER_REJECT;
534574
}

0 commit comments

Comments
 (0)