From d2d9d2fb5619175099a03005f90d867b746f8419 Mon Sep 17 00:00:00 2001
From: Bob Weinand <bob.weinand@datadoghq.com>
Date: Wed, 7 Feb 2024 13:59:05 +0100
Subject: [PATCH] Have datadog.trace.db_client_split_by_instance also affect
 mysqli and sqlsrv (#2508)

Signed-off-by: Bob Weinand <bob.weinand@datadoghq.com>
---
 .../Integrations/Mysqli/MysqliIntegration.php     |  6 ++++++
 .../Integrations/SQLSRV/SQLSRVIntegration.php     |  7 +++++++
 tests/Integrations/Mysqli/MysqliTest.php          | 14 ++++++++++++++
 tests/Integrations/PDO/PDOTest.php                | 15 +++++++++++++++
 4 files changed, 42 insertions(+)

diff --git a/src/Integrations/Integrations/Mysqli/MysqliIntegration.php b/src/Integrations/Integrations/Mysqli/MysqliIntegration.php
index 33eb8c30ba..d0ebf0ac87 100644
--- a/src/Integrations/Integrations/Mysqli/MysqliIntegration.php
+++ b/src/Integrations/Integrations/Mysqli/MysqliIntegration.php
@@ -424,6 +424,12 @@ public function setConnectionInfo(SpanData $span, $mysqli)
         foreach ($hostInfo as $tagName => $value) {
             $span->meta[$tagName] = $value;
         }
+        if (\DDTrace\Util\Runtime::getBoolIni("datadog.trace.db_client_split_by_instance")) {
+            if (isset($hostInfo[Tag::TARGET_HOST])) {
+                $span->service .=
+                    '-' . \DDTrace\Util\Normalizer::normalizeHostUdsAsService($hostInfo[Tag::TARGET_HOST]);
+            }
+        }
         $dbName = ObjectKVStore::get($mysqli, MysqliIntegration::KEY_DATABASE_NAME);
         if ($dbName) {
             $span->meta[Tag::DB_NAME] = $dbName;
diff --git a/src/Integrations/Integrations/SQLSRV/SQLSRVIntegration.php b/src/Integrations/Integrations/SQLSRV/SQLSRVIntegration.php
index f2717f5406..e44c8a2510 100644
--- a/src/Integrations/Integrations/SQLSRV/SQLSRVIntegration.php
+++ b/src/Integrations/Integrations/SQLSRV/SQLSRVIntegration.php
@@ -218,6 +218,13 @@ public static function setDefaultAttributes($source, SpanData $span, $name, $que
         foreach ($storedConnectionInfo as $tag => $value) {
             $span->meta[$tag] = $value;
         }
+
+        $targetName = $storedConnectionInfo[Tag::DB_INSTANCE] ?? $storedConnectionInfo[Tag::TARGET_HOST] ?? "<default>";
+        if (\DDTrace\Util\Runtime::getBoolIni("datadog.trace.db_client_split_by_instance")) {
+            if ($targetName !== "<default>") {
+                $span->service .= '-' . \DDTrace\Util\Normalizer::normalizeHostUdsAsService($targetName);
+            }
+        }
     }
 
     public static function detectError($SQLSRVRetval, SpanData $span)
diff --git a/tests/Integrations/Mysqli/MysqliTest.php b/tests/Integrations/Mysqli/MysqliTest.php
index 8faa53b58f..4ec7d3c74c 100644
--- a/tests/Integrations/Mysqli/MysqliTest.php
+++ b/tests/Integrations/Mysqli/MysqliTest.php
@@ -34,6 +34,7 @@ protected function envsToCleanUpAtTearDown()
             'DD_TRACE_PEER_SERVICE_DEFAULTS_ENABLED',
             'DD_TRACE_REMOVE_INTEGRATION_SERVICE_NAMES_ENABLED',
             'DD_SERVICE',
+            'DD_SERVICE_MAPPING',
         ];
     }
 
@@ -606,6 +607,19 @@ public function testNoFakeServices()
         ], true, false);
     }
 
+    public function testServiceMappedSplitByDomain()
+    {
+        self::putEnv('DD_TRACE_DB_CLIENT_SPLIT_BY_INSTANCE=true');
+        self::putEnv('DD_SERVICE_MAPPING=mysqli:my-mysqli');
+        $traces = $this->isolateTracer(function () {
+            new \mysqli(self::$host, self::$user, self::$password, self::$db);
+        });
+
+        $this->assertSpans($traces, [
+            SpanAssertion::build('mysqli.__construct', 'my-mysqli-mysql_integration', 'sql', 'mysqli.__construct', SpanAssertion::NOT_TESTED)
+        ]);
+    }
+
     private function baseTags($expectDbName = true, $expectPeerService = false)
     {
         $tags = [
diff --git a/tests/Integrations/PDO/PDOTest.php b/tests/Integrations/PDO/PDOTest.php
index 773c98a0ea..10ef2357d4 100644
--- a/tests/Integrations/PDO/PDOTest.php
+++ b/tests/Integrations/PDO/PDOTest.php
@@ -50,6 +50,7 @@ protected function envsToCleanUpAtTearDown()
             'DD_TRACE_DB_CLIENT_SPLIT_BY_INSTANCE',
             'DD_TRACE_PEER_SERVICE_DEFAULTS_ENABLED',
             'DD_TRACE_REMOVE_INTEGRATION_SERVICE_NAMES_ENABLED',
+            'DD_SERVICE_MAPPING',
             'DD_SERVICE',
         ];
     }
@@ -166,6 +167,20 @@ public function testPDOSplitByDomain()
         ]);
     }
 
+    public function testPDOServiceMappedSplitByDomain()
+    {
+        self::putEnv('DD_TRACE_DB_CLIENT_SPLIT_BY_INSTANCE=true');
+        self::putEnv('DD_SERVICE_MAPPING=pdo:my-pdo');
+        $traces = $this->isolateTracer(function () {
+            $this->pdoInstance();
+        });
+
+        $this->assertSpans($traces, [
+            SpanAssertion::build('PDO.__construct', 'my-pdo-mysql_integration', 'sql', 'PDO.__construct')
+                ->withExactTags($this->baseTags()),
+        ]);
+    }
+
     public function testPDOConstructError()
     {
         $traces = $this->isolateTracer(function () {