From 3c476be23b88cca5bc5b19680456c4bad448d1a4 Mon Sep 17 00:00:00 2001 From: forsakenMystery Date: Mon, 26 Jan 2026 10:20:01 +0330 Subject: [PATCH] fix:Rollback SQLAlchemy session on asyncio cancellation --- .../helpers/decorators/sqlalchemy_atomic.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/archipy/helpers/decorators/sqlalchemy_atomic.py b/archipy/helpers/decorators/sqlalchemy_atomic.py index 9e836ea..7d55778 100644 --- a/archipy/helpers/decorators/sqlalchemy_atomic.py +++ b/archipy/helpers/decorators/sqlalchemy_atomic.py @@ -31,6 +31,7 @@ DatabaseTransactionError, InternalError, ) +import asyncio logger = logging.getLogger(__name__) @@ -243,6 +244,19 @@ async def async_wrapper(*args: Any, **kwargs: Any) -> R: async with session.begin(): result = await func(*args, **kwargs) return result + except asyncio.CancelledError: + # Real scenarios: request cancelled (client disconnect), timeouts (asyncio/AnyIO), + # TaskGroup/structured concurrency cancellation, or graceful shutdown. + # + # In Python 3.14, CancelledError inherits from BaseException (not Exception), + # so it bypasses `except Exception:` unless handled explicitly. + # We rollback for DB safety and re-raise to preserve cancellation semantics. + await asyncio.shield(session.rollback()) + raise + except (SystemExit, KeyboardInterrupt): + # Process is exiting (shutdown / Ctrl+C). Roll back to avoid leaving a transaction open, happens in dev mostly i think. + await asyncio.shield(session.rollback()) + raise except Exception as exception: await session.rollback() func_name = getattr(func, "__name__", "unknown") @@ -291,6 +305,10 @@ def sync_wrapper(*args: Any, **kwargs: Any) -> R: else: with session.begin(): return func(*args, **kwargs) + except (KeyboardInterrupt, SystemExit): + # Ctrl+C / shutdown: rollback for DB safety, then re-raise. + session.rollback() + raise except Exception as exception: session.rollback() func_name = getattr(func, "__name__", "unknown")