From 810d6deb45ca9306cfea53bfb5a2d5cfd2b696c7 Mon Sep 17 00:00:00 2001
From: Alexandre Choura <alexandre.choura@datadoghq.com>
Date: Thu, 7 Dec 2023 11:29:15 +0100
Subject: [PATCH] fix: dddbs service mapping

---
 .../DatabaseIntegrationHelper.php             |  4 ++
 tests/Integration/DatabaseMonitoringTest.php  | 61 +++++++++++++++++++
 2 files changed, 65 insertions(+)

diff --git a/src/Integrations/Integrations/DatabaseIntegrationHelper.php b/src/Integrations/Integrations/DatabaseIntegrationHelper.php
index 084c0b1bf94..ce5221cd672 100644
--- a/src/Integrations/Integrations/DatabaseIntegrationHelper.php
+++ b/src/Integrations/Integrations/DatabaseIntegrationHelper.php
@@ -63,6 +63,10 @@ public static function propagateViaSqlComments($query, $databaseService, $mode =
         $tags = [];
 
         if ($databaseService != "") {
+            $mapping = dd_trace_env_config('DD_SERVICE_MAPPING');
+            if (isset($mapping[$databaseService])) {
+                $databaseService = $mapping[$databaseService];
+            }
             $tags["dddbs"] = $databaseService;
         }
 
diff --git a/tests/Integration/DatabaseMonitoringTest.php b/tests/Integration/DatabaseMonitoringTest.php
index 74c10d5b2d2..6e36b3dd025 100644
--- a/tests/Integration/DatabaseMonitoringTest.php
+++ b/tests/Integration/DatabaseMonitoringTest.php
@@ -57,6 +57,67 @@ public function testInjection()
         ]);
     }
 
+    public function testInjectionServiceMappingOnce()
+    {
+        try {
+            $hook = \DDTrace\install_hook(self::class . "::instrumented", function (HookData $hook) {
+                $span = $hook->span();
+                $span->service = "pdo";
+                $span->name = "instrumented";
+                DatabaseIntegrationHelper::injectDatabaseIntegrationData($hook, 'mysql', 1);
+            });
+            self::putEnv("DD_TRACE_DEBUG_PRNG_SEED=42");
+            self::putEnv("DD_DBM_PROPAGATION_MODE=full");
+            self::putEnv("DD_SERVICE_MAPPING=pdo:mapped-service");
+            $traces = $this->isolateTracer(function () use (&$commentedQuery) {
+                \DDTrace\start_trace_span();
+                $commentedQuery = $this->instrumented(0, "SELECT 1");
+                \DDTrace\close_span();
+            });
+        } finally {
+            \DDTrace\remove_hook($hook);
+        }
+
+        $this->assertRegularExpression('/^\/\*dddbs=\'mapped-service\',ddps=\'phpunit\',traceparent=\'00-[0-9a-f]{16}c151df7d6ee5e2d6-a3978fb9b92502a8-01\'\*\/ SELECT 1$/', $commentedQuery);
+        $this->assertFlameGraph($traces, [
+            SpanAssertion::exists("phpunit")->withChildren([
+                SpanAssertion::exists('instrumented')->withExactTags([
+                    "_dd.dbm_trace_injected" => "true",
+                    "_dd.base_service" => "mapped-service",
+                ])
+            ])
+        ]);
+    }
+
+    public function testInjectionServiceMappingTwice()
+    {
+        try {
+            $hook = \DDTrace\install_hook(self::class . "::instrumented", function (HookData $hook) {
+                $span = $hook->span();
+                $span->service = "pdo";
+                $span->name = "instrumented";
+                DatabaseIntegrationHelper::injectDatabaseIntegrationData($hook, 'mysql', 1);
+            });
+            self::putEnv("DD_TRACE_DEBUG_PRNG_SEED=42");
+            self::putEnv("DD_DBM_PROPAGATION_MODE=full");
+            self::putEnv("DD_SERVICE_MAPPING=pdo:mapped-service");
+            // Note that here, we don't start a new trace, hence the service mapping should apply to both dddbs & ddps
+            $traces = $this->isolateTracer(function () use (&$commentedQuery) {
+                $commentedQuery = $this->instrumented(0, "SELECT 1");
+            });
+        } finally {
+            \DDTrace\remove_hook($hook);
+        }
+
+        $this->assertRegularExpression('/^\/\*dddbs=\'mapped-service\',ddps=\'mapped-service\',traceparent=\'00-[0-9a-f]{16}c151df7d6ee5e2d6-c151df7d6ee5e2d6-01\'\*\/ SELECT 1$/', $commentedQuery);
+        $this->assertFlameGraph($traces, [
+            SpanAssertion::exists('instrumented')->withExactTags([
+                "_dd.dbm_trace_injected" => "true",
+                "_dd.base_service" => "mapped-service",
+            ])
+        ]);
+    }
+
     public function testInjectionPeerService()
     {
         try {