From dc9ab8a45caf5b459125f052800ce45984273c84 Mon Sep 17 00:00:00 2001 From: disqualifier Date: Fri, 16 Oct 2026 04:32:47 -0400 Subject: [PATCH] refactor: transaction() as an asynccontextmanager over aiomysql's acquire context T3-21: the _Transaction class hand-rolled acquire / start / commit-or-rollback / release plus a leak guard for a failed start. an asynccontextmanager over the driver's own pool.acquire() context does the same, and the acquire context releases on every path including a failed start. verified on a real backend: commit, rollback + re-raise, cancellation rollback, failed-start leak guard (pool max 1) and release after 50 transactions. Signed-off-by: disqualifier --- src/mysql/mysql.py | 42 +++++++++++++----------------------------- 1 file changed, 13 insertions(+), 29 deletions(-) diff --git a/src/mysql/mysql.py b/src/mysql/mysql.py index 20c359f..3732000 100644 --- a/src/mysql/mysql.py +++ b/src/mysql/mysql.py @@ -26,6 +26,7 @@ for anything not wrapped, use `.pool`. import asyncio import logging +from contextlib import asynccontextmanager from typing import Any, Dict, List, Optional, Sequence import aiomysql @@ -271,7 +272,7 @@ class MysqlDB: commits on clean exit, rolls back and re-raises on any error. `conn` is a raw aiomysql connection (autocommit disabled for the block). """ - return _Transaction(self.pool) + return _transaction(self.pool) # ------------------------------------------------------------------------- # layer 1 - friendly portable verbs (identical to the psql lib) @@ -400,39 +401,22 @@ class MysqlDB: return rowcount -class _Transaction: - """async ctx mgr: acquire a pooled connection + run a transaction on it +@asynccontextmanager +async def _transaction(pool: aiomysql.Pool): + """acquire a pooled connection and run the block in a transaction on it - autocommit is disabled for the block so the statements commit together on clean - exit (or roll back on error). + begin() disables autocommit for the block; commit on clean exit, rollback + re-raise + on any error. the acquire context releases the connection on every path - including + a failed begin(), so a failed start never burns a pool slot """ - - def __init__(self, pool: aiomysql.Pool): - self._pool = pool - self._conn = None - - async def __aenter__(self): - self._conn = await self._pool.acquire() + async with pool.acquire() as conn: + await conn.begin() try: - await self._conn.begin() + yield conn except BaseException: - # __aexit__ never runs if __aenter__ raises, so begin() failing after acquire - # would otherwise leak the pooled connection - release it here instead. - await self._pool.release(self._conn) - self._conn = None + await conn.rollback() raise - return self._conn - - async def __aexit__(self, exc_type, exc, tb) -> None: - try: - if exc_type is None: - await self._conn.commit() - else: - await self._conn.rollback() - finally: - # release() returns a task (schedules a _wakeup); await it to avoid leaving - # it un-awaited each transaction, matching psql's and aiomysql's own idiom. - await self._pool.release(self._conn) + await conn.commit() def _where(conditions: Optional[Dict[str, Any]]) -> tuple: