Re-use loaders after they have been performed, but clear on mutation (#…

…63)

lib/graphql/batch.rb

@@ -20,10 +20,15 @@ def self.batch
     end
 
     def self.use(schema_defn)
+      schema = schema_defn.target
       if GraphQL::VERSION >= "1.6.0"
-        schema_defn.instrument(:multiplex, GraphQL::Batch::SetupMultiplex)
+        instrumentation = GraphQL::Batch::SetupMultiplex.new(schema)
+        schema_defn.instrument(:multiplex, instrumentation)
+        schema_defn.instrument(:field, instrumentation)
       else
-        schema_defn.instrument(:query, GraphQL::Batch::Setup)
+        instrumentation = GraphQL::Batch::Setup.new(schema)
+        schema_defn.instrument(:query, instrumentation)
+        schema_defn.instrument(:field, instrumentation)
       end
       schema_defn.lazy_resolve(::Promise, :sync)
     end

lib/graphql/batch/loader.rb

@@ -58,7 +58,6 @@ def resolve #:nodoc:
     # For Promise#sync
     def wait #:nodoc:
       if executor
-        executor.loaders.delete(loader_key)
         executor.resolve(self)
       else
         resolve

lib/graphql/batch/mutation_execution_strategy.rb

@@ -9,11 +9,13 @@ def get_finished_value(raw_value)
         strategy = execution_context.strategy
         return super if strategy.enable_batching
 
+        GraphQL::Batch::Executor.current.clear
         begin
           strategy.enable_batching = true
           strategy.deep_sync(Promise.sync(super))
         ensure
           strategy.enable_batching = false
+          GraphQL::Batch::Executor.current.clear
         end
       end
     end

lib/graphql/batch/setup.rb

@@ -1,14 +1,54 @@
 module GraphQL::Batch
-  module Setup
-    extend self
+  class Setup
+    class << self
+      def start_batching
+        raise NestedError if GraphQL::Batch::Executor.current
+        GraphQL::Batch::Executor.current = GraphQL::Batch::Executor.new
+      end
+
+      def end_batching
+        GraphQL::Batch::Executor.current = nil
+      end
+
+      def instrument_field(schema, type, field)
+        return field unless type == schema.mutation
+        old_resolve_proc = field.resolve_proc
+        field.redefine do
+          resolve ->(obj, args, ctx) {
+            GraphQL::Batch::Executor.current.clear
+            begin
+              Promise.sync(old_resolve_proc.call(obj, args, ctx))
+            ensure
+              GraphQL::Batch::Executor.current.clear
+            end
+          }
+        end
+      end
+
+      def before_query(query)
+        warn "Deprecated graphql-batch setup `instrument(:query, GraphQL::Batch::Setup)`, replace with `use GraphQL::Batch`"
+        start_batching
+      end
+
+      def after_query(query)
+        end_batching
+      end
+    end
+
+    def initialize(schema)
+      @schema = schema
+    end
 
     def before_query(query)
-      raise NestedError if GraphQL::Batch::Executor.current
-      GraphQL::Batch::Executor.current = GraphQL::Batch::Executor.new
+      Setup.start_batching
     end
 
     def after_query(query)
-      GraphQL::Batch::Executor.current = nil
+      Setup.end_batching
+    end
+
+    def instrument(type, field)
+      Setup.instrument_field(@schema, type, field)
     end
   end
 end

lib/graphql/batch/setup_multiplex.rb

@@ -1,14 +1,19 @@
 module GraphQL::Batch
-  module SetupMultiplex
-    extend self
+  class SetupMultiplex
+    def initialize(schema)
+      @schema = schema
+    end
 
     def before_multiplex(multiplex)
-      raise NestedError if GraphQL::Batch::Executor.current
-      GraphQL::Batch::Executor.current = GraphQL::Batch::Executor.new
+      Setup.start_batching
     end
 
     def after_multiplex(multiplex)
-      GraphQL::Batch::Executor.current = nil
+      Setup.end_batching
+    end
+
+    def instrument(type, field)
+      Setup.instrument_field(@schema, type, field)
     end
   end
 end

test/graphql_test.rb

@@ -279,6 +279,36 @@ def test_load_list_of_objects_with_loaded_field
     assert_equal ["Product?limit=2", "Product/1,2/variants", "ProductVariant/1,2,4,5,6/images"], queries
   end
 
+  def test_loader_reused_after_loading
+    query_string = <<-GRAPHQL
+      {
+        product(id: "2") {
+          variants {
+            id
+            product {
+              id
+              title
+            }
+          }
+        }
+      }
+    GRAPHQL
+    result = schema_execute(query_string)
+    expected = {
+      "data" => {
+        "product" => {
+          "variants" => [
+            { "id" => "4", "product" => { "id" => "2", "title" => "Pants" } },
+            { "id" => "5", "product" => { "id" => "2", "title" => "Pants" } },
+            { "id" => "6", "product" => { "id" => "2", "title" => "Pants" } },
+          ],
+        }
+      }
+    }
+    assert_equal expected, result
+    assert_equal ["Product/2", "Product/2/variants"], queries
+  end
+
   def test_load_error
     query_string = <<-GRAPHQL
       {

test/support/schema.rb

@@ -17,6 +17,11 @@
       end
     }
   end
+  field :product, !ProductType do
+    resolve ->(variant, _, _) {
+      RecordLoader.for(Product).load(variant.product_id)
+    }
+  end
 end
 
 ProductType = GraphQL::ObjectType.define do