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 

38_ERROR_MESSAGE = ( 

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

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

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

42 "'pip install sqlalchemy[asyncio]'" 

43) 

44 

45 

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

47 raise ImportError(_ERROR_MESSAGE) 

48 

49 

50class _concurrency_shim_cls: 

51 """Late import shim for greenlet""" 

52 

53 __slots__ = ( 

54 "_has_greenlet", 

55 "greenlet", 

56 "_AsyncIoGreenlet", 

57 "getcurrent", 

58 ) 

59 

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

61 """Import greenlet and initialize the class""" 

62 if "greenlet" in globals(): 

63 return 

64 

65 if not TYPE_CHECKING: 

66 global getcurrent, greenlet, _AsyncIoGreenlet 

67 global _has_gr_context 

68 

69 try: 

70 from greenlet import getcurrent 

71 from greenlet import greenlet 

72 except ImportError as e: 

73 if not TYPE_CHECKING: 

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

75 greenlet = None 

76 self._has_greenlet = False 

77 self._initialize_no_greenlet() 

78 if raise_: 

79 raise ImportError(_ERROR_MESSAGE) from e 

80 else: 

81 self._has_greenlet = True 

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

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

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

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

86 

87 # implementation based on snaury gist at 

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

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

90 

91 class _AsyncIoGreenlet(greenlet): 

92 dead: bool 

93 

94 __sqlalchemy_greenlet_provider__ = True 

95 

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

97 greenlet.__init__(self, fn, driver) 

98 if _has_gr_context: 

99 self.gr_context = driver.gr_context 

100 

101 self.greenlet = greenlet 

102 self.getcurrent = getcurrent 

103 self._AsyncIoGreenlet = _AsyncIoGreenlet 

104 

105 def _initialize_no_greenlet(self): 

106 self.getcurrent = _not_implemented 

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

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

109 

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

111 if key in self.__slots__: 

112 self._initialize() 

113 return getattr(self, key) 

114 else: 

115 raise AttributeError(key) 

116 

117 

118_concurrency_shim = _concurrency_shim_cls() 

119 

120if TYPE_CHECKING: 

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

122 

123 def iscoroutine( 

124 awaitable: Awaitable[_T_co], 

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

126 

127else: 

128 iscoroutine = asyncio.iscoroutine 

129 

130 

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

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

133 

134 if iscoroutine(awaitable): 

135 awaitable.close() 

136 

137 

138def in_greenlet() -> bool: 

139 current = _concurrency_shim.getcurrent() 

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

141 

142 

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

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

145 

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

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

148 

149 :param awaitable: The coroutine to call. 

150 

151 """ 

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

153 current = _concurrency_shim.getcurrent() 

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

155 _safe_cancel_awaitable(awaitable) 

156 

157 raise exc.MissingGreenlet( 

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

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

160 ) 

161 

162 # returns the control to the driver greenlet passing it 

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

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

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

166 assert current.parent 

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

168 

169 

170await_only = await_ # old name. deprecated on 2.2 

171 

172 

173async def greenlet_spawn( 

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

175 *args: Any, 

176 _require_await: bool = False, 

177 **kwargs: Any, 

178) -> _T: 

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

180 

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

182 functions. 

183 

184 :param fn: The sync callable to call. 

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

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

187 """ 

188 

189 result: Any 

190 context = _concurrency_shim._AsyncIoGreenlet( 

191 fn, _concurrency_shim.getcurrent() 

192 ) 

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

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

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

196 # returned, and its result can be returned. 

197 switch_occurred = False 

198 

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

200 while not context.dead: 

201 switch_occurred = True 

202 try: 

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

204 # result back to it. 

205 value = await result 

206 except BaseException: 

207 # this allows an exception to be raised within 

208 # the moderated greenlet so that it can continue 

209 # its expected flow. 

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

211 else: 

212 result = context.switch(value) 

213 

214 if _require_await and not switch_occurred: 

215 raise exc.AwaitRequired( 

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

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

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

219 ) 

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

221 

222 

223class AsyncAdaptedLock: 

224 @memoized_property 

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

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

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

228 return asyncio.Lock() 

229 

230 def __enter__(self) -> bool: 

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

232 # coroutine has created the mutex. 

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

234 

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

236 self.mutex.release() 

237 

238 

239class _AsyncUtil: 

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

241 

242 def __init__(self) -> None: 

243 # runner it lazy so it can be created here 

244 self.runner = asyncio.Runner() 

245 

246 def run( 

247 self, 

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

249 *args: Any, 

250 **kwargs: Any, 

251 ) -> _T: 

252 """Run coroutine on the loop""" 

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

254 

255 def run_in_greenlet( 

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

257 ) -> _T: 

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

259 _concurrency_shim._initialize(raise_=False) 

260 

261 if _concurrency_shim._has_greenlet: 

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

263 # allow for a wrapped test function to call another 

264 assert in_greenlet() 

265 return fn(*args, **kwargs) 

266 else: 

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

268 else: 

269 return fn(*args, **kwargs) 

270 

271 def close(self) -> None: 

272 self.runner.close()