From 3ac27d865b4dba2d2f925ba35ff70b8b8b5b2577 Mon Sep 17 00:00:00 2001
From: Tony Hsu <tony.hsu@datadoghq.com>
Date: Sat, 27 Jan 2024 00:46:07 +0100
Subject: [PATCH] Pin to transport

---
 .../tracing/contrib/elasticsearch/patcher.rb  | 20 +++++++++++++++++++
 .../contrib/elasticsearch/transport_spec.rb   |  9 +--------
 2 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/lib/datadog/tracing/contrib/elasticsearch/patcher.rb b/lib/datadog/tracing/contrib/elasticsearch/patcher.rb
index d2ab99764d9..0b3396da445 100644
--- a/lib/datadog/tracing/contrib/elasticsearch/patcher.rb
+++ b/lib/datadog/tracing/contrib/elasticsearch/patcher.rb
@@ -24,6 +24,7 @@ def patch
             require 'json'
             require_relative 'quantize'
 
+            transport_module::Client.prepend(DatadogPin)
             transport_module::Client.prepend(Client)
           end
 
@@ -126,6 +127,25 @@ def datadog_configuration
           # rubocop:enable Metrics/MethodLength
           # rubocop:enable Metrics/AbcSize
 
+          # Patch to support both `elasticsearch` and `elastic-transport` versions
+          module DatadogPin
+            def datadog_pin=(pin)
+              pin.onto(pin_candidate)
+            end
+
+            def datadog_pin
+              Datadog.configuration_for(pin_candidate)
+            end
+
+            def pin_candidate(candidate = self)
+              if candidate.respond_to?(:transport)
+                pin_candidate(candidate.transport)
+              else
+                candidate
+              end
+            end
+          end
+
           # `Elasticsearch` namespace renamed to `Elastic` in version 8.0.0 of the transport gem:
           # @see https://github.com/elastic/elastic-transport-ruby/commit/ef804cbbd284f2a82d825221f87124f8b5ff823c
           def transport_module
diff --git a/spec/datadog/tracing/contrib/elasticsearch/transport_spec.rb b/spec/datadog/tracing/contrib/elasticsearch/transport_spec.rb
index e2f6f2c0921..6e0561e9e43 100644
--- a/spec/datadog/tracing/contrib/elasticsearch/transport_spec.rb
+++ b/spec/datadog/tracing/contrib/elasticsearch/transport_spec.rb
@@ -177,16 +177,9 @@ def call(env)
 
   describe 'transport configuration override' do
     context 'when #service is overridden' do
-      let(:version_greater_than_8) { Gem::Version.new(::Elasticsearch::VERSION) >= Gem::Version.new('8.0.0') }
       let(:service_name) { 'bar' }
 
-      before do
-        if version_greater_than_8
-          Datadog.configure_onto(client.transport, service_name: service_name)
-        else
-          Datadog.configure_onto(client.transport.transport, service_name: service_name)
-        end
-      end
+      before { Datadog.configure_onto(client.transport, service_name: service_name) }
 
       describe 'then a GET request' do
         subject(:response) { client.perform_request(method, path) }