Add HasDataloader

lib/graphql/dataloader/active_record_association_source.rb

@@ -7,8 +7,9 @@ class Dataloader
     class ActiveRecordAssociationSource < GraphQL::Dataloader::Source
       RECORD_SOURCE_CLASS = ActiveRecordSource
 
-      def initialize(association)
+      def initialize(association, scope: nil)
         @association = association
+        @scope = nil
       end
 
       def fetch(records)
@@ -29,7 +30,7 @@ def fetch(records)
           available_records.concat(already_loaded_records)
         end
 
-        ::ActiveRecord::Associations::Preloader.new(records: records, associations: @association, available_records: available_records).call
+        ::ActiveRecord::Associations::Preloader.new(records: records, associations: @association, available_records: available_records, scope: @scope).call
 
         loaded_associated_records = records.map { |r| r.public_send(@association) }
         records_by_model = {}
@@ -40,8 +41,12 @@ def fetch(records)
           end
         end
 
-        records_by_model.each do |model_class, updates|
-          dataloader.with(RECORD_SOURCE_CLASS, model_class).merge(updates)
+        if @scope.nil?
+          # Don't cache records loaded via scope because they might have reduced `SELECT`s
+          # Could check .select_values here?
+          records_by_model.each do |model_class, updates|
+            dataloader.with(RECORD_SOURCE_CLASS, model_class).merge(updates)
+          end
         end
 
         loaded_associated_records

lib/graphql/schema/member/has_dataloader.rb

@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+module GraphQL
+  class Schema
+    class Member
+      module HasDataloader
+        # @return [GraphQL::Dataloader] The dataloader for the currently-running query
+        def dataloader
+          context.dataloader
+        end
+
+        # Find an object with ActiveRecord via {Dataloader::ActiveRecordSource}.
+        # @param model [Class<ActiveRecord::Base>]
+        # @param find_by_value [Object] Usually an `id`, might be another value if `find_by:` is also provided
+        # @param find_by [Symbol, String] A column name to look the record up by. (Defaults to the model's primary key.)
+        # @return [ActiveRecord::Base, nil]
+        def dataload_record(model, find_by_value, find_by: nil)
+          source = if find_by
+            dataloader.with(Dataloader::ActiveRecordSource, model, find_by: find_by)
+          else
+            dataloader.with(Dataloader::ActiveRecordSource, model)
+          end
+
+          source.load(find_by_value)
+        end
+
+        # Look up an associated record using a Rails association.
+        # @param association_name [Symbol] A `belongs_to` or `has_one` association. (If a `has_many` association is named here, it will be selected without pagination.)
+        # @param record [ActiveRecord::Base] The object that the association belongs to.
+        # @param scope [ActiveRecord::Relation] A scope to look up the associated record in
+        # @return [ActiveRecord::Base, nil] The associated record, if there is one
+        # @example Looking up a belongs_to on the current object
+        #    dataload_association(:parent) # Equivalent to `object.parent`, but dataloaded
+        # @example Looking up an associated record on some other object
+        #    dataload_association(:post, comment) # Equivalent to `comment.post`, but dataloaded
+        def dataload_association(association_name, record = object, scope: nil)
+          source = if scope
+            dataloader.with(Dataloader::ActiveRecordAssociationSource, association_name, scope: scope)
+          else
+            dataloader.with(Dataloader::ActiveRecordAssociationSource, association_name)
+          end
+          source.load(record)
+        end
+      end
+    end
+  end
+end

lib/graphql/schema/resolver.rb

@@ -28,6 +28,7 @@ class Resolver
       include Schema::Member::HasPath
       extend Schema::Member::HasPath
       extend Schema::Member::HasDirectives
+      include Schema::Member::HasDataloader
 
       # @param object [Object] The application object that this field is being resolved on
       # @param context [GraphQL::Query::Context]
@@ -50,11 +51,6 @@ def initialize(object:, context:, field:)
       # @return [GraphQL::Query::Context]
       attr_reader :context
 
-      # @return [GraphQL::Dataloader]
-      def dataloader
-        context.dataloader
-      end
-
       # @return [GraphQL::Schema::Field]
       attr_reader :field