GraphQL Ruby Chain Loader

association_loader.rb

# frozen_string_literal: true

# https://github.com/Shopify/graphql-batch/blob/master/examples/association_loader.rb
require 'graphql/batch'

class AssociationLoader < GraphQL::Batch::Loader

  def initialize(association_name)
    @association_name = association_name
  end

  def load(record)
    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 preload_association(records)
    ::ActiveRecord::Associations::Preloader.new.preload(records, @association_name)
  end

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

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

end

chain_loader.rb

# frozen_string_literal: true

# A proxy that turns chained association calls into batch loaded promises.
# For brevity.
#
# In BaseObject there's a method `chain_load(x)` that does `ChainLoader.for(x)`
#
#   # This:
#   chain_load(tsb).offer.outlet.then do |tsb, offer, outlet|
#     routes.compass_outlet_offer_time_slot_url(
#       outlet.slug, offer.slug, tsb.id
#     )
#   end
#
#   # Gets turned into this (roughly):
#   RecordLoader.for(:offer).load(tsb).then do |offer|
#     RecordLoader.for(:outlet).load(offer).then do |outlet|
#       routes.compass_outlet_offer_time_slot_url(
#         outlet.slug, offer.slug, tsb.id
#       )
#     end
#   end
#
# Not really a Loader but using the same naming to show intent.
#
# As it's a proxy it inherits BasicObject which has less methods than Object,
# so all chained method calls are caught by method_missing.
class ChainLoader < BasicObject

  # underscore prefix so we don't clash with potential association names.
  attr_reader :_records, :_associations, :_block

  def self.for(record)
    new(record)
  end

  def initialize(record)
    @_associations = []
    @_records = [record]
  end

  # Each chained association hits this method_missing and is added to an array.
  def method_missing(association, *args)
    if args.empty?
      _associations << association # store association for later
      self # return self so we can chain
    else
      super
    end
  end

  # Be a good citizen.
  def respond_to_missing?(_association, *args)
    args.empty? # not an association
  end

  # First time called from user code with block and no `idx`.
  # Subsequent recursion calls increment `idx`.
  def then(idx = 0, &block)
    if idx >= _associations.size # recursion base case
      yield(*_records.drop(1)) # call the user's block with all the records
    else
      loader_for _associations[idx], load: _records[idx] do |record|
        _records << record
        self.then(idx + 1, &block) # recursion
      end
    end
  end

  private

  # Find the correct type of loader: single record or multi record association.
  def loader_for(association, load:, &block)
    if load.association(association).reflection.collection?
      ::AssociationLoader
    else
      ::RecordLoader
    end.for(association).load(load).then(&block)
  end

end

record_loader.rb

# frozen_string_literal: true

# https://github.com/Shopify/graphql-batch/blob/master/examples/record_loader.rb
require 'graphql/batch'

class RecordLoader < GraphQL::Batch::Loader

  def initialize(association)
    @association = association
  end

  # # Docs have this but seems to cause issues rather than solve them.
  # def cache_key(record)
  #   record.object_id
  # end

  def load(record)
    final_setup(record) unless parent # only needs to be done once

    super(record.send(foreign_key)) # Eg: offer.outlet_id
  end

  def perform(ids)
    child.where(primary_key => ids).each { |record| fulfill(record[primary_key], record) }
    # A has_one relationship can be nil but still needs fulfilling:
    ids.each { |id| fulfill(id, nil) unless fulfilled?(id) }
  end

  private

  attr_reader :association, :parent, :child, :primary_key, :foreign_key

  # Figure things out about the association
  def final_setup(record)
    @parent = record.class
    @child = parent.reflect_on_association(association).klass
    @foreign_key = parent.reflect_on_association(association).join_foreign_key
    @primary_key = parent.reflect_on_association(association).join_primary_key
  end

end