From 1acaf0b2e4859a274e753b5054dcde3d5c7ca10e Mon Sep 17 00:00:00 2001 From: Federico Caselli Date: Fri, 17 Jun 2022 23:12:39 +0200 Subject: add shield() in aexit Added ``asyncio.shield()`` to the connection and session release process specifically within the ``__aexit__()`` context manager exit, when using :class:`.AsyncConnection` or :class:`.AsyncSession` as a context manager that releases the object when the context manager is complete. This appears to help with task cancellation when using alternate concurrency libraries such as ``anyio``, ``uvloop`` that otherwise don't provide an async context for the connection pool to release the connection properly during task cancellation. Fixes: #8145 Change-Id: I0b1ea9c3a22a18619341cbb8591225fcd339042c --- lib/sqlalchemy/ext/asyncio/session.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) (limited to 'lib/sqlalchemy/ext/asyncio/session.py') diff --git a/lib/sqlalchemy/ext/asyncio/session.py b/lib/sqlalchemy/ext/asyncio/session.py index be3414cef..14e08b5c5 100644 --- a/lib/sqlalchemy/ext/asyncio/session.py +++ b/lib/sqlalchemy/ext/asyncio/session.py @@ -6,6 +6,7 @@ # the MIT License: https://www.opensource.org/licenses/mit-license.php from __future__ import annotations +import asyncio from typing import Any from typing import Dict from typing import Generic @@ -837,7 +838,7 @@ class AsyncSession(ReversibleProxy[Session]): :meth:`_asyncio.AsyncSession.close` """ - return await greenlet_spawn(self.sync_session.close) + await greenlet_spawn(self.sync_session.close) async def invalidate(self) -> None: """Close this Session, using connection invalidation. @@ -855,7 +856,7 @@ class AsyncSession(ReversibleProxy[Session]): return self async def __aexit__(self, type_: Any, value: Any, traceback: Any) -> None: - await self.close() + await asyncio.shield(self.close()) def _maker_context_manager(self: _AS) -> _AsyncSessionContextManager[_AS]: return _AsyncSessionContextManager(self) @@ -1516,8 +1517,11 @@ class _AsyncSessionContextManager(Generic[_AS]): return self.async_session async def __aexit__(self, type_: Any, value: Any, traceback: Any) -> None: - await self.trans.__aexit__(type_, value, traceback) - await self.async_session.__aexit__(type_, value, traceback) + async def go() -> None: + await self.trans.__aexit__(type_, value, traceback) + await self.async_session.__aexit__(type_, value, traceback) + + await asyncio.shield(go()) class AsyncSessionTransaction( -- cgit v1.2.1