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)