Skip to content

Commit 0cf28b7

Browse files
committed
Inject headers when patching http requests.
1 parent 3312a30 commit 0cf28b7

File tree

2 files changed

+181
-135
lines changed

2 files changed

+181
-135
lines changed

datadog_lambda/tracing.py

Lines changed: 27 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -72,25 +72,21 @@ def _convert_xray_trace_id(xray_trace_id):
7272
"""
7373
Convert X-Ray trace id (hex)'s last 63 bits to a Datadog trace id (int).
7474
"""
75-
return str(0x7FFFFFFFFFFFFFFF & int(xray_trace_id[-16:], 16))
75+
return 0x7FFFFFFFFFFFFFFF & int(xray_trace_id[-16:], 16)
7676

7777

7878
def _convert_xray_entity_id(xray_entity_id):
7979
"""
8080
Convert X-Ray (sub)segement id (hex) to a Datadog span id (int).
8181
"""
82-
return str(int(xray_entity_id, 16))
82+
return int(xray_entity_id, 16)
8383

8484

8585
def _convert_xray_sampling(xray_sampled):
8686
"""
8787
Convert X-Ray sampled (True/False) to its Datadog counterpart.
8888
"""
89-
return (
90-
str(SamplingPriority.USER_KEEP)
91-
if xray_sampled
92-
else str(SamplingPriority.USER_REJECT)
93-
)
89+
return SamplingPriority.USER_KEEP if xray_sampled else SamplingPriority.USER_REJECT
9490

9591

9692
def _get_xray_trace_context():
@@ -332,16 +328,16 @@ def extract_context_from_kinesis_event(event, lambda_context):
332328
return extract_context_from_lambda_context(lambda_context)
333329

334330

335-
def _deterministic_md5_hash(s: str) -> str:
331+
def _deterministic_md5_hash(s: str) -> int:
336332
"""MD5 here is to generate trace_id, not for any encryption."""
337333
hex_number = hashlib.md5(s.encode("ascii")).hexdigest()
338334
binary = bin(int(hex_number, 16))
339335
binary_str = str(binary)
340336
binary_str_remove_0b = binary_str[2:].rjust(128, "0")
341337
most_significant_64_bits_without_leading_1 = "0" + binary_str_remove_0b[1:-64]
342-
result = str(int(most_significant_64_bits_without_leading_1, 2))
343-
if result == "0" * 64:
344-
return "1"
338+
result = int(most_significant_64_bits_without_leading_1, 2)
339+
if result == 0:
340+
return 1
345341
return result
346342

347343

@@ -378,7 +374,9 @@ def extract_context_custom_extractor(extractor, event, lambda_context):
378374
sampling_priority,
379375
) = extractor(event, lambda_context)
380376
return Context(
381-
trace_id=trace_id, span_id=parent_id, sampling_priority=sampling_priority
377+
trace_id=int(trace_id),
378+
span_id=int(parent_id),
379+
sampling_priority=int(sampling_priority),
382380
)
383381
except Exception as e:
384382
logger.debug("The trace extractor returned with error %s", e)
@@ -480,7 +478,7 @@ def extract_dd_trace_context(
480478
return dd_trace_context, trace_context_source, event_source
481479

482480

483-
def get_dd_trace_context():
481+
def get_dd_trace_context_obj():
484482
"""
485483
Return the Datadog trace context to be propagated on the outgoing requests.
486484
@@ -517,9 +515,24 @@ def get_dd_trace_context():
517515
trace_id=dd_trace_context.trace_id,
518516
span_id=xray_context.span_id,
519517
sampling_priority=dd_trace_context.sampling_priority,
518+
meta=dd_trace_context._meta.copy(),
519+
metrics=dd_trace_context._metrics.copy(),
520520
)
521521

522522

523+
def get_dd_trace_context():
524+
"""
525+
Return the Datadog trace context to be propagated on the outgoing requests,
526+
as a dict of headers.
527+
"""
528+
headers = {}
529+
context = get_dd_trace_context_obj()
530+
if not _is_context_complete(context):
531+
return headers
532+
propagator.inject(context, headers)
533+
return headers
534+
535+
523536
def set_correlation_ids():
524537
"""
525538
Create a dummy span, and overrides its trace_id and span_id, to make
@@ -535,7 +548,7 @@ def set_correlation_ids():
535548
logger.debug("using ddtrace implementation for spans")
536549
return
537550

538-
context = get_dd_trace_context()
551+
context = get_dd_trace_context_obj()
539552
if not _is_context_complete(context):
540553
return
541554

0 commit comments

Comments
 (0)