diff --git a/datadog_lambda/durable.py b/datadog_lambda/durable.py new file mode 100644 index 00000000..e9443f92 --- /dev/null +++ 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 index e7dca1f3..bd138976 100644 --- 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 index b2f83f13..aff4a44e 100644 --- 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 new file mode 100644 index 00000000..60914934 --- /dev/null +++ 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, {})