lib/graphql/batch.rb in graphql-batch-0.5.3 vs lib/graphql/batch.rb in graphql-batch-0.5.4

@@ -14,11 +14,15 @@
         GraphQL::Batch::Executor.end_batch
       end
     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"
 
         schema_defn.mutation.fields.each do |name, field|