emit enhanced error metric and create span when an exception is raised outside of the handler function

Author: TalUsvyatsky

datadog_lambda/handler.py

@@ -7,6 +7,9 @@
 from importlib import import_module
 
 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.module_name import modify_module_name
 
@@ -15,6 +18,27 @@ class HandlerError(Exception):
     pass
 
 
+class _ErrorOutsideHandlerDecorator(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(
@@ -27,5 +51,11 @@ class HandlerError(Exception):
 
 (mod_name, handler_name) = parts
 modified_mod_name = modify_module_name(mod_name)
-handler_module = import_module(modified_mod_name)
-handler = datadog_lambda_wrapper(getattr(handler_module, handler_name))
+
+try:
+    start_time_ns = time.time_ns()
+    handler_module = import_module(modified_mod_name)
+    handler_func = getattr(handler_module, handler_name)
+    handler = datadog_lambda_wrapper(handler_func)
+except Exception as e:
+    handler = _ErrorOutsideHandlerDecorator(e, modified_mod_name, start_time_ns)

datadog_lambda/metric.py

@@ -100,7 +100,7 @@ def submit_enhanced_metric(metric_name, lambda_context):
 
     Args:
         metric_name (str): metric name w/o enhanced prefix i.e. "invocations" or "errors"
-        lambda_context (dict): Lambda context dict passed to the function by AWS
+        lambda_context (object): Lambda context dict passed to the function by AWS
     """
     if not enhanced_metrics_enabled:
         logger.debug(
@@ -118,7 +118,7 @@ def submit_invocations_metric(lambda_context):
     """Increment aws.lambda.enhanced.invocations by 1, applying runtime, layer, and cold_start tags
 
     Args:
-        lambda_context (dict): Lambda context dict passed to the function by AWS
+        lambda_context (object): Lambda context dict passed to the function by AWS
     """
     submit_enhanced_metric("invocations", lambda_context)
 
@@ -127,6 +127,6 @@ def submit_errors_metric(lambda_context):
     """Increment aws.lambda.enhanced.errors by 1, applying runtime, layer, and cold_start tags
 
     Args:
-        lambda_context (dict): Lambda context dict passed to the function by AWS
+        lambda_context (object): Lambda context dict passed to the function by AWS
     """
     submit_enhanced_metric("errors", lambda_context)

datadog_lambda/tracing.py

@@ -6,6 +6,7 @@
 import logging
 import os
 import base64
+import traceback
 import ujson as json
 from datetime import datetime, timezone
 from typing import Optional, Dict
@@ -1320,3 +1321,38 @@ def is_async(span: Span) -> bool:
                 e,
             )
             return False
+
+
+def emit_telemetry_on_exception_outside_of_handler(
+    context, exception, resource_name, start_time_ns
+):
+    """
+    Emit an enhanced error metric and create a span for exceptions occuring outside of the handler
+    """
+    submit_errors_metric(context)
+
+    span = tracer.trace(
+        "aws.lambda",
+        service="aws.lambda",
+        resource=resource_name,
+        span_type="serverless",
+    )
+    span.start_ns = start_time_ns
+    tags = {
+        "error.status": 500,
+        "error.type": type(exception).__name__,
+        "error.message": exception,
+        "error.stack": "".join(
+            traceback.format_exception(
+                type(exception), exception, exception.__traceback__
+            )
+        ),
+        "resource_names": resource_name,
+        "resource.name": resource_name,
+        "operation_name": "aws.lambda",
+        "status": "error",
+        "request_id": context.aws_request_id,
+    }
+    span.set_tags(tags)
+    span.error = 1
+    span.finish()

tests/test_handler.py

@@ -0,0 +1,80 @@
+import os
+import sys
+import unittest
+from unittest.mock import patch
+
+from tests.utils import get_mock_context
+
+
+class TestHandler(unittest.TestCase):
+    def tearDown(self):
+        for mod in sys.modules.copy():
+            if mod.startswith("datadog_lambda.handler"):
+                del sys.modules[mod]
+
+    def test_dd_lambda_handler_env_var_none(self):
+        with self.assertRaises(Exception) as context:
+            import datadog_lambda.handler as handler
+
+            assert context.exception == handler.HandlerError(
+                "DD_LAMBDA_HANDLER is not defined. Can't use prebuilt datadog handler"
+            )
+
+    @patch.dict(os.environ, {"DD_LAMBDA_HANDLER": "malformed"}, clear=True)
+    def test_dd_lambda_handler_env_var_malformed(self):
+        with self.assertRaises(Exception) as context:
+            import datadog_lambda.handler as handler
+
+            assert context.exception == handler.HandlerError(
+                "Value malformed for DD_LAMBDA_HANDLER has invalid format."
+            )
+
+    @patch.dict(os.environ, {"DD_LAMBDA_HANDLER": "nonsense.nonsense"}, clear=True)
+    @patch("datadog_lambda.tracing.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:
+            import datadog_lambda.handler
+
+            lambda_context = get_mock_context()
+            datadog_lambda.handler.handler.__call__(None, lambda_context)
+        mock_emit_telemetry.assert_called_once_with(
+            lambda_context, test_context.exception, "nonsense", 42
+        )
+
+    @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("time.time_ns", return_value=42)
+    def test_exception_getting_handler_func(
+        self, mock_time, mock_emit_telemetry, mock_import
+    ):
+        with self.assertRaises(AttributeError) as test_context:
+            import datadog_lambda.handler
+
+            lambda_context = get_mock_context()
+            datadog_lambda.handler.handler.__call__(None, lambda_context)
+        mock_emit_telemetry.assert_called_once_with(
+            lambda_context, test_context.exception, "nonsense", 42
+        )
+
+    @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("time.time_ns", return_value=42)
+    @patch("datadog_lambda.wrapper.datadog_lambda_wrapper")
+    def test_handler_success(
+        self, mock_lambda_wrapper, mock_time, mock_emit_telemetry, mock_import
+    ):
+        def nonsense():
+            pass
+
+        mock_import.nonsense.return_value = nonsense
+
+        import datadog_lambda.handler
+
+        lambda_context = get_mock_context()
+        datadog_lambda.handler.handler.__call__(None, lambda_context)
+
+        mock_emit_telemetry.assert_not_called()
+        mock_lambda_wrapper.assert_called_once_with(mock_import().nonsense)

tests/test_tracing.py

@@ -1,6 +1,7 @@
 import copy
 import functools
 import json
+import traceback
 import pytest
 import os
 import unittest
@@ -36,6 +37,7 @@
     determine_service_name,
     service_mapping as global_service_mapping,
     propagator,
+    emit_telemetry_on_exception_outside_of_handler,
 )
 from datadog_lambda.trigger import EventTypes
 
@@ -1999,3 +2001,48 @@ def test_deterministic_m5_hash__always_leading_with_zero(self):
             # Leading zeros will be omitted, so only test for full 64 bits present
             if len(result_in_binary) == 66:  # "0b" + 64 bits.
                 self.assertTrue(result_in_binary.startswith("0b0"))
+
+
+class TestExceptionOutsideHandler(unittest.TestCase):
+    @patch("datadog_lambda.tracing.submit_errors_metric")
+    def test_exception_outside_handler(self, mock_submit_errors_metric):
+        fake_error = ValueError("Some error message")
+        resource_name = "my_handler"
+        span_type = "aws.lambda"
+        mock_span = Mock()
+        context = get_mock_context()
+        with patch(
+            "datadog_lambda.tracing.tracer.trace", return_value=mock_span
+        ) as mock_trace:
+            emit_telemetry_on_exception_outside_of_handler(
+                context, fake_error, resource_name, 42
+            )
+
+        mock_submit_errors_metric.assert_called_once_with(context)
+
+        mock_trace.assert_called_once_with(
+            span_type,
+            service="aws.lambda",
+            resource=resource_name,
+            span_type="serverless",
+        )
+        mock_span.set_tags.assert_called_once_with(
+            {
+                "error.status": 500,
+                "error.type": "ValueError",
+                "error.message": fake_error,
+                "error.stack": "".join(
+                    traceback.format_exception(
+                        type(fake_error), fake_error, fake_error.__traceback__
+                    )
+                ),
+                "resource_names": resource_name,
+                "resource.name": resource_name,
+                "operation_name": span_type,
+                "status": "error",
+                "request_id": context.aws_request_id,
+            }
+        )
+        mock_span.finish.assert_called_once()
+        assert mock_span.error == 1
+        assert mock_span.start_ns == 42