def wrapper_use_tracer(*args, **kwargs):
# type: (Any, Any) -> Any
passed_in_parent = kwargs.pop("parent_span", None)
orig_wrapped_span = tracing_context.current_span.get()
wrapper_class = settings.tracing_implementation()
original_span_instance = None
if wrapper_class is not None:
original_span_instance = wrapper_class.get_current_span()
parent_span = common.get_parent_span(passed_in_parent)
ans = None
if parent_span is not None and orig_wrapped_span is None:
common.set_span_contexts(parent_span)
name = name_of_span or common.get_function_and_class_name(func, *args) # type: ignore
child = parent_span.span(name=name)
child.start()
common.set_span_contexts(child)
ans = func(*args, **kwargs) # type: ignore
child.finish()
common.set_span_contexts(parent_span)
if orig_wrapped_span is None and passed_in_parent is None and original_span_instance is None:
parent_span.finish()
common.set_span_contexts(orig_wrapped_span, span_instance=original_span_instance)
else:
ans = func(*args, **kwargs) # type: ignore
return ans
def on_request(self, request):
# type: (PipelineRequest[HttpRequest], Any) -> None
parent_span = tracing_context.current_span.get()
wrapper_class = settings.tracing_implementation()
original_context = [parent_span, None]
if parent_span is None and wrapper_class is not None:
current_span_instance = wrapper_class.get_current_span()
original_context[1] = current_span_instance
parent_span = wrapper_class(current_span_instance)
if parent_span is None:
return
only_propagate = settings.tracing_should_only_propagate()
if only_propagate:
self.set_header(request, parent_span)
return
path = urlparse(request.http_request.url).path
if not path:
path = "/"
child = parent_span.span(name=path)
child.start()
set_span_contexts(child)
self.parent_span_dict[child] = original_context
self.set_header(request, child)
def trace_message(event, parent_span=None):
# type: (EventData, Optional[AbstractSpan]) -> None
"""Add tracing information to this event.
Will open and close a "Azure.EventHubs.message" span, and
add the "DiagnosticId" as app properties of the event.
"""
try:
span_impl_type = settings.tracing_implementation(
) # type: Type[AbstractSpan]
if span_impl_type is not None:
current_span = parent_span or span_impl_type(
span_impl_type.get_current_span())
link = Link({'traceparent': current_span.get_trace_parent()})
with current_span.span(name="Azure.EventHubs.message",
kind=SpanKind.PRODUCER,
links=[link]) as message_span:
message_span.add_attribute("az.namespace",
"Microsoft.EventHub")
if not event.properties:
event.properties = dict()
event.properties.setdefault(
b"Diagnostic-Id",
message_span.get_trace_parent().encode("ascii"))
except Exception as exp: # pylint:disable=broad-except
_LOGGER.warning("trace_message had an exception %r", exp)
def trace_message(message, parent_span=None):
# type: (ServiceBusMessage, Optional[AbstractSpan]) -> None
"""Add tracing information to this message.
Will open and close a "Azure.Servicebus.message" span, and
add the "DiagnosticId" as app properties of the message.
"""
try:
span_impl_type = settings.tracing_implementation(
) # type: Type[AbstractSpan]
if span_impl_type is not None:
current_span = parent_span or span_impl_type(
span_impl_type.get_current_span())
with current_span.span(name=SPAN_NAME_MESSAGE) as message_span:
message_span.kind = SpanKind.PRODUCER
message_span.add_attribute(TRACE_NAMESPACE_PROPERTY,
TRACE_NAMESPACE)
# TODO: Remove intermediary message; this is standin while this var is being renamed in a concurrent PR
if not message.message.application_properties:
message.message.application_properties = dict()
message.message.application_properties.setdefault(
TRACE_PARENT_PROPERTY,
message_span.get_trace_parent().encode(
TRACE_PROPERTY_ENCODING),
)
except Exception as exp: # pylint:disable=broad-except
_log.warning("trace_message had an exception %r", exp)
def wrapper_use_tracer(self, *args, **kwargs):
# type: (Any) -> Any
passed_in_parent = kwargs.pop("parent_span", None)
orig_wrapped_span = tracing_context.current_span.get()
wrapper_class = settings.tracing_implementation()
original_span_instance = None
if wrapper_class is not None:
original_span_instance = wrapper_class.get_current_span()
parent_span = common.get_parent_span(passed_in_parent)
ans = None
if common.should_use_trace(parent_span):
common.set_span_contexts(parent_span)
name = name_of_span or self.__class__.__name__ + "." + func.__name__
child = parent_span.span(name=name)
child.start()
common.set_span_contexts(child)
ans = func(self, *args, **kwargs)
child.finish()
common.set_span_contexts(parent_span)
if orig_wrapped_span is None and passed_in_parent is None and original_span_instance is None:
parent_span.finish()
common.set_span_contexts(orig_wrapped_span,
span_instance=original_span_instance)
else:
ans = func(self, *args, **kwargs)
return ans
def trace_link_message(messages, parent_span=None):
# type: (Union[ServiceBusMessage, Iterable[ServiceBusMessage]], Optional[AbstractSpan]) -> None
"""Link the current message(s) to current span or provided parent span.
Will extract DiagnosticId if available.
"""
trace_messages = (
messages if isinstance(messages, Iterable) # pylint:disable=isinstance-second-argument-not-valid-type
else (messages, ))
try: # pylint:disable=too-many-nested-blocks
span_impl_type = settings.tracing_implementation(
) # type: Type[AbstractSpan]
if span_impl_type is not None:
current_span = parent_span or span_impl_type(
span_impl_type.get_current_span())
if current_span:
for message in trace_messages: # type: ignore
if message.message.application_properties:
traceparent = message.message.application_properties.get(
TRACE_PARENT_PROPERTY,
"").decode(TRACE_PROPERTY_ENCODING)
if traceparent:
current_span.link(
traceparent,
attributes={
SPAN_ENQUEUED_TIME_PROPERTY:
message.message.annotations.get(
TRACE_ENQUEUED_TIME_PROPERTY)
},
)
except Exception as exp: # pylint:disable=broad-except
_log.warning("trace_link_message had an exception %r", exp)
def get_parent_span(parent_span):
# type: (Any) -> Optional[AbstractSpan]
"""
Returns the current span so that the function's span will be its child. It will create a new span if there is
no current span in any of the context.
:param parent_span: The parent_span arg that the user passes into the top level function
:returns: the parent_span of the function to be traced
:rtype: `azure.core.tracing.abstract_span.AbstractSpan`
"""
wrapper_class = settings.tracing_implementation()
if wrapper_class is None:
return None
orig_wrapped_span = tracing_context.current_span.get()
# parent span is given, get from my context, get from the implementation context or make our own
parent_span = orig_wrapped_span if parent_span is None else wrapper_class(parent_span)
if parent_span is None:
current_span = wrapper_class.get_current_span()
parent_span = (
wrapper_class(span=current_span)
if current_span
else wrapper_class(name="azure-sdk-for-python-first_parent_span")
)
return parent_span
def trace_link_message(events, parent_span=None):
# type: (Union[EventData, Iterable[EventData]], Optional[AbstractSpan]) -> None
"""Link the current event(s) to current span or provided parent span.
Will extract DiagnosticId if available.
"""
trace_events = events if isinstance(events, Iterable) else (events, )
try: # pylint:disable=too-many-nested-blocks
span_impl_type = settings.tracing_implementation(
) # type: Type[AbstractSpan]
if span_impl_type is not None:
current_span = parent_span or span_impl_type(
span_impl_type.get_current_span())
if current_span:
for event in trace_events: # type: ignore
if event.properties:
traceparent = event.properties.get(
b"Diagnostic-Id", "").decode("ascii")
if traceparent:
current_span.link(
traceparent,
attributes={
"enqueuedTime":
event.message.annotations.get(
PROP_TIMESTAMP)
})
except Exception as exp: # pylint:disable=broad-except
_LOGGER.warning("trace_link_message had an exception %r", exp)
def send_context_manager():
span_impl_type = settings.tracing_implementation() # type: Type[AbstractSpan]
if span_impl_type is not None:
with span_impl_type(name="Azure.EventHubs.send", kind=SpanKind.CLIENT) as child:
yield child
else:
yield None
def send_trace_context_manager(span_name=SPAN_NAME_SEND):
span_impl_type = settings.tracing_implementation() # type: Type[AbstractSpan]
if span_impl_type is not None:
with span_impl_type(name=span_name, kind=SpanKind.CLIENT) as child:
yield child
else:
yield None
def test_set_span_context(self):
with ContextHelper(environ={"AZURE_SDK_TRACING_IMPLEMENTATION": "opencensus"}):
wrapper = settings.tracing_implementation()
assert wrapper is OpenCensusSpan
assert tracing_context.current_span.get() is None
assert wrapper.get_current_span() is None
parent = OpenCensusSpan()
common.set_span_contexts(parent)
assert parent.span_instance == wrapper.get_current_span()
assert tracing_context.current_span.get() == parent
async def send(
self,
event_data: Union[EventData, EventDataBatch, Iterable[EventData]],
*,
partition_key: Optional[AnyStr] = None,
timeout: Optional[float] = None
) -> None:
"""
Sends an event data and blocks until acknowledgement is
received or operation times out.
:param event_data: The event to be sent. It can be an EventData object, or iterable of EventData objects
:type event_data: ~azure.eventhub.common.EventData, Iterator, Generator, list
:param partition_key: With the given partition_key, event data will land to
a particular partition of the Event Hub decided by the service. partition_key
could be omitted if event_data is of type ~azure.eventhub.EventDataBatch.
:type partition_key: str
:param timeout: The maximum wait time to send the event data.
If not specified, the default wait time specified when the producer was created will be used.
:type timeout: float
:raises: ~azure.eventhub.exceptions.AuthenticationError,
~azure.eventhub.exceptions.ConnectError,
~azure.eventhub.exceptions.ConnectionLostError,
~azure.eventhub.exceptions.EventDataError,
~azure.eventhub.exceptions.EventDataSendError,
~azure.eventhub.exceptions.EventHubError
:return: None
:rtype: None
"""
# Tracing code
async with self._lock:
span_impl_type = (
settings.tracing_implementation()
) # type: Type[AbstractSpan]
child = None
if span_impl_type is not None:
child = span_impl_type(name="Azure.EventHubs.send")
child.kind = SpanKind.CLIENT # Should be PRODUCER
self._check_closed()
wrapper_event_data = self._wrap_eventdata(event_data, child, partition_key)
self._unsent_events = [wrapper_event_data.message]
if span_impl_type is not None and child is not None:
with child:
self._client._add_span_request_attributes( # pylint: disable=protected-access
child
)
await self._send_event_data_with_retry(
timeout=timeout
) # pylint:disable=unexpected-keyword-arg
else:
await self._send_event_data_with_retry(
timeout=timeout
) # pylint:disable=unexpected-keyword-arg
async def send(
self, event_data: Union[EventData, EventDataBatch, Iterable[EventData]],
*, partition_key: Union[str, bytes] = None, timeout: float = None):
"""
Sends an event data and blocks until acknowledgement is
received or operation times out.
:param event_data: The event to be sent. It can be an EventData object, or iterable of EventData objects
:type event_data: ~azure.eventhub.common.EventData, Iterator, Generator, list
:param partition_key: With the given partition_key, event data will land to
a particular partition of the Event Hub decided by the service. partition_key
could be omitted if event_data is of type ~azure.eventhub.EventDataBatch.
:type partition_key: str
:param timeout: The maximum wait time to send the event data.
If not specified, the default wait time specified when the producer was created will be used.
:type timeout: float
:raises: ~azure.eventhub.AuthenticationError, ~azure.eventhub.ConnectError, ~azure.eventhub.ConnectionLostError,
~azure.eventhub.EventDataError, ~azure.eventhub.EventDataSendError, ~azure.eventhub.EventHubError
:return: None
:rtype: None
"""
# Tracing code
span_impl_type = settings.tracing_implementation() # type: Type[AbstractSpan]
child = None
if span_impl_type is not None:
child = span_impl_type(name="Azure.EventHubs.send")
child.kind = SpanKind.CLIENT # Should be PRODUCER
self._check_closed()
if isinstance(event_data, EventData):
if partition_key:
event_data._set_partition_key(partition_key) # pylint: disable=protected-access
wrapper_event_data = event_data
wrapper_event_data._trace_message(child) # pylint: disable=protected-access
else:
if isinstance(event_data, EventDataBatch):
if partition_key and partition_key != event_data._partition_key: # pylint: disable=protected-access
raise EventDataError('The partition_key does not match the one of the EventDataBatch')
wrapper_event_data = event_data #type: ignore
else:
if partition_key:
event_data = _set_partition_key(event_data, partition_key)
event_data = _set_trace_message(event_data)
wrapper_event_data = EventDataBatch._from_batch(event_data, partition_key) # pylint: disable=protected-access
wrapper_event_data.message.on_send_complete = self._on_outcome
self._unsent_events = [wrapper_event_data.message]
if span_impl_type is not None:
with child:
self._client._add_span_request_attributes(child) # pylint: disable=protected-access
await self._send_event_data_with_retry(timeout=timeout) # pylint:disable=unexpected-keyword-arg # TODO: to refactor
else:
await self._send_event_data_with_retry(timeout=timeout) # pylint:disable=unexpected-keyword-arg # TODO: to refactor
def with_current_context(func):
# type: (Callable) -> Any
"""Passes the current spans to the new context the function will be run in.
:param func: The function that will be run in the new context
:return: The target the pass in instead of the function
"""
span_impl_type = settings.tracing_implementation() # type: Type[AbstractSpan]
if span_impl_type is None:
return func
return span_impl_type.with_current_context(func)
def _trace_link_message(self, parent_span=None):
"""Link the current message to current span.
Will extract DiagnosticId if available.
"""
span_impl_type = settings.tracing_implementation() # type: Type[AbstractSpan]
if span_impl_type is not None:
current_span = parent_span or span_impl_type(span_impl_type.get_current_span())
if current_span and self.application_properties:
traceparent = self.application_properties.get(b"Diagnostic-Id", "").decode('ascii')
if traceparent:
current_span.link(traceparent)
def receive_trace_context_manager(receiver, span_name=SPAN_NAME_RECEIVE, links=None):
# type: (ReceiverMixin, str, List[Link]) -> Iterator[None]
"""Tracing"""
span_impl_type = settings.tracing_implementation() # type: Type[AbstractSpan]
if span_impl_type is None:
yield
else:
receive_span = span_impl_type(name=span_name, kind=SpanKind.CONSUMER, links=links)
receiver._add_span_request_attributes(receive_span) # type: ignore # pylint: disable=protected-access
with receive_span:
yield
def _context(self, links=None):
# type: (List[Link]) -> Iterator[None]
"""Tracing"""
span_impl_type = settings.tracing_implementation(
) # type: Type[AbstractSpan]
if span_impl_type is None:
yield
else:
child = span_impl_type(name="Azure.EventHubs.process",
kind=SpanKind.CONSUMER,
links=links)
self._eventhub_client._add_span_request_attributes(child) # type: ignore # pylint: disable=protected-access
with child:
yield
def _context(self, event):
# Tracing
span_impl_type = settings.tracing_implementation(
) # type: Type[AbstractSpan]
if span_impl_type is None:
yield
else:
child = span_impl_type(name="Azure.EventHubs.process")
self._eventhub_client._add_span_request_attributes(child) # pylint: disable=protected-access
child.kind = SpanKind.SERVER
event._trace_link_message(child) # pylint: disable=protected-access
with child:
yield
def _context(self, event):
# type: (Union[EventData, Iterable[EventData]]) -> Iterator[None]
"""Tracing"""
span_impl_type = settings.tracing_implementation(
) # type: Type[AbstractSpan]
if span_impl_type is None:
yield
else:
child = span_impl_type(name="Azure.EventHubs.process")
self._eventhub_client._add_span_request_attributes(child) # type: ignore # pylint: disable=protected-access
child.kind = SpanKind.CONSUMER
trace_link_message(event, child)
with child:
yield
def _trace_message(self, parent_span=None):
"""Add tracing information to this message.
Will open and close a "Azure.EventHubs.message" span, and
add the "DiagnosticId" as app properties of the message.
"""
span_impl_type = settings.tracing_implementation() # type: Type[AbstractSpan]
if span_impl_type is not None:
current_span = parent_span or span_impl_type(span_impl_type.get_current_span())
message_span = current_span.span(name="Azure.EventHubs.message")
message_span.start()
app_prop = dict(self.application_properties)
app_prop.setdefault(b"Diagnostic-Id", message_span.get_trace_parent().encode('ascii'))
self.application_properties = app_prop
message_span.finish()
def set_span_contexts(wrapped_span, span_instance=None):
# type: (Union[AbstractSpan, None], Optional[AbstractSpan]) -> None
"""
Set the sdk context and the implementation context. `span_instance` will be used to set the implementation context
if passed in else will use `wrapped_span.span_instance`.
:param wrapped_span: The `AbstractSpan` to set as the sdk context
:type wrapped_span: `azure.core.tracing.abstract_span.AbstractSpan`
:param span_instance: The span to set as the current span for the implementation context
"""
tracing_context.current_span.set(wrapped_span)
impl_wrapper = settings.tracing_implementation()
if wrapped_span is not None:
span_instance = wrapped_span.span_instance
if impl_wrapper is not None:
impl_wrapper.set_current_span(span_instance)
def wrapper_use_tracer(*args, **kwargs):
# type: (Any, Any) -> Any
merge_span = kwargs.pop('merge_span', False)
passed_in_parent = kwargs.pop("parent_span", None)
span_impl_type = settings.tracing_implementation()
if span_impl_type is None:
return func(*args, **kwargs) # type: ignore
# Merge span is parameter is set, but only if no explicit parent are passed
if merge_span and not passed_in_parent:
return func(*args, **kwargs) # type: ignore
with common.change_context(passed_in_parent):
name = name_of_span or common.get_function_and_class_name(
func, *args) # type: ignore
with span_impl_type(name=name):
return func(*args, **kwargs) # type: ignore
def trace_link_message(event, parent_span=None):
# type: (EventData, Optional[AbstractSpan]) -> None
"""Link the current event to current span or provided parent span.
Will extract DiagnosticId if available.
"""
try:
span_impl_type = settings.tracing_implementation() # type: Type[AbstractSpan]
if span_impl_type is not None:
current_span = parent_span or span_impl_type(
span_impl_type.get_current_span()
)
if current_span and event.properties:
traceparent = event.properties.get(b"Diagnostic-Id", "").decode("ascii")
if traceparent:
current_span.link(traceparent)
except Exception as exp: # pylint:disable=broad-except
_LOGGER.warning("trace_link_message had an exception %r", exp)
def on_request(self, request):
# type: (PipelineRequest) -> None
try:
span_impl_type = settings.tracing_implementation()
if span_impl_type is None:
return
path = urllib.parse.urlparse(request.http_request.url).path
if not path:
path = "/"
span = span_impl_type(name=path)
span.start()
self.set_header(request, span)
request.context[self.TRACING_CONTEXT] = span
except Exception as err: # pylint: disable=broad-except
_LOGGER.warning("Unable to start network span: %s", err)
def on_request(self, request):
# type: (PipelineRequest) -> None
ctxt = request.context.options
try:
span_impl_type = settings.tracing_implementation()
if span_impl_type is None:
return
namer = ctxt.pop('network_span_namer', self._network_span_namer)
span_name = namer(request.http_request)
span = span_impl_type(name=span_name)
span.start()
headers = span.to_header()
request.http_request.headers.update(headers)
request.context[self.TRACING_CONTEXT] = span
except Exception as err: # pylint: disable=broad-except
_LOGGER.warning("Unable to start network span: %s", err)
def receive_trace_context_manager(receiver,
message=None,
span_name=SPAN_NAME_RECEIVE):
# type: (ReceiverMixin, Optional[Union[ServiceBusMessage, Iterable[ServiceBusMessage]]], str) -> Iterator[None]
"""Tracing"""
span_impl_type = settings.tracing_implementation(
) # type: Type[AbstractSpan]
if span_impl_type is None:
yield
else:
receive_span = span_impl_type(name=span_name)
receiver._add_span_request_attributes(receive_span) # type: ignore # pylint: disable=protected-access
receive_span.kind = SpanKind.CONSUMER
# If it is desired to create link before span open
if message:
trace_link_message(message, receive_span)
with receive_span:
yield
def with_current_context(self, func):
# type: (Callable[[Any], Any]) -> Any
"""
Passes the current spans to the new context the function will be run in.
:param func: The function that will be run in the new context
:return: The target the pass in instead of the function
"""
wrapped_span = tracing_context.current_span.get()
wrapper_class = settings.tracing_implementation()
if wrapper_class is not None:
current_impl_span = wrapper_class.get_current_span()
current_impl_tracer = wrapper_class.get_current_tracer()
def call_with_current_context(*args, **kwargs):
if wrapper_class is not None:
wrapper_class.set_current_span(current_impl_span)
wrapper_class.set_current_tracer(current_impl_tracer)
current_span = wrapped_span or wrapper_class(current_impl_span)
self.current_span.set(current_span)
return func(*args, **kwargs)
return call_with_current_context
def change_context(span):
# type: (Optional[AbstractSpan]) -> Generator
"""Execute this block inside the given context and restore it afterwards.
This does not start and ends the span, but just make sure all code is executed within
that span.
If span is None, no-op.
:param span: A span
:type span: AbstractSpan
"""
span_impl_type = settings.tracing_implementation() # type: Type[AbstractSpan]
if span_impl_type is None or span is None:
yield
else:
original_span = span_impl_type.get_current_span()
try:
span_impl_type.set_current_span(span)
yield
finally:
span_impl_type.set_current_span(original_span)
def trace_message(event, parent_span=None):
# type: (EventData, Optional[AbstractSpan]) -> None
"""Add tracing information to this event.
Will open and close a "Azure.EventHubs.message" span, and
add the "DiagnosticId" as app properties of the event.
"""
try:
span_impl_type = settings.tracing_implementation(
) # type: Type[AbstractSpan]
if span_impl_type is not None:
current_span = parent_span or span_impl_type(
span_impl_type.get_current_span())
message_span = current_span.span(name="Azure.EventHubs.message")
message_span.start()
app_prop = dict(event.properties) if event.properties else dict()
app_prop.setdefault(
b"Diagnostic-Id",
message_span.get_trace_parent().encode("ascii"))
event.properties = app_prop
message_span.finish()
except Exception as exp: # pylint:disable=broad-except
_LOGGER.warning("trace_message had an exception %r", exp)
