diff --git a/bugsnag/__init__.py b/bugsnag/__init__.py index 5a74a852..9c023b2f 100644 --- a/bugsnag/__init__.py +++ b/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') diff --git a/bugsnag/client.py b/bugsnag/client.py index 391d1784..961f63b8 100644 --- a/bugsnag/client.py +++ b/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) diff --git a/bugsnag/legacy.py b/bugsnag/legacy.py index df793db4..018e6735 100644 --- a/bugsnag/legacy.py +++ b/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) diff --git a/tests/test_client.py b/tests/test_client.py index bfc63dd8..cb392feb 100644 --- a/tests/test_client.py +++ b/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' diff --git a/tests/utils.py b/tests/utils.py index 24b01b90..d17e138e 100644 --- a/tests/utils.py +++ b/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()