Automatically applying Dataloader to GraphQL fields

auto_dataloader.rb

require "bundler/inline"

gemfile do
  gem "graphql"
  gem "sqlite3"
  gem "activerecord"
end

# Set up the database for the example 
require "active_record"
ActiveRecord::Base.establish_connection({ adapter: "sqlite3", database: ":memory:" })
ActiveRecord::Schema.define do
  create_table :books, force: true do |t|
    t.string :title
    t.integer :author_id
  end
  create_table :authors, force: true do |t|
    t.string :name
  end
end

class Author < ActiveRecord::Base
end

class Book < ActiveRecord::Base
  belongs_to :author
end

Author.has_many(:books)
author = Author.create!(name: "Sandra Boynton")
Book.create!(title: "Barnyard Dance", author: author)
Book.create!(title: "Birthday Monsters", author: author)

author2 = Author.create!(name: "Beatrix Potter")
Book.create!(title: "The Tale of Mrs. Tiggywinkle", author: author2)
Book.create!(title: "The Tale of Ginger and Pickles", author: author2)

ActiveRecord::Base.logger = Logger.new(STDOUT)

# A basic dataloader source to fetch an object
module Sources
  class ActiveRecordObject < GraphQL::Dataloader::Source
    def initialize(model_class)
      @model_class = model_class
    end

    def fetch(ids)
      records = @model_class.where(id: ids)
      ids.map { |id| records.find { |r| r.id == id }}
    end
  end
end

class BaseField < GraphQL::Schema::Field
  # This resolver class looks at the field to determine the arguments to pass to dataloader. 
  # You could add custom options to `def initialize` below and save them as instance variables 
  # to customize or override this lookup.
  class AutoDataloadResolver < GraphQL::Schema::Resolver
    def resolve
      id_method = field.original_name.to_s + "_id"
      model_class = field.original_name.to_s.camelize.constantize
      record_id = object.public_send(id_method)
      dataloader.with(Sources::ActiveRecordObject, model_class).load(record_id)
    end
  end
  
  # Here, convert `dataloader: true` to `resolver_class: AutoDataloadResolver`
  def initialize(*args, dataloader: false, **kwargs, &block)
    if dataloader
      kwargs[:resolver_class] = AutoDataloadResolver
    end
    super(*args, **kwargs, &block)
  end

  private
end


class MySchema < GraphQL::Schema
  class BaseObject < GraphQL::Schema::Object
    field_class(BaseField)
  end

  class Author < BaseObject
    field :name, String
  end

  class Book < BaseObject
    field :title, String
    # Here it is in use: 
    field :author, Author, dataloader: true
  end

  class Query < BaseObject
    field :books, [Book]

    def books
      ::Book.all
    end
  end

  query(Query)
  use GraphQL::Dataloader
end

# An example query to show AutoDataloadResolver in action:
pp MySchema.execute("{ books { title author { name } } }").to_h
# Book Load (0.1ms)  SELECT "books".* FROM "books"
#  Author Load (0.1ms)  SELECT "authors".* FROM "authors" WHERE "authors"."id" IN (?, ?)  [["id", 1], ["id", 2]]
# {"data"=>
#   {"books"=>
#     [{"title"=>"Barnyard Dance", "author"=>{"name"=>"Sandra Boynton"}},
#      {"title"=>"Birthday Monsters", "author"=>{"name"=>"Sandra Boynton"}},
#      {"title"=>"The Tale of Mrs. Tiggywinkle", "author"=>{"name"=>"Beatrix Potter"}},
#      {"title"=>"The Tale of Ginger and Pickles", "author"=>{"name"=>"Beatrix Potter"}}]}}