diff --git a/datadog_lambda/wrapper.py b/datadog_lambda/wrapper.py index d955480c..73057983 100644 --- a/datadog_lambda/wrapper.py +++ b/datadog_lambda/wrapper.py @@ -9,7 +9,6 @@ from importlib import import_module from time import time_ns -from ddtrace.internal._exceptions import BlockingException from datadog_lambda.extension import should_use_extension, flush_extension from datadog_lambda.cold_start import ( set_cold_start, @@ -191,8 +190,6 @@ def __call__(self, event, context, **kwargs): if self.blocking_response: return self.blocking_response self.response = self.func(event, context, **kwargs) - except BlockingException: - self.blocking_response = get_asm_blocked_response(self.event_source) except Exception: from datadog_lambda.metric import submit_errors_metric @@ -201,6 +198,9 @@ def __call__(self, event, context, **kwargs): if self.span: self.span.set_traceback() raise + except BaseException as e: + if "BlockingException" in type(e).__name__: + self.blocking_response = get_asm_blocked_response(self.event_source) finally: self._after(event, context) if self.blocking_response: diff --git a/tests/test_wrapper.py b/tests/test_wrapper.py index ba2cce5e..35b5b9fa 100644 --- a/tests/test_wrapper.py +++ b/tests/test_wrapper.py @@ -724,11 +724,14 @@ def test_blocking_before(self): def test_blocking_during(self): self.mock_get_asm_blocking_response.return_value = None + class BlockingException(BaseException): + pass + def lambda_handler(event, context): self.mock_get_asm_blocking_response.return_value = ( self.fake_blocking_response ) - raise wrapper.BlockingException() + raise BlockingException() lambda_handler = wrapper.datadog_lambda_wrapper(lambda_handler)