Preloading Associations with graphql-batch

association_loader.rb

# frozen_string_literal: true
class AssociationLoader < GraphQL::Batch::Loader
  attr_reader :klass, :association

  def initialize(klass, association)
    raise ArgumentError, "association to load must be a symbol (got #{association.inspect})" unless association.is_a?(Symbol)
    raise ArgumentError, "cannot load associations for class #{klass.name}" unless klass < ActiveRecord::Base
    raise TypeError, "association #{association} does not exist on #{klass.name}" unless klass.reflect_on_association(association)

    @klass = klass
    @association = association
  end

  def load(model)
    raise TypeError, "loader for #{klass.name} can't load associations for objects of type #{model.class.name}" unless model.is_a?(klass)
    model.association(@association).loaded? ? Promise.resolve(model) : super
  end

  def perform(models)
    ActiveRecord::Associations::Preloader.new.preload(models, association)
    models.each { |m| fulfill(m, m) }
  end
end

association_preload_instrumentation.rb

# frozen_string_literal: true
class AssociationPreloadInstrumentation
  def instrument(_type, field)
    return field unless field.metadata.include?(:preload)

    old_resolver = field.resolve_proc
    new_resolver = -> (object, args, ctx) do
      preload(object, field.metadata[:preload]).then do
        old_resolver.call(object, args, ctx)
      end
    end

    field.redefine do
      resolve(new_resolver)
    end
  end

  private

  def preload(object, associations)
    if associations.is_a?(Symbol)
      preload_single_association(object, associations)
    else
      promises = []

      Array.wrap(associations).each do |value|
        case value
        when Symbol
          promises << preload_single_association(object, value)
        when Array
          value.each { |sub_value| promises << preload(object, sub_value) }
        when Hash
          value.each do |key, sub_value|
            promises << preload_single_association(object, key).then do
              next_value = object.public_send(key)
              case next_value
              when ActiveRecord::Base
                preload(next_value, sub_value)
              else
                Promise.all(Array.wrap(next_value).map { |next_model| preload(next_model, sub_value) })
              end
            end
          end
        end
      end

      Promise.all(promises)
    end
  end

  def preload_single_association(object, association)
    return Promise.resolve(object) if object.association(association).loaded?
    AssociationLoader.for(object.class, association).load(object)
  end
end

initializer.rb

# frozen_string_literal: true
# This goes in a Rails initializer, or some other code that runs before your schema/types are required
GraphQL::Field.accepts_definitions(
  preload: -> (type, *args) do
    type.metadata[:preload] ||= []
    type.metadata[:preload].concat(args)
  end
)

post_type.rb

PostType = GraphQL::ObjectType.define do
  name "Post"
  
  field :comments, !types[!CommentType] do
    preload comments: :author
    resolve -> (post, args, ctx) { post.comments }
  end
  
  # Or you can use the more terse syntax:
  field :comments, !types[!CommentType], preload: { comments: :author }, property: :comments
end

schema.rb

# frozen_string_literal: true
Schema = GraphQL::Schema.define do
   # configuration stuff
  instrument(:field, AssociationPreloadInstrumentation.new)
end