ネストした SQLAlchemy オブジェクトを dict に変換する

SQLAlchemy で遊んでいるときに作ったもののメモ。

SQLAlchemy (v2.x) を使用してデータ取得して、結果を Pydantic や TypedDict にマッピングしたい。このとき、以下の制約があるということにする。

• 関連テーブルの取得は selectinload などを使用し、lazy loading はしない
  • マッピングする前提なので
• SQLModel は使わない
  • それ使うと全然違う話になってくるので

最初 Pydantic の .model_validate(obj, from_attribute=True) を使おうとしたら lazy loading で無限ループが発生して死んだ。

いったん取得結果を dict に変換できればあとはどうとでもなりそうと思ったが、ネストされた属性まで含めて変換してくれる便利関数は SQLAlchemy からは提供されてなさそうだったので、書いた。

実装

from typing import Any, overload

from sqlalchemy import inspect
from sqlalchemy.orm import DeclarativeBase


@overload
def sa_to_dict(obj: DeclarativeBase) -> dict[str, Any]: ...


@overload
def sa_to_dict(obj: None) -> None: ...


def sa_to_dict(obj: DeclarativeBase | None) -> dict[str, Any] | None:
    if obj is None:
        return None

    info = inspect(obj)
    data: dict[str, Any] = {}

    for col in info.mapper.column_attrs:
        data[col.key] = getattr(obj, col.key)

    for rel in info.mapper.relationships:
        if rel.key in info.unloaded:
            continue
        value = getattr(obj, rel.key)
        if value is None:
            data[rel.key] = None
        elif rel.uselist:
            data[rel.key] = [sa_to_dict(x) for x in value]
        else:
            data[rel.key] = sa_to_dict(value)

    return data

使用例

stmt = (
    select(Article)
    .where(Article.published_at.is_not(None))
    .order_by(Article.created_at.desc())
    .options(
        selectinload(Article.author),
        selectinload(Article.comments).selectinload(Comment.author),
    )
)

with Session(engine) as session:
    article = sa_to_dict(session.scalars(stmt).first())

Pydantic 型にマッピングする

ここまでできていると、あとは Pydantic の .model_validate() だけでネストされたデータも含めてマッピングできる。

Article.model_validate(articles[0])

「SQLAlchemy モデルと Pydantic モデルを両方定義しないといけないじゃんダルい」という話はあるが、それについてはまた別の話で。