fix: support ALB multiValueHeaders for inferred aws.alb span

Handle both ALB event subtypes via resolve_alb_request_headers, fix
inbound trace propagation from multiValueHeaders, and expand unit/wrapper
coverage for HTTPS URLs, peer.service, and dd_resource_key.

FRSLES-851

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: chemystery09

datadog_lambda/tracing.py

@@ -43,6 +43,7 @@
     is_step_function_event,
     EventTypes,
     EventSubtypes,
+    resolve_alb_request_headers,
 )
 from datadog_lambda.durable import extract_context_from_durable_execution
 
@@ -197,6 +198,11 @@ def extract_context_from_http_event_or_context(
             return context
 
     headers = event.get("headers")
+    if not isinstance(headers, dict) or not headers:
+        if isinstance(event.get("multiValueHeaders"), dict):
+            headers = resolve_alb_request_headers(event)
+        else:
+            headers = {}
     context = propagator.extract(headers)
 
     if not _is_context_complete(context):
@@ -658,7 +664,9 @@ def extract_dd_trace_context(
         context = extract_context_from_request_header_or_context(
             event, lambda_context, event_source
         )
-    elif isinstance(event, (set, dict)) and "headers" in event:
+    elif isinstance(event, (set, dict)) and (
+        "headers" in event or "multiValueHeaders" in event
+    ):
         context = extract_context_from_http_event_or_context(
             event, lambda_context, event_source, decode_authorizer_context
         )
@@ -837,7 +845,7 @@ def create_inferred_span(
         elif event_source.equals(EventTypes.LAMBDA_FUNCTION_URL):
             logger.debug("Function URL event detected. Inferring a span")
             return create_inferred_span_from_lambda_function_url_event(event, context)
-        elif event_source.equals(EventTypes.ALB, subtype=EventSubtypes.ALB):
+        elif event_source.event_type == EventTypes.ALB:
             logger.debug("ALB event detected. Inferring a span")
             return create_inferred_span_from_alb_event(event, context)
         elif event_source.equals(
@@ -960,9 +968,7 @@ def create_inferred_span_from_alb_event(event, context):
     elb = request_context.get("elb") or {}
     target_group_arn = elb.get("targetGroupArn")
 
-    headers = event.get("headers")
-    if not isinstance(headers, dict):
-        headers = {}
+    headers = resolve_alb_request_headers(event)
     host = headers.get("host")
     method = event.get("httpMethod")
     path = event.get("path")
@@ -972,7 +978,9 @@ def create_inferred_span_from_alb_event(event, context):
     # fall back to it when DD_TRACE_AWS_SERVICE_REPRESENTATION_ENABLED is on.
     service_name = determine_service_name(service_mapping, host, "lambda_alb", host)
 
-    http_url = f"{proto}://{host}{path}" if host and path is not None else None
+    http_url = (
+        "%s://%s%s" % (proto, host, path) if host and path is not None else None
+    )
     if method and path is not None:
         resource = f"{method} {path}"
     else:
@@ -992,13 +1000,13 @@ def create_inferred_span_from_alb_event(event, context):
     # Drop tags we couldn't derive so the span never carries malformed values.
     tags = {key: value for key, value in tags.items() if value is not None}
 
-    InferredSpanInfo.set_tags(tags, tag_source="self", synchronicity="sync")
     tracer.set_tags(_dd_origin)
     # ALB events carry no request timestamp (unlike API GW requestTimeEpoch /
     # Function URL timeEpoch), so the span starts at handler time.
     span = tracer.trace(
         "aws.alb", service=service_name, resource=resource, span_type="http"
     )
+    InferredSpanInfo.set_tags(tags, tag_source="self", synchronicity="sync")
     if span:
         span.set_tags(tags)
         span.set_metric(InferredSpanInfo.METRIC, 1.0)

datadog_lambda/trigger.py

@@ -293,6 +293,29 @@ def get_event_source_arn(source: _EventSource, event: dict, context: Any) -> str
     return event_source_arn
 
 
+def resolve_alb_request_headers(event):
+    """
+    Resolve ALB request headers from single-value ``headers`` or
+    ``multiValueHeaders`` (first value per key, matching datadog-lambda-js).
+    """
+    headers = event.get("headers")
+    if isinstance(headers, dict) and headers:
+        return headers
+
+    multi_value = event.get("multiValueHeaders")
+    if not isinstance(multi_value, dict):
+        return {}
+
+    resolved = {}
+    for key, value in multi_value.items():
+        if isinstance(value, list):
+            if value:
+                resolved[key] = value[0]
+        elif isinstance(value, str):
+            resolved[key] = value
+    return resolved
+
+
 def extract_http_tags(event):
     """
     Extracts HTTP facet tags from the triggering event
@@ -327,13 +350,11 @@ def extract_http_tags(event):
     elif request_context and request_context.get("elb"):
         # ALB events have no requestContext.stage; derive the URL from the
         # forwarded host/proto headers and the top-level path.
-        alb_headers = event.get("headers")
-        if not isinstance(alb_headers, dict):
-            alb_headers = {}
+        alb_headers = resolve_alb_request_headers(event)
         host = alb_headers.get("host")
         if host:
             proto = alb_headers.get("x-forwarded-proto", "http")
-            http_tags["http.url"] = f"{proto}://{host}"
+            http_tags["http.url"] = proto + "://" + host
 
         user_agent = alb_headers.get("user-agent")
         if user_agent:

tests/test_tracing.py

@@ -2023,7 +2023,8 @@ def test_create_inferred_span_from_alb_event(self):
         self.assertEqual(span.get_tag("span.kind"), "server")
         self.assertEqual(span.get_tag("http.method"), "GET")
         self.assertEqual(
-            span.get_tag("http.url"), f"http://{self.ALB_HOST}/lambda"
+            span.get_tag("http.url"),
+            "http://%s/lambda" % self.ALB_HOST,
         )
         self.assertEqual(span.get_tag("http.useragent"), self.ALB_USER_AGENT)
         self.assertEqual(span.get_tag("endpoint"), "/lambda")
@@ -2050,10 +2051,32 @@ def test_create_inferred_span_omits_tags_when_headers_missing(self):
         self.assertNotIn("http.method", span.get_tags())
         self.assertNotIn("http.useragent", span.get_tags())
 
-    def test_multivalue_headers_subtype_returns_none(self):
+    def test_multivalue_headers_subtype_emits_inferred_span(self):
         event = self._load_event(self.ALB_MULTIVALUE)
         span = create_inferred_span(event, get_mock_context())
-        self.assertIsNone(span)
+        self.assertIsNotNone(span)
+        self.assertEqual(span.name, "aws.alb")
+        self.assertEqual(span.get_tag("http.method"), "GET")
+        self.assertEqual(
+            span.get_tag("http.url"),
+            "http://%s/lambda" % self.ALB_HOST,
+        )
+        self.assertEqual(span.get_tag("http.useragent"), self.ALB_USER_AGENT)
+
+    @with_trace_propagation_style("datadog")
+    def test_inbound_datadog_context_from_multivalue_headers(self):
+        event = self._load_event(self.ALB_MULTIVALUE)
+        ctx = get_mock_context()
+
+        parent_ctx, source, _ = extract_dd_trace_context(event, ctx)
+        self.assertIsNotNone(parent_ctx)
+        self.assertEqual(parent_ctx.trace_id, 12345)
+        self.assertEqual(parent_ctx.span_id, 67890)
+
+        set_dd_trace_py_root(source, merge_xray_traces=False)
+        span = create_inferred_span(event, ctx)
+        self.assertEqual(span.trace_id, parent_ctx.trace_id)
+        self.assertEqual(span.parent_id, parent_ctx.span_id)
 
     @with_trace_propagation_style("datadog")
     def test_inbound_datadog_context_parents_inferred_span(self):
@@ -2084,6 +2107,28 @@ def test_inbound_w3c_context_extracted_from_alb_event(self):
         self.assertEqual(ctx.trace_id, 0xABCD)
         self.assertEqual(ctx.span_id, 0x4D)
 
+    def test_http_url_uses_https_when_forwarded_proto_is_https(self):
+        event = self._load_event(self.ALB_SAMPLE)
+        event["headers"]["x-forwarded-proto"] = "https"
+
+        span = create_inferred_span(event, get_mock_context())
+
+        self.assertEqual(
+            span.get_tag("http.url"),
+            "https://%s/lambda" % self.ALB_HOST,
+        )
+
+    def test_http_url_excludes_query_string(self):
+        event = self._load_event(self.ALB_SAMPLE)
+
+        span = create_inferred_span(event, get_mock_context())
+
+        self.assertEqual(
+            span.get_tag("http.url"),
+            "http://%s/lambda" % self.ALB_HOST,
+        )
+        self.assertNotIn("query=", span.get_tag("http.url") or "")
+
 
 class _Span(object):
     def __init__(self, service, start, span_type, parent_name=None, tags=None):

tests/test_trigger.py

@@ -445,9 +445,10 @@ def test_extract_trigger_tags_application_load_balancer_multivalue_headers(self)
         assert tags.get("function_trigger.event_source") == "application-load-balancer"
         assert tags.get("http.method") == "GET"
         assert tags.get("http.route") == "/lambda"
-        # multi-value subtype has no single-value ``headers`` map
-        assert "http.url" not in tags
-        assert "http.useragent" not in tags
+        assert tags.get("http.url") == (
+            "http://lambda-alb-123578498.us-east-2.elb.amazonaws.com/lambda"
+        )
+        assert tags.get("http.useragent").startswith("Mozilla/5.0")
 
     def test_extract_trigger_tags_cloudfront(self):
         event_sample_source = "cloudfront"

tests/test_wrapper.py

@@ -1057,7 +1057,36 @@ def lambda_handler(event, context):
         self.assertEqual(inferred.get_tag("http.status_code"), "200")
         self.assertEqual(inferred.get_tag("http.route"), "/lambda")
         self.assertEqual(execution.parent_id, inferred.span_id)
-        self.assertEqual(execution.get_tag("http.status_code"), "200")
+
+    @patch("datadog_lambda.config.Config.trace_enabled", True)
+    @patch("datadog_lambda.config.Config.make_inferred_span", True)
+    @patch("datadog_lambda.config.Config.service", "alb-demo-downstream")
+    def test_wrapper_sets_peer_service_and_dd_resource_key(self):
+        @wrapper.datadog_lambda_wrapper
+        def lambda_handler(event, context):
+            return self._alb_response(200)
+
+        lambda_handler(self.alb_event, get_mock_context())
+
+        inferred = lambda_handler.inferred_span
+
+        self.assertEqual(inferred.get_tag("peer.service"), "alb-demo-downstream")
+        self.assertEqual(
+            inferred.get_tag("dd_resource_key"),
+            "arn:aws:elasticloadbalancing:us-east-2:123456789012:targetgroup/lambda-xyz/123abc",
+        )
+
+    @patch("datadog_lambda.config.Config.trace_enabled", True)
+    @patch("datadog_lambda.config.Config.make_inferred_span", False)
+    def test_wrapper_skips_inferred_alb_span_when_disabled(self):
+        @wrapper.datadog_lambda_wrapper
+        def lambda_handler(event, context):
+            return self._alb_response(200)
+
+        lambda_handler(self.alb_event, get_mock_context())
+
+        self.assertIsNone(lambda_handler.inferred_span)
+        self.assertIsNotNone(lambda_handler.span)
 
     @patch("datadog_lambda.config.Config.trace_enabled", True)
     @patch("datadog_lambda.config.Config.make_inferred_span", True)
@@ -1093,7 +1122,7 @@ def lambda_handler(event, context):
 
     @patch("datadog_lambda.config.Config.trace_enabled", True)
     @patch("datadog_lambda.config.Config.make_inferred_span", True)
-    def test_wrapper_multivalue_alb_event_has_no_inferred_span(self):
+    def test_wrapper_emits_inferred_alb_span_for_multivalue_headers(self):
         with open(
             "tests/event_samples/application-load-balancer-multivalue-headers.json"
         ) as f:
@@ -1105,5 +1134,16 @@ def lambda_handler(event, context):
 
         lambda_handler(event, get_mock_context())
 
-        self.assertIsNone(lambda_handler.inferred_span)
-        self.assertIsNotNone(lambda_handler.span)
+        inferred = lambda_handler.inferred_span
+        execution = lambda_handler.span
+
+        self.assertIsNotNone(inferred)
+        self.assertEqual(inferred.name, "aws.alb")
+        self.assertEqual(inferred.get_tag("http.method"), "GET")
+        self.assertEqual(
+            inferred.get_tag("http.url"),
+            "http://lambda-alb-123578498.us-east-2.elb.amazonaws.com/lambda",
+        )
+        self.assertEqual(inferred.get_tag("http.status_code"), "200")
+        self.assertEqual(inferred.get_tag("http.route"), "/lambda")
+        self.assertEqual(execution.parent_id, inferred.span_id)