From a716aa63eb73945e6509da65a43767fd489185ae Mon Sep 17 00:00:00 2001
From: Robert Mosolgo <rdmosolgo@gmail.com>
Date: Tue, 30 Jul 2024 11:41:06 -0400
Subject: [PATCH] Add super calls to GraphqlTrace

Fix test to work with schema setup

Write a failing test
---
 .../graphql/tracers/graphql_trace.rb          | 24 +++++++------
 .../graphql/tracers/graphql_trace_test.rb     | 35 +++++++++++++++++++
 2 files changed, 48 insertions(+), 11 deletions(-)

diff --git a/instrumentation/graphql/lib/opentelemetry/instrumentation/graphql/tracers/graphql_trace.rb b/instrumentation/graphql/lib/opentelemetry/instrumentation/graphql/tracers/graphql_trace.rb
index 70c9666235..cb82bf2c7f 100644
--- a/instrumentation/graphql/lib/opentelemetry/instrumentation/graphql/tracers/graphql_trace.rb
+++ b/instrumentation/graphql/lib/opentelemetry/instrumentation/graphql/tracers/graphql_trace.rb
@@ -60,15 +60,15 @@ def initialize(trace_scalars: false, **_options)
           end
 
           def execute_multiplex(multiplex:, &block)
-            tracer.in_span('graphql.execute_multiplex', &block)
+            tracer.in_span('graphql.execute_multiplex') { super }
           end
 
           def lex(query_string:, &block)
-            tracer.in_span('graphql.lex', &block)
+            tracer.in_span('graphql.lex') { super }
           end
 
           def parse(query_string:, &block)
-            tracer.in_span('graphql.parse', &block)
+            tracer.in_span('graphql.parse') { super }
           end
 
           def validate(query:, validate:, &block)
@@ -89,11 +89,11 @@ def validate(query:, validate:, &block)
           end
 
           def analyze_multiplex(multiplex:, &block)
-            tracer.in_span('graphql.analyze_multiplex', &block)
+            tracer.in_span('graphql.analyze_multiplex') { super }
           end
 
           def analyze_query(query:, &block)
-            tracer.in_span('graphql.analyze_query', &block)
+            tracer.in_span('graphql.analyze_query') { super }
           end
 
           def execute_query(query:, &block)
@@ -102,11 +102,13 @@ def execute_query(query:, &block)
             attributes['graphql.operation.type'] = query.selected_operation.operation_type
             attributes['graphql.document'] = query.query_string
 
-            tracer.in_span('graphql.execute_query', attributes: attributes, &block)
+            tracer.in_span('graphql.execute_query', attributes: attributes) do
+              super
+            end
           end
 
           def execute_query_lazy(query:, multiplex:, &block)
-            tracer.in_span('graphql.execute_query_lazy', &block)
+            tracer.in_span('graphql.execute_query_lazy') { super }
           end
 
           def execute_field(field:, query:, ast_node:, arguments:, object:, &block)
@@ -133,7 +135,7 @@ def authorized(query:, type:, object:, &block)
 
             attributes = @_otel_type_attrs_cache[type]
 
-            tracer.in_span(platform_key, attributes: attributes, &block)
+            tracer.in_span(platform_key, attributes: attributes) { super }
           end
 
           def authorized_lazy(query:, type:, object:, &block)
@@ -141,19 +143,19 @@ def authorized_lazy(query:, type:, object:, &block)
             return super unless platform_key
 
             attributes = @_otel_lazy_type_attrs_cache[type]
-            tracer.in_span(platform_key, attributes: attributes, &block)
+            tracer.in_span(platform_key, attributes: attributes) { super }
           end
 
           def resolve_type(query:, type:, object:, &block)
             platform_key = @_otel_resolve_type_key_cache[type]
             attributes = @_otel_type_attrs_cache[type]
-            tracer.in_span(platform_key, attributes: attributes, &block)
+            tracer.in_span(platform_key, attributes: attributes) { super }
           end
 
           def resolve_type_lazy(query:, type:, object:, &block)
             platform_key = @_otel_resolve_type_key_cache[type]
             attributes = @_otel_lazy_type_attrs_cache[type]
-            tracer.in_span(platform_key, attributes: attributes, &block)
+            tracer.in_span(platform_key, attributes: attributes) { super }
           end
 
           private
diff --git a/instrumentation/graphql/test/instrumentation/graphql/tracers/graphql_trace_test.rb b/instrumentation/graphql/test/instrumentation/graphql/tracers/graphql_trace_test.rb
index 22e194036f..4fb0e9d85a 100644
--- a/instrumentation/graphql/test/instrumentation/graphql/tracers/graphql_trace_test.rb
+++ b/instrumentation/graphql/test/instrumentation/graphql/tracers/graphql_trace_test.rb
@@ -361,6 +361,41 @@ def platform_field_key(field)
           _(custom_events).must_equal(true)
         end
       end
+
+      module CustomTrace
+        def execute_multiplex(multiplex:)
+          multiplex.context[:custom_trace_execute_multiplex_ran] = true
+          super
+        end
+
+        def execute_query(query:)
+          query.context[:custom_trace_execute_query_ran] = true
+          super
+        end
+      end
+
+      it 'works with other trace modules' do
+        [GraphQL::Schema, SomeOtherGraphQLAppSchema, SomeGraphQLAppSchema].each(&:_reset_tracer_for_testing)
+        instrumentation.instance_variable_set(:@installed, false)
+
+        custom_trace_schema = Class.new(GraphQL::Schema) do
+          query(Class.new(GraphQL::Schema::Object) do
+            graphql_name 'Query'
+            field :int, Integer, fallback_value: 5
+          end)
+
+          trace_with(CustomTrace)
+        end
+
+        instrumentation.install({ schemas: [custom_trace_schema] })
+        res = custom_trace_schema.execute('{ int }')
+        assert_equal 5, res['data']['int'], 'The query ran successfully'
+
+        assert res.context[:custom_trace_execute_query_ran], 'The custom execute_query hook ran'
+        assert res.query.multiplex.context[:custom_trace_execute_multiplex_ran], 'The custom execute_multiplex hook ran'
+
+        assert spans.find { |s| s.name == 'graphql.execute_query' }, 'OpenTelementry ran'
+      end
     end
   end
 end