gql_batch_performance_tests.rb

require "bundler/inline"

# allows for declaring a Gemfile inline in a ruby script
# optionally installing any gems that aren't already installed
gemfile(true) do
  source "https://rubygems.org"

  gem "rails", "6.1.4.1"
  gem "sqlite3"
  gem "graphql", "~> 1.12"
  gem "rspec-rails", "~> 4.0.1"
  gem "db-query-matchers"
  gem "graphql-batch"
end

# we don't need logging
class App < Rails::Application
  config.logger = Logger.new('/dev/null')
end

App.initialize!

# db
# sqlite3 gql_demo => to create db

require "active_record"

ActiveRecord::Base.establish_connection(adapter: "sqlite3", database: "gql_demo")
ActiveRecord::Schema.define do
  create_table "tasks", force: :cascade do |t|
    t.string "name", null: false
    t.bigint "person_id", null: false
  end

  create_table "task_subscriptions", force: :cascade do |t|
    t.integer "task_id", null: false
    t.integer "person_id", null: false
  end

  create_table "people", force: :cascade do |t|
    t.string "name", null: false
  end
end

# models

class Task < ActiveRecord::Base
  has_many :task_subscriptions
end

class TaskSubscription < ActiveRecord::Base
  belongs_to :task
  belongs_to :person
end

class Person < ActiveRecord::Base
  has_many :tasks
  has_many :task_subscriptions, through: :tasks, source: :task_subscriptions
end

# # batch loader

class AssociationLoader < GraphQL::Batch::Loader
  def initialize(model, association_name)
    @model, @association_name = model, association_name
  end

  def load(record)
    return Promise.resolve(read_association(record)) if record.association(@association_name).loaded?

    super
  end

  def perform(records)
    ActiveRecord::Associations::Preloader.new.preload(records, @association_name)
    records.each { |record| fulfill(record, read_association(record)) }
  end

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

# gql types

class BaseObject < GraphQL::Schema::Object
  def current_person
    context[:current_person]
  end
end

module Types
  # Zeitwerk sigh :)
  class Person < BaseObject; end
  class Task < BaseObject; end
  class TaskSubscription < BaseObject; end
  class Query < BaseObject; end

  class Person
    field :name, String, null: false
    field :tasks, [Task], null: true
    field :task_subscriptions, [TaskSubscription], null: true
  end

  class Task
    field :name, String, null: false
    field :person, Types::Person, null: false

    field :task_subscriptions, [TaskSubscription], null: true
    field :task_subscriptions_batch, [TaskSubscription], null: true

    def task_subscriptions_batch
      AssociationLoader.for(Task, :task_subscriptions).load(object)
    end
  end

  class TaskSubscription
    field :id, Int, null: false
    field :person, Types::Person, null: false
    field :task, Types::Task, null: false
  end

  class Query
    field :session_owner, Types::Person, null: true, resolver_method: :current_person
  end
end

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

  use GraphQL::Batch
end

# specs

require "action_controller"
require 'rspec/rails'

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

RSpec.describe 'task subscriptions' do
  let(:picard)  { Person.create(name: "Jean Luc") }
  let(:riker)   { Person.create(name: "Number 1") }
  let(:troi)    { Person.create(name: "Counselor") }
  let(:beverly) { Person.create(name: "Beverly Crusher") }

  let(:picard_notes) { troi.tasks.create(name: "Jean Luc's hallucinations") }
  let(:riker_notes) { troi.tasks.create(name: "Will's daddy issues") }

  subject do
    GraphqlSchema.execute(query, context: { current_person: troi })
  end

  before do
    picard_notes.task_subscriptions.create(person: picard)
    riker_notes.task_subscriptions.create(person: riker)
    picard_notes.task_subscriptions.create(person: beverly)
    riker_notes.task_subscriptions.create(person: beverly)
  end

  context "with N+1" do
    let(:query) do
      <<~GQL
        query {
          sessionOwner {
            tasks {
              taskSubscriptions {
                id
              }
            }
          }
        }
      GQL
    end

    it "performs 3 queries" do
      # definitely an N+1 issues here with task subscriptions
      # SELECT "tasks".* FROM "tasks" WHERE "tasks"."person_id" = ?
      # SELECT "task_subscriptions".* FROM "task_subscriptions" WHERE "task_subscriptions"."task_id" = ?
      # SELECT "task_subscriptions".* FROM "task_subscriptions" WHERE "task_subscriptions"."task_id" = ?

      expect { subject }.to make_database_queries(count: 3)
    end
  end

  context "with GraphQL batch (no N+1)" do
    let(:query) do
      <<~GQL
        query {
          sessionOwner {
            tasks {
              taskSubscriptionsBatch {
                id
              }
            }
          }
        }
      GQL
    end

    it "performs 2 queries" do
      # definitely no more N+1 issues here with task subscriptions
      # SELECT "tasks".* FROM "tasks" WHERE "tasks"."person_id" = ?
      # SELECT "task_subscriptions".* FROM "task_subscriptions" WHERE "task_subscriptions"."task_id" IN (?, ?)
      expect { subject }.to make_database_queries(count: 2)
    end
  end

  context "with redesigned schema" do
    let(:query) do
      <<~GQL
        query {
          sessionOwner {
            taskSubscriptions {
              id
            }
          }
        }
      GQL
    end

    it "performs 1 query" do
      # definitely no more N+1 issues here with task subscriptions
      # SELECT "task_subscriptions".* FROM "task_subscriptions" INNER JOIN "tasks" ON "task_subscriptions"."task_id" = "tasks"."id" WHERE "tasks"."person_id" = ?
      expect { subject }.to make_database_queries(count: 1)
    end
  end
end