273 lines
8.6 KiB
Python
273 lines
8.6 KiB
Python
# util/concurrency.py
|
|
# Copyright (C) 2005-2026 the SQLAlchemy authors and contributors
|
|
# <see AUTHORS file>
|
|
#
|
|
# This module is part of SQLAlchemy and is released under
|
|
# the MIT License: https://www.opensource.org/licenses/mit-license.php
|
|
# mypy: allow-untyped-defs, allow-untyped-calls
|
|
|
|
"""asyncio-related concurrency functions."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import sys
|
|
from typing import Any
|
|
from typing import Awaitable
|
|
from typing import Callable
|
|
from typing import Coroutine
|
|
from typing import NoReturn
|
|
from typing import TYPE_CHECKING
|
|
from typing import TypeGuard
|
|
from typing import TypeVar
|
|
|
|
from .langhelpers import memoized_property
|
|
from .. import exc
|
|
|
|
_T = TypeVar("_T")
|
|
|
|
|
|
def is_exit_exception(e: BaseException) -> bool:
|
|
# note asyncio.CancelledError is already BaseException
|
|
# so was an exit exception in any case
|
|
return not isinstance(e, Exception) or isinstance(
|
|
e, (asyncio.TimeoutError, asyncio.CancelledError)
|
|
)
|
|
|
|
|
|
_ERROR_MESSAGE = (
|
|
"The SQLAlchemy asyncio module requires that the Python 'greenlet' "
|
|
"library is installed. In order to ensure this dependency is "
|
|
"available, use the 'sqlalchemy[asyncio]' install target: "
|
|
"'pip install sqlalchemy[asyncio]'"
|
|
)
|
|
|
|
|
|
def _not_implemented(*arg: Any, **kw: Any) -> NoReturn:
|
|
raise ImportError(_ERROR_MESSAGE)
|
|
|
|
|
|
class _concurrency_shim_cls:
|
|
"""Late import shim for greenlet"""
|
|
|
|
__slots__ = (
|
|
"_has_greenlet",
|
|
"greenlet",
|
|
"_AsyncIoGreenlet",
|
|
"getcurrent",
|
|
)
|
|
|
|
def _initialize(self, *, raise_: bool = True) -> None:
|
|
"""Import greenlet and initialize the class"""
|
|
if "greenlet" in globals():
|
|
return
|
|
|
|
if not TYPE_CHECKING:
|
|
global getcurrent, greenlet, _AsyncIoGreenlet
|
|
global _has_gr_context
|
|
|
|
try:
|
|
from greenlet import getcurrent
|
|
from greenlet import greenlet
|
|
except ImportError as e:
|
|
if not TYPE_CHECKING:
|
|
# set greenlet in the global scope to prevent re-init
|
|
greenlet = None
|
|
self._has_greenlet = False
|
|
self._initialize_no_greenlet()
|
|
if raise_:
|
|
raise ImportError(_ERROR_MESSAGE) from e
|
|
else:
|
|
self._has_greenlet = True
|
|
# If greenlet.gr_context is present in current version of greenlet,
|
|
# it will be set with the current context on creation.
|
|
# Refs: https://github.com/python-greenlet/greenlet/pull/198
|
|
_has_gr_context = hasattr(getcurrent(), "gr_context")
|
|
|
|
# implementation based on snaury gist at
|
|
# https://gist.github.com/snaury/202bf4f22c41ca34e56297bae5f33fef
|
|
# Issue for context: https://github.com/python-greenlet/greenlet/issues/173 # noqa: E501
|
|
|
|
class _AsyncIoGreenlet(greenlet):
|
|
dead: bool
|
|
|
|
__sqlalchemy_greenlet_provider__ = True
|
|
|
|
def __init__(self, fn: Callable[..., Any], driver: greenlet):
|
|
greenlet.__init__(self, fn, driver)
|
|
if _has_gr_context:
|
|
self.gr_context = driver.gr_context
|
|
|
|
self.greenlet = greenlet
|
|
self.getcurrent = getcurrent
|
|
self._AsyncIoGreenlet = _AsyncIoGreenlet
|
|
|
|
def _initialize_no_greenlet(self):
|
|
self.getcurrent = _not_implemented
|
|
self.greenlet = _not_implemented # type: ignore[assignment]
|
|
self._AsyncIoGreenlet = _not_implemented # type: ignore[assignment]
|
|
|
|
def __getattr__(self, key: str) -> Any:
|
|
if key in self.__slots__:
|
|
self._initialize()
|
|
return getattr(self, key)
|
|
else:
|
|
raise AttributeError(key)
|
|
|
|
|
|
_concurrency_shim = _concurrency_shim_cls()
|
|
|
|
if TYPE_CHECKING:
|
|
_T_co = TypeVar("_T_co", covariant=True)
|
|
|
|
def iscoroutine(
|
|
awaitable: Awaitable[_T_co],
|
|
) -> TypeGuard[Coroutine[Any, Any, _T_co]]: ...
|
|
|
|
else:
|
|
iscoroutine = asyncio.iscoroutine
|
|
|
|
|
|
def _safe_cancel_awaitable(awaitable: Awaitable[Any]) -> None:
|
|
# https://docs.python.org/3/reference/datamodel.html#coroutine.close
|
|
|
|
if iscoroutine(awaitable):
|
|
awaitable.close()
|
|
|
|
|
|
def in_greenlet() -> bool:
|
|
current = _concurrency_shim.getcurrent()
|
|
return getattr(current, "__sqlalchemy_greenlet_provider__", False)
|
|
|
|
|
|
def await_(awaitable: Awaitable[_T]) -> _T:
|
|
"""Awaits an async function in a sync method.
|
|
|
|
The sync method must be inside a :func:`greenlet_spawn` context.
|
|
:func:`await_` calls cannot be nested.
|
|
|
|
:param awaitable: The coroutine to call.
|
|
|
|
"""
|
|
# this is called in the context greenlet while running fn
|
|
current = _concurrency_shim.getcurrent()
|
|
if not getattr(current, "__sqlalchemy_greenlet_provider__", False):
|
|
_safe_cancel_awaitable(awaitable)
|
|
|
|
raise exc.MissingGreenlet(
|
|
"greenlet_spawn has not been called; can't call await_() "
|
|
"here. Was IO attempted in an unexpected place?"
|
|
)
|
|
|
|
# returns the control to the driver greenlet passing it
|
|
# a coroutine to run. Once the awaitable is done, the driver greenlet
|
|
# switches back to this greenlet with the result of awaitable that is
|
|
# then returned to the caller (or raised as error)
|
|
assert current.parent
|
|
return current.parent.switch(awaitable) # type: ignore[no-any-return]
|
|
|
|
|
|
await_only = await_ # old name. deprecated on 2.2
|
|
|
|
|
|
async def greenlet_spawn(
|
|
fn: Callable[..., _T],
|
|
*args: Any,
|
|
_require_await: bool = False,
|
|
**kwargs: Any,
|
|
) -> _T:
|
|
"""Runs a sync function ``fn`` in a new greenlet.
|
|
|
|
The sync function can then use :func:`await_` to wait for async
|
|
functions.
|
|
|
|
:param fn: The sync callable to call.
|
|
:param \\*args: Positional arguments to pass to the ``fn`` callable.
|
|
:param \\*\\*kwargs: Keyword arguments to pass to the ``fn`` callable.
|
|
"""
|
|
|
|
result: Any
|
|
context = _concurrency_shim._AsyncIoGreenlet(
|
|
fn, _concurrency_shim.getcurrent()
|
|
)
|
|
# runs the function synchronously in gl greenlet. If the execution
|
|
# is interrupted by await_, context is not dead and result is a
|
|
# coroutine to wait. If the context is dead the function has
|
|
# returned, and its result can be returned.
|
|
switch_occurred = False
|
|
|
|
result = context.switch(*args, **kwargs)
|
|
while not context.dead:
|
|
switch_occurred = True
|
|
try:
|
|
# wait for a coroutine from await_ and then return its
|
|
# result back to it.
|
|
value = await result
|
|
except BaseException:
|
|
# this allows an exception to be raised within
|
|
# the moderated greenlet so that it can continue
|
|
# its expected flow.
|
|
result = context.throw(*sys.exc_info())
|
|
else:
|
|
result = context.switch(value)
|
|
|
|
if _require_await and not switch_occurred:
|
|
raise exc.AwaitRequired(
|
|
"The current operation required an async execution but none was "
|
|
"detected. This will usually happen when using a non compatible "
|
|
"DBAPI driver. Please ensure that an async DBAPI is used."
|
|
)
|
|
return result # type: ignore[no-any-return]
|
|
|
|
|
|
class AsyncAdaptedLock:
|
|
@memoized_property
|
|
def mutex(self) -> asyncio.Lock:
|
|
# there should not be a race here for coroutines creating the
|
|
# new lock as we are not using await, so therefore no concurrency
|
|
return asyncio.Lock()
|
|
|
|
def __enter__(self) -> bool:
|
|
# await is used to acquire the lock only after the first calling
|
|
# coroutine has created the mutex.
|
|
return await_(self.mutex.acquire())
|
|
|
|
def __exit__(self, *arg: Any, **kw: Any) -> None:
|
|
self.mutex.release()
|
|
|
|
|
|
class _AsyncUtil:
|
|
"""Asyncio util for test suite/ util only"""
|
|
|
|
def __init__(self) -> None:
|
|
# runner it lazy so it can be created here
|
|
self.runner = asyncio.Runner()
|
|
|
|
def run(
|
|
self,
|
|
fn: Callable[..., Coroutine[Any, Any, _T]],
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> _T:
|
|
"""Run coroutine on the loop"""
|
|
return self.runner.run(fn(*args, **kwargs))
|
|
|
|
def run_in_greenlet(
|
|
self, fn: Callable[..., _T], *args: Any, **kwargs: Any
|
|
) -> _T:
|
|
"""Run sync function in greenlet. Support nested calls"""
|
|
_concurrency_shim._initialize(raise_=False)
|
|
|
|
if _concurrency_shim._has_greenlet:
|
|
if self.runner.get_loop().is_running():
|
|
# allow for a wrapped test function to call another
|
|
assert in_greenlet()
|
|
return fn(*args, **kwargs)
|
|
else:
|
|
return self.runner.run(greenlet_spawn(fn, *args, **kwargs))
|
|
else:
|
|
return fn(*args, **kwargs)
|
|
|
|
def close(self) -> None:
|
|
self.runner.close()
|