Use trace_with to start and end batches

Shopify · Jan 15, 2024 · 81bb4dc · 81bb4dc
1 parent 8f5b102
commit 81bb4dc
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 2 deletions.
diff --git a/lib/graphql/batch.rb b/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"

diff --git a/lib/graphql/batch/setup_multiplex.rb b/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