Add LoadedFieldExtension for batch loading individual fields

lib/graphql/batch.rb

@@ -39,6 +39,8 @@ def self.use(schema_defn, executor_class: GraphQL::Batch::Executor)
       end
       schema_defn.lazy_resolve(::Promise, :sync)
     end
+
+    autoload :LoadedFieldExtension, 'graphql/batch/loaded_field_extension'
   end
 end
 

lib/graphql/batch/loaded_field_extension.rb

@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+if Gem::Version.new(GraphQL::VERSION) < Gem::Version.new("1.9.0")
+  raise "GraphQL::Batch::LoadedFieldExtension is not supported on graphql gem versions less than 1.9"
+end
+
+require_relative 'loaded_field_extension/loader'
+
+module GraphQL::Batch
+  # Resolve the field using a class method on the GraphQL::Schema::Object
+  # for multiple instances. This avoids the need to extract the logic
+  # out into a {GraphQL::Batch::Loader} and automatically groups selections
+  # to load together.
+  #
+  # The class method must set the value on all given instances using an attribute
+  # writer of the same name as the resolver method.
+  #
+  # @example
+  #   class Product < GraphQL::Schema::Object
+  #     field :inventory_quantity, Int, null: false do
+  #       extension GraphQL::Batch::LoadedFieldExtension
+  #     end
+  #     def self.inventory_quantity(instances)
+  #       product_ids = instances.map { |instance| instance.object.id }
+  #       quantities = ProductVariant.group(:product_id).where(product_id: product_ids).sum(:inventory_quantity)
+  #       instances.each do |instance|
+  #         instance.inventory_quantity = quantities.fetch(instance.object.id, 0)
+  #       end
+  #     end
+  #
+  # For field selections to be loaded together, they must be given the same
+  # arguments. If the lookahead extra is used on the field, then it will group
+  # objects for the same selection set.
+  class LoadedFieldExtension < GraphQL::Schema::FieldExtension
+    def apply
+      @iv_name = iv_name = :"@#{field.resolver_method}"
+      resolver_method = field.resolver_method
+      field.owner.class_eval do
+        attr_writer(resolver_method)
+      end
+    end
+
+    def resolve(object:, arguments:, context:)
+      selections = if field.extras.include?(:lookahead)
+        arguments.delete(:lookahead)
+      elsif field.extras.include?(:irep_node)
+        arguments.delete(:irep_node)
+      end
+      Loader.for(selections, object.class, field.resolver_method, arguments, @iv_name).load(object)
+    end
+  end
+end

lib/graphql/batch/loaded_field_extension/loader.rb

@@ -0,0 +1,48 @@
+# frozen_string_literal: true
+
+require 'graphql/batch'
+
+module GraphQL::Batch
+  class LoadedFieldExtension < GraphQL::Schema::FieldExtension
+    class Loader < GraphQL::Batch::Loader
+      def self.loader_key_for(selections, object_class, resolver_method, arguments, iv_name)
+        [self, selections&.ast_nodes, object_class, resolver_method, arguments]
+      end
+
+      def initialize(selections, object_class, resolver_method, arguments, iv_name)
+        @selections = selections
+        @object_class = object_class
+        @resolver_method = resolver_method
+        @arguments = arguments
+        @iv_name = iv_name
+      end
+
+      def perform(instances)
+        arguments = @arguments
+        case @selections
+        when nil
+        when GraphQL::Execution::Lookahead
+          arguments = arguments.merge(lookahead: @selections)
+        when GraphQL::InternalRepresentation::Node
+          arguments = arguments.merge(irep_node: @selections)
+        end
+        if arguments.empty?
+          @object_class.public_send(@resolver_method, instances)
+        else
+          @object_class.public_send(@resolver_method, instances, arguments)
+        end
+        instances.each do |instance|
+          if instance.instance_variable_defined?(@iv_name)
+            value = instance.remove_instance_variable(@iv_name)
+            fulfill(instance, value)
+          else
+            message = "Attribute #{@resolver_method} wasn't set by " \
+              "#{@object_class}.#{@resolver_method} for object #{instance.object.inspect}"
+            reject(instance, ::Promise::BrokenError.new(message))
+          end
+        end
+      end
+    end
+    private_constant :Loader
+  end
+end

test/graphql_test.rb

@@ -1,16 +1,7 @@
 require_relative 'test_helper'
 
 class GraphQL::GraphQLTest < Minitest::Test
-  attr_reader :queries
-
-  def setup
-    @queries = []
-    QueryNotifier.subscriber = ->(query) { @queries << query }
-  end
-
-  def teardown
-    QueryNotifier.subscriber = nil
-  end
+  include QueryCollector
 
   def test_no_queries
     query_string = '{ constant }'

test/loaded_field_extension_test.rb

@@ -0,0 +1,213 @@
+# frozen_string_literal: true
+
+require_relative 'test_helper'
+
+# Only run tests for supported graphql gem versions
+if Gem::Version.new(GraphQL::VERSION) >= Gem::Version.new("1.9.0")
+  class GraphQL::Batch::LoadedFieldExtensionTest < Minitest::Test
+    include QueryCollector
+
+    class ImageType < GraphQL::Schema::Object
+      field :id, ID, null: false
+      field :filename, String, null: false
+    end
+
+    class ProductVariantType < GraphQL::Schema::Object
+      field :id, ID, null: false
+      field :title, String, null: false
+    end
+
+    class ProductType < GraphQL::Schema::Object
+      field :id, ID, null: false
+
+      field :variants, [ProductVariantType], null: true do
+        argument :first, Int, required: true
+        extension GraphQL::Batch::LoadedFieldExtension
+      end
+      def self.variants(instances, first:)
+        products = instances.map(&:object)
+        Product.preload_association(products, :variants)
+        instances.each do |instance|
+          instance.variants = instance.object.variants.first(first)
+        end
+      end
+
+      field :variants_count, Int, null: false do
+        extension GraphQL::Batch::LoadedFieldExtension
+      end
+      def self.variants_count(instances)
+        products = instances.map(&:object)
+        Product.preload_association(products, :variants)
+        instances.each do |instance|
+          instance.variants_count = instance.object.variants.length
+        end
+      end
+
+      field :image, ImageType, null: false, extras: [:lookahead] do
+        extension GraphQL::Batch::LoadedFieldExtension
+      end
+      def self.image(instances, lookahead:)
+        if lookahead.selections.any? { |s| s.name != :id }
+          ids = instances.map(&:object).map(&:image_id)
+          images = Image.find(ids)
+          instances.each do |instance|
+            instance.image = images.detect { |image| image.id == instance.object.image_id }
+          end
+        else
+          instances.each do |instance|
+            product = instance.object
+            instance.image = Image.new.tap { |image| image.id = product.image_id }
+          end
+        end
+      end
+
+      field :legacy_image, ImageType, null: false, extras: [:irep_node] do
+        extension GraphQL::Batch::LoadedFieldExtension
+      end
+      def self.legacy_image(instances, irep_node:)
+        if irep_node.scoped_children.values.flat_map(&:keys).any? { |key| key != 'id' }
+          ids = instances.map(&:object).map(&:image_id)
+          images = Image.find(ids)
+          instances.each do |instance|
+            instance.legacy_image = images.detect { |image| image.id == instance.object.image_id }
+          end
+        else
+          instances.each do |instance|
+            product = instance.object
+            instance.legacy_image = Image.new.tap { |image| image.id = product.image_id }
+          end
+        end
+      end
+
+      field :buggy, Int, null: false do
+        extension GraphQL::Batch::LoadedFieldExtension
+      end
+      def self.buggy(instances)
+        instances.first.buggy = 1
+      end
+    end
+
+    class QueryType < GraphQL::Schema::Object
+      field :products, [ProductType], null: false do
+        argument :first, Int, required: true
+      end
+      def products(first:)
+        Product.first(first)
+      end
+
+      field :product, ProductType, null: true do
+        argument :id, ID, required: true
+      end
+      def product(id:)
+        Product.find(Integer(id)).first
+      end
+    end
+
+    class Schema < GraphQL::Schema
+      query QueryType
+
+      if ENV["TESTING_INTERPRETER"] == "true"
+        use GraphQL::Execution::Interpreter
+      end
+
+      use GraphQL::Batch
+    end
+
+    def test_scalar_field
+      query_string = '{ products(first: 2) { id, variantsCount } }'
+      result = Schema.execute(query_string).to_h
+      expected = {
+        "data" => {
+          "products" => [
+            { "id" => '1', "variantsCount" => 2 },
+            { "id" => '2', "variantsCount" => 3 },
+          ]
+        }
+      }
+      assert_equal expected, result
+      assert_equal ["Product?limit=2", "Product/1,2/variants"], queries
+    end
+
+    def test_selections_with_same_arguments
+      query_string = <<~GRAPHQL
+        {
+          product1: product(id: "1") { variants(first: 1) { id } }
+          product2: product(id: "2") { variants(first: 1) { title } }
+        }
+      GRAPHQL
+      result = Schema.execute(query_string).to_h
+      expected = {
+        "data" => {
+          "product1" => { "variants" => [{ "id" => '1' }] },
+          "product2" => { "variants" => [{ "title" => 'Small' }] },
+        }
+      }
+      assert_equal expected, result
+      assert_equal ["Product/1", "Product/2", "Product/1,2/variants"], queries
+    end
+
+    def test_selections_with_different_arguments
+      query_string = <<~GRAPHQL
+        {
+          product1: product(id: "1") { variants(first: 1) { id } }
+          product2: product(id: "2") { variants(first: 2) { title } }
+        }
+      GRAPHQL
+      result = Schema.execute(query_string).to_h
+      expected = {
+        "data" => {
+          "product1" => { "variants" => [{ "id" => '1' }] },
+          "product2" => { "variants" => [{ "title" => 'Small' }, { "title" => 'Medium' }] },
+        }
+      }
+      assert_equal expected, result
+      assert_equal ["Product/1", "Product/2", "Product/1/variants", "Product/2/variants"], queries
+    end
+
+    def test_lookahead_with_different_nested_selections
+      query_string = <<~GRAPHQL
+        {
+          product1: product(id: "1") { image { filename } }
+          product2: product(id: "2") { image { id } }
+        }
+      GRAPHQL
+      result = Schema.execute(query_string).to_h
+      expected = {
+        "data" => {
+          "product1" => { "image" => { "filename" => 'shirt.jpg' } },
+          "product2" => { "image" => { "id" => '2' } },
+        }
+      }
+      assert_equal expected, result
+      assert_equal ["Product/1", "Product/2", "Image/1"], queries
+    end
+
+    def test_lookahead_with_shared_ast_nodes
+      query_string = <<~GRAPHQL
+        query {
+          product1: product(id: "1") { ...ProductFields }
+          product2: product(id: "2") { ...ProductFields }
+        }
+        fragment ProductFields on Product { image: legacyImage { filename } }
+      GRAPHQL
+      result = Schema.execute(query_string).to_h
+      expected = {
+        "data" => {
+          "product1" => { "image" => { "filename" => 'shirt.jpg' } },
+          "product2" => { "image" => { "filename" => 'pants.jpg' } },
+        }
+      }
+      assert_equal expected, result
+      assert_equal ["Product/1", "Product/2", "Image/1,2"], queries
+    end
+
+    def test_unset_value_error
+      query_string = '{ products(first: 2) { buggy } }'
+      error = assert_raises(::Promise::BrokenError) do
+        Schema.execute(query_string).to_h
+      end
+      product = Product.first(2)[1]
+      assert_equal error.message, "Attribute buggy wasn't set by #{ProductType}.buggy for object #{product.inspect}"
+    end
+  end
+end

test/support/query_collector.rb

@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+module QueryCollector
+  attr_reader :queries
+
+  def setup
+    @queries = []
+    QueryNotifier.subscriber = ->(query) { @queries << query }
+  end
+
+  def teardown
+    QueryNotifier.subscriber = nil
+  end
+end

test/test_helper.rb

@@ -4,5 +4,6 @@
 require_relative 'support/loaders'
 require_relative 'support/schema'
 require_relative 'support/db'
+require_relative 'support/query_collector'
 
 require 'minitest/autorun'