feat(PDO): correct a bug on prepared statement regarding DBM correlation

src/DDTrace/Integrations/DatabaseIntegrationHelper.php

@@ -21,7 +21,7 @@ class DatabaseIntegrationHelper
         Tag::TARGET_HOST,
     ];
 
-    public static function injectDatabaseIntegrationData(HookData $hook, $backend, $argNum = 0)
+    public static function injectDatabaseIntegrationData(HookData $hook, $backend, $argNum = 0, $forcedMode = null)
     {
         $allowedBackends = [
             "sqlsrv" => true,
@@ -32,7 +32,7 @@ public static function injectDatabaseIntegrationData(HookData $hook, $backend, $
             "odbc" => true,
         ];
 
-        $propagationMode = dd_trace_env_config("DD_DBM_PROPAGATION_MODE");
+        $propagationMode = $forcedMode ?? dd_trace_env_config("DD_DBM_PROPAGATION_MODE");
         if ($propagationMode != \DDTrace\DBM_PROPAGATION_DISABLED && isset($allowedBackends[$backend])) {
             $fullPropagationBackends = [
                 "mysql" => true,

src/DDTrace/Integrations/Mysqli/MysqliIntegration.php

@@ -152,7 +152,8 @@ function (SpanData $span, $args) {
             $span = $hook->span();
             self::setDefaultAttributes($span, 'mysqli_prepare', $query);
 
-            DatabaseIntegrationHelper::injectDatabaseIntegrationData($hook, 'mysql', 1);
+            // For prepared statements, downgrade to service mode
+            DatabaseIntegrationHelper::injectDatabaseIntegrationData($hook, 'mysql', 1, \DDTrace\DBM_PROPAGATION_SERVICE);
             self::handleRasp($span);
         }, static function (HookData $hook) {
             list($mysqli, $query) = $hook->args;
@@ -222,7 +223,8 @@ function (SpanData $span, $args) {
             $span = $hook->span();
             self::setDefaultAttributes($span, 'mysqli.prepare', $query);
 
-            DatabaseIntegrationHelper::injectDatabaseIntegrationData($hook, 'mysql');
+            // For prepared statements, downgrade to service mode
+            DatabaseIntegrationHelper::injectDatabaseIntegrationData($hook, 'mysql', 0, \DDTrace\DBM_PROPAGATION_SERVICE);
             self::handleRasp($span);
         }, static function (HookData $hook) {
             list($query) = $hook->args;

src/DDTrace/Integrations/PDO/PDOIntegration.php

@@ -125,7 +125,7 @@ public static function init(): int
             $instance = $hook->instance;
             PDOIntegration::setCommonSpanInfo($instance, $span);
 
-            PDOIntegration::injectDBIntegration($instance, $hook);
+            PDOIntegration::injectDBIntegration($instance, $hook, \DDTrace\DBM_PROPAGATION_SERVICE);
             PDOIntegration::handleRasp($instance, $span);
         }, static function (HookData $hook) {
             ObjectKVStore::propagate($hook->instance, $hook->returned, PDOIntegration::CONNECTION_TAGS_KEY);
@@ -265,7 +265,7 @@ private static function parseDsn($dsn)
         return $tags;
     }
 
-    public static function injectDBIntegration($pdo, $hook)
+    public static function injectDBIntegration($pdo, $hook, $forcedMode = null)
     {
         $driver = $pdo->getAttribute(\PDO::ATTR_DRIVER_NAME);
         if ($driver === "odbc") {
@@ -275,7 +275,7 @@ public static function injectDBIntegration($pdo, $hook)
                 return;
             }
         }
-        DatabaseIntegrationHelper::injectDatabaseIntegrationData($hook, $driver);
+        DatabaseIntegrationHelper::injectDatabaseIntegrationData($hook, $driver, 0, $forcedMode);
     }
 
     public static function extractConnectionMetadata(array $constructorArgs)

src/DDTrace/Integrations/SQLSRV/SQLSRVIntegration.php

@@ -69,7 +69,8 @@ public static function init(): int
             $span = $hook->span();
             self::setDefaultAttributes($conn, $span, 'sqlsrv_prepare', $query);
 
-            DatabaseIntegrationHelper::injectDatabaseIntegrationData($hook, 'sqlsrv', 1);
+            // For prepared statements, downgrade to service mode
+            DatabaseIntegrationHelper::injectDatabaseIntegrationData($hook, 'sqlsrv', 1, \DDTrace\DBM_PROPAGATION_SERVICE);
         }, static function (HookData $hook) {
             list($conn, $query) = $hook->args;
             $span = $hook->span();

tests/Integrations/Mysqli/MysqliTest.php

@@ -577,6 +577,58 @@ public function testProceduralPreparedStatementPeerServiceEnabled()
         ]);
     }
 
+    public function testPreparedStatementUsesServiceModeForDBM()
+    {
+        $query = "SELECT * FROM tests WHERE id = ?";
+        $traces = $this->isolateTracer(function () use ($query) {
+            $mysqli = new \mysqli(self::$host, self::$user, self::$password, self::$database);
+            $stmt = $mysqli->prepare($query);
+            $id = 1;
+            $stmt->bind_param('i', $id);
+            $stmt->execute();
+            $result = $stmt->get_result();
+            $results = $result->fetch_all();
+            $this->assertNotEmpty($results);
+            $mysqli->close();
+        });
+
+        // Get the raw spans
+        $spans = $traces[0];
+
+        // Find prepare and execute spans
+        $constructSpan = null;
+        $prepareSpan = null;
+        $executeSpan = null;
+
+        foreach ($spans as $span) {
+            if ($span['name'] === 'mysqli.__construct') {
+                $constructSpan = $span;
+            } elseif ($span['name'] === 'mysqli.prepare') {
+                $prepareSpan = $span;
+            } elseif ($span['name'] === 'mysqli_stmt.execute') {
+                $executeSpan = $span;
+            }
+        }
+
+        $this->assertNotNull($constructSpan, 'mysqli.__construct span should exist');
+        $this->assertNotNull($prepareSpan, 'mysqli.prepare span should exist');
+        $this->assertNotNull($executeSpan, 'mysqli_stmt.execute span should exist');
+
+        // Verify that execute and prepare span are siblings
+        $this->assertEquals(
+            $prepareSpan['parent_id'],
+            $executeSpan['parent_id'],
+            'mysqli_stmt.execute should be a sibling of mysqli.prepare'
+        );
+
+        // Verify that SERVICE mode is used for the prepare span
+        $this->assertArrayNotHasKey(
+            '_dd.dbm_trace_injected',
+            $prepareSpan['meta'] ?? [],
+            'mysqli.prepare should use SERVICE mode'
+        );
+    }
+
     public function testConstructorConnectError()
     {
         $traces = $this->isolateTracer(function () {

tests/Integrations/PDO/PDOTest.php

@@ -680,6 +680,104 @@ public function testPDOStatementExceptionPeerServiceEnabled()
         ]);
     }
 
+    public function testPreparedStatementUsesServiceModeForDBM()
+    {
+        $query = "SELECT * FROM tests WHERE id = ?";
+        $traces = $this->isolateTracer(function () use ($query) {
+            $pdo = $this->pdoInstance();
+            $stmt = $pdo->prepare($query);
+            $stmt->execute([1]);
+            $results = $stmt->fetchAll();
+            $this->assertEquals('Tom', $results[0]['name']);
+            $stmt->closeCursor();
+            $stmt = null;
+            $pdo = null;
+        });
+
+        // Get the raw spans
+        $spans = $traces[0];
+
+        // Find prepare and execute spans
+        $constructSpan = null;
+        $prepareSpan = null;
+        $executeSpan = null;
+
+        foreach ($spans as $span) {
+            if ($span['name'] === 'PDO.__construct') {
+                $constructSpan = $span;
+            } elseif ($span['name'] === 'PDO.prepare') {
+                $prepareSpan = $span;
+            } elseif ($span['name'] === 'PDOStatement.execute') {
+                $executeSpan = $span;
+            }
+        }
+
+        $this->assertNotNull($constructSpan, 'PDO.__construct span should exist');
+        $this->assertNotNull($prepareSpan, 'PDO.prepare span should exist');
+        $this->assertNotNull($executeSpan, 'PDOStatement.execute span should exist');
+
+        // Verify that execute and prepare span are siblings
+        $this->assertEquals(
+            $prepareSpan['parent_id'],
+            $executeSpan['parent_id'],
+            'PDOStatement.execute should be a sibling of PDO.prepare'
+        );
+
+        // Verify that SERVICE mode is used for the prepare span
+        $this->assertArrayNotHasKey(
+            '_dd.dbm_trace_injected',
+            $prepareSpan['meta'] ?? [],
+            'PDO.prepare should use SERVICE mode'
+        );
+    }
+
+    public function testDirectQueryHasNoParentIssues()
+    {
+        $query = "SELECT * FROM tests WHERE id=1";
+        $traces = $this->isolateTracer(function () use ($query) {
+            $pdo = $this->pdoInstance();
+            $pdo->query($query);
+            $pdo = null;
+        });
+
+        // Get the raw spans
+        $spans = $traces[0];
+
+        // Find construct and query spans
+        $constructSpan = null;
+        $querySpan = null;
+
+        foreach ($spans as $span) {
+            if ($span['name'] === 'PDO.__construct') {
+                $constructSpan = $span;
+            } elseif ($span['name'] === 'PDO.query') {
+                $querySpan = $span;
+            }
+        }
+
+        $this->assertNotNull($constructSpan, 'PDO.__construct span should exist');
+        $this->assertNotNull($querySpan, 'PDO.query span should exist');
+
+        // Verify that query span has a parent (should be root or construct)
+        $this->assertTrue(
+            isset($querySpan['parent_id']),
+            'PDO.query should have a parent_id'
+        );
+
+        // Verify spans are created correctly with proper structure
+        $this->assertSpans($traces, [
+            SpanAssertion::exists('PDO.__construct'),
+            SpanAssertion::build('PDO.query', 'pdo', 'sql', $query)
+                ->withExactTags($this->baseTags())
+                ->withExactMetrics([
+                    Tag::DB_ROW_COUNT => 1.0,
+                    Tag::ANALYTICS_KEY => 1.0,
+                    '_dd.agent_psr' => 1.0,
+                    '_sampling_priority_v1' => 1.0,
+                ]),
+        ]);
+    }
+
     public function testLimitedTracerPDO()
     {
         $query = "SELECT * FROM tests WHERE id = ?";

tests/Integrations/SQLSRV/SQLSRVTest.php

@@ -323,6 +323,53 @@ public function testPrepareErrorPeerServiceEnabled()
         ]);
     }
 
+    public function testPreparedStatementUsesServiceModeForDBM()
+    {
+        $query = "SELECT * FROM tests WHERE id = ?";
+        $traces = $this->isolateTracer(function () use ($query) {
+            $conn = $this->createConnection();
+            $stmt = sqlsrv_prepare($conn, $query, [1], ['Scrollable' => 'buffered']);
+            sqlsrv_execute($stmt);
+            sqlsrv_close($conn);
+        });
+
+        // Get the raw spans
+        $spans = $traces[0];
+
+        // Find prepare and execute spans
+        $connectSpan = null;
+        $prepareSpan = null;
+        $executeSpan = null;
+
+        foreach ($spans as $span) {
+            if ($span['name'] === 'sqlsrv_connect') {
+                $connectSpan = $span;
+            } elseif ($span['name'] === 'sqlsrv_prepare') {
+                $prepareSpan = $span;
+            } elseif ($span['name'] === 'sqlsrv_execute') {
+                $executeSpan = $span;
+            }
+        }
+
+        $this->assertNotNull($connectSpan, 'sqlsrv_connect span should exist');
+        $this->assertNotNull($prepareSpan, 'sqlsrv_prepare span should exist');
+        $this->assertNotNull($executeSpan, 'sqlsrv_execute span should exist');
+
+        // Verify that execute and prepare span are siblings
+        $this->assertEquals(
+            $prepareSpan['parent_id'],
+            $executeSpan['parent_id'],
+            'sqlsrv_execute should be a sibling of sqlsrv_prepare'
+        );
+
+        // Verify that SERVICE mode is used for the prepare span
+        $this->assertArrayNotHasKey(
+            '_dd.dbm_trace_injected',
+            $prepareSpan['meta'] ?? [],
+            'sqlsrv_prepare should use SERVICE mode'
+        );
+    }
+
     public function testExecError()
     {
         $query = "SELECT * FROM non_existing_table";