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