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

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 

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

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

11# sell 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 

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

23# IN THE SOFTWARE. 

24# 

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

26from __future__ import annotations 

27from types import TracebackType 

28from typing import Any, Union, Generic, TypeVar, List, Dict, Optional, Iterable, Type 

29from typing_extensions import AsyncContextManager 

30 

31from azure.core.pipeline import PipelineRequest, PipelineResponse, PipelineContext 

32from azure.core.pipeline.policies import AsyncHTTPPolicy, SansIOHTTPPolicy 

33from ._tools_async import await_result as _await_result 

34from ._base import cleanup_kwargs_for_transport 

35from .transport import AsyncHttpTransport 

36 

37AsyncHTTPResponseType = TypeVar("AsyncHTTPResponseType") 

38HTTPRequestType = TypeVar("HTTPRequestType") 

39 

40 

41class _SansIOAsyncHTTPPolicyRunner( 

42 AsyncHTTPPolicy[HTTPRequestType, AsyncHTTPResponseType] 

43): # pylint: disable=unsubscriptable-object 

44 """Async implementation of the SansIO policy. 

45 

46 Modifies the request and sends to the next policy in the chain. 

47 

48 :param policy: A SansIO policy. 

49 :type policy: ~azure.core.pipeline.policies.SansIOHTTPPolicy 

50 """ 

51 

52 def __init__(self, policy: SansIOHTTPPolicy[HTTPRequestType, AsyncHTTPResponseType]) -> None: 

53 super(_SansIOAsyncHTTPPolicyRunner, self).__init__() 

54 self._policy = policy 

55 

56 async def send( 

57 self, request: PipelineRequest[HTTPRequestType] 

58 ) -> PipelineResponse[HTTPRequestType, AsyncHTTPResponseType]: 

59 """Modifies the request and sends to the next policy in the chain. 

60 

61 :param request: The PipelineRequest object. 

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

63 :return: The PipelineResponse object. 

64 :rtype: ~azure.core.pipeline.PipelineResponse 

65 """ 

66 await _await_result(self._policy.on_request, request) 

67 response: PipelineResponse[HTTPRequestType, AsyncHTTPResponseType] 

68 try: 

69 response = await self.next.send(request) 

70 except Exception: # pylint: disable=broad-except 

71 await _await_result(self._policy.on_exception, request) 

72 raise 

73 else: 

74 await _await_result(self._policy.on_response, request, response) 

75 return response 

76 

77 

78class _AsyncTransportRunner( 

79 AsyncHTTPPolicy[HTTPRequestType, AsyncHTTPResponseType] 

80): # pylint: disable=unsubscriptable-object 

81 """Async Transport runner. 

82 

83 Uses specified HTTP transport type to send request and returns response. 

84 

85 :param sender: The async Http Transport instance. 

86 :type sender: ~azure.core.pipeline.transport.AsyncHttpTransport 

87 """ 

88 

89 def __init__(self, sender: AsyncHttpTransport[HTTPRequestType, AsyncHTTPResponseType]) -> None: 

90 super(_AsyncTransportRunner, self).__init__() 

91 self._sender = sender 

92 

93 async def send( 

94 self, request: PipelineRequest[HTTPRequestType] 

95 ) -> PipelineResponse[HTTPRequestType, AsyncHTTPResponseType]: 

96 """Async HTTP transport send method. 

97 

98 :param request: The PipelineRequest object. 

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

100 :return: The PipelineResponse object. 

101 :rtype: ~azure.core.pipeline.PipelineResponse 

102 """ 

103 cleanup_kwargs_for_transport(request.context.options) 

104 return PipelineResponse( 

105 request.http_request, 

106 await self._sender.send(request.http_request, **request.context.options), 

107 request.context, 

108 ) 

109 

110 

111class AsyncPipeline(AsyncContextManager["AsyncPipeline"], Generic[HTTPRequestType, AsyncHTTPResponseType]): 

112 """Async pipeline implementation. 

113 

114 This is implemented as a context manager, that will activate the context 

115 of the HTTP sender. 

116 

117 :param transport: The async Http Transport instance. 

118 :type transport: ~azure.core.pipeline.transport.AsyncHttpTransport 

119 :param list policies: List of configured policies. 

120 

121 .. admonition:: Example: 

122 

123 .. literalinclude:: ../samples/test_example_async.py 

124 :start-after: [START build_async_pipeline] 

125 :end-before: [END build_async_pipeline] 

126 :language: python 

127 :dedent: 4 

128 :caption: Builds the async pipeline for asynchronous transport. 

129 """ 

130 

131 def __init__( 

132 self, 

133 transport: AsyncHttpTransport[HTTPRequestType, AsyncHTTPResponseType], 

134 policies: Optional[ 

135 Iterable[ 

136 Union[ 

137 AsyncHTTPPolicy[HTTPRequestType, AsyncHTTPResponseType], 

138 SansIOHTTPPolicy[HTTPRequestType, AsyncHTTPResponseType], 

139 ] 

140 ] 

141 ] = None, 

142 ) -> None: 

143 self._impl_policies: List[AsyncHTTPPolicy[HTTPRequestType, AsyncHTTPResponseType]] = [] 

144 self._transport = transport 

145 

146 for policy in policies or []: 

147 if isinstance(policy, SansIOHTTPPolicy): 

148 self._impl_policies.append(_SansIOAsyncHTTPPolicyRunner(policy)) 

149 elif policy: 

150 self._impl_policies.append(policy) 

151 for index in range(len(self._impl_policies) - 1): 

152 self._impl_policies[index].next = self._impl_policies[index + 1] 

153 if self._impl_policies: 

154 self._impl_policies[-1].next = _AsyncTransportRunner(self._transport) 

155 

156 async def __aenter__(self) -> AsyncPipeline[HTTPRequestType, AsyncHTTPResponseType]: 

157 await self._transport.__aenter__() 

158 return self 

159 

160 async def __aexit__( 

161 self, 

162 exc_type: Optional[Type[BaseException]] = None, 

163 exc_value: Optional[BaseException] = None, 

164 traceback: Optional[TracebackType] = None, 

165 ) -> None: 

166 await self._transport.__aexit__(exc_type, exc_value, traceback) 

167 

168 async def _prepare_multipart_mixed_request(self, request: HTTPRequestType) -> None: 

169 """Will execute the multipart policies. 

170 

171 Does nothing if "set_multipart_mixed" was never called. 

172 

173 :param request: The HTTP request object. 

174 :type request: ~azure.core.rest.HttpRequest 

175 """ 

176 multipart_mixed_info = request.multipart_mixed_info # type: ignore 

177 if not multipart_mixed_info: 

178 return 

179 

180 requests: List[HTTPRequestType] = multipart_mixed_info[0] 

181 policies: List[SansIOHTTPPolicy] = multipart_mixed_info[1] 

182 pipeline_options: Dict[str, Any] = multipart_mixed_info[3] 

183 

184 async def prepare_requests(req): 

185 if req.multipart_mixed_info: 

186 # Recursively update changeset "sub requests" 

187 await self._prepare_multipart_mixed_request(req) 

188 context = PipelineContext(None, **pipeline_options) 

189 pipeline_request = PipelineRequest(req, context) 

190 for policy in policies: 

191 await _await_result(policy.on_request, pipeline_request) 

192 

193 # Not happy to make this code asyncio specific, but that's multipart only for now 

194 # If we need trio and multipart, let's reinvesitgate that later 

195 import asyncio 

196 

197 await asyncio.gather(*[prepare_requests(req) for req in requests]) 

198 

199 async def _prepare_multipart(self, request: HTTPRequestType) -> None: 

200 # This code is fine as long as HTTPRequestType is actually 

201 # azure.core.pipeline.transport.HTTPRequest, bu we don't check it in here 

202 # since we didn't see (yet) pipeline usage where it's not this actual instance 

203 # class used 

204 await self._prepare_multipart_mixed_request(request) 

205 request.prepare_multipart_body() # type: ignore 

206 

207 async def run( 

208 self, request: HTTPRequestType, **kwargs: Any 

209 ) -> PipelineResponse[HTTPRequestType, AsyncHTTPResponseType]: 

210 """Runs the HTTP Request through the chained policies. 

211 

212 :param request: The HTTP request object. 

213 :type request: ~azure.core.pipeline.transport.HttpRequest 

214 :return: The PipelineResponse object. 

215 :rtype: ~azure.core.pipeline.PipelineResponse 

216 """ 

217 await self._prepare_multipart(request) 

218 context = PipelineContext(self._transport, **kwargs) 

219 pipeline_request = PipelineRequest(request, context) 

220 first_node = self._impl_policies[0] if self._impl_policies else _AsyncTransportRunner(self._transport) 

221 return await first_node.send(pipeline_request)