diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index cc751e76..8704eeab 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -38,13 +38,39 @@ def my_lambda_handle(event, context): """ +class _NoopDecorator(object): + + def __init__(self, func): + self.func = func + + def __call__(self, *args, **kwargs): + return self.func(*args, **kwargs) + + class _LambdaDecorator(object): """ Decorator to automatically initialize Datadog API client, flush metrics, and extracts/injects trace context. """ + _instance = None + _force_new = False + + def __new__(cls, func): + """ + If the decorator is accidentally applied to multiple functions or + the same function multiple times, only the first one takes effect. + + If _force_new, always return a real decorator, useful for unit tests. + """ + if cls._force_new or cls._instance is None: + cls._instance = super(_LambdaDecorator, cls).__new__(cls) + return cls._instance + else: + return _NoopDecorator(func) + def __init__(self, func): + """Executes when the wrapped function gets wrapped""" self.func = func self.flush_to_log = os.environ.get("DD_FLUSH_TO_LOG", "").lower() == "true" self.logs_injection = ( @@ -59,6 +85,17 @@ def __init__(self, func): patch_all() logger.debug("datadog_lambda_wrapper initialized") + def __call__(self, event, context, **kwargs): + """Executes when the wrapped function gets called""" + self._before(event, context) + try: + return self.func(event, context, **kwargs) + except Exception: + submit_errors_metric(context) + raise + finally: + self._after(event, context) + def _before(self, event, context): set_cold_start() try: @@ -78,15 +115,5 @@ def _after(self, event, context): except Exception: traceback.print_exc() - def __call__(self, event, context, **kwargs): - self._before(event, context) - try: - return self.func(event, context, **kwargs) - except Exception: - submit_errors_metric(context) - raise - finally: - self._after(event, context) - datadog_lambda_wrapper = _LambdaDecorator diff --git a/tests/test_wrapper.py b/tests/test_wrapper.py index ec4b4a32..50f5e563 100644 --- a/tests/test_wrapper.py +++ b/tests/test_wrapper.py @@ -24,6 +24,10 @@ def get_mock_context( class TestDatadogLambdaWrapper(unittest.TestCase): def setUp(self): + # Force @datadog_lambda_wrapper to always create a real + # (not no-op) wrapper. + datadog_lambda_wrapper._force_new = True + patcher = patch("datadog_lambda.metric.lambda_stats") self.mock_metric_lambda_stats = patcher.start() self.addCleanup(patcher.stop) @@ -251,3 +255,28 @@ def lambda_handler(event, context): self.mock_write_metric_point_to_stdout.assert_not_called() del os.environ["DD_ENHANCED_METRICS"] + + def test_only_one_wrapper_in_use(self): + patcher = patch("datadog_lambda.wrapper.submit_invocations_metric") + self.mock_submit_invocations_metric = patcher.start() + self.addCleanup(patcher.stop) + + @datadog_lambda_wrapper + def lambda_handler(event, context): + lambda_metric("test.metric", 100) + + # Turn off _force_new to emulate the nested wrapper scenario, + # the second @datadog_lambda_wrapper should actually be no-op. + datadog_lambda_wrapper._force_new = False + + @datadog_lambda_wrapper + def lambda_handler_wrapper(event, context): + lambda_handler(event, context) + + lambda_event = {} + + lambda_handler_wrapper(lambda_event, get_mock_context()) + + self.mock_patch_all.assert_called_once() + self.mock_wrapper_lambda_stats.flush.assert_called_once() + self.mock_submit_invocations_metric.assert_called_once()