From a1e98e50a9ef37cb073097a41ba4630cc2e63b8a Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Fri, 16 Sep 2022 17:09:27 -0400 Subject: [PATCH 01/26] authorizer spans checkin --- datadog_lambda/constants.py | 4 + datadog_lambda/tracing.py | 156 ++++++++++++++++-- datadog_lambda/wrapper.py | 29 +++- .../authorizer-request-api-gateway-v1.json | 93 +++++++++++ tests/test_tracing.py | 50 ++++++ 5 files changed, 320 insertions(+), 12 deletions(-) create mode 100644 tests/event_samples/authorizer-request-api-gateway-v1.json diff --git a/datadog_lambda/constants.py b/datadog_lambda/constants.py index 9bd49ce3..61945f42 100644 --- a/datadog_lambda/constants.py +++ b/datadog_lambda/constants.py @@ -42,3 +42,7 @@ class XrayDaemon(object): XRAY_TRACE_ID_HEADER_NAME = "_X_AMZN_TRACE_ID" XRAY_DAEMON_ADDRESS = "AWS_XRAY_DAEMON_ADDRESS" FUNCTION_NAME_HEADER_NAME = "AWS_LAMBDA_FUNCTION_NAME" + + +class OtherConsts(object): + parentSpanFinishTimeHeader = "x-datadog-parent-span-finish-time" diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 4db7db98..16761b8f 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -8,6 +8,7 @@ import json import base64 from datetime import datetime, timezone +import traceback from typing import Optional, Dict from datadog_lambda.metric import submit_errors_metric @@ -23,12 +24,13 @@ TraceHeader, TraceContextSource, XrayDaemon, + OtherConsts, ) from datadog_lambda.xray import ( send_segment, parse_xray_header, ) -from ddtrace import tracer, patch +from ddtrace import tracer, patch, Span from ddtrace import __version__ as ddtrace_version from ddtrace.propagation.http import HTTPPropagator from datadog_lambda import __version__ as datadog_lambda_version @@ -174,6 +176,12 @@ def extract_context_from_http_event_or_context(event, lambda_context): Falls back to lambda context if no trace data is found in the `headers` """ + if is_non_cached_authorizer_invocation(event): + try: + return extract_context_from_authorizer_event(event) + except Exception as e: + logger.debug("extract_context_from_authorizer_event returned with error. \ + Continue without injecting the authorizer span %s", e) headers = event.get("headers", {}) lowercase_headers = {k.lower(): v for k, v in headers.items()} @@ -317,6 +325,26 @@ def extract_context_custom_extractor(extractor, event, lambda_context): return None, None, None +def extract_context_from_authorizer_event(event): + # fail fast if any KeyError until '_datadog' + dd_data = json.loads(event['requestContext']['authorizer']['_datadog']) + trace_id = dd_data.get(TraceHeader.TRACE_ID) + parent_id = dd_data.get(TraceHeader.PARENT_ID) + sampling_priority = dd_data.get(TraceHeader.SAMPLING_PRIORITY) + return trace_id, parent_id, sampling_priority + + +def is_non_cached_authorizer_invocation(event) -> bool: + try: + # has '_datadog' header and not cached because integrationLatency > 0 + injected_authorizer_headers = event.get('requestContext', {}).get('authorizer', {}) + return injected_authorizer_headers.get('_datadog') and \ + int(injected_authorizer_headers.get('integrationLatency', 0)) > 0 # non-cached + except Exception as e: + logger.debug("Failed to check if invocated by an authorizer. error %s", e) + return False + + def extract_dd_trace_context(event, lambda_context, extractor=None): """ Extract Datadog trace context from the Lambda `event` object. @@ -543,7 +571,7 @@ def create_inferred_span(event, context): return create_inferred_span_from_eventbridge_event(event, context) except Exception as e: logger.debug( - "Unable to infer span. Detected type: {}. Reason: {}", + "Unable to infer span. Detected type: %s. Reason: %s", event_source.to_string(), e, ) @@ -588,7 +616,42 @@ def is_api_gateway_invocation_async(event): return event.get("headers", {}).get("X-Amz-Invocation-Type") == "Event" +def insert_upstream_authorizer_span( + kwargs_to_start_span, + other_tags_for_span, + start_time_s, + finish_time_s=None): + """ Insert the authorizer span. + Without this: parent span --child-> inferred span + With this insertion: parent span --child-> upstreamAuthorizerSpan --child-> inferred span + + Args: + kwargs_to_start_span (Dict): the same keyword arguments used for the inferred span + other_tags_for_span (Dict): the same tag keyword arguments used for the inferred span + start_time_s (int): the start time of the span in seconds + finish_time_s (int): the finish time of the sapn in seconds + """ + + upstream_authorizer_span = tracer.trace("aws.apigateway.authorizer", **kwargs_to_start_span) + upstream_authorizer_span.set_tags(other_tags_for_span) + upstream_authorizer_span.set_tag("operation_name", "aws.apigateway.authorizer") + # always sync for the authorizer invocation + InferredSpanInfo.set_tags_to_span(upstream_authorizer_span, synchronicity="sync") + upstream_authorizer_span.start = start_time_s + upstream_authorizer_span.finish(finish_time=finish_time_s) + return upstream_authorizer_span + + +def get_start_and_finish_time_for_authorizer_span(request_context): + parsed_upstream_context = json.loads(request_context.get('authorizer', {}).get('_datadog')) + start_time_s = int(parsed_upstream_context.get(OtherConsts.parentSpanFinishTimeHeader)) / 1000 + finish_time_s = (int(request_context.get('requestTimeEpoch')) + int(request_context.get( + 'authorizer', {}).get('integrationLatency', 0))) / 1000 + return start_time_s, finish_time_s + + def create_inferred_span_from_api_gateway_websocket_event(event, context): + trace_ctx = tracer.current_trace_context() request_context = event.get("requestContext") domain = request_context.get("domainName") endpoint = request_context.get("routeKey") @@ -616,14 +679,31 @@ def create_inferred_span_from_api_gateway_websocket_event(event, context): "span_type": "web", } tracer.set_tags({"_dd.origin": "lambda"}) + upstream_authorizer_span = None + if is_non_cached_authorizer_invocation(event): + try: + start_time_s, finish_time_s = \ + get_start_and_finish_time_for_authorizer_span(request_context) + upstream_authorizer_span = \ + insert_upstream_authorizer_span(args, tags, start_time_s, finish_time_s) + except Exception as e: + traceback.print_exc() + logger.debug("Unable to insert authorizer span. Continue to generate the main span.\ + Event %s, Reason: %s", event, e) + # trace context needs to be set again as it is reset by upstream_authorizer_span.finish + tracer.context_provider.activate(trace_ctx) span = tracer.trace("aws.apigateway.websocket", **args) if span: span.set_tags(tags) - span.start = request_time_epoch / 1000 + logger.debug(f"request_time_epoch: {request_time_epoch}, type: websocket") + span.start = request_time_epoch / 1000 + if upstream_authorizer_span: + span.parent_id = upstream_authorizer_span.span_id return span def create_inferred_span_from_api_gateway_event(event, context): + trace_ctx = tracer.current_trace_context() request_context = event.get("requestContext") domain = request_context.get("domainName", "") method = event.get("httpMethod") @@ -640,7 +720,7 @@ def create_inferred_span_from_api_gateway_event(event, context): "stage": request_context.get("stage"), "request_id": request_context.get("requestId"), } - request_time_epoch = request_context.get("requestTimeEpoch") + request_time_epoch = int(request_context.get("requestTimeEpoch")) / 1000 if is_api_gateway_invocation_async(event): InferredSpanInfo.set_tags(tags, tag_source="self", synchronicity="async") else: @@ -651,14 +731,32 @@ def create_inferred_span_from_api_gateway_event(event, context): "span_type": "http", } tracer.set_tags({"_dd.origin": "lambda"}) + upstream_authorizer_span = None + finish_time_s = None + if is_non_cached_authorizer_invocation(event): + try: + start_time_s, finish_time_s = \ + get_start_and_finish_time_for_authorizer_span(request_context) + upstream_authorizer_span = \ + insert_upstream_authorizer_span(args, tags, start_time_s, finish_time_s) + except Exception as e: + traceback.print_exc() + logger.debug("Unable to insert authorizer span. Continue to generate the main span. \ + Event %s, Reason: %s", event, e) + # trace context needs to be set again as it is reset by upstream_authorizer_span.finish + tracer.context_provider.activate(trace_ctx) span = tracer.trace("aws.apigateway", **args) if span: span.set_tags(tags) - span.start = request_time_epoch / 1000 + # start time pushed by the inserted authorizer span + span.start = finish_time_s if finish_time_s is not None else request_time_epoch + if upstream_authorizer_span: + span.parent_id = upstream_authorizer_span.span_id return span def create_inferred_span_from_http_api_event(event, context): + trace_ctx = tracer.current_trace_context() request_context = event.get("requestContext") domain = request_context.get("domainName") method = request_context.get("http", {}).get("method") @@ -682,6 +780,7 @@ def create_inferred_span_from_http_api_event(event, context): if is_api_gateway_invocation_async(event): InferredSpanInfo.set_tags(tags, tag_source="self", synchronicity="async") else: + traceback.print_exc() InferredSpanInfo.set_tags(tags, tag_source="self", synchronicity="sync") args = { "service": domain, @@ -689,10 +788,26 @@ def create_inferred_span_from_http_api_event(event, context): "span_type": "http", } tracer.set_tags({"_dd.origin": "lambda"}) + upstream_authorizer_span = None + if is_non_cached_authorizer_invocation(event): + try: + start_time_s, finish_time_s = \ + get_start_and_finish_time_for_authorizer_span(request_context) + upstream_authorizer_span = \ + insert_upstream_authorizer_span(args, tags, start_time_s, finish_time_s) + except Exception as e: + traceback.print_exc() + logger.debug("Unable to insert authorizer span. \ + Continue to generate the main span. Event %s, Reason: %s", event, e) + # trace context needs to be set again as it is reset by upstream_authorizer_span.finish + tracer.context_provider.activate(trace_ctx) span = tracer.trace("aws.httpapi", **args) if span: span.set_tags(tags) - span.start = request_time_epoch / 1000 + logger.debug(f"request_time_epoch: {request_time_epoch}, type: http_api") + span.start = request_time_epoch / 1e3 + if upstream_authorizer_span: + span.parent_id = upstream_authorizer_span.span_id return span @@ -973,14 +1088,35 @@ class InferredSpanInfo(object): SYNCHRONICITY = f"{BASE_NAME}.synchronicity" TAG_SOURCE = f"{BASE_NAME}.tag_source" - @classmethod + @staticmethod def set_tags( - cls, tags: Dict[str, str], synchronicity: Optional[Literal["sync", "async"]] = None, tag_source: Optional[Literal["labmda", "self"]] = None, ): if synchronicity is not None: - tags[cls.SYNCHRONICITY] = str(synchronicity) + tags[InferredSpanInfo.SYNCHRONICITY] = str(synchronicity) + if tag_source is not None: + tags[InferredSpanInfo.TAG_SOURCE] = str(tag_source) + + @staticmethod + def set_tags_to_span( + span: Span, + synchronicity: Optional[Literal["sync", "async"]] = None, + tag_source: Optional[Literal["labmda", "self"]] = None, + ): + if synchronicity is not None: + span.set_tags({InferredSpanInfo.SYNCHRONICITY: synchronicity}) if tag_source is not None: - tags[cls.TAG_SOURCE] = str(tag_source) + span.set_tags({InferredSpanInfo.TAG_SOURCE: str(tag_source)}) + + @staticmethod + def is_async(span: Span) -> bool: + if not span: + return False + try: + return span.get_tag(InferredSpanInfo.SYNCHRONICITY) == 'async' + except Exception as e: + logger.debug("Unabled to read the %s tag, returning False. \ + Reason: %s.", InferredSpanInfo.SYNCHRONICITY, e) + return False diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 88508ab3..3050b9e7 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -7,12 +7,16 @@ import logging import traceback from importlib import import_module +import json +from time import time_ns +from ddtrace.propagation.http import HTTPPropagator from datadog_lambda.extension import should_use_extension, flush_extension from datadog_lambda.cold_start import set_cold_start, is_cold_start from datadog_lambda.constants import ( TraceContextSource, XraySubsegment, + OtherConsts, ) from datadog_lambda.metric import ( flush_stats, @@ -116,6 +120,9 @@ def __init__(self, func): self.make_inferred_span = ( os.environ.get("DD_TRACE_MANAGED_SERVICES", "true").lower() == "true" ) + self.make_authorizer_span = ( + os.environ.get("DD_LINK_LAMBDA_AUTHORIZER", "true").lower() == "true" + ) self.response = None if profiling_env_var: self.prof = profiler.Profiler(env=env_env_var, service=service_env_var) @@ -148,6 +155,14 @@ def __call__(self, event, context, **kwargs): self._before(event, context) try: self.response = self.func(event, context, **kwargs) + try: + if self.make_authorizer_span and self.response and \ + self.response.get('principalId') and self.response.get('policyDocument'): + self._inject_authorizer_span_headers() + except Exception as e: + traceback.print_exc() + logger.debug("Unable to inject the authorizer headers. \ + Continue to return the original response. Reason: %s", e) return self.response except Exception: submit_errors_metric(context) @@ -157,6 +172,17 @@ def __call__(self, event, context, **kwargs): finally: self._after(event, context) + def _inject_authorizer_span_headers(self): + finish_time_ns = self.span.start_ns if InferredSpanInfo.is_async(self.inferred_span) \ + and self.span else time_ns() + self.response.setdefault('context', {}) + self.response['context'].setdefault('_datadog', {}) + injected_headers = {} + source_span = self.inferred_span if self.inferred_span else self.span + HTTPPropagator.inject(source_span.context, injected_headers) + injected_headers[OtherConsts.parentSpanFinishTimeHeader] = finish_time_ns / 1e6 + self.response['context']['_datadog'] = json.dumps(injected_headers) + def _before(self, event, context): try: self.response = None @@ -222,8 +248,7 @@ def _after(self, event, context): self.inferred_span.set_tag("http.status_code", status_code) if ( - self.inferred_span.get_tag(InferredSpanInfo.SYNCHRONICITY) - == "async" + InferredSpanInfo.is_async(self.inferred_span) and self.span ): self.inferred_span.finish(finish_time=self.span.start) diff --git a/tests/event_samples/authorizer-request-api-gateway-v1.json b/tests/event_samples/authorizer-request-api-gateway-v1.json new file mode 100644 index 00000000..fc16e7fc --- /dev/null +++ b/tests/event_samples/authorizer-request-api-gateway-v1.json @@ -0,0 +1,93 @@ +{ + "resource": "/hello", + "path": "/hello", + "httpMethod": "GET", + "headers": { + "Accept": "*/*", + "Accept-Encoding": "gzip, deflate, br", + "Authorization": "secretT0k3n", + "authorizationToken": "secretT0k3n", + "Cache-Control": "no-cache", + "CloudFront-Forwarded-Proto": "https", + "CloudFront-Is-Desktop-Viewer": "true", + "CloudFront-Is-Mobile-Viewer": "false", + "CloudFront-Is-SmartTV-Viewer": "false", + "CloudFront-Is-Tablet-Viewer": "false", + "CloudFront-Viewer-ASN": "12271", + "CloudFront-Viewer-Country": "US", + "Host": "amddr1rix9.execute-api.sa-east-1.amazonaws.com", + "Postman-Token": "75a5e2b0-d8e8-4370-ad2c-7209761c3f43", + "User-Agent": "PostmanRuntime/7.29.2", + "Via": "1.1 ed4584f7c263c11cf4adf75ba3a25764.cloudfront.net (CloudFront)", + "X-Amz-Cf-Id": "wbruuv68N6jFldH7oUaE9kWABzHsvPi3zzByABVGXCVsnGYOhcRjow==", + "X-Amzn-Trace-Id": "Root=1-6323de2b-0155211e49422a532809a10e", + "X-Forwarded-For": "24.193.182.233, 15.158.35.8", + "X-Forwarded-Port": "443", + "X-Forwarded-Proto": "https" + }, + "multiValueHeaders": { + "Accept": ["*/*"], + "Accept-Encoding": ["gzip, deflate, br"], + "Authorization": ["secretT0k3n"], + "authorizationToken": ["secretT0k3n"], + "Cache-Control": ["no-cache"], + "CloudFront-Forwarded-Proto": ["https"], + "CloudFront-Is-Desktop-Viewer": ["true"], + "CloudFront-Is-Mobile-Viewer": ["false"], + "CloudFront-Is-SmartTV-Viewer": ["false"], + "CloudFront-Is-Tablet-Viewer": ["false"], + "CloudFront-Viewer-ASN": ["12271"], + "CloudFront-Viewer-Country": ["US"], + "Host": ["amddr1rix9.execute-api.sa-east-1.amazonaws.com"], + "Postman-Token": ["75a5e2b0-d8e8-4370-ad2c-7209761c3f43"], + "User-Agent": ["PostmanRuntime/7.29.2"], + "Via": ["1.1 ed4584f7c263c11cf4adf75ba3a25764.cloudfront.net (CloudFront)"], + "X-Amz-Cf-Id": ["wbruuv68N6jFldH7oUaE9kWABzHsvPi3zzByABVGXCVsnGYOhcRjow=="], + "X-Amzn-Trace-Id": ["Root=1-6323de2b-0155211e49422a532809a10e"], + "X-Forwarded-For": ["24.193.182.233, 15.158.35.8"], + "X-Forwarded-Port": ["443"], + "X-Forwarded-Proto": ["https"] + }, + "queryStringParameters": null, + "multiValueQueryStringParameters": null, + "pathParameters": null, + "stageVariables": null, + "requestContext": { + "resourceId": "0et54l", + "authorizer": { + "_datadog": "{\"x-datadog-trace-id\": \"13478705995797221209\", \"x-datadog-parent-id\": \"8471288263384216896\", \"x-datadog-sampling-priority\": \"1\", \"x-datadog-tags\": \"_dd.p.dm=-0\", \"x-datadog-parent-span-finish-time\": 1663295021827.521}", + "scope": "this is just a string", + "principalId": "foo", + "integrationLatency": 1897 + }, + "resourcePath": "/hello", + "httpMethod": "GET", + "extendedRequestId": "Yh-m5G_wGjQFftQ=", + "requestTime": "16/Sep/2022:02:23:39 +0000", + "path": "/dev/hello", + "accountId": "601427279990", + "protocol": "HTTP/1.1", + "stage": "dev", + "domainPrefix": "amddr1rix9", + "requestTimeEpoch": 1663295019935, + "requestId": "e4816b6c-8954-4510-8771-837988d7f485", + "identity": { + "cognitoIdentityPoolId": null, + "accountId": null, + "cognitoIdentityId": null, + "caller": null, + "sourceIp": "24.193.182.233", + "principalOrgId": null, + "accessKey": null, + "cognitoAuthenticationType": null, + "cognitoAuthenticationProvider": null, + "userArn": null, + "userAgent": "PostmanRuntime/7.29.2", + "user": null + }, + "domainName": "amddr1rix9.execute-api.sa-east-1.amazonaws.com", + "apiId": "amddr1rix9" + }, + "body": null, + "isBase64Encoded": false +} diff --git a/tests/test_tracing.py b/tests/test_tracing.py index c578de07..045435c2 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -579,6 +579,56 @@ def test_mixed_parent_context_when_merging(self): self.mock_activate.assert_has_calls([call(expected_context)]) +class TestAuthorizerInferredSpans(unittest.TestCase): + + def setUp(self): + patcher = patch('ddtrace.Span.finish', autospec=True) + self.mock_span_stop = patcher.start() + self.addCleanup(patcher.stop) + + def test_create_inferred_span_from_authorizer_request_api_gateway_v1_event(self): + event_sample_source = "authorizer-request-api-gateway-v1" + test_file = event_samples + event_sample_source + ".json" + with open(test_file, "r") as event: + event = json.load(event) + ctx = get_mock_context() + ctx.aws_request_id = "123" + span = create_inferred_span(event, ctx) + self.assertEqual(span.get_tag("operation_name"), "aws.apigateway.rest") + self.assertEqual( + span.service, + "amddr1rix9.execute-api.sa-east-1.amazonaws.com", + ) + self.assertEqual( + span.get_tag("http.url"), + "amddr1rix9.execute-api.sa-east-1.amazonaws.com/hello", + ) + self.assertEqual(span.get_tag("endpoint"), "/hello") + self.assertEqual(span.get_tag("http.method"), "GET") + self.assertEqual( + span.get_tag("resource_names"), + "GET /hello", + ) + self.assertEqual( + span.get_tag("request_id"), "e4816b6c-8954-4510-8771-837988d7f485" + ) + self.assertEqual(span.get_tag("apiid"), "amddr1rix9") + self.assertEqual(span.get_tag("apiname"), "amddr1rix9") + self.assertEqual(span.get_tag("stage"), "dev") + self.assertEqual(span.start, 1663295021.832) # request_time_epoch + integrationLatency + self.assertEqual(span.span_type, "http") + self.assertEqual(span.get_tag(InferredSpanInfo.TAG_SOURCE), "self") + self.assertEqual(span.get_tag(InferredSpanInfo.SYNCHRONICITY), "sync") + + # checking the upstream_authorizer_span + self.mock_span_stop.assert_called_once() + args, kwargs = self.mock_span_stop.call_args_list[0] + self.assertEqual(kwargs['finish_time'], 1663295021.832) + authorizer_span = args[0] + self.assertEqual(authorizer_span.name, 'aws.apigateway.authorizer') + self.assertEqual(span.parent_id, authorizer_span.span_id) + + class TestInferredSpans(unittest.TestCase): def test_create_inferred_span_from_api_gateway_event(self): event_sample_source = "api-gateway" From 992369b7e129b9b6aa6f2b54ff01298e1f70c799 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Fri, 16 Sep 2022 17:09:27 -0400 Subject: [PATCH 02/26] authorizer spans checkin --- datadog_lambda/tracing.py | 2 -- datadog_lambda/wrapper.py | 2 +- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 16761b8f..23ded5c9 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -695,7 +695,6 @@ def create_inferred_span_from_api_gateway_websocket_event(event, context): span = tracer.trace("aws.apigateway.websocket", **args) if span: span.set_tags(tags) - logger.debug(f"request_time_epoch: {request_time_epoch}, type: websocket") span.start = request_time_epoch / 1000 if upstream_authorizer_span: span.parent_id = upstream_authorizer_span.span_id @@ -804,7 +803,6 @@ def create_inferred_span_from_http_api_event(event, context): span = tracer.trace("aws.httpapi", **args) if span: span.set_tags(tags) - logger.debug(f"request_time_epoch: {request_time_epoch}, type: http_api") span.start = request_time_epoch / 1e3 if upstream_authorizer_span: span.parent_id = upstream_authorizer_span.span_id diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 3050b9e7..b578edb6 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -157,7 +157,7 @@ def __call__(self, event, context, **kwargs): self.response = self.func(event, context, **kwargs) try: if self.make_authorizer_span and self.response and \ - self.response.get('principalId') and self.response.get('policyDocument'): + self.response.get('principalId') and self.response.get('policyDocument'): self._inject_authorizer_span_headers() except Exception as e: traceback.print_exc() From 9b564536a44feca1cab276099e4c6dabb327b0d2 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Tue, 20 Sep 2022 11:26:06 -0400 Subject: [PATCH 03/26] fix integration tests --- datadog_lambda/tracing.py | 1 - 1 file changed, 1 deletion(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 23ded5c9..87a8eac2 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -779,7 +779,6 @@ def create_inferred_span_from_http_api_event(event, context): if is_api_gateway_invocation_async(event): InferredSpanInfo.set_tags(tags, tag_source="self", synchronicity="async") else: - traceback.print_exc() InferredSpanInfo.set_tags(tags, tag_source="self", synchronicity="sync") args = { "service": domain, From 59df10b45269ca7fcffb1a8752c5cde346fdd4f8 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Tue, 20 Sep 2022 11:44:37 -0400 Subject: [PATCH 04/26] reformatting --- datadog_lambda/tracing.py | 109 +++++++++++++++++++++++++------------- datadog_lambda/wrapper.py | 33 +++++++----- 2 files changed, 92 insertions(+), 50 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 87a8eac2..4fc19891 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -180,8 +180,11 @@ def extract_context_from_http_event_or_context(event, lambda_context): try: return extract_context_from_authorizer_event(event) except Exception as e: - logger.debug("extract_context_from_authorizer_event returned with error. \ - Continue without injecting the authorizer span %s", e) + logger.debug( + "extract_context_from_authorizer_event returned with error. \ + Continue without injecting the authorizer span %s", + e, + ) headers = event.get("headers", {}) lowercase_headers = {k.lower(): v for k, v in headers.items()} @@ -327,7 +330,7 @@ def extract_context_custom_extractor(extractor, event, lambda_context): def extract_context_from_authorizer_event(event): # fail fast if any KeyError until '_datadog' - dd_data = json.loads(event['requestContext']['authorizer']['_datadog']) + dd_data = json.loads(event["requestContext"]["authorizer"]["_datadog"]) trace_id = dd_data.get(TraceHeader.TRACE_ID) parent_id = dd_data.get(TraceHeader.PARENT_ID) sampling_priority = dd_data.get(TraceHeader.SAMPLING_PRIORITY) @@ -337,9 +340,13 @@ def extract_context_from_authorizer_event(event): def is_non_cached_authorizer_invocation(event) -> bool: try: # has '_datadog' header and not cached because integrationLatency > 0 - injected_authorizer_headers = event.get('requestContext', {}).get('authorizer', {}) - return injected_authorizer_headers.get('_datadog') and \ - int(injected_authorizer_headers.get('integrationLatency', 0)) > 0 # non-cached + injected_authorizer_headers = event.get("requestContext", {}).get( + "authorizer", {} + ) + return ( + injected_authorizer_headers.get("_datadog") + and int(injected_authorizer_headers.get("integrationLatency", 0)) > 0 + ) # non-cached except Exception as e: logger.debug("Failed to check if invocated by an authorizer. error %s", e) return False @@ -617,11 +624,9 @@ def is_api_gateway_invocation_async(event): def insert_upstream_authorizer_span( - kwargs_to_start_span, - other_tags_for_span, - start_time_s, - finish_time_s=None): - """ Insert the authorizer span. + kwargs_to_start_span, other_tags_for_span, start_time_s, finish_time_s=None +): + """Insert the authorizer span. Without this: parent span --child-> inferred span With this insertion: parent span --child-> upstreamAuthorizerSpan --child-> inferred span @@ -632,7 +637,9 @@ def insert_upstream_authorizer_span( finish_time_s (int): the finish time of the sapn in seconds """ - upstream_authorizer_span = tracer.trace("aws.apigateway.authorizer", **kwargs_to_start_span) + upstream_authorizer_span = tracer.trace( + "aws.apigateway.authorizer", **kwargs_to_start_span + ) upstream_authorizer_span.set_tags(other_tags_for_span) upstream_authorizer_span.set_tag("operation_name", "aws.apigateway.authorizer") # always sync for the authorizer invocation @@ -643,10 +650,16 @@ def insert_upstream_authorizer_span( def get_start_and_finish_time_for_authorizer_span(request_context): - parsed_upstream_context = json.loads(request_context.get('authorizer', {}).get('_datadog')) - start_time_s = int(parsed_upstream_context.get(OtherConsts.parentSpanFinishTimeHeader)) / 1000 - finish_time_s = (int(request_context.get('requestTimeEpoch')) + int(request_context.get( - 'authorizer', {}).get('integrationLatency', 0))) / 1000 + parsed_upstream_context = json.loads( + request_context.get("authorizer", {}).get("_datadog") + ) + start_time_s = ( + int(parsed_upstream_context.get(OtherConsts.parentSpanFinishTimeHeader)) / 1000 + ) + finish_time_s = ( + int(request_context.get("requestTimeEpoch")) + + int(request_context.get("authorizer", {}).get("integrationLatency", 0)) + ) / 1000 return start_time_s, finish_time_s @@ -682,14 +695,20 @@ def create_inferred_span_from_api_gateway_websocket_event(event, context): upstream_authorizer_span = None if is_non_cached_authorizer_invocation(event): try: - start_time_s, finish_time_s = \ - get_start_and_finish_time_for_authorizer_span(request_context) - upstream_authorizer_span = \ - insert_upstream_authorizer_span(args, tags, start_time_s, finish_time_s) + start_time_s, finish_time_s = get_start_and_finish_time_for_authorizer_span( + request_context + ) + upstream_authorizer_span = insert_upstream_authorizer_span( + args, tags, start_time_s, finish_time_s + ) except Exception as e: traceback.print_exc() - logger.debug("Unable to insert authorizer span. Continue to generate the main span.\ - Event %s, Reason: %s", event, e) + logger.debug( + "Unable to insert authorizer span. Continue to generate the main span.\ + Event %s, Reason: %s", + event, + e, + ) # trace context needs to be set again as it is reset by upstream_authorizer_span.finish tracer.context_provider.activate(trace_ctx) span = tracer.trace("aws.apigateway.websocket", **args) @@ -734,14 +753,20 @@ def create_inferred_span_from_api_gateway_event(event, context): finish_time_s = None if is_non_cached_authorizer_invocation(event): try: - start_time_s, finish_time_s = \ - get_start_and_finish_time_for_authorizer_span(request_context) - upstream_authorizer_span = \ - insert_upstream_authorizer_span(args, tags, start_time_s, finish_time_s) + start_time_s, finish_time_s = get_start_and_finish_time_for_authorizer_span( + request_context + ) + upstream_authorizer_span = insert_upstream_authorizer_span( + args, tags, start_time_s, finish_time_s + ) except Exception as e: traceback.print_exc() - logger.debug("Unable to insert authorizer span. Continue to generate the main span. \ - Event %s, Reason: %s", event, e) + logger.debug( + "Unable to insert authorizer span. Continue to generate the main span. \ + Event %s, Reason: %s", + event, + e, + ) # trace context needs to be set again as it is reset by upstream_authorizer_span.finish tracer.context_provider.activate(trace_ctx) span = tracer.trace("aws.apigateway", **args) @@ -789,14 +814,20 @@ def create_inferred_span_from_http_api_event(event, context): upstream_authorizer_span = None if is_non_cached_authorizer_invocation(event): try: - start_time_s, finish_time_s = \ - get_start_and_finish_time_for_authorizer_span(request_context) - upstream_authorizer_span = \ - insert_upstream_authorizer_span(args, tags, start_time_s, finish_time_s) + start_time_s, finish_time_s = get_start_and_finish_time_for_authorizer_span( + request_context + ) + upstream_authorizer_span = insert_upstream_authorizer_span( + args, tags, start_time_s, finish_time_s + ) except Exception as e: traceback.print_exc() - logger.debug("Unable to insert authorizer span. \ - Continue to generate the main span. Event %s, Reason: %s", event, e) + logger.debug( + "Unable to insert authorizer span. \ + Continue to generate the main span. Event %s, Reason: %s", + event, + e, + ) # trace context needs to be set again as it is reset by upstream_authorizer_span.finish tracer.context_provider.activate(trace_ctx) span = tracer.trace("aws.httpapi", **args) @@ -1112,8 +1143,12 @@ def is_async(span: Span) -> bool: if not span: return False try: - return span.get_tag(InferredSpanInfo.SYNCHRONICITY) == 'async' + return span.get_tag(InferredSpanInfo.SYNCHRONICITY) == "async" except Exception as e: - logger.debug("Unabled to read the %s tag, returning False. \ - Reason: %s.", InferredSpanInfo.SYNCHRONICITY, e) + logger.debug( + "Unabled to read the %s tag, returning False. \ + Reason: %s.", + InferredSpanInfo.SYNCHRONICITY, + e, + ) return False diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index b578edb6..0e41a98d 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -156,13 +156,20 @@ def __call__(self, event, context, **kwargs): try: self.response = self.func(event, context, **kwargs) try: - if self.make_authorizer_span and self.response and \ - self.response.get('principalId') and self.response.get('policyDocument'): + if ( + self.make_authorizer_span + and self.response + and self.response.get("principalId") + and self.response.get("policyDocument") + ): self._inject_authorizer_span_headers() except Exception as e: traceback.print_exc() - logger.debug("Unable to inject the authorizer headers. \ - Continue to return the original response. Reason: %s", e) + logger.debug( + "Unable to inject the authorizer headers. \ + Continue to return the original response. Reason: %s", + e, + ) return self.response except Exception: submit_errors_metric(context) @@ -173,15 +180,18 @@ def __call__(self, event, context, **kwargs): self._after(event, context) def _inject_authorizer_span_headers(self): - finish_time_ns = self.span.start_ns if InferredSpanInfo.is_async(self.inferred_span) \ - and self.span else time_ns() - self.response.setdefault('context', {}) - self.response['context'].setdefault('_datadog', {}) + finish_time_ns = ( + self.span.start_ns + if InferredSpanInfo.is_async(self.inferred_span) and self.span + else time_ns() + ) + self.response.setdefault("context", {}) + self.response["context"].setdefault("_datadog", {}) injected_headers = {} source_span = self.inferred_span if self.inferred_span else self.span HTTPPropagator.inject(source_span.context, injected_headers) injected_headers[OtherConsts.parentSpanFinishTimeHeader] = finish_time_ns / 1e6 - self.response['context']['_datadog'] = json.dumps(injected_headers) + self.response["context"]["_datadog"] = json.dumps(injected_headers) def _before(self, event, context): try: @@ -247,10 +257,7 @@ def _after(self, event, context): if status_code: self.inferred_span.set_tag("http.status_code", status_code) - if ( - InferredSpanInfo.is_async(self.inferred_span) - and self.span - ): + if InferredSpanInfo.is_async(self.inferred_span) and self.span: self.inferred_span.finish(finish_time=self.span.start) else: self.inferred_span.finish() From 479353616274f4249925955678cda8d95f7c6f37 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Tue, 20 Sep 2022 13:07:40 -0400 Subject: [PATCH 05/26] tests reformat --- tests/test_tracing.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 045435c2..06e299e0 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -580,9 +580,8 @@ def test_mixed_parent_context_when_merging(self): class TestAuthorizerInferredSpans(unittest.TestCase): - def setUp(self): - patcher = patch('ddtrace.Span.finish', autospec=True) + patcher = patch("ddtrace.Span.finish", autospec=True) self.mock_span_stop = patcher.start() self.addCleanup(patcher.stop) @@ -615,7 +614,9 @@ def test_create_inferred_span_from_authorizer_request_api_gateway_v1_event(self) self.assertEqual(span.get_tag("apiid"), "amddr1rix9") self.assertEqual(span.get_tag("apiname"), "amddr1rix9") self.assertEqual(span.get_tag("stage"), "dev") - self.assertEqual(span.start, 1663295021.832) # request_time_epoch + integrationLatency + self.assertEqual( + span.start, 1663295021.832 + ) # request_time_epoch + integrationLatency self.assertEqual(span.span_type, "http") self.assertEqual(span.get_tag(InferredSpanInfo.TAG_SOURCE), "self") self.assertEqual(span.get_tag(InferredSpanInfo.SYNCHRONICITY), "sync") @@ -623,9 +624,9 @@ def test_create_inferred_span_from_authorizer_request_api_gateway_v1_event(self) # checking the upstream_authorizer_span self.mock_span_stop.assert_called_once() args, kwargs = self.mock_span_stop.call_args_list[0] - self.assertEqual(kwargs['finish_time'], 1663295021.832) + self.assertEqual(kwargs["finish_time"], 1663295021.832) authorizer_span = args[0] - self.assertEqual(authorizer_span.name, 'aws.apigateway.authorizer') + self.assertEqual(authorizer_span.name, "aws.apigateway.authorizer") self.assertEqual(span.parent_id, authorizer_span.span_id) From ae237c7db85fcab17698ee0f113ddf1b5ce4b616 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Wed, 21 Sep 2022 13:36:47 -0400 Subject: [PATCH 06/26] add httpapi case --- datadog_lambda/tracing.py | 125 +++++++++--------- datadog_lambda/wrapper.py | 4 +- .../authorizer-request-api-gateway-v1.json | 2 +- .../authorizer-request-api-gateway-v2.json | 46 +++++++ .../authorizer-token-api-gateway-v1.json | 93 +++++++++++++ tests/test_tracing.py | 91 +++++++++---- tests/test_wrapper.py | 2 +- 7 files changed, 270 insertions(+), 93 deletions(-) create mode 100644 tests/event_samples/authorizer-request-api-gateway-v2.json create mode 100644 tests/event_samples/authorizer-token-api-gateway-v1.json diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 4fc19891..ef7a7913 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -35,6 +35,7 @@ from ddtrace.propagation.http import HTTPPropagator from datadog_lambda import __version__ as datadog_lambda_version from datadog_lambda.trigger import ( + _EventSource, parse_event_source, get_first_record, EventTypes, @@ -169,16 +170,21 @@ def extract_context_from_lambda_context(lambda_context): return trace_id, parent_id, sampling_priority -def extract_context_from_http_event_or_context(event, lambda_context): +def extract_context_from_http_event_or_context(event, lambda_context, event_source: _EventSource): """ Extract Datadog trace context from the `headers` key in from the Lambda `event` object. Falls back to lambda context if no trace data is found in the `headers` """ - if is_non_cached_authorizer_invocation(event): + injected_authorizer_data = get_injected_authorizer_data(event, event_source) + if injected_authorizer_data: try: - return extract_context_from_authorizer_event(event) + # fail fast on any KeyError here + trace_id = injected_authorizer_data[TraceHeader.TRACE_ID] + parent_id = injected_authorizer_data[TraceHeader.PARENT_ID] + sampling_priority = injected_authorizer_data[TraceHeader.SAMPLING_PRIORITY] + return trace_id, parent_id, sampling_priority except Exception as e: logger.debug( "extract_context_from_authorizer_event returned with error. \ @@ -328,28 +334,31 @@ def extract_context_custom_extractor(extractor, event, lambda_context): return None, None, None -def extract_context_from_authorizer_event(event): - # fail fast if any KeyError until '_datadog' - dd_data = json.loads(event["requestContext"]["authorizer"]["_datadog"]) - trace_id = dd_data.get(TraceHeader.TRACE_ID) - parent_id = dd_data.get(TraceHeader.PARENT_ID) - sampling_priority = dd_data.get(TraceHeader.SAMPLING_PRIORITY) - return trace_id, parent_id, sampling_priority - - -def is_non_cached_authorizer_invocation(event) -> bool: +def get_injected_authorizer_data(event, event_source: _EventSource) -> dict: try: + dd_data_raw = None # has '_datadog' header and not cached because integrationLatency > 0 - injected_authorizer_headers = event.get("requestContext", {}).get( - "authorizer", {} - ) - return ( - injected_authorizer_headers.get("_datadog") - and int(injected_authorizer_headers.get("integrationLatency", 0)) > 0 - ) # non-cached + authorizer_headers = event.get("requestContext", {}).get("authorizer", {}) + + # [API_GATEWAY V1]if integrationLatency == 0, it's cached and no trace for the authorizer func + if event_source.equals( + EventTypes.API_GATEWAY, subtype=EventSubtypes.API_GATEWAY + ) and int(authorizer_headers.get("integrationLatency", 0)) == 0: + return None + + if event_source.equals(EventTypes.API_GATEWAY, subtype=EventSubtypes.HTTP_API): + dd_data_raw = authorizer_headers.get("lambda", {}).get("_datadog") + else: + dd_data_raw = authorizer_headers.get("_datadog") + + if dd_data_raw: + dd_data = json.loads(dd_data_raw) + return dd_data + return None + except Exception as e: logger.debug("Failed to check if invocated by an authorizer. error %s", e) - return False + return None def extract_dd_trace_context(event, lambda_context, extractor=None): @@ -374,7 +383,7 @@ def extract_dd_trace_context(event, lambda_context, extractor=None): trace_id, parent_id, sampling_priority, - ) = extract_context_from_http_event_or_context(event, lambda_context) + ) = extract_context_from_http_event_or_context(event, lambda_context, event_source) elif event_source.equals(EventTypes.SNS) or event_source.equals(EventTypes.SQS): ( trace_id, @@ -414,7 +423,7 @@ def extract_dd_trace_context(event, lambda_context, extractor=None): if dd_trace_context: trace_context_source = TraceContextSource.XRAY logger.debug("extracted dd trace context %s", dd_trace_context) - return dd_trace_context, trace_context_source + return dd_trace_context, trace_context_source, event_source def get_dd_trace_context(): @@ -537,8 +546,9 @@ def set_dd_trace_py_root(trace_context_source, merge_xray_traces): ) -def create_inferred_span(event, context): - event_source = parse_event_source(event) +def create_inferred_span(event, context, event_source: _EventSource = None): + if event_source is None: + event_source = parse_event_source(event) try: if event_source.equals( EventTypes.API_GATEWAY, subtype=EventSubtypes.API_GATEWAY @@ -649,20 +659,6 @@ def insert_upstream_authorizer_span( return upstream_authorizer_span -def get_start_and_finish_time_for_authorizer_span(request_context): - parsed_upstream_context = json.loads( - request_context.get("authorizer", {}).get("_datadog") - ) - start_time_s = ( - int(parsed_upstream_context.get(OtherConsts.parentSpanFinishTimeHeader)) / 1000 - ) - finish_time_s = ( - int(request_context.get("requestTimeEpoch")) - + int(request_context.get("authorizer", {}).get("integrationLatency", 0)) - ) / 1000 - return start_time_s, finish_time_s - - def create_inferred_span_from_api_gateway_websocket_event(event, context): trace_ctx = tracer.current_trace_context() request_context = event.get("requestContext") @@ -681,7 +677,7 @@ def create_inferred_span_from_api_gateway_websocket_event(event, context): "event_type": request_context.get("eventType"), "message_direction": request_context.get("messageDirection"), } - request_time_epoch = request_context.get("requestTimeEpoch") + request_time_epoch_s = int(request_context.get("requestTimeEpoch")) / 1000 if is_api_gateway_invocation_async(event): InferredSpanInfo.set_tags(tags, tag_source="self", synchronicity="async") else: @@ -693,14 +689,18 @@ def create_inferred_span_from_api_gateway_websocket_event(event, context): } tracer.set_tags({"_dd.origin": "lambda"}) upstream_authorizer_span = None - if is_non_cached_authorizer_invocation(event): + injected_authorizer_data = get_injected_authorizer_data(event, _EventSource(EventTypes.API_GATEWAY, EventSubtypes.WEBSOCKET)) + if injected_authorizer_data: try: - start_time_s, finish_time_s = get_start_and_finish_time_for_authorizer_span( - request_context - ) + start_time_s = int(injected_authorizer_data.get(OtherConsts.parentSpanFinishTimeHeader)) / 1000 + finish_time_s = request_time_epoch_s + \ + (int(request_context.get("authorizer", {}).get("integrationLatency", 0)) + ) / 1000 upstream_authorizer_span = insert_upstream_authorizer_span( args, tags, start_time_s, finish_time_s ) + # trace context needs to be set again as it is reset by upstream_authorizer_span.finish + tracer.context_provider.activate(trace_ctx) except Exception as e: traceback.print_exc() logger.debug( @@ -714,7 +714,7 @@ def create_inferred_span_from_api_gateway_websocket_event(event, context): span = tracer.trace("aws.apigateway.websocket", **args) if span: span.set_tags(tags) - span.start = request_time_epoch / 1000 + span.start = finish_time_s if finish_time_s is not None else request_time_epoch_s if upstream_authorizer_span: span.parent_id = upstream_authorizer_span.span_id return span @@ -738,7 +738,7 @@ def create_inferred_span_from_api_gateway_event(event, context): "stage": request_context.get("stage"), "request_id": request_context.get("requestId"), } - request_time_epoch = int(request_context.get("requestTimeEpoch")) / 1000 + request_time_epoch_s = int(request_context.get("requestTimeEpoch")) / 1000 if is_api_gateway_invocation_async(event): InferredSpanInfo.set_tags(tags, tag_source="self", synchronicity="async") else: @@ -751,14 +751,17 @@ def create_inferred_span_from_api_gateway_event(event, context): tracer.set_tags({"_dd.origin": "lambda"}) upstream_authorizer_span = None finish_time_s = None - if is_non_cached_authorizer_invocation(event): + injected_authorizer_data = get_injected_authorizer_data(event, _EventSource(EventTypes.API_GATEWAY, EventSubtypes.API_GATEWAY)) + if injected_authorizer_data: try: - start_time_s, finish_time_s = get_start_and_finish_time_for_authorizer_span( - request_context - ) + start_time_s = int(injected_authorizer_data.get(OtherConsts.parentSpanFinishTimeHeader)) / 1000 + finish_time_s = request_time_epoch_s + \ + (int(request_context.get("authorizer", {}).get("integrationLatency", 0))) / 1000 upstream_authorizer_span = insert_upstream_authorizer_span( args, tags, start_time_s, finish_time_s ) + # trace context needs to be set again as it is reset by upstream_authorizer_span.finish + tracer.context_provider.activate(trace_ctx) except Exception as e: traceback.print_exc() logger.debug( @@ -767,13 +770,12 @@ def create_inferred_span_from_api_gateway_event(event, context): event, e, ) - # trace context needs to be set again as it is reset by upstream_authorizer_span.finish - tracer.context_provider.activate(trace_ctx) + span = tracer.trace("aws.apigateway", **args) if span: span.set_tags(tags) # start time pushed by the inserted authorizer span - span.start = finish_time_s if finish_time_s is not None else request_time_epoch + span.start = finish_time_s if finish_time_s is not None else request_time_epoch_s if upstream_authorizer_span: span.parent_id = upstream_authorizer_span.span_id return span @@ -800,7 +802,7 @@ def create_inferred_span_from_http_api_event(event, context): "apiname": request_context.get("apiId"), "stage": request_context.get("stage"), } - request_time_epoch = request_context.get("timeEpoch") + request_time_epoch_s = int(request_context.get("timeEpoch")) / 1000 if is_api_gateway_invocation_async(event): InferredSpanInfo.set_tags(tags, tag_source="self", synchronicity="async") else: @@ -812,14 +814,17 @@ def create_inferred_span_from_http_api_event(event, context): } tracer.set_tags({"_dd.origin": "lambda"}) upstream_authorizer_span = None - if is_non_cached_authorizer_invocation(event): + finish_time_s = None + injected_authorizer_data = get_injected_authorizer_data(event, _EventSource(EventTypes.API_GATEWAY, EventSubtypes.HTTP_API)) + if injected_authorizer_data: try: - start_time_s, finish_time_s = get_start_and_finish_time_for_authorizer_span( - request_context - ) + start_time_s = int(injected_authorizer_data.get(OtherConsts.parentSpanFinishTimeHeader)) / 1000 + finish_time_s = start_time_s # we don't have the integrationLatency info for the authorizer upstream_authorizer_span = insert_upstream_authorizer_span( args, tags, start_time_s, finish_time_s ) + # trace context needs to be set again as it is reset by upstream_authorizer_span.finish + tracer.context_provider.activate(trace_ctx) except Exception as e: traceback.print_exc() logger.debug( @@ -828,12 +833,10 @@ def create_inferred_span_from_http_api_event(event, context): event, e, ) - # trace context needs to be set again as it is reset by upstream_authorizer_span.finish - tracer.context_provider.activate(trace_ctx) span = tracer.trace("aws.httpapi", **args) if span: span.set_tags(tags) - span.start = request_time_epoch / 1e3 + span.start = finish_time_s if finish_time_s is not None else request_time_epoch_s if upstream_authorizer_span: span.parent_id = upstream_authorizer_span.span_id return span diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 0e41a98d..4269275c 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -200,7 +200,7 @@ def _before(self, event, context): submit_invocations_metric(context) self.trigger_tags = extract_trigger_tags(event, context) # Extract Datadog trace context and source from incoming requests - dd_context, trace_context_source = extract_dd_trace_context( + dd_context, trace_context_source, event_source = extract_dd_trace_context( event, context, extractor=self.trace_extractor ) # Create a Datadog X-Ray subsegment with the trace context @@ -212,7 +212,7 @@ def _before(self, event, context): if dd_tracing_enabled: set_dd_trace_py_root(trace_context_source, self.merge_xray_traces) if self.make_inferred_span: - self.inferred_span = create_inferred_span(event, context) + self.inferred_span = create_inferred_span(event, context, event_source) self.span = create_function_execution_span( context, self.function_name, diff --git a/tests/event_samples/authorizer-request-api-gateway-v1.json b/tests/event_samples/authorizer-request-api-gateway-v1.json index fc16e7fc..16dfbd2a 100644 --- a/tests/event_samples/authorizer-request-api-gateway-v1.json +++ b/tests/event_samples/authorizer-request-api-gateway-v1.json @@ -70,7 +70,7 @@ "stage": "dev", "domainPrefix": "amddr1rix9", "requestTimeEpoch": 1663295019935, - "requestId": "e4816b6c-8954-4510-8771-837988d7f485", + "requestId": "1234567", "identity": { "cognitoIdentityPoolId": null, "accountId": null, diff --git a/tests/event_samples/authorizer-request-api-gateway-v2.json b/tests/event_samples/authorizer-request-api-gateway-v2.json new file mode 100644 index 00000000..09834155 --- /dev/null +++ b/tests/event_samples/authorizer-request-api-gateway-v2.json @@ -0,0 +1,46 @@ +{ + "version": "2.0", + "routeKey": "GET /hello", + "rawPath": "/hello", + "rawQueryString": "", + "headers": { + "accept": "*/*", + "accept-encoding": "gzip, deflate, br", + "authorization": "secretT0k3n", + "authorizationtoken": "secretT0k3n", + "cache-control": "no-cache", + "content-length": "0", + "host": "amddr1rix9.execute-api.sa-east-1.amazonaws.com", + "postman-token": "a9231288-1486-4753-a589-7ad3ac853c78", + "user-agent": "curl/7.64.1", + "x-amzn-trace-id": "Root=1-632a6082-6d32606a1e652ce274e8f055", + "x-forwarded-for": "38.122.226.210", + "x-forwarded-port": "443", + "x-forwarded-proto": "https" + }, + "requestContext": { + "accountId": "601427279990", + "apiId": "amddr1rix9", + "authorizer": { + "lambda": { + "_datadog": "{\"x-datadog-trace-id\": \"12056652649662348062\", \"x-datadog-parent-id\": \"12647339963998804799\", \"x-datadog-sampling-priority\": \"1\", \"x-datadog-tags\": \"_dd.p.dm=-0\", \"x-datadog-parent-span-finish-time\": 1663721602440.8286}", + "scope": "this is just a string" + } + }, + "domainName": "amddr1rix9.execute-api.sa-east-1.amazonaws.com", + "domainPrefix": "amddr1rix9", + "http": { + "method": "GET", + "path": "/hello", + "protocol": "HTTP/1.1", + "sourceIp": "38.122.226.210", + "userAgent": "curl/7.64.1" + }, + "requestId": "1234567", + "routeKey": "GET /hello", + "stage": "$default", + "time": "21/Sep/2022:00:53:22 +0000", + "timeEpoch": 1663295021832 + }, + "isBase64Encoded": false +} diff --git a/tests/event_samples/authorizer-token-api-gateway-v1.json b/tests/event_samples/authorizer-token-api-gateway-v1.json new file mode 100644 index 00000000..0c28772a --- /dev/null +++ b/tests/event_samples/authorizer-token-api-gateway-v1.json @@ -0,0 +1,93 @@ +{ + "resource": "/hello", + "path": "/hello", + "httpMethod": "GET", + "headers": { + "Accept": "*/*", + "Accept-Encoding": "gzip, deflate, br", + "Authorization": "secretT0k3n", + "authorizationToken": "secretT0k3n", + "Cache-Control": "no-cache", + "CloudFront-Forwarded-Proto": "https", + "CloudFront-Is-Desktop-Viewer": "true", + "CloudFront-Is-Mobile-Viewer": "false", + "CloudFront-Is-SmartTV-Viewer": "false", + "CloudFront-Is-Tablet-Viewer": "false", + "CloudFront-Viewer-ASN": "174", + "CloudFront-Viewer-Country": "US", + "Host": "amddr1rix9.execute-api.sa-east-1.amazonaws.com", + "Postman-Token": "791b2f79-1732-4877-abc1-5c538daf4642", + "User-Agent": "PostmanRuntime/7.29.2", + "Via": "1.1 d58463d219ef6ca0331e7200a6667c18.cloudfront.net (CloudFront)", + "X-Amz-Cf-Id": "tqU0SVyX_dKKI0XIN4-kwzCHrROBFECLXtAKCtjKr1uRQ-uijj99sQ==", + "X-Amzn-Trace-Id": "Root=1-632b4732-3e69c6a34f21f00d7f9b5599", + "X-Forwarded-For": "38.122.226.210, 64.252.135.71", + "X-Forwarded-Port": "443", + "X-Forwarded-Proto": "https" + }, + "multiValueHeaders": { + "Accept": ["*/*"], + "Accept-Encoding": ["gzip, deflate, br"], + "Authorization": ["secretT0k3n"], + "authorizationToken": ["secretT0k3n"], + "Cache-Control": ["no-cache"], + "CloudFront-Forwarded-Proto": ["https"], + "CloudFront-Is-Desktop-Viewer": ["true"], + "CloudFront-Is-Mobile-Viewer": ["false"], + "CloudFront-Is-SmartTV-Viewer": ["false"], + "CloudFront-Is-Tablet-Viewer": ["false"], + "CloudFront-Viewer-ASN": ["174"], + "CloudFront-Viewer-Country": ["US"], + "Host": ["amddr1rix9.execute-api.sa-east-1.amazonaws.com"], + "Postman-Token": ["791b2f79-1732-4877-abc1-5c538daf4642"], + "User-Agent": ["PostmanRuntime/7.29.2"], + "Via": ["1.1 d58463d219ef6ca0331e7200a6667c18.cloudfront.net (CloudFront)"], + "X-Amz-Cf-Id": ["tqU0SVyX_dKKI0XIN4-kwzCHrROBFECLXtAKCtjKr1uRQ-uijj99sQ=="], + "X-Amzn-Trace-Id": ["Root=1-632b4732-3e69c6a34f21f00d7f9b5599"], + "X-Forwarded-For": ["38.122.226.210, 64.252.135.71"], + "X-Forwarded-Port": ["443"], + "X-Forwarded-Proto": ["https"] + }, + "queryStringParameters": null, + "multiValueQueryStringParameters": null, + "pathParameters": null, + "stageVariables": null, + "requestContext": { + "resourceId": "elem3u", + "authorizer": { + "_datadog": "{\"x-datadog-trace-id\": \"17874798268144902712\", \"x-datadog-parent-id\": \"16184667399315372101\", \"x-datadog-sampling-priority\": \"1\", \"x-datadog-tags\": \"_dd.p.dm=-0\", \"x-datadog-parent-span-finish-time\": 1663721602440.8286}", + "scope": "this is just a string", + "principalId": "foo", + "integrationLatency": 1897 + }, + "resourcePath": "/hello", + "httpMethod": "GET", + "extendedRequestId": "Y0gP8FG9mjQFZTQ=", + "requestTime": "21/Sep/2022:17:17:38 +0000", + "path": "/dev/hello", + "accountId": "601427279990", + "protocol": "HTTP/1.1", + "stage": "dev", + "domainPrefix": "amddr1rix9", + "requestTimeEpoch": 1663295019935, + "requestId": "1234567", + "identity": { + "cognitoIdentityPoolId": null, + "accountId": null, + "cognitoIdentityId": null, + "caller": null, + "sourceIp": "38.122.226.210", + "principalOrgId": null, + "accessKey": null, + "cognitoAuthenticationType": null, + "cognitoAuthenticationProvider": null, + "userArn": null, + "userAgent": "PostmanRuntime/7.29.2", + "user": null + }, + "domainName": "amddr1rix9.execute-api.sa-east-1.amazonaws.com", + "apiId": "amddr1rix9" + }, + "body": null, + "isBase64Encoded": false +} diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 06e299e0..ead79940 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -29,6 +29,7 @@ InferredSpanInfo, extract_context_from_eventbridge_event, ) +from datadog_lambda.trigger import EventTypes function_arn = "arn:aws:lambda:us-west-1:123457598159:function:python-layer-test" @@ -84,7 +85,7 @@ def tearDown(self): def test_without_datadog_trace_headers(self): lambda_ctx = get_mock_context() - ctx, source = extract_dd_trace_context({}, lambda_ctx) + ctx, source, event_source = extract_dd_trace_context({}, lambda_ctx) self.assertEqual(source, "xray") self.assertDictEqual( ctx, @@ -106,7 +107,7 @@ def test_without_datadog_trace_headers(self): def test_with_non_object_event(self): lambda_ctx = get_mock_context() - ctx, source = extract_dd_trace_context(b"", lambda_ctx) + ctx, source, event_source = extract_dd_trace_context(b"", lambda_ctx) self.assertEqual(source, "xray") self.assertDictEqual( ctx, @@ -128,7 +129,7 @@ def test_with_non_object_event(self): def test_with_incomplete_datadog_trace_headers(self): lambda_ctx = get_mock_context() - ctx, source = extract_dd_trace_context( + ctx, source, event_source = extract_dd_trace_context( {"headers": {TraceHeader.TRACE_ID: "123", TraceHeader.PARENT_ID: "321"}}, lambda_ctx, ) @@ -152,7 +153,7 @@ def test_with_incomplete_datadog_trace_headers(self): def test_with_complete_datadog_trace_headers(self): lambda_ctx = get_mock_context() - ctx, source = extract_dd_trace_context( + ctx, source, event_source = extract_dd_trace_context( { "headers": { TraceHeader.TRACE_ID: "123", @@ -193,7 +194,7 @@ def extractor_foo(event, context): return trace_id, parent_id, sampling_priority lambda_ctx = get_mock_context() - ctx, ctx_source = extract_dd_trace_context( + ctx, ctx_source, event_source = extract_dd_trace_context( { "foo": { TraceHeader.TRACE_ID: "123", @@ -227,7 +228,7 @@ def extractor_raiser(event, context): raise Exception("kreator") lambda_ctx = get_mock_context() - ctx, ctx_source = extract_dd_trace_context( + ctx, ctx_source, event_source = extract_dd_trace_context( { "foo": { TraceHeader.TRACE_ID: "123", @@ -289,7 +290,7 @@ def test_with_sqs_distributed_datadog_trace_data(self): } ] } - ctx, source = extract_dd_trace_context(sqs_event, lambda_ctx) + ctx, source, event_source = extract_dd_trace_context(sqs_event, lambda_ctx) self.assertEqual(source, "event") self.assertDictEqual( ctx, @@ -323,7 +324,7 @@ def test_with_legacy_client_context_datadog_trace_data(self): } } ) - ctx, source = extract_dd_trace_context({}, lambda_ctx) + ctx, source, event_source = extract_dd_trace_context({}, lambda_ctx) self.assertEqual(source, "event") self.assertDictEqual( ctx, @@ -356,7 +357,7 @@ def test_with_new_client_context_datadog_trace_data(self): TraceHeader.SAMPLING_PRIORITY: "1", } ) - ctx, source = extract_dd_trace_context({}, lambda_ctx) + ctx, source, event_source = extract_dd_trace_context({}, lambda_ctx) self.assertEqual(source, "event") self.assertDictEqual( ctx, @@ -556,7 +557,7 @@ def test_mixed_parent_context_when_merging(self): # use the dd-trace trace-id and the x-ray parent-id # This allows parenting relationships like dd-trace -> x-ray -> dd-trace lambda_ctx = get_mock_context() - ctx, source = extract_dd_trace_context( + ctx, source, event_type = extract_dd_trace_context( { "headers": { TraceHeader.TRACE_ID: "123", @@ -587,13 +588,51 @@ def setUp(self): def test_create_inferred_span_from_authorizer_request_api_gateway_v1_event(self): event_sample_source = "authorizer-request-api-gateway-v1" + span = self._http_common_testing_items(event_sample_source, "aws.apigateway.rest", 1663295021.832) + + # api-gateway-v1 specific checks: + self.assertEqual( + span.start, 1663295021.832 + ) # request_time_epoch + integrationLatency in api-gateway-v1 case + self.assertEqual(span.get_tag("apiid"), "amddr1rix9") + self.assertEqual(span.get_tag("apiname"), "amddr1rix9") + self.assertEqual(span.get_tag("stage"), "dev") + + def test_create_inferred_span_from_authorizer_token_api_gateway_v1_event(self): + event_sample_source = "authorizer-token-api-gateway-v1" + span = self._http_common_testing_items(event_sample_source, "aws.apigateway.rest", 1663295021.832) + + # api-gateway-v1 specific checks: + self.assertEqual( + span.start, 1663295021.832 + ) # request_time_epoch + integrationLatency in api-gateway-v1 case + self.assertEqual(span.get_tag("apiid"), "amddr1rix9") + self.assertEqual(span.get_tag("apiname"), "amddr1rix9") + self.assertEqual(span.get_tag("stage"), "dev") + + def test_create_inferred_span_from_authorizer_request_api_gateway_v2_event(self): + event_sample_source = "authorizer-request-api-gateway-v2" + span = self._http_common_testing_items(event_sample_source, "aws.httpapi", 1663721602.44) + + # http-api specific checks: + self.assertEqual(span.get_tag("apiid"), "amddr1rix9") + self.assertEqual(span.get_tag("apiname"), "amddr1rix9") + self.assertEqual(span.get_tag("stage"), "$default") + self.assertEqual(span.get_tag("http.protocol"), "HTTP/1.1") + self.assertEqual(span.get_tag("http.source_ip"), "38.122.226.210") + self.assertEqual(span.get_tag("http.user_agent"), "curl/7.64.1") + self.assertEqual( + span.start, 1663721602.44 + ) # use the injected parent span finish time as an approximation + + def _http_common_testing_items(self, event_sample_source, operation_name, finish_time): test_file = event_samples + event_sample_source + ".json" with open(test_file, "r") as event: event = json.load(event) ctx = get_mock_context() - ctx.aws_request_id = "123" + ctx.aws_request_id = "1234567" span = create_inferred_span(event, ctx) - self.assertEqual(span.get_tag("operation_name"), "aws.apigateway.rest") + self.assertEqual(span.get_tag("operation_name"), operation_name) self.assertEqual( span.service, "amddr1rix9.execute-api.sa-east-1.amazonaws.com", @@ -609,14 +648,9 @@ def test_create_inferred_span_from_authorizer_request_api_gateway_v1_event(self) "GET /hello", ) self.assertEqual( - span.get_tag("request_id"), "e4816b6c-8954-4510-8771-837988d7f485" + span.get_tag("request_id"), "1234567" ) - self.assertEqual(span.get_tag("apiid"), "amddr1rix9") - self.assertEqual(span.get_tag("apiname"), "amddr1rix9") - self.assertEqual(span.get_tag("stage"), "dev") - self.assertEqual( - span.start, 1663295021.832 - ) # request_time_epoch + integrationLatency + self.assertEqual(span.span_type, "http") self.assertEqual(span.get_tag(InferredSpanInfo.TAG_SOURCE), "self") self.assertEqual(span.get_tag(InferredSpanInfo.SYNCHRONICITY), "sync") @@ -624,10 +658,11 @@ def test_create_inferred_span_from_authorizer_request_api_gateway_v1_event(self) # checking the upstream_authorizer_span self.mock_span_stop.assert_called_once() args, kwargs = self.mock_span_stop.call_args_list[0] - self.assertEqual(kwargs["finish_time"], 1663295021.832) + self.assertEqual(kwargs["finish_time"], finish_time) authorizer_span = args[0] self.assertEqual(authorizer_span.name, "aws.apigateway.authorizer") self.assertEqual(span.parent_id, authorizer_span.span_id) + return span class TestInferredSpans(unittest.TestCase): @@ -1151,7 +1186,7 @@ def test_extract_dd_trace_context_for_eventbridge(self): with open(test_file, "r") as event: event = json.load(event) ctx = get_mock_context() - context, source = extract_dd_trace_context(event, ctx) + context, source, event_type = extract_dd_trace_context(event, ctx) self.assertEqual(context["trace-id"], "12345") self.assertEqual(context["parent-id"], "67890") @@ -1161,7 +1196,7 @@ def test_extract_context_from_sqs_event_with_string_msg_attr(self): with open(test_file, "r") as event: event = json.load(event) ctx = get_mock_context() - context, source = extract_dd_trace_context(event, ctx) + context, source, event_type = extract_dd_trace_context(event, ctx) self.assertEqual(context["trace-id"], "2684756524522091840") self.assertEqual(context["parent-id"], "7431398482019833808") self.assertEqual(context["sampling-priority"], "1") @@ -1172,7 +1207,7 @@ def test_extract_context_from_sqs_batch_event(self): with open(test_file, "r") as event: event = json.load(event) ctx = get_mock_context() - context, source = extract_dd_trace_context(event, ctx) + context, source, event_source = extract_dd_trace_context(event, ctx) self.assertEqual(context["trace-id"], "2684756524522091840") self.assertEqual(context["parent-id"], "7431398482019833808") self.assertEqual(context["sampling-priority"], "1") @@ -1183,7 +1218,7 @@ def test_extract_context_from_sns_event_with_string_msg_attr(self): with open(test_file, "r") as event: event = json.load(event) ctx = get_mock_context() - context, source = extract_dd_trace_context(event, ctx) + context, source, event_source = extract_dd_trace_context(event, ctx) self.assertEqual(context["trace-id"], "4948377316357291421") self.assertEqual(context["parent-id"], "6746998015037429512") self.assertEqual(context["sampling-priority"], "1") @@ -1194,7 +1229,7 @@ def test_extract_context_from_sns_event_with_b64_msg_attr(self): with open(test_file, "r") as event: event = json.load(event) ctx = get_mock_context() - context, source = extract_dd_trace_context(event, ctx) + context, source, event_source = extract_dd_trace_context(event, ctx) self.assertEqual(context["trace-id"], "4948377316357291421") self.assertEqual(context["parent-id"], "6746998015037429512") self.assertEqual(context["sampling-priority"], "1") @@ -1205,7 +1240,7 @@ def test_extract_context_from_sns_batch_event(self): with open(test_file, "r") as event: event = json.load(event) ctx = get_mock_context() - context, source = extract_dd_trace_context(event, ctx) + context, source, event_source = extract_dd_trace_context(event, ctx) self.assertEqual(context["trace-id"], "4948377316357291421") self.assertEqual(context["parent-id"], "6746998015037429512") self.assertEqual(context["sampling-priority"], "1") @@ -1216,7 +1251,7 @@ def test_extract_context_from_kinesis_event(self): with open(test_file, "r") as event: event = json.load(event) ctx = get_mock_context() - context, source = extract_dd_trace_context(event, ctx) + context, source, event_source = extract_dd_trace_context(event, ctx) self.assertEqual(context["trace-id"], "4948377316357291421") self.assertEqual(context["parent-id"], "2876253380018681026") self.assertEqual(context["sampling-priority"], "1") @@ -1227,7 +1262,7 @@ def test_extract_context_from_kinesis_batch_event(self): with open(test_file, "r") as event: event = json.load(event) ctx = get_mock_context() - context, source = extract_dd_trace_context(event, ctx) + context, source, event_source = extract_dd_trace_context(event, ctx) self.assertEqual(context["trace-id"], "4948377316357291421") self.assertEqual(context["parent-id"], "2876253380018681026") self.assertEqual(context["sampling-priority"], "1") diff --git a/tests/test_wrapper.py b/tests/test_wrapper.py index a78cddfc..6240ec4f 100644 --- a/tests/test_wrapper.py +++ b/tests/test_wrapper.py @@ -38,7 +38,7 @@ def setUp(self): patcher = patch("datadog_lambda.wrapper.extract_dd_trace_context") self.mock_extract_dd_trace_context = patcher.start() - self.mock_extract_dd_trace_context.return_value = ({}, None) + self.mock_extract_dd_trace_context.return_value = ({}, None, None) self.addCleanup(patcher.stop) patcher = patch("datadog_lambda.wrapper.set_correlation_ids") From 5313cc650736b0ba7d094b314293787becd7f2ea Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Wed, 21 Sep 2022 13:55:42 -0400 Subject: [PATCH 07/26] formatting fixes --- .vscode/settings.json | 5 ++- datadog_lambda/tracing.py | 94 ++++++++++++++++++++++++++++++--------- datadog_lambda/wrapper.py | 4 +- tests/test_tracing.py | 20 ++++++--- 4 files changed, 93 insertions(+), 30 deletions(-) diff --git a/.vscode/settings.json b/.vscode/settings.json index db92ccf3..e58c4de2 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -1,4 +1,7 @@ { "python.pythonPath": "/usr/local/bin/python3", - "python.formatting.provider": "black" + "python.formatting.provider": "black", + "python.testing.unittestArgs": ["-v", "-s", "./tests", "-p", "test*.py"], + "python.testing.pytestEnabled": false, + "python.testing.unittestEnabled": true } diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index ef7a7913..2de3352e 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -170,7 +170,9 @@ def extract_context_from_lambda_context(lambda_context): return trace_id, parent_id, sampling_priority -def extract_context_from_http_event_or_context(event, lambda_context, event_source: _EventSource): +def extract_context_from_http_event_or_context( + event, lambda_context, event_source: _EventSource +): """ Extract Datadog trace context from the `headers` key in from the Lambda `event` object. @@ -340,10 +342,13 @@ def get_injected_authorizer_data(event, event_source: _EventSource) -> dict: # has '_datadog' header and not cached because integrationLatency > 0 authorizer_headers = event.get("requestContext", {}).get("authorizer", {}) - # [API_GATEWAY V1]if integrationLatency == 0, it's cached and no trace for the authorizer func - if event_source.equals( - EventTypes.API_GATEWAY, subtype=EventSubtypes.API_GATEWAY - ) and int(authorizer_headers.get("integrationLatency", 0)) == 0: + # [API_GATEWAY V1]if integrationLatency == 0, it's cached and no trace for the authorizer + if ( + event_source.equals( + EventTypes.API_GATEWAY, subtype=EventSubtypes.API_GATEWAY + ) + and int(authorizer_headers.get("integrationLatency", 0)) == 0 + ): return None 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): trace_id, parent_id, sampling_priority, - ) = extract_context_from_http_event_or_context(event, lambda_context, event_source) + ) = extract_context_from_http_event_or_context( + event, lambda_context, event_source + ) elif event_source.equals(EventTypes.SNS) or event_source.equals(EventTypes.SQS): ( trace_id, @@ -689,13 +696,29 @@ def create_inferred_span_from_api_gateway_websocket_event(event, context): } tracer.set_tags({"_dd.origin": "lambda"}) upstream_authorizer_span = None - injected_authorizer_data = get_injected_authorizer_data(event, _EventSource(EventTypes.API_GATEWAY, EventSubtypes.WEBSOCKET)) + finish_time_s = None + injected_authorizer_data = get_injected_authorizer_data( + event, _EventSource(EventTypes.API_GATEWAY, EventSubtypes.WEBSOCKET) + ) if injected_authorizer_data: try: - start_time_s = int(injected_authorizer_data.get(OtherConsts.parentSpanFinishTimeHeader)) / 1000 - finish_time_s = request_time_epoch_s + \ - (int(request_context.get("authorizer", {}).get("integrationLatency", 0)) - ) / 1000 + start_time_s = ( + int( + injected_authorizer_data.get(OtherConsts.parentSpanFinishTimeHeader) + ) + / 1000 + ) + finish_time_s = ( + request_time_epoch_s + + ( + int( + request_context.get("authorizer", {}).get( + "integrationLatency", 0 + ) + ) + ) + / 1000 + ) upstream_authorizer_span = insert_upstream_authorizer_span( args, tags, start_time_s, finish_time_s ) @@ -714,7 +737,9 @@ def create_inferred_span_from_api_gateway_websocket_event(event, context): span = tracer.trace("aws.apigateway.websocket", **args) if span: span.set_tags(tags) - span.start = finish_time_s if finish_time_s is not None else request_time_epoch_s + span.start = ( + finish_time_s if finish_time_s is not None else request_time_epoch_s + ) if upstream_authorizer_span: span.parent_id = upstream_authorizer_span.span_id return span @@ -751,12 +776,28 @@ def create_inferred_span_from_api_gateway_event(event, context): tracer.set_tags({"_dd.origin": "lambda"}) upstream_authorizer_span = None finish_time_s = None - injected_authorizer_data = get_injected_authorizer_data(event, _EventSource(EventTypes.API_GATEWAY, EventSubtypes.API_GATEWAY)) + injected_authorizer_data = get_injected_authorizer_data( + event, _EventSource(EventTypes.API_GATEWAY, EventSubtypes.API_GATEWAY) + ) if injected_authorizer_data: try: - start_time_s = int(injected_authorizer_data.get(OtherConsts.parentSpanFinishTimeHeader)) / 1000 - finish_time_s = request_time_epoch_s + \ - (int(request_context.get("authorizer", {}).get("integrationLatency", 0))) / 1000 + start_time_s = ( + int( + injected_authorizer_data.get(OtherConsts.parentSpanFinishTimeHeader) + ) + / 1000 + ) + finish_time_s = ( + request_time_epoch_s + + ( + int( + request_context.get("authorizer", {}).get( + "integrationLatency", 0 + ) + ) + ) + / 1000 + ) upstream_authorizer_span = insert_upstream_authorizer_span( args, tags, start_time_s, finish_time_s ) @@ -775,7 +816,9 @@ def create_inferred_span_from_api_gateway_event(event, context): if span: span.set_tags(tags) # start time pushed by the inserted authorizer span - span.start = finish_time_s if finish_time_s is not None else request_time_epoch_s + span.start = ( + finish_time_s if finish_time_s is not None else request_time_epoch_s + ) if upstream_authorizer_span: span.parent_id = upstream_authorizer_span.span_id return span @@ -815,11 +858,18 @@ def create_inferred_span_from_http_api_event(event, context): tracer.set_tags({"_dd.origin": "lambda"}) upstream_authorizer_span = None finish_time_s = None - injected_authorizer_data = get_injected_authorizer_data(event, _EventSource(EventTypes.API_GATEWAY, EventSubtypes.HTTP_API)) + injected_authorizer_data = get_injected_authorizer_data( + event, _EventSource(EventTypes.API_GATEWAY, EventSubtypes.HTTP_API) + ) if injected_authorizer_data: try: - start_time_s = int(injected_authorizer_data.get(OtherConsts.parentSpanFinishTimeHeader)) / 1000 - finish_time_s = start_time_s # we don't have the integrationLatency info for the authorizer + start_time_s = ( + int( + injected_authorizer_data.get(OtherConsts.parentSpanFinishTimeHeader) + ) + / 1000 + ) + finish_time_s = start_time_s # no integrationLatency info in this case upstream_authorizer_span = insert_upstream_authorizer_span( args, tags, start_time_s, finish_time_s ) @@ -836,7 +886,9 @@ def create_inferred_span_from_http_api_event(event, context): span = tracer.trace("aws.httpapi", **args) if span: span.set_tags(tags) - span.start = finish_time_s if finish_time_s is not None else request_time_epoch_s + span.start = ( + finish_time_s if finish_time_s is not None else request_time_epoch_s + ) if upstream_authorizer_span: span.parent_id = upstream_authorizer_span.span_id return span diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 4269275c..3ab09147 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -212,7 +212,9 @@ def _before(self, event, context): if dd_tracing_enabled: set_dd_trace_py_root(trace_context_source, self.merge_xray_traces) if self.make_inferred_span: - self.inferred_span = create_inferred_span(event, context, event_source) + self.inferred_span = create_inferred_span( + event, context, event_source + ) self.span = create_function_execution_span( context, self.function_name, diff --git a/tests/test_tracing.py b/tests/test_tracing.py index ead79940..30c2b3f3 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -588,7 +588,9 @@ def setUp(self): def test_create_inferred_span_from_authorizer_request_api_gateway_v1_event(self): event_sample_source = "authorizer-request-api-gateway-v1" - span = self._http_common_testing_items(event_sample_source, "aws.apigateway.rest", 1663295021.832) + span = self._http_common_testing_items( + event_sample_source, "aws.apigateway.rest", 1663295021.832 + ) # api-gateway-v1 specific checks: self.assertEqual( @@ -600,7 +602,9 @@ def test_create_inferred_span_from_authorizer_request_api_gateway_v1_event(self) def test_create_inferred_span_from_authorizer_token_api_gateway_v1_event(self): event_sample_source = "authorizer-token-api-gateway-v1" - span = self._http_common_testing_items(event_sample_source, "aws.apigateway.rest", 1663295021.832) + span = self._http_common_testing_items( + event_sample_source, "aws.apigateway.rest", 1663295021.832 + ) # api-gateway-v1 specific checks: self.assertEqual( @@ -612,7 +616,9 @@ def test_create_inferred_span_from_authorizer_token_api_gateway_v1_event(self): def test_create_inferred_span_from_authorizer_request_api_gateway_v2_event(self): event_sample_source = "authorizer-request-api-gateway-v2" - span = self._http_common_testing_items(event_sample_source, "aws.httpapi", 1663721602.44) + span = self._http_common_testing_items( + event_sample_source, "aws.httpapi", 1663721602.44 + ) # http-api specific checks: self.assertEqual(span.get_tag("apiid"), "amddr1rix9") @@ -625,7 +631,9 @@ def test_create_inferred_span_from_authorizer_request_api_gateway_v2_event(self) span.start, 1663721602.44 ) # use the injected parent span finish time as an approximation - def _http_common_testing_items(self, event_sample_source, operation_name, finish_time): + def _http_common_testing_items( + self, event_sample_source, operation_name, finish_time + ): test_file = event_samples + event_sample_source + ".json" with open(test_file, "r") as event: event = json.load(event) @@ -647,9 +655,7 @@ def _http_common_testing_items(self, event_sample_source, operation_name, finish span.get_tag("resource_names"), "GET /hello", ) - self.assertEqual( - span.get_tag("request_id"), "1234567" - ) + self.assertEqual(span.get_tag("request_id"), "1234567") self.assertEqual(span.span_type, "http") self.assertEqual(span.get_tag(InferredSpanInfo.TAG_SOURCE), "self") From 4a915522d3079f8c846e732f98c6f407769ef9ba Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Mon, 26 Sep 2022 20:45:32 -0400 Subject: [PATCH 08/26] fixing http api caching cases --- datadog_lambda/constants.py | 1 + datadog_lambda/tracing.py | 14 +++-- datadog_lambda/wrapper.py | 58 ++++++++++++------- .../authorizer-request-api-gateway-v2.json | 11 ++-- tests/test_tracing.py | 4 +- 5 files changed, 56 insertions(+), 32 deletions(-) diff --git a/datadog_lambda/constants.py b/datadog_lambda/constants.py index 61945f42..acc6a75c 100644 --- a/datadog_lambda/constants.py +++ b/datadog_lambda/constants.py @@ -46,3 +46,4 @@ class XrayDaemon(object): class OtherConsts(object): parentSpanFinishTimeHeader = "x-datadog-parent-span-finish-time" + originalAuthorizerRequestEpoch = "x-datadog-original-epoch" diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 2de3352e..98fa66e5 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -353,13 +353,17 @@ def get_injected_authorizer_data(event, event_source: _EventSource) -> dict: if event_source.equals(EventTypes.API_GATEWAY, subtype=EventSubtypes.HTTP_API): dd_data_raw = authorizer_headers.get("lambda", {}).get("_datadog") + dd_data = json.loads(base64.b64decode(dd_data_raw)) + # [API_GATEWAY V2]if original authorizer request epoch is different, then it's cached + if event.get("requestContext", {}).get("timeEpoch") == dd_data.get( + OtherConsts.originalAuthorizerRequestEpoch + ): + return dd_data + else: + return None else: dd_data_raw = authorizer_headers.get("_datadog") - - if dd_data_raw: - dd_data = json.loads(dd_data_raw) - return dd_data - return None + return json.loads(dd_data_raw) except Exception as e: logger.debug("Failed to check if invocated by an authorizer. error %s", e) diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 3ab09147..02d2485c 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -3,6 +3,7 @@ # This product includes software developed at Datadog (https://www.datadoghq.com/). # Copyright 2019 Datadog, Inc. +import base64 import os import logging import traceback @@ -37,7 +38,12 @@ create_inferred_span, InferredSpanInfo, ) -from datadog_lambda.trigger import extract_trigger_tags, extract_http_status_code_tag +from datadog_lambda.trigger import ( + extract_trigger_tags, + extract_http_status_code_tag, + EventTypes, + EventSubtypes, +) from datadog_lambda.tag_object import tag_object profiling_env_var = os.environ.get("DD_PROFILING_ENABLED", "false").lower() == "true" @@ -155,21 +161,7 @@ def __call__(self, event, context, **kwargs): self._before(event, context) try: self.response = self.func(event, context, **kwargs) - try: - if ( - self.make_authorizer_span - and self.response - and self.response.get("principalId") - and self.response.get("policyDocument") - ): - self._inject_authorizer_span_headers() - except Exception as e: - traceback.print_exc() - logger.debug( - "Unable to inject the authorizer headers. \ - Continue to return the original response. Reason: %s", - e, - ) + self._ending(event) return self.response except Exception: submit_errors_metric(context) @@ -179,19 +171,31 @@ def __call__(self, event, context, **kwargs): finally: self._after(event, context) - def _inject_authorizer_span_headers(self): + def _inject_authorizer_span_headers(self, event): finish_time_ns = ( self.span.start_ns if InferredSpanInfo.is_async(self.inferred_span) and self.span else time_ns() ) - self.response.setdefault("context", {}) - self.response["context"].setdefault("_datadog", {}) injected_headers = {} source_span = self.inferred_span if self.inferred_span else self.span HTTPPropagator.inject(source_span.context, injected_headers) injected_headers[OtherConsts.parentSpanFinishTimeHeader] = finish_time_ns / 1e6 - self.response["context"]["_datadog"] = json.dumps(injected_headers) + encode_datadog_data = False + if self.event_source and self.event_source.equals( + EventTypes.API_GATEWAY, EventSubtypes.HTTP_API + ): + injected_headers[OtherConsts.originalAuthorizerRequestEpoch] = event.get( + "requestContext", {} + ).get("timeEpoch") + encode_datadog_data = True + datadog_data = json.dumps(injected_headers).encode() + if encode_datadog_data: + datadog_data = base64.b64encode( + datadog_data + ) # special characters can corrupt the response + self.response.setdefault("context", {}) + self.response["context"]["_datadog"] = datadog_data def _before(self, event, context): try: @@ -203,6 +207,7 @@ def _before(self, event, context): dd_context, trace_context_source, event_source = extract_dd_trace_context( event, context, extractor=self.trace_extractor ) + self.event_source = event_source # Create a Datadog X-Ray subsegment with the trace context if dd_context and trace_context_source == TraceContextSource.EVENT: create_dd_dummy_metadata_subsegment( @@ -232,6 +237,19 @@ def _before(self, event, context): except Exception: traceback.print_exc() + def _ending(self, event): + try: + if ( + self.make_authorizer_span + and self.response + and self.response.get("principalId") + and self.response.get("policyDocument") + and self.event_source + ): + self._inject_authorizer_span_headers(event) + except Exception: + traceback.print_exc() + def _after(self, event, context): try: status_code = extract_http_status_code_tag(self.trigger_tags, self.response) diff --git a/tests/event_samples/authorizer-request-api-gateway-v2.json b/tests/event_samples/authorizer-request-api-gateway-v2.json index 09834155..4beba794 100644 --- a/tests/event_samples/authorizer-request-api-gateway-v2.json +++ b/tests/event_samples/authorizer-request-api-gateway-v2.json @@ -11,9 +11,10 @@ "cache-control": "no-cache", "content-length": "0", "host": "amddr1rix9.execute-api.sa-east-1.amazonaws.com", - "postman-token": "a9231288-1486-4753-a589-7ad3ac853c78", + "postman-token": "5e0efdb2-d93f-4d44-8792-45404d40dae2", + "user": "12345", "user-agent": "curl/7.64.1", - "x-amzn-trace-id": "Root=1-632a6082-6d32606a1e652ce274e8f055", + "x-amzn-trace-id": "Root=1-63321d1e-6b80ef0d2c2cf49600b9c28e", "x-forwarded-for": "38.122.226.210", "x-forwarded-port": "443", "x-forwarded-proto": "https" @@ -23,7 +24,7 @@ "apiId": "amddr1rix9", "authorizer": { "lambda": { - "_datadog": "{\"x-datadog-trace-id\": \"12056652649662348062\", \"x-datadog-parent-id\": \"12647339963998804799\", \"x-datadog-sampling-priority\": \"1\", \"x-datadog-tags\": \"_dd.p.dm=-0\", \"x-datadog-parent-span-finish-time\": 1663721602440.8286}", + "_datadog": "eyJ4LWRhdGFkb2ctdHJhY2UtaWQiOiAiMTQzNTY5ODM2MTk4NTI5MzMzNTQiLCAieC1kYXRhZG9nLXBhcmVudC1pZCI6ICIxMjY1ODYyMTA4MzUwNTQxMzgwOSIsICJ4LWRhdGFkb2ctc2FtcGxpbmctcHJpb3JpdHkiOiAiMSIsICJ4LWRhdGFkb2ctdGFncyI6ICJfZGQucC5kbT0tMCIsICJ4LWRhdGFkb2ctcGFyZW50LXNwYW4tZmluaXNoLXRpbWUiOiAxNjY0MjI4NjM5NTMzLjc3NTQsICJ4LWRhdGFkb2ctb3JpZ2luYWwtZXBvY2giOiAxNjY0MjI4NjM4MDg0fQ==", "scope": "this is just a string" } }, @@ -39,8 +40,8 @@ "requestId": "1234567", "routeKey": "GET /hello", "stage": "$default", - "time": "21/Sep/2022:00:53:22 +0000", - "timeEpoch": 1663295021832 + "time": "26/Sep/2022:21:43:58 +0000", + "timeEpoch": 1664228638084 }, "isBase64Encoded": false } diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 30c2b3f3..bb889b67 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -617,7 +617,7 @@ def test_create_inferred_span_from_authorizer_token_api_gateway_v1_event(self): def test_create_inferred_span_from_authorizer_request_api_gateway_v2_event(self): event_sample_source = "authorizer-request-api-gateway-v2" span = self._http_common_testing_items( - event_sample_source, "aws.httpapi", 1663721602.44 + event_sample_source, "aws.httpapi", 1664228639.533 ) # http-api specific checks: @@ -628,7 +628,7 @@ def test_create_inferred_span_from_authorizer_request_api_gateway_v2_event(self) self.assertEqual(span.get_tag("http.source_ip"), "38.122.226.210") self.assertEqual(span.get_tag("http.user_agent"), "curl/7.64.1") self.assertEqual( - span.start, 1663721602.44 + span.start, 1664228639.533 ) # use the injected parent span finish time as an approximation def _http_common_testing_items( From eeaadb3c069ec99cbbcee2d625c39290cf30c28e Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Mon, 3 Oct 2022 11:51:46 -0400 Subject: [PATCH 09/26] support websocket cases --- datadog_lambda/tracing.py | 6 +- .../authorizer-request-api-gateway-v2.json | 2 +- ...request-api-gateway-websocket-connect.json | 48 ++++++++ ...er-request-api-gateway-websocket-main.json | 27 +++++ tests/test_tracing.py | 106 +++++++++--------- 5 files changed, 130 insertions(+), 59 deletions(-) create mode 100644 tests/event_samples/authorizer-request-api-gateway-websocket-connect.json create mode 100644 tests/event_samples/authorizer-request-api-gateway-websocket-main.json diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 98fa66e5..b1a024ae 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -347,8 +347,10 @@ def get_injected_authorizer_data(event, event_source: _EventSource) -> dict: event_source.equals( EventTypes.API_GATEWAY, subtype=EventSubtypes.API_GATEWAY ) - and int(authorizer_headers.get("integrationLatency", 0)) == 0 - ): + or event_source.equals( + EventTypes.API_GATEWAY, subtype=EventSubtypes.WEBSOCKET + ) + ) and int(authorizer_headers.get("integrationLatency", 0)) == 0: return None if event_source.equals(EventTypes.API_GATEWAY, subtype=EventSubtypes.HTTP_API): diff --git a/tests/event_samples/authorizer-request-api-gateway-v2.json b/tests/event_samples/authorizer-request-api-gateway-v2.json index 4beba794..ac11c6ed 100644 --- a/tests/event_samples/authorizer-request-api-gateway-v2.json +++ b/tests/event_samples/authorizer-request-api-gateway-v2.json @@ -39,7 +39,7 @@ }, "requestId": "1234567", "routeKey": "GET /hello", - "stage": "$default", + "stage": "dev", "time": "26/Sep/2022:21:43:58 +0000", "timeEpoch": 1664228638084 }, diff --git a/tests/event_samples/authorizer-request-api-gateway-websocket-connect.json b/tests/event_samples/authorizer-request-api-gateway-websocket-connect.json new file mode 100644 index 00000000..a325a3f7 --- /dev/null +++ b/tests/event_samples/authorizer-request-api-gateway-websocket-connect.json @@ -0,0 +1,48 @@ +{ + "headers": { + "Auth": "secretT0k3n", + "Host": "amddr1rix9.execute-api.sa-east-1.amazonaws.com", + "Sec-WebSocket-Extensions": "permessage-deflate; client_max_window_bits", + "Sec-WebSocket-Key": "j9g8xXvniFqyyA+34vQ/Lg==", + "Sec-WebSocket-Version": "13", + "X-Amzn-Trace-Id": "Root=1-63348d21-727f66307696916d50f93f1b", + "X-Forwarded-For": "38.122.226.210", + "X-Forwarded-Port": "443", + "X-Forwarded-Proto": "https" + }, + "multiValueHeaders": { + "Auth": ["secretT0k3n"], + "Host": ["amddr1rix9.execute-api.sa-east-1.amazonaws.com"], + "Sec-WebSocket-Extensions": ["permessage-deflate; client_max_window_bits"], + "Sec-WebSocket-Key": ["j9g8xXvniFqyyA+34vQ/Lg=="], + "Sec-WebSocket-Version": ["13"], + "X-Amzn-Trace-Id": ["Root=1-63348d21-727f66307696916d50f93f1b"], + "X-Forwarded-For": ["38.122.226.210"], + "X-Forwarded-Port": ["443"], + "X-Forwarded-Proto": ["https"] + }, + "requestContext": { + "routeKey": "$connect", + "authorizer": { + "_datadog": "{\"x-datadog-trace-id\": \"5351047404834723189\", \"x-datadog-parent-id\": \"18230460631156161837\", \"x-datadog-sampling-priority\": \"1\", \"x-datadog-tags\": \"_dd.p.dm=-0\", \"x-datadog-parent-span-finish-time\": 1664388386890}", + "scope": "this is just a string", + "principalId": "foo", + "integrationLatency": 1440 + }, + "eventType": "CONNECT", + "extendedRequestId": "1234567", + "requestTime": "28/Sep/2022:18:06:25 +0000", + "messageDirection": "IN", + "stage": "dev", + "connectedAt": 1664388385450, + "requestTimeEpoch": 1664388385452, + "identity": { + "sourceIp": "38.122.226.210" + }, + "requestId": "1234567", + "domainName": "amddr1rix9.execute-api.sa-east-1.amazonaws.com", + "connectionId": "ZLr9QeNLmjQCIZA=", + "apiId": "amddr1rix9" + }, + "isBase64Encoded": false +} diff --git a/tests/event_samples/authorizer-request-api-gateway-websocket-main.json b/tests/event_samples/authorizer-request-api-gateway-websocket-main.json new file mode 100644 index 00000000..0cda1a6e --- /dev/null +++ b/tests/event_samples/authorizer-request-api-gateway-websocket-main.json @@ -0,0 +1,27 @@ +{ + "requestContext": { + "routeKey": "main", + "authorizer": { + "_datadog": "{\"x-datadog-trace-id\": \"17374069065025434920\", \"x-datadog-parent-id\": \"7915186751914176106\", \"x-datadog-sampling-priority\": \"1\", \"x-datadog-tags\": \"_dd.p.dm=-0\", \"x-datadog-parent-span-finish-time\": 1664390333526.5068}", + "scope": "this is just a string", + "principalId": "foo" + }, + "messageId": "ZLw3leP1GjQCI8Q=", + "eventType": "MESSAGE", + "extendedRequestId": "ZLw3lFowmjQFvSA=", + "requestTime": "28/Sep/2022:18:39:57 +0000", + "messageDirection": "IN", + "stage": "dev", + "connectedAt": 1664390332261, + "requestTimeEpoch": 1664390397117, + "identity": { + "sourceIp": "38.122.226.210" + }, + "requestId": "1234567", + "domainName": "amddr1rix9.execute-api.sa-east-1.amazonaws.com", + "connectionId": "ZLwtceO1mjQCI8Q=", + "apiId": "amddr1rix9" + }, + "body": "{\n \"action\":\"main\",\n \"message\":\"hi\"\n}", + "isBase64Encoded": false +} diff --git a/tests/test_tracing.py b/tests/test_tracing.py index bb889b67..d2909832 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -588,88 +588,82 @@ def setUp(self): def test_create_inferred_span_from_authorizer_request_api_gateway_v1_event(self): event_sample_source = "authorizer-request-api-gateway-v1" - span = self._http_common_testing_items( - event_sample_source, "aws.apigateway.rest", 1663295021.832 - ) - - # api-gateway-v1 specific checks: - self.assertEqual( - span.start, 1663295021.832 - ) # request_time_epoch + integrationLatency in api-gateway-v1 case - self.assertEqual(span.get_tag("apiid"), "amddr1rix9") - self.assertEqual(span.get_tag("apiname"), "amddr1rix9") - self.assertEqual(span.get_tag("stage"), "dev") + finish_time = 1663295021.832 #request_time_epoch + integrationLatency for api-gateway-v1 + span = self._authorizer_span_testing_items(event_sample_source, finish_time) + self._basic_common_checks(span, "aws.apigateway.rest") def test_create_inferred_span_from_authorizer_token_api_gateway_v1_event(self): event_sample_source = "authorizer-token-api-gateway-v1" - span = self._http_common_testing_items( - event_sample_source, "aws.apigateway.rest", 1663295021.832 - ) - - # api-gateway-v1 specific checks: - self.assertEqual( - span.start, 1663295021.832 - ) # request_time_epoch + integrationLatency in api-gateway-v1 case - self.assertEqual(span.get_tag("apiid"), "amddr1rix9") - self.assertEqual(span.get_tag("apiname"), "amddr1rix9") - self.assertEqual(span.get_tag("stage"), "dev") + finish_time = 1663295021.832 #request_time_epoch + integrationLatency for api-gateway-v1 + span = self._authorizer_span_testing_items(event_sample_source, finish_time) + self._basic_common_checks(span, "aws.apigateway.rest") def test_create_inferred_span_from_authorizer_request_api_gateway_v2_event(self): event_sample_source = "authorizer-request-api-gateway-v2" - span = self._http_common_testing_items( - event_sample_source, "aws.httpapi", 1664228639.533 - ) - - # http-api specific checks: - self.assertEqual(span.get_tag("apiid"), "amddr1rix9") - self.assertEqual(span.get_tag("apiname"), "amddr1rix9") - self.assertEqual(span.get_tag("stage"), "$default") - self.assertEqual(span.get_tag("http.protocol"), "HTTP/1.1") - self.assertEqual(span.get_tag("http.source_ip"), "38.122.226.210") - self.assertEqual(span.get_tag("http.user_agent"), "curl/7.64.1") - self.assertEqual( - span.start, 1664228639.533 - ) # use the injected parent span finish time as an approximation + finish_time = 1664228639.533 # use the injected parent span finish time as an approximation + span = self._authorizer_span_testing_items(event_sample_source, finish_time) + self._basic_common_checks(span, "aws.httpapi") + + def test_create_inferred_span_from_authorizer_request_api_gateway_websocket_connect_event(self): + event_sample_source = "authorizer-request-api-gateway-websocket-connect" + finish_time = 1664388386.892 # request_time_epoch + integrationLatency in websocket case + span = self._authorizer_span_testing_items(event_sample_source, finish_time) + self._basic_common_checks(span, "aws.apigateway.websocket", "web", "$connect", None) + + def test_create_inferred_span_from_authorizer_request_api_gateway_websocket_main_event(self): + event_sample_source = "authorizer-request-api-gateway-websocket-main" + test_file = event_samples + event_sample_source + ".json" + with open(test_file, "r") as event: + event = json.load(event) + ctx = get_mock_context() + ctx.aws_request_id = "1234567" + span = create_inferred_span(event, ctx) + self.mock_span_stop.assert_not_called() # NO authorizer span is injected + self._basic_common_checks(span, "aws.apigateway.websocket", "web", "main", None) - def _http_common_testing_items( - self, event_sample_source, operation_name, finish_time - ): + def _authorizer_span_testing_items(self, event_sample_source, finish_time): test_file = event_samples + event_sample_source + ".json" with open(test_file, "r") as event: event = json.load(event) ctx = get_mock_context() ctx.aws_request_id = "1234567" span = create_inferred_span(event, ctx) + self.assertEqual(span.get_tag(InferredSpanInfo.TAG_SOURCE), "self") + self.assertEqual(span.get_tag(InferredSpanInfo.SYNCHRONICITY), "sync") + + # checking the upstream_authorizer_span + self.mock_span_stop.assert_called_once() + args, kwargs = self.mock_span_stop.call_args_list[0] + self.assertEqual(kwargs["finish_time"], finish_time) + self.assertEqual(span.start, finish_time) + authorizer_span = args[0] + self.assertEqual(authorizer_span.name, "aws.apigateway.authorizer") + self.assertEqual(span.parent_id, authorizer_span.span_id) + return span + + + def _basic_common_checks(self, span, operation_name, span_type="http", route_key="/hello", http_method="GET"): + self.assertEqual(span.get_tag("apiid"), "amddr1rix9") + self.assertEqual(span.get_tag("apiname"), "amddr1rix9") + self.assertEqual(span.get_tag("stage"), "dev") self.assertEqual(span.get_tag("operation_name"), operation_name) + self.assertEqual(span.span_type, span_type) self.assertEqual( span.service, "amddr1rix9.execute-api.sa-east-1.amazonaws.com", ) self.assertEqual( span.get_tag("http.url"), - "amddr1rix9.execute-api.sa-east-1.amazonaws.com/hello", + "amddr1rix9.execute-api.sa-east-1.amazonaws.com"+route_key, ) - self.assertEqual(span.get_tag("endpoint"), "/hello") - self.assertEqual(span.get_tag("http.method"), "GET") + self.assertEqual(span.get_tag("endpoint"), route_key) + self.assertEqual(span.get_tag("http.method"), http_method) self.assertEqual( span.get_tag("resource_names"), - "GET /hello", + f"{http_method} {route_key}" if http_method else route_key, ) self.assertEqual(span.get_tag("request_id"), "1234567") - self.assertEqual(span.span_type, "http") - self.assertEqual(span.get_tag(InferredSpanInfo.TAG_SOURCE), "self") - self.assertEqual(span.get_tag(InferredSpanInfo.SYNCHRONICITY), "sync") - - # checking the upstream_authorizer_span - self.mock_span_stop.assert_called_once() - args, kwargs = self.mock_span_stop.call_args_list[0] - self.assertEqual(kwargs["finish_time"], finish_time) - authorizer_span = args[0] - self.assertEqual(authorizer_span.name, "aws.apigateway.authorizer") - self.assertEqual(span.parent_id, authorizer_span.span_id) - return span - class TestInferredSpans(unittest.TestCase): def test_create_inferred_span_from_api_gateway_event(self): From 1d6d58b617b30fdbb4144f8e36e739d524d4acb5 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Mon, 3 Oct 2022 12:02:59 -0400 Subject: [PATCH 10/26] reformat tests --- tests/test_tracing.py | 38 ++++++++++++++++++++++++++++---------- 1 file changed, 28 insertions(+), 10 deletions(-) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index d2909832..4537f630 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -588,29 +588,41 @@ def setUp(self): def test_create_inferred_span_from_authorizer_request_api_gateway_v1_event(self): event_sample_source = "authorizer-request-api-gateway-v1" - finish_time = 1663295021.832 #request_time_epoch + integrationLatency for api-gateway-v1 + finish_time = ( + 1663295021.832 # request_time_epoch + integrationLatency for api-gateway-v1 + ) span = self._authorizer_span_testing_items(event_sample_source, finish_time) self._basic_common_checks(span, "aws.apigateway.rest") def test_create_inferred_span_from_authorizer_token_api_gateway_v1_event(self): event_sample_source = "authorizer-token-api-gateway-v1" - finish_time = 1663295021.832 #request_time_epoch + integrationLatency for api-gateway-v1 + finish_time = ( + 1663295021.832 # request_time_epoch + integrationLatency for api-gateway-v1 + ) span = self._authorizer_span_testing_items(event_sample_source, finish_time) self._basic_common_checks(span, "aws.apigateway.rest") def test_create_inferred_span_from_authorizer_request_api_gateway_v2_event(self): event_sample_source = "authorizer-request-api-gateway-v2" - finish_time = 1664228639.533 # use the injected parent span finish time as an approximation + finish_time = 1664228639.533 # use the injected parent span finish time as an approximation span = self._authorizer_span_testing_items(event_sample_source, finish_time) self._basic_common_checks(span, "aws.httpapi") - def test_create_inferred_span_from_authorizer_request_api_gateway_websocket_connect_event(self): + def test_create_inferred_span_from_authorizer_request_api_gateway_websocket_connect_event( + self, + ): event_sample_source = "authorizer-request-api-gateway-websocket-connect" - finish_time = 1664388386.892 # request_time_epoch + integrationLatency in websocket case + finish_time = ( + 1664388386.892 # request_time_epoch + integrationLatency in websocket case + ) span = self._authorizer_span_testing_items(event_sample_source, finish_time) - self._basic_common_checks(span, "aws.apigateway.websocket", "web", "$connect", None) + self._basic_common_checks( + span, "aws.apigateway.websocket", "web", "$connect", None + ) - def test_create_inferred_span_from_authorizer_request_api_gateway_websocket_main_event(self): + def test_create_inferred_span_from_authorizer_request_api_gateway_websocket_main_event( + self, + ): event_sample_source = "authorizer-request-api-gateway-websocket-main" test_file = event_samples + event_sample_source + ".json" with open(test_file, "r") as event: @@ -641,8 +653,14 @@ def _authorizer_span_testing_items(self, event_sample_source, finish_time): self.assertEqual(span.parent_id, authorizer_span.span_id) return span - - def _basic_common_checks(self, span, operation_name, span_type="http", route_key="/hello", http_method="GET"): + def _basic_common_checks( + self, + span, + operation_name, + span_type="http", + route_key="/hello", + http_method="GET", + ): self.assertEqual(span.get_tag("apiid"), "amddr1rix9") self.assertEqual(span.get_tag("apiname"), "amddr1rix9") self.assertEqual(span.get_tag("stage"), "dev") @@ -654,7 +672,7 @@ def _basic_common_checks(self, span, operation_name, span_type="http", route_key ) self.assertEqual( span.get_tag("http.url"), - "amddr1rix9.execute-api.sa-east-1.amazonaws.com"+route_key, + "amddr1rix9.execute-api.sa-east-1.amazonaws.com" + route_key, ) self.assertEqual(span.get_tag("endpoint"), route_key) self.assertEqual(span.get_tag("http.method"), http_method) From a65787a1a3f7e25c3767bfb567d969d5e9201cd7 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Wed, 26 Oct 2022 13:13:49 -0400 Subject: [PATCH 11/26] encode authorizer injected data and other fixes --- datadog_lambda/constants.py | 9 +- datadog_lambda/tracing.py | 55 +++++------ datadog_lambda/wrapper.py | 62 +++++-------- ...horizer-request-api-gateway-v1-cached.json | 93 +++++++++++++++++++ .../authorizer-request-api-gateway-v1.json | 4 +- ...horizer-request-api-gateway-v2-cached.json | 47 ++++++++++ .../authorizer-request-api-gateway-v2.json | 4 +- ...request-api-gateway-websocket-connect.json | 6 +- ...equest-api-gateway-websocket-message.json} | 4 +- ...uthorizer-token-api-gateway-v1-cached.json | 93 +++++++++++++++++++ .../authorizer-token-api-gateway-v1.json | 4 +- tests/test_tracing.py | 49 +++++++++- 12 files changed, 341 insertions(+), 89 deletions(-) create mode 100644 tests/event_samples/authorizer-request-api-gateway-v1-cached.json create mode 100644 tests/event_samples/authorizer-request-api-gateway-v2-cached.json rename tests/event_samples/{authorizer-request-api-gateway-websocket-main.json => authorizer-request-api-gateway-websocket-message.json} (66%) create mode 100644 tests/event_samples/authorizer-token-api-gateway-v1-cached.json diff --git a/datadog_lambda/constants.py b/datadog_lambda/constants.py index acc6a75c..216def0d 100644 --- a/datadog_lambda/constants.py +++ b/datadog_lambda/constants.py @@ -44,6 +44,9 @@ class XrayDaemon(object): FUNCTION_NAME_HEADER_NAME = "AWS_LAMBDA_FUNCTION_NAME" -class OtherConsts(object): - parentSpanFinishTimeHeader = "x-datadog-parent-span-finish-time" - originalAuthorizerRequestEpoch = "x-datadog-original-epoch" +class Headers(object): + Parent_Span_Finish_Time = "x-datadog-parent-span-finish-time" + Authorizing_Request_Id = "x-datadog-authorizing-requestid" + TAGS_HEADER_TO_DELETE = ( + "x-datadog-tags" # injected by the HTTPPropagator.inject but no use + ) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index b1a024ae..ffbda2fd 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -24,7 +24,7 @@ TraceHeader, TraceContextSource, XrayDaemon, - OtherConsts, + Headers, ) from datadog_lambda.xray import ( send_segment, @@ -338,34 +338,29 @@ def extract_context_custom_extractor(extractor, event, lambda_context): def get_injected_authorizer_data(event, event_source: _EventSource) -> dict: try: - dd_data_raw = None - # has '_datadog' header and not cached because integrationLatency > 0 - authorizer_headers = event.get("requestContext", {}).get("authorizer", {}) + authorizer_headers = event.get("requestContext", {}).get("authorizer") + if not authorizer_headers: + return None - # [API_GATEWAY V1]if integrationLatency == 0, it's cached and no trace for the authorizer - if ( - event_source.equals( - EventTypes.API_GATEWAY, subtype=EventSubtypes.API_GATEWAY - ) - or event_source.equals( - EventTypes.API_GATEWAY, subtype=EventSubtypes.WEBSOCKET + dd_data_raw = ( + authorizer_headers.get("lambda", {}).get("_datadog") + if event_source.equals( + EventTypes.API_GATEWAY, subtype=EventSubtypes.HTTP_API ) - ) and int(authorizer_headers.get("integrationLatency", 0)) == 0: + else authorizer_headers.get("_datadog") + ) + + if not dd_data_raw: return None - if event_source.equals(EventTypes.API_GATEWAY, subtype=EventSubtypes.HTTP_API): - dd_data_raw = authorizer_headers.get("lambda", {}).get("_datadog") - dd_data = json.loads(base64.b64decode(dd_data_raw)) - # [API_GATEWAY V2]if original authorizer request epoch is different, then it's cached - if event.get("requestContext", {}).get("timeEpoch") == dd_data.get( - OtherConsts.originalAuthorizerRequestEpoch - ): - return dd_data - else: - return None + injected_data = json.loads(base64.b64decode(dd_data_raw)) + # use the integrationLatency (needed for token type authorizer) or the injected requestId to tell if it's the authorizing invocation (not cached) + if authorizer_headers.get("integrationLatency", 0) > 0 or event.get( + "requestContext", {} + ).get("requestId") == injected_data.get(Headers.Authorizing_Request_Id): + return injected_data else: - dd_data_raw = authorizer_headers.get("_datadog") - return json.loads(dd_data_raw) + return None except Exception as e: logger.debug("Failed to check if invocated by an authorizer. error %s", e) @@ -709,9 +704,7 @@ def create_inferred_span_from_api_gateway_websocket_event(event, context): if injected_authorizer_data: try: start_time_s = ( - int( - injected_authorizer_data.get(OtherConsts.parentSpanFinishTimeHeader) - ) + int(injected_authorizer_data.get(Headers.Parent_Span_Finish_Time)) / 1000 ) finish_time_s = ( @@ -788,9 +781,7 @@ def create_inferred_span_from_api_gateway_event(event, context): if injected_authorizer_data: try: start_time_s = ( - int( - injected_authorizer_data.get(OtherConsts.parentSpanFinishTimeHeader) - ) + int(injected_authorizer_data.get(Headers.Parent_Span_Finish_Time)) / 1000 ) finish_time_s = ( @@ -870,9 +861,7 @@ def create_inferred_span_from_http_api_event(event, context): if injected_authorizer_data: try: start_time_s = ( - int( - injected_authorizer_data.get(OtherConsts.parentSpanFinishTimeHeader) - ) + int(injected_authorizer_data.get(Headers.Parent_Span_Finish_Time)) / 1000 ) finish_time_s = start_time_s # no integrationLatency info in this case diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 02d2485c..4d0401eb 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -17,7 +17,7 @@ from datadog_lambda.constants import ( TraceContextSource, XraySubsegment, - OtherConsts, + Headers, ) from datadog_lambda.metric import ( flush_stats, @@ -126,8 +126,8 @@ def __init__(self, func): self.make_inferred_span = ( os.environ.get("DD_TRACE_MANAGED_SERVICES", "true").lower() == "true" ) - self.make_authorizer_span = ( - os.environ.get("DD_LINK_LAMBDA_AUTHORIZER", "true").lower() == "true" + self.encode_authorizer_context = ( + os.environ.get("DD_ENCODE_AUTHORIZER_CONTEXT", "true").lower() == "true" ) self.response = None if profiling_env_var: @@ -161,7 +161,6 @@ def __call__(self, event, context, **kwargs): self._before(event, context) try: self.response = self.func(event, context, **kwargs) - self._ending(event) return self.response except Exception: submit_errors_metric(context) @@ -171,29 +170,17 @@ def __call__(self, event, context, **kwargs): finally: self._after(event, context) - def _inject_authorizer_span_headers(self, event): - finish_time_ns = ( - self.span.start_ns - if InferredSpanInfo.is_async(self.inferred_span) and self.span - else time_ns() - ) + def _inject_authorizer_span_headers(self, request_id, finish_time_ns): injected_headers = {} source_span = self.inferred_span if self.inferred_span else self.span HTTPPropagator.inject(source_span.context, injected_headers) - injected_headers[OtherConsts.parentSpanFinishTimeHeader] = finish_time_ns / 1e6 - encode_datadog_data = False - if self.event_source and self.event_source.equals( - EventTypes.API_GATEWAY, EventSubtypes.HTTP_API - ): - injected_headers[OtherConsts.originalAuthorizerRequestEpoch] = event.get( - "requestContext", {} - ).get("timeEpoch") - encode_datadog_data = True - datadog_data = json.dumps(injected_headers).encode() - if encode_datadog_data: - datadog_data = base64.b64encode( - datadog_data - ) # special characters can corrupt the response + injected_headers.pop( + Headers.TAGS_HEADER_TO_DELETE, None + ) # remove unused header + injected_headers[Headers.Parent_Span_Finish_Time] = finish_time_ns / 1e6 + if request_id is not None: + injected_headers[Headers.Authorizing_Request_Id] = request_id + datadog_data = base64.b64encode(json.dumps(injected_headers).encode()) self.response.setdefault("context", {}) self.response["context"]["_datadog"] = datadog_data @@ -237,19 +224,6 @@ def _before(self, event, context): except Exception: traceback.print_exc() - def _ending(self, event): - try: - if ( - self.make_authorizer_span - and self.response - and self.response.get("principalId") - and self.response.get("policyDocument") - and self.event_source - ): - self._inject_authorizer_span_headers(event) - except Exception: - traceback.print_exc() - def _after(self, event, context): try: status_code = extract_http_status_code_tag(self.trigger_tags, self.response) @@ -282,6 +256,20 @@ def _after(self, event, context): else: self.inferred_span.finish() + if ( + self.encode_authorizer_context + and self.response + and self.response.get("principalId") + and self.response.get("policyDocument") + ): + finish_time_ns = ( + self.span.start_ns + if InferredSpanInfo.is_async(self.inferred_span) and self.span + else time_ns() + ) + self._inject_authorizer_span_headers( + event.get("requestContext", {}).get("requestId"), finish_time_ns + ) if not self.flush_to_log or should_use_extension: flush_stats() if should_use_extension: diff --git a/tests/event_samples/authorizer-request-api-gateway-v1-cached.json b/tests/event_samples/authorizer-request-api-gateway-v1-cached.json new file mode 100644 index 00000000..937e94b6 --- /dev/null +++ b/tests/event_samples/authorizer-request-api-gateway-v1-cached.json @@ -0,0 +1,93 @@ +{ + "resource": "/hello", + "path": "/hello", + "httpMethod": "GET", + "headers": { + "Accept": "*/*", + "Accept-Encoding": "gzip, deflate, br", + "Authorization": "secretT0k3n", + "authorizationToken": "secretT0k3n", + "Cache-Control": "no-cache", + "CloudFront-Forwarded-Proto": "https", + "CloudFront-Is-Desktop-Viewer": "true", + "CloudFront-Is-Mobile-Viewer": "false", + "CloudFront-Is-SmartTV-Viewer": "false", + "CloudFront-Is-Tablet-Viewer": "false", + "CloudFront-Viewer-ASN": "12271", + "CloudFront-Viewer-Country": "US", + "Host": "amddr1rix9.execute-api.sa-east-1.amazonaws.com", + "Postman-Token": "5ddd364a-4191-4dbc-9831-09f1bf46bca8", + "User-Agent": "PostmanRuntime/7.29.2", + "Via": "1.1 x.cloudfront.net (CloudFront)", + "X-Amz-Cf-Id": "ToqncSIka6Ae1pZlyxEOtyqqcns8cwL3hVxwo0ZZqMC60zTSSXW7yA==", + "X-Amzn-Trace-Id": "Root=1-63580c1d-12aa2dcf30cd2d3476a6f18f", + "X-Forwarded-For": "24.193.182.233, 15.158.35.232", + "X-Forwarded-Port": "443", + "X-Forwarded-Proto": "https" + }, + "multiValueHeaders": { + "Accept": ["*/*"], + "Accept-Encoding": ["gzip, deflate, br"], + "Authorization": ["secretT0k3n"], + "authorizationToken": ["secretT0k3n"], + "Cache-Control": ["no-cache"], + "CloudFront-Forwarded-Proto": ["https"], + "CloudFront-Is-Desktop-Viewer": ["true"], + "CloudFront-Is-Mobile-Viewer": ["false"], + "CloudFront-Is-SmartTV-Viewer": ["false"], + "CloudFront-Is-Tablet-Viewer": ["false"], + "CloudFront-Viewer-ASN": ["12271"], + "CloudFront-Viewer-Country": ["US"], + "Host": ["amddr1rix9.execute-api.sa-east-1.amazonaws.com"], + "Postman-Token": ["5ddd364a-4191-4dbc-9831-09f1bf46bca8"], + "User-Agent": ["PostmanRuntime/7.29.2"], + "Via": ["1.1 x.cloudfront.net (CloudFront)"], + "X-Amz-Cf-Id": ["ToqncSIka6Ae1pZlyxEOtyqqcns8cwL3hVxwo0ZZqMC60zTSSXW7yA=="], + "X-Amzn-Trace-Id": ["Root=1-63580c1d-12aa2dcf30cd2d3476a6f18f"], + "X-Forwarded-For": ["24.193.182.233, 15.158.35.232"], + "X-Forwarded-Port": ["443"], + "X-Forwarded-Proto": ["https"] + }, + "queryStringParameters": null, + "multiValueQueryStringParameters": null, + "pathParameters": null, + "stageVariables": null, + "requestContext": { + "resourceId": "var16m", + "authorizer": { + "_datadog": "eyJ4LWRhdGFkb2ctdHJhY2UtaWQiOiAiMzQyMDA1NDg3OTQxNzk5OTg3NyIsICJ4LWRhdGFkb2ctcGFyZW50LWlkIjogIjE1NDY2ODY0ODU4OTY3MTY5NjM2IiwgIngtZGF0YWRvZy1zYW1wbGluZy1wcmlvcml0eSI6ICIxIiwgIngtZGF0YWRvZy1wYXJlbnQtc3Bhbi1maW5pc2gtdGltZSI6IDE2NjY3MTQ2NDg1OTAuMTA5NCwgIngtZGF0YWRvZy1hdXRob3JpemluZy1yZXF1ZXN0aWQiOiAiMjRkNjk4OTItNjcwMi00OTYxLThlNDgtMDg1NWNjNTU0MjQwIn0=", + "scope": "this is just a string", + "principalId": "foo", + "integrationLatency": 0 + }, + "resourcePath": "/hello", + "httpMethod": "GET", + "extendedRequestId": "akbUqFBZmjQFbJg=", + "requestTime": "25/Oct/2022:16:17:33 +0000", + "path": "/dev/hello", + "accountId": "425362996713", + "protocol": "HTTP/1.1", + "stage": "dev", + "domainPrefix": "amddr1rix9", + "requestTimeEpoch": 1666714653636, + "requestId": "abc123", + "identity": { + "cognitoIdentityPoolId": null, + "accountId": null, + "cognitoIdentityId": null, + "caller": null, + "sourceIp": "24.193.182.233", + "principalOrgId": null, + "accessKey": null, + "cognitoAuthenticationType": null, + "cognitoAuthenticationProvider": null, + "userArn": null, + "userAgent": "PostmanRuntime/7.29.2", + "user": null + }, + "domainName": "amddr1rix9.execute-api.sa-east-1.amazonaws.com", + "apiId": "amddr1rix9" + }, + "body": null, + "isBase64Encoded": false +} diff --git a/tests/event_samples/authorizer-request-api-gateway-v1.json b/tests/event_samples/authorizer-request-api-gateway-v1.json index 16dfbd2a..00818eac 100644 --- a/tests/event_samples/authorizer-request-api-gateway-v1.json +++ b/tests/event_samples/authorizer-request-api-gateway-v1.json @@ -55,7 +55,7 @@ "requestContext": { "resourceId": "0et54l", "authorizer": { - "_datadog": "{\"x-datadog-trace-id\": \"13478705995797221209\", \"x-datadog-parent-id\": \"8471288263384216896\", \"x-datadog-sampling-priority\": \"1\", \"x-datadog-tags\": \"_dd.p.dm=-0\", \"x-datadog-parent-span-finish-time\": 1663295021827.521}", + "_datadog": "eyJ4LWRhdGFkb2ctdHJhY2UtaWQiOiAiMTM0Nzg3MDU5OTU3OTcyMjEyMDkiLCAieC1kYXRhZG9nLXBhcmVudC1pZCI6ICI4NDcxMjg4MjYzMzg0MjE2ODk2IiwgIngtZGF0YWRvZy1zYW1wbGluZy1wcmlvcml0eSI6ICIxIiwgIngtZGF0YWRvZy1wYXJlbnQtc3Bhbi1maW5pc2gtdGltZSI6IDE2NjMyOTUwMjE4MjcuNTIxLCAieC1kYXRhZG9nLWF1dGhvcml6aW5nLXJlcXVlc3RpZCI6ICJhYmMxMjMifQ==", "scope": "this is just a string", "principalId": "foo", "integrationLatency": 1897 @@ -70,7 +70,7 @@ "stage": "dev", "domainPrefix": "amddr1rix9", "requestTimeEpoch": 1663295019935, - "requestId": "1234567", + "requestId": "abc123", "identity": { "cognitoIdentityPoolId": null, "accountId": null, diff --git a/tests/event_samples/authorizer-request-api-gateway-v2-cached.json b/tests/event_samples/authorizer-request-api-gateway-v2-cached.json new file mode 100644 index 00000000..53ee1624 --- /dev/null +++ b/tests/event_samples/authorizer-request-api-gateway-v2-cached.json @@ -0,0 +1,47 @@ +{ + "version": "2.0", + "routeKey": "GET /hello", + "rawPath": "/hello", + "rawQueryString": "", + "headers": { + "accept": "*/*", + "accept-encoding": "gzip, deflate, br", + "authorization": "secretT0k3n", + "authorizationtoken": "secretT0k3n", + "cache-control": "no-cache", + "content-length": "0", + "host": "amddr1rix9.execute-api.sa-east-1.amazonaws.com", + "postman-token": "f901530c-cc2f-475f-8634-d9156e9cd9d7", + "user-agent": "PostmanRuntime/7.29.2", + "userid": "1236", + "x-amzn-trace-id": "Root=1-63580f25-5f01511071135d5966c8c63d", + "x-forwarded-for": "24.193.182.233", + "x-forwarded-port": "443", + "x-forwarded-proto": "https" + }, + "requestContext": { + "accountId": "425362996713", + "apiId": "amddr1rix9", + "authorizer": { + "lambda": { + "_datadog": "eyJ4LWRhdGFkb2ctdHJhY2UtaWQiOiAiMTc4MDIxMTQ1NTc4ODI1OTY0MTYiLCAieC1kYXRhZG9nLXBhcmVudC1pZCI6ICIxMjg3ODYzNTg0NTE5MTY5NDQ5NCIsICJ4LWRhdGFkb2ctc2FtcGxpbmctcHJpb3JpdHkiOiAiMSIsICJ4LWRhdGFkb2ctcGFyZW50LXNwYW4tZmluaXNoLXRpbWUiOiAxNjY2NzE1NDI5Mzk1Ljc3ODgsICJ4LWRhdGFkb2ctYXV0aG9yaXppbmctcmVxdWVzdGlkIjogImFrZE4zaEdER2pRRU1Ddz0ifQ==", + "scope": "this is just a string" + } + }, + "domainName": "amddr1rix9.execute-api.sa-east-1.amazonaws.com", + "domainPrefix": "amddr1rix9", + "http": { + "method": "GET", + "path": "/hello", + "protocol": "HTTP/1.1", + "sourceIp": "24.193.182.233", + "userAgent": "PostmanRuntime/7.29.2" + }, + "requestId": "abc123", + "routeKey": "GET /hello", + "stage": "dev", + "time": "25/Oct/2022:16:30:29 +0000", + "timeEpoch": 1666715429349 + }, + "isBase64Encoded": false +} diff --git a/tests/event_samples/authorizer-request-api-gateway-v2.json b/tests/event_samples/authorizer-request-api-gateway-v2.json index ac11c6ed..2be965e7 100644 --- a/tests/event_samples/authorizer-request-api-gateway-v2.json +++ b/tests/event_samples/authorizer-request-api-gateway-v2.json @@ -24,7 +24,7 @@ "apiId": "amddr1rix9", "authorizer": { "lambda": { - "_datadog": "eyJ4LWRhdGFkb2ctdHJhY2UtaWQiOiAiMTQzNTY5ODM2MTk4NTI5MzMzNTQiLCAieC1kYXRhZG9nLXBhcmVudC1pZCI6ICIxMjY1ODYyMTA4MzUwNTQxMzgwOSIsICJ4LWRhdGFkb2ctc2FtcGxpbmctcHJpb3JpdHkiOiAiMSIsICJ4LWRhdGFkb2ctdGFncyI6ICJfZGQucC5kbT0tMCIsICJ4LWRhdGFkb2ctcGFyZW50LXNwYW4tZmluaXNoLXRpbWUiOiAxNjY0MjI4NjM5NTMzLjc3NTQsICJ4LWRhdGFkb2ctb3JpZ2luYWwtZXBvY2giOiAxNjY0MjI4NjM4MDg0fQ==", + "_datadog": "eyJ4LWRhdGFkb2ctdHJhY2UtaWQiOiAiMTQzNTY5ODM2MTk4NTI5MzMzNTQiLCAieC1kYXRhZG9nLXBhcmVudC1pZCI6ICIxMjY1ODYyMTA4MzUwNTQxMzgwOSIsICJ4LWRhdGFkb2ctc2FtcGxpbmctcHJpb3JpdHkiOiAiMSIsICJ4LWRhdGFkb2ctcGFyZW50LXNwYW4tZmluaXNoLXRpbWUiOiAxNjY0MjI4NjM5NTMzLjc3NTQsICJ4LWRhdGFkb2ctYXV0aG9yaXppbmctcmVxdWVzdGlkIjogImFiYzEyMyJ9", "scope": "this is just a string" } }, @@ -37,7 +37,7 @@ "sourceIp": "38.122.226.210", "userAgent": "curl/7.64.1" }, - "requestId": "1234567", + "requestId": "abc123", "routeKey": "GET /hello", "stage": "dev", "time": "26/Sep/2022:21:43:58 +0000", diff --git a/tests/event_samples/authorizer-request-api-gateway-websocket-connect.json b/tests/event_samples/authorizer-request-api-gateway-websocket-connect.json index a325a3f7..1875097d 100644 --- a/tests/event_samples/authorizer-request-api-gateway-websocket-connect.json +++ b/tests/event_samples/authorizer-request-api-gateway-websocket-connect.json @@ -24,13 +24,13 @@ "requestContext": { "routeKey": "$connect", "authorizer": { - "_datadog": "{\"x-datadog-trace-id\": \"5351047404834723189\", \"x-datadog-parent-id\": \"18230460631156161837\", \"x-datadog-sampling-priority\": \"1\", \"x-datadog-tags\": \"_dd.p.dm=-0\", \"x-datadog-parent-span-finish-time\": 1664388386890}", + "_datadog": "eyJ4LWRhdGFkb2ctdHJhY2UtaWQiOiAiNTM1MTA0NzQwNDgzNDcyMzE4OSIsICJ4LWRhdGFkb2ctcGFyZW50LWlkIjogIjE4MjMwNDYwNjMxMTU2MTYxODM3IiwgIngtZGF0YWRvZy1zYW1wbGluZy1wcmlvcml0eSI6ICIxIiwgIngtZGF0YWRvZy1wYXJlbnQtc3Bhbi1maW5pc2gtdGltZSI6IDE2NjQzODgzODY4OTAsICJ4LWRhdGFkb2ctYXV0aG9yaXppbmctcmVxdWVzdGlkIjogImFiYzEyMyJ9", "scope": "this is just a string", "principalId": "foo", "integrationLatency": 1440 }, "eventType": "CONNECT", - "extendedRequestId": "1234567", + "extendedRequestId": "abc123", "requestTime": "28/Sep/2022:18:06:25 +0000", "messageDirection": "IN", "stage": "dev", @@ -39,7 +39,7 @@ "identity": { "sourceIp": "38.122.226.210" }, - "requestId": "1234567", + "requestId": "abc123", "domainName": "amddr1rix9.execute-api.sa-east-1.amazonaws.com", "connectionId": "ZLr9QeNLmjQCIZA=", "apiId": "amddr1rix9" diff --git a/tests/event_samples/authorizer-request-api-gateway-websocket-main.json b/tests/event_samples/authorizer-request-api-gateway-websocket-message.json similarity index 66% rename from tests/event_samples/authorizer-request-api-gateway-websocket-main.json rename to tests/event_samples/authorizer-request-api-gateway-websocket-message.json index 0cda1a6e..8ac5d2e9 100644 --- a/tests/event_samples/authorizer-request-api-gateway-websocket-main.json +++ b/tests/event_samples/authorizer-request-api-gateway-websocket-message.json @@ -2,7 +2,7 @@ "requestContext": { "routeKey": "main", "authorizer": { - "_datadog": "{\"x-datadog-trace-id\": \"17374069065025434920\", \"x-datadog-parent-id\": \"7915186751914176106\", \"x-datadog-sampling-priority\": \"1\", \"x-datadog-tags\": \"_dd.p.dm=-0\", \"x-datadog-parent-span-finish-time\": 1664390333526.5068}", + "_datadog": "eyJ4LWRhdGFkb2ctdHJhY2UtaWQiOiAiMTczNzQwNjkwNjUwMjU0MzQ5MjAiLCAieC1kYXRhZG9nLXBhcmVudC1pZCI6ICI3OTE1MTg2NzUxOTE0MTc2MTA2IiwgIngtZGF0YWRvZy1zYW1wbGluZy1wcmlvcml0eSI6ICIxIiwgIngtZGF0YWRvZy1wYXJlbnQtc3Bhbi1maW5pc2gtdGltZSI6IDE2NjQzOTAzMzM1MjYuNTA2OCwgIngtZGF0YWRvZy1hdXRob3JpemluZy1yZXF1ZXN0aWQiOiAiYWJjMzIxIn0=", "scope": "this is just a string", "principalId": "foo" }, @@ -17,7 +17,7 @@ "identity": { "sourceIp": "38.122.226.210" }, - "requestId": "1234567", + "requestId": "abc123", "domainName": "amddr1rix9.execute-api.sa-east-1.amazonaws.com", "connectionId": "ZLwtceO1mjQCI8Q=", "apiId": "amddr1rix9" diff --git a/tests/event_samples/authorizer-token-api-gateway-v1-cached.json b/tests/event_samples/authorizer-token-api-gateway-v1-cached.json new file mode 100644 index 00000000..18ade248 --- /dev/null +++ b/tests/event_samples/authorizer-token-api-gateway-v1-cached.json @@ -0,0 +1,93 @@ +{ + "resource": "/hello", + "path": "/hello", + "httpMethod": "GET", + "headers": { + "Accept": "*/*", + "Accept-Encoding": "gzip, deflate, br", + "Authorization": "secretT0k3n", + "authorizationToken": "secretT0k3n", + "Cache-Control": "no-cache", + "CloudFront-Forwarded-Proto": "https", + "CloudFront-Is-Desktop-Viewer": "true", + "CloudFront-Is-Mobile-Viewer": "false", + "CloudFront-Is-SmartTV-Viewer": "false", + "CloudFront-Is-Tablet-Viewer": "false", + "CloudFront-Viewer-ASN": "174", + "CloudFront-Viewer-Country": "US", + "Host": "amddr1rix9.execute-api.sa-east-1.amazonaws.com", + "Postman-Token": "bb30e418-d4ee-45b6-a659-89f4dbd41c84", + "User-Agent": "PostmanRuntime/7.29.2", + "Via": "1.1 x.cloudfront.net (CloudFront)", + "X-Amz-Cf-Id": "iFM3jhMRNr0jJn8b4A7lgYVLUhUb6pOxFQm1YA4JDTDylUHJOOl-Jw==", + "X-Amzn-Trace-Id": "Root=1-635967a6-2b0999ba248345b32a8ff5ce", + "X-Forwarded-For": "38.122.226.210, 130.176.179.111", + "X-Forwarded-Port": "443", + "X-Forwarded-Proto": "https" + }, + "multiValueHeaders": { + "Accept": ["*/*"], + "Accept-Encoding": ["gzip, deflate, br"], + "Authorization": ["secretT0k3n"], + "authorizationToken": ["secretT0k3n"], + "Cache-Control": ["no-cache"], + "CloudFront-Forwarded-Proto": ["https"], + "CloudFront-Is-Desktop-Viewer": ["true"], + "CloudFront-Is-Mobile-Viewer": ["false"], + "CloudFront-Is-SmartTV-Viewer": ["false"], + "CloudFront-Is-Tablet-Viewer": ["false"], + "CloudFront-Viewer-ASN": ["174"], + "CloudFront-Viewer-Country": ["US"], + "Host": ["amddr1rix9.execute-api.sa-east-1.amazonaws.com"], + "Postman-Token": ["bb30e418-d4ee-45b6-a659-89f4dbd41c84"], + "User-Agent": ["PostmanRuntime/7.29.2"], + "Via": ["1.1 x.cloudfront.net (CloudFront)"], + "X-Amz-Cf-Id": ["iFM3jhMRNr0jJn8b4A7lgYVLUhUb6pOxFQm1YA4JDTDylUHJOOl-Jw=="], + "X-Amzn-Trace-Id": ["Root=1-635967a6-2b0999ba248345b32a8ff5ce"], + "X-Forwarded-For": ["38.122.226.210, 130.176.179.111"], + "X-Forwarded-Port": ["443"], + "X-Forwarded-Proto": ["https"] + }, + "queryStringParameters": null, + "multiValueQueryStringParameters": null, + "pathParameters": null, + "stageVariables": null, + "requestContext": { + "resourceId": "mi7nzd", + "authorizer": { + "_datadog": "eyJ4LWRhdGFkb2ctdHJhY2UtaWQiOiAiMTM3NTkyNDk3NDA5MDA1OTMyNjEiLCAieC1kYXRhZG9nLXBhcmVudC1pZCI6ICIxNTkzNTAyMzY1NTYxMjM0MTgxOSIsICJ4LWRhdGFkb2ctc2FtcGxpbmctcHJpb3JpdHkiOiAiMSIsICJ4LWRhdGFkb2ctcGFyZW50LXNwYW4tZmluaXNoLXRpbWUiOiAxNjY2ODAzNjE2OTMzLjIwNzh9", + "scope": "this is just a string", + "principalId": "foo", + "integrationLatency": 0 + }, + "resourcePath": "/hello", + "httpMethod": "GET", + "extendedRequestId": "an0iHGmHmjQFjLw=", + "requestTime": "26/Oct/2022:17:00:22 +0000", + "path": "/dev/hello", + "accountId": "425362996713", + "protocol": "HTTP/1.1", + "stage": "dev", + "domainPrefix": "amddr1rix9", + "requestTimeEpoch": 1666803622990, + "requestId": "abc123", + "identity": { + "cognitoIdentityPoolId": null, + "accountId": null, + "cognitoIdentityId": null, + "caller": null, + "sourceIp": "38.122.226.210", + "principalOrgId": null, + "accessKey": null, + "cognitoAuthenticationType": null, + "cognitoAuthenticationProvider": null, + "userArn": null, + "userAgent": "PostmanRuntime/7.29.2", + "user": null + }, + "domainName": "amddr1rix9.execute-api.sa-east-1.amazonaws.com", + "apiId": "amddr1rix9" + }, + "body": null, + "isBase64Encoded": false +} diff --git a/tests/event_samples/authorizer-token-api-gateway-v1.json b/tests/event_samples/authorizer-token-api-gateway-v1.json index 0c28772a..1edba05e 100644 --- a/tests/event_samples/authorizer-token-api-gateway-v1.json +++ b/tests/event_samples/authorizer-token-api-gateway-v1.json @@ -55,7 +55,7 @@ "requestContext": { "resourceId": "elem3u", "authorizer": { - "_datadog": "{\"x-datadog-trace-id\": \"17874798268144902712\", \"x-datadog-parent-id\": \"16184667399315372101\", \"x-datadog-sampling-priority\": \"1\", \"x-datadog-tags\": \"_dd.p.dm=-0\", \"x-datadog-parent-span-finish-time\": 1663721602440.8286}", + "_datadog": "eyJ4LWRhdGFkb2ctdHJhY2UtaWQiOiAiMTc4NzQ3OTgyNjgxNDQ5MDI3MTIiLCAieC1kYXRhZG9nLXBhcmVudC1pZCI6ICIxNjE4NDY2NzM5OTMxNTM3MjEwMSIsICJ4LWRhdGFkb2ctc2FtcGxpbmctcHJpb3JpdHkiOiAiMSIsICJ4LWRhdGFkb2ctcGFyZW50LXNwYW4tZmluaXNoLXRpbWUiOiAxNjYzNzIxNjAyNDQwLjgyODZ9", "scope": "this is just a string", "principalId": "foo", "integrationLatency": 1897 @@ -70,7 +70,7 @@ "stage": "dev", "domainPrefix": "amddr1rix9", "requestTimeEpoch": 1663295019935, - "requestId": "1234567", + "requestId": "abc123", "identity": { "cognitoIdentityPoolId": null, "accountId": null, diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 4537f630..13481be3 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -594,6 +594,19 @@ def test_create_inferred_span_from_authorizer_request_api_gateway_v1_event(self) span = self._authorizer_span_testing_items(event_sample_source, finish_time) self._basic_common_checks(span, "aws.apigateway.rest") + def test_create_inferred_span_from_authorizer_request_api_gateway_v1_cached_event( + self, + ): + event_sample_source = "authorizer-request-api-gateway-v1-cached" + test_file = event_samples + event_sample_source + ".json" + with open(test_file, "r") as event: + event = json.load(event) + ctx = get_mock_context() + ctx.aws_request_id = "abc123" # injected data's requestId is abc321 + span = create_inferred_span(event, ctx) + self.mock_span_stop.assert_not_called() # NO authorizer span is injected + self._basic_common_checks(span, "aws.apigateway.rest") + def test_create_inferred_span_from_authorizer_token_api_gateway_v1_event(self): event_sample_source = "authorizer-token-api-gateway-v1" finish_time = ( @@ -602,12 +615,38 @@ def test_create_inferred_span_from_authorizer_token_api_gateway_v1_event(self): span = self._authorizer_span_testing_items(event_sample_source, finish_time) self._basic_common_checks(span, "aws.apigateway.rest") + def test_create_inferred_span_from_authorizer_token_api_gateway_v2_cached_event( + self, + ): + event_sample_source = "authorizer-token-api-gateway-v1-cached" + test_file = event_samples + event_sample_source + ".json" + with open(test_file, "r") as event: + event = json.load(event) + ctx = get_mock_context() + ctx.aws_request_id = "abc123" # injected data's requestId is abc321 + span = create_inferred_span(event, ctx) + self.mock_span_stop.assert_not_called() # NO authorizer span is injected + self._basic_common_checks(span, "aws.apigateway.rest") + def test_create_inferred_span_from_authorizer_request_api_gateway_v2_event(self): event_sample_source = "authorizer-request-api-gateway-v2" finish_time = 1664228639.533 # use the injected parent span finish time as an approximation span = self._authorizer_span_testing_items(event_sample_source, finish_time) self._basic_common_checks(span, "aws.httpapi") + def test_create_inferred_span_from_authorizer_request_api_gateway_v2_cached_event( + self, + ): + event_sample_source = "authorizer-request-api-gateway-v2-cached" + test_file = event_samples + event_sample_source + ".json" + with open(test_file, "r") as event: + event = json.load(event) + ctx = get_mock_context() + ctx.aws_request_id = "abc123" # injected data's requestId is abc321 + span = create_inferred_span(event, ctx) + self.mock_span_stop.assert_not_called() # NO authorizer span is injected + self._basic_common_checks(span, "aws.httpapi") + def test_create_inferred_span_from_authorizer_request_api_gateway_websocket_connect_event( self, ): @@ -620,15 +659,15 @@ def test_create_inferred_span_from_authorizer_request_api_gateway_websocket_conn span, "aws.apigateway.websocket", "web", "$connect", None ) - def test_create_inferred_span_from_authorizer_request_api_gateway_websocket_main_event( + def test_create_inferred_span_from_authorizer_request_api_gateway_websocket_message_event( self, ): - event_sample_source = "authorizer-request-api-gateway-websocket-main" + event_sample_source = "authorizer-request-api-gateway-websocket-message" test_file = event_samples + event_sample_source + ".json" with open(test_file, "r") as event: event = json.load(event) ctx = get_mock_context() - ctx.aws_request_id = "1234567" + ctx.aws_request_id = "abc123" # injected data's requestId is abc321 span = create_inferred_span(event, ctx) self.mock_span_stop.assert_not_called() # NO authorizer span is injected self._basic_common_checks(span, "aws.apigateway.websocket", "web", "main", None) @@ -638,7 +677,7 @@ def _authorizer_span_testing_items(self, event_sample_source, finish_time): with open(test_file, "r") as event: event = json.load(event) ctx = get_mock_context() - ctx.aws_request_id = "1234567" + ctx.aws_request_id = "abc123" span = create_inferred_span(event, ctx) self.assertEqual(span.get_tag(InferredSpanInfo.TAG_SOURCE), "self") self.assertEqual(span.get_tag(InferredSpanInfo.SYNCHRONICITY), "sync") @@ -680,7 +719,7 @@ def _basic_common_checks( span.get_tag("resource_names"), f"{http_method} {route_key}" if http_method else route_key, ) - self.assertEqual(span.get_tag("request_id"), "1234567") + self.assertEqual(span.get_tag("request_id"), "abc123") class TestInferredSpans(unittest.TestCase): From 7b672e076107f73bcc3cd43e0ca3205f9f3c0404 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Wed, 26 Oct 2022 13:26:39 -0400 Subject: [PATCH 12/26] format fixes --- datadog_lambda/constants.py | 6 +++--- datadog_lambda/tracing.py | 3 ++- datadog_lambda/wrapper.py | 7 ++----- 3 files changed, 7 insertions(+), 9 deletions(-) diff --git a/datadog_lambda/constants.py b/datadog_lambda/constants.py index 216def0d..23bce904 100644 --- a/datadog_lambda/constants.py +++ b/datadog_lambda/constants.py @@ -47,6 +47,6 @@ class XrayDaemon(object): class Headers(object): Parent_Span_Finish_Time = "x-datadog-parent-span-finish-time" Authorizing_Request_Id = "x-datadog-authorizing-requestid" - TAGS_HEADER_TO_DELETE = ( - "x-datadog-tags" # injected by the HTTPPropagator.inject but no use - ) + + # injected by the HTTPPropagator.inject but no use + TAGS_HEADER_TO_DELETE = "x-datadog-tags" diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index ffbda2fd..51205bcb 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -354,7 +354,8 @@ def get_injected_authorizer_data(event, event_source: _EventSource) -> dict: return None injected_data = json.loads(base64.b64decode(dd_data_raw)) - # use the integrationLatency (needed for token type authorizer) or the injected requestId to tell if it's the authorizing invocation (not cached) + # use the integrationLatency (needed for token type authorizer) or the injected requestId + # to tell if it's the authorizing invocation (not cached) if authorizer_headers.get("integrationLatency", 0) > 0 or event.get( "requestContext", {} ).get("requestId") == injected_data.get(Headers.Authorizing_Request_Id): diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 4d0401eb..680834a9 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -41,8 +41,6 @@ from datadog_lambda.trigger import ( extract_trigger_tags, extract_http_status_code_tag, - EventTypes, - EventSubtypes, ) from datadog_lambda.tag_object import tag_object @@ -174,9 +172,8 @@ def _inject_authorizer_span_headers(self, request_id, finish_time_ns): injected_headers = {} source_span = self.inferred_span if self.inferred_span else self.span HTTPPropagator.inject(source_span.context, injected_headers) - injected_headers.pop( - Headers.TAGS_HEADER_TO_DELETE, None - ) # remove unused header + # remove unused header + injected_headers.pop(Headers.TAGS_HEADER_TO_DELETE, None) injected_headers[Headers.Parent_Span_Finish_Time] = finish_time_ns / 1e6 if request_id is not None: injected_headers[Headers.Authorizing_Request_Id] = request_id From 8c130f247d0ffdd27315c33e308d93e1deebcf40 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Wed, 26 Oct 2022 14:58:59 -0400 Subject: [PATCH 13/26] fix token-type authorizer span's startTime --- datadog_lambda/wrapper.py | 21 ++++++++++++++++----- 1 file changed, 16 insertions(+), 5 deletions(-) diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 680834a9..c4a6c5fc 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -259,11 +259,22 @@ def _after(self, event, context): and self.response.get("principalId") and self.response.get("policyDocument") ): - finish_time_ns = ( - self.span.start_ns - if InferredSpanInfo.is_async(self.inferred_span) and self.span - else time_ns() - ) + finish_time_ns = None + # the finish_time_ns should be set as the end of the inferred span if it exist + # or the end of the current span + if not self.inferred_span: + finish_time_ns = ( + self.span.start_ns + self.span.duration_ns + if self.span and self.span.duration_ns + else time_ns() + ) + else: + finish_time_ns = ( + self.span.start_ns + if InferredSpanInfo.is_async(self.inferred_span) and self.span + else time_ns() + ) + self._inject_authorizer_span_headers( event.get("requestContext", {}).get("requestId"), finish_time_ns ) From 63d1777dc3dcddbe1f58804fa0189ff583863d75 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Wed, 26 Oct 2022 15:29:21 -0400 Subject: [PATCH 14/26] more precise on the parent span finish time --- datadog_lambda/wrapper.py | 31 +++++++++++++------------------ 1 file changed, 13 insertions(+), 18 deletions(-) diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index c4a6c5fc..3ce71d29 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -253,35 +253,30 @@ def _after(self, event, context): else: self.inferred_span.finish() + if not self.flush_to_log or should_use_extension: + flush_stats() + if should_use_extension: + flush_extension() + if ( self.encode_authorizer_context and self.response and self.response.get("principalId") and self.response.get("policyDocument") ): - finish_time_ns = None # the finish_time_ns should be set as the end of the inferred span if it exist # or the end of the current span - if not self.inferred_span: - finish_time_ns = ( - self.span.start_ns + self.span.duration_ns - if self.span and self.span.duration_ns - else time_ns() - ) - else: - finish_time_ns = ( - self.span.start_ns - if InferredSpanInfo.is_async(self.inferred_span) and self.span - else time_ns() - ) - + reference_span = self.inferred_span if self.inferred_span else self.span + finish_time_ns = ( + reference_span.start_ns + reference_span.duration_ns + if reference_span is not None + and hasattr(reference_span, "start_ns") + and hasattr(reference_span, "duration_ns") + else time_ns() + ) self._inject_authorizer_span_headers( event.get("requestContext", {}).get("requestId"), finish_time_ns ) - if not self.flush_to_log or should_use_extension: - flush_stats() - if should_use_extension: - flush_extension() logger.debug("datadog_lambda_wrapper _after() done") except Exception: traceback.print_exc() From efce8eb2f7ed42d7d985709cf7dbe85de4b95d0a Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Wed, 26 Oct 2022 17:02:02 -0400 Subject: [PATCH 15/26] fix the authorizer span's start time --- datadog_lambda/tracing.py | 31 +++++++++++++++---------------- datadog_lambda/wrapper.py | 2 +- 2 files changed, 16 insertions(+), 17 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 51205bcb..b500333d 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -643,7 +643,7 @@ def is_api_gateway_invocation_async(event): def insert_upstream_authorizer_span( - kwargs_to_start_span, other_tags_for_span, start_time_s, finish_time_s=None + kwargs_to_start_span, other_tags_for_span, start_time_ns, finish_time_s=None ): """Insert the authorizer span. Without this: parent span --child-> inferred span @@ -652,7 +652,7 @@ def insert_upstream_authorizer_span( Args: kwargs_to_start_span (Dict): the same keyword arguments used for the inferred span other_tags_for_span (Dict): the same tag keyword arguments used for the inferred span - start_time_s (int): the start time of the span in seconds + start_time_ns (int): the start time of the span in nanoseconds finish_time_s (int): the finish time of the sapn in seconds """ @@ -663,7 +663,7 @@ def insert_upstream_authorizer_span( upstream_authorizer_span.set_tag("operation_name", "aws.apigateway.authorizer") # always sync for the authorizer invocation InferredSpanInfo.set_tags_to_span(upstream_authorizer_span, synchronicity="sync") - upstream_authorizer_span.start = start_time_s + upstream_authorizer_span.start_ns = start_time_ns upstream_authorizer_span.finish(finish_time=finish_time_s) return upstream_authorizer_span @@ -704,9 +704,8 @@ def create_inferred_span_from_api_gateway_websocket_event(event, context): ) if injected_authorizer_data: try: - start_time_s = ( - int(injected_authorizer_data.get(Headers.Parent_Span_Finish_Time)) - / 1000 + start_time_ns = int( + injected_authorizer_data.get(Headers.Parent_Span_Finish_Time) ) finish_time_s = ( request_time_epoch_s @@ -720,7 +719,7 @@ def create_inferred_span_from_api_gateway_websocket_event(event, context): / 1000 ) upstream_authorizer_span = insert_upstream_authorizer_span( - args, tags, start_time_s, finish_time_s + args, tags, start_time_ns, finish_time_s ) # trace context needs to be set again as it is reset by upstream_authorizer_span.finish tracer.context_provider.activate(trace_ctx) @@ -781,9 +780,8 @@ def create_inferred_span_from_api_gateway_event(event, context): ) if injected_authorizer_data: try: - start_time_s = ( - int(injected_authorizer_data.get(Headers.Parent_Span_Finish_Time)) - / 1000 + start_time_ns = int( + injected_authorizer_data.get(Headers.Parent_Span_Finish_Time) ) finish_time_s = ( request_time_epoch_s @@ -797,7 +795,7 @@ def create_inferred_span_from_api_gateway_event(event, context): / 1000 ) upstream_authorizer_span = insert_upstream_authorizer_span( - args, tags, start_time_s, finish_time_s + args, tags, start_time_ns, finish_time_s ) # trace context needs to be set again as it is reset by upstream_authorizer_span.finish tracer.context_provider.activate(trace_ctx) @@ -861,13 +859,14 @@ def create_inferred_span_from_http_api_event(event, context): ) if injected_authorizer_data: try: - start_time_s = ( - int(injected_authorizer_data.get(Headers.Parent_Span_Finish_Time)) - / 1000 + start_time_ns = int( + injected_authorizer_data.get(Headers.Parent_Span_Finish_Time) ) - finish_time_s = start_time_s # no integrationLatency info in this case + finish_time_s = ( + start_time_ns / 1000 + ) # no integrationLatency info in this case upstream_authorizer_span = insert_upstream_authorizer_span( - args, tags, start_time_s, finish_time_s + args, tags, start_time_ns, finish_time_s ) # trace context needs to be set again as it is reset by upstream_authorizer_span.finish tracer.context_provider.activate(trace_ctx) diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 3ce71d29..e45292b0 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -174,7 +174,7 @@ def _inject_authorizer_span_headers(self, request_id, finish_time_ns): HTTPPropagator.inject(source_span.context, injected_headers) # remove unused header injected_headers.pop(Headers.TAGS_HEADER_TO_DELETE, None) - injected_headers[Headers.Parent_Span_Finish_Time] = finish_time_ns / 1e6 + injected_headers[Headers.Parent_Span_Finish_Time] = finish_time_ns if request_id is not None: injected_headers[Headers.Authorizing_Request_Id] = request_id datadog_data = base64.b64encode(json.dumps(injected_headers).encode()) From cddc19411dc3fd99569c432e0b7a3d2cd0226ee9 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Wed, 26 Oct 2022 17:16:43 -0400 Subject: [PATCH 16/26] fix httpapi finish time case --- datadog_lambda/tracing.py | 2 +- tests/event_samples/authorizer-request-api-gateway-v2.json | 2 +- tests/test_tracing.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index b500333d..8c61d39d 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -863,7 +863,7 @@ def create_inferred_span_from_http_api_event(event, context): injected_authorizer_data.get(Headers.Parent_Span_Finish_Time) ) finish_time_s = ( - start_time_ns / 1000 + start_time_ns / 1e9 ) # no integrationLatency info in this case upstream_authorizer_span = insert_upstream_authorizer_span( args, tags, start_time_ns, finish_time_s diff --git a/tests/event_samples/authorizer-request-api-gateway-v2.json b/tests/event_samples/authorizer-request-api-gateway-v2.json index 2be965e7..92ce09ba 100644 --- a/tests/event_samples/authorizer-request-api-gateway-v2.json +++ b/tests/event_samples/authorizer-request-api-gateway-v2.json @@ -24,7 +24,7 @@ "apiId": "amddr1rix9", "authorizer": { "lambda": { - "_datadog": "eyJ4LWRhdGFkb2ctdHJhY2UtaWQiOiAiMTQzNTY5ODM2MTk4NTI5MzMzNTQiLCAieC1kYXRhZG9nLXBhcmVudC1pZCI6ICIxMjY1ODYyMTA4MzUwNTQxMzgwOSIsICJ4LWRhdGFkb2ctc2FtcGxpbmctcHJpb3JpdHkiOiAiMSIsICJ4LWRhdGFkb2ctcGFyZW50LXNwYW4tZmluaXNoLXRpbWUiOiAxNjY0MjI4NjM5NTMzLjc3NTQsICJ4LWRhdGFkb2ctYXV0aG9yaXppbmctcmVxdWVzdGlkIjogImFiYzEyMyJ9", + "_datadog": "eyJ4LWRhdGFkb2ctdHJhY2UtaWQiOiAiMTQzNTY5ODM2MTk4NTI5MzMzNTQiLCAieC1kYXRhZG9nLXBhcmVudC1pZCI6ICIxMjY1ODYyMTA4MzUwNTQxMzgwOSIsICJ4LWRhdGFkb2ctc2FtcGxpbmctcHJpb3JpdHkiOiAiMSIsICJ4LWRhdGFkb2ctcGFyZW50LXNwYW4tZmluaXNoLXRpbWUiOiAxNjY0MjI4NjM5NTMzNzc1NDAwLCAieC1kYXRhZG9nLWF1dGhvcml6aW5nLXJlcXVlc3RpZCI6ICJhYmMxMjMifQ==", "scope": "this is just a string" } }, diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 13481be3..cec2d31f 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -630,7 +630,7 @@ def test_create_inferred_span_from_authorizer_token_api_gateway_v2_cached_event( def test_create_inferred_span_from_authorizer_request_api_gateway_v2_event(self): event_sample_source = "authorizer-request-api-gateway-v2" - finish_time = 1664228639.533 # use the injected parent span finish time as an approximation + finish_time = 1664228639.5337753 # use the injected parent span finish time as an approximation span = self._authorizer_span_testing_items(event_sample_source, finish_time) self._basic_common_checks(span, "aws.httpapi") From 39919432831e7f098ef71ec67a280e31f1b355ca Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Thu, 27 Oct 2022 12:09:23 -0400 Subject: [PATCH 17/26] fix .vscode/settings.json --- .vscode/settings.json | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/.vscode/settings.json b/.vscode/settings.json index e58c4de2..db92ccf3 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -1,7 +1,4 @@ { "python.pythonPath": "/usr/local/bin/python3", - "python.formatting.provider": "black", - "python.testing.unittestArgs": ["-v", "-s", "./tests", "-p", "test*.py"], - "python.testing.pytestEnabled": false, - "python.testing.unittestEnabled": true + "python.formatting.provider": "black" } From 6120317bf86e418b2f4b04e53f667f4f628d07fc Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Mon, 7 Nov 2022 12:04:17 -0500 Subject: [PATCH 18/26] finish_time accuracey update to nanosecond --- datadog_lambda/tracing.py | 62 ++++++++++++++++++++------------------- tests/test_tracing.py | 6 ++-- 2 files changed, 35 insertions(+), 33 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 8c61d39d..2c23ef37 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -643,7 +643,7 @@ def is_api_gateway_invocation_async(event): def insert_upstream_authorizer_span( - kwargs_to_start_span, other_tags_for_span, start_time_ns, finish_time_s=None + kwargs_to_start_span, other_tags_for_span, start_time_ns, finish_time_ns ): """Insert the authorizer span. Without this: parent span --child-> inferred span @@ -653,7 +653,7 @@ def insert_upstream_authorizer_span( kwargs_to_start_span (Dict): the same keyword arguments used for the inferred span other_tags_for_span (Dict): the same tag keyword arguments used for the inferred span start_time_ns (int): the start time of the span in nanoseconds - finish_time_s (int): the finish time of the sapn in seconds + finish_time_ns (int): the finish time of the sapn in nanoseconds """ upstream_authorizer_span = tracer.trace( @@ -663,8 +663,8 @@ def insert_upstream_authorizer_span( upstream_authorizer_span.set_tag("operation_name", "aws.apigateway.authorizer") # always sync for the authorizer invocation InferredSpanInfo.set_tags_to_span(upstream_authorizer_span, synchronicity="sync") - upstream_authorizer_span.start_ns = start_time_ns - upstream_authorizer_span.finish(finish_time=finish_time_s) + upstream_authorizer_span.start_ns = int(start_time_ns) + upstream_authorizer_span.finish(finish_time_ns / 1e9) return upstream_authorizer_span @@ -686,7 +686,7 @@ def create_inferred_span_from_api_gateway_websocket_event(event, context): "event_type": request_context.get("eventType"), "message_direction": request_context.get("messageDirection"), } - request_time_epoch_s = int(request_context.get("requestTimeEpoch")) / 1000 + request_time_epoch_ms = int(request_context.get("requestTimeEpoch")) if is_api_gateway_invocation_async(event): InferredSpanInfo.set_tags(tags, tag_source="self", synchronicity="async") else: @@ -698,7 +698,7 @@ def create_inferred_span_from_api_gateway_websocket_event(event, context): } tracer.set_tags({"_dd.origin": "lambda"}) upstream_authorizer_span = None - finish_time_s = None + finish_time_ns = None injected_authorizer_data = get_injected_authorizer_data( event, _EventSource(EventTypes.API_GATEWAY, EventSubtypes.WEBSOCKET) ) @@ -707,8 +707,8 @@ def create_inferred_span_from_api_gateway_websocket_event(event, context): start_time_ns = int( injected_authorizer_data.get(Headers.Parent_Span_Finish_Time) ) - finish_time_s = ( - request_time_epoch_s + finish_time_ns = ( + request_time_epoch_ms + ( int( request_context.get("authorizer", {}).get( @@ -716,10 +716,9 @@ def create_inferred_span_from_api_gateway_websocket_event(event, context): ) ) ) - / 1000 - ) + ) * 1e6 upstream_authorizer_span = insert_upstream_authorizer_span( - args, tags, start_time_ns, finish_time_s + args, tags, start_time_ns, finish_time_ns ) # trace context needs to be set again as it is reset by upstream_authorizer_span.finish tracer.context_provider.activate(trace_ctx) @@ -736,8 +735,10 @@ def create_inferred_span_from_api_gateway_websocket_event(event, context): span = tracer.trace("aws.apigateway.websocket", **args) if span: span.set_tags(tags) - span.start = ( - finish_time_s if finish_time_s is not None else request_time_epoch_s + span.start_ns = int( + finish_time_ns + if finish_time_ns is not None + else request_time_epoch_ms * 1e6 ) if upstream_authorizer_span: span.parent_id = upstream_authorizer_span.span_id @@ -762,7 +763,7 @@ def create_inferred_span_from_api_gateway_event(event, context): "stage": request_context.get("stage"), "request_id": request_context.get("requestId"), } - request_time_epoch_s = int(request_context.get("requestTimeEpoch")) / 1000 + request_time_epoch_ms = int(request_context.get("requestTimeEpoch")) if is_api_gateway_invocation_async(event): InferredSpanInfo.set_tags(tags, tag_source="self", synchronicity="async") else: @@ -774,7 +775,7 @@ def create_inferred_span_from_api_gateway_event(event, context): } tracer.set_tags({"_dd.origin": "lambda"}) upstream_authorizer_span = None - finish_time_s = None + finish_time_ns = None injected_authorizer_data = get_injected_authorizer_data( event, _EventSource(EventTypes.API_GATEWAY, EventSubtypes.API_GATEWAY) ) @@ -783,8 +784,8 @@ def create_inferred_span_from_api_gateway_event(event, context): start_time_ns = int( injected_authorizer_data.get(Headers.Parent_Span_Finish_Time) ) - finish_time_s = ( - request_time_epoch_s + finish_time_ns = ( + request_time_epoch_ms + ( int( request_context.get("authorizer", {}).get( @@ -792,10 +793,9 @@ def create_inferred_span_from_api_gateway_event(event, context): ) ) ) - / 1000 - ) + ) * 1e6 upstream_authorizer_span = insert_upstream_authorizer_span( - args, tags, start_time_ns, finish_time_s + args, tags, start_time_ns, finish_time_ns ) # trace context needs to be set again as it is reset by upstream_authorizer_span.finish tracer.context_provider.activate(trace_ctx) @@ -812,8 +812,10 @@ def create_inferred_span_from_api_gateway_event(event, context): if span: span.set_tags(tags) # start time pushed by the inserted authorizer span - span.start = ( - finish_time_s if finish_time_s is not None else request_time_epoch_s + span.start_ns = int( + finish_time_ns + if finish_time_ns is not None + else request_time_epoch_ms * 1e6 ) if upstream_authorizer_span: span.parent_id = upstream_authorizer_span.span_id @@ -841,7 +843,7 @@ def create_inferred_span_from_http_api_event(event, context): "apiname": request_context.get("apiId"), "stage": request_context.get("stage"), } - request_time_epoch_s = int(request_context.get("timeEpoch")) / 1000 + request_time_epoch_ms = int(request_context.get("timeEpoch")) if is_api_gateway_invocation_async(event): InferredSpanInfo.set_tags(tags, tag_source="self", synchronicity="async") else: @@ -853,7 +855,7 @@ def create_inferred_span_from_http_api_event(event, context): } tracer.set_tags({"_dd.origin": "lambda"}) upstream_authorizer_span = None - finish_time_s = None + finish_time_ns = None injected_authorizer_data = get_injected_authorizer_data( event, _EventSource(EventTypes.API_GATEWAY, EventSubtypes.HTTP_API) ) @@ -862,11 +864,9 @@ def create_inferred_span_from_http_api_event(event, context): start_time_ns = int( injected_authorizer_data.get(Headers.Parent_Span_Finish_Time) ) - finish_time_s = ( - start_time_ns / 1e9 - ) # no integrationLatency info in this case + finish_time_ns = start_time_ns upstream_authorizer_span = insert_upstream_authorizer_span( - args, tags, start_time_ns, finish_time_s + args, tags, start_time_ns, finish_time_ns ) # trace context needs to be set again as it is reset by upstream_authorizer_span.finish tracer.context_provider.activate(trace_ctx) @@ -881,8 +881,10 @@ def create_inferred_span_from_http_api_event(event, context): span = tracer.trace("aws.httpapi", **args) if span: span.set_tags(tags) - span.start = ( - finish_time_s if finish_time_s is not None else request_time_epoch_s + span.start_ns = int( + finish_time_ns + if finish_time_ns is not None + else request_time_epoch_ms * 1e6 ) if upstream_authorizer_span: span.parent_id = upstream_authorizer_span.span_id diff --git a/tests/test_tracing.py b/tests/test_tracing.py index cec2d31f..c512a135 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -685,7 +685,7 @@ def _authorizer_span_testing_items(self, event_sample_source, finish_time): # checking the upstream_authorizer_span self.mock_span_stop.assert_called_once() args, kwargs = self.mock_span_stop.call_args_list[0] - self.assertEqual(kwargs["finish_time"], finish_time) + self.assertEqual(kwargs.get("finish_time", args[1]), finish_time) self.assertEqual(span.start, finish_time) authorizer_span = args[0] self.assertEqual(authorizer_span.name, "aws.apigateway.authorizer") @@ -786,7 +786,7 @@ def test_create_inferred_span_from_api_gateway_non_proxy_event_async(self): self.assertEqual(span.get_tag("apiid"), "lgxbo6a518") self.assertEqual(span.get_tag("apiname"), "lgxbo6a518") self.assertEqual(span.get_tag("stage"), "dev") - self.assertEqual(span.start, 1631210915.2509997) + self.assertEqual(span.start, 1631210915.2510002) self.assertEqual(span.span_type, "http") self.assertEqual(span.get_tag(InferredSpanInfo.TAG_SOURCE), "self") self.assertEqual(span.get_tag(InferredSpanInfo.SYNCHRONICITY), "async") @@ -820,7 +820,7 @@ def test_create_inferred_span_from_api_gateway_non_proxy_event_sync(self): self.assertEqual(span.get_tag("apiid"), "lgxbo6a518") self.assertEqual(span.get_tag("apiname"), "lgxbo6a518") self.assertEqual(span.get_tag("stage"), "dev") - self.assertEqual(span.start, 1631210915.2509997) + self.assertEqual(span.start, 1631210915.2510002) self.assertEqual(span.span_type, "http") self.assertEqual(span.get_tag(InferredSpanInfo.TAG_SOURCE), "self") self.assertEqual(span.get_tag(InferredSpanInfo.SYNCHRONICITY), "sync") From 11de76dea01ffa22ce8047484a2c1215aff191ac Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Mon, 7 Nov 2022 15:08:39 -0500 Subject: [PATCH 19/26] test_no_error_with_nonetype_headers fix --- datadog_lambda/tracing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 2c23ef37..17cfa1e1 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -193,7 +193,7 @@ def extract_context_from_http_event_or_context( Continue without injecting the authorizer span %s", e, ) - headers = event.get("headers", {}) + headers = event.get("headers", {}) or {} lowercase_headers = {k.lower(): v for k, v in headers.items()} trace_id = lowercase_headers.get(TraceHeader.TRACE_ID) From 75fc3270d9f8f52b327ec2c50844860eeb290247 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Mon, 7 Nov 2022 15:25:29 -0500 Subject: [PATCH 20/26] fix test_no_error_with_nonetype_headers --- tests/test_tracing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index f3e6eb1b..619ae154 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -1383,7 +1383,7 @@ def test_mark_trace_as_error_for_5xx_responses_sends_error_metric_and_set_error_ def test_no_error_with_nonetype_headers(self): lambda_ctx = get_mock_context() - ctx, source = extract_dd_trace_context( + ctx, source, event_type = extract_dd_trace_context( {"headers": None}, lambda_ctx, ) From 11ea7728728530efed41bcceaf218bbf5c6844c8 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Wed, 9 Nov 2022 13:56:15 -0500 Subject: [PATCH 21/26] add decode_authorizer_context option and no zero-length span for httpapi case --- datadog_lambda/tracing.py | 209 +++++++++++++++++--------------------- datadog_lambda/wrapper.py | 5 +- tests/test_tracing.py | 13 ++- 3 files changed, 108 insertions(+), 119 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 17cfa1e1..0dc4740b 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -171,7 +171,7 @@ def extract_context_from_lambda_context(lambda_context): def extract_context_from_http_event_or_context( - event, lambda_context, event_source: _EventSource + event, lambda_context, event_source: _EventSource, decode_authorizer_context: bool = True ): """ Extract Datadog trace context from the `headers` key in from the Lambda @@ -179,20 +179,22 @@ def extract_context_from_http_event_or_context( Falls back to lambda context if no trace data is found in the `headers` """ - injected_authorizer_data = get_injected_authorizer_data(event, event_source) - if injected_authorizer_data: - try: - # fail fast on any KeyError here - trace_id = injected_authorizer_data[TraceHeader.TRACE_ID] - parent_id = injected_authorizer_data[TraceHeader.PARENT_ID] - sampling_priority = injected_authorizer_data[TraceHeader.SAMPLING_PRIORITY] - return trace_id, parent_id, sampling_priority - except Exception as e: - logger.debug( - "extract_context_from_authorizer_event returned with error. \ - Continue without injecting the authorizer span %s", - e, - ) + if decode_authorizer_context: + injected_authorizer_data = get_injected_authorizer_data(event, event_source) + if injected_authorizer_data: + try: + # fail fast on any KeyError here + trace_id = injected_authorizer_data[TraceHeader.TRACE_ID] + parent_id = injected_authorizer_data[TraceHeader.PARENT_ID] + sampling_priority = injected_authorizer_data[TraceHeader.SAMPLING_PRIORITY] + return trace_id, parent_id, sampling_priority + except Exception as e: + logger.debug( + "extract_context_from_authorizer_event returned with error. \ + Continue without injecting the authorizer span %s", + e, + ) + headers = event.get("headers", {}) or {} lowercase_headers = {k.lower(): v for k, v in headers.items()} @@ -368,7 +370,7 @@ def get_injected_authorizer_data(event, event_source: _EventSource) -> dict: return None -def extract_dd_trace_context(event, lambda_context, extractor=None): +def extract_dd_trace_context(event, lambda_context, extractor=None, decode_authorizer_context: bool= True ): """ Extract Datadog trace context from the Lambda `event` object. @@ -391,7 +393,7 @@ def extract_dd_trace_context(event, lambda_context, extractor=None): parent_id, sampling_priority, ) = extract_context_from_http_event_or_context( - event, lambda_context, event_source + event, lambda_context, event_source, decode_authorizer_context ) elif event_source.equals(EventTypes.SNS) or event_source.equals(EventTypes.SQS): ( @@ -555,7 +557,7 @@ def set_dd_trace_py_root(trace_context_source, merge_xray_traces): ) -def create_inferred_span(event, context, event_source: _EventSource = None): +def create_inferred_span(event, context, event_source: _EventSource = None, decode_authorizer_context: bool = True): if event_source is None: event_source = parse_event_source(event) try: @@ -563,7 +565,7 @@ def create_inferred_span(event, context, event_source: _EventSource = None): EventTypes.API_GATEWAY, subtype=EventSubtypes.API_GATEWAY ): logger.debug("API Gateway event detected. Inferring a span") - return create_inferred_span_from_api_gateway_event(event, context) + return create_inferred_span_from_api_gateway_event(event, decode_authorizer_context) elif event_source.equals(EventTypes.LAMBDA_FUNCTION_URL): logger.debug("Function URL event detected. Inferring a span") return create_inferred_span_from_lambda_function_url_event(event, context) @@ -571,12 +573,12 @@ def create_inferred_span(event, context, event_source: _EventSource = None): EventTypes.API_GATEWAY, subtype=EventSubtypes.HTTP_API ): logger.debug("HTTP API event detected. Inferring a span") - return create_inferred_span_from_http_api_event(event, context) + return create_inferred_span_from_http_api_event(event, context, decode_authorizer_context) elif event_source.equals( EventTypes.API_GATEWAY, subtype=EventSubtypes.WEBSOCKET ): logger.debug("API Gateway Websocket event detected. Inferring a span") - return create_inferred_span_from_api_gateway_websocket_event(event, context) + return create_inferred_span_from_api_gateway_websocket_event(event, decode_authorizer_context) elif event_source.equals(EventTypes.SQS): logger.debug("SQS event detected. Inferring a span") return create_inferred_span_from_sqs_event(event, context) @@ -668,7 +670,7 @@ def insert_upstream_authorizer_span( return upstream_authorizer_span -def create_inferred_span_from_api_gateway_websocket_event(event, context): +def create_inferred_span_from_api_gateway_websocket_event(event, decode_authorizer_context: bool = True): trace_ctx = tracer.current_trace_context() request_context = event.get("requestContext") domain = request_context.get("domainName") @@ -699,39 +701,38 @@ def create_inferred_span_from_api_gateway_websocket_event(event, context): tracer.set_tags({"_dd.origin": "lambda"}) upstream_authorizer_span = None finish_time_ns = None - injected_authorizer_data = get_injected_authorizer_data( - event, _EventSource(EventTypes.API_GATEWAY, EventSubtypes.WEBSOCKET) - ) - if injected_authorizer_data: - try: - start_time_ns = int( - injected_authorizer_data.get(Headers.Parent_Span_Finish_Time) - ) - finish_time_ns = ( - request_time_epoch_ms - + ( - int( - request_context.get("authorizer", {}).get( - "integrationLatency", 0 + if decode_authorizer_context: + injected_authorizer_data = get_injected_authorizer_data( + event, _EventSource(EventTypes.API_GATEWAY, EventSubtypes.WEBSOCKET) + ) + if injected_authorizer_data: + try: + start_time_ns = int( + injected_authorizer_data.get(Headers.Parent_Span_Finish_Time) + ) + finish_time_ns = ( + request_time_epoch_ms + + ( + int( + request_context.get("authorizer", {}).get( + "integrationLatency", 0 + ) ) ) + ) * 1e6 + upstream_authorizer_span = insert_upstream_authorizer_span( + args, tags, start_time_ns, finish_time_ns + ) + # trace context needs to be set again as it is reset by upstream_authorizer_span.finish + tracer.context_provider.activate(trace_ctx) + except Exception as e: + traceback.print_exc() + logger.debug( + "Unable to insert authorizer span. Continue to generate the main span.\ + Event %s, Reason: %s", + event, + e, ) - ) * 1e6 - upstream_authorizer_span = insert_upstream_authorizer_span( - args, tags, start_time_ns, finish_time_ns - ) - # trace context needs to be set again as it is reset by upstream_authorizer_span.finish - tracer.context_provider.activate(trace_ctx) - except Exception as e: - traceback.print_exc() - logger.debug( - "Unable to insert authorizer span. Continue to generate the main span.\ - Event %s, Reason: %s", - event, - e, - ) - # trace context needs to be set again as it is reset by upstream_authorizer_span.finish - tracer.context_provider.activate(trace_ctx) span = tracer.trace("aws.apigateway.websocket", **args) if span: span.set_tags(tags) @@ -745,7 +746,7 @@ def create_inferred_span_from_api_gateway_websocket_event(event, context): return span -def create_inferred_span_from_api_gateway_event(event, context): +def create_inferred_span_from_api_gateway_event(event, decode_authorizer_context: bool = True): trace_ctx = tracer.current_trace_context() request_context = event.get("requestContext") domain = request_context.get("domainName", "") @@ -776,38 +777,38 @@ def create_inferred_span_from_api_gateway_event(event, context): tracer.set_tags({"_dd.origin": "lambda"}) upstream_authorizer_span = None finish_time_ns = None - injected_authorizer_data = get_injected_authorizer_data( - event, _EventSource(EventTypes.API_GATEWAY, EventSubtypes.API_GATEWAY) - ) - if injected_authorizer_data: - try: - start_time_ns = int( - injected_authorizer_data.get(Headers.Parent_Span_Finish_Time) - ) - finish_time_ns = ( - request_time_epoch_ms - + ( - int( - request_context.get("authorizer", {}).get( - "integrationLatency", 0 + if decode_authorizer_context: + injected_authorizer_data = get_injected_authorizer_data( + event, _EventSource(EventTypes.API_GATEWAY, EventSubtypes.API_GATEWAY) + ) + if injected_authorizer_data: + try: + start_time_ns = int( + injected_authorizer_data.get(Headers.Parent_Span_Finish_Time) + ) + finish_time_ns = ( + request_time_epoch_ms + + ( + int( + request_context.get("authorizer", {}).get( + "integrationLatency", 0 + ) ) ) + ) * 1e6 + upstream_authorizer_span = insert_upstream_authorizer_span( + args, tags, start_time_ns, finish_time_ns + ) + # trace context needs to be set again as it is reset by upstream_authorizer_span.finish + tracer.context_provider.activate(trace_ctx) + except Exception as e: + traceback.print_exc() + logger.debug( + "Unable to insert authorizer span. Continue to generate the main span. \ + Event %s, Reason: %s", + event, + e, ) - ) * 1e6 - upstream_authorizer_span = insert_upstream_authorizer_span( - args, tags, start_time_ns, finish_time_ns - ) - # trace context needs to be set again as it is reset by upstream_authorizer_span.finish - tracer.context_provider.activate(trace_ctx) - except Exception as e: - traceback.print_exc() - logger.debug( - "Unable to insert authorizer span. Continue to generate the main span. \ - Event %s, Reason: %s", - event, - e, - ) - span = tracer.trace("aws.apigateway", **args) if span: span.set_tags(tags) @@ -817,15 +818,15 @@ def create_inferred_span_from_api_gateway_event(event, context): if finish_time_ns is not None else request_time_epoch_ms * 1e6 ) + print(f"HTTP START_NS: {span.start_ns}") if upstream_authorizer_span: span.parent_id = upstream_authorizer_span.span_id return span -def create_inferred_span_from_http_api_event(event, context): - trace_ctx = tracer.current_trace_context() +def create_inferred_span_from_http_api_event(event, context, decode_authorizer_context: bool = True): request_context = event.get("requestContext") - domain = request_context.get("domainName") + domain = request_context.get("domainName", "") method = request_context.get("http", {}).get("method") path = event.get("rawPath") resource = "{0} {1}".format(method, path) @@ -854,40 +855,16 @@ def create_inferred_span_from_http_api_event(event, context): "span_type": "http", } tracer.set_tags({"_dd.origin": "lambda"}) - upstream_authorizer_span = None - finish_time_ns = None - injected_authorizer_data = get_injected_authorizer_data( - event, _EventSource(EventTypes.API_GATEWAY, EventSubtypes.HTTP_API) - ) - if injected_authorizer_data: - try: - start_time_ns = int( - injected_authorizer_data.get(Headers.Parent_Span_Finish_Time) - ) - finish_time_ns = start_time_ns - upstream_authorizer_span = insert_upstream_authorizer_span( - args, tags, start_time_ns, finish_time_ns - ) - # trace context needs to be set again as it is reset by upstream_authorizer_span.finish - tracer.context_provider.activate(trace_ctx) - except Exception as e: - traceback.print_exc() - logger.debug( - "Unable to insert authorizer span. \ - Continue to generate the main span. Event %s, Reason: %s", - event, - e, - ) span = tracer.trace("aws.httpapi", **args) if span: span.set_tags(tags) - span.start_ns = int( - finish_time_ns - if finish_time_ns is not None - else request_time_epoch_ms * 1e6 - ) - if upstream_authorizer_span: - span.parent_id = upstream_authorizer_span.span_id + span.start_ns = request_time_epoch_ms * 1e6 + if decode_authorizer_context: + injected_authorizer_data = get_injected_authorizer_data( + event, _EventSource(EventTypes.API_GATEWAY, EventSubtypes.HTTP_API) + ) + if injected_authorizer_data and injected_authorizer_data.get(Headers.Parent_Span_Finish_Time): + span.start_ns = int(injected_authorizer_data.get(Headers.Parent_Span_Finish_Time)) return span diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index e45292b0..4a12a429 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -127,6 +127,9 @@ def __init__(self, func): self.encode_authorizer_context = ( os.environ.get("DD_ENCODE_AUTHORIZER_CONTEXT", "true").lower() == "true" ) + self.decode_authorizer_context = ( + os.environ.get("DD_DECODE_AUTHORIZER_CONTEXT", "true").lower() == "true" + ) self.response = None if profiling_env_var: self.prof = profiler.Profiler(env=env_env_var, service=service_env_var) @@ -189,7 +192,7 @@ def _before(self, event, context): self.trigger_tags = extract_trigger_tags(event, context) # Extract Datadog trace context and source from incoming requests dd_context, trace_context_source, event_source = extract_dd_trace_context( - event, context, extractor=self.trace_extractor + event, context, extractor=self.trace_extractor, decode_authorizer_context=self.decode_authorizer_context ) self.event_source = event_source # Create a Datadog X-Ray subsegment with the trace context diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 619ae154..e728988c 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -630,8 +630,17 @@ def test_create_inferred_span_from_authorizer_token_api_gateway_v2_cached_event( def test_create_inferred_span_from_authorizer_request_api_gateway_v2_event(self): event_sample_source = "authorizer-request-api-gateway-v2" - finish_time = 1664228639.5337753 # use the injected parent span finish time as an approximation - span = self._authorizer_span_testing_items(event_sample_source, finish_time) + finish_time = 1664228639533775400 # use the injected parent span finish time as an approximation + test_file = event_samples + event_sample_source + ".json" + with open(test_file, "r") as event: + event = json.load(event) + ctx = get_mock_context() + ctx.aws_request_id = "abc123" + span = create_inferred_span(event, ctx) + self.assertEqual(span.get_tag(InferredSpanInfo.TAG_SOURCE), "self") + self.assertEqual(span.get_tag(InferredSpanInfo.SYNCHRONICITY), "sync") + self.mock_span_stop.assert_not_called() + self.assertEqual(span.start_ns, finish_time) self._basic_common_checks(span, "aws.httpapi") def test_create_inferred_span_from_authorizer_request_api_gateway_v2_cached_event( From 0976510f680db9efdc2de96780e38cb3620deee4 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Wed, 9 Nov 2022 14:33:04 -0500 Subject: [PATCH 22/26] format and unittest fix --- datadog_lambda/tracing.py | 53 +++++++++++++++++++++++++++++---------- datadog_lambda/wrapper.py | 5 +++- tests/test_wrapper.py | 2 +- 3 files changed, 45 insertions(+), 15 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 0dc4740b..a6d9973e 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -171,7 +171,10 @@ def extract_context_from_lambda_context(lambda_context): def extract_context_from_http_event_or_context( - event, lambda_context, event_source: _EventSource, decode_authorizer_context: bool = True + event, + lambda_context, + event_source: _EventSource, + decode_authorizer_context: bool = True, ): """ Extract Datadog trace context from the `headers` key in from the Lambda @@ -186,7 +189,9 @@ def extract_context_from_http_event_or_context( # fail fast on any KeyError here trace_id = injected_authorizer_data[TraceHeader.TRACE_ID] parent_id = injected_authorizer_data[TraceHeader.PARENT_ID] - sampling_priority = injected_authorizer_data[TraceHeader.SAMPLING_PRIORITY] + sampling_priority = injected_authorizer_data[ + TraceHeader.SAMPLING_PRIORITY + ] return trace_id, parent_id, sampling_priority except Exception as e: logger.debug( @@ -370,7 +375,9 @@ def get_injected_authorizer_data(event, event_source: _EventSource) -> dict: return None -def extract_dd_trace_context(event, lambda_context, extractor=None, decode_authorizer_context: bool= True ): +def extract_dd_trace_context( + event, lambda_context, extractor=None, decode_authorizer_context: bool = True +): """ Extract Datadog trace context from the Lambda `event` object. @@ -557,7 +564,12 @@ def set_dd_trace_py_root(trace_context_source, merge_xray_traces): ) -def create_inferred_span(event, context, event_source: _EventSource = None, decode_authorizer_context: bool = True): +def create_inferred_span( + event, + context, + event_source: _EventSource = None, + decode_authorizer_context: bool = True, +): if event_source is None: event_source = parse_event_source(event) try: @@ -565,7 +577,9 @@ def create_inferred_span(event, context, event_source: _EventSource = None, deco EventTypes.API_GATEWAY, subtype=EventSubtypes.API_GATEWAY ): logger.debug("API Gateway event detected. Inferring a span") - return create_inferred_span_from_api_gateway_event(event, decode_authorizer_context) + return create_inferred_span_from_api_gateway_event( + event, decode_authorizer_context + ) elif event_source.equals(EventTypes.LAMBDA_FUNCTION_URL): logger.debug("Function URL event detected. Inferring a span") return create_inferred_span_from_lambda_function_url_event(event, context) @@ -573,12 +587,16 @@ def create_inferred_span(event, context, event_source: _EventSource = None, deco EventTypes.API_GATEWAY, subtype=EventSubtypes.HTTP_API ): logger.debug("HTTP API event detected. Inferring a span") - return create_inferred_span_from_http_api_event(event, context, decode_authorizer_context) + return create_inferred_span_from_http_api_event( + event, context, decode_authorizer_context + ) elif event_source.equals( EventTypes.API_GATEWAY, subtype=EventSubtypes.WEBSOCKET ): logger.debug("API Gateway Websocket event detected. Inferring a span") - return create_inferred_span_from_api_gateway_websocket_event(event, decode_authorizer_context) + return create_inferred_span_from_api_gateway_websocket_event( + event, decode_authorizer_context + ) elif event_source.equals(EventTypes.SQS): logger.debug("SQS event detected. Inferring a span") return create_inferred_span_from_sqs_event(event, context) @@ -670,7 +688,9 @@ def insert_upstream_authorizer_span( return upstream_authorizer_span -def create_inferred_span_from_api_gateway_websocket_event(event, decode_authorizer_context: bool = True): +def create_inferred_span_from_api_gateway_websocket_event( + event, decode_authorizer_context: bool = True +): trace_ctx = tracer.current_trace_context() request_context = event.get("requestContext") domain = request_context.get("domainName") @@ -746,7 +766,9 @@ def create_inferred_span_from_api_gateway_websocket_event(event, decode_authoriz return span -def create_inferred_span_from_api_gateway_event(event, decode_authorizer_context: bool = True): +def create_inferred_span_from_api_gateway_event( + event, decode_authorizer_context: bool = True +): trace_ctx = tracer.current_trace_context() request_context = event.get("requestContext") domain = request_context.get("domainName", "") @@ -818,13 +840,14 @@ def create_inferred_span_from_api_gateway_event(event, decode_authorizer_context if finish_time_ns is not None else request_time_epoch_ms * 1e6 ) - print(f"HTTP START_NS: {span.start_ns}") if upstream_authorizer_span: span.parent_id = upstream_authorizer_span.span_id return span -def create_inferred_span_from_http_api_event(event, context, decode_authorizer_context: bool = True): +def create_inferred_span_from_http_api_event( + event, context, decode_authorizer_context: bool = True +): request_context = event.get("requestContext") domain = request_context.get("domainName", "") method = request_context.get("http", {}).get("method") @@ -863,8 +886,12 @@ def create_inferred_span_from_http_api_event(event, context, decode_authorizer_c injected_authorizer_data = get_injected_authorizer_data( event, _EventSource(EventTypes.API_GATEWAY, EventSubtypes.HTTP_API) ) - if injected_authorizer_data and injected_authorizer_data.get(Headers.Parent_Span_Finish_Time): - span.start_ns = int(injected_authorizer_data.get(Headers.Parent_Span_Finish_Time)) + if injected_authorizer_data and injected_authorizer_data.get( + Headers.Parent_Span_Finish_Time + ): + span.start_ns = int( + injected_authorizer_data.get(Headers.Parent_Span_Finish_Time) + ) return span diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 4a12a429..d807f536 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -192,7 +192,10 @@ def _before(self, event, context): self.trigger_tags = extract_trigger_tags(event, context) # Extract Datadog trace context and source from incoming requests dd_context, trace_context_source, event_source = extract_dd_trace_context( - event, context, extractor=self.trace_extractor, decode_authorizer_context=self.decode_authorizer_context + event, + context, + extractor=self.trace_extractor, + decode_authorizer_context=self.decode_authorizer_context, ) self.event_source = event_source # Create a Datadog X-Ray subsegment with the trace context diff --git a/tests/test_wrapper.py b/tests/test_wrapper.py index 6240ec4f..250891c4 100644 --- a/tests/test_wrapper.py +++ b/tests/test_wrapper.py @@ -109,7 +109,7 @@ def lambda_handler(event, context): ] ) self.mock_extract_dd_trace_context.assert_called_with( - lambda_event, lambda_context, extractor=None + lambda_event, lambda_context, extractor=None, decode_authorizer_context=True ) self.mock_set_correlation_ids.assert_called() self.mock_inject_correlation_ids.assert_called() From 079f5942d597687e57e4977516dce8dff3b46c86 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Wed, 9 Nov 2022 14:40:00 -0500 Subject: [PATCH 23/26] lint tests --- tests/test_tracing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_tracing.py b/tests/test_tracing.py index e728988c..47a028da 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -630,7 +630,7 @@ def test_create_inferred_span_from_authorizer_token_api_gateway_v2_cached_event( def test_create_inferred_span_from_authorizer_request_api_gateway_v2_event(self): event_sample_source = "authorizer-request-api-gateway-v2" - finish_time = 1664228639533775400 # use the injected parent span finish time as an approximation + finish_time = 1664228639533775400 # use the injected parent span finish time as an approximation test_file = event_samples + event_sample_source + ".json" with open(test_file, "r") as event: event = json.load(event) From bfddf6e27ff1927e3c0ed5499b4b4860c02e11e5 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Wed, 9 Nov 2022 15:00:06 -0500 Subject: [PATCH 24/26] small refactor and reformat --- datadog_lambda/tracing.py | 10 +++------- .../snapshots/logs/sync-metrics_python39.log | 2 +- 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index a6d9973e..acbc03b9 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -675,7 +675,7 @@ def insert_upstream_authorizer_span( start_time_ns (int): the start time of the span in nanoseconds finish_time_ns (int): the finish time of the sapn in nanoseconds """ - + trace_ctx = tracer.current_trace_context() upstream_authorizer_span = tracer.trace( "aws.apigateway.authorizer", **kwargs_to_start_span ) @@ -685,13 +685,14 @@ def insert_upstream_authorizer_span( InferredSpanInfo.set_tags_to_span(upstream_authorizer_span, synchronicity="sync") upstream_authorizer_span.start_ns = int(start_time_ns) upstream_authorizer_span.finish(finish_time_ns / 1e9) + # trace context needs to be set again as it is reset by finish() + tracer.context_provider.activate(trace_ctx) return upstream_authorizer_span def create_inferred_span_from_api_gateway_websocket_event( event, decode_authorizer_context: bool = True ): - trace_ctx = tracer.current_trace_context() request_context = event.get("requestContext") domain = request_context.get("domainName") endpoint = request_context.get("routeKey") @@ -743,8 +744,6 @@ def create_inferred_span_from_api_gateway_websocket_event( upstream_authorizer_span = insert_upstream_authorizer_span( args, tags, start_time_ns, finish_time_ns ) - # trace context needs to be set again as it is reset by upstream_authorizer_span.finish - tracer.context_provider.activate(trace_ctx) except Exception as e: traceback.print_exc() logger.debug( @@ -769,7 +768,6 @@ def create_inferred_span_from_api_gateway_websocket_event( def create_inferred_span_from_api_gateway_event( event, decode_authorizer_context: bool = True ): - trace_ctx = tracer.current_trace_context() request_context = event.get("requestContext") domain = request_context.get("domainName", "") method = event.get("httpMethod") @@ -821,8 +819,6 @@ def create_inferred_span_from_api_gateway_event( upstream_authorizer_span = insert_upstream_authorizer_span( args, tags, start_time_ns, finish_time_ns ) - # trace context needs to be set again as it is reset by upstream_authorizer_span.finish - tracer.context_provider.activate(trace_ctx) except Exception as e: traceback.print_exc() logger.debug( diff --git a/tests/integration/snapshots/logs/sync-metrics_python39.log b/tests/integration/snapshots/logs/sync-metrics_python39.log index d4f3c9c7..fb7dd5b7 100644 --- a/tests/integration/snapshots/logs/sync-metrics_python39.log +++ b/tests/integration/snapshots/logs/sync-metrics_python39.log @@ -1478,7 +1478,6 @@ HTTP GET https://www.datadoghq.com/ Headers: ["Accept-Encoding:gzip, deflate", " ] } HTTP POST https://api.datadoghq.com/api/v1/distribution_points?api_key=XXXX Headers: ["Accept-Encoding:gzip, deflate", "Accept:*/*", "Connection:keep-alive", "Content-Encoding:deflate", "Content-Length:XXXX", "Content-Type:application/json", "User-Agent:datadogpy/XX (python XX; os linux; arch x86_64)", "x-datadog-parent-id:XXXX", "x-datadog-sampling-priority:1", "x-datadog-tags:_dd.p.dm=-0", "x-datadog-trace-id:XXXX"] Data: {"series": [{"metric": "hello.dog", "points": [[XXXX, [1.0]]], "type": "distribution", "host": null, "device": null, "tags": ["team:serverless", "role:hello", "dd_lambda_layer:datadog-python39_X.X.X"], "interval": 10}, {"metric": "tests.integration.count", "points": [[XXXX, [21.0]]], "type": "distribution", "host": null, "device": null, "tags": ["test:integration", "role:hello", "dd_lambda_layer:datadog-python39_X.X.X"], "interval": 10}]} -END Duration: XXXX ms Memory Used: XXXX MB { "traces": [ [ @@ -1513,3 +1512,4 @@ END Duration: XXXX ms Memory Used: XXXX MB ] ] } +END Duration: XXXX ms Memory Used: XXXX MB From 7108dfcfc956e72e3f60ff378c077d5ceec9c56d Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Thu, 10 Nov 2022 19:20:02 -0500 Subject: [PATCH 25/26] add an unittest case for the encoding part and fix some bugs --- datadog_lambda/tracing.py | 24 +++++++++++------------ datadog_lambda/wrapper.py | 2 +- tests/test_wrapper.py | 40 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 53 insertions(+), 13 deletions(-) diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index acbc03b9..47fc969d 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -845,7 +845,7 @@ def create_inferred_span_from_http_api_event( event, context, decode_authorizer_context: bool = True ): request_context = event.get("requestContext") - domain = request_context.get("domainName", "") + domain = request_context.get("domainName") method = request_context.get("http", {}).get("method") path = event.get("rawPath") resource = "{0} {1}".format(method, path) @@ -874,20 +874,20 @@ def create_inferred_span_from_http_api_event( "span_type": "http", } tracer.set_tags({"_dd.origin": "lambda"}) + inferred_span_start_ns = None + if decode_authorizer_context: + injected_authorizer_data = get_injected_authorizer_data( + event, _EventSource(EventTypes.API_GATEWAY, EventSubtypes.HTTP_API) + ) + if injected_authorizer_data: + inferred_span_start_ns = int( + injected_authorizer_data.get(Headers.Parent_Span_Finish_Time) + ) + inferred_span_start_ns = inferred_span_start_ns or request_time_epoch_ms * 1e6 span = tracer.trace("aws.httpapi", **args) if span: span.set_tags(tags) - span.start_ns = request_time_epoch_ms * 1e6 - if decode_authorizer_context: - injected_authorizer_data = get_injected_authorizer_data( - event, _EventSource(EventTypes.API_GATEWAY, EventSubtypes.HTTP_API) - ) - if injected_authorizer_data and injected_authorizer_data.get( - Headers.Parent_Span_Finish_Time - ): - span.start_ns = int( - injected_authorizer_data.get(Headers.Parent_Span_Finish_Time) - ) + span.start_ns = int(inferred_span_start_ns) return span diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index d807f536..028b7410 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -208,7 +208,7 @@ def _before(self, event, context): set_dd_trace_py_root(trace_context_source, self.merge_xray_traces) if self.make_inferred_span: self.inferred_span = create_inferred_span( - event, context, event_source + event, context, event_source, self.decode_authorizer_context ) self.span = create_function_execution_span( context, diff --git a/tests/test_wrapper.py b/tests/test_wrapper.py index 250891c4..4efadd96 100644 --- a/tests/test_wrapper.py +++ b/tests/test_wrapper.py @@ -1,7 +1,10 @@ +import base64 +import json import os import unittest from unittest.mock import patch, call, ANY, MagicMock +from datadog_lambda.constants import TraceHeader from datadog_lambda.wrapper import datadog_lambda_wrapper from datadog_lambda.metric import lambda_metric @@ -504,3 +507,40 @@ def lambda_handler(event, context): self.assertEqual(os.environ.get("DD_REQUESTS_SERVICE_NAME"), "myAwesomeService") del os.environ["DD_SERVICE"] + + def test_encode_authorizer_span(self): + @datadog_lambda_wrapper + def lambda_handler(event, context): + return { + "principalId": "foo", + "policyDocument": { + "Version": "2012-10-17", + "Statement": [ + { + "Action": "execute-api:Invoke", + "Effect": "Allow", + "Resource": "dummy", + }, + ], + }, + "context": {"scope": "still here"}, + } + + lambda_event = {} + + lambda_context = get_mock_context() + mock_span = MagicMock() + mock_span.context.span_id = "123" + mock_span.context.trace_id = "456" + mock_span.context.sampling_priority = "1" + mock_span.context.dd_origin = None + mock_span.start_ns = 1668127541671386817 + mock_span.duration_ns = 1e8 + lambda_handler.inferred_span = mock_span + + result = lambda_handler(lambda_event, lambda_context) + inject_data = json.loads(base64.b64decode(result["context"]["_datadog"])) + self.assertEquals(inject_data[TraceHeader.PARENT_ID], "123") + self.assertEquals(inject_data[TraceHeader.TRACE_ID], "456") + self.assertEquals(inject_data[TraceHeader.SAMPLING_PRIORITY], "1") + self.assertEquals(result["context"]["scope"], "still here") From af5d190e9ff6183c9b3ecd37d12461f68015a7c7 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Tue, 15 Nov 2022 11:42:55 -0500 Subject: [PATCH 26/26] Addressing comments. --- datadog_lambda/constants.py | 4 ++ datadog_lambda/tracing.py | 128 ++++++++++++++++-------------------- datadog_lambda/wrapper.py | 25 +++---- 3 files changed, 74 insertions(+), 83 deletions(-) diff --git a/datadog_lambda/constants.py b/datadog_lambda/constants.py index 23bce904..5fb6683d 100644 --- a/datadog_lambda/constants.py +++ b/datadog_lambda/constants.py @@ -46,6 +46,10 @@ class XrayDaemon(object): class Headers(object): Parent_Span_Finish_Time = "x-datadog-parent-span-finish-time" + + # For one request from the client, the event.requestContext.requestIds in the authorizer lambda + # invocation and the main function invocation are IDENTICAL. Therefore we can use it to tell + # whether current invocation is the actual original authorizing request or a cached request. Authorizing_Request_Id = "x-datadog-authorizing-requestid" # injected by the HTTPPropagator.inject but no use diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 47fc969d..62c64c61 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -8,7 +8,6 @@ import json import base64 from datetime import datetime, timezone -import traceback from typing import Optional, Dict from datadog_lambda.metric import submit_errors_metric @@ -183,7 +182,10 @@ def extract_context_from_http_event_or_context( Falls back to lambda context if no trace data is found in the `headers` """ if decode_authorizer_context: - injected_authorizer_data = get_injected_authorizer_data(event, event_source) + is_http_api = event_source.equals( + EventTypes.API_GATEWAY, subtype=EventSubtypes.HTTP_API + ) + injected_authorizer_data = get_injected_authorizer_data(event, is_http_api) if injected_authorizer_data: try: # fail fast on any KeyError here @@ -343,7 +345,7 @@ def extract_context_custom_extractor(extractor, event, lambda_context): return None, None, None -def get_injected_authorizer_data(event, event_source: _EventSource) -> dict: +def get_injected_authorizer_data(event, is_http_api) -> dict: try: authorizer_headers = event.get("requestContext", {}).get("authorizer") if not authorizer_headers: @@ -351,9 +353,7 @@ def get_injected_authorizer_data(event, event_source: _EventSource) -> dict: dd_data_raw = ( authorizer_headers.get("lambda", {}).get("_datadog") - if event_source.equals( - EventTypes.API_GATEWAY, subtype=EventSubtypes.HTTP_API - ) + if is_http_api else authorizer_headers.get("_datadog") ) @@ -361,8 +361,14 @@ def get_injected_authorizer_data(event, event_source: _EventSource) -> dict: return None injected_data = json.loads(base64.b64decode(dd_data_raw)) - # use the integrationLatency (needed for token type authorizer) or the injected requestId - # to tell if it's the authorizing invocation (not cached) + + # Lambda authorizer's results can be cached. But the payload will still have the injected + # data in cached requests. How to distinguish cached case and ignore the injected data ? + # APIGateway automatically injects a integrationLatency data in some cases. If it's >0 we + # know that it's not cached. But integrationLatency is not available for Http API case. In + # that case, we use the injected Authorizing_Request_Id to tell if it's cached. But token + # authorizers don't pass on the requestId. The Authorizing_Request_Id can't work for all + # cases neither. As a result, we combine both methods as shown below. if authorizer_headers.get("integrationLatency", 0) > 0 or event.get( "requestContext", {} ).get("requestId") == injected_data.get(Headers.Authorizing_Request_Id): @@ -690,6 +696,41 @@ def insert_upstream_authorizer_span( return upstream_authorizer_span +def process_injected_data(event, request_time_epoch_ms, args, tags): + """ + This covers the ApiGateway RestAPI and Websocket cases. It doesn't cover Http API cases. + """ + injected_authorizer_data = get_injected_authorizer_data(event, False) + if injected_authorizer_data: + try: + start_time_ns = int( + injected_authorizer_data.get(Headers.Parent_Span_Finish_Time) + ) + finish_time_ns = ( + request_time_epoch_ms + + ( + int( + event["requestContext"]["authorizer"].get( + "integrationLatency", 0 + ) + ) + ) + ) * 1e6 + upstream_authorizer_span = insert_upstream_authorizer_span( + args, tags, start_time_ns, finish_time_ns + ) + return upstream_authorizer_span, finish_time_ns + except Exception as e: + logger.debug( + "Unable to insert authorizer span. Continue to generate the main span.\ + Reason: %s", + e, + ) + return None, None + else: + return None, None + + def create_inferred_span_from_api_gateway_websocket_event( event, decode_authorizer_context: bool = True ): @@ -723,35 +764,9 @@ def create_inferred_span_from_api_gateway_websocket_event( upstream_authorizer_span = None finish_time_ns = None if decode_authorizer_context: - injected_authorizer_data = get_injected_authorizer_data( - event, _EventSource(EventTypes.API_GATEWAY, EventSubtypes.WEBSOCKET) + upstream_authorizer_span, finish_time_ns = process_injected_data( + event, request_time_epoch_ms, args, tags ) - if injected_authorizer_data: - try: - start_time_ns = int( - injected_authorizer_data.get(Headers.Parent_Span_Finish_Time) - ) - finish_time_ns = ( - request_time_epoch_ms - + ( - int( - request_context.get("authorizer", {}).get( - "integrationLatency", 0 - ) - ) - ) - ) * 1e6 - upstream_authorizer_span = insert_upstream_authorizer_span( - args, tags, start_time_ns, finish_time_ns - ) - except Exception as e: - traceback.print_exc() - logger.debug( - "Unable to insert authorizer span. Continue to generate the main span.\ - Event %s, Reason: %s", - event, - e, - ) span = tracer.trace("aws.apigateway.websocket", **args) if span: span.set_tags(tags) @@ -798,35 +813,9 @@ def create_inferred_span_from_api_gateway_event( upstream_authorizer_span = None finish_time_ns = None if decode_authorizer_context: - injected_authorizer_data = get_injected_authorizer_data( - event, _EventSource(EventTypes.API_GATEWAY, EventSubtypes.API_GATEWAY) + upstream_authorizer_span, finish_time_ns = process_injected_data( + event, request_time_epoch_ms, args, tags ) - if injected_authorizer_data: - try: - start_time_ns = int( - injected_authorizer_data.get(Headers.Parent_Span_Finish_Time) - ) - finish_time_ns = ( - request_time_epoch_ms - + ( - int( - request_context.get("authorizer", {}).get( - "integrationLatency", 0 - ) - ) - ) - ) * 1e6 - upstream_authorizer_span = insert_upstream_authorizer_span( - args, tags, start_time_ns, finish_time_ns - ) - except Exception as e: - traceback.print_exc() - logger.debug( - "Unable to insert authorizer span. Continue to generate the main span. \ - Event %s, Reason: %s", - event, - e, - ) span = tracer.trace("aws.apigateway", **args) if span: span.set_tags(tags) @@ -874,16 +863,13 @@ def create_inferred_span_from_http_api_event( "span_type": "http", } tracer.set_tags({"_dd.origin": "lambda"}) - inferred_span_start_ns = None + inferred_span_start_ns = request_time_epoch_ms * 1e6 if decode_authorizer_context: - injected_authorizer_data = get_injected_authorizer_data( - event, _EventSource(EventTypes.API_GATEWAY, EventSubtypes.HTTP_API) - ) + injected_authorizer_data = get_injected_authorizer_data(event, True) if injected_authorizer_data: - inferred_span_start_ns = int( - injected_authorizer_data.get(Headers.Parent_Span_Finish_Time) + inferred_span_start_ns = injected_authorizer_data.get( + Headers.Parent_Span_Finish_Time ) - inferred_span_start_ns = inferred_span_start_ns or request_time_epoch_ms * 1e6 span = tracer.trace("aws.httpapi", **args) if span: span.set_tags(tags) diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index 028b7410..f0d56a3d 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -171,7 +171,18 @@ def __call__(self, event, context, **kwargs): finally: self._after(event, context) - def _inject_authorizer_span_headers(self, request_id, finish_time_ns): + def _inject_authorizer_span_headers(self, request_id): + reference_span = self.inferred_span if self.inferred_span else self.span + assert reference_span.finished + # the finish_time_ns should be set as the end of the inferred span if it exist + # or the end of the current span + finish_time_ns = ( + reference_span.start_ns + reference_span.duration_ns + if reference_span is not None + and hasattr(reference_span, "start_ns") + and hasattr(reference_span, "duration_ns") + else time_ns() + ) injected_headers = {} source_span = self.inferred_span if self.inferred_span else self.span HTTPPropagator.inject(source_span.context, injected_headers) @@ -270,18 +281,8 @@ def _after(self, event, context): and self.response.get("principalId") and self.response.get("policyDocument") ): - # the finish_time_ns should be set as the end of the inferred span if it exist - # or the end of the current span - reference_span = self.inferred_span if self.inferred_span else self.span - finish_time_ns = ( - reference_span.start_ns + reference_span.duration_ns - if reference_span is not None - and hasattr(reference_span, "start_ns") - and hasattr(reference_span, "duration_ns") - else time_ns() - ) self._inject_authorizer_span_headers( - event.get("requestContext", {}).get("requestId"), finish_time_ns + event.get("requestContext", {}).get("requestId") ) logger.debug("datadog_lambda_wrapper _after() done") except Exception: