Async connections not returned to pool if task is cancelled

Describe the bug

The AsyncConnection context manager does not reliably close the connection when exiting due to a cancellation. This results in resources leaking, which is clearly surfaced by the warning during garbage collection. Garbage collection can’t perform the cleanup itself, but tells us that something was missed earlier.

To Reproduce

import anyio
from sqlalchemy.ext.asyncio import create_async_engine
from sqlalchemy.sql import text

async def main():
    engine = create_async_engine(
        "postgresql+asyncpg://username:password@localhost/db_name",
        echo=True,
        connect_args={"timeout": 5}
        )

    async def sleep_in_db():
        print(engine.pool.status()) # 0 connections
        async with engine.connect() as connection:
            print(engine.pool.status()) # 1 connections
            while True:
                await connection.execute(text("select now();"))
                await anyio.sleep(1)

    async with anyio.create_task_group() as db_tasks:
        db_tasks.start_soon(sleep_in_db)

        await anyio.sleep(2.5)

        db_tasks.cancel_scope.cancel() # Will raise a CancelledError under the hood


    print(engine.pool.status()) # 1 connection????

    await engine.dispose()

anyio.run(main)

Error

The garbage collector is trying to clean up connection <AdaptedConnection <asyncpg.connection.Connection object at 0x7f1cf4760660>>. This feature is unsupported on async dbapi, since no IO can be performed at this stage to reset the connection. Please close out all connections when they are no longer used, calling ``close()`` or using a context manager to manage their lifetime.
sys:1: SAWarning: The garbage collector is trying to clean up connection <AdaptedConnection <asyncpg.connection.Connection object at 0x7f1cf4760660>>. This feature is unsupported on async dbapi, since no IO can be performed at this stage to reset the connection. Please close out all connections when they are no longer used, calling ``close()`` or using a context manager to manage their lifetime.

Versions

• OS: Ubuntu on WSL2 in Windows 10
• Python: 3.10.4
• SQLAlchemy: 1.4.37
• Database: Postgres 14
• DBAPI (eg: psycopg, cx_oracle, mysqlclient): asyncpg 0.25.0

Additional context

I think the problem is that the way cancellation works is that the event loop (asyncio in this case) raises CancelledError when you do await, so if any await is needed while actually processing a cancellation (ex. to cleanup), that needs to be “shielded” from the cancellation we’re trying to process. We can see in AsyncConnection that on __aexit__() it’s trying to await self.close(), but it’s unable to (self.close() is never run) because of the cancellation: https://github.com/sqlalchemy/sqlalchemy/blob/main/lib/sqlalchemy/ext/asyncio/engine.py#L700

    async def __aexit__(self, type_: Any, value: Any, traceback: Any) -> None:
        await self.close()

Here’s a good description of this concept in the trio docs: https://trio.readthedocs.io/en/stable/reference-core.html?highlight=shield#cancellation-semantics (it takes a bit to get to the relevant part, but read the whole block).

I don’t know asyncio well enough to code up a robust solution but I believe these are the relevant docs: https://docs.python.org/3/library/asyncio-task.html#shielding-from-cancellation

Using anyio to achieve trio-like structured concurrency on top of asyncio, I am able to patch this __aexit__() method and then the connection is properly closed. I’m not sure if you’d consider using anyio in sqlalchemy. Could be used in just this one place, throughout the codebase in order to add compatibility with trio, or piecemeal to get the advantages of structured concurrency without a big refactor. Here’s the anyio docs for shielding: https://anyio.readthedocs.io/en/stable/cancellation.html?highlight=shield#shielding

import anyio

class PatchedAsyncConnection(AsyncConnection):
    async def __aexit__(self, type_, value, traceback):
        with anyio.open_cancel_scope(shield=True):
            await self.close()