diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index f032059b..5e338253 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -68,6 +68,8 @@ propagator = HTTPPropagator() +DD_TRACE_JAVA_TRACE_ID_PADDING = "00000000" + def _convert_xray_trace_id(xray_trace_id): """ @@ -248,28 +250,52 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context): first_record.get("Sns", {}).get("MessageAttributes", {}), ) dd_payload = msg_attributes.get("_datadog", {}) - # SQS uses dataType and binaryValue/stringValue - # SNS uses Type and Value - dd_json_data_type = dd_payload.get("Type", dd_payload.get("dataType", "")) - if dd_json_data_type == "Binary": - dd_json_data = dd_payload.get( - "binaryValue", - dd_payload.get("Value", r"{}"), - ) - dd_json_data = base64.b64decode(dd_json_data) - elif dd_json_data_type == "String": - dd_json_data = dd_payload.get( - "stringValue", - dd_payload.get("Value", r"{}"), - ) + if dd_payload: + # SQS uses dataType and binaryValue/stringValue + # SNS uses Type and Value + dd_json_data = None + dd_json_data_type = dd_payload.get("Type", dd_payload.get("dataType", "")) + if dd_json_data_type == "Binary": + dd_json_data = dd_payload.get( + "binaryValue", + dd_payload.get("Value", r"{}"), + ) + dd_json_data = base64.b64decode(dd_json_data) + elif dd_json_data_type == "String": + dd_json_data = dd_payload.get( + "stringValue", + dd_payload.get("Value", r"{}"), + ) + else: + logger.debug( + "Datadog Lambda Python only supports extracting trace" + "context from String or Binary SQS/SNS message attributes" + ) + + if dd_json_data: + dd_data = json.loads(dd_json_data) + return propagator.extract(dd_data) else: - logger.debug( - "Datadog Lambda Python only supports extracting trace" - "context from String or Binary SQS/SNS message attributes" - ) - return extract_context_from_lambda_context(lambda_context) - dd_data = json.loads(dd_json_data) - return propagator.extract(dd_data) + # Handle case where trace context is injected into attributes.AWSTraceHeader + # example: Root=1-654321ab-000000001234567890abcdef;Parent=0123456789abcdef;Sampled=1 + x_ray_header = first_record.get("attributes", {}).get("AWSTraceHeader") + if x_ray_header: + x_ray_context = parse_xray_header(x_ray_header) + trace_id_parts = x_ray_context.get("trace_id", "").split("-") + if len(trace_id_parts) > 2 and trace_id_parts[2].startswith( + DD_TRACE_JAVA_TRACE_ID_PADDING + ): + # If it starts with eight 0's padding, + # then this AWSTraceHeader contains Datadog injected trace context + logger.debug( + "Found dd-trace injected trace context from AWSTraceHeader" + ) + return Context( + trace_id=int(trace_id_parts[2][8:], 16), + span_id=int(int(x_ray_context["parent_id"], 16)), + sampling_priority=float(x_ray_context["sampled"]), + ) + return extract_context_from_lambda_context(lambda_context) except Exception as e: logger.debug("The trace extractor returned with error %s", e) return extract_context_from_lambda_context(lambda_context) diff --git a/tests/event_samples/sqs-java-upstream.json b/tests/event_samples/sqs-java-upstream.json new file mode 100644 index 00000000..23ab6881 --- /dev/null +++ b/tests/event_samples/sqs-java-upstream.json @@ -0,0 +1,22 @@ +{ + "Records": [ + { + "messageId": "f7e888aa-1368-484c-8e15-fc3f0f7c6fea", + "receiptHandle": "AQEBN1aYTQ1c5huZh9bkhBYqcMMnqTUMRh8MfUPyGXkEolcn23rvM9saGEg3wTK/7JnJ1s3Uk107uLjaP6yV6+zS3oQRU0vMG2LfyTgHovWhYQ8TnrpC7XpYL+Uf+oc9KoILQopiYi4wsFnOWQqy82yQmlOA3W+CZ3Rvq8N6rNcmyaZEXVdozHG+FyMCMQ8QdTcCHhzR9YKnkZ87Y40+LhysUR57VNPVtRwENI8H1uMEfmxaCkW+CAkdCGoXeX+KioT7pHJDZaEutXM3VRmGXDDzCXvfUJQ9JQIlP5xe66JO8/cpCyl5sDoHsCjLy6X/XCmfG2+XclPObGHBzcMSjG1RQtHsEGTOAJrLREucqf/oj0Ab4svpxz6lR4UXrICygZ2x0NZcNFXcZx3GV2QL9nHmJxzrO2lnNTEOMuYB4SnqtIhsaDTcmkYHumaAJdRHl5BksFcU5qpS7BQrnRvXn5Sz3hYdR2KuYKN5Oq6W1vuT16o=", + "body": "{\"hello\":\"world\"}", + "attributes": { + "ApproximateReceiveCount": "1", + "AWSTraceHeader": "Root=1-65eb7350-000000006dfd06bf489aa4e5;Parent=48cc02b6aafae897;Sampled=1", + "SentTimestamp": "1709929297382", + "SenderId": "AROAWGCM4HXUTHYJKOM7M:DdTraceXLambda-sqsjavache-sqsjavaproducerforPython-moi7s7Hu7Ppy", + "ApproximateFirstReceiveTimestamp": "1709929297387" + }, + "messageAttributes": {}, + "md5OfMessageAttributes": "", + "md5OfBody": "fbc24bcc7a1794758fc1327fcfebdaf6", + "eventSource": "aws:sqs", + "eventSourceARN": "arn:aws:sqs:us-west-2:444442222111:DdTraceXLambda-sqsjavachecksNeste-sqsJavaProducer2sqsjavaproducerfo-dwpHQF6fcZT4", + "awsRegion": "us-west-2" + } + ] +} diff --git a/tests/test_cold_start.py b/tests/test_cold_start.py index 2ce37e7c..c7444c49 100644 --- a/tests/test_cold_start.py +++ b/tests/test_cold_start.py @@ -240,7 +240,6 @@ def test_trace_ignore_libs(self): def test_lazy_loaded_package_imports(monkeypatch): - spans = [] def finish(span): diff --git a/tests/test_tracing.py b/tests/test_tracing.py index d51f7b3f..e30d202c 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -1963,6 +1963,17 @@ def test_extract_context_from_sqs_batch_event(self): self.assertEqual(context.span_id, 7431398482019833808) self.assertEqual(context.sampling_priority, 1) + def test_extract_context_from_sqs_java_upstream_event(self): + event_sample_source = "sqs-java-upstream" + test_file = event_samples + event_sample_source + ".json" + with open(test_file, "r") as event: + event = json.load(event) + ctx = get_mock_context() + context, source, event_type = extract_dd_trace_context(event, ctx) + self.assertEqual(context.trace_id, 7925498337868555493) + self.assertEqual(context.span_id, 5245570649555658903) + self.assertEqual(context.sampling_priority, 1) + def test_extract_context_from_sns_event_with_string_msg_attr(self): event_sample_source = "sns-string-msg-attribute" test_file = event_samples + event_sample_source + ".json"