From: Federico Caselli <cfederico87@gmail.com>
Date: Fri, 17 Jun 2022 21:12:39 +0000 (+0200)
Subject: add shield() in aexit
X-Git-Tag: rel_1_4_40~24
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=6d0570de2b43463b82fbf6628610bf3c89befd59;p=thirdparty%2Fsqlalchemy%2Fsqlalchemy.git

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
(cherry picked from commit 1acaf0b2e4859a274e753b5054dcde3d5c7ca10e)
---

diff --git a/doc/build/changelog/unreleased_14/8145.rst b/doc/build/changelog/unreleased_14/8145.rst
new file mode 100644
index 0000000000..4cd6c12a58
--- /dev/null
+++ b/doc/build/changelog/unreleased_14/8145.rst
@@ -0,0 +1,14 @@
+.. change::
+    :tags: bug, asyncio
+    :tickets: 8145
+
+    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.
+
+
diff --git a/lib/sqlalchemy/ext/asyncio/engine.py b/lib/sqlalchemy/ext/asyncio/engine.py
index 63d148c975..4fbe4f7a59 100644
--- a/lib/sqlalchemy/ext/asyncio/engine.py
+++ b/lib/sqlalchemy/ext/asyncio/engine.py
@@ -4,6 +4,8 @@
 #
 # This module is part of SQLAlchemy and is released under
 # the MIT License: https://www.opensource.org/licenses/mit-license.php
+import asyncio
+
 from . import exc as async_exc
 from .base import ProxyComparable
 from .base import StartableContext
@@ -549,7 +551,7 @@ class AsyncConnection(ProxyComparable, StartableContext, AsyncConnectable):
         return self.start().__await__()
 
     async def __aexit__(self, type_, value, traceback):
-        await self.close()
+        await asyncio.shield(self.close())
 
 
 @util.create_proxy_methods(
@@ -600,8 +602,11 @@ class AsyncEngine(ProxyComparable, AsyncConnectable):
             return self.conn
 
         async def __aexit__(self, type_, value, traceback):
-            await self.transaction.__aexit__(type_, value, traceback)
-            await self.conn.close()
+            async def go():
+                await self.transaction.__aexit__(type_, value, traceback)
+                await self.conn.close()
+
+            await asyncio.shield(go())
 
     def __init__(self, sync_engine):
         if not sync_engine.dialect.is_async:
@@ -698,7 +703,7 @@ class AsyncEngine(ProxyComparable, AsyncConnectable):
 
         """
 
-        return await greenlet_spawn(self.sync_engine.dispose)
+        await greenlet_spawn(self.sync_engine.dispose)
 
 
 class AsyncTransaction(ProxyComparable, StartableContext):
diff --git a/lib/sqlalchemy/ext/asyncio/session.py b/lib/sqlalchemy/ext/asyncio/session.py
index ce6a0db090..378cbcbf2f 100644
--- a/lib/sqlalchemy/ext/asyncio/session.py
+++ b/lib/sqlalchemy/ext/asyncio/session.py
@@ -4,6 +4,9 @@
 #
 # This module is part of SQLAlchemy and is released under
 # the MIT License: https://www.opensource.org/licenses/mit-license.php
+
+import asyncio
+
 from . import engine
 from . import result as _result
 from .base import ReversibleProxy
@@ -607,7 +610,7 @@ class AsyncSession(ReversibleProxy):
             :meth:`_asyncio.AsyncSession.close`
 
         """
-        return await greenlet_spawn(self.sync_session.close)
+        await greenlet_spawn(self.sync_session.close)
 
     async def invalidate(self):
         """Close this Session, using connection invalidation.
@@ -625,7 +628,7 @@ class AsyncSession(ReversibleProxy):
         return self
 
     async def __aexit__(self, type_, value, traceback):
-        await self.close()
+        await asyncio.shield(self.close())
 
     def _maker_context_manager(self):
         # no @contextlib.asynccontextmanager until python3.7, gr
@@ -642,8 +645,11 @@ class _AsyncSessionContextManager:
         return self.async_session
 
     async def __aexit__(self, type_, value, traceback):
-        await self.trans.__aexit__(type_, value, traceback)
-        await self.async_session.__aexit__(type_, value, traceback)
+        async def go():
+            await self.trans.__aexit__(type_, value, traceback)
+            await self.async_session.__aexit__(type_, value, traceback)
+
+        await asyncio.shield(go())
 
 
 class AsyncSessionTransaction(ReversibleProxy, StartableContext):