From 3db6575801835c3331fb57f89121accefde795ab Mon Sep 17 00:00:00 2001
From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com>
Date: Tue, 22 Sep 2026 19:21:51 -0400
Subject: [PATCH] fix: parent lazily imported module spans to the function span

Modules imported during the invocation were parented to whatever span
was active at import time. When that is a client span (e.g. botocore
lazily importing helpers on the first DynamoDB call), the backend peer
service pipeline learns the client's host -> the import span's service
and sets peer.service=aws.lambda on DynamoDB spans (APMS-20524).

Parent them to the function span instead. Imports during init are
unchanged, and span service, names and metrics are not affected.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 datadog_lambda/cold_start.py |  6 +++++-
 tests/test_cold_start.py     | 42 ++++++++++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/datadog_lambda/cold_start.py b/datadog_lambda/cold_start.py
index 260d491b..2581dc2e 100644
--- a/datadog_lambda/cold_start.py
+++ b/datadog_lambda/cold_start.py
@@ -202,7 +202,11 @@ def trace(self, root_nodes: List[ImportNode] = root_nodes):
         cold_start_span = self.create_cold_start_span(cold_start_span_start_time_ns)
         while root_nodes:
             root_node = root_nodes.pop()
-            parent = root_node.context or cold_start_span
+            # Modules lazily imported during the invocation hang off the function
+            # span, not the span active at import time: nesting them under a
+            # client span (e.g. a DynamoDB call) makes the backend peer service
+            # pipeline map that client's host to the import span's service.
+            parent = self.trace_ctx if root_node.context else cold_start_span
             self.trace_tree(root_node, parent)
         self.finish_span(cold_start_span, cold_start_span_end_time_ns)
 
diff --git a/tests/test_cold_start.py b/tests/test_cold_start.py
index 5fc11e2c..7ca83eaf 100644
--- a/tests/test_cold_start.py
+++ b/tests/test_cold_start.py
@@ -325,3 +325,45 @@ def handler(event, context):
     assert import_span.trace_id == function_span.trace_id
     assert load_span is not None
     assert load_span.trace_id == function_span.trace_id
+
+
+def test_lazy_loaded_package_imports_inside_client_span(monkeypatch):
+    spans = []
+
+    def finish(span):
+        spans.append(span)
+
+    monkeypatch.setattr(wrapper.tracer, "_on_span_finish", finish)
+    monkeypatch.setattr(wrapper, "is_new_sandbox", lambda: True)
+    monkeypatch.setattr("datadog_lambda.config.Config.trace_enabled", True)
+    monkeypatch.setenv(
+        "DD_COLD_START_TRACE_SKIP_LIB", "ddtrace.contrib.logging,datadog_lambda.wrapper"
+    )
+    monkeypatch.setenv("DD_MIN_COLD_START_DURATION", "0")
+    # ensure the import below is not served from the module cache
+    monkeypatch.delitem(modules, "colorsys", raising=False)
+
+    @wrapper.datadog_lambda_wrapper
+    def handler(event, context):
+        # e.g. botocore lazily importing modules on the first DynamoDB call
+        with wrapper.tracer.trace("dynamodb.command", service="aws.dynamodb") as span:
+            span.set_tag("span.kind", "client")
+            import colorsys  # noqa: F401
+
+    handler({}, get_mock_context())
+
+    function_span = client_span = import_span = None
+    for span in spans:
+        if span.resource == "colorsys":
+            import_span = span
+        elif span.name == "aws.lambda":
+            function_span = span
+        elif span.name == "dynamodb.command":
+            client_span = span
+
+    assert function_span is not None
+    assert client_span is not None
+    assert import_span is not None
+    assert import_span.parent_id == function_span.span_id
+    assert import_span.parent_id != client_span.span_id
+    assert import_span.trace_id == function_span.trace_id