사이드 프로젝트를 FastAPI와 MongoDB를 조합해서z 진행하고 있는데, 오늘 DB 연결과 세션 관리를 하다가 @asynccontextmanager 어노테이션을 쓰게 됐다. 원래는 그냥 평소처럼 함수 하나 만들어서 처리하려고 했는데, 이걸 쓰면 좀 더 깔끔하게 정리될 것 같아서 시도해봤다. ( 그동안은 try ~ except ~ finally 형태를 주로 썼다. )
@asynccontextmanager로 정의한 함수를 어떻게 사용했는지 부터 살펴보자. ( get_db 함수 )
routes/users.py
from beanie import PydanticObjectId
from fastapi import APIRouter, Depends
from app.db.mongodb import get_db
from app.models.models import Users
...
@router.get("/{id}")
async def get_users(id: PydanticObjectId, db=Depends(get_db)): <- 이 부분
"""특정 유저 조회"""
user = await Users.get(id)
return user
routes/users.py 소스코드에서 특정 유저를 조회하는 API를 정의할 때, get_db() 함수를 자동으로 호출하여 get_db의 반환값이 자동으로 db 변수에 할당될 수 있도록 했다.
그 결과, 별도로 데이터베이스에 연결하는 코드 없이도 바로 유저 정보를 조회할 수 있다.
하지만 조회 후에는 session을 close해야 하는데, 이는 API 호출이 끝난 후 전체 컨텍스트가 종료될 때 자동으로 정리되어야 한다.
이 과정을 직접 코드로 구현하려면 꽤 번거롭겠지만, @asynccontextmanager가 그 역할을 대신해 준다.
@asynccontextmanager 로 구현된 get_db 함수 VS 직접 코드로 구현
@asynccontextmanager 로 구현된 get_db 함수
class MongoDB:
def __init__(self) -> None:
self._client = AsyncIOMotorClient(
f"mongodb://{MONGO_USER}:{MONGO_PASSWORD}@{MONGO_URI}?authSource={MONGO_USER}"
)
self._db = self._client[MONGO_NAME]
@asynccontextmanager
async def start_session(self):
"""MongoDB 연결 및 session 관리"""
session = await client.start_session()
await init_beanie(db, document_models=[Users, EasterEggs, Chats])
yield session
await session.end_session()
mongo_client = MongoDB()
async def get_db() -> AsyncGenerator:
async with mongo_client.start_session() as session:
db = mongo_client._db
yield db
직접 코드로 구현
from motor.motor_asyncio import AsyncIOMotorClient
from typing import AsyncGenerator
class MongoDB:
def __init__(self) -> None:
self._client = AsyncIOMotorClient(
f"mongodb://{MONGO_USER}:{MONGO_PASSWORD}@{MONGO_URI}?authSource={MONGO_USER}"
)
self._db = self._client[MONGO_NAME]
self._session = None
async def __aenter__(self):
"""세션을 시작하고 DB를 초기화"""
self._session = await self._client.start_session()
await init_beanie(self._db, document_models=[Users, EasterEggs, Chats])
return self._session # `async with`에서 반환되는 값
async def __aexit__(self, exc_type, exc_value, traceback):
"""세션 종료"""
if self._session:
await self._session.end_session()
print("end session")
if exc_type:
print(f"Error: {exc_value}")
async def get_db() -> AsyncGenerator:
async with MongoDB(mongo_client._client) as session:
db = mongo_client._db
yield db
육안으로 보더라도 @asynccontextmanager를 사용하면 비동기 컨텍스트를 훨씬 간결하고 가독성 있게 구현할 수 있다.
매번 try~except~finally 블록을 작성할 필요가 없어, 앞으로도 이 어노테이션을 자주 사용할 것 같다.
또한, 어노테이션을 통해 한 번에 컨텍스트 관리가 이루어지기 때문에 여러 개의 DB를 사용할 때도 세션 관리나 연결에 있어서 중복되는 코드 없이 손쉽게 활용할 수 있을 것 같다.
왜 진작에 이 어노테이션을 그동안 쓰지 않았는지 의문이다. ( 회사에서 조차... )
참고 링크
https://docs.python.org/3/library/contextlib.html#contextlib.asynccontextmanager