1# ------------------------------------
2# Copyright (c) Microsoft Corporation.
3# Licensed under the MIT License.
4# ------------------------------------
5"""Protocol that defines what functions wrappers of tracing libraries should implement."""
6from __future__ import annotations
7from enum import Enum
8from urllib.parse import urlparse
9
10from typing import Any, Sequence, Optional, Union, Callable, Dict, Type, Generic, TypeVar
11from types import TracebackType
12from typing_extensions import Protocol, ContextManager, runtime_checkable
13from azure.core.pipeline.transport import HttpRequest, HttpResponse, AsyncHttpResponse
14from azure.core.rest import (
15 HttpResponse as RestHttpResponse,
16 AsyncHttpResponse as AsyncRestHttpResponse,
17 HttpRequest as RestHttpRequest,
18)
19
20HttpResponseType = Union[HttpResponse, AsyncHttpResponse, RestHttpResponse, AsyncRestHttpResponse]
21HttpRequestType = Union[HttpRequest, RestHttpRequest]
22
23AttributeValue = Union[
24 str,
25 bool,
26 int,
27 float,
28 Sequence[str],
29 Sequence[bool],
30 Sequence[int],
31 Sequence[float],
32]
33Attributes = Dict[str, AttributeValue]
34SpanType = TypeVar("SpanType")
35
36
37class SpanKind(Enum):
38 UNSPECIFIED = 1
39 SERVER = 2
40 CLIENT = 3
41 PRODUCER = 4
42 CONSUMER = 5
43 INTERNAL = 6
44
45
46@runtime_checkable
47class AbstractSpan(Protocol, Generic[SpanType]):
48 """Wraps a span from a distributed tracing implementation.
49
50 If a span is given wraps the span. Else a new span is created.
51 The optional argument name is given to the new span.
52
53 :param span: The span to wrap
54 :type span: Any
55 :param name: The name of the span
56 :type name: str
57 """
58
59 def __init__( # pylint: disable=super-init-not-called
60 self, span: Optional[SpanType] = None, name: Optional[str] = None, **kwargs: Any
61 ) -> None:
62 pass
63
64 def span(self, name: str = "child_span", **kwargs: Any) -> AbstractSpan[SpanType]:
65 """
66 Create a child span for the current span and append it to the child spans list.
67 The child span must be wrapped by an implementation of AbstractSpan
68
69 :param name: The name of the child span
70 :type name: str
71 :return: The child span
72 :rtype: AbstractSpan
73 """
74 ...
75
76 @property
77 def kind(self) -> Optional[SpanKind]:
78 """Get the span kind of this span.
79
80 :rtype: SpanKind
81 :return: The span kind of this span
82 """
83 ...
84
85 @kind.setter
86 def kind(self, value: SpanKind) -> None:
87 """Set the span kind of this span.
88
89 :param value: The span kind of this span
90 :type value: SpanKind
91 """
92 ...
93
94 def __enter__(self) -> AbstractSpan[SpanType]:
95 """Start a span."""
96 ...
97
98 def __exit__(
99 self,
100 exception_type: Optional[Type[BaseException]],
101 exception_value: Optional[BaseException],
102 traceback: TracebackType,
103 ) -> None:
104 """Finish a span.
105
106 :param exception_type: The type of the exception
107 :type exception_type: type
108 :param exception_value: The value of the exception
109 :type exception_value: Exception
110 :param traceback: The traceback of the exception
111 :type traceback: Traceback
112 """
113 ...
114
115 def start(self) -> None:
116 """Set the start time for a span."""
117 ...
118
119 def finish(self) -> None:
120 """Set the end time for a span."""
121 ...
122
123 def to_header(self) -> Dict[str, str]:
124 """Returns a dictionary with the header labels and values.
125
126 :return: A dictionary with the header labels and values
127 :rtype: dict
128 """
129 ...
130
131 def add_attribute(self, key: str, value: Union[str, int]) -> None:
132 """
133 Add attribute (key value pair) to the current span.
134
135 :param key: The key of the key value pair
136 :type key: str
137 :param value: The value of the key value pair
138 :type value: Union[str, int]
139 """
140 ...
141
142 def set_http_attributes(self, request: HttpRequestType, response: Optional[HttpResponseType] = None) -> None:
143 """
144 Add correct attributes for a http client span.
145
146 :param request: The request made
147 :type request: azure.core.rest.HttpRequest
148 :param response: The response received by the server. Is None if no response received.
149 :type response: ~azure.core.pipeline.transport.HttpResponse or ~azure.core.pipeline.transport.AsyncHttpResponse
150 """
151 ...
152
153 def get_trace_parent(self) -> str:
154 """Return traceparent string.
155
156 :return: a traceparent string
157 :rtype: str
158 """
159 ...
160
161 @property
162 def span_instance(self) -> SpanType:
163 """
164 Returns the span the class is wrapping.
165 """
166 ...
167
168 @classmethod
169 def link(cls, traceparent: str, attributes: Optional[Attributes] = None) -> None:
170 """
171 Given a traceparent, extracts the context and links the context to the current tracer.
172
173 :param traceparent: A string representing a traceparent
174 :type traceparent: str
175 :param attributes: Any additional attributes that should be added to link
176 :type attributes: dict
177 """
178 ...
179
180 @classmethod
181 def link_from_headers(cls, headers: Dict[str, str], attributes: Optional[Attributes] = None) -> None:
182 """
183 Given a dictionary, extracts the context and links the context to the current tracer.
184
185 :param headers: A dictionary of the request header as key value pairs.
186 :type headers: dict
187 :param attributes: Any additional attributes that should be added to link
188 :type attributes: dict
189 """
190 ...
191
192 @classmethod
193 def get_current_span(cls) -> SpanType:
194 """
195 Get the current span from the execution context. Return None otherwise.
196
197 :return: The current span
198 :rtype: AbstractSpan
199 """
200 ...
201
202 @classmethod
203 def get_current_tracer(cls) -> Any:
204 """
205 Get the current tracer from the execution context. Return None otherwise.
206
207 :return: The current tracer
208 :rtype: Any
209 """
210 ...
211
212 @classmethod
213 def set_current_span(cls, span: SpanType) -> None:
214 """Set the given span as the current span in the execution context.
215
216 :param span: The span to set as the current span
217 :type span: Any
218 """
219 ...
220
221 @classmethod
222 def set_current_tracer(cls, tracer: Any) -> None:
223 """Set the given tracer as the current tracer in the execution context.
224
225 :param tracer: The tracer to set as the current tracer
226 :type tracer: Any
227 """
228 ...
229
230 @classmethod
231 def change_context(cls, span: SpanType) -> ContextManager[SpanType]:
232 """Change the context for the life of this context manager.
233
234 :param span: The span to run in the new context
235 :type span: Any
236 :rtype: contextmanager
237 :return: A context manager that will run the given span in the new context
238 """
239 ...
240
241 @classmethod
242 def with_current_context(cls, func: Callable) -> Callable:
243 """Passes the current spans to the new context the function will be run in.
244
245 :param func: The function that will be run in the new context
246 :type func: callable
247 :return: The target the pass in instead of the function
248 :rtype: callable
249 """
250 ...
251
252
253class HttpSpanMixin:
254 """Can be used to get HTTP span attributes settings for free."""
255
256 _SPAN_COMPONENT = "component"
257 _HTTP_USER_AGENT = "http.user_agent"
258 _HTTP_METHOD = "http.method"
259 _HTTP_URL = "http.url"
260 _HTTP_STATUS_CODE = "http.status_code"
261 _NET_PEER_NAME = "net.peer.name"
262 _NET_PEER_PORT = "net.peer.port"
263 _ERROR_TYPE = "error.type"
264
265 def set_http_attributes(
266 self: AbstractSpan, request: HttpRequestType, response: Optional[HttpResponseType] = None
267 ) -> None:
268 """
269 Add correct attributes for a http client span.
270
271 :param request: The request made
272 :type request: azure.core.rest.HttpRequest
273 :param response: The response received from the server. Is None if no response received.
274 :type response: ~azure.core.pipeline.transport.HttpResponse or ~azure.core.pipeline.transport.AsyncHttpResponse
275 """
276 # Also see https://github.com/python/mypy/issues/5837
277 self.kind = SpanKind.CLIENT
278 self.add_attribute(HttpSpanMixin._SPAN_COMPONENT, "http")
279 self.add_attribute(HttpSpanMixin._HTTP_METHOD, request.method)
280 self.add_attribute(HttpSpanMixin._HTTP_URL, request.url)
281
282 parsed_url = urlparse(request.url)
283 if parsed_url.hostname:
284 self.add_attribute(HttpSpanMixin._NET_PEER_NAME, parsed_url.hostname)
285 if parsed_url.port and parsed_url.port not in [80, 443]:
286 self.add_attribute(HttpSpanMixin._NET_PEER_PORT, parsed_url.port)
287
288 user_agent = request.headers.get("User-Agent")
289 if user_agent:
290 self.add_attribute(HttpSpanMixin._HTTP_USER_AGENT, user_agent)
291 if response and response.status_code:
292 self.add_attribute(HttpSpanMixin._HTTP_STATUS_CODE, response.status_code)
293 if response.status_code >= 400:
294 self.add_attribute(HttpSpanMixin._ERROR_TYPE, str(response.status_code))
295 else:
296 self.add_attribute(HttpSpanMixin._HTTP_STATUS_CODE, 504)
297 self.add_attribute(HttpSpanMixin._ERROR_TYPE, "504")
298
299
300class Link:
301 """
302 This is a wrapper class to link the context to the current tracer.
303 :param headers: A dictionary of the request header as key value pairs.
304 :type headers: dict
305 :param attributes: Any additional attributes that should be added to link
306 :type attributes: dict
307 """
308
309 def __init__(self, headers: Dict[str, str], attributes: Optional[Attributes] = None) -> None:
310 self.headers = headers
311 self.attributes = attributes