GraphQL Ruby で複合主キーなレコードをバッチロードする

概要

複合主キーなレコードをバッチロードしてみます。

自分の所属では数年前に GraphQL を導入したこともあり graphql-batch が使われていますが、今だと graphql-ruby の dataloder を使うのが一般的と思うので両方で実装してみます。

前提

以下のようなテーブルがあるとし、

viewer がユーザをフォローしているか否かを以下のクエリで取得するシチュエーションを考えます。

query ($id: ID!) {
  user(id: $id) { isFollowedByViewer }
}

GraphQL Batch の場合

Loaders::MultiColumnRecordLoader として実装します。

module Loaders
  class MultiColumnRecordLoader < GraphQL::Batch::Loader
    def initialize(model, columns:)
      @model = model
      @columns = columns
    end

    def perform(keys)
      conditions = keys.map {|key| @columns.zip(key).to_h }

      scoped = conditions.inject(@model.none) {|s, condition| s.or(@model.where(condition)) }

      scoped.each do |record|
        key = @columns.map {|column| record.public_send(column) }

        fulfill(key, record)
      end

      keys.each do |key|
        next if fulfilled?(key)

        fulfill(key, nil)
      end
    end
  end
end

呼び出しは以下のようになるでしょう

module Types
  class UserType < Types::BaseObject
    field :is_followed_by_viewer, Boolean, resolver_method: :followed_by_viewer?

    def followed_by_viewer?
      return false unless viewer

      loader = Loaders::MultiColumnRecordLoader.for(Follow, columns: %i[follower_id followee_id])

      loader([viewer.id, object.id]).then(&:present?)
    end
  end
end

GraphQL Ruby - Dataloader の場合

Sources::MultiColumnRecordSource として実装します。

module Sources
  class MultiColumnRecordSource < GraphQL::Dataloader::Source
    def initialize(model, columns:)
      @model = model
      @columns = columns
    end

    def fetch(keys)
      conditions = keys.map {|key| @columns.zip(key).to_h }

      scoped = conditions.inject(@model.none) {|s, condition| s.or(@model.where(condition)) }

      mapping = scoped.inject({}) do |record, m|
        key = @columns.map {|column| record.public_send(column) }

        m[key] = record
      end

      keys.map {|key| mapping.fetch(key, nil) }
    end
  end
end

呼び出しは以下のようになるでしょう

module Types
  class UserType < Types::BaseObject
    field :is_followed_by_viewer, Boolean, resolver_method: :followed_by_viewer?

    def followed_by_viewer?
      return false unless viewer

      loader = dataloader.with(Sources::MultiColumnRecordSource, Follow, %i[follower_id followee_id])

      follow = loader.load([viewer.id, object.id])

      follow.present?
    end
  end
end

所感

GraphQL Ruby の Dataloader の方が好きです。