diff --git a/datadog_lambda/durable.py b/datadog_lambda/durable.py index e9443f92..3ae0ebc6 100644 --- a/datadog_lambda/durable.py +++ b/datadog_lambda/durable.py @@ -2,9 +2,12 @@ # under the Apache License Version 2.0. # This product includes software developed at Datadog (https://www.datadoghq.com/). # Copyright 2019 Datadog, Inc. +import functools import logging import re +from ddtrace import tracer + logger = logging.getLogger(__name__) @@ -47,3 +50,27 @@ def extract_durable_function_tags(event): "durable_function_execution_name": execution_name, "durable_function_execution_id": execution_id, } + + +def durable_execution(func): + """ + Decorator for AWS Lambda durable execution orchestration functions. + Sets the durable_function_first_invocation tag on the current span + based on whether this is the first invocation (not replaying history). + """ + + @functools.wraps(func) + def wrapper(context, *args, **kwargs): + try: + is_first_invocation = not context.state.is_replaying() + span = tracer.current_span() + if span: + span.set_tag( + "durable_function_first_invocation", + str(is_first_invocation).lower(), + ) + except Exception: + logger.debug("Failed to set durable_function_first_invocation tag") + return func(context, *args, **kwargs) + + return wrapper diff --git a/tests/test_durable.py b/tests/test_durable.py index 60914934..3f3f5c71 100644 --- a/tests/test_durable.py +++ b/tests/test_durable.py @@ -3,9 +3,11 @@ # This product includes software developed at Datadog (https://www.datadoghq.com/). # Copyright 2019 Datadog, Inc. import unittest +from unittest.mock import MagicMock, patch from datadog_lambda.durable import ( _parse_durable_execution_arn, + durable_execution, extract_durable_function_tags, ) @@ -89,3 +91,76 @@ def test_returns_empty_dict_when_durable_execution_arn_cannot_be_parsed(self): def test_returns_empty_dict_when_event_is_empty(self): result = extract_durable_function_tags({}) self.assertEqual(result, {}) + + +class TestDurableExecution(unittest.TestCase): + def _make_durable_context(self, is_replaying): + ctx = MagicMock() + ctx.state.is_replaying.return_value = is_replaying + return ctx + + def test_sets_first_invocation_true_when_not_replaying(self): + mock_span = MagicMock() + with patch("datadog_lambda.durable.tracer") as mock_tracer: + mock_tracer.current_span.return_value = mock_span + ctx = self._make_durable_context(is_replaying=False) + + @durable_execution + def handler(context): + return "result" + + result = handler(ctx) + + self.assertEqual(result, "result") + mock_span.set_tag.assert_called_once_with( + "durable_function_first_invocation", "true" + ) + + def test_sets_first_invocation_false_when_replaying(self): + mock_span = MagicMock() + with patch("datadog_lambda.durable.tracer") as mock_tracer: + mock_tracer.current_span.return_value = mock_span + ctx = self._make_durable_context(is_replaying=True) + + @durable_execution + def handler(context): + return "result" + + result = handler(ctx) + + self.assertEqual(result, "result") + mock_span.set_tag.assert_called_once_with( + "durable_function_first_invocation", "false" + ) + + def test_does_not_set_tag_when_no_active_span(self): + with patch("datadog_lambda.durable.tracer") as mock_tracer: + mock_tracer.current_span.return_value = None + ctx = self._make_durable_context(is_replaying=False) + + @durable_execution + def handler(context): + return "result" + + result = handler(ctx) + + self.assertEqual(result, "result") + + def test_does_not_raise_when_context_has_no_state(self): + ctx = MagicMock(spec=[]) # no attributes + with patch("datadog_lambda.durable.tracer"): + + @durable_execution + def handler(context): + return "result" + + result = handler(ctx) + + self.assertEqual(result, "result") + + def test_preserves_function_name(self): + @durable_execution + def my_handler(context): + pass + + self.assertEqual(my_handler.__name__, "my_handler")