Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/google/api_core/future/async_future.py: 45%

1# Copyright 2020, Google LLC 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14 

15"""AsyncIO implementation of the abstract base Future class.""" 

16 

17import asyncio 

18 

19from google.api_core import exceptions 

20from google.api_core import retry 

21from google.api_core import retry_async 

22from google.api_core.future import base 

23 

24 

25class _OperationNotComplete(Exception): 

26 """Private exception used for polling via retry.""" 

27 

28 pass 

29 

30 

31RETRY_PREDICATE = retry.if_exception_type( 

32 _OperationNotComplete, 

33 exceptions.TooManyRequests, 

34 exceptions.InternalServerError, 

35 exceptions.BadGateway, 

36) 

37DEFAULT_RETRY = retry_async.AsyncRetry(predicate=RETRY_PREDICATE) 

38 

39 

40class AsyncFuture(base.Future): 

41 """A Future that polls peer service to self-update. 

42 

43 The :meth:`done` method should be implemented by subclasses. The polling 

44 behavior will repeatedly call ``done`` until it returns True. 

45 

46 .. note:: 

47 

48 Privacy here is intended to prevent the final class from 

49 overexposing, not to prevent subclasses from accessing methods. 

50 

51 Args: 

52 retry (google.api_core.retry.Retry): The retry configuration used 

53 when polling. This can be used to control how often :meth:`done` 

54 is polled. Regardless of the retry's ``deadline``, it will be 

55 overridden by the ``timeout`` argument to :meth:`result`. 

56 """ 

57 

58 def __init__(self, retry=DEFAULT_RETRY): 

59 super().__init__() 

60 self._retry = retry 

61 self._future = asyncio.get_event_loop().create_future() 

62 self._background_task = None 

63 

64 async def done(self, retry=DEFAULT_RETRY): 

65 """Checks to see if the operation is complete. 

66 

67 Args: 

68 retry (google.api_core.retry.Retry): (Optional) How to retry the RPC. 

69 

70 Returns: 

71 bool: True if the operation is complete, False otherwise. 

72 """ 

73 # pylint: disable=redundant-returns-doc, missing-raises-doc 

74 raise NotImplementedError() 

75 

76 async def _done_or_raise(self): 

77 """Check if the future is done and raise if it's not.""" 

78 result = await self.done() 

79 if not result: 

80 raise _OperationNotComplete() 

81 

82 async def running(self): 

83 """True if the operation is currently running.""" 

84 result = await self.done() 

85 return not result 

86 

87 async def _blocking_poll(self, timeout=None): 

88 """Poll and await for the Future to be resolved. 

89 

90 Args: 

91 timeout (int): 

92 How long (in seconds) to wait for the operation to complete. 

93 If None, wait indefinitely. 

94 """ 

95 if self._future.done(): 

96 return 

97 

98 retry_ = self._retry.with_timeout(timeout) 

99 

100 try: 

101 await retry_(self._done_or_raise)() 

102 except exceptions.RetryError: 

103 raise asyncio.TimeoutError( 

104 "Operation did not complete within the designated " "timeout." 

105 ) 

106 

107 async def result(self, timeout=None): 

108 """Get the result of the operation. 

109 

110 Args: 

111 timeout (int): 

112 How long (in seconds) to wait for the operation to complete. 

113 If None, wait indefinitely. 

114 

115 Returns: 

116 google.protobuf.Message: The Operation's result. 

117 

118 Raises: 

119 google.api_core.GoogleAPICallError: If the operation errors or if 

120 the timeout is reached before the operation completes. 

121 """ 

122 await self._blocking_poll(timeout=timeout) 

123 return self._future.result() 

124 

125 async def exception(self, timeout=None): 

126 """Get the exception from the operation. 

127 

128 Args: 

129 timeout (int): How long to wait for the operation to complete. 

130 If None, wait indefinitely. 

131 

132 Returns: 

133 Optional[google.api_core.GoogleAPICallError]: The operation's 

134 error. 

135 """ 

136 await self._blocking_poll(timeout=timeout) 

137 return self._future.exception() 

138 

139 def add_done_callback(self, fn): 

140 """Add a callback to be executed when the operation is complete. 

141 

142 If the operation is completed, the callback will be scheduled onto the 

143 event loop. Otherwise, the callback will be stored and invoked when the 

144 future is done. 

145 

146 Args: 

147 fn (Callable[Future]): The callback to execute when the operation 

148 is complete. 

149 """ 

150 if self._background_task is None: 

151 self._background_task = asyncio.get_event_loop().create_task( 

152 self._blocking_poll() 

153 ) 

154 self._future.add_done_callback(fn) 

155 

156 def set_result(self, result): 

157 """Set the Future's result.""" 

158 self._future.set_result(result) 

159 

160 def set_exception(self, exception): 

161 """Set the Future's exception.""" 

162 self._future.set_exception(exception)