FastAPIで型付けをサボるとORMのlazy loadでハマる

14 min読了の目安(約8600字TECH技術記事

概要

FastAPIでデータベースを利用するアプリケーションを書いたときに、APIの返り値の中の要素が足りないことがありました。例えば、User has many Itemsの関係があったときに、GET /usersで返り値のjsonにItemsが含まれていない場合です。
この記事ではその原因と解決策を調査します。

環境

$ pip freeze
click==7.1.2
fastapi==0.61.1
h11==0.9.0
httptools==0.1.1
pydantic==1.6.1
SQLAlchemy==1.3.19
starlette==0.13.6
uvicorn==0.11.8
uvloop==0.14.0
websockets==8.1

コード

今回のサンプル用のコードはgithubにも置きます。

データベースはSQLiteを利用。ORMはSQLAlchemy

from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker

SQLALCHEMY_DATABASE_URL = "sqlite:///./sample.db"

engine = create_engine(
    SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False}
)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

Base = declarative_base()

モデルはUserとItemのみ。
Userがitemsを持っているので、User.itemsの形でItemのリストを取得することができます。

from sqlalchemy import Boolean, Column, ForeignKey, Integer, String
from sqlalchemy.orm import relationship

from database import Base


class User(Base):
    __tablename__ = "users"

    id = Column(Integer, primary_key=True, index=True)
    name = Column(String, unique=True, index=True)

    items = relationship("Item")


class Item(Base):
    __tablename__ = "items"

    id = Column(Integer, primary_key=True, index=True)
    title = Column(String, index=True)
    owner_id = Column(Integer, ForeignKey("users.id"))

CRUD

from sqlalchemy.orm import Session

import models

def get_user(db: Session):
    return db.query(models.User).first()

main.pyで初期データを入れてサーバを起動します。
API endpointは一つで GET /one にアクセスするとUserが返ります。

import uvicorn
from fastapi import FastAPI, Depends
from typing import List
from sqlalchemy.orm import Session
import crud, models
from database import SessionLocal, engine

models.Base.metadata.create_all(bind=engine)

app = FastAPI()

def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


@app.get('/one')
def read_root(db: Session = Depends(get_db)):
    user = crud.get_user(db)
    return user

if __name__ == '__main__':
    db = SessionLocal()
    if db.query(models.User).count() == 0:
        user = models.User(name='test_user')
        item = models.Item(title='test_item')
        user.items = [item]
        db.add(user)
        db.commit()
    db.close()
    uvicorn.run(app, host="0.0.0.0", port=8000)

再現

上記のコードを動かしてサーバを起動します。

 python app/main.py 
INFO:     Started server process [18366]
INFO:     Waiting for application startup.
INFO:     Application startup complete.
INFO:     Uvicorn running on http://0.0.0.0:8000 (Press CTRL+C to quit)

別のターミナルを開いてAPI callしてみます。

$ curl 'http://localhost:8000/one'
{"name":"test_user","id":1}

Userモデルが帰ってきましたが、items要素がありません。
これを今回の問題とします。

原因

先のコードではORMとしてSQLAlchemyを利用しています。
SQLAlchemyではデータベースアクセス時に全ての関係を取得するのではなく、必要なときに必要なレコードを取得するlazy loadという読み方をデフォルトとしています。
itemsが読み込まれていないのはこれが原因の可能性が高いので、事実を確認していきます。

発行されるSQLの確認

まずは発行されるSQLを確認して、itemsがDBから読み出されているかを確認します。
lazy loadが原因だとするとそもそもDBから読み出されていませんし、逆にDBからitemsが読み出されている場合は別の所に原因があることになります。

SQLAlchemyのログを標準出力に出すには次の設定を書きます。

import logging
 
logging.basicConfig(level=logging.INFO)
logging.getLogger('sqlalchemy.engine').setLevel(logging.INFO)

リクエストを投げたところ次のログが標準出力に書き出されました。

INFO:sqlalchemy.engine.base.Engine:BEGIN (implicit)
INFO:sqlalchemy.engine.base.Engine:SELECT users.id AS users_id, users.name AS users_name 
FROM users
 LIMIT ? OFFSET ?
INFO:sqlalchemy.engine.base.Engine:(1, 0)
INFO:     127.0.0.1:64938 - "GET /one HTTP/1.1" 200 OK
INFO:sqlalchemy.engine.base.Engine:ROLLBACK

SELECT users.id AS users_id, users.name AS users_name FROM users LIMIT ? OFFSET ? でUsersテーブルからレコードを取得していますが、Itemsテーブルを触っている気配はありません。怠惰ですね。

Lazy Loadさせてみる

Lazy Loadが原因だとすると、itemsにアクセスしたらSQLAlchemyがItemsテーブルからレコードを引っ張ってくるはずです。その様子もSQLクエリで見てみます。

GET /one endpointの中でuser.itemsを評価するように書き換えます。SQLクエリ発行のタイミングを見るため、評価の前後でログ出力をしています。

@app.get('/one')
def read_root(db: Session = Depends(get_db)):
    user = crud.get_user(db)
    logging.info('before lazy load')
    user.items # ここでSQLクエリが発行されてitemsに値が入るはず
    logging.info('after lazy load')
    return user

この状態でAPI callしたところ意図通りitemsを含んだUserレコードが返ってきました!

$ curl 'http://localhost:8000/one'
{"name":"test_user","id":1,"items":[{"title":"test_item","id":1,"owner_id":1}]}

ログの方でもbefore lazy loadとafter lazy loadの間でItemsを取りに行っていることが確認できます。

INFO:sqlalchemy.engine.base.Engine:BEGIN (implicit)
INFO:sqlalchemy.engine.base.Engine:SELECT users.id AS users_id, users.name AS users_name 
FROM users
 LIMIT ? OFFSET ?
INFO:sqlalchemy.engine.base.Engine:(1, 0)
INFO:root:before lazy load
INFO:sqlalchemy.engine.base.Engine:SELECT items.id AS items_id, items.title AS items_title, items.owner_id AS items_owner_id 
FROM items 
WHERE ? = items.owner_id
INFO:sqlalchemy.engine.base.Engine:(1,)
INFO:root:after lazy load
INFO:     127.0.0.1:50448 - "GET /one HTTP/1.1" 200 OK
INFO:sqlalchemy.engine.base.Engine:ROLLBACK

解決策

原因がSQLAlchemyのlazy loadによるものと確認できました。
それでは関係(items)を含む結果を返したいときはどうしたら良いかというと

Lazy Loadさせる

先程確認したように、items属性を評価することでitemsを読み込むことができます。
しかし、コードに意味のない行が入り後々バグの原因となるためこの方法を取るべきではありません。

SQLAlchemyでLazy Loadしない設定にする

SQLAlchemyはデフォルトでlazy loadする設定になっていますが、モデルのrelationshipにlazy="joined"などを渡すことでしない設定にもできます。

class User(Base):
    __tablename__ = "users"

    id = Column(Integer, primary_key=True, index=True)
    name = Column(String, unique=True, index=True)

    items = relationship("Item", lazy="joined")

このとき、ログから次のことがわかります。

  • 1回のSQLでUsersテーブルとItemsテーブルをjoinして返している
  • before lazy loadのタイミング以前にitemsが読み込まれている
INFO:sqlalchemy.engine.base.Engine:BEGIN (implicit)
INFO:sqlalchemy.engine.base.Engine:SELECT anon_1.users_id AS anon_1_users_id, anon_1.users_name AS anon_1_users_name, items_1.id AS items_1_id, items_1.title AS items_1_title, items_1.owner_id AS items_1_owner_id 
FROM (SELECT users.id AS users_id, users.name AS users_name 
FROM users
 LIMIT ? OFFSET ?) AS anon_1 LEFT OUTER JOIN items AS items_1 ON anon_1.users_id = items_1.owner_id
INFO:sqlalchemy.engine.base.Engine:(1, 0)
INFO:root:before lazy load
INFO:root:after lazy load
INFO:     127.0.0.1:53589 - "GET /one HTTP/1.1" 200 OK
INFO:sqlalchemy.engine.base.Engine:ROLLBACK

ちなみにSQLAlchemyの公式ドキュメントよりオプションは次の種類があります

  • Lazy Loadする
    • lazy='select'
    • lazy='raise'
  • Lazy Loadしない(一度に読み込む)
    • lazy='joined'
    • lazy='subquery'
    • lazy='selectin'
  • Lazy Loadしない(そもそもitemsを読み込まない)
    • lazy='noload'

FastAPIの型Validationを利用する

FastAPIでは型によるValidationの機能を利用してこの問題を解決しています。
SQLAlchemyのモデルを再度lazy loadにして、GET /one endpointを次のように書き換えます。response_model=schemas.Userがポイント。
確認のためloggingを入れています。

import schemas
@app.get('/one', response_model=schemas.User)
def read_root(db: Session = Depends(get_db)):
    user = crud.get_user(db)
    logging.info('before lazy load')
    return user

schemasは何かというと、型定義です。

from typing import List

from pydantic import BaseModel

class Item(BaseModel):
    id: int
    title: str

    class Config:
        orm_mode = True

class User(BaseModel):
    id: int
    name: str
    items: List[Item] = []

    class Config:
        orm_mode = True

FastAPIではresponse_model=schemas.Userをつけることによって返り値の型Validationを実行してくれます。その際、(pydanticが)itemsにアクセスしようとするので明示的にitemsにアクセスしなくてもlazy loadが発生することになります。

これを実行するとitemsが入った状態

$ curl 'http://localhost:8000/one'
{"id":1,"name":"test_user","items":[{"id":1,"title":"test_item"}]}

before lazy loadの後にitemsがloadされているのがわかります。

INFO:sqlalchemy.engine.base.Engine:BEGIN (implicit)
INFO:sqlalchemy.engine.base.Engine:SELECT users.id AS users_id, users.name AS users_name 
FROM users
 LIMIT ? OFFSET ?
INFO:sqlalchemy.engine.base.Engine:(1, 0)
INFO:root:before lazy load
INFO:sqlalchemy.engine.base.Engine:SELECT items.id AS items_id, items.title AS items_title, items.owner_id AS items_owner_id 
FROM items 
WHERE ? = items.owner_id
INFO:sqlalchemy.engine.base.Engine:(1,)
INFO:     127.0.0.1:57512 - "GET /one HTTP/1.1" 200 OK
INFO:sqlalchemy.engine.base.Engine:ROLLBACK

型Validationによる恩恵もあるので、FastAPIを使う場合はこの方法を取るのが良いでしょう。

まとめ

FastAPIを利用するときは型付けをサボらずにやりましょう。
ちなみにこの話は公式ドキュメントにちゃんと書いてあります。FastAPIはドキュメントが丁寧。