Skip to content

Commit 737c420

Browse files
majorgreysmajorgreys
committed
Add Datadog propagator
1 parent 9e58b8a commit 737c420

File tree

5 files changed

+355
-0
lines changed

5 files changed

+355
-0
lines changed

Diff for: ‎ext/opentelemetry-ext-datadog/src/opentelemetry/ext/datadog/constants.py

+4
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
DD_ORIGIN = "_dd_origin"
2+
AUTO_REJECT = 0
3+
AUTO_KEEP = 1
4+
USER_KEEP = 2

Diff for: ‎ext/opentelemetry-ext-datadog/src/opentelemetry/ext/datadog/exporter.py

+14
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,9 @@
2424
from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
2525
from opentelemetry.trace.status import StatusCanonicalCode
2626

27+
# pylint:disable=relative-beyond-top-level
28+
from .constants import DD_ORIGIN
29+
2730
logger = logging.getLogger(__name__)
2831

2932

@@ -128,6 +131,11 @@ def _translate_to_datadog(self, spans):
128131

129132
datadog_span.set_tags(span.attributes)
130133

134+
# add origin to root span
135+
origin = _get_origin(span)
136+
if origin and parent_id == 0:
137+
datadog_span.set_tag(DD_ORIGIN, origin)
138+
131139
# span events and span links are not supported
132140

133141
datadog_spans.append(datadog_span)
@@ -202,3 +210,9 @@ def _get_exc_info(span):
202210
"""Parse span status description for exception type and value"""
203211
exc_type, exc_val = span.status.description.split(":", 1)
204212
return exc_type, exc_val.strip()
213+
214+
215+
def _get_origin(span):
216+
ctx = span.get_context()
217+
origin = ctx.trace_state.get(DD_ORIGIN)
218+
return origin

Diff for: ‎ext/opentelemetry-ext-datadog/src/opentelemetry/ext/datadog/propagator.py

+127
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
# Copyright The OpenTelemetry Authors
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import typing
16+
17+
from opentelemetry import trace
18+
from opentelemetry.context import Context
19+
from opentelemetry.trace.propagation import (
20+
get_span_from_context,
21+
set_span_in_context,
22+
)
23+
from opentelemetry.trace.propagation.httptextformat import (
24+
Getter,
25+
HTTPTextFormat,
26+
HTTPTextFormatT,
27+
Setter,
28+
)
29+
30+
# pylint:disable=relative-beyond-top-level
31+
from . import constants
32+
33+
34+
class DatadogFormat(HTTPTextFormat):
35+
"""Propagator for the Datadog HTTP header format.
36+
"""
37+
38+
TRACE_ID_KEY = "x-datadog-trace-id"
39+
PARENT_ID_KEY = "x-datadog-parent-id"
40+
SAMPLING_PRIORITY_KEY = "x-datadog-sampling-priority"
41+
ORIGIN_KEY = "x-datadog-origin"
42+
43+
def extract(
44+
self,
45+
get_from_carrier: Getter[HTTPTextFormatT],
46+
carrier: HTTPTextFormatT,
47+
context: typing.Optional[Context] = None,
48+
) -> Context:
49+
trace_id = extract_first_element(
50+
get_from_carrier(carrier, self.TRACE_ID_KEY)
51+
)
52+
53+
span_id = extract_first_element(
54+
get_from_carrier(carrier, self.PARENT_ID_KEY)
55+
)
56+
57+
sampled = extract_first_element(
58+
get_from_carrier(carrier, self.SAMPLING_PRIORITY_KEY)
59+
)
60+
61+
origin = extract_first_element(
62+
get_from_carrier(carrier, self.ORIGIN_KEY)
63+
)
64+
65+
trace_flags = trace.TraceFlags()
66+
if sampled and int(sampled) in (
67+
constants.AUTO_KEEP,
68+
constants.USER_KEEP,
69+
):
70+
trace_flags |= trace.TraceFlags.SAMPLED
71+
72+
if trace_id is None or span_id is None:
73+
return set_span_in_context(trace.INVALID_SPAN, context)
74+
75+
span_context = trace.SpanContext(
76+
trace_id=int(trace_id),
77+
span_id=int(span_id),
78+
is_remote=True,
79+
trace_flags=trace_flags,
80+
trace_state=trace.TraceState({constants.DD_ORIGIN: origin}),
81+
)
82+
83+
return set_span_in_context(trace.DefaultSpan(span_context), context)
84+
85+
def inject(
86+
self,
87+
set_in_carrier: Setter[HTTPTextFormatT],
88+
carrier: HTTPTextFormatT,
89+
context: typing.Optional[Context] = None,
90+
) -> None:
91+
span = get_span_from_context(context=context)
92+
sampled = (trace.TraceFlags.SAMPLED & span.context.trace_flags) != 0
93+
set_in_carrier(
94+
carrier, self.TRACE_ID_KEY, format_trace_id(span.context.trace_id),
95+
)
96+
set_in_carrier(
97+
carrier, self.PARENT_ID_KEY, format_span_id(span.context.span_id)
98+
)
99+
set_in_carrier(
100+
carrier,
101+
self.SAMPLING_PRIORITY_KEY,
102+
str(constants.AUTO_KEEP if sampled else constants.AUTO_REJECT),
103+
)
104+
if constants.DD_ORIGIN in span.context.trace_state:
105+
set_in_carrier(
106+
carrier,
107+
self.ORIGIN_KEY,
108+
span.context.trace_state[constants.DD_ORIGIN],
109+
)
110+
111+
112+
def format_trace_id(trace_id: int) -> str:
113+
"""Format the trace id for Datadog."""
114+
return str(trace_id & 0xFFFFFFFFFFFFFFFF)
115+
116+
117+
def format_span_id(span_id: int) -> str:
118+
"""Format the span id for Datadog."""
119+
return str(span_id)
120+
121+
122+
def extract_first_element(
123+
items: typing.Iterable[HTTPTextFormatT],
124+
) -> typing.Optional[HTTPTextFormatT]:
125+
if items is None:
126+
return None
127+
return next(iter(items), None)

Diff for: ‎ext/opentelemetry-ext-datadog/tests/test_datadog_exporter.py

+37
Original file line numberDiff line numberDiff line change
@@ -403,3 +403,40 @@ def test_span_processor_scheduled_delay(self):
403403
self.assertEqual(len(datadog_spans), 1)
404404

405405
tracer_provider.shutdown()
406+
407+
def test_origin(self):
408+
context = trace_api.SpanContext(
409+
trace_id=0x000000000000000000000000DEADBEEF,
410+
span_id=trace_api.INVALID_SPAN,
411+
is_remote=True,
412+
trace_state=trace_api.TraceState(
413+
{datadog.constants.DD_ORIGIN: "origin-service"}
414+
),
415+
)
416+
417+
root_span = trace.Span(name="root", context=context, parent=None)
418+
child_span = trace.Span(
419+
name="child", context=context, parent=root_span
420+
)
421+
root_span.start()
422+
child_span.start()
423+
child_span.end()
424+
root_span.end()
425+
426+
# pylint: disable=protected-access
427+
exporter = datadog.DatadogSpanExporter()
428+
datadog_spans = [
429+
span.to_dict()
430+
for span in exporter._translate_to_datadog([root_span, child_span])
431+
]
432+
433+
self.assertEqual(len(datadog_spans), 2)
434+
435+
actual = [
436+
span["meta"].get(datadog.constants.DD_ORIGIN)
437+
if "meta" in span
438+
else None
439+
for span in datadog_spans
440+
]
441+
expected = ["origin-service", None]
442+
self.assertListEqual(actual, expected)

Diff for: ‎ext/opentelemetry-ext-datadog/tests/test_datadog_format.py

+173
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,173 @@
1+
# Copyright The OpenTelemetry Authors
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import unittest
16+
17+
from opentelemetry import trace as trace_api
18+
from opentelemetry.ext.datadog import constants, propagator
19+
from opentelemetry.sdk import trace
20+
from opentelemetry.trace.propagation import (
21+
get_span_from_context,
22+
set_span_in_context,
23+
)
24+
25+
FORMAT = propagator.DatadogFormat()
26+
27+
28+
def get_as_list(dict_object, key):
29+
value = dict_object.get(key)
30+
return [value] if value is not None else []
31+
32+
33+
class TestDatadogFormat(unittest.TestCase):
34+
@classmethod
35+
def setUpClass(cls):
36+
cls.serialized_trace_id = propagator.format_trace_id(
37+
trace.generate_trace_id()
38+
)
39+
cls.serialized_parent_id = propagator.format_span_id(
40+
trace.generate_span_id()
41+
)
42+
cls.serialized_origin = "origin-service"
43+
44+
def test_malformed_headers(self):
45+
"""Test with no Datadog headers"""
46+
malformed_trace_id_key = FORMAT.TRACE_ID_KEY + "-x"
47+
malformed_parent_id_key = FORMAT.PARENT_ID_KEY + "-x"
48+
context = get_span_from_context(
49+
FORMAT.extract(
50+
get_as_list,
51+
{
52+
malformed_trace_id_key: self.serialized_trace_id,
53+
malformed_parent_id_key: self.serialized_parent_id,
54+
},
55+
)
56+
).get_context()
57+
58+
self.assertNotEqual(context.trace_id, int(self.serialized_trace_id))
59+
self.assertNotEqual(context.span_id, int(self.serialized_parent_id))
60+
self.assertFalse(context.is_remote)
61+
62+
def test_missing_trace_id(self):
63+
"""If a trace id is missing, populate an invalid trace id."""
64+
carrier = {
65+
FORMAT.PARENT_ID_KEY: self.serialized_parent_id,
66+
}
67+
68+
ctx = FORMAT.extract(get_as_list, carrier)
69+
span_context = get_span_from_context(ctx).get_context()
70+
self.assertEqual(span_context.trace_id, trace_api.INVALID_TRACE_ID)
71+
72+
def test_missing_parent_id(self):
73+
"""If a parent id is missing, populate an invalid trace id."""
74+
carrier = {
75+
FORMAT.TRACE_ID_KEY: self.serialized_trace_id,
76+
}
77+
78+
ctx = FORMAT.extract(get_as_list, carrier)
79+
span_context = get_span_from_context(ctx).get_context()
80+
self.assertEqual(span_context.span_id, trace_api.INVALID_SPAN_ID)
81+
82+
def test_context_propagation(self):
83+
"""Test the propagation of Datadog headers."""
84+
parent_context = get_span_from_context(
85+
FORMAT.extract(
86+
get_as_list,
87+
{
88+
FORMAT.TRACE_ID_KEY: self.serialized_trace_id,
89+
FORMAT.PARENT_ID_KEY: self.serialized_parent_id,
90+
FORMAT.SAMPLING_PRIORITY_KEY: str(constants.AUTO_KEEP),
91+
FORMAT.ORIGIN_KEY: self.serialized_origin,
92+
},
93+
)
94+
).get_context()
95+
96+
self.assertEqual(
97+
parent_context.trace_id, int(self.serialized_trace_id)
98+
)
99+
self.assertEqual(
100+
parent_context.span_id, int(self.serialized_parent_id)
101+
)
102+
self.assertEqual(parent_context.trace_flags, constants.AUTO_KEEP)
103+
self.assertEqual(
104+
parent_context.trace_state.get(constants.DD_ORIGIN),
105+
self.serialized_origin,
106+
)
107+
self.assertTrue(parent_context.is_remote)
108+
109+
child = trace.Span(
110+
"child",
111+
trace_api.SpanContext(
112+
parent_context.trace_id,
113+
trace.generate_span_id(),
114+
is_remote=False,
115+
trace_flags=parent_context.trace_flags,
116+
trace_state=parent_context.trace_state,
117+
),
118+
parent=parent_context,
119+
)
120+
121+
child_carrier = {}
122+
child_context = set_span_in_context(child)
123+
FORMAT.inject(dict.__setitem__, child_carrier, context=child_context)
124+
125+
self.assertEqual(
126+
child_carrier[FORMAT.TRACE_ID_KEY], self.serialized_trace_id
127+
)
128+
self.assertEqual(
129+
child_carrier[FORMAT.PARENT_ID_KEY], str(child.context.span_id)
130+
)
131+
self.assertEqual(
132+
child_carrier[FORMAT.SAMPLING_PRIORITY_KEY],
133+
str(constants.AUTO_KEEP),
134+
)
135+
self.assertEqual(
136+
child_carrier.get(FORMAT.ORIGIN_KEY), self.serialized_origin
137+
)
138+
139+
def test_sampling_priority_auto_reject(self):
140+
"""Test sampling priority rejected."""
141+
parent_context = get_span_from_context(
142+
FORMAT.extract(
143+
get_as_list,
144+
{
145+
FORMAT.TRACE_ID_KEY: self.serialized_trace_id,
146+
FORMAT.PARENT_ID_KEY: self.serialized_parent_id,
147+
FORMAT.SAMPLING_PRIORITY_KEY: str(constants.AUTO_REJECT),
148+
},
149+
)
150+
).get_context()
151+
152+
self.assertEqual(parent_context.trace_flags, constants.AUTO_REJECT)
153+
154+
child = trace.Span(
155+
"child",
156+
trace_api.SpanContext(
157+
parent_context.trace_id,
158+
trace.generate_span_id(),
159+
is_remote=False,
160+
trace_flags=parent_context.trace_flags,
161+
trace_state=parent_context.trace_state,
162+
),
163+
parent=parent_context,
164+
)
165+
166+
child_carrier = {}
167+
child_context = set_span_in_context(child)
168+
FORMAT.inject(dict.__setitem__, child_carrier, context=child_context)
169+
170+
self.assertEqual(
171+
child_carrier[FORMAT.SAMPLING_PRIORITY_KEY],
172+
str(constants.AUTO_REJECT),
173+
)

0 commit comments

Comments
 (0)