From ade85857190cf706e92d32737394d5b46083114d Mon Sep 17 00:00:00 2001
From: Robert Mosolgo <rdmosolgo@gmail.com>
Date: Fri, 2 Feb 2024 09:43:46 -0500
Subject: [PATCH] Use GraphQL::Tracing::DataDogTrace when it's available

---
 docs/GettingStarted.md                        | 30 ++++---------------
 .../tracing/contrib/graphql/patcher.rb        | 12 +++++++-
 .../contrib/graphql/integration_spec.rb       | 16 +++++++++-
 3 files changed, 31 insertions(+), 27 deletions(-)

diff --git a/docs/GettingStarted.md b/docs/GettingStarted.md
index eeb64012c50..c6afe585274 100644
--- a/docs/GettingStarted.md
+++ b/docs/GettingStarted.md
@@ -902,18 +902,8 @@ If you prefer to individually configure the tracer settings for a schema (e.g. y
 ```ruby
 # Class-based schema
 class YourSchema < GraphQL::Schema
-  use(
-    GraphQL::Tracing::DataDogTracing,
-    service: 'graphql'
-  )
-end
-```
-
-```ruby
-# .define-style schema
-YourSchema = GraphQL::Schema.define do
-  use(
-    GraphQL::Tracing::DataDogTracing,
+  trace_with(
+    GraphQL::Tracing::DatadogTrace,
     service: 'graphql'
   )
 end
@@ -923,22 +913,12 @@ Or you can modify an already defined schema:
 
 ```ruby
 # Class-based schema
-YourSchema.use(
-    GraphQL::Tracing::DataDogTracing,
-    service: 'graphql'
+YourSchema.trace_with(
+  GraphQL::Tracing::DatadogTrace,
+  service: 'graphql'
 )
 ```
 
-```ruby
-# .define-style schema
-YourSchema.define do
-  use(
-    GraphQL::Tracing::DataDogTracing,
-    service: 'graphql'
-  )
-end
-```
-
 Do *NOT* `instrument :graphql` in `Datadog.configure` if you choose to configure manually, as to avoid double tracing. These two means of configuring GraphQL tracing are considered mutually exclusive.
 
 ### gRPC
diff --git a/lib/datadog/tracing/contrib/graphql/patcher.rb b/lib/datadog/tracing/contrib/graphql/patcher.rb
index 0bf5771bc4e..bcac23863e0 100644
--- a/lib/datadog/tracing/contrib/graphql/patcher.rb
+++ b/lib/datadog/tracing/contrib/graphql/patcher.rb
@@ -28,7 +28,17 @@ def patch_schema!(schema)
             analytics_enabled = Contrib::Analytics.enabled?(get_option(:analytics_enabled))
             analytics_sample_rate = get_option(:analytics_sample_rate)
 
-            if schema.respond_to?(:use)
+            if schema.respond_to?(:trace_with)
+              schema.trace_with(
+                ::GraphQL::Tracing::DataDogTrace,
+                # By default, Tracing::DataDogTrace holds a reference to a tracer.
+                # If we provide a tracer argument here it will be eagerly cached,
+                # and Tracing::DataDogTracing will send traces to a stale tracer instance.
+                service: service_name,
+                analytics_enabled: analytics_enabled,
+                analytics_sample_rate: analytics_sample_rate
+              )
+            elsif schema.respond_to?(:use)
               schema.use(
                 ::GraphQL::Tracing::DataDogTracing,
                 # By default, Tracing::DataDogTracing holds a reference to a tracer.
diff --git a/spec/datadog/tracing/contrib/graphql/integration_spec.rb b/spec/datadog/tracing/contrib/graphql/integration_spec.rb
index ce1faff27da..11ff1e174db 100644
--- a/spec/datadog/tracing/contrib/graphql/integration_spec.rb
+++ b/spec/datadog/tracing/contrib/graphql/integration_spec.rb
@@ -22,10 +22,11 @@
   describe '.loaded?' do
     subject(:loaded?) { described_class.loaded? }
 
-    context 'when neither GraphQL or GraphQL::Tracing::DataDogTracing are defined' do
+    context 'when neither GraphQL, GraphQL::Tracing::DataDogTracing, nor GraphQL::Tracing::DataDogTrace are defined' do
       before do
         hide_const('GraphQL')
         hide_const('GraphQL::Tracing::DataDogTracing')
+        hide_const('GraphQL::Tracing::DataDogTrace')
       end
 
       it { is_expected.to be false }
@@ -35,6 +36,7 @@
       before do
         stub_const('GraphQL', Class.new)
         hide_const('GraphQL::Tracing::DataDogTracing')
+        hide_const('GraphQL::Tracing::DataDogTrace')
       end
 
       it { is_expected.to be false }
@@ -45,6 +47,18 @@
 
       it { is_expected.to be true }
     end
+
+    context 'when GraphQL::Tracing::DataDogTrace is defined' do
+      context 'when only GraphQL is defined' do
+        before do
+          stub_const('GraphQL', Class.new)
+          hide_const('GraphQL::Tracing::DataDogTracing')
+          stub_const('GraphQL::Tracing::DataDogTrace', Class.new)
+        end
+
+        it { is_expected.to be true }
+      end
+    end
   end
 
   describe '.compatible?' do