Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/azure/core/tracing/_abstract_span.py: 62%

Shortcuts on this page

r m x   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

108 statements  

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