Add aws_lambda_handler decorator to Client

bugsnag/__init__.py

@@ -15,7 +15,8 @@
                             auto_notify_exc_info, logger, leave_breadcrumb,
                             add_on_breadcrumb, remove_on_breadcrumb,
                             add_feature_flag, add_feature_flags,
-                            clear_feature_flag, clear_feature_flags)
+                            clear_feature_flag, clear_feature_flags,
+                            aws_lambda_handler)
 
 __all__ = ('Client', 'Event', 'Configuration', 'RequestConfiguration',
            'configuration', 'configure', 'configure_request',
@@ -25,4 +26,5 @@
            'BreadcrumbType', 'Breadcrumb', 'Breadcrumbs',
            'OnBreadcrumbCallback', 'leave_breadcrumb', 'add_on_breadcrumb',
            'remove_on_breadcrumb', 'FeatureFlag', 'add_feature_flag',
-           'add_feature_flags', 'clear_feature_flag', 'clear_feature_flags')
+           'add_feature_flags', 'clear_feature_flag', 'clear_feature_flags',
+           'aws_lambda_handler')

bugsnag/client.py

@@ -2,9 +2,9 @@
 import sys
 import threading
 import warnings
+import functools
 
 from datetime import datetime, timezone
-from functools import wraps
 from typing import Union, Tuple, Callable, Optional, List, Type, Dict, Any
 
 from bugsnag.breadcrumbs import (
@@ -366,6 +366,75 @@ def block_until_no_requests():
 
             raise Exception("flush timed out after %dms" % timeout_ms)
 
+    def add_metadata_tab(self, tab_name: str, data: Dict[str, Any]) -> None:
+        metadata = RequestConfiguration.get_instance().metadata
+
+        if tab_name not in metadata:
+            metadata[tab_name] = {}
+
+        metadata[tab_name].update(data)
+
+    def aws_lambda_handler(
+        self,
+        real_handler: Optional[Callable] = None,
+        flush_timeout_ms: int = 2000,
+    ) -> Callable:
+        # handle being called with just 'flush_timeout_ms'
+        if real_handler is None:
+            return functools.partial(
+                self.aws_lambda_handler,
+                flush_timeout_ms=flush_timeout_ms,
+            )
+
+        # attributes from the aws context that we want to capture as metadata
+        # the context is an instance of LambdaContext, which isn't iterable and
+        # so can't be added to metadata as-is
+        aws_context_attributes = [
+            'function_name',
+            'function_version',
+            'invoked_function_arn',
+            'memory_limit_in_mb',
+            'aws_request_id',
+            'log_group_name',
+            'log_stream_name',
+            'identity',
+            'client_context',
+        ]
+
+        @functools.wraps(real_handler)
+        def wrapped_handler(aws_event, aws_context):
+            try:
+                aws_context_metadata = {
+                    attribute:
+                        getattr(aws_context, attribute, None)
+                        for attribute in aws_context_attributes
+                }
+
+                self.add_metadata_tab('AWS Lambda Event', aws_event)
+                self.add_metadata_tab(
+                    'AWS Lambda Context',
+                    aws_context_metadata
+                )
+
+                if self.configuration.auto_capture_sessions:
+                    self.session_tracker.start_session()
+
+                return real_handler(aws_event, aws_context)
+            except Exception as exception:
+                if self.configuration.auto_notify:
+                    self.notify(exception)
+
+                raise
+            finally:
+                try:
+                    self.flush(flush_timeout_ms)
+                except Exception as exception:
+                    warnings.warn(
+                        'Delivery may be unsuccessful: ' + str(exception)
+                    )
+
+        return wrapped_handler
+
 
 class ClientContext:
     def __init__(self, client,
@@ -378,7 +447,7 @@ def __init__(self, client,
         self.exception_types = exception_types or (Exception,)
 
     def __call__(self, function: Callable):
-        @wraps(function)
+        @functools.wraps(function)
         def decorate(*args, **kwargs):
             try:
                 return function(*args, **kwargs)

bugsnag/legacy.py

@@ -1,4 +1,4 @@
-from typing import Dict, Any, Tuple, Type, Optional, Union, List
+from typing import Dict, Any, Tuple, Type, Optional, Union, List, Callable
 import types
 import sys
 
@@ -38,11 +38,7 @@ def add_metadata_tab(tab_name: str, data: Dict[str, Any]):
 
     bugsnag.add_metadata_tab("user", {"id": "1", "name": "Conrad"})
     """
-    metadata = RequestConfiguration.get_instance().metadata
-    if tab_name not in metadata:
-        metadata[tab_name] = {}
-
-    metadata[tab_name].update(data)
+    default_client.add_metadata_tab(tab_name, data)
 
 
 def clear_request_config():
@@ -180,3 +176,10 @@ def clear_feature_flag(name: Union[str, bytes]) -> None:
 
 def clear_feature_flags() -> None:
     default_client.clear_feature_flags()
+
+
+def aws_lambda_handler(
+    real_handler: Optional[Callable] = None,
+    flush_timeout_ms: int = 2000,
+) -> Callable:
+    return default_client.aws_lambda_handler(real_handler, flush_timeout_ms)

tests/test_client.py

@@ -1822,6 +1822,100 @@ def flush_request_queue():
         # the thread should have stopped before flush could exit
         assert not thread.is_alive()
 
+    def test_aws_lambda_handler_decorator(self):
+        aws_lambda_context = LambdaContext(function_name='abcdef')
+
+        @self.client.aws_lambda_handler
+        def my_handler(event, context):
+            assert event == {'a': 1}
+            assert context == aws_lambda_context
+
+            raise Exception('oh dear')
+
+        with pytest.raises(Exception) as exception:
+            my_handler({'a': 1}, aws_lambda_context)
+
+            assert str(exception) == 'Exception: oh dear'
+
+        assert self.sent_report_count == 1
+        assert self.sent_session_count == 1
+
+        payload = self.server.events_received[0]['json_body']
+        event = payload['events'][0]
+
+        assert event['exceptions'][0]['message'] == 'oh dear'
+        assert event['metaData']['AWS Lambda Event'] == {'a': 1}
+        assert event['metaData']['AWS Lambda Context'] == {
+            'function_name': 'abcdef',
+            'function_version': 'function_version',
+            'invoked_function_arn': 'invoked_function_arn',
+            'memory_limit_in_mb': 'memory_limit_in_mb',
+            'aws_request_id': 'aws_request_id',
+            'log_group_name': 'log_group_name',
+            'log_stream_name': 'log_stream_name',
+            'identity': 'identity',
+            'client_context': 'client_context',
+        }
+
+    def test_aws_lambda_handler_decorator_accepts_flush_timeout(self):
+        aws_lambda_context = LambdaContext(function_version='$LATEST')
+
+        @self.client.aws_lambda_handler(flush_timeout_ms=1000)
+        def my_handler(event, context):
+            assert event == {'z': 9}
+            assert context == aws_lambda_context
+
+            raise Exception('oh dear')
+
+        with pytest.raises(Exception) as exception:
+            my_handler({'z': 9}, aws_lambda_context)
+
+            assert str(exception) == 'Exception: oh dear'
+
+        assert self.sent_report_count == 1
+        assert self.sent_session_count == 1
+
+        payload = self.server.events_received[0]['json_body']
+        event = payload['events'][0]
+
+        assert event['exceptions'][0]['message'] == 'oh dear'
+        assert event['metaData']['AWS Lambda Event'] == {'z': 9}
+        assert event['metaData']['AWS Lambda Context'] == {
+            'function_name': 'function_name',
+            'function_version': '$LATEST',
+            'invoked_function_arn': 'invoked_function_arn',
+            'memory_limit_in_mb': 'memory_limit_in_mb',
+            'aws_request_id': 'aws_request_id',
+            'log_group_name': 'log_group_name',
+            'log_stream_name': 'log_stream_name',
+            'identity': 'identity',
+            'client_context': 'client_context',
+        }
+
+    def test_aws_lambda_handler_decorator_warns_after_timeout(self):
+        aws_lambda_context = LambdaContext()
+        client = Client(delivery=QueueingDelivery(), api_key='abc')
+
+        @client.aws_lambda_handler(flush_timeout_ms=50)
+        def my_handler(event, context):
+            assert event == {'z': 9}
+            assert context == aws_lambda_context
+
+            raise Exception('oh dear')
+
+        with pytest.warns(UserWarning) as warnings:
+            with pytest.raises(Exception) as exception:
+                my_handler({'z': 9}, aws_lambda_context)
+
+                assert str(exception) == 'Exception: oh dear'
+
+            assert len(warnings) == 1
+            assert warnings[0].message.args[0] == \
+                'Delivery may be unsuccessful: flush timed out after 50ms'
+
+        assert self.sent_report_count == 0
+        assert self.sent_session_count == 0
+
 
 @pytest.mark.parametrize("metadata,type", [
     (1234, 'int'),
@@ -1852,3 +1946,28 @@ def test_breadcrumb_metadata_is_coerced_to_dict(metadata, type):
     assert breadcrumb.metadata == {}
     assert breadcrumb.type == BreadcrumbType.MANUAL
     assert is_valid_timestamp(breadcrumb.timestamp)
+
+
+class LambdaContext:
+    def __init__(
+        self,
+        function_name='function_name',
+        function_version='function_version',
+        invoked_function_arn='invoked_function_arn',
+        memory_limit_in_mb='memory_limit_in_mb',
+        aws_request_id='aws_request_id',
+        log_group_name='log_group_name',
+        log_stream_name='log_stream_name',
+        identity='identity',
+        client_context='client_context',
+    ):
+        self.function_name = function_name
+        self.function_version = function_version
+        self.invoked_function_arn = invoked_function_arn
+        self.memory_limit_in_mb = memory_limit_in_mb
+        self.aws_request_id = aws_request_id
+        self.log_group_name = log_group_name
+        self.log_stream_name = log_stream_name
+        self.identity = identity
+        self.client_context = client_context
+        self.another_attribute = 'another_attribute'

tests/utils.py

@@ -8,6 +8,7 @@
 
 import bugsnag
 from bugsnag.delivery import Delivery
+from bugsnag.configuration import RequestConfiguration
 
 
 try:
@@ -32,6 +33,7 @@ def setUp(self):
         self.server.sessions_received = []
 
     def tearDown(self):
+        RequestConfiguration.get_instance().clear()
         previous_client = bugsnag.legacy.default_client
         previous_client.uninstall_sys_hook()