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

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

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

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

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

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

モデル定義

SQLAlchemy と Pydantic を併用すると定義するモデルクラス名が被ってしまうので、今回は SQLAlchemy のクラス名に Table という suffix をつけている。また、紛らわしいので SQLAlchemy の方では「モデル」という用語は使用しないようにして Pydantic の方を「モデル」と呼ぶことにする。

tables.py

import uuid

from sqlalchemy.orm import DeclarativeBase


class BaseTable(DeclarativeBase):
    id: Mapped[uuid.UUID] = ...  # 略
    # 略


class UserTable(BaseTable):
    __tablename__ = "users"
    # 略


class ArticleTable(BaseTable):
    __tablename__ = "articles"

    author_id: Mapped[uuid.UUID] = ...  # 略
    # 略


class CommentTable(BaseTable):
    __tablename__ = "comments"

    article_id: Mapped[uuid.UUID] = ...  # 略
    author_id: Mapped[uuid.UUID] = ...  # 略
    # 略

models.py

from pydantic import BaseModel


class User(BaseModel):
    # 略


class Article(BaseModel):
    # 略


class Comment(BaseModel):
    # 略
    

変換処理の実装

SQLAlchemy オブジェクトの column_attrs と relationships を取得して dict に変換していく。

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

使用例

selectinload または joinedload を使用して Eager Loading した結果を先ほど作った関数に渡して dict に変換する。

stmt = (
    select(ArticleTable)
    .where(ArticleTable.published_at.is_not(None))
    .order_by(ArticleTable.created_at.desc())
    .options(
        selectinload(ArticleTable.author),
        selectinload(ArticleTable.comments).selectinload(CommentTable.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 モデルを両方定義しないといけないじゃんダルい」という話はあるが、それについてはまた別の話で。