try/except 지옥을 그만 쓰세요: SQLAlchemy와 Unit Of Work로 깨끗한 데이터베이스 트랜잭션

Source: Dev.to

async def create_order(user_id: int, items_in_basket: list[dict]) -> Order:
    session = get_session()
    try:
        user: User = await session.get(User, user_id)
        if not user:
            await session.rollback()
            raise ValueError("User not found")

        order: Order = Order(user_id=user_id)
        session.add(order)

        for item in items_in_basket:
            line: OrderLineItem = OrderLineItem(**item, order=order)
            session.add(line)

        await session.commit()
        return order
    except DuplicateError:
        await session.rollback()
        raise
    except Exception:
        await session.rollback()
        raise
    finally:
        await session.close()

위 패턴은 rollback 호출이 여기저기 흩어지고 finally 블록이 필요하며 비즈니스 로직과 보일러플레이트가 뒤섞입니다. 더 깔끔한 접근법은 Unit of Work(UoW)를 사용해 세션 수명 주기를 자동으로 관리하는 것입니다.

Unit of Work란?

Martin Fowler는 이를 다음과 같이 정의했습니다:

“Unit of Work는 데이터베이스에 영향을 줄 수 있는 비즈니스 트랜잭션 동안 수행한 모든 작업을 추적합니다. 작업이 끝나면, 작업 결과에 따라 데이터베이스를 변경하기 위해 해야 할 모든 일을 파악합니다.”

온라인 쇼핑 카트를 떠올려 보세요: 아이템을 추가·제거·변경하지만 구매 버튼을 누를 때까지는 아무 것도 영구 저장되지 않습니다. 카트를 버리면 변경 사항이 저장되지 않죠. 데이터베이스 작업도 같은 원리입니다:

• 메모리에서 변경 사항을 추적
• 한 번에 모두 커밋하거나
• 문제가 발생하면 모두 롤백

구현

from typing import Self
from types import TracebackType
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker

class UnitOfWork:
    """Async context manager for database transactions.

    Commits on success, rolls back on exception, always cleans up.
    """

    def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None:
        self._session_factory = session_factory

    async def __aenter__(self) -> Self:
        self.session = self._session_factory()
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        if exc_type is not None:
            await self.rollback()
        await self.session.close()

    async def commit(self) -> None:
        await self.session.commit()

    async def rollback(self) -> None:
        await self.session.rollback()

작동 방식

• **__aenter__**는 각 async with 블록마다 새로운 AsyncSession을 생성해, 트랜잭션마다 독립된 세션을 제공합니다.
• **__aexit__**는 블록을 빠져나올 때 예외가 있으면 자동으로 롤백하고, 이후 세션을 닫습니다.
• **명시적인 commit()**을 사용하면 언제 변경 사항을 영구 저장할지 직접 제어할 수 있습니다.

Explicit is better than implicit – commit()을 수동으로 호출하면 (예: 조기 반환이나 읽기 전용 쿼리를 수행할 때) 의도치 않은 쓰기를 방지할 수 있습니다.

사용 방법

애플리케이션 시작 시 한 번 세션 팩토리 설정

from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncEngine

engine: AsyncEngine = create_async_engine(
    "postgresql+asyncpg://user:pass@localhost/db"
)

SessionFactory = async_sessionmaker(engine, expire_on_commit=False)

Unit of Work 생성 후 작업 수행

async def main() -> None:
    uow = UnitOfWork(session_factory=SessionFactory)

    async with uow:
        uow.session.add(User(name="Alice", email="alice@example.com"))
        await uow.commit()
        print(f"Created user with ID: {user.id}")

시나리오

정상 흐름 – 커밋 성공

async with uow:
    uow.session.add(User(name="Bob"))
    await uow.commit()   # changes are saved

커밋 전 예외 발생 – 자동 롤백

async with uow:
    uow.session.add(User(name="Charlie"))
    raise ValueError("Something went wrong")
    # No commit called; __aexit__ rolls back automatically

서비스 함수에서 Unit of Work 사용

async def create_order(user_id: int, items_in_basket: list[dict], uow: UnitOfWork) -> Order:
    async with uow:
        user = await uow.session.get(User, user_id)
        if not user:
            raise ValueError("User not found")

        order = Order(user_id=user_id)
        uow.session.add(order)

        for item in items_in_basket:
            uow.session.add(OrderLine(**item, order=order))

        await uow.commit()
        return order

이 패턴은 흩어져 있던 rollback 호출을 없애고 finally 블록의 필요성을 제거하며, 비즈니스 로직과 영속성 로직을 깔끔하게 분리합니다.