Pydantic SQLAlchemy relationship conditional load

1. 개요

현재 서버 프레임워크로는 FastAPI를 사용하고있고 ORM은 SQLAlchemy를 사용하고 있다. FastAPI는 내부적으로 Pydantic이라는 라이브러리와 강결합되어있는 형태이며 해당 라이브러리를 디폴트로 Serialize등의 행위를 수행한다. 서버의 아키텍처는 Layered architecture를 사용하고 있었고, Repository에서의 반환값을 ORM모델이 아닌 Pydantic모델로 통일하자는 이야기가 있었다. (Entity와 비슷한 개념이라고 볼 수 있겠다) 

일반적으로 ORM 모델을 정의할 때 ORM의 장점을 살리기위해 relationship을 정의하고 해당 property를 통해 관계된 테이블의 요소에 접근한다. 여기서 문제가 하나 발생했는데, SQLAlchemy 모델 -> Pydantic 변환 시 lazy로 정의한 relationship도 모두 eager하게 가져오는 문제가 발생했다. 이러한 문제점을 해결하는 방법에 대해 좋은 코드가 있길래 간단하게 정리한다.

2. 해결 방법

다들 알다시피 Pydantic모델을 정의할 때 class Config내부에 orm_mode = True속성을 주게되면 from_orm() 메소드를 통해 손쉽게 모델 변환을 진행할 수 있다. Config 클래스 내부에는 사용자의 편의성과 커스터마이징을 돕기 위한 많은 요소들이 존재하는데, 그 중 getter_dict를 통해 우리의 문제를 해결할 수 있다.

먼저 아래와 같은 클래스를 하나 작성한다.

from typing import Any

from pydantic import BaseModel
from pydantic.utils import GetterDict
import sqlalchemy

class IgnoreLazyGetterDict(GetterDict):
    def __getitem__(self, key: str) -> Any:
        try:
            if self._is_lazy_loaded(key):
                return None

            return getattr(self._obj, key)
        except AttributeError as e:
            raise KeyError(key) from e

    def get(self, key: Any, default: Any = None) -> Any:
        # if a relationship field is not already loaded as a result of explicit join, ignore it,
        # so that pydantic schema doesn't trigger a lazy-load
        if self._is_lazy_loaded(key):
            return None

        return getattr(self._obj, key, default)

    def _is_lazy_loaded(self, key: Any) -> bool:
        return key in sqlalchemy.orm.attributes.instance_state(self._obj).unloaded

Pydantic은 내부적으로 GetterDict라는 클래스를 사용하여 값을 처리한다. 위 코드에서 _is_lazy_loaded() 함수를 보면 알겠지만, 요소에 접근할 때 해당 요소가 로드된 상태인지를 판단한다. (unloaded) 따라서 각 필드들을 변환할 때 만약 로드되지않은 상태라면 Pydantic모델로도 변환하지 않는다.

마지막으로 Pydantic모델의 Config클래스를 아래와 같이 구성한다.

class IgnoreLazyBaseModel(BaseModel):
    class Config:
        orm_mode = True
        getter_dict = IgnoreLazyGetterDict

위처럼 getter_dict에 우리가 작성한 Custom class를 넣어주면 내부적으로 변환 작업을 진행할 때 Pydantic에서 사용하고있는 GetterDict를 사용하는것이 아닌 해당 클래스를 상속받고 커스터마이징하여 새롭게 작성한 IgnoreLazyGetterDict를 통해 변환하게 되고 결과적으로 현재 로드된 필드만 변환을 진행한다.

본문 : https://github.com/tiangolo/fastapi/issues/2010