Skip to content

Commit 6521850

Browse files
joeyzhao2018joeyzhao2018
committed
formatting fixes
1 parent 9b25540 commit 6521850

File tree

4 files changed

+93
-30
lines changed

4 files changed

+93
-30
lines changed

‎.vscode/settings.json

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,7 @@
11
{
22
"python.pythonPath": "/usr/local/bin/python3",
3-
"python.formatting.provider": "black"
3+
"python.formatting.provider": "black",
4+
"python.testing.unittestArgs": ["-v", "-s", "./tests", "-p", "test*.py"],
5+
"python.testing.pytestEnabled": false,
6+
"python.testing.unittestEnabled": true
47
}

‎datadog_lambda/tracing.py

Lines changed: 73 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -170,7 +170,9 @@ def extract_context_from_lambda_context(lambda_context):
170170
return trace_id, parent_id, sampling_priority
171171

172172

173-
def extract_context_from_http_event_or_context(event, lambda_context, event_source: _EventSource):
173+
def extract_context_from_http_event_or_context(
174+
event, lambda_context, event_source: _EventSource
175+
):
174176
"""
175177
Extract Datadog trace context from the `headers` key in from the Lambda
176178
`event` object.
@@ -340,10 +342,13 @@ def get_injected_authorizer_data(event, event_source: _EventSource) -> dict:
340342
# has '_datadog' header and not cached because integrationLatency > 0
341343
authorizer_headers = event.get("requestContext", {}).get("authorizer", {})
342344

343-
# [API_GATEWAY V1]if integrationLatency == 0, it's cached and no trace for the authorizer func
344-
if event_source.equals(
345-
EventTypes.API_GATEWAY, subtype=EventSubtypes.API_GATEWAY
346-
) and int(authorizer_headers.get("integrationLatency", 0)) == 0:
345+
# [API_GATEWAY V1]if integrationLatency == 0, it's cached and no trace for the authorizer
346+
if (
347+
event_source.equals(
348+
EventTypes.API_GATEWAY, subtype=EventSubtypes.API_GATEWAY
349+
)
350+
and int(authorizer_headers.get("integrationLatency", 0)) == 0
351+
):
347352
return None
348353

349354
if event_source.equals(EventTypes.API_GATEWAY, subtype=EventSubtypes.HTTP_API):
@@ -383,7 +388,9 @@ def extract_dd_trace_context(event, lambda_context, extractor=None):
383388
trace_id,
384389
parent_id,
385390
sampling_priority,
386-
) = extract_context_from_http_event_or_context(event, lambda_context, event_source)
391+
) = extract_context_from_http_event_or_context(
392+
event, lambda_context, event_source
393+
)
387394
elif event_source.equals(EventTypes.SNS) or event_source.equals(EventTypes.SQS):
388395
(
389396
trace_id,
@@ -689,13 +696,29 @@ def create_inferred_span_from_api_gateway_websocket_event(event, context):
689696
}
690697
tracer.set_tags({"_dd.origin": "lambda"})
691698
upstream_authorizer_span = None
692-
injected_authorizer_data = get_injected_authorizer_data(event, _EventSource(EventTypes.API_GATEWAY, EventSubtypes.WEBSOCKET))
699+
finish_time_s = None
700+
injected_authorizer_data = get_injected_authorizer_data(
701+
event, _EventSource(EventTypes.API_GATEWAY, EventSubtypes.WEBSOCKET)
702+
)
693703
if injected_authorizer_data:
694704
try:
695-
start_time_s = int(injected_authorizer_data.get(OtherConsts.parentSpanFinishTimeHeader)) / 1000
696-
finish_time_s = request_time_epoch_s + \
697-
(int(request_context.get("authorizer", {}).get("integrationLatency", 0))
698-
) / 1000
705+
start_time_s = (
706+
int(
707+
injected_authorizer_data.get(OtherConsts.parentSpanFinishTimeHeader)
708+
)
709+
/ 1000
710+
)
711+
finish_time_s = (
712+
request_time_epoch_s
713+
+ (
714+
int(
715+
request_context.get("authorizer", {}).get(
716+
"integrationLatency", 0
717+
)
718+
)
719+
)
720+
/ 1000
721+
)
699722
upstream_authorizer_span = insert_upstream_authorizer_span(
700723
args, tags, start_time_s, finish_time_s
701724
)
@@ -714,7 +737,9 @@ def create_inferred_span_from_api_gateway_websocket_event(event, context):
714737
span = tracer.trace("aws.apigateway.websocket", **args)
715738
if span:
716739
span.set_tags(tags)
717-
span.start = finish_time_s if finish_time_s is not None else request_time_epoch_s
740+
span.start = (
741+
finish_time_s if finish_time_s is not None else request_time_epoch_s
742+
)
718743
if upstream_authorizer_span:
719744
span.parent_id = upstream_authorizer_span.span_id
720745
return span
@@ -751,12 +776,28 @@ def create_inferred_span_from_api_gateway_event(event, context):
751776
tracer.set_tags({"_dd.origin": "lambda"})
752777
upstream_authorizer_span = None
753778
finish_time_s = None
754-
injected_authorizer_data = get_injected_authorizer_data(event, _EventSource(EventTypes.API_GATEWAY, EventSubtypes.API_GATEWAY))
779+
injected_authorizer_data = get_injected_authorizer_data(
780+
event, _EventSource(EventTypes.API_GATEWAY, EventSubtypes.API_GATEWAY)
781+
)
755782
if injected_authorizer_data:
756783
try:
757-
start_time_s = int(injected_authorizer_data.get(OtherConsts.parentSpanFinishTimeHeader)) / 1000
758-
finish_time_s = request_time_epoch_s + \
759-
(int(request_context.get("authorizer", {}).get("integrationLatency", 0))) / 1000
784+
start_time_s = (
785+
int(
786+
injected_authorizer_data.get(OtherConsts.parentSpanFinishTimeHeader)
787+
)
788+
/ 1000
789+
)
790+
finish_time_s = (
791+
request_time_epoch_s
792+
+ (
793+
int(
794+
request_context.get("authorizer", {}).get(
795+
"integrationLatency", 0
796+
)
797+
)
798+
)
799+
/ 1000
800+
)
760801
upstream_authorizer_span = insert_upstream_authorizer_span(
761802
args, tags, start_time_s, finish_time_s
762803
)
@@ -775,7 +816,9 @@ def create_inferred_span_from_api_gateway_event(event, context):
775816
if span:
776817
span.set_tags(tags)
777818
# start time pushed by the inserted authorizer span
778-
span.start = finish_time_s if finish_time_s is not None else request_time_epoch_s
819+
span.start = (
820+
finish_time_s if finish_time_s is not None else request_time_epoch_s
821+
)
779822
if upstream_authorizer_span:
780823
span.parent_id = upstream_authorizer_span.span_id
781824
return span
@@ -815,11 +858,18 @@ def create_inferred_span_from_http_api_event(event, context):
815858
tracer.set_tags({"_dd.origin": "lambda"})
816859
upstream_authorizer_span = None
817860
finish_time_s = None
818-
injected_authorizer_data = get_injected_authorizer_data(event, _EventSource(EventTypes.API_GATEWAY, EventSubtypes.HTTP_API))
861+
injected_authorizer_data = get_injected_authorizer_data(
862+
event, _EventSource(EventTypes.API_GATEWAY, EventSubtypes.HTTP_API)
863+
)
819864
if injected_authorizer_data:
820865
try:
821-
start_time_s = int(injected_authorizer_data.get(OtherConsts.parentSpanFinishTimeHeader)) / 1000
822-
finish_time_s = start_time_s # we don't have the integrationLatency info for the authorizer
866+
start_time_s = (
867+
int(
868+
injected_authorizer_data.get(OtherConsts.parentSpanFinishTimeHeader)
869+
)
870+
/ 1000
871+
)
872+
finish_time_s = start_time_s # no integrationLatency info in this case
823873
upstream_authorizer_span = insert_upstream_authorizer_span(
824874
args, tags, start_time_s, finish_time_s
825875
)
@@ -836,7 +886,9 @@ def create_inferred_span_from_http_api_event(event, context):
836886
span = tracer.trace("aws.httpapi", **args)
837887
if span:
838888
span.set_tags(tags)
839-
span.start = finish_time_s if finish_time_s is not None else request_time_epoch_s
889+
span.start = (
890+
finish_time_s if finish_time_s is not None else request_time_epoch_s
891+
)
840892
if upstream_authorizer_span:
841893
span.parent_id = upstream_authorizer_span.span_id
842894
return span

‎datadog_lambda/wrapper.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -212,7 +212,9 @@ def _before(self, event, context):
212212
if dd_tracing_enabled:
213213
set_dd_trace_py_root(trace_context_source, self.merge_xray_traces)
214214
if self.make_inferred_span:
215-
self.inferred_span = create_inferred_span(event, context, event_source)
215+
self.inferred_span = create_inferred_span(
216+
event, context, event_source
217+
)
216218
self.span = create_function_execution_span(
217219
context,
218220
self.function_name,

‎tests/test_tracing.py

Lines changed: 13 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -588,7 +588,9 @@ def setUp(self):
588588

589589
def test_create_inferred_span_from_authorizer_request_api_gateway_v1_event(self):
590590
event_sample_source = "authorizer-request-api-gateway-v1"
591-
span = self._http_common_testing_items(event_sample_source, "aws.apigateway.rest", 1663295021.832)
591+
span = self._http_common_testing_items(
592+
event_sample_source, "aws.apigateway.rest", 1663295021.832
593+
)
592594

593595
# api-gateway-v1 specific checks:
594596
self.assertEqual(
@@ -600,7 +602,9 @@ def test_create_inferred_span_from_authorizer_request_api_gateway_v1_event(self)
600602

601603
def test_create_inferred_span_from_authorizer_token_api_gateway_v1_event(self):
602604
event_sample_source = "authorizer-token-api-gateway-v1"
603-
span = self._http_common_testing_items(event_sample_source, "aws.apigateway.rest", 1663295021.832)
605+
span = self._http_common_testing_items(
606+
event_sample_source, "aws.apigateway.rest", 1663295021.832
607+
)
604608

605609
# api-gateway-v1 specific checks:
606610
self.assertEqual(
@@ -612,7 +616,9 @@ def test_create_inferred_span_from_authorizer_token_api_gateway_v1_event(self):
612616

613617
def test_create_inferred_span_from_authorizer_request_api_gateway_v2_event(self):
614618
event_sample_source = "authorizer-request-api-gateway-v2"
615-
span = self._http_common_testing_items(event_sample_source, "aws.httpapi", 1663721602.44)
619+
span = self._http_common_testing_items(
620+
event_sample_source, "aws.httpapi", 1663721602.44
621+
)
616622

617623
# http-api specific checks:
618624
self.assertEqual(span.get_tag("apiid"), "amddr1rix9")
@@ -625,7 +631,9 @@ def test_create_inferred_span_from_authorizer_request_api_gateway_v2_event(self)
625631
span.start, 1663721602.44
626632
) # use the injected parent span finish time as an approximation
627633

628-
def _http_common_testing_items(self, event_sample_source, operation_name, finish_time):
634+
def _http_common_testing_items(
635+
self, event_sample_source, operation_name, finish_time
636+
):
629637
test_file = event_samples + event_sample_source + ".json"
630638
with open(test_file, "r") as event:
631639
event = json.load(event)
@@ -647,9 +655,7 @@ def _http_common_testing_items(self, event_sample_source, operation_name, finish
647655
span.get_tag("resource_names"),
648656
"GET /hello",
649657
)
650-
self.assertEqual(
651-
span.get_tag("request_id"), "1234567"
652-
)
658+
self.assertEqual(span.get_tag("request_id"), "1234567")
653659

654660
self.assertEqual(span.span_type, "http")
655661
self.assertEqual(span.get_tag(InferredSpanInfo.TAG_SOURCE), "self")

0 commit comments

Comments
 (0)