Merge pull request #5213 from rmosolgo/rails-dataloader

Improve built-in Rails support for Dataloader

Author: rmosolgo

lib/graphql/dataloader.rb

@@ -4,6 +4,8 @@
 require "graphql/dataloader/request"
 require "graphql/dataloader/request_all"
 require "graphql/dataloader/source"
+require "graphql/dataloader/active_record_association_source"
+require "graphql/dataloader/active_record_source"
 
 module GraphQL
   # This plugin supports Fiber-based concurrency, along with {GraphQL::Dataloader::Source}.
@@ -255,6 +257,22 @@ def spawn_fiber
       }
     end
 
+    # Pre-warm the Dataloader cache with ActiveRecord objects which were loaded elsewhere.
+    # These will be used by {Dataloader::ActiveRecordSource}, {Dataloader::ActiveRecordAssociationSource} and their helper
+    # methods, `dataload_record` and `dataload_association`.
+    # @param records [Array<ActiveRecord::Base>] Already-loaded records to warm the cache with
+    # @param index_by [Symbol] The attribute to use as the cache key. (Should match `find_by:` when using {ActiveRecordSource})
+    # @return [void]
+    def merge_records(records, index_by: :id)
+      records_by_class = Hash.new { |h, k| h[k] = {} }
+      records.each do |r|
+        records_by_class[r.class][r.public_send(index_by)] = r
+      end
+      records_by_class.each do |r_class, records|
+        with(ActiveRecordSource, r_class).merge(records)
+      end
+    end
+
     private
 
     def calculate_fiber_limit

lib/graphql/dataloader/active_record_association_source.rb

@@ -0,0 +1,64 @@
+# frozen_string_literal: true
+require "graphql/dataloader/source"
+require "graphql/dataloader/active_record_source"
+
+module GraphQL
+  class Dataloader
+    class ActiveRecordAssociationSource < GraphQL::Dataloader::Source
+      RECORD_SOURCE_CLASS = ActiveRecordSource
+
+      def initialize(association, scope = nil)
+        @association = association
+        @scope = scope
+      end
+
+      def load(record)
+        if (assoc = record.association(@association)).loaded?
+          assoc.target
+        else
+          super
+        end
+      end
+
+      def fetch(records)
+        record_classes = Set.new.compare_by_identity
+        associated_classes = Set.new.compare_by_identity
+        records.each do |record|
+          if record_classes.add?(record.class)
+            reflection = record.class.reflect_on_association(@association)
+            if !reflection.polymorphic? && reflection.klass
+              associated_classes.add(reflection.klass)
+            end
+          end
+        end
+
+        available_records = []
+        associated_classes.each do |assoc_class|
+          already_loaded_records = dataloader.with(RECORD_SOURCE_CLASS, assoc_class).results.values
+          available_records.concat(already_loaded_records)
+        end
+
+        ::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 = {}
+        loaded_associated_records.each do |record|
+          if record
+            updates = records_by_model[record.class] ||= {}
+            updates[record.id] = record
+          end
+        end
+
+        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
+      end
+    end
+  end
+end

lib/graphql/dataloader/active_record_source.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+require "graphql/dataloader/source"
+
+module GraphQL
+  class Dataloader
+    class ActiveRecordSource < GraphQL::Dataloader::Source
+      def initialize(model_class, find_by: model_class.primary_key)
+        @model_class = model_class
+        @find_by = find_by
+        @type_for_column = @model_class.type_for_attribute(@find_by)
+      end
+
+      def load(requested_key)
+        casted_key = @type_for_column.cast(requested_key)
+        super(casted_key)
+      end
+
+      def fetch(record_ids)
+        records = @model_class.where(@find_by => record_ids)
+        record_lookup = {}
+        records.each { |r| record_lookup[r.public_send(@find_by)] = r }
+        record_ids.map { |id| record_lookup[id] }
+      end
+    end
+  end
+end

lib/graphql/schema/interface.rb

@@ -13,6 +13,7 @@ module DefinitionMethods
         include GraphQL::Schema::Member::Scoped
         include GraphQL::Schema::Member::HasAstNode
         include GraphQL::Schema::Member::HasUnresolvedTypeError
+        include GraphQL::Schema::Member::HasDataloader
         include GraphQL::Schema::Member::HasDirectives
         include GraphQL::Schema::Member::HasInterfaces
 

lib/graphql/schema/member.rb

@@ -2,6 +2,7 @@
 require 'graphql/schema/member/base_dsl_methods'
 require 'graphql/schema/member/graphql_type_names'
 require 'graphql/schema/member/has_ast_node'
+require 'graphql/schema/member/has_dataloader'
 require 'graphql/schema/member/has_directives'
 require 'graphql/schema/member/has_deprecation_reason'
 require 'graphql/schema/member/has_interfaces'

lib/graphql/schema/member/has_dataloader.rb

@@ -0,0 +1,56 @@
+# 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
+
+        # A shortcut method for loading a key from a source.
+        # Identical to `dataloader.with(source_class, *source_args).load(load_key)`
+        # @param source_class [Class<GraphQL::Dataloader::Source>]
+        # @param source_args [Array<Object>] Any extra parameters defined in `source_class`'s `initialize` method
+        # @param load_key [Object] The key to look up using `def fetch`
+        def dataload(source_class, *source_args, load_key)
+          dataloader.with(source_class, *source_args).load(load_key)
+        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(record = object, association_name, scope: nil)
+          source = if scope
+            dataloader.with(Dataloader::ActiveRecordAssociationSource, association_name, scope)
+          else
+            dataloader.with(Dataloader::ActiveRecordAssociationSource, association_name)
+          end
+          source.load(record)
+        end
+      end
+    end
+  end
+end

lib/graphql/schema/object.rb

@@ -7,6 +7,7 @@ class Schema
     class Object < GraphQL::Schema::Member
       extend GraphQL::Schema::Member::HasFields
       extend GraphQL::Schema::Member::HasInterfaces
+      include Member::HasDataloader
 
       # Raised when an Object doesn't have any field defined and hasn't explicitly opted out of this requirement
       class FieldsAreRequiredError < GraphQL::Error

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
 

spec/graphql/dataloader/active_record_association_source_spec.rb

@@ -0,0 +1,109 @@
+# frozen_string_literal: true
+require "spec_helper"
+
+describe GraphQL::Dataloader::ActiveRecordAssociationSource do
+  if testing_rails?
+    it_dataloads "queries for associated records when the association isn't already loaded" do |d|
+      my_first_car = ::Album.find(2)
+      homey = ::Album.find(4)
+      log = with_active_record_log(colorize: false) do
+        vulfpeck, chon = d.with(GraphQL::Dataloader::ActiveRecordAssociationSource, :band).load_all([my_first_car, homey])
+        assert_equal "Vulfpeck", vulfpeck.name
+        assert_equal "Chon", chon.name
+      end
+
+      assert_includes log, '[["id", 1], ["id", 3]]'
+
+      toms_story = ::Album.find(3)
+      log = with_active_record_log(colorize: false) do
+        vulfpeck, chon, toms_story_band = d.with(GraphQL::Dataloader::ActiveRecordAssociationSource, :band).load_all([my_first_car, homey, toms_story])
+        assert_equal "Vulfpeck", vulfpeck.name
+        assert_equal "Chon", chon.name
+        assert_equal "Tom's Story", toms_story_band.name
+      end
+
+      assert_includes log, '[["id", 2]]'
+    end
+
+    it_dataloads "doesn't load records that are already cached by ActiveRecordSource" do |d|
+      d.with(GraphQL::Dataloader::ActiveRecordSource, Band).load_all([1,2,3])
+
+      my_first_car = ::Album.find(2)
+      homey = ::Album.find(4)
+      toms_story = ::Album.find(3)
+
+      log = with_active_record_log(colorize: false) do
+        vulfpeck, chon, toms_story_band = d.with(GraphQL::Dataloader::ActiveRecordAssociationSource, :band).load_all([my_first_car, homey, toms_story])
+        assert_equal "Vulfpeck", vulfpeck.name
+        assert_equal "Chon", chon.name
+        assert_equal "Tom's Story", toms_story_band.name
+      end
+
+      assert_equal "", log
+    end
+
+    it_dataloads "warms the cache for ActiveRecordSource" do |d|
+      my_first_car = ::Album.find(2)
+      homey = ::Album.find(4)
+      toms_story = ::Album.find(3)
+      d.with(GraphQL::Dataloader::ActiveRecordAssociationSource, :band).load_all([my_first_car, homey, toms_story])
+
+      log = with_active_record_log(colorize: false) do
+        d.with(GraphQL::Dataloader::ActiveRecordSource, Band).load_all([1,2,3])
+      end
+
+      assert_equal "", log
+    end
+
+    it_dataloads "doesn't warm the cache when a scope is given" do |d|
+      my_first_car = ::Album.find(2)
+      homey = ::Album.find(4)
+      summerteeth = ::Album.find(6)
+      results = d.with(GraphQL::Dataloader::ActiveRecordAssociationSource, :band, ::Band.country).load_all([my_first_car, homey, summerteeth])
+      assert_equal [nil, nil, ::Band.find(4)], results
+
+      log = with_active_record_log(colorize: false) do
+        d.with(GraphQL::Dataloader::ActiveRecordSource, Band).load_all([1,2,4])
+      end
+
+      assert_includes log, "SELECT \"bands\".* FROM \"bands\" WHERE \"bands\".\"id\" IN (?, ?, ?)  [[\"id\", 1], [\"id\", 2], [\"id\", 4]]"
+    end
+
+    it_dataloads "doesn't pause when the association is already loaded" do |d|
+      source = d.with(GraphQL::Dataloader::ActiveRecordAssociationSource, :band)
+      assert_equal 0, source.results.size
+      assert_equal 0, source.pending.size
+
+      my_first_car = ::Album.find(2)
+      vulfpeck = my_first_car.band
+
+      vulfpeck2 = source.load(my_first_car)
+
+      assert_equal vulfpeck, vulfpeck2
+
+      assert_equal 0, source.results.size
+      assert_equal 0, source.pending.size
+
+      my_first_car.reload
+      vulfpeck3 = source.load(my_first_car)
+      assert_equal vulfpeck, vulfpeck3
+
+      assert_equal 1, source.results.size
+      assert_equal 0, source.pending.size
+    end
+
+    it_dataloads "raises an error with a non-existent association" do |d|
+      my_first_car = ::Album.find(2)
+      source = d.with(GraphQL::Dataloader::ActiveRecordAssociationSource, :tour_bus)
+      assert_raises ActiveRecord::AssociationNotFoundError do
+        source.load(my_first_car)
+      end
+    end
+
+    it_dataloads "works with polymorphic associations" do |d|
+      wilco = ::Band.find(4)
+      vulfpeck = d.with(GraphQL::Dataloader::ActiveRecordAssociationSource, :thing).load(wilco)
+      assert_equal ::Band.find(1), vulfpeck
+    end
+  end
+end