feat: add durable function execution tags to Lambda spans (#728)

lym953 · claude · web-flow · commit 1fafd5edd164 · 2026-02-23T10:13:08.000-05:00
* feat: add durable function execution tags to Lambda spans

Extract DurableExecutionArn from the Lambda event payload and add
durable_function_execution_name and durable_function_execution_id
as span tags, matching the equivalent feature in datadog-lambda-js#730.

Co-Authored-By: Claude Sonnet 4.6 &lt;noreply@anthropic.com&gt;

* Change log level from debug to error

* Format

* fmt

* Fix long line

---------

Co-authored-by: Claude Sonnet 4.6 &lt;noreply@anthropic.com&gt;
diff --git a/datadog_lambda/durable.py b/datadog_lambda/durable.py
@@ -0,0 +1,49 @@
+# Unless explicitly stated otherwise all files in this repository are licensed
+# under the Apache License Version 2.0.
+# This product includes software developed at Datadog (https://www.datadoghq.com/).
+# Copyright 2019 Datadog, Inc.
+import logging
+import re
+
+logger = logging.getLogger(__name__)
+
+
+def _parse_durable_execution_arn(arn):
+    """
+    Parses a DurableExecutionArn to extract execution name and ID.
+    ARN format:
+        arn:aws:lambda:{region}:{account}:function:{func}:{version}/durable-execution/{name}/{id}
+    Returns (execution_name, execution_id) or None if parsing fails.
+    """
+    match = re.search(r"/durable-execution/([^/]+)/([^/]+)$", arn)
+    if not match:
+        return None
+    execution_name, execution_id = match.group(1), match.group(2)
+    if not execution_name or not execution_id:
+        return None
+    return execution_name, execution_id
+
+
+def extract_durable_function_tags(event):
+    """
+    Extracts durable function tags from the Lambda event payload.
+    Returns a dict with durable function tags, or an empty dict if the event
+    is not a durable function invocation.
+    """
+    if not isinstance(event, dict):
+        return {}
+
+    durable_execution_arn = event.get("DurableExecutionArn")
+    if not isinstance(durable_execution_arn, str):
+        return {}
+
+    parsed = _parse_durable_execution_arn(durable_execution_arn)
+    if not parsed:
+        logger.error("Failed to parse DurableExecutionArn: %s", durable_execution_arn)
+        return {}
+
+    execution_name, execution_id = parsed
+    return {
+        "durable_function_execution_name": execution_name,
+        "durable_function_execution_id": execution_id,
+    }
diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py
@@ -1449,6 +1449,7 @@ def create_function_execution_span(
     trace_context_source,
     merge_xray_traces,
     trigger_tags,
+    durable_function_tags=None,
     parent_span=None,
     span_pointers=None,
 ):
@@ -1477,6 +1478,8 @@ def create_function_execution_span(
     if trace_context_source == TraceContextSource.XRAY and merge_xray_traces:
         tags["_dd.parent_source"] = trace_context_source
     tags.update(trigger_tags)
+    if durable_function_tags:
+        tags.update(durable_function_tags)
     tracer.set_tags(_dd_origin)
     # Determine service name based on config and env var
     if config.service:
diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py
@@ -42,6 +42,7 @@
     tracer,
     propagator,
 )
+from datadog_lambda.durable import extract_durable_function_tags
 from datadog_lambda.trigger import (
     extract_trigger_tags,
     extract_http_status_code_tag,
@@ -243,6 +244,7 @@ def _before(self, event, context):
                 submit_invocations_metric(context)
 
             self.trigger_tags = extract_trigger_tags(event, context)
+            self.durable_function_tags = extract_durable_function_tags(event)
             # Extract Datadog trace context and source from incoming requests
             dd_context, trace_context_source, event_source = extract_dd_trace_context(
                 event,
@@ -280,6 +282,7 @@ def _before(self, event, context):
                     trace_context_source=trace_context_source,
                     merge_xray_traces=config.merge_xray_traces,
                     trigger_tags=self.trigger_tags,
+                    durable_function_tags=self.durable_function_tags,
                     parent_span=self.inferred_span,
                     span_pointers=calculate_span_pointers(event_source, event),
                 )
diff --git a/tests/test_durable.py b/tests/test_durable.py
@@ -0,0 +1,91 @@
+# Unless explicitly stated otherwise all files in this repository are licensed
+# under the Apache License Version 2.0.
+# This product includes software developed at Datadog (https://www.datadoghq.com/).
+# Copyright 2019 Datadog, Inc.
+import unittest
+
+from datadog_lambda.durable import (
+    _parse_durable_execution_arn,
+    extract_durable_function_tags,
+)
+
+
+class TestParseDurableExecutionArn(unittest.TestCase):
+    def test_returns_name_and_id_for_valid_arn(self):
+        arn = "arn:aws:lambda:us-east-1:123456789012:function:my-func:$LATEST/durable-execution/order-123/550e8400-e29b-41d4-a716-446655440001"
+        result = _parse_durable_execution_arn(arn)
+        self.assertEqual(result, ("order-123", "550e8400-e29b-41d4-a716-446655440001"))
+
+    def test_returns_none_for_arn_without_durable_execution_marker(self):
+        arn = "arn:aws:lambda:us-east-1:123456789012:function:my-func:$LATEST"
+        result = _parse_durable_execution_arn(arn)
+        self.assertIsNone(result)
+
+    def test_returns_none_for_malformed_arn_with_only_execution_name(self):
+        arn = "arn:aws:lambda:us-east-1:123456789012:function:my-func:$LATEST/durable-execution/order-123"
+        result = _parse_durable_execution_arn(arn)
+        self.assertIsNone(result)
+
+    def test_returns_none_for_malformed_arn_with_empty_execution_name(self):
+        arn = "arn:aws:lambda:us-east-1:123456789012:function:my-func:$LATEST/durable-execution//550e8400-e29b-41d4-a716-446655440002"
+        result = _parse_durable_execution_arn(arn)
+        self.assertIsNone(result)
+
+    def test_returns_none_for_malformed_arn_with_empty_execution_id(self):
+        arn = "arn:aws:lambda:us-east-1:123456789012:function:my-func:$LATEST/durable-execution/order-123/"
+        result = _parse_durable_execution_arn(arn)
+        self.assertIsNone(result)
+
+    def test_works_with_numeric_version_qualifier(self):
+        arn = "arn:aws:lambda:us-east-1:123456789012:function:my-func:1/durable-execution/my-execution/550e8400-e29b-41d4-a716-446655440004"
+        result = _parse_durable_execution_arn(arn)
+        self.assertEqual(
+            result, ("my-execution", "550e8400-e29b-41d4-a716-446655440004")
+        )
+
+
+class TestExtractDurableFunctionTags(unittest.TestCase):
+    def test_extracts_tags_from_event_with_durable_execution_arn(self):
+        event = {
+            "DurableExecutionArn": "arn:aws:lambda:us-east-1:123456789012:function:my-func:1/durable-execution/my-execution/550e8400-e29b-41d4-a716-446655440004",
+            "CheckpointToken": "some-token",
+            "InitialExecutionState": {"Operations": []},
+        }
+        result = extract_durable_function_tags(event)
+        self.assertEqual(
+            result,
+            {
+                "durable_function_execution_name": "my-execution",
+                "durable_function_execution_id": "550e8400-e29b-41d4-a716-446655440004",
+            },
+        )
+
+    def test_returns_empty_dict_for_regular_lambda_event(self):
+        event = {
+            "body": '{"key": "value"}',
+            "headers": {"Content-Type": "application/json"},
+        }
+        result = extract_durable_function_tags(event)
+        self.assertEqual(result, {})
+
+    def test_returns_empty_dict_when_event_is_none(self):
+        result = extract_durable_function_tags(None)
+        self.assertEqual(result, {})
+
+    def test_returns_empty_dict_when_event_is_not_a_dict(self):
+        result = extract_durable_function_tags("not-a-dict")
+        self.assertEqual(result, {})
+
+    def test_returns_empty_dict_when_durable_execution_arn_is_not_a_string(self):
+        event = {"DurableExecutionArn": 12345}
+        result = extract_durable_function_tags(event)
+        self.assertEqual(result, {})
+
+    def test_returns_empty_dict_when_durable_execution_arn_cannot_be_parsed(self):
+        event = {"DurableExecutionArn": "invalid-arn-without-durable-execution-marker"}
+        result = extract_durable_function_tags(event)
+        self.assertEqual(result, {})
+
+    def test_returns_empty_dict_when_event_is_empty(self):
+        result = extract_durable_function_tags({})
+        self.assertEqual(result, {})
