move fallback handler to wrapper.py

Author: TalUsvyatsky

datadog_lambda/handler.py

@@ -9,36 +9,14 @@
 import os
 import time
 
-from datadog_lambda.tracing import emit_telemetry_on_exception_outside_of_handler
-from datadog_lambda.wrapper import datadog_lambda_wrapper
+from datadog_lambda.wrapper import datadog_lambda_wrapper, error_fallback_handler
 from datadog_lambda.module_name import modify_module_name
 
 
 class HandlerError(Exception):
     pass
 
 
-class _ErrorFallbackHandler(object):
-    """
-    Decorator for when an exception occurs outside of the handler function.
-    Emits telemetry and re-raises the exception.
-    """
-
-    def __init__(self, exception, modified_mod_name, start_time_ns):
-        self.exception = exception
-        self.modified_mod_name = modified_mod_name
-        self.start_time_ns = start_time_ns
-
-    def __call__(self, event, context, **kwargs):
-        emit_telemetry_on_exception_outside_of_handler(
-            context,
-            self.exception,
-            self.modified_mod_name,
-            self.start_time_ns,
-        )
-        raise self.exception
-
-
 path = os.environ.get("DD_LAMBDA_HANDLER", None)
 if path is None:
     raise HandlerError(
@@ -58,4 +36,4 @@ def __call__(self, event, context, **kwargs):
     handler_func = getattr(handler_module, handler_name)
     handler = datadog_lambda_wrapper(handler_func)
 except Exception as e:
-    handler = _ErrorFallbackHandler(e, modified_mod_name, start_time_ns)
+    handler = error_fallback_handler(e, modified_mod_name, start_time_ns)

datadog_lambda/wrapper.py

@@ -44,6 +44,7 @@
     is_authorizer_response,
     tracer,
     propagator,
+    emit_telemetry_on_exception_outside_of_handler,
 )
 from datadog_lambda.trigger import (
     extract_trigger_tags,
@@ -382,9 +383,31 @@ def _after(self, event, context):
             logger.error(format_err_with_traceback(e))
 
 
+class _ErrorFallbackHandler(object):
+    """
+    Fallback handler for when an exception occurs outside of the handler function.
+    Emits telemetry and re-raises the exception.
+    """
+
+    def __init__(self, exception, modified_mod_name, start_time_ns):
+        self.exception = exception
+        self.modified_mod_name = modified_mod_name
+        self.start_time_ns = start_time_ns
+
+    def __call__(self, event, context, **kwargs):
+        emit_telemetry_on_exception_outside_of_handler(
+            context,
+            self.exception,
+            self.modified_mod_name,
+            self.start_time_ns,
+        )
+        raise self.exception
+
+
 def format_err_with_traceback(e):
     tb = traceback.format_exc().replace("\n", "\r")
     return f"Error {e}. Traceback: {tb}"
 
 
 datadog_lambda_wrapper = _LambdaDecorator
+error_fallback_handler = _ErrorFallbackHandler

tests/test_handler.py

@@ -30,7 +30,7 @@ def test_dd_lambda_handler_env_var_malformed(self):
             )
 
     @patch.dict(os.environ, {"DD_LAMBDA_HANDLER": "nonsense.nonsense"}, clear=True)
-    @patch("datadog_lambda.tracing.emit_telemetry_on_exception_outside_of_handler")
+    @patch("datadog_lambda.wrapper.emit_telemetry_on_exception_outside_of_handler")
     @patch("time.time_ns", return_value=42)
     def test_exception_importing_module(self, mock_time, mock_emit_telemetry):
         with self.assertRaises(ModuleNotFoundError) as test_context:
@@ -44,7 +44,7 @@ def test_exception_importing_module(self, mock_time, mock_emit_telemetry):
 
     @patch.dict(os.environ, {"DD_LAMBDA_HANDLER": "nonsense.nonsense"}, clear=True)
     @patch("importlib.import_module", return_value=None)
-    @patch("datadog_lambda.tracing.emit_telemetry_on_exception_outside_of_handler")
+    @patch("datadog_lambda.wrapper.emit_telemetry_on_exception_outside_of_handler")
     @patch("time.time_ns", return_value=42)
     def test_exception_getting_handler_func(
         self, mock_time, mock_emit_telemetry, mock_import
@@ -60,7 +60,7 @@ def test_exception_getting_handler_func(
 
     @patch.dict(os.environ, {"DD_LAMBDA_HANDLER": "nonsense.nonsense"}, clear=True)
     @patch("importlib.import_module")
-    @patch("datadog_lambda.tracing.emit_telemetry_on_exception_outside_of_handler")
+    @patch("datadog_lambda.wrapper.emit_telemetry_on_exception_outside_of_handler")
     @patch("time.time_ns", return_value=42)
     @patch("datadog_lambda.wrapper.datadog_lambda_wrapper")
     def test_handler_success(