N+1 in graphql-ruby

1_graphql_ruby_n_plus_one.md

How to run examples:

1. Run $ createdb nplusonedb to create DB
2. Run specs $ rspec demo.rb

2_demo.rb

begin
  require "bundler/inline"
rescue LoadError => e
  $stderr.puts "Bundler version 1.10 or later is required. Please update your Bundler"
  raise e
end

gemfile(true) do
  source "https://rubygems.org"

  gem "rails", "6.0.3"
  gem "pg"
  gem "graphql", "~> 1.12"
  gem "rspec-rails", "~> 4.0.1"
  gem "db-query-matchers"
  gem "ar_lazy_preload"
  gem "graphql-batch"
end

# app initialization

require "active_record"

class App < Rails::Application
  config.logger = Logger.new('/dev/null')
end

App.initialize!

ActiveRecord::Base.establish_connection(adapter: "postgresql", database: "nplusonedb")

ActiveRecord::Schema.define do
  enable_extension "plpgsql"

  create_table "users", force: :cascade do |t|
    t.string "nickname", null: false
    t.datetime "created_at", precision: 6, null: false
    t.datetime "updated_at", precision: 6, null: false
  end

  create_table "user_connections", force: :cascade do |t|
    t.bigint "user_id", null: false
    t.bigint "follower_id", null: false
    t.datetime "created_at", precision: 6, null: false
    t.datetime "updated_at", precision: 6, null: false
  end

  create_table "tweets", force: :cascade do |t|
    t.text "content", null: false
    t.bigint "author_id", null: false
    t.datetime "created_at", precision: 6, null: false
    t.datetime "updated_at", precision: 6, null: false
  end

  add_foreign_key "tweets", "users", column: "author_id"
  add_foreign_key "user_connections", "users", column: "user_id"
  add_foreign_key "user_connections", "users", column: "follower_id"
end

# models

class UserConnection < ActiveRecord::Base
  belongs_to :user
  belongs_to :follower, class_name: "User"
end

class User < ActiveRecord::Base
  has_many :tweets, foreign_key: :author_id

  has_many :followers_connections, class_name: "UserConnection", foreign_key: :user_id
  has_many :followers, through: :followers_connections, source: :follower, class_name: "User"

  has_many :followed_connections, class_name: "UserConnection", foreign_key: :follower_id
  has_many :followed_users, through: :followed_connections, source: :user, class_name: "User"
end

class Tweet < ActiveRecord::Base
  belongs_to :author, class_name: "User"
end

# Batch loader

module Batch
  class RecordLoader < GraphQL::Batch::Loader
    def initialize(model)
      @model = model
    end

    def perform(ids)
      @model.where(id: ids).each { |record| fulfill(record.id, record) }
      ids.each { |id| fulfill(id, nil) unless fulfilled?(id) }
    end
  end
end

# Dataloader

module Sources
  class ActiveRecord < GraphQL::Dataloader::Source
    def initialize(model_class)
      @model_class = model_class
    end

    def fetch(ids)
      records = @model_class.where(id: ids).index_by(&:id)
      records.slice(*ids).values
    end
  end
end

# Feed query

module FeedBuilder
  module_function

  def for(user)
    Tweet.where(author: user.followed_users)
         .order(created_at: :desc)
         .limit(10)
  end
end

# GraphQL types part 1

module WithCurrentUser
  def current_user
    context[:current_user]
  end
end

class BaseObject < GraphQL::Schema::Object
  include WithCurrentUser
end

class BaseResolver < GraphQL::Schema::Resolver
  include WithCurrentUser
end

module Types
  class User < BaseObject
    field :nickname, String, null: false
    field :is_followed, Boolean, null: false, resolver_method: :followed?
    field :followers, [User], "cannot be fetched for multiple users at once", null: false do
      argument :limit, Integer, required: true, default_value: 2
      argument :cursor, Integer, required: false
    end

    def followed?
      Resolvers::LazyFollowedResolver.new(context, object.id)
    end

    def followers(limit:, cursor: nil)
      scope = object.followers.order(id: :desc).limit(limit)
      scope = scope.where("id < cursor", cursor) if cursor
      scope
    end
  end

  class Tweet < BaseObject
    field :content, String, null: false
    field :author, Types::User, null: false
    field :author_lazy, Types::User, null: false
    field :author_batch, Types::User, null: false
    field :author_dataloader, Types::User, null: false

    def author_lazy
      Resolvers::LazyUserResolver.new(context, object.author_id)
    end

    def author_batch
      Batch::RecordLoader.for(::User).load(object.author_id)
    end

    def author_dataloader
      dataloader.with(Sources::ActiveRecord, ::User).load(object.author_id)
    end
  end
end

# GraphQL resolvers

module Resolvers
  class BaseFeedResolver < BaseResolver
    type [Types::Tweet], null: false

    def resolve
      raise NotImplementedError
    end
  end

  class FeedResolver < BaseFeedResolver
    def resolve
      FeedBuilder.for(current_user)
    end
  end

  class FeedResolverPreload < BaseFeedResolver
    def resolve
      FeedBuilder.for(current_user).includes(:author)
    end
  end

  class FeedResolverLookahead < BaseFeedResolver
    extras [:lookahead]

    def resolve(lookahead:)
      FeedBuilder.for(current_user)
                 .merge(relation_with_includes(lookahead))
    end

    private

    def relation_with_includes(lookahead)
      return Tweet.all unless lookahead.selects?(:author)

      Tweet.includes(:author)
    end
  end

  class FeedResolverLazyPreload < BaseFeedResolver
    def resolve
      FeedBuilder.for(current_user).lazy_preload(:author)
    end
  end

  class LazyUserResolver
    def initialize(context, user_id)
      @user_id = user_id

      @lazy_state = context[:lazy_user_resolver] ||= {
        user_ids: Set.new,
        users_cache: nil,
      }

      @lazy_state[:user_ids] << user_id
    end

    def user
      users_cache[@user_id]
    end

    private

    def users_cache
      @lazy_state[:users_cache] ||=
        begin
          user_ids = @lazy_state[:user_ids].to_a
          @lazy_state[:user_ids].clear
          User.where(id: user_ids).index_by(&:id)
        end
    end
  end

  class LazyFollowedResolver
    def initialize(context, user_id)
      @user_id = user_id
      @context = context

      @lazy_state = context[:lazy_followed_resolver] ||= {
        user_ids: Set.new,
        users_cache: nil,
      }

      @lazy_state[:user_ids] << user_id
    end

    def is_followed?
      return false unless current_user
      users_cache[@user_id]
    end

    private

    def current_user
      @context[:current_user]
    end

    def users_cache
      @lazy_state[:users_cache] ||=
        begin
          user_ids = @lazy_state[:user_ids].to_a
          @lazy_state[:user_ids].clear
          User.where(id: user_ids.to_a)
              .pluck(Arel.sql("id, EXISTS(#{followed_query.to_sql}) AS is_followed"))
              .to_h
        end
    end

    def followed_query
      Arel::Table.new(:user_connections).then do |table|
        table.where(table[:follower_id].eq(current_user.id))
             .where(table[:user_id].eq(User.arel_table[:id]))
             .project("1")
      end
    end
  end
end

# More GraphQL types

module Types
  class Viewer < BaseObject
    field :feed, resolver: Resolvers::FeedResolver
    field :feed_with_preload, resolver: Resolvers::FeedResolverPreload
    field :feed_with_lookahead, resolver: Resolvers::FeedResolverLookahead
    field :feed_with_lazy_preload, resolver: Resolvers::FeedResolverLazyPreload
  end

  class Query < BaseObject
    field :viewer, Types::Viewer, null: true, resolver_method: :current_user
    field :users, [User], null: false, extras: [:lookahead]
    field :user, User, null: true do
      argument :user_id, ID, required: true
    end

    def users(lookahead:)
      if lookahead.selects?(:followers)
        raise GraphQL::ExecutionError, "followers can be accessed in singular association only"
      end

      ::User.all
    end

    def user(user_id:)
      ::User.find(user_id)
    rescue ActiveRecord::RecordNotFound
      nil
    end
  end
end

class GraphqlSchema < GraphQL::Schema
  lazy_resolve(Resolvers::LazyUserResolver, :user)
  lazy_resolve(Resolvers::LazyFollowedResolver, :is_followed?)

  query Types::Query

  use GraphQL::Batch
end

class GraphqlDataloaderSchema < GraphQL::Schema
  query Types::Query

  use GraphQL::Dataloader
end

# Specs

require "action_controller"
require 'rspec/rails'

RSpec.configure do |config|
  config.use_transactional_fixtures = true
end

RSpec.describe 'feed' do
  let(:john) { User.create(nickname: "John") }
  let(:jane) { User.create(nickname: "Jane") }
  let(:max) { User.create(nickname: "Max") }

  let(:query) do
    <<~GQL
      query {
        viewer {
          feed {
            content
            author {
              nickname
            }
          }
        }
      }
    GQL
  end

  subject do
    GraphqlSchema.execute(query, context: { current_user: john }).to_h
  end

  before do
    jane.followers << john
    max.followers << john

    jane.tweets.create(content: "Hi!", created_at: Time.new(2020, 6, 12, 10))
    max.tweets.create(content: "Hello!", created_at: Time.new(2020, 6, 13, 10))
    jane.tweets.create(content: "My second tweet is here", created_at: Time.new(2020, 6, 15, 10))
    max.tweets.create(content: "The weather is nice", created_at: Time.new(2020, 6, 17, 10))
  end

  shared_examples "feed loading" do
    it "returns feed data" do
      expect(subject["data"]).to eq(
        "viewer" => {
          "feed" => [
            { "content" => "The weather is nice", "author" => { "nickname" => "Max" } },
            { "content" => "My second tweet is here", "author" => { "nickname" => "Jane" } },
            { "content" => "Hello!", "author" => { "nickname" => "Max" } },
            { "content" => "Hi!", "author" => { "nickname" => "Jane" } }
          ]
        }
      )
    end
  end

  context "without N+1" do
    include_examples "feed loading"

    it "performs 5 queries" do
      expect { subject }.to make_database_queries(count: 5)
    end
  end

  context "with preloading" do
    let(:query) do
      <<~GQL
        query {
          viewer {
            feed: feedWithPreload {
              content
              author {
                nickname
              }
            }
          }
        }
      GQL
    end

    include_examples "feed loading"

    it "performs 2 queries" do
      expect { subject }.to make_database_queries(count: 2)
    end

    context "when author is not requested" do
      let(:query) do
        <<~GQL
          query {
            viewer {
              feed: feedWithPreload {
                content
              }
            }
          }
        GQL
      end

      it "performs 2 queries" do
        expect { subject }.to make_database_queries(count: 2)
      end
    end
  end

  context "with lookahead" do
    let(:query) do
      <<~GQL
        query {
          viewer {
            feed: feedWithLookahead {
              content
              author {
                nickname
              }
            }
          }
        }
      GQL
    end

    include_examples "feed loading"

    it "performs 2 queries" do
      expect { subject }.to make_database_queries(count: 2)
    end

    context "when author is not requested" do
      let(:query) do
        <<~GQL
          query {
            viewer {
              feed: feedWithLookahead {
                content
              }
            }
          }
        GQL
      end

      it "performs 1 query" do
        expect { subject }.to make_database_queries(count: 1)
      end
    end
  end

  context "with lazy preloading" do
    let(:query) do
      <<~GQL
        query {
          viewer {
            feed: feedWithLazyPreload {
              content
              author {
                nickname
              }
            }
          }
        }
      GQL
    end

    include_examples "feed loading"

    it "performs 2 queries" do
      expect { subject }.to make_database_queries(count: 2)
    end

    context "when author is not requested" do
      let(:query) do
        <<~GQL
          query {
            viewer {
              feed: feedWithLazyPreload {
                content
              }
            }
          }
        GQL
      end

      it "performs 1 query" do
        expect { subject }.to make_database_queries(count: 1)
      end
    end
  end

  context "with lazy resolver" do
    let(:query) do
      <<~GQL
        query {
          viewer {
            feed {
              content
              author: authorLazy {
                nickname
              }
            }
          }
        }
      GQL
    end

    include_examples "feed loading"

    it "performs 2 queries" do
      expect { subject }.to make_database_queries(count: 2)
    end
  end

  context "with batch loading" do
    let(:query) do
      <<~GQL
        query {
          viewer {
            feed {
              content
              author: authorBatch {
                nickname
              }
            }
          }
        }
      GQL
    end

    include_examples "feed loading"

    it "performs 2 queries" do
      expect { subject }.to make_database_queries(count: 2)
    end
  end

  context "with dataloader" do
    subject do
      GraphqlDataloaderSchema.execute(query, context: { current_user: john }).to_h
    end

    let(:query) do
      <<~GQL
        query {
          viewer {
            feed {
              content
              author: authorDataloader {
                nickname
              }
            }
          }
        }
      GQL
    end

    include_examples "feed loading"

    it "performs 2 queries" do
      expect { subject }.to make_database_queries(count: 2)
    end
  end
end

RSpec.describe 'followed' do
  let(:john) { User.create(nickname: "John") }
  let(:jane) { User.create(nickname: "Jane") }
  let(:max) { User.create(nickname: "Max") }

  let(:query) do
    <<~GQL
      query {
        viewer {
          feed {
            content
            author: authorLazy {
              nickname
              isFollowed
            }
          }
        }
      }
    GQL
  end

  subject do
    GraphqlSchema.execute(query, context: { current_user: john }).to_h
  end

  before do
    jane.followers << john
    max.followers << john

    jane.tweets.create(content: "Hi!", created_at: Time.new(2020, 6, 12, 10))
    max.tweets.create(content: "Hello!", created_at: Time.new(2020, 6, 13, 10))
  end

  it "returns followed data" do
    expect(subject["data"]).to eq(
      "viewer" => {
        "feed" => [
          { "content" => "Hello!", "author" => { "nickname" => "Max", "isFollowed" => true } },
          { "content" => "Hi!", "author" => { "nickname" => "Jane", "isFollowed" => true } }
        ]
      }
    )
  end

  it "performs 3 queries" do
    expect { subject }.to make_database_queries(count: 3)
  end
end

RSpec.describe 'followers' do
  let(:john) { User.create(nickname: "John") }
  let(:jane) { User.create(nickname: "Jane") }
  let(:max) { User.create(nickname: "Max") }

  subject do
    GraphqlSchema.execute(query, context: { current_user: john }, variables: { userId: john.id }).to_h
  end

  before do
    john.followers << jane
    john.followers << max

    jane.followers << john
    jane.followers << max
  end

  context "when followers of one user are requested" do
    let(:query) do
      <<~GQL
        query GetUser($userId: ID!) {
          user(userId: $userId) {
            followers {
              nickname
            }
          }
        }
      GQL
    end

    it "returns followed data" do
      expect(subject["data"]).to eq(
        "user" => {
          "followers" => [
            { "nickname" => "Max" },
            { "nickname" => "Jane" }
          ]
        }
      )
    end

    it "performs 2 queries" do
      expect { subject }.to make_database_queries(count: 2)
    end
  end

  context "when users with followers are requested" do
    let(:query) do
      <<~GQL
        query {
          users {
            nickname
            followers {
              nickname
            }
          }
        }
      GQL
    end

    it "returns error" do
      expect(subject["errors"].map { |error| error["message"] }).to \
        include("followers can be accessed in singular association only")
    end
  end
end

Thanks a lot for your article! It was very interesting and useful to read.

There is a little typo on the line 277 in the comment for the code:

Mpre GraphQL types

Oh, my bad, thanks @georgiybykov!