Skip to content

Check for the new direct invocation trace format and use that #146

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
Jun 17, 2021
Merged

Check for the new direct invocation trace format and use that #146

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
64eb98f
Check for the new direct invocation trace format and use that
Jun 16, 2021
6719312
run Black
Jun 16, 2021
535988a
noqa Optional
Jun 16, 2021
275903d
Do what the lint build step is telling me to do
Jun 16, 2021
72aaa2d
installed the correct version of Black
Jun 16, 2021
25e3efd
Black against the tests now
Jun 16, 2021
a10a1fe
Remove type hints
Jun 17, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 20 additions & 8 deletions datadog_lambda/tracing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import logging
import os
import json
from typing import Optional # noqa: F401

from aws_xray_sdk.core import xray_recorder
from aws_xray_sdk.core.lambda_launcher import LambdaContext
Expand Down Expand Up @@ -129,16 +130,27 @@ def extract_context_from_lambda_context(lambda_context):
dd_trace libraries inject this trace context on synchronous invocations
"""
client_context = lambda_context.client_context

trace_id = None # type: Optional[str]
parent_id = None # type: Optional[str]
sampling_priority = None # type: Optional[str]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the "str" type is a bit complicated, because it should be 'unicode' in python 2.7, and 'str' in 3+. It seems like the correct type would be 'Text' https://mypy.readthedocs.io/en/stable/cheat_sheet.html#built-in-types

Not sure it's worth introducing type hinting in this PR. I think it would make sense to start at the same time we drop 2.7 support, so we can use the more convenient 3+ syntax

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That makes sense. My instinct is to scatter type hints in as many places as possible, but it's probably better to wait until the 2 -> 3 conversion.

if client_context and client_context.custom:
dd_data = client_context.custom.get("_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)
if "_datadog" in client_context.custom:
# Legacy trace propagation dict
dd_data = client_context.custom.get("_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)
elif (
TraceHeader.TRACE_ID in client_context.custom
and TraceHeader.PARENT_ID in client_context.custom
and TraceHeader.SAMPLING_PRIORITY in client_context.custom
):
# New trace propagation keys
trace_id = client_context.custom.get(TraceHeader.TRACE_ID)
parent_id = client_context.custom.get(TraceHeader.PARENT_ID)
sampling_priority = client_context.custom.get(TraceHeader.SAMPLING_PRIORITY)

return trace_id, parent_id, sampling_priority

return None, None, None
return trace_id, parent_id, sampling_priority


def extract_context_from_http_event_or_context(event, lambda_context):
Expand Down
37 changes: 36 additions & 1 deletion tests/test_tracing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -272,7 +272,7 @@ def test_with_sqs_distributed_datadog_trace_data(self):
XraySubsegment.NAMESPACE,
)

def test_with_client_context_datadog_trace_data(self):
def test_with_legacy_client_context_datadog_trace_data(self):
lambda_ctx = get_mock_context(
custom={
"_datadog": {
Expand Down Expand Up @@ -309,6 +309,41 @@ def test_with_client_context_datadog_trace_data(self):
XraySubsegment.NAMESPACE,
)

def test_with_new_client_context_datadog_trace_data(self):
lambda_ctx = get_mock_context(
custom={
TraceHeader.TRACE_ID: "666",
TraceHeader.PARENT_ID: "777",
TraceHeader.SAMPLING_PRIORITY: "1",
}
)
ctx, source = extract_dd_trace_context({}, lambda_ctx)
self.assertEqual(source, "event")
self.assertDictEqual(
ctx,
{
"trace-id": "666",
"parent-id": "777",
"sampling-priority": "1",
},
)
self.assertDictEqual(
get_dd_trace_context(),
{
TraceHeader.TRACE_ID: "666",
TraceHeader.PARENT_ID: "65535",
TraceHeader.SAMPLING_PRIORITY: "1",
},
)
create_dd_dummy_metadata_subsegment(ctx, XraySubsegment.TRACE_KEY)
self.mock_xray_recorder.begin_subsegment.assert_called()
self.mock_xray_recorder.end_subsegment.assert_called()
self.mock_current_subsegment.put_metadata.assert_called_with(
XraySubsegment.TRACE_KEY,
{"trace-id": "666", "parent-id": "777", "sampling-priority": "1"},
XraySubsegment.NAMESPACE,
)

def test_with_complete_datadog_trace_headers_with_mixed_casing(self):
lambda_ctx = get_mock_context()
extract_dd_trace_context(
Expand Down