Merge pull request #164 from rmosolgo/use-trace-with

Use trace_with to start and end batches

lib/graphql/batch.rb

@@ -16,8 +16,12 @@ def self.batch(executor_class: GraphQL::Batch::Executor)
     end
 
     def self.use(schema_defn, executor_class: GraphQL::Batch::Executor)
-      instrumentation = GraphQL::Batch::SetupMultiplex.new(schema_defn, executor_class: executor_class)
-      schema_defn.instrument(:multiplex, instrumentation)
+      if schema_defn.respond_to?(:trace_with)
+        schema_defn.trace_with(GraphQL::Batch::SetupMultiplex::Trace, executor_class: executor_class)
+      else
+        instrumentation = GraphQL::Batch::SetupMultiplex.new(schema_defn, executor_class: executor_class)
+        schema_defn.instrument(:multiplex, instrumentation)
+      end
 
       if schema_defn.mutation
         require_relative "batch/mutation_field_extension"

lib/graphql/batch/setup_multiplex.rb

@@ -12,5 +12,19 @@ def before_multiplex(multiplex)
     def after_multiplex(multiplex)
       GraphQL::Batch::Executor.end_batch
     end
+
+    module Trace
+      def initialize(executor_class:, **_rest)
+        @executor_class = executor_class
+        super
+      end
+
+      def execute_multiplex(multiplex:)
+        GraphQL::Batch::Executor.start_batch(@executor_class)
+        super
+      ensure
+        GraphQL::Batch::Executor.end_batch
+      end
+    end
   end
 end