diff --git a/py_zipkin/zipkin.py b/py_zipkin/zipkin.py index 91b99ea..791ea49 100644 --- a/py_zipkin/zipkin.py +++ b/py_zipkin/zipkin.py @@ -617,6 +617,29 @@ def override_span_name(self, name): if self.logging_context: self.logging_context.span_name = name + def create_http_headers_for_my_span(self): + """ + Generate the headers for sharing this context object's zipkin_attrs + with a shared span on another host. + + If this instance doesn't have zipkin_attrs set, for some reason, an + empty dict is returned. + + :returns: dict containing (X-B3-TraceId, X-B3-SpanId, X-B3-ParentSpanId, + X-B3-Flags and X-B3-Sampled) or an empty dict. + """ + zipkin_attrs = self.zipkin_attrs + if not zipkin_attrs: + return {} + + return { + 'X-B3-TraceId': zipkin_attrs.trace_id, + 'X-B3-SpanId': zipkin_attrs.span_id, + 'X-B3-ParentSpanId': zipkin_attrs.parent_span_id, + 'X-B3-Flags': '0', + 'X-B3-Sampled': '1' if zipkin_attrs.is_sampled else '0', + } + def _validate_args(kwargs): if 'kind' in kwargs: diff --git a/tests/zipkin_test.py b/tests/zipkin_test.py index ebb8282..feb3344 100644 --- a/tests/zipkin_test.py +++ b/tests/zipkin_test.py @@ -844,6 +844,48 @@ def test_override_span_name(self): assert span.span_name == 'new_name' assert span.logging_context.span_name == 'new_name' + def test_create_http_headers_for_my_span_no_zipkin_attrs(self): + with zipkin.zipkin_client_span( + service_name='test_service', + span_name='test_span', + transport_handler=MockTransportHandler(), + sample_rate=100.0, + ) as span: + # Not sure how this could ever happen in real life, but if it + # did... + span.zipkin_attrs = None + assert {} == span.create_http_headers_for_my_span() + + def test_create_http_headers_for_my_span_is_sampled(self): + with zipkin.zipkin_client_span( + service_name='test_service', + span_name='test_span', + transport_handler=MockTransportHandler(), + sample_rate=100.0, + ) as span: + assert { + 'X-B3-TraceId': span.zipkin_attrs.trace_id, + 'X-B3-SpanId': span.zipkin_attrs.span_id, + 'X-B3-ParentSpanId': span.zipkin_attrs.parent_span_id, + 'X-B3-Flags': '0', + 'X-B3-Sampled': '1', + } == span.create_http_headers_for_my_span() + + def test_create_http_headers_for_my_span_is_NOT_sampled(self): + with zipkin.zipkin_client_span( + service_name='test_service', + span_name='test_span', + transport_handler=MockTransportHandler(), + sample_rate=0.0, + ) as span: + assert { + 'X-B3-TraceId': span.zipkin_attrs.trace_id, + 'X-B3-SpanId': span.zipkin_attrs.span_id, + 'X-B3-ParentSpanId': span.zipkin_attrs.parent_span_id, + 'X-B3-Flags': '0', + 'X-B3-Sampled': '0', + } == span.create_http_headers_for_my_span() + def test_zipkin_client_span(): context = zipkin.zipkin_client_span('test_service', 'test_span')