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',
@@ -118,17 +119,37 @@ public static function init(): int
         \DDTrace\install_hook('PDO::prepare', static function (HookData $hook) {
             list($query) = $hook->args;
 
-            $span = $hook->span();
-            Integration::handleOrphan($span);
-            $span->name = 'PDO.prepare';
-            $span->resource = Integration::toString($query);
+            // Create the prepare span
+            $prepareSpan = $hook->span();
+            Integration::handleOrphan($prepareSpan);
+            $prepareSpan->name = 'PDO.prepare';
+            $prepareSpan->resource = Integration::toString($query);
             $instance = $hook->instance;
-            PDOIntegration::setCommonSpanInfo($instance, $span);
+            PDOIntegration::setCommonSpanInfo($instance, $prepareSpan);
+
+            // Create the execute span as a child of prepare for DBM injection
+            // This ensures traceparent points to the execute span
+            $executeSpan = \DDTrace\start_span();
+            $executeSpan->name = 'PDOStatement.execute';
+            $executeSpan->resource = Integration::toString($query);
+            PDOIntegration::setCommonSpanInfo($instance, $executeSpan);
 
+            // Inject DBM with the execute span's ID
             PDOIntegration::injectDBIntegration($instance, $hook);
-            PDOIntegration::handleRasp($instance, $span);
+            PDOIntegration::handleRasp($instance, $executeSpan);
+
+            // Close the execute span (removes from stack) but don't finish it yet
+            // It will be resumed during execute
+            \DDTrace\close_span();
+
+            // Store the execute span to be resumed by execute
+            $hook->data = ['execute_span' => $executeSpan];
         }, static function (HookData $hook) {
             ObjectKVStore::propagate($hook->instance, $hook->returned, PDOIntegration::CONNECTION_TAGS_KEY);
+            // Store the execute span in the PDOStatement for later use
+            if ($hook->returned instanceof \PDOStatement && !$hook->exception && isset($hook->data['execute_span'])) {
+                ObjectKVStore::put($hook->returned, PDOIntegration::PREPARE_SPAN_KEY, $hook->data['execute_span']);
+            }
         });
 
         // public bool PDO::commit ( void )
@@ -143,12 +164,29 @@ public static function init(): int
         \DDTrace\install_hook(
             'PDOStatement::execute',
             static function (HookData $hook) {
-                $hook->span();
+                // Check if we have a stored execute span from prepare
+                $storedSpan = ObjectKVStore::get($hook->instance, PDOIntegration::PREPARE_SPAN_KEY);
+                if (!$storedSpan) {
+                    // No stored span (e.g., from PDO::query()), create normally
+                    $hook->span();
+                }
+                // If we have a stored span, don't create a new one - we'll use the stored one
+                // The stored span was already created during prepare with the correct traceparent
             },
             static function (HookData $hook) {
-                $span = $hook->span();
                 $instance = $hook->instance;
+                $storedSpan = ObjectKVStore::get($instance, PDOIntegration::PREPARE_SPAN_KEY);
+
+                if ($storedSpan) {
+                    // Use the execute span that was created during prepare
+                    $span = $storedSpan;
+                } else {
+                    // No stored span (e.g., from PDO::query()), use the hook's span
+                    $span = $hook->span();
+                }
+
                 Integration::handleOrphan($span);
+                // Update span properties now that execute has completed
                 $span->name = 'PDOStatement.execute';
                 Integration::handleInternalSpanServiceName($span, PDOIntegration::NAME);
                 $span->type = Type::SQL;

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 = ?";