Свой контекстный менеджер для транзакций в Python: рабочие приёмы

Типичная сцена: функция открывает транзакцию, что-то пишет в базу, а между execute и commit вылетает исключение. Если rollback забыли или он сам оказался внутри try, соединение остаётся с незакрытой транзакцией, блокировки висят, а следующий запрос ловит таймаут. Контекстный менеджер with решает это структурно: выход из блока гарантированно вызовет фиксацию или откат, что бы ни случилось внутри.

Разберём, как написать такой менеджер под чистый DB-API (psycopg, sqlite3), чем отличаются два способа его реализации и какие подводные камни ждут во вложенных транзакциях и async-коде. Речь про Python 3.10+ и стандартный протокол DB-API 2.0 (PEP 249).

Что даёт with и чего от него ждут

Объект-контекстный менеджер обязан реализовать два метода: __enter__ (вызывается на входе в блок, его результат попадает в переменную после as) и __exit__(exc_type, exc_val, exc_tb) (вызывается на выходе — и при нормальном завершении, и при исключении).

Ключевой момент — семантика возврата из __exit__:

• если __exit__ вернул истинное значение, исключение считается обработанным и подавляется — наружу не пробрасывается;
• если вернул ложное (в том числе None) — исключение пробрасывается дальше после выполнения тела __exit__.

Для транзакций нам почти всегда нужно откатить и пробросить исключение дальше, а не проглотить его. Значит, __exit__ должен возвращать ложь. Тихо подавленная ошибка при откате — это баг, из-за которого вызывающий код думает, что запись прошла.

Контекстный менеджер транзакции решает ровно одну задачу — гарантировать, что на выходе из блока соединение окажется либо зафиксированным, либо откатанным. Всё остальное (пул, ретраи, логирование) — отдельные слои, не смешивайте их в одном объекте.

Способ 1: класс с __enter__ и __exit__

Прямая реализация протокола. Подходит, когда нужен объект с состоянием — например, держать курсор и отдавать его в блок.

# Python 3.10+, psycopg 3.x
import psycopg


class Transaction:
    def __init__(self, conn: psycopg.Connection):
        self.conn = conn
        self.cursor: psycopg.Cursor | None = None

    def __enter__(self) -> psycopg.Cursor:
        # Открываем курсор и отдаём его в блок через as
        self.cursor = self.conn.cursor()
        return self.cursor

    def __exit__(self, exc_type, exc_val, exc_tb) -> bool:
        if self.cursor is not None:
            self.cursor.close()
        if exc_type is None:
            # Тело блока отработало без ошибок — фиксируем
            self.conn.commit()
        else:
            # Было исключение — откатываем
            self.conn.rollback()
        # Возвращаем False: исключение пробрасывается дальше
        return False


# Использование
conn = psycopg.connect("YOUR_DSN")
with Transaction(conn) as cur:
    cur.execute("INSERT INTO orders (sku, qty) VALUES (%s, %s)", ("A1", 3))
    cur.execute("UPDATE stock SET qty = qty - %s WHERE sku = %s", (3, "A1"))
# Здесь уже либо commit, либо rollback

Проверка, что работает: оберните два execute, где второй заведомо падает (например, нарушает ограничение), и убедитесь, что первый INSERT не остался в базе.

Способ 2: генератор с @contextmanager

Когда состояние не нужно и объект — лишний, короче написать функцию-генератор и декоратор contextmanager из contextlib. Всё до yield — это __enter__, всё после — __exit__.

# Python 3.10+, sqlite3 из стандартной библиотеки
import sqlite3
from contextlib import contextmanager


@contextmanager
def transaction(conn: sqlite3.Connection):
    try:
        yield conn
    except Exception:
        conn.rollback()
        raise  # обязателен: без него исключение проглотится
    else:
        conn.commit()


# Использование
conn = sqlite3.connect("shop.db")
with transaction(conn) as c:
    c.execute("INSERT INTO orders (sku, qty) VALUES (?, ?)", ("A1", 3))

Здесь raise без аргументов заменяет собой «вернуть False» из класса. Забудете его — и ошибка исчезнет молча. Это самая частая ошибка в этом варианте.

Что выбрать

Правило простое: нужен объект с методами и состоянием — класс. Нужно обернуть один линейный сценарий «сделай — зафиксируй — при ошибке откати» — генератор.

Где это ломается

Ошибка при самом commit или rollback

Сеть отвалилась ровно в момент commit — и он бросает исключение уже из __exit__. Если внутри блока тоже была ошибка, вы получите исключение из __exit__, а исходное потеряется как контекст. В Python исходная ошибка сохранится в цепочке (__context__), но новая станет главной. Не глотайте ошибки rollback в try/except: pass — вы скроете реальную причину.

Вложенные with на одном соединении

Два вложенных with transaction(conn) на одном соединении — ловушка. У большинства драйверов одно соединение держит одну транзакцию. Внутренний блок сделает commit, и внешний окажется вне транзакции, хотя думает, что всё ещё в ней. Для настоящей вложенности нужны SAVEPOINT¹, а не наивная обёртка. Если вложенность нужна — считайте её отдельной задачей и реализуйте через savepoint явно.

Автокоммит и уже открытая транзакция

Поведение зависит от драйвера. Например, у psycopg 3 подключение по умолчанию не в autocommit, и транзакция стартует с первого запроса. Если вы вручную включили autocommit=True, ваши commit/rollback станут бессмысленными. Проверьте режим соединения до того, как оборачивать его менеджером.

async-код

Обычный with не умеет ждать корутины. Для asyncpg или psycopg в async-режиме нужен асинхронный контекстный менеджер: методы __aenter__/__aexit__ либо декоратор asynccontextmanager, и вызов через async with. Синхронный менеджер в async-коде либо не скомпилируется, либо заблокирует цикл событий.

Что сделать дальше

Готовый менеджер стоит проверить под нагрузкой: под тем нагрузочным сценарием, где часть транзакций падает, посмотрите в мониторинге СУБД число активных сессий в состоянии idle in transaction (для PostgreSQL — через pg_stat_activity). Если оно растёт — где-то транзакция не закрывается, и ваша обёртка не срабатывает на всех путях. Дальше — вынести получение соединения из пула в тот же with, чтобы соединение возвращалось в пул автоматически, и отдельно разобрать ретраи при deadlock: это уже логика поверх транзакции, а не часть менеджера.

¹ SAVEPOINT — точка внутри транзакции, к которой можно откатиться, не отменяя всю транзакцию целиком. Позволяет имитировать вложенные транзакции в СУБД, которые их не поддерживают напрямую.