Merge pull request #4497 from rmosolgo/context-trace-instance

Accept trace instances during query execution

guides/queries/tracing.md

@@ -32,6 +32,8 @@ end
 
 For a full list of methods and their arguments, see {{ "GraphQL::Tracing::Trace" | api_doc }}.
 
+By default, GraphQL-Ruby makes a new trace instance when it runs a query. You can pass an existing instance as `context: { trace: ... }`. Also, `GraphQL.parse( ..., trace: ...)` accepts a trace instance.
+
 ## ActiveSupport::Notifications
 
 You can emit events to `ActiveSupport::Notifications` with an experimental tracer, `ActiveSupportNotificationsTrace`.

lib/graphql/query.rb

@@ -168,7 +168,7 @@ def interpreter?
 
     # @return [GraphQL::Tracing::Trace]
     def current_trace
-      @current_trace ||= multiplex ? multiplex.current_trace : schema.new_trace(multiplex: multiplex, query: self)
+      @current_trace ||= context[:trace] || (multiplex ? multiplex.current_trace : schema.new_trace(multiplex: multiplex, query: self))
     end
 
     def subscription_update?

lib/graphql/schema.rb

@@ -1047,6 +1047,7 @@ def execute(query_str = nil, **kwargs)
           {
             backtrace: ctx[:backtrace],
             tracers: ctx[:tracers],
+            trace: ctx[:trace],
             dataloader: ctx[:dataloader],
           }
         else

spec/graphql/execution/multiplex_spec.rb

@@ -242,4 +242,41 @@ def raise_exception
       assert_equal unhandled_err_json, InspectQueryInstrumentation.last_json
     end
   end
+
+  describe "context[:trace]" do
+    class MultiplexTraceSchema < GraphQL::Schema
+      class Query < GraphQL::Schema::Object
+        field :int, Integer
+        def int; 1; end
+      end
+
+      class Trace < GraphQL::Tracing::Trace
+        def execute_multiplex(multiplex:)
+          @execute_multiplex_count ||= 0
+          @execute_multiplex_count += 1
+          super
+        end
+
+        def execute_query(query:)
+          @execute_query_count ||= 0
+          @execute_query_count += 1
+          super
+        end
+
+        attr_reader :execute_multiplex_count, :execute_query_count
+      end
+
+      query(Query)
+    end
+
+    it "uses it instead of making a new trace" do
+      query_str = "{ int }"
+      trace_instance = MultiplexTraceSchema::Trace.new
+      res = MultiplexTraceSchema.multiplex([{query: query_str}, {query: query_str}], context: { trace: trace_instance })
+      assert_equal [1, 1], res.map { |r| r["data"]["int"]}
+
+      assert_equal 1, trace_instance.execute_multiplex_count
+      assert_equal 2, trace_instance.execute_query_count
+    end
+  end
 end

spec/graphql/query_spec.rb

@@ -1016,4 +1016,49 @@ def self.parse(query_str, trace: nil, filename: nil)
       assert_equal "Query", res["data"]["__typename"]
     end
   end
+
+  describe "context[:trace]" do
+    class QueryTraceSchema < GraphQL::Schema
+      class Query < GraphQL::Schema::Object
+        field :int, Integer
+        def int; 1; end
+      end
+
+      class Trace < GraphQL::Tracing::Trace
+        def execute_multiplex(multiplex:)
+          @execute_multiplex_count ||= 0
+          @execute_multiplex_count += 1
+          super
+        end
+
+        def execute_query(query:)
+          @execute_query_count ||= 0
+          @execute_query_count += 1
+          super
+        end
+
+        def execute_field(**rest)
+          @execute_field_count ||= 0
+          @execute_field_count += 1
+          super
+        end
+
+        attr_reader :execute_multiplex_count, :execute_query_count, :execute_field_count
+      end
+
+      query(Query)
+    end
+
+    it "uses it instead of making a new trace" do
+      query_str = "{ int __typename }"
+      trace_instance = QueryTraceSchema::Trace.new
+      res = QueryTraceSchema.execute(query_str, context: { trace: trace_instance })
+
+      assert_equal 1, res["data"]["int"]
+
+      assert_equal 1, trace_instance.execute_multiplex_count
+      assert_equal 1, trace_instance.execute_query_count
+      assert_equal 2, trace_instance.execute_field_count
+    end
+  end
 end