Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/sqlalchemy/util/concurrency.py: 37%

1# util/concurrency.py 

2# Copyright (C) 2005-2026 the SQLAlchemy authors and contributors 

3# <see AUTHORS file> 

4# 

5# This module is part of SQLAlchemy and is released under 

6# the MIT License: https://www.opensource.org/licenses/mit-license.php 

7# mypy: allow-untyped-defs, allow-untyped-calls 

8 

9"""asyncio-related concurrency functions.""" 

10 

11from __future__ import annotations 

12 

13import asyncio 

14import sys 

15from typing import Any 

16from typing import Awaitable 

17from typing import Callable 

18from typing import Coroutine 

19from typing import NoReturn 

20from typing import TYPE_CHECKING 

21from typing import TypeGuard 

22from typing import TypeVar 

23 

24from .langhelpers import memoized_property 

25from .. import exc 

26 

27_T = TypeVar("_T") 

28 

29 

30def is_exit_exception(e: BaseException) -> bool: 

31 # note asyncio.CancelledError is already BaseException 

32 # so was an exit exception in any case 

33 return not isinstance(e, Exception) or isinstance( 

34 e, (asyncio.TimeoutError, asyncio.CancelledError) 

35 ) 

36 

37 

38def is_cancellation_exception(e: BaseException) -> bool: 

39 return isinstance(e, asyncio.CancelledError) 

40 

41 

42_ERROR_MESSAGE = ( 

43 "The SQLAlchemy asyncio module requires that the Python 'greenlet' " 

44 "library is installed. In order to ensure this dependency is " 

45 "available, use the 'sqlalchemy[asyncio]' install target: " 

46 "'pip install sqlalchemy[asyncio]'" 

47) 

48 

49 

50def _not_implemented(*arg: Any, **kw: Any) -> NoReturn: 

51 raise ImportError(_ERROR_MESSAGE) 

52 

53 

54class _concurrency_shim_cls: 

55 """Late import shim for greenlet""" 

56 

57 __slots__ = ( 

58 "_has_greenlet", 

59 "greenlet", 

60 "_AsyncIoGreenlet", 

61 "getcurrent", 

62 ) 

63 

64 def _initialize(self, *, raise_: bool = True) -> None: 

65 """Import greenlet and initialize the class""" 

66 if "greenlet" in globals(): 

67 return 

68 

69 if not TYPE_CHECKING: 

70 global getcurrent, greenlet, _AsyncIoGreenlet 

71 global _has_gr_context 

72 

73 try: 

74 from greenlet import getcurrent 

75 from greenlet import greenlet 

76 except ImportError as e: 

77 if not TYPE_CHECKING: 

78 # set greenlet in the global scope to prevent re-init 

79 greenlet = None 

80 self._has_greenlet = False 

81 self._initialize_no_greenlet() 

82 if raise_: 

83 raise ImportError(_ERROR_MESSAGE) from e 

84 else: 

85 self._has_greenlet = True 

86 # If greenlet.gr_context is present in current version of greenlet, 

87 # it will be set with the current context on creation. 

88 # Refs: https://github.com/python-greenlet/greenlet/pull/198 

89 _has_gr_context = hasattr(getcurrent(), "gr_context") 

90 

91 # implementation based on snaury gist at 

92 # https://gist.github.com/snaury/202bf4f22c41ca34e56297bae5f33fef 

93 # Issue for context: https://github.com/python-greenlet/greenlet/issues/173 # noqa: E501 

94 

95 class _AsyncIoGreenlet(greenlet): 

96 dead: bool 

97 

98 __sqlalchemy_greenlet_provider__ = True 

99 

100 def __init__(self, fn: Callable[..., Any], driver: greenlet): 

101 greenlet.__init__(self, fn, driver) 

102 if _has_gr_context: 

103 self.gr_context = driver.gr_context 

104 

105 self.greenlet = greenlet 

106 self.getcurrent = getcurrent 

107 self._AsyncIoGreenlet = _AsyncIoGreenlet 

108 

109 def _initialize_no_greenlet(self): 

110 self.getcurrent = _not_implemented 

111 self.greenlet = _not_implemented # type: ignore[assignment] 

112 self._AsyncIoGreenlet = _not_implemented # type: ignore[assignment] 

113 

114 def __getattr__(self, key: str) -> Any: 

115 if key in self.__slots__: 

116 self._initialize() 

117 return getattr(self, key) 

118 else: 

119 raise AttributeError(key) 

120 

121 

122_concurrency_shim = _concurrency_shim_cls() 

123 

124if TYPE_CHECKING: 

125 _T_co = TypeVar("_T_co", covariant=True) 

126 

127 def iscoroutine( 

128 awaitable: Awaitable[_T_co], 

129 ) -> TypeGuard[Coroutine[Any, Any, _T_co]]: ... 

130 

131else: 

132 iscoroutine = asyncio.iscoroutine 

133 

134 

135def _safe_cancel_awaitable(awaitable: Awaitable[Any]) -> None: 

136 # https://docs.python.org/3/reference/datamodel.html#coroutine.close 

137 

138 if iscoroutine(awaitable): 

139 awaitable.close() 

140 

141 

142def in_greenlet() -> bool: 

143 current = _concurrency_shim.getcurrent() 

144 return getattr(current, "__sqlalchemy_greenlet_provider__", False) 

145 

146 

147def await_(awaitable: Awaitable[_T]) -> _T: 

148 """Awaits an async function in a sync method. 

149 

150 The sync method must be inside a :func:`greenlet_spawn` context. 

151 :func:`await_` calls cannot be nested. 

152 

153 :param awaitable: The coroutine to call. 

154 

155 """ 

156 # this is called in the context greenlet while running fn 

157 current = _concurrency_shim.getcurrent() 

158 if not getattr(current, "__sqlalchemy_greenlet_provider__", False): 

159 _safe_cancel_awaitable(awaitable) 

160 

161 raise exc.MissingGreenlet( 

162 "greenlet_spawn has not been called; can't call await_() " 

163 "here. Was IO attempted in an unexpected place?" 

164 ) 

165 

166 # returns the control to the driver greenlet passing it 

167 # a coroutine to run. Once the awaitable is done, the driver greenlet 

168 # switches back to this greenlet with the result of awaitable that is 

169 # then returned to the caller (or raised as error) 

170 assert current.parent 

171 return current.parent.switch(awaitable) # type: ignore[no-any-return] 

172 

173 

174await_only = await_ # old name. deprecated on 2.2 

175 

176 

177async def greenlet_spawn( 

178 fn: Callable[..., _T], 

179 *args: Any, 

180 _require_await: bool = False, 

181 **kwargs: Any, 

182) -> _T: 

183 """Runs a sync function ``fn`` in a new greenlet. 

184 

185 The sync function can then use :func:`await_` to wait for async 

186 functions. 

187 

188 :param fn: The sync callable to call. 

189 :param \\*args: Positional arguments to pass to the ``fn`` callable. 

190 :param \\*\\*kwargs: Keyword arguments to pass to the ``fn`` callable. 

191 """ 

192 

193 result: Any 

194 context = _concurrency_shim._AsyncIoGreenlet( 

195 fn, _concurrency_shim.getcurrent() 

196 ) 

197 # runs the function synchronously in gl greenlet. If the execution 

198 # is interrupted by await_, context is not dead and result is a 

199 # coroutine to wait. If the context is dead the function has 

200 # returned, and its result can be returned. 

201 switch_occurred = False 

202 

203 result = context.switch(*args, **kwargs) 

204 while not context.dead: 

205 switch_occurred = True 

206 try: 

207 # wait for a coroutine from await_ and then return its 

208 # result back to it. 

209 value = await result 

210 except BaseException: 

211 # this allows an exception to be raised within 

212 # the moderated greenlet so that it can continue 

213 # its expected flow. 

214 result = context.throw(*sys.exc_info()) 

215 else: 

216 result = context.switch(value) 

217 

218 if _require_await and not switch_occurred: 

219 raise exc.AwaitRequired( 

220 "The current operation required an async execution but none was " 

221 "detected. This will usually happen when using a non compatible " 

222 "DBAPI driver. Please ensure that an async DBAPI is used." 

223 ) 

224 return result # type: ignore[no-any-return] 

225 

226 

227class AsyncAdaptedLock: 

228 @memoized_property 

229 def mutex(self) -> asyncio.Lock: 

230 # there should not be a race here for coroutines creating the 

231 # new lock as we are not using await, so therefore no concurrency 

232 return asyncio.Lock() 

233 

234 def __enter__(self) -> bool: 

235 # await is used to acquire the lock only after the first calling 

236 # coroutine has created the mutex. 

237 return await_(self.mutex.acquire()) 

238 

239 def __exit__(self, *arg: Any, **kw: Any) -> None: 

240 self.mutex.release() 

241 

242 

243class _AsyncUtil: 

244 """Asyncio util for test suite/ util only""" 

245 

246 def __init__(self) -> None: 

247 # runner it lazy so it can be created here 

248 self.runner = asyncio.Runner() 

249 

250 def run( 

251 self, 

252 fn: Callable[..., Coroutine[Any, Any, _T]], 

253 *args: Any, 

254 **kwargs: Any, 

255 ) -> _T: 

256 """Run coroutine on the loop""" 

257 return self.runner.run(fn(*args, **kwargs)) 

258 

259 def run_in_greenlet( 

260 self, fn: Callable[..., _T], *args: Any, **kwargs: Any 

261 ) -> _T: 

262 """Run sync function in greenlet. Support nested calls""" 

263 _concurrency_shim._initialize(raise_=False) 

264 

265 if _concurrency_shim._has_greenlet: 

266 if self.runner.get_loop().is_running(): 

267 # allow for a wrapped test function to call another 

268 assert in_greenlet() 

269 return fn(*args, **kwargs) 

270 else: 

271 return self.runner.run(greenlet_spawn(fn, *args, **kwargs)) 

272 else: 

273 return fn(*args, **kwargs) 

274 

275 def close(self) -> None: 

276 self.runner.close()