diff --git a/src/psql/psql.py b/src/psql/psql.py index 52b5dc5..cd84b89 100644 --- a/src/psql/psql.py +++ b/src/psql/psql.py @@ -44,6 +44,7 @@ errors (FAIL LOUD - unlike the mongo lib's swallow-and-default): import asyncio import logging +from contextlib import asynccontextmanager from typing import Any, Dict, List, Optional, Sequence import asyncpg @@ -236,7 +237,7 @@ class PsqlDB: commits on clean exit, rolls back and re-raises on any error. `conn` is a raw asyncpg connection ($1-placeholder execute/fetch/... directly). """ - return _Transaction(self.pool) + return _transaction(self.pool) # ------------------------------------------------------------------------- # layer 1 - friendly portable verbs (identical across psql/mysql) @@ -365,36 +366,17 @@ class PsqlDB: return _status_count(status) -class _Transaction: - """async ctx mgr: acquire a pooled connection + open a transaction on it""" +@asynccontextmanager +async def _transaction(pool: asyncpg.Pool): + """acquire a pooled connection and run the block in a transaction on it - def __init__(self, pool: asyncpg.Pool): - self._pool = pool - self._conn = None - self._tx = None - - async def __aenter__(self): - self._conn = await self._pool.acquire() - try: - self._tx = self._conn.transaction() - await self._tx.start() - except BaseException: - # start() (or transaction()) failing after acquire would otherwise leak the - # pooled connection - __aexit__ is not called when __aenter__ raises. release - # it and reset so a failed transaction start never burns a pool slot. - await self._pool.release(self._conn) - self._conn = None - raise - return self._conn - - async def __aexit__(self, exc_type, exc, tb) -> None: - try: - if exc_type is None: - await self._tx.commit() - else: - await self._tx.rollback() - finally: - await self._pool.release(self._conn) + asyncpg's own context managers do the work: the transaction commits on clean exit + and rolls back on any error, and the acquire releases the connection on every path - + including a failed transaction start, so a failed start never burns a pool slot + """ + async with pool.acquire() as conn: + async with conn.transaction(): + yield conn def _where(conditions: Optional[Dict[str, Any]]) -> tuple: