Skip to content

Commit d9a2c8c

Browse files
NathanielRNNathanielRN
committed
More pythonic propagator + package clean up
1 parent 73702bc commit d9a2c8c

File tree

3 files changed

+132
-98
lines changed

3 files changed

+132
-98
lines changed

‎sdk-extension/opentelemetry-sdk-extension-aws/setup.cfg

+2-2
Original file line numberDiff line numberDiff line change
@@ -39,15 +39,15 @@ package_dir=
3939
=src
4040
packages=find_namespace:
4141
install_requires =
42-
opentelemetry-api == 0.15.b0
42+
opentelemetry-api == 0.16.dev0
4343

4444
[options.entry_points]
4545
opentelemetry_propagator =
4646
aws_xray = opentelemetry.sdk.extension.aws.trace.propagation.aws_xray_format:AwsXRayFormat
4747

4848
[options.extras_require]
4949
test =
50-
opentelemetry-test == 0.15.b0
50+
opentelemetry-test == 0.16.dev0
5151

5252
[options.packages.find]
5353
where = src

‎sdk-extension/opentelemetry-sdk-extension-aws/src/opentelemetry/sdk/extension/aws/trace/propagation/aws_xray_format.py

+95-96
Original file line numberDiff line numberDiff line change
@@ -55,10 +55,6 @@ class AwsXRayFormat(TextMapPropagator):
5555
IS_SAMPLED = "1"
5656
NOT_SAMPLED = "0"
5757

58-
# pylint: disable=too-many-locals
59-
# pylint: disable=too-many-return-statements
60-
# pylint: disable=too-many-branches
61-
# pylint: disable=too-many-statements
6258
def extract(
6359
self,
6460
getter: Getter[TextMapPropagatorT],
@@ -79,73 +75,78 @@ def extract(
7975
trace.INVALID_SPAN, context=context
8076
)
8177

78+
trace_id, span_id, sampled, err = self.extract_span_properties(
79+
trace_header
80+
)
81+
82+
if err is not None:
83+
return trace.set_span_in_context(
84+
trace.INVALID_SPAN, context=context
85+
)
86+
87+
options = 0
88+
if sampled:
89+
options |= trace.TraceFlags.SAMPLED
90+
91+
span_context = trace.SpanContext(
92+
trace_id=trace_id,
93+
span_id=span_id,
94+
is_remote=True,
95+
trace_flags=trace.TraceFlags(options),
96+
trace_state=trace.TraceState(),
97+
)
98+
99+
if not span_context.is_valid:
100+
_logger.error(
101+
"Invalid Span Extracted. Insertting INVALID span into provided context."
102+
)
103+
return trace.set_span_in_context(
104+
trace.INVALID_SPAN, context=context
105+
)
106+
107+
return trace.set_span_in_context(
108+
trace.DefaultSpan(span_context), context=context
109+
)
110+
111+
def extract_span_properties(self, trace_header):
82112
trace_id = trace.INVALID_TRACE_ID
83113
span_id = trace.INVALID_SPAN_ID
84114
sampled = False
85115

86-
next_kv_pair_start = 0
116+
extract_err = None
87117

88-
while next_kv_pair_start < len(trace_header):
89-
try:
90-
kv_pair_delimiter_index = trace_header.index(
91-
self.KV_PAIR_DELIMITER, next_kv_pair_start
92-
)
93-
kv_pair_subset = trace_header[
94-
next_kv_pair_start:kv_pair_delimiter_index
95-
]
96-
next_kv_pair_start = kv_pair_delimiter_index + 1
97-
except ValueError:
98-
kv_pair_subset = trace_header[next_kv_pair_start:]
99-
next_kv_pair_start = len(trace_header)
100-
101-
stripped_kv_pair = kv_pair_subset.strip()
118+
for kv_pair_str in trace_header.split(self.KV_PAIR_DELIMITER):
119+
if extract_err:
120+
break
102121

103122
try:
104-
key_and_value_delimiter_index = stripped_kv_pair.index(
123+
key_str, value_str = kv_pair_str.split(
105124
self.KEY_AND_VALUE_DELIMITER
106125
)
126+
key, value = key_str.strip(), value_str.strip()
107127
except ValueError:
108128
_logger.error(
109129
(
110130
"Error parsing X-Ray trace header. Invalid key value pair: %s. Returning INVALID span context.",
111-
kv_pair_subset,
131+
kv_pair_str,
112132
)
113133
)
114-
return trace.set_span_in_context(
115-
trace.INVALID_SPAN, context=context
116-
)
134+
return trace_id, span_id, sampled, extract_err
117135

118-
value = stripped_kv_pair[key_and_value_delimiter_index + 1 :]
119-
120-
if stripped_kv_pair.startswith(self.TRACE_ID_KEY):
121-
if (
122-
len(value) != self.TRACE_ID_LENGTH
123-
or not value.startswith(self.TRACE_ID_VERSION)
124-
or value[self.TRACE_ID_DELIMITER_INDEX_1]
125-
!= self.TRACE_ID_DELIMITER
126-
or value[self.TRACE_ID_DELIMITER_INDEX_2]
127-
!= self.TRACE_ID_DELIMITER
128-
):
136+
if key == self.TRACE_ID_KEY:
137+
if not self.validate_trace_id(value):
129138
_logger.error(
130139
(
131140
"Invalid TraceId in X-Ray trace header: '%s' with value '%s'. Returning INVALID span context.",
132141
self.TRACE_HEADER_KEY,
133142
trace_header,
134143
)
135144
)
136-
return trace.set_span_in_context(
137-
trace.INVALID_SPAN, context=context
138-
)
145+
extract_err = True
146+
break
139147

140-
timestamp_subset = value[
141-
self.TRACE_ID_DELIMITER_INDEX_1
142-
+ 1 : self.TRACE_ID_DELIMITER_INDEX_2
143-
]
144-
unique_id_subset = value[
145-
self.TRACE_ID_DELIMITER_INDEX_2 + 1 : self.TRACE_ID_LENGTH
146-
]
147148
try:
148-
trace_id = int(timestamp_subset + unique_id_subset, 16)
149+
trace_id = self.parse_trace_id(value)
149150
except ValueError:
150151
_logger.error(
151152
(
@@ -154,24 +155,21 @@ def extract(
154155
trace_header,
155156
)
156157
)
157-
return trace.set_span_in_context(
158-
trace.INVALID_SPAN, context=context
159-
)
160-
elif stripped_kv_pair.startswith(self.PARENT_ID_KEY):
161-
if len(value) != self.PARENT_ID_LENGTH:
158+
extract_err = True
159+
elif key == self.PARENT_ID_KEY:
160+
if not self.validate_span_id(value):
162161
_logger.error(
163162
(
164163
"Invalid ParentId in X-Ray trace header: '%s' with value '%s'. Returning INVALID span context.",
165164
self.TRACE_HEADER_KEY,
166165
trace_header,
167166
)
168167
)
169-
return trace.set_span_in_context(
170-
trace.INVALID_SPAN, context=context
171-
)
168+
extract_err = True
169+
break
172170

173171
try:
174-
span_id = int(value, 16)
172+
span_id = AwsXRayFormat.parse_span_id(value)
175173
except ValueError:
176174
_logger.error(
177175
(
@@ -180,60 +178,61 @@ def extract(
180178
trace_header,
181179
)
182180
)
183-
return trace.set_span_in_context(
184-
trace.INVALID_SPAN, context=context
185-
)
186-
elif stripped_kv_pair.startswith(self.SAMPLED_FLAG_KEY):
187-
is_sampled_flag_valid = True
188-
189-
if len(value) != self.SAMPLED_FLAG_LENGTH:
190-
is_sampled_flag_valid = False
191-
192-
if is_sampled_flag_valid:
193-
sampled_flag = value[0]
194-
if sampled_flag == self.IS_SAMPLED:
195-
sampled = True
196-
elif sampled_flag == self.NOT_SAMPLED:
197-
sampled = False
198-
else:
199-
is_sampled_flag_valid = False
200-
201-
if not is_sampled_flag_valid:
181+
extract_err = True
182+
elif key == self.SAMPLED_FLAG_KEY:
183+
if not self.validate_sampled_flag(value):
202184
_logger.error(
203185
(
204186
"Invalid Sampling flag in X-Ray trace header: '%s' with value '%s'. Returning INVALID span context.",
205187
self.TRACE_HEADER_KEY,
206188
trace_header,
207189
)
208190
)
209-
return trace.set_span_in_context(
210-
trace.INVALID_SPAN, context=context
211-
)
191+
extract_err = True
192+
break
212193

213-
options = 0
214-
if sampled:
215-
options |= trace.TraceFlags.SAMPLED
194+
sampled = self.parse_sampled_flag(value)
216195

217-
span_context = trace.SpanContext(
218-
trace_id=trace_id,
219-
span_id=span_id,
220-
is_remote=True,
221-
trace_flags=trace.TraceFlags(options),
222-
trace_state=trace.TraceState(),
223-
)
196+
return trace_id, span_id, sampled, extract_err
224197

225-
if not span_context.is_valid:
226-
_logger.error(
227-
"Invalid Span Extracted. Insertting INVALID span into provided context."
228-
)
229-
return trace.set_span_in_context(
230-
trace.INVALID_SPAN, context=context
231-
)
198+
def validate_trace_id(self, trace_id_str):
199+
return (
200+
len(trace_id_str) == self.TRACE_ID_LENGTH
201+
and trace_id_str.startswith(self.TRACE_ID_VERSION)
202+
and trace_id_str[self.TRACE_ID_DELIMITER_INDEX_1]
203+
== self.TRACE_ID_DELIMITER
204+
and trace_id_str[self.TRACE_ID_DELIMITER_INDEX_2]
205+
== self.TRACE_ID_DELIMITER
206+
)
232207

233-
return trace.set_span_in_context(
234-
trace.DefaultSpan(span_context), context=context
208+
def parse_trace_id(self, trace_id_str):
209+
timestamp_subset = trace_id_str[
210+
self.TRACE_ID_DELIMITER_INDEX_1
211+
+ 1 : self.TRACE_ID_DELIMITER_INDEX_2
212+
]
213+
unique_id_subset = trace_id_str[
214+
self.TRACE_ID_DELIMITER_INDEX_2 + 1 : self.TRACE_ID_LENGTH
215+
]
216+
return int(timestamp_subset + unique_id_subset, 16)
217+
218+
def validate_span_id(self, span_id_str):
219+
return len(span_id_str) == self.PARENT_ID_LENGTH
220+
221+
@staticmethod
222+
def parse_span_id(span_id_str):
223+
return int(span_id_str, 16)
224+
225+
def validate_sampled_flag(self, sampled_flag_str):
226+
return len(
227+
sampled_flag_str
228+
) == self.SAMPLED_FLAG_LENGTH and sampled_flag_str in (
229+
self.IS_SAMPLED,
230+
self.NOT_SAMPLED,
235231
)
236232

233+
def parse_sampled_flag(self, sampled_flag_str):
234+
return sampled_flag_str[0] == self.IS_SAMPLED
235+
237236
def inject(
238237
self,
239238
set_in_carrier: Setter[TextMapPropagatorT],

‎sdk-extension/opentelemetry-sdk-extension-aws/tests/trace/propagation/test_aws_xray_format.py

+35
Original file line numberDiff line numberDiff line change
@@ -247,6 +247,41 @@ def test_extract_with_additional_fields(self):
247247
get_extracted_span_context(build_test_context()),
248248
)
249249

250+
def test_extract_with_extra_whitespace(self):
251+
default_xray_trace_header_dict = build_dict_with_xray_trace_header()
252+
trace_header_components = default_xray_trace_header_dict[
253+
AwsXRayFormat.TRACE_HEADER_KEY
254+
].split(AwsXRayFormat.KV_PAIR_DELIMITER)
255+
xray_trace_header_dict_with_extra_whitespace = CaseInsensitiveDict(
256+
{
257+
AwsXRayFormat.TRACE_HEADER_KEY: AwsXRayFormat.KV_PAIR_DELIMITER.join(
258+
[
259+
AwsXRayFormat.KEY_AND_VALUE_DELIMITER.join(
260+
[
261+
" " + key + " ",
262+
" " + value + " ",
263+
]
264+
)
265+
for kv_pair_str in trace_header_components
266+
for key, value in [
267+
kv_pair_str.split(
268+
AwsXRayFormat.KEY_AND_VALUE_DELIMITER
269+
)
270+
]
271+
]
272+
)
273+
}
274+
)
275+
actual_context_encompassing_extracted = AwsXRayPropagatorTest.XRAY_PROPAGATOR.extract(
276+
AwsXRayPropagatorTest.carrier_getter,
277+
xray_trace_header_dict_with_extra_whitespace,
278+
)
279+
280+
self.assertEqual(
281+
get_extracted_span_context(actual_context_encompassing_extracted),
282+
get_extracted_span_context(build_test_context()),
283+
)
284+
250285
def test_extract_invalid_xray_trace_header(self):
251286
actual_context_encompassing_extracted = AwsXRayPropagatorTest.XRAY_PROPAGATOR.extract(
252287
AwsXRayPropagatorTest.carrier_getter,

0 commit comments

Comments
 (0)