SQLAlchemy(automap) + GraphQLでN+1問題を解決する

何も考えずにgrapheneを使いSQLAlchemy(automap)でGraphQLを使うとN+1問題に遭遇すると思います。
回避しようとGraphQLのSchemaでresolveでいちいち書いたりするのは非常に生産性が悪く、また対応漏れを防ぐのが難しいと思います。
それを定義せずに回避する組み合わせがあったので書いてみます。

欲しい物

• 既に存在するRDBMSから自動でModelを生成できる
• Grapheneのスキーマにバインドできる
• relationshipで関係するリレーションを取得するときにN+1問題を起こさない

環境

• Python 3.7.6

requirements

• Flask==1.1.1
• Flask-GraphQL==2.0.1
• SQLAlchemy==1.3.11
• graphene==2.1.8
• SQLAlchemy-bulk-lazy-loader==0.9.9

解決法

SQLAlchemy-bulk-lazy-loader を使いましょう。

1. ライブラリの読み込み

automapをする前に読み込んでおきます。
記述すると、SQLAlchemyのModelに定義するrelationshipに lazy='bulk' パラメータを指定することができます。

# 1:N問題回避
from sqlalchemy_bulk_lazy_loader import BulkLazyLoader
BulkLazyLoader.register_loader()

2. automapで生成されるrelationshipへのキーワードを変更

そのままautomapしてもN+1問題は解決されません。
automapでrelationshipに渡すパラメータを変更できるよう、automap_baseが返すオブジェクトのprepareメソッドに渡すハンドラを先に定義しておきます。

from sqlalchemy.ext.automap import automap_base, generate_relationship

def _generate_relationship(base, direction, return_fn, attrname, local_cls, referred_cls, **kw):
    kw['lazy'] = 'bulk'
    return generate_relationship(base, direction, return_fn, attrname, local_cls, referred_cls, **kw)

3. automap_baseのprepareを呼ぶ

1. 1. のコードを含め、全コードは下記のようになります。

import urllib.parse
from sqlalchemy import create_engine, MetaData
from sqlalchemy.orm import scoped_session, sessionmaker
from sqlalchemy.ext.automap import automap_base, generate_relationship

# 1:N問題回避
from sqlalchemy_bulk_lazy_loader import BulkLazyLoader


def _generate_relationship(base, direction, return_fn, attrname, local_cls, referred_cls, **kw):
    kw['lazy'] = 'bulk'
    return generate_relationship(base, direction, return_fn, attrname, local_cls, referred_cls, **kw)


BulkLazyLoader.register_loader()

# DBへの接続作るよ
dsn = 'Driver=hogehoge'
engine = create_engine('hogehoge+pyodbc:///?odbc_connect=%s' % urllib.parse.quote_plus(dsn))
db_session = scoped_session(sessionmaker(bind=engine))

# automapするテーブルを制限するよ
metadata = MetaData()
metadata.reflect(engine, only=('employee', 'department'))
Base = automap_base(metadata=metadata)

# 先程定義した_generate_relationshipを指定してね
Base.prepare(generate_relationship=_generate_relationship)

Employee = Base.classes.employee
Department = Base.classes.department
Base.query = db_session.query_property()

4. N+1問題が解消される

GraphQL例

{
  department {
    edges {
      node {
         employees {
           name
         }
      }
    }
  }
}

発行されるSQL(プリペアドステートメント付き)

SELECT id, name from department

SELECT id, name from employee where id in (?, ?, ?, ?, ?, ?, ?, ?)

BulkLazyLoaderを使わないと関連するemployee全てのレコードに対してクエリが発行されます。

SELECT id, name from employee where id = ?;
SELECT id, name from employee where id = ?;
SELECT id, name from employee where id = ?;
SELECT id, name from employee where id = ?;
SELECT id, name from employee where id = ?;
SELECT id, name from employee where id = ?;
SELECT id, name from employee where id = ?;
SELECT id, name from employee where id = ?;

まとめ

BulkLazyLoader最高。
SQLAlchemyを使ったGraphQL開発はなかなか情報がなくて辛いです。