Merge branch 'main' into feat/support-eb-dsm-extraction

Author: jeastham1993

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:

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)