Python @contextmanager로 트랜잭션 롤백 누락 방지

서버 코드에서 트랜잭션 버그는 대부분 “로직이 틀려서”가 아니라 “예외 경로에서 정리 코드가 실행되지 않아서” 발생합니다. 특히 Python에서는 함수가 길어지거나 return 이 여러 군데에 생기면, commit() 과 rollback() 호출이 쉽게 흩어지고, 일부 경로에서 빠집니다.

이 글은 contextlib.contextmanager 를 이용해 트랜잭션 경계를 언어 레벨의 구조로 고정해 롤백 누락을 원천 차단하는 패턴을 설명합니다. DB-API(예: psycopg2, pymysql)부터 SQLAlchemy까지 동일한 원칙으로 적용할 수 있습니다.

또한 장애가 나면 “원인을 찾기 전에 증상이 번진다”는 점에서, 운영 환경에서의 방어적 코딩이 중요합니다. 비슷한 맥락의 운영 진단 글로 Redis 핫키로 QPS 폭주? LFU로 5분 진단, 리눅스 디스크 100% - 삭제해도 용량 안 줄 때도 함께 참고하면 좋습니다.

왜 롤백 누락이 생길까

전형적인 안티패턴은 다음과 같습니다.

성공 시에는 commit() 을 호출하지만
실패 시에는 rollback() 이 호출되지 않거나
rollback() 은 호출되지만 예외가 삼켜져 상위 레이어가 실패를 모르는 상태가 되거나
커넥션/커서 close가 누락되어 풀 고갈로 이어짐

예를 들어 아래 코드는 얼핏 정상처럼 보이지만, 중간에 예외가 나면 롤백이 누락될 수 있습니다.

def transfer(conn, from_id, to_id, amount):
    cur = conn.cursor()
    cur.execute("UPDATE accounts SET balance = balance - %s WHERE id = %s", (amount, from_id))

    # 여기서 예외가 나면? (네트워크, 데드락, 제약조건 등)
    cur.execute("UPDATE accounts SET balance = balance + %s WHERE id = %s", (amount, to_id))

    conn.commit()
    cur.close()

이 문제를 try/except/finally 로 고치려 해도, 코드가 길어질수록 “정리 로직의 중복”과 “예외 처리의 일관성”이 무너집니다. 그래서 트랜잭션은 문법 구조로 강제하는 편이 안정적입니다.

목표: 트랜잭션 경계를 한 곳에 고정하기

원하는 형태는 아래처럼 “업무 로직은 with 블록 안”에만 존재하고, 커밋/롤백/close는 공통 유틸리티가 책임지는 것입니다.

with transaction(conn):
    # 업무 로직
    ...

이렇게 만들면 다음이 자동으로 보장됩니다.

블록이 정상 종료되면 commit()
블록에서 예외가 발생하면 rollback() 후 예외 재전파
커서/커넥션 정리(필요 시)
로깅/메트릭/트레이싱을 한 곳에서 수행

@contextmanager로 트랜잭션 래퍼 만들기 (DB-API 공통)

Python 표준 라이브러리 contextlib.contextmanager 는 제너레이터 기반으로 컨텍스트 매니저를 쉽게 구현하게 해줍니다.

핵심은 yield 이전은 진입 로직, yield 이후는 정상 종료 로직이며, 예외는 except 블록에서 잡을 수 있다는 점입니다.

from contextlib import contextmanager

@contextmanager
def transaction(conn):
    """DB-API connection 기반 트랜잭션 컨텍스트.

    - 정상 종료: commit
    - 예외 발생: rollback 후 예외 재전파
    """
    try:
        yield conn
    except Exception:
        # 롤백 실패 자체도 중요한 장애 신호이므로, 여기서 예외를 삼키지 않는다.
        conn.rollback()
        raise
    else:
        conn.commit()

사용 예시는 아래와 같습니다.

def transfer(conn, from_id, to_id, amount):
    with transaction(conn):
        cur = conn.cursor()
        try:
            cur.execute(
                "UPDATE accounts SET balance = balance - %s WHERE id = %s",
                (amount, from_id),
            )
            cur.execute(
                "UPDATE accounts SET balance = balance + %s WHERE id = %s",
                (amount, to_id),
            )
        finally:
            cur.close()

여기서 중요한 포인트는 두 가지입니다.

트랜잭션 경계(커밋/롤백)는 transaction() 하나로 통일
커서 close 같은 리소스 정리는 finally 로 별도 처리(혹은 커서도 컨텍스트로 감싸기)

커서까지 함께 안전하게 만들기

커서 close도 누락되기 쉬우니, 커서 역시 컨텍스트 매니저로 감싸면 더 단단해집니다.

@contextmanager
def cursor(conn):
    cur = conn.cursor()
    try:
        yield cur
    finally:
        cur.close()

이제 호출부는 다음처럼 읽기 쉬워집니다.

def transfer(conn, from_id, to_id, amount):
    with transaction(conn):
        with cursor(conn) as cur:
            cur.execute(
                "UPDATE accounts SET balance = balance - %s WHERE id = %s",
                (amount, from_id),
            )
            cur.execute(
                "UPDATE accounts SET balance = balance + %s WHERE id = %s",
                (amount, to_id),
            )

업무 로직에서 “정리 코드”가 사라지고, 예외가 나도 롤백과 close가 구조적으로 실행됩니다.

실전에서 자주 놓치는 디테일

1) 예외를 잡고 삼키지 말 것

롤백 후 raise 를 하지 않으면 상위 레이어는 성공으로 오해할 수 있습니다. 예외를 로깅하고 싶다면 로깅 후 재전파가 기본입니다.

import logging
logger = logging.getLogger(__name__)

@contextmanager
def transaction(conn):
    try:
        yield conn
    except Exception as e:
        logger.exception("transaction failed, rollback")
        conn.rollback()
        raise
    else:
        conn.commit()

2) 롤백 자체가 실패할 수 있음

네트워크 단절이나 커넥션 종료 등으로 rollback() 이 예외를 던질 수 있습니다. 이 경우 “원래 예외”와 “롤백 예외”가 섞이는데, 운영에서는 둘 다 중요합니다.

가장 단순한 전략은 롤백 예외도 그대로 올리는 것입니다. 더 고급으로는 원래 예외를 보존하면서 롤백 실패를 로깅만 하는 방식이 있습니다.

@contextmanager
def transaction(conn, *, rollback_log_only=True):
    try:
        yield conn
    except Exception:
        try:
            conn.rollback()
        except Exception:
            if rollback_log_only:
                import logging
                logging.getLogger(__name__).exception("rollback failed")
            else:
                raise
        raise
    else:
        conn.commit()

3) 자동 커밋(autocommit) 설정 확인

드라이버별로 autocommit 기본값이 다릅니다. autocommit이 켜져 있으면 “트랜잭션처럼 보이지만 실제로는 statement 단위로 커밋”됩니다.

psycopg2 는 기본적으로 autocommit이 꺼져 있는 편이지만, 코드에서 켜는 경우가 많습니다.
pymysql 등은 설정에 따라 달라질 수 있습니다.

따라서 transaction() 를 쓰는 코드에서는 커넥션 생성 시점에 autocommit 정책을 명확히 해두는 것을 권장합니다.

# 예시: psycopg2
# conn.autocommit = False

4) 중첩 호출: 이미 트랜잭션 안에서 또 transaction을 열면?

서비스 레이어가 커지면 아래처럼 중첩이 생깁니다.

service_a() 가 with transaction(conn) 을 열고
내부에서 service_b() 도 with transaction(conn) 을 열어버림

이 경우 “이중 커밋/롤백”이 되거나, 드라이버에 따라 예상치 못한 동작을 할 수 있습니다.

대응 전략은 크게 3가지입니다.

규칙으로 금지하고 코드 리뷰로 잡기
스레드 로컬이나 contextvars 로 “현재 트랜잭션 깊이”를 추적해 최상위에서만 commit/rollback
DB savepoint를 사용해 중첩을 지원

아래는 contextvars 로 깊이를 추적해 최상위에서만 커밋/롤백하는 간단한 예시입니다.

from contextlib import contextmanager
from contextvars import ContextVar

_tx_depth: ContextVar[int] = ContextVar("tx_depth", default=0)

@contextmanager
def transaction(conn):
    depth = _tx_depth.get()
    token = _tx_depth.set(depth + 1)

    is_outermost = depth == 0
    try:
        yield conn
    except Exception:
        if is_outermost:
            conn.rollback()
        raise
    else:
        if is_outermost:
            conn.commit()
    finally:
        _tx_depth.reset(token)

이 패턴은 “한 요청 처리 흐름”에서만 유효하며, 멀티스레드/비동기 환경에서도 contextvars 는 비교적 안전하게 동작합니다.

SQLAlchemy에서는 어떻게 적용할까

SQLAlchemy는 이미 세션/트랜잭션 컨텍스트를 제공합니다. 그럼에도 @contextmanager 를 쓰는 이유는 다음과 같습니다.

프로젝트 표준 로깅/메트릭/에러 변환을 한 곳에 모으기
세션 생성/반납을 강제하기
서비스 레이어의 사용성을 단순화하기

SQLAlchemy 2.x 스타일 예시

from contextlib import contextmanager
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

engine = create_engine("postgresql+psycopg2://user:pass@host/db")
Session = sessionmaker(bind=engine, autoflush=False, autocommit=False)

@contextmanager
def session_scope():
    session = Session()
    try:
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()

사용:

def create_order(order_repo, payload):
    with session_scope() as session:
        order_repo.create(session, payload)
        order_repo.add_audit_log(session, payload)

SQLAlchemy 자체에도 Session.begin() 컨텍스트가 있지만, 위 래퍼는 “프로젝트 표준”을 강제하는 역할을 합니다.

테스트로 롤백 누락을 재현하고 막기

트랜잭션 래퍼는 테스트로 안전성을 확인하는 게 좋습니다. 예를 들어 “두 번째 쿼리에서 예외가 나면 첫 번째 업데이트가 롤백되었는지”를 검증합니다.

아래는 개념 예시(의사 코드에 가깝습니다)입니다.

import pytest

class Boom(Exception):
    pass

def test_transaction_rolls_back_on_error(conn):
    # 초기값 세팅
    with transaction(conn):
        with cursor(conn) as cur:
            cur.execute("UPDATE accounts SET balance = 100 WHERE id = 1")

    # 실패 유도
    with pytest.raises(Boom):
        with transaction(conn):
            with cursor(conn) as cur:
                cur.execute("UPDATE accounts SET balance = 0 WHERE id = 1")
                raise Boom("fail after update")

    # 롤백 확인
    with cursor(conn) as cur:
        cur.execute("SELECT balance FROM accounts WHERE id = 1")
        (balance,) = cur.fetchone()

    assert balance == 100

이 테스트가 통과하면 “예외 경로에서도 롤백이 실행된다”는 계약을 지속적으로 보장할 수 있습니다.

운영 관점 체크리스트

트랜잭션 롤백 누락은 데이터 정합성 문제로 끝나지 않고, 락 경합과 커넥션 고갈로 이어져 장애를 키울 수 있습니다. 운영에서 함께 점검할 포인트입니다.

트랜잭션이 너무 오래 열려 있지 않은가(긴 I/O, 외부 API 호출을 트랜잭션 안에서 수행)
예외 로깅에 쿼리 컨텍스트가 포함되는가(요청 ID, 사용자 ID, 핵심 파라미터)
커넥션 풀이 고갈될 때의 증상과 알람이 있는가

성능/장애가 “증폭”되는 구조는 다른 영역에서도 자주 보입니다. 예를 들어 트래픽이 몰릴 때 캐시 키 하나가 병목이 되면 전체 QPS가 튀는 문제는 Redis 핫키로 QPS 폭주? LFU로 5분 진단에서 다룬 방식으로 빠르게 확인할 수 있습니다.

결론

@contextmanager 로 트랜잭션을 감싸면 커밋/롤백/리소스 정리가 “개발자 습관”이 아니라 “코드 구조”로 강제됩니다. 특히 다음 효과가 큽니다.

예외 경로에서 롤백 누락 방지
커밋/롤백 정책의 중앙집중화
중복 try/finally 제거로 가독성 향상
로깅/메트릭/트레이싱을 한 곳에서 일관되게 처리

트랜잭션은 애플리케이션의 데이터 안전벨트입니다. 안전벨트를 매번 손으로 매는 대신, with transaction(...) 이라는 구조를 도입해 “안 매면 출발이 안 되게” 만드는 것이 가장 확실한 예방책입니다.