GraphQL::Batchを使ってよくあるユースケースをTwitterを例に実装してみる

こんにちは、プレックスの 石塚 です。

今回のブログでは、Ruby向けのデータローダーである GraphQL::Batch を使用して、よくアプリケーションで実装されるユースケースの実装方法をTwitterを例に紹介します。

データローダーとは？

データローダーとはGraphQLに用意されているデータを一括取得するための仕組みです。

例えば以下のようにユーザーを5件取得して、そのユーザーのツイート一覧をデータローダーを使用しないで取得しようとすると下記のように6回（5 + 1回）のクエリが発生することになります。

{
  users(first: 5) {
    id
    tweets {
      id
      content
    }
  }
}

SELECT id FROM users limit 5;
SELECT id, content FROM tweets WHERE id = 1;
SELECT id, content FROM tweets WHERE id = 2;
SELECT id, content FROM tweets WHERE id = 3;
SELECT id, content FROM tweets WHERE id = 4;
SELECT id, content FROM tweets WHERE id = 5;

これをデータローダーを使用することにより、以下のように2回のクエリでデータを取得できるようになります。

SELECT id FROM users limit 5;
SELECT id, content FROM tweets WHERE id IN (1, 2, 3, 4, 5);

GraphQLにおけるデータローダーの必要性

GraphQLではオブジェクトがネストしており、再帰的に処理を読み出せる構造になっているため、N+1が起きやすいという特徴があります。

Railsに慣れている方ならば、それってpreloadやincludesを使用すればN+1を避けることができるのではないのか？と思うかもしれません。 しかしその方法で対処してしまうと関連するモデルの呼び出しが必要のないところでも、不必要にpreloadやincludesなどが走ってしまったり、再帰的に呼び出された場合にN+1を避けることが難しくなってしまいます。そのため、基本的にはデータローダーを使うのがベターです。

本エントリーで実装したいこと

以下の3つのユースケースを実装していきます。Twitterを例にしていますが、サブテーマとして書いているように、アプリケーションを作っていると似たような機能を作ることがよくあるのではないかと思います。

• 複数のユーザーのツイート一覧を取得する
  • ActiveRecordのリレーションをデータローダーに置き換える
• あるユーザーのツイート一覧をいいね数と一緒に取得する
  • GROUP BYでのカウントの実装
• あるユーザーのツイート一覧を自身がいいねしたかどうかと一緒に取得する
  • WHERE句などの条件指定をできるように拡張する

前準備

動作環境やライブラリのバージョンは下記のとおりです。

	バージョン
ruby	3.0.2
rails	6.1.4
graphql-ruby	1.12.17
graphql-batch	0.4.3

Rubyにはgraphql-ruby標準のデータローダーもありますが、弊社では、スターの数が多く、使いやすいと感じた GraphQL::Batch を使っています。

以下のモデルファイルと対応するテーブルが存在すると仮定します。

class User < ApplicationRecord
  has_many :tweets
  has_many :favorites
end

class Tweet < ApplicationRecord
  has_many :favorites
  belongs_to :user
end

class Favorite < ApplicationRecord
  belongs_to :user
  belongs_to :tweet
end

実装

複数のユーザーのツイート一覧を取得する（ActiveRecordのリレーションをデータローダーに置き換える）

まずは一番基本的なユースケースとして起こり得そうな、複数のユーザーのツイート一覧を取得するQueryから実装していきます。

N+1を引き起こすコードを書いてから、それをデータローダーを使った実装にリファクタリングしていくという流れで書いていきます。

N+1を引き起こすコード

# app/graphql/types/query_type.rb

module Types
  class QueryType < Types::BaseObject
    include GraphQL::Types::Relay::HasNodeField
    include GraphQL::Types::Relay::HasNodesField

    field :users, [UserType], null: false do
      description 'ユーザー一覧をツイート一覧と共に取得する'
    end

    def users
      User.limit(5)
    end
  end
end

# app/graphql/types/user_type.rb

module Types
  class UserType < Types::BaseObject
    field :id, ID, null: false
    field :tweets, [TweetType], null: false
  end
end

# app/graphql/types/tweet_type.rb

module Types
  class TweetType < Types::BaseObject
    field :id, ID, null: false
    field :content, String, null: false
  end
end

これでN+1を引き起こすQueryの実装は完了です。幸か不幸かRailsにはアソシエーションという機能があるため、たったこれだけのコードでツイートとともにユーザーを取得するAPIを書けるのですが、N+1を引き起こしてしまいます。

データローダーを使ったコード

これをデータローダーを使ったコードに書き換えていきます。まずは GraphQL::Batch のインストール。

$ gem install graphql-batch

次にスキーマで GraphQL::Batch を使うように宣言します。

# app/graphql/my_schema.rb

class MySchema < GraphQL::Schema
  query MyQueryType
  mutation MyMutationType
  use GraphQL::Batch
end

GraphQL::Batch にはRailsのアソシエーションをよしなに読み込んでくれる AssociationLoader というクラスがサンプルで用意されています。 これと同じコードを app/graphql/loaders/association_loader.rb に記述します。

# app/graphql/loaders/association_loader.rb

module Loaders
  class class AssociationLoader < GraphQL::Batch::Loader
    def self.validate(model, association_name)
      new(model, association_name)
      nil
    end

    def initialize(model, association_name)
      super()
      @model = model
      @association_name = association_name
      validate
    end

    def load(record)
      raise TypeError, "#{@model} loader can't load association for #{record.class}" unless record.is_a?(@model)
      return Promise.resolve(read_association(record)) if association_loaded?(record)
      super
    end

    # We want to load the associations on all records, even if they have the same id
    def cache_key(record)
      record.object_id
    end

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

    private

    def validate
      unless @model.reflect_on_association(@association_name)
        raise ArgumentError, "No association #{@association_name} on #{@model}"
      end
    end

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

    def read_association(record)
      record.public_send(@association_name)
    end

    def association_loaded?(record)
      record.association(@association_name).loaded?
    end
  end
end

そして UserType から作成した AssociationLoader を使うように変更します。

# app/graphql/types/user_type.rb

module Types
  class UserType < Types::BaseObject
    field :id, ID, null: false
    field :tweets, [TweetType], null: false

    def tweets
      Loaders::AssociationLoader.for(User, :tweets).load(object)
    end
  end
end

これでリレーションではなく、データローダーを使ってツイートを読み込んでくれるため、N+1を解消できます。

あるユーザーのツイート一覧をいいね数と一緒に取得する（GROUP BYでのカウントの実装）

次にいいね数をカウントしてツイート一覧を返す実装です。

N+1を引き起こすコード

こちらもN+1を起こす実装から作っていきましょう。

# app/graphql/types/query_type.rb

module Types
  class QueryType < Types::BaseObject
    field :tweets, [TweetType], null: false do
      description 'ツイート一覧をいいね数と共に取得する'
    end

    def tweets
      Tweet.limit(5)
    end
  end
end

# app/graphql/types/tweet_type.rb

module Types
  class TweetType < Types::BaseObject
    field :id, ID, null: false
    field :content, String, null: false
    field :favorite_count, Int, null: false

    def favorite_count
      object.favorites.size
    end
  end
end

こちらで無事に？N+1クエリが発生します。

データローダーを使ったコード

次にデータローダーを作ってN+1を解消していきます。アソシエーション用のデータローダーは GraphQL::Batch で用意されていましたが、集計用のデータローダーは用意されていないため、アソシエーションを参考に自前で作っていきます。

# app/graphql/loaders/record_count_loader.rb

module Loaders
  class RecordCountLoader < GraphQL::Batch::Loader
    def initialize(model, column: model.primary_key, distinct_column: nil)
      super()
      @model = model
      @column = column.to_s
      @column_type = model.type_for_attribute(@column)
      @distinct_column = distinct_column
    end

    def load(key)
      super(@column_type.cast(key))
    end

    def perform(keys)
      query(keys).each { |key, count| fulfill(key, count) }
      keys.each { |key| fulfill(key, 0) unless fulfilled?(key) }
    end

    private

      def query(keys)
        scope = @model
        scope.where(@column => keys)
             .group(@column)
             .distinct(@distinct_column)
             .count
      end
  end
end

そして TweetType で実装したデータローダーを使うようにします。

# app/graphql/types/tweet_type.rb

module Types
  class TweetType < Types::BaseObject
    field :id, ID, null: false
    field :content, String, null: false
    field :favorite_count, Int, null: false

    def favorite_count
      Loaders::RecordCountLoader.for(Favorite,
                                     column: 'tweet_id',
                                     distinct_column: :tweet_id).load(object.id)
    end
  end
end

これでN+1を回避できます。データローダーを汎用的に使い回せる形で実装しておくと、タイプから呼び出すだけで実装が済むので便利ですね。

あるユーザーのツイート一覧を自身がいいねしたかどうかと一緒に取得する（WHERE句などの条件指定をできるように拡張する）

最後に特定のユーザーのツイート一覧を自身がいいねしたかどうかと一緒に取得するクエリを考えます。

N+1を引き起こすコード

こちらもまずはN+1を起こすクエリから書いていきます。

# app/graphql/types/query_type.rb

module Types
  class QueryType < Types::BaseObject
    field :user, UserType, null: false do
      description '特定のユーザーのツイート一覧を自身がいいねしたかどうかと一緒に取得する'
      argument :id, ID, required: true
    end

    def tweet(id:)
      Tweet.find(id)
    end
  end
end

UserType には先ほど作成した AssociationLoader でツイート一覧を読み込みます。

# app/graphql/types/user_type.rb

module Types
  class UserType < Types::BaseObject
    field :id, ID, null: false
    field :tweets, [TweetType], null: false

    def tweets
      Loaders::AssociationLoader.for(User, :tweets).load(object)
    end
  end
end

TweetType に新しく is_favorited という自身がいいねしたかどうかを表すフィールドを追加します。

認証情報は context に入っていると仮定して、 自身のユーザーIDをそこから取得しています。

# app/graphql/types/tweet_type.rb

module Types
  class TweetType < Types::BaseObject
    field :id, ID, null: false
    field :content, String, null: false
    field :favorite_count, Int, null: false
    field :is_favorited, Boolean, null: false

    def favorite_count
      Loaders::RecordCountLoader.for(Favorite,
                                     column: 'tweet_id',
                                     distinct_column: :tweet_id).load(object.id)
    end

    def is_favorited
      object.favorites.exists?(user_id: context[:user].id)
    end
  end
end

これでいいねしたかどうかの情報が取れるようになり、N+1が発生します。

データローダーを使ったコード

RecordPresenceLoader というデータローダーを新しく作成します。

RecordCountLoader との違いとしては引数として where を受け取れるようにしている点です。 GraphQL::Batch では ActiveRecord をベースにデータローダーを拡張していけるので、このような形でWHERE句を追加したり、JOIN句を追加していくことができます。

# app/graphql/loaders/record_presence_loader.rb

module Loaders
  class RecordPresenceLoader < GraphQL::Batch::Loader
    def initialize(model, column: model.primary_key, where: nil)
      super()
      @model = model
      @column = column.to_s
      @column_type = model.type_for_attribute(@column)
      @where = where
    end

    def load(key)
      super(@column_type.cast(key))
    end

    def perform(keys)
      query(keys).each { |record| fulfill(record.public_send(@column), record.present?) }
      keys.each { |key| fulfill(key, false) unless fulfilled?(key) }
    end

    private

      def query(keys)
        scope = @model
        scope = scope.where(@where) if @where
        scope.where(@column => keys)
      end
  end
end

そして RecordPresenceLoader をタイプから使用します。

# app/graphql/types/tweet_type.rb

module Types
  class TweetType < Types::BaseObject
    field :id, ID, null: false
    field :content, String, null: false
    field :favorite_count, Int, null: false
    field :is_favorited, Boolean, null: false

    def favorite_count
      Loaders::RecordCountLoader.for(Favorite, column: 'tweet_id', distinct_column: :tweet_id).load(object.id)
    end

    def is_favorited
      Loaders::RecordPresenceLoader.for(Favorite,
                                        column: 'tweet_id',
                                        where: {
                                          user_id: context[:user].id
                                        }).load(object.id)
    end
  end
end

これで自身がいいねしたかどうかの情報もN+1を起こさずに取得できるようになりました。

終わりに

今回の記事はデータローダーのライブラリの使い方という小ネタでしたが、普段GraphQLの実装をしていると、日本語で書かれている情報がまだまだ少ないと感じることが多いため、少しでも参考になればと思って書いてみました。

最後になりますが、プレックスではソフトウェアエンジニア、フロントエンドエンジニアを募集しています。

少しでも興味を持っていただけた方は業務委託や副業からでも、ぜひご応募いただけると幸いです。

dev.plex.co.jp