From 64498b918c7bec0fd934348f735e82a8891818fe Mon Sep 17 00:00:00 2001
From: Robert Mosolgo <rdmosolgo@gmail.com>
Date: Mon, 16 May 2022 14:14:38 -0400
Subject: [PATCH 1/2] Add #prepare_span hook for custom Datadog tags

---
 lib/graphql/tracing/data_dog_tracing.rb       |  9 ++++++
 lib/graphql/tracing/platform_tracing.rb       |  4 +++
 spec/graphql/tracing/data_dog_tracing_spec.rb | 28 +++++++++++++++++++
 3 files changed, 41 insertions(+)

diff --git a/lib/graphql/tracing/data_dog_tracing.rb b/lib/graphql/tracing/data_dog_tracing.rb
index e710a8e087..6f582ff7e0 100644
--- a/lib/graphql/tracing/data_dog_tracing.rb
+++ b/lib/graphql/tracing/data_dog_tracing.rb
@@ -45,6 +45,8 @@ def platform_trace(platform_key, key, data)
             span.set_tag(:query_string, data[:query].query_string)
           end
 
+          prepare_span(key, data, span)
+
           yield
         end
       end
@@ -53,6 +55,13 @@ def service_name
         options.fetch(:service, 'ruby-graphql')
       end
 
+      # Implement this method in a subclass to apply custom tags to datadog spans
+      # @param key [String] The event being traced
+      # @param data [Hash] The runtime data for this event (@see GraphQL::Tracing for keys for each event)
+      # @param span [Datadog::Tracing::SpanOperation] The datadog span for this event
+      def prepare_span(key, data, span)
+      end
+
       def tracer
         default_tracer = defined?(Datadog::Tracing) ? Datadog::Tracing : Datadog.tracer
 
diff --git a/lib/graphql/tracing/platform_tracing.rb b/lib/graphql/tracing/platform_tracing.rb
index 77663b8a7b..66d062bba4 100644
--- a/lib/graphql/tracing/platform_tracing.rb
+++ b/lib/graphql/tracing/platform_tracing.rb
@@ -10,6 +10,10 @@ module Tracing
     class PlatformTracing
       class << self
         attr_accessor :platform_keys
+
+        def inherited(child_class)
+          child_class.platform_keys = self.platform_keys
+        end
       end
 
       def initialize(options = {})
diff --git a/spec/graphql/tracing/data_dog_tracing_spec.rb b/spec/graphql/tracing/data_dog_tracing_spec.rb
index 61d530b2c2..de7c10a007 100644
--- a/spec/graphql/tracing/data_dog_tracing_spec.rb
+++ b/spec/graphql/tracing/data_dog_tracing_spec.rb
@@ -18,6 +18,16 @@ class TestSchema < GraphQL::Schema
       query(Query)
       use(GraphQL::Tracing::DataDogTracing)
     end
+
+    class CustomTracerTestSchema < GraphQL::Schema
+      class CustomDataDogTracing < GraphQL::Tracing::DataDogTracing
+        def prepare_span(trace_key, data, span)
+          span.set_tag("custom:#{trace_key}", data.keys.join(","))
+        end
+      end
+      query(Query)
+      use(CustomDataDogTracing)
+    end
   end
 
   before do
@@ -47,4 +57,22 @@ class TestSchema < GraphQL::Schema
     assert_includes Datadog::SPAN_TAGS, ['component', 'graphql']
     assert_includes Datadog::SPAN_TAGS, ['operation', 'execute_multiplex']
   end
+
+  it "sets custom tags tags" do
+    DataDogTest::CustomTracerTestSchema.execute("{ int }")
+    expected_custom_tags = [
+      ["custom:lex", "query_string"],
+      ["custom:parse", "query_string"],
+      ["custom:execute_multiplex", "multiplex"],
+      ["custom:analyze_multiplex", "multiplex"],
+      ["custom:validate", "validate,query"],
+      ["custom:analyze_query", "query"],
+      ["custom:execute_query", "query"],
+      ["custom:authorized", "context,type,object,path"],
+      ["custom:execute_query_lazy", "multiplex,query"],
+    ]
+
+    actual_custom_tags = Datadog::SPAN_TAGS.reject { |t| t[0] == "operation" || t[0] == "component" || t[0].is_a?(Symbol) }
+    assert_equal expected_custom_tags, actual_custom_tags
+  end
 end

From d60dae3883c0ead1827082c2ef403631edb7bf50 Mon Sep 17 00:00:00 2001
From: Robert Mosolgo <rdmosolgo@gmail.com>
Date: Mon, 16 May 2022 14:44:04 -0400
Subject: [PATCH 2/2] fix lint error

---
 spec/graphql/schema/introspection_system_spec.rb | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/spec/graphql/schema/introspection_system_spec.rb b/spec/graphql/schema/introspection_system_spec.rb
index f232643594..342be2beb3 100644
--- a/spec/graphql/schema/introspection_system_spec.rb
+++ b/spec/graphql/schema/introspection_system_spec.rb
@@ -247,10 +247,10 @@ class HidingIntrospectionSchema < GraphQL::Schema
       module HideIntrospectionByContext
         def visible?(ctx)
           super && if introspection?
-            !ctx[:hide_introspection]
+                     !ctx[:hide_introspection]
           else
             true
-          end
+                   end
         end
       end