始めに

Pythonのコンテキストマネージャは、with文と組み合わせてリソース（ファイル、DB接続、ロックなど）の確保と解放を自動化する仕組みです。enterメソッドでセットアップ、exitメソッドでクリーンアップを行うことで、例外発生時も安全にリソース管理できます。

知っておけば便利なことが多いので、ブログに便利な箇所をまとめます。

環境

Python
- 3.13.3

実装

基本的には明示的にリソースの開放が必要な時に使用します。ただ、DB自体のリソース開放はライブラリ側が自動的に行ってくれますので、インスタンス生成時にコンテキストマネジャーを使用しましょう。

async def get_reader_db():
    async with async_session() as session:
        yield session
        
SessionReaderDep = Annotated[AsyncSession, Depends(get_reader_db)]

@router.get("/sample")
async def sample(
    db: SessionReaderDep,
) -> Any:
    try:
        // 適当な処理
        await sample2(db);
        await session.commit()
    except HTTPException:
        await session.rollback()
        raise
    except Exception as e:
        await session.rollback()
        raise e

ただし、通常時のコミット処理や、エラーが発生したときのロールバックやログの出力等々はアプリケーション側の責任となります。どの箇所でも同じ記載をすることになりますので、それもまたコンテキストマネジャーで共通化しましょう。

from contextlib import asynccontextmanager

from fastapi import HTTPException
from sqlalchemy.ext.asyncio import AsyncSession

# コンテキストマネジャーの定義
@asynccontextmanager
async def session_context(session: AsyncSession) -> None:
    try:
        yield session
        await session.commit()
    except HTTPException:
        await session.rollback()
        raise
    except Exception as e:
        await session.rollback()
        raise e
 
 
@router.get("/sample")
async def sample(
    param_db: SessionReaderDep,
) -> Any:
    async with session_context(param_db) as db:
        // 適当な処理
        await sample2(db);

さらに重要なのが__exit__または__aexit__処理でFalseかNoneを返却することです。もし定義しない場合、エラー発生行がコンテキストマネジャーで発生したことになり、何行目でエラーが発生したかがわからなくなります。ライブラリとしては意図的にエラーをまとめていると思いますが、アプリケーションの場合にはエラー行を特定させない行為にメリットは薄いので、必ず追加で定義したほうが良いです。

https://docs.python.org/ja/3.8/reference/compound_stmts.html#id1

例外は、別の例外を送出するような finally 節が無い場合にのみ呼び出しスタックへ伝わります。新しい例外によって、古い例外は失われます。 `

@asynccontextmanager
async def session_context(session: AsyncSession) -> None:
    try:
        yield session
        await session.commit()
    except HTTPException:
        await session.rollback()
        raise
    except Exception as e:
        await session.rollback()
        raise e
        
    async def __aexit__(self, exc_type, exc_val, exc_tb):  # type: ignore
        """
        呼出元にエラーが発生したときに伝播させたいので、Falseを常に返却する。
        発生しないと、全てのエラーがcontext_managerの行で発生したことになる。
        """
        return False

@router.get("/sample")
async def sample(
    param_db: SessionReaderDep,
) -> Any:
    async with session_context(param_db) as db: // この行でエラーになったことになる
        // 適当な処理
        await sample2(db);
        
async def sample2(db):
    raise ValueError("") // 本当はこの行でエラーになったのに。

また、コンテキストマネジャーはデコレータとしても使用できます。例えば、バッチ処理で処理開始前と処理開始後にログを出力して、処理時間を計りたい等々があったとします。それも簡単に作成できます。

@asynccontextmanager
async def timeit(operation_name: str = "Operation"):
    """
    処理時間を計測し、開始と終了時にログを出力するコンテキストマネージャー

    Args:
        operation_name: 操作名（ログに表示される）
    """
    start_time = time.time()
    logger.info(f"[{operation_name}] 処理開始")

    try:
        yield
    finally:
        end_time = time.time()
        elapsed_time = end_time - start_time
        logger.info(f"[{operation_name}] 処理完了 - 実行時間: {elapsed_time:.4f}秒")
        
async def test_decolator():
    @timeit("TEST DECOLATOR")
    async def test_func():
        await asyncio.sleep(0.1)

    # WHEN
    await test_func()
    # ログとして、次のログが出力される
    # [TEST DECOLATOR] 処理開始
    # [TEST DECOLATOR] 処理完了 - 実行時間: 0.1022秒