N+1 問題を解決する GraphQL::Batch の使い方とその仕組み

こんにちは。アルダグラムでエンジニアしている前山です。
本記事は株式会社アルダグラム Advent Calendar 2023 2日目の記事です。

弊社で開発している KANNA では、バックエンドは Rails で GraphQL Ruby を使用しています。

GraphQL を使用する開発の過程で、N+1 問題に直面することは珍しくありません。graphql-ruby には、N+1 問題に対処するための GraphQL::DataLoader が組み込まれていますが、KANNA では GraphQL::Batch を使用しています。

本記事では、GraphQL::Batch を用いた N+1 問題の解決方法と、その背後にある仕組みについて紹介できればと思います。

GraphQL::Batch とは

GraphQL::Batch は、Shopify によって開発された Rails のライブラリで GraphQL を用いた開発における、N+1 問題に対処するためのツールです。N+1 問題は、GraphQL クエリが複数の関連レコードを取得する際に、各レコードに対して個別のクエリが発行されることにより発生します。これにより、データベースへのクエリ数が増大し、結果としてアプリケーションのレスポンスタイムが低下することがあります。

GraphQL::Batch は、これらの個別のクエリを効率的に一つにまとめるバッチ処理を提供します。このプロセスは、複数のデータ取得要求を集約し、単一のクエリでデータを取得することで、データベースへの負荷を減少させることができます。

GraphQL::Batch を使った N+1 問題の解消方法

本記事では、Project モデルと Assignee モデルを例にあげます。Project は複数の Assignee (担当) を持つという関係です。

class Project < ApplicationRecord
  has_many :assignees
end

class Assignee < ApplicationRecord
  belongs_to :project
end

N+1 問題が発生する例

ProjectType で assignees フィールドを定義し、各 Project に関連する Assignee を取得します。

class ProjectType < GraphQL::Schema::Object
  field :assignees, [AssigneeType], null: false

  def assignees
    object.assignees
  end
end

この定義に基づいて、GraphQL クライアントが以下のようなクエリを実行したとします。

{
  projects {
    id
    name
    assignees {
      id
      name
    }
  }
}

上記の実装では、各プロジェクトに対して、関連する担当を取得するための個別のクエリが実行されます。プロジェクトの数だけ以下のようなクエリが実行されるため、N+1 問題が発生します。

SELECT "assignees".* FROM "assignees" WHERE "assignees"."project_id" = 1
SELECT "assignees".* FROM "assignees" WHERE "assignees"."project_id" = 2
...
SELECT "assignees".* FROM "assignees" WHERE "assignees"."project_id" = N

GraphQL::Batch を使った解決策

今回のケースの場合、Project が複数の Assignee を持つというシナリオにおいて、AssociationLoader が適しています。この Loader は関連する複数のオブジェクトを一度のクエリで効率的にロードし、各 Project に対して個別のクエリを発行する必要性を排除することで、N+1問題を解決するためです。

https://github.com/Shopify/graphql-batch/blob/main/examples/association_loader.rb

上記の Loader の実装を Rails プロジェクトに実装し、ProjectType を以下のように実装します。

class ProjectType < GraphQL::Schema::Object
  field :assignees, [AssigneeType], null: false

  def assignees
    AssociationLoader.for(Project, :assignees).load(object)
  end
end

この実装では、各 Project インスタンスに対して関連する Assignee を AssociationLoader を通じて効率的にロードします。その結果、SQL クエリは以下のように、単一の最適化されたクエリが生成され N+1 問題を解決します。

SELECT "assignees".* FROM "assignees" WHERE "assignees"."project_id" IN (1, 2, ..., N)

GraphQL::Batch の仕組み

GraphQL::Batch では GraphQL::Batch::Loader というクラスを提供し、これを継承したカスタムローダーを作成することで、特定のデータ取得要求に最適化されたクエリを実行することができます。

キーの収集

Loader クラスは、キーをキャッシュに格納し、キューに追加することで、データ取得要求を管理します。
具体的には、load メソッドを通じて各キー（データベースのレコードIDやオブジェクト等）をキューに追加し、それぞれに対して Promise を作成します。このプロセスは、以下のコードになります。

def load(key)
  cache[cache_key(key)] ||= begin
    queue << key
    ::Promise.new.tap { |promise| promise.source = self }
  end
end

上記のコードでは、まず与えられたキーがキャッシュに存在するかを確認し、存在しない場合にはキューに追加して新しい Promise オブジェクトを生成します。Promise オブジェクトは、非同期操作の完了を表し、データが利用可能になると解決されます。

Promise

上記の load メソッドで使用されている Promise オブジェクトは、非同期操作の完了を表すために用いられます。この Promise は、Rubyの promise.rb ライブラリに基づいています。Promise は、非同期処理の結果を表すプレースホルダーのようなもので、処理が完了すると「解決」(fulfilled) され、結果が利用可能になります。GraphQL::Batch では、各キーに対して非同期にデータを取得する際に Promise を使用し、データが利用可能になった時点でそのPromiseを「解決」しています。

Promise の使用により、GraphQL::Batch は非同期処理を効率的に管理し、データの取得が完了するまでの間、その他の処理をブロックしないで待機できます。これにより、リソースを効率的に活用し、パフォーマンスを向上させることができます。

SQL クエリの最適化

すべてのキーがキューに追加された後、perform メソッドがこれらのキーを基に SQL クエリを構築します。Loader クラスの perform メソッドは、データ取得のための主要なメソッドであり、AssociationLoader などのカスタムローダーでは、このメソッドをオーバーライドして、特定のデータを効率的に取得するロジックを実装します。

AssociationLoader の場合、perform メソッドから実行される ActiveRecord::Associations::Preloader が内部的にIN句を使用したクエリを生成し、関連データを一括で事前にロードします。

def perform(records)
  preload_association(records)
  records.each { |record| fulfill(record, read_association(record)) }
end

def preload_association(records)
  ::ActiveRecord::Associations::Preloader.new(records: records, associations: @association_name).call
end

データの取得とマッピング

次に、ロードした各レコードに対して関連するデータをマッピングします。このステップで各レコードに対して fulfill メソッドを使用して、特定のキーに対して取得したデータを割り当てます。

def perform(records)
  preload_association(records)
  records.each { |record| fulfill(record, read_association(record)) }
end

この方法により、各オブジェクトに対して必要な関連データが効率的に関連付けられ、N+1問題が解決されます。

これまでの処理の概要は以下になります。

最後に

以上で、GraphQL::Batch の使い方とその仕組みを紹介しました。本記事では主に AssociationLoader の実装を中心に紹介しましたが、実際のプロジェクトの特定の要件に合わせて柔軟にカスタムローダーを作成することが可能です。

また、GraphQL::Batch のライブラリのコードを読むことで、このツールの機能を深く理解する上で役に立ちました。コードを読むことにより、どのように効率的なデータ処理の実装をすれば良いか、より具体的かつ実践的な理解が得られてよかったです。

---

もっとアルダグラムエンジニア組織を知りたい人、ぜひ下記の情報をチェックしてみてください！