# ------------------------------------
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
# ------------------------------------
"""Protocol that defines what functions wrappers of tracing libraries should implement."""
from __future__ import annotations
from enum import Enum
from urllib.parse import urlparse
from typing import Any, Sequence, Optional, Union, Callable, Dict, Type, Generic, TypeVar
from types import TracebackType
from typing_extensions import Protocol, ContextManager, runtime_checkable
from azure.core.pipeline.transport import HttpRequest, HttpResponse, AsyncHttpResponse
from azure.core.rest import (
HttpResponse as RestHttpResponse,
AsyncHttpResponse as AsyncRestHttpResponse,
HttpRequest as RestHttpRequest,
)
HttpResponseType = Union[HttpResponse, AsyncHttpResponse, RestHttpResponse, AsyncRestHttpResponse]
HttpRequestType = Union[HttpRequest, RestHttpRequest]
AttributeValue = Union[
str,
bool,
int,
float,
Sequence[str],
Sequence[bool],
Sequence[int],
Sequence[float],
]
Attributes = Dict[str, AttributeValue]
SpanType = TypeVar("SpanType")
[docs]class SpanKind(Enum):
UNSPECIFIED = 1
SERVER = 2
CLIENT = 3
PRODUCER = 4
CONSUMER = 5
INTERNAL = 6
[docs]@runtime_checkable
class AbstractSpan(Protocol, Generic[SpanType]):
"""Wraps a span from a distributed tracing implementation.
If a span is given wraps the span. Else a new span is created.
The optional argument name is given to the new span.
:param span: The span to wrap
:type span: Any
:param name: The name of the span
:type name: str
"""
def __init__( # pylint: disable=super-init-not-called
self, span: Optional[SpanType] = None, name: Optional[str] = None, **kwargs: Any
) -> None:
pass
[docs] def span(self, name: str = "child_span", **kwargs: Any) -> AbstractSpan[SpanType]:
"""
Create a child span for the current span and append it to the child spans list.
The child span must be wrapped by an implementation of AbstractSpan
:param name: The name of the child span
:type name: str
:return: The child span
:rtype: AbstractSpan
"""
...
@property
def kind(self) -> Optional[SpanKind]:
"""Get the span kind of this span.
:rtype: SpanKind
:return: The span kind of this span
"""
...
@kind.setter
def kind(self, value: SpanKind) -> None:
"""Set the span kind of this span.
:param value: The span kind of this span
:type value: SpanKind
"""
...
def __enter__(self) -> AbstractSpan[SpanType]:
"""Start a span."""
...
def __exit__(
self,
exception_type: Optional[Type[BaseException]],
exception_value: Optional[BaseException],
traceback: TracebackType,
) -> None:
"""Finish a span.
:param exception_type: The type of the exception
:type exception_type: type
:param exception_value: The value of the exception
:type exception_value: Exception
:param traceback: The traceback of the exception
:type traceback: Traceback
"""
...
[docs] def start(self) -> None:
"""Set the start time for a span."""
...
[docs] def finish(self) -> None:
"""Set the end time for a span."""
...
[docs] def add_attribute(self, key: str, value: Union[str, int]) -> None:
"""
Add attribute (key value pair) to the current span.
:param key: The key of the key value pair
:type key: str
:param value: The value of the key value pair
:type value: Union[str, int]
"""
...
[docs] def set_http_attributes(self, request: HttpRequestType, response: Optional[HttpResponseType] = None) -> None:
"""
Add correct attributes for a http client span.
:param request: The request made
:type request: azure.core.rest.HttpRequest
:param response: The response received by the server. Is None if no response received.
:type response: ~azure.core.pipeline.transport.HttpResponse or ~azure.core.pipeline.transport.AsyncHttpResponse
"""
...
[docs] def get_trace_parent(self) -> str:
"""Return traceparent string.
:return: a traceparent string
:rtype: str
"""
...
@property
def span_instance(self) -> SpanType:
"""
Returns the span the class is wrapping.
"""
...
[docs] @classmethod
def link(cls, traceparent: str, attributes: Optional[Attributes] = None) -> None:
"""
Given a traceparent, extracts the context and links the context to the current tracer.
:param traceparent: A string representing a traceparent
:type traceparent: str
:param attributes: Any additional attributes that should be added to link
:type attributes: dict
"""
...
[docs] @classmethod
def get_current_span(cls) -> SpanType:
"""
Get the current span from the execution context. Return None otherwise.
:return: The current span
:rtype: AbstractSpan
"""
...
[docs] @classmethod
def get_current_tracer(cls) -> Any:
"""
Get the current tracer from the execution context. Return None otherwise.
:return: The current tracer
:rtype: Any
"""
...
[docs] @classmethod
def set_current_span(cls, span: SpanType) -> None:
"""Set the given span as the current span in the execution context.
:param span: The span to set as the current span
:type span: Any
"""
...
[docs] @classmethod
def set_current_tracer(cls, tracer: Any) -> None:
"""Set the given tracer as the current tracer in the execution context.
:param tracer: The tracer to set as the current tracer
:type tracer: Any
"""
...
[docs] @classmethod
def change_context(cls, span: SpanType) -> ContextManager[SpanType]:
"""Change the context for the life of this context manager.
:param span: The span to run in the new context
:type span: Any
:rtype: contextmanager
:return: A context manager that will run the given span in the new context
"""
...
[docs] @classmethod
def with_current_context(cls, func: Callable) -> Callable:
"""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
:type func: callable
:return: The target the pass in instead of the function
:rtype: callable
"""
...
[docs]class HttpSpanMixin:
"""Can be used to get HTTP span attributes settings for free."""
_SPAN_COMPONENT = "component"
_HTTP_USER_AGENT = "http.user_agent"
_HTTP_METHOD = "http.method"
_HTTP_URL = "http.url"
_HTTP_STATUS_CODE = "http.status_code"
_NET_PEER_NAME = "net.peer.name"
_NET_PEER_PORT = "net.peer.port"
[docs] def set_http_attributes(
self: AbstractSpan, request: HttpRequestType, response: Optional[HttpResponseType] = None
) -> None:
"""
Add correct attributes for a http client span.
:param request: The request made
:type request: azure.core.rest.HttpRequest
:param response: The response received from the server. Is None if no response received.
:type response: ~azure.core.pipeline.transport.HttpResponse or ~azure.core.pipeline.transport.AsyncHttpResponse
"""
# Also see https://github.com/python/mypy/issues/5837
self.kind = SpanKind.CLIENT
self.add_attribute(HttpSpanMixin._SPAN_COMPONENT, "http")
self.add_attribute(HttpSpanMixin._HTTP_METHOD, request.method)
self.add_attribute(HttpSpanMixin._HTTP_URL, request.url)
parsed_url = urlparse(request.url)
if parsed_url.hostname:
self.add_attribute(HttpSpanMixin._NET_PEER_NAME, parsed_url.hostname)
if parsed_url.port and parsed_url.port not in [80, 443]:
self.add_attribute(HttpSpanMixin._NET_PEER_PORT, parsed_url.port)
user_agent = request.headers.get("User-Agent")
if user_agent:
self.add_attribute(HttpSpanMixin._HTTP_USER_AGENT, user_agent)
if response and response.status_code:
self.add_attribute(HttpSpanMixin._HTTP_STATUS_CODE, response.status_code)
else:
self.add_attribute(HttpSpanMixin._HTTP_STATUS_CODE, 504)
[docs]class Link:
"""
This is a wrapper class to link the context to the current tracer.
:param headers: A dictionary of the request header as key value pairs.
:type headers: dict
:param attributes: Any additional attributes that should be added to link
:type attributes: dict
"""
def __init__(self, headers: Dict[str, str], attributes: Optional[Attributes] = None) -> None:
self.headers = headers
self.attributes = attributes