Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/azure/core/pipeline/policies/_distributed_tracing.py: 42%

1# -------------------------------------------------------------------------- 

2# 

3# Copyright (c) Microsoft Corporation. All rights reserved. 

4# 

5# The MIT License (MIT) 

6# 

7# Permission is hereby granted, free of charge, to any person obtaining a copy 

8# of this software and associated documentation files (the ""Software""), to deal 

9# in the Software without restriction, including without limitation the rights 

10# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 

11# copies of the Software, and to permit persons to whom the Software is 

12# furnished to do so, subject to the following conditions: 

13# 

14# The above copyright notice and this permission notice shall be included in 

15# all copies or substantial portions of the Software. 

16# 

17# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 

18# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 

19# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 

20# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 

21# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 

22# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 

23# THE SOFTWARE. 

24# 

25# -------------------------------------------------------------------------- 

26"""Traces network calls using the implementation library from the settings.""" 

27import logging 

28import sys 

29import urllib.parse 

30from typing import TYPE_CHECKING, Optional, Tuple, TypeVar, Union, Any, Type 

31from types import TracebackType 

32 

33from azure.core.pipeline import PipelineRequest, PipelineResponse 

34from azure.core.pipeline.policies import SansIOHTTPPolicy 

35from azure.core.pipeline.transport import HttpResponse as LegacyHttpResponse, HttpRequest as LegacyHttpRequest 

36from azure.core.rest import HttpResponse, HttpRequest 

37from azure.core.settings import settings 

38from azure.core.tracing import SpanKind 

39 

40if TYPE_CHECKING: 

41 from azure.core.tracing._abstract_span import ( 

42 AbstractSpan, 

43 ) 

44 

45HTTPResponseType = TypeVar("HTTPResponseType", HttpResponse, LegacyHttpResponse) 

46HTTPRequestType = TypeVar("HTTPRequestType", HttpRequest, LegacyHttpRequest) 

47ExcInfo = Tuple[Type[BaseException], BaseException, TracebackType] 

48OptExcInfo = Union[ExcInfo, Tuple[None, None, None]] 

49 

50_LOGGER = logging.getLogger(__name__) 

51 

52 

53def _default_network_span_namer(http_request: HTTPRequestType) -> str: 

54 """Extract the path to be used as network span name. 

55 

56 :param http_request: The HTTP request 

57 :type http_request: ~azure.core.pipeline.transport.HttpRequest 

58 :returns: The string to use as network span name 

59 :rtype: str 

60 """ 

61 path = urllib.parse.urlparse(http_request.url).path 

62 if not path: 

63 path = "/" 

64 return path 

65 

66 

67class DistributedTracingPolicy(SansIOHTTPPolicy[HTTPRequestType, HTTPResponseType]): 

68 """The policy to create spans for Azure calls. 

69 

70 :keyword network_span_namer: A callable to customize the span name 

71 :type network_span_namer: callable[[~azure.core.pipeline.transport.HttpRequest], str] 

72 :keyword tracing_attributes: Attributes to set on all created spans 

73 :type tracing_attributes: dict[str, str] 

74 """ 

75 

76 TRACING_CONTEXT = "TRACING_CONTEXT" 

77 _REQUEST_ID = "x-ms-client-request-id" 

78 _RESPONSE_ID = "x-ms-request-id" 

79 

80 def __init__(self, **kwargs: Any): 

81 self._network_span_namer = kwargs.get("network_span_namer", _default_network_span_namer) 

82 self._tracing_attributes = kwargs.get("tracing_attributes", {}) 

83 

84 def on_request(self, request: PipelineRequest[HTTPRequestType]) -> None: 

85 ctxt = request.context.options 

86 try: 

87 span_impl_type = settings.tracing_implementation() 

88 if span_impl_type is None: 

89 return 

90 

91 namer = ctxt.pop("network_span_namer", self._network_span_namer) 

92 span_name = namer(request.http_request) 

93 

94 span = span_impl_type(name=span_name, kind=SpanKind.CLIENT) 

95 for attr, value in self._tracing_attributes.items(): 

96 span.add_attribute(attr, value) 

97 span.start() 

98 

99 headers = span.to_header() 

100 request.http_request.headers.update(headers) 

101 

102 request.context[self.TRACING_CONTEXT] = span 

103 except Exception as err: # pylint: disable=broad-except 

104 _LOGGER.warning("Unable to start network span: %s", err) 

105 

106 def end_span( 

107 self, 

108 request: PipelineRequest[HTTPRequestType], 

109 response: Optional[HTTPResponseType] = None, 

110 exc_info: Optional[OptExcInfo] = None, 

111 ) -> None: 

112 """Ends the span that is tracing the network and updates its status. 

113 

114 :param request: The PipelineRequest object 

115 :type request: ~azure.core.pipeline.PipelineRequest 

116 :param response: The HttpResponse object 

117 :type response: ~azure.core.rest.HTTPResponse or ~azure.core.pipeline.transport.HttpResponse 

118 :param exc_info: The exception information 

119 :type exc_info: tuple 

120 """ 

121 if self.TRACING_CONTEXT not in request.context: 

122 return 

123 

124 span: "AbstractSpan" = request.context[self.TRACING_CONTEXT] 

125 http_request: Union[HttpRequest, LegacyHttpRequest] = request.http_request 

126 if span is not None: 

127 span.set_http_attributes(http_request, response=response) 

128 request_id = http_request.headers.get(self._REQUEST_ID) 

129 if request_id is not None: 

130 span.add_attribute(self._REQUEST_ID, request_id) 

131 if response and self._RESPONSE_ID in response.headers: 

132 span.add_attribute(self._RESPONSE_ID, response.headers[self._RESPONSE_ID]) 

133 if exc_info: 

134 span.__exit__(*exc_info) 

135 else: 

136 span.finish() 

137 

138 def on_response( 

139 self, request: PipelineRequest[HTTPRequestType], response: PipelineResponse[HTTPRequestType, HTTPResponseType] 

140 ) -> None: 

141 self.end_span(request, response=response.http_response) 

142 

143 def on_exception(self, request: PipelineRequest[HTTPRequestType]) -> None: 

144 self.end_span(request, exc_info=sys.exc_info())