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

Signed-off-by: Alexandre Rulleau <alexandre.rulleau@datadoghq.com>

Author: Leiyks

src/DDTrace/Integrations/PDO/PDOIntegration.php

@@ -15,6 +15,7 @@ class PDOIntegration extends Integration
     const NAME = 'pdo';
 
     const CONNECTION_TAGS_KEY = 'connection_tags';
+    const PREPARE_SPAN_KEY = 'prepare_span';
 
     const DB_DRIVER_TO_SYSTEM = [
         'cubrid' => 'other_sql',
@@ -129,6 +130,10 @@ public static function init(): int
             PDOIntegration::handleRasp($instance, $span);
         }, static function (HookData $hook) {
             ObjectKVStore::propagate($hook->instance, $hook->returned, PDOIntegration::CONNECTION_TAGS_KEY);
+            // Store the prepare span so execute can use it as parent
+            if ($hook->returned instanceof \PDOStatement) {
+                ObjectKVStore::put($hook->returned, PDOIntegration::PREPARE_SPAN_KEY, $hook->span());
+            }
         });
 
         // public bool PDO::commit ( void )
@@ -143,7 +148,9 @@ public static function init(): int
         \DDTrace\install_hook(
             'PDOStatement::execute',
             static function (HookData $hook) {
-                $hook->span();
+                // If this statement was prepared with PDO::prepare, use that span as parent
+                $prepareSpan = ObjectKVStore::get($hook->instance, PDOIntegration::PREPARE_SPAN_KEY);
+                $hook->span($prepareSpan);
             },
             static function (HookData $hook) {
                 $span = $hook->span();

tests/Integrations/PDO/PDOTest.php

@@ -680,6 +680,97 @@ public function testPDOStatementExceptionPeerServiceEnabled()
         ]);
     }
 
+    public function testPreparedStatementExecuteIsChildOfPrepare()
+    {
+        $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 to check parent-child relationship
+        $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 span's parent is the prepare span
+        $this->assertEquals(
+            $prepareSpan['span_id'],
+            $executeSpan['parent_id'],
+            'PDOStatement.execute should be a child of PDO.prepare'
+        );
+    }
+
+    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 = ?";