Skip to content

Commit 3c8060e

Browse files
maxdaymaxday
committed
lint
1 parent ab20caf commit 3c8060e

File tree

3 files changed

+58
-23
lines changed

3 files changed

+58
-23
lines changed

‎datadog_lambda/constants.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,4 +37,4 @@ class TraceContextSource(object):
3737

3838

3939
XRAY_TRACE_ID_HEADER_NAME = "_X_AMZN_TRACE_ID"
40-
FUNCTION_NAME_HEADER_NAME = "AWS_LAMBDA_FUNCTION_NAME"
40+
FUNCTION_NAME_HEADER_NAME = "AWS_LAMBDA_FUNCTION_NAME"

‎datadog_lambda/tracing.py

Lines changed: 19 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
TraceHeader,
1313
XraySubsegment,
1414
TraceContextSource,
15-
XRAY_TRACE_ID_HEADER_NAME
15+
XRAY_TRACE_ID_HEADER_NAME,
1616
)
1717
from ddtrace import tracer, patch
1818
from ddtrace import __version__ as ddtrace_version
@@ -53,7 +53,9 @@ def _convert_xray_sampling(xray_sampled):
5353

5454

5555
def _get_xray_trace_context():
56-
xray_trace_entity = _get_context_from_raw_xray_trace_id(os.environ.get(XRAY_TRACE_ID_HEADER_NAME,""))
56+
xray_trace_entity = _get_context_from_raw_xray_trace_id(
57+
os.environ.get(XRAY_TRACE_ID_HEADER_NAME, "")
58+
)
5759
if xray_trace_entity is None:
5860
return None
5961
trace_context = {
@@ -64,7 +66,11 @@ def _get_xray_trace_context():
6466
logger.debug(
6567
"Converted trace context %s from X-Ray segment %s",
6668
trace_context,
67-
(xray_trace_entity.get("trace_id"), xray_trace_entity.get("id"), xray_trace_entity.get("sampled")),
69+
(
70+
xray_trace_entity.get("trace_id"),
71+
xray_trace_entity.get("id"),
72+
xray_trace_entity.get("sampled"),
73+
),
6874
)
6975
return trace_context
7076

@@ -95,8 +101,9 @@ def _context_obj_to_headers(obj):
95101
TraceHeader.SAMPLING_PRIORITY: str(obj.get("sampling-priority")),
96102
}
97103

104+
98105
def _get_context_from_raw_xray_trace_id(raw_trace_id):
99-
#example : Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1
106+
# example : Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1
100107
if len(raw_trace_id) == 0:
101108
return None
102109
parts = raw_trace_id.split(";")
@@ -105,14 +112,19 @@ def _get_context_from_raw_xray_trace_id(raw_trace_id):
105112
root = parts[0].replace("Root=", "")
106113
parent = parts[1].replace("Parent=", "")
107114
sampled = parts[2].replace("Sampled=", "")
108-
if len(root) == len(parts[0]) or len(parent) == len(parts[1]) or len(sampled) == len(parts[2]):
115+
if (
116+
len(root) == len(parts[0])
117+
or len(parent) == len(parts[1])
118+
or len(sampled) == len(parts[2])
119+
):
109120
return None
110121
return {
111122
"id": parent,
112123
"trace_id": root,
113124
"sampled": sampled,
114125
}
115126

127+
116128
def extract_context_from_lambda_context(lambda_context):
117129
"""
118130
Extract Datadog trace context from the `client_context` attr
@@ -293,7 +305,6 @@ def get_dd_trace_context():
293305
return _context_obj_to_headers(context) if context is not None else {}
294306

295307

296-
297308
def set_correlation_ids():
298309
"""
299310
Create a dummy span, and overrides its trace_id and span_id, to make
@@ -346,9 +357,11 @@ def inject_correlation_ids():
346357

347358
logger.debug("logs injection configured")
348359

360+
349361
def is_lambda_context():
350362
return os.environ.get("FUNCTION_NAME_HEADER_NAME", "") != ""
351363

364+
352365
def set_dd_trace_py_root(trace_context_source, merge_xray_traces):
353366
if trace_context_source == TraceContextSource.EVENT or merge_xray_traces:
354367
headers = _context_obj_to_headers(dd_trace_context)

‎tests/test_tracing.py

Lines changed: 38 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,12 @@
99

1010
from ddtrace.helpers import get_correlation_ids
1111

12-
from datadog_lambda.constants import SamplingPriority, TraceHeader, XraySubsegment, XRAY_TRACE_ID_HEADER_NAME
12+
from datadog_lambda.constants import (
13+
SamplingPriority,
14+
TraceHeader,
15+
XraySubsegment,
16+
XRAY_TRACE_ID_HEADER_NAME,
17+
)
1318
from datadog_lambda.tracing import (
1419
extract_dd_trace_context,
1520
create_function_execution_span,
@@ -18,12 +23,14 @@
1823
_convert_xray_trace_id,
1924
_convert_xray_entity_id,
2025
_convert_xray_sampling,
21-
_get_context_from_raw_xray_trace_id
26+
_get_context_from_raw_xray_trace_id,
2227
)
2328

2429
function_arn = "arn:aws:lambda:us-west-1:123457598159:function:python-layer-test"
2530

26-
fake_xray_header_value = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1";
31+
fake_xray_header_value = (
32+
"Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1"
33+
)
2734
fake_xray_header_value_parent_decimal = "10713633173203262661"
2835
fake_xray_header_value_root_decimal = "3995693151288333088"
2936

@@ -72,7 +79,11 @@ def test_without_datadog_trace_headers(self):
7279
self.assertEqual(source, "xray")
7380
self.assertDictEqual(
7481
ctx,
75-
{"trace-id": fake_xray_header_value_root_decimal, "parent-id": fake_xray_header_value_parent_decimal, "sampling-priority": "2"},
82+
{
83+
"trace-id": fake_xray_header_value_root_decimal,
84+
"parent-id": fake_xray_header_value_parent_decimal,
85+
"sampling-priority": "2",
86+
},
7687
)
7788
self.assertDictEqual(
7889
get_dd_trace_context(),
@@ -93,7 +104,11 @@ def test_with_incomplete_datadog_trace_headers(self):
93104
self.assertEqual(source, "xray")
94105
self.assertDictEqual(
95106
ctx,
96-
{"trace-id": fake_xray_header_value_root_decimal, "parent-id": fake_xray_header_value_parent_decimal, "sampling-priority": "2"},
107+
{
108+
"trace-id": fake_xray_header_value_root_decimal,
109+
"parent-id": fake_xray_header_value_parent_decimal,
110+
"sampling-priority": "2",
111+
},
97112
)
98113
self.assertDictEqual(
99114
get_dd_trace_context(),
@@ -122,15 +137,14 @@ def test_with_complete_datadog_trace_headers(self):
122137
{"trace-id": "123", "parent-id": "321", "sampling-priority": "1"},
123138
)
124139
self.assertDictEqual(
125-
get_d,
140+
get_dd_trace_context(),
126141
{
127142
TraceHeader.TRACE_ID: "123",
128143
TraceHeader.PARENT_ID: fake_xray_header_value_parent_decimal,
129144
TraceHeader.SAMPLING_PRIORITY: "1",
130145
},
131146
)
132147

133-
134148
def test_with_extractor_function(self):
135149
def extractor_foo(event, context):
136150
foo = event.get("foo", {})
@@ -174,6 +188,7 @@ def extractor_foo(event, context):
174188
def test_graceful_fail_of_extractor_function(self):
175189
def extractor_raiser(event, context):
176190
raise Exception("kreator")
191+
177192
lambda_ctx = get_mock_context()
178193
ctx, ctx_source = extract_dd_trace_context(
179194
{
@@ -369,20 +384,27 @@ def test_convert_xray_sampling(self):
369384
)
370385

371386
def test_get_context_from_raw_xray_trace_id(self):
372-
#example : Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1
373-
raw_trace_id = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1"
374-
self.assertEqual(_get_context_from_raw_xray_trace_id(raw_trace_id), {
375-
"trace_id": "1-5e272390-8c398be037738dc042009320",
376-
"id": "94ae789b969f1cc5",
377-
"sampled": "1"
378-
})
387+
# example : Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1
388+
raw_trace_id = (
389+
"Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1"
390+
)
391+
self.assertEqual(
392+
_get_context_from_raw_xray_trace_id(raw_trace_id),
393+
{
394+
"trace_id": "1-5e272390-8c398be037738dc042009320",
395+
"id": "94ae789b969f1cc5",
396+
"sampled": "1",
397+
},
398+
)
379399

380400
def test_get_context_from_raw_xray_trace_id_truncated_trace(self):
381-
raw_trace_id = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5"
401+
raw_trace_id = (
402+
"Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5"
403+
)
382404
self.assertIsNone(_get_context_from_raw_xray_trace_id(raw_trace_id))
383405

384406
def test_get_context_from_raw_xray_trace_id_truncated_trace(self):
385-
#example : Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1
407+
# example : Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1
386408
raw_trace_id = "Ro_xxx_ot=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1"
387409
self.assertIsNone(_get_context_from_raw_xray_trace_id(raw_trace_id))
388410

0 commit comments

Comments
 (0)