fix(aws): Don't crash if event isn't a single dict (#915)

Per https://docs.aws.amazon.com/lambda/latest/dg/python-handler.html, the `event` argument passed to the lambda function handler can be any jsonifiable type - string, int, list, etc - rather than just the dictionary we've previously assumed it to be. (This is particularly relevant for batch requests, which come in as a list of event dictionaries.)

When faced with such an `event`, our current integration crashes, because it tries to run `.get()` on it. This fixes that, by introducing the following behavior:

- If `event` is a list, tag the transaction as a batch and with the batch size.
- If `event` is a list, take the first entry as representative for the purposes of grabbing request data.
- If `event` (or the representative) isn't a dictionary, handle it gracefully and move on without request data.

Author: lobsterkatie

sentry_sdk/integrations/aws_lambda.py

@@ -51,12 +51,12 @@ def sentry_init_error(*args, **kwargs):
 
             exc_info = sys.exc_info()
             if exc_info and all(exc_info):
-                event, hint = event_from_exception(
+                sentry_event, hint = event_from_exception(
                     exc_info,
                     client_options=client.options,
                     mechanism={"type": "aws_lambda", "handled": False},
                 )
-                hub.capture_event(event, hint=hint)
+                hub.capture_event(sentry_event, hint=hint)
 
         return init_error(*args, **kwargs)
 
@@ -65,12 +65,36 @@ def sentry_init_error(*args, **kwargs):
 
 def _wrap_handler(handler):
     # type: (F) -> F
-    def sentry_handler(event, context, *args, **kwargs):
+    def sentry_handler(aws_event, context, *args, **kwargs):
         # type: (Any, Any, *Any, **Any) -> Any
+
+        # Per https://docs.aws.amazon.com/lambda/latest/dg/python-handler.html,
+        # `event` here is *likely* a dictionary, but also might be a number of
+        # other types (str, int, float, None).
+        #
+        # In some cases, it is a list (if the user is batch-invoking their
+        # function, for example), in which case we'll use the first entry as a
+        # representative from which to try pulling request data. (Presumably it
+        # will be the same for all events in the list, since they're all hitting
+        # the lambda in the same request.)
+
+        if isinstance(aws_event, list):
+            request_data = aws_event[0]
+            batch_size = len(aws_event)
+        else:
+            request_data = aws_event
+            batch_size = 1
+
+        if not isinstance(request_data, dict):
+            # If we're not dealing with a dictionary, we won't be able to get
+            # headers, path, http method, etc in any case, so it's fine that
+            # this is empty
+            request_data = {}
+
         hub = Hub.current
         integration = hub.get_integration(AwsLambdaIntegration)
         if integration is None:
-            return handler(event, context, *args, **kwargs)
+            return handler(aws_event, context, *args, **kwargs)
 
         # If an integration is there, a client has to be there.
         client = hub.client  # type: Any
@@ -80,9 +104,14 @@ def sentry_handler(event, context, *args, **kwargs):
             with capture_internal_exceptions():
                 scope.clear_breadcrumbs()
                 scope.add_event_processor(
-                    _make_request_event_processor(event, context, configured_time)
+                    _make_request_event_processor(
+                        request_data, context, configured_time
+                    )
                 )
                 scope.set_tag("aws_region", context.invoked_function_arn.split(":")[3])
+                if batch_size > 1:
+                    scope.set_tag("batch_request", True)
+                    scope.set_tag("batch_size", batch_size)
 
                 timeout_thread = None
                 # Starting the Timeout thread only if the configured time is greater than Timeout warning
@@ -103,21 +132,21 @@ def sentry_handler(event, context, *args, **kwargs):
                     # Starting the thread to raise timeout warning exception
                     timeout_thread.start()
 
-            headers = event.get("headers", {})
+            headers = request_data.get("headers", {})
             transaction = Transaction.continue_from_headers(
                 headers, op="serverless.function", name=context.function_name
             )
             with hub.start_transaction(transaction):
                 try:
-                    return handler(event, context, *args, **kwargs)
+                    return handler(aws_event, context, *args, **kwargs)
                 except Exception:
                     exc_info = sys.exc_info()
-                    event, hint = event_from_exception(
+                    sentry_event, hint = event_from_exception(
                         exc_info,
                         client_options=client.options,
                         mechanism={"type": "aws_lambda", "handled": False},
                     )
-                    hub.capture_event(event, hint=hint)
+                    hub.capture_event(sentry_event, hint=hint)
                     reraise(*exc_info)
                 finally:
                     if timeout_thread:
@@ -255,12 +284,12 @@ def _make_request_event_processor(aws_event, aws_context, configured_timeout):
     # type: (Any, Any, Any) -> EventProcessor
     start_time = datetime.utcnow()
 
-    def event_processor(event, hint, start_time=start_time):
+    def event_processor(sentry_event, hint, start_time=start_time):
         # type: (Event, Hint, datetime) -> Optional[Event]
         remaining_time_in_milis = aws_context.get_remaining_time_in_millis()
         exec_duration = configured_timeout - remaining_time_in_milis
 
-        extra = event.setdefault("extra", {})
+        extra = sentry_event.setdefault("extra", {})
         extra["lambda"] = {
             "function_name": aws_context.function_name,
             "function_version": aws_context.function_version,
@@ -276,7 +305,7 @@ def event_processor(event, hint, start_time=start_time):
             "log_stream": aws_context.log_stream_name,
         }
 
-        request = event.get("request", {})
+        request = sentry_event.get("request", {})
 
         if "httpMethod" in aws_event:
             request["method"] = aws_event["httpMethod"]
@@ -290,7 +319,7 @@ def event_processor(event, hint, start_time=start_time):
             request["headers"] = _filter_headers(aws_event["headers"])
 
         if _should_send_default_pii():
-            user_info = event.setdefault("user", {})
+            user_info = sentry_event.setdefault("user", {})
 
             id = aws_event.get("identity", {}).get("userArn")
             if id is not None:
@@ -308,31 +337,31 @@ def event_processor(event, hint, start_time=start_time):
                 # event. Meaning every body is unstructured to us.
                 request["data"] = AnnotatedValue("", {"rem": [["!raw", "x", 0, 0]]})
 
-        event["request"] = request
+        sentry_event["request"] = request
 
-        return event
+        return sentry_event
 
     return event_processor
 
 
-def _get_url(event, context):
+def _get_url(aws_event, aws_context):
     # type: (Any, Any) -> str
-    path = event.get("path", None)
-    headers = event.get("headers", {})
+    path = aws_event.get("path", None)
+    headers = aws_event.get("headers", {})
     host = headers.get("Host", None)
     proto = headers.get("X-Forwarded-Proto", None)
     if proto and host and path:
         return "{}://{}{}".format(proto, host, path)
-    return "awslambda:///{}".format(context.function_name)
+    return "awslambda:///{}".format(aws_context.function_name)
 
 
-def _get_cloudwatch_logs_url(context, start_time):
+def _get_cloudwatch_logs_url(aws_context, start_time):
     # type: (Any, datetime) -> str
     """
     Generates a CloudWatchLogs console URL based on the context object
 
     Arguments:
-        context {Any} -- context from lambda handler
+        aws_context {Any} -- context from lambda handler
 
     Returns:
         str -- AWS Console URL to logs.
@@ -345,8 +374,8 @@ def _get_cloudwatch_logs_url(context, start_time):
         ";start={start_time};end={end_time}"
     ).format(
         region=environ.get("AWS_REGION"),
-        log_group=context.log_group_name,
-        log_stream=context.log_stream_name,
+        log_group=aws_context.log_group_name,
+        log_stream=aws_context.log_stream_name,
         start_time=(start_time - timedelta(seconds=1)).strftime(formatstring),
         end_time=(datetime.utcnow() + timedelta(seconds=2)).strftime(formatstring),
     )

tests/integrations/aws_lambda/test_aws.py

@@ -35,21 +35,37 @@
 from sentry_sdk.transport import HttpTransport
 
 def event_processor(event):
-    # AWS Lambda truncates the log output to 4kb. If you only need a
-    # subsection of the event, override this function in your test
-    # to print less to logs.
-    return event
+    # AWS Lambda truncates the log output to 4kb, which is small enough to miss
+    # parts of even a single error-event/transaction-envelope pair if considered
+    # in full, so only grab the data we need.
+
+    event_data = {}
+    event_data["contexts"] = {}
+    event_data["contexts"]["trace"] = event.get("contexts", {}).get("trace")
+    event_data["exception"] = event.get("exception")
+    event_data["extra"] = event.get("extra")
+    event_data["level"] = event.get("level")
+    event_data["request"] = event.get("request")
+    event_data["tags"] = event.get("tags")
+    event_data["transaction"] = event.get("transaction")
+
+    return event_data
 
 def envelope_processor(envelope):
+    # AWS Lambda truncates the log output to 4kb, which is small enough to miss
+    # parts of even a single error-event/transaction-envelope pair if considered
+    # in full, so only grab the data we need.
+
     (item,) = envelope.items
     envelope_json = json.loads(item.get_bytes())
 
     envelope_data = {}
-    envelope_data[\"contexts\"] = {}
-    envelope_data[\"type\"] = envelope_json[\"type\"]
-    envelope_data[\"transaction\"] = envelope_json[\"transaction\"]
-    envelope_data[\"contexts\"][\"trace\"] = envelope_json[\"contexts\"][\"trace\"]
-    envelope_data[\"request\"] = envelope_json[\"request\"]
+    envelope_data["contexts"] = {}
+    envelope_data["type"] = envelope_json["type"]
+    envelope_data["transaction"] = envelope_json["transaction"]
+    envelope_data["contexts"]["trace"] = envelope_json["contexts"]["trace"]
+    envelope_data["request"] = envelope_json["request"]
+    envelope_data["tags"] = envelope_json["tags"]
 
     return envelope_data
 
@@ -107,10 +123,15 @@ def inner(code, payload, timeout=30, syntax_check=True):
             syntax_check=syntax_check,
         )
 
+        # for better debugging
+        response["LogResult"] = base64.b64decode(response["LogResult"]).splitlines()
+        response["Payload"] = response["Payload"].read()
+        del response["ResponseMetadata"]
+
         events = []
         envelopes = []
 
-        for line in base64.b64decode(response["LogResult"]).splitlines():
+        for line in response["LogResult"]:
             print("AWS:", line)
             if line.startswith(b"EVENT: "):
                 line = line[len(b"EVENT: ") :]
@@ -362,3 +383,128 @@ def test_handler(event, context):
     assert envelope["contexts"]["trace"]["op"] == "serverless.function"
     assert envelope["transaction"].startswith("test_function_")
     assert envelope["transaction"] in envelope["request"]["url"]
+
+
+@pytest.mark.parametrize(
+    "aws_event, has_request_data, batch_size",
+    [
+        (b"1231", False, 1),
+        (b"11.21", False, 1),
+        (b'"Good dog!"', False, 1),
+        (b"true", False, 1),
+        (
+            b"""
+            [
+                {"good dog": "Maisey"},
+                {"good dog": "Charlie"},
+                {"good dog": "Cory"},
+                {"good dog": "Bodhi"}
+            ]
+            """,
+            False,
+            4,
+        ),
+        (
+            b"""
+            [
+                {
+                    "headers": {
+                        "Host": "dogs.are.great",
+                        "X-Forwarded-Proto": "http"
+                    },
+                    "httpMethod": "GET",
+                    "path": "/tricks/kangaroo",
+                    "queryStringParameters": {
+                        "completed_successfully": "true",
+                        "treat_provided": "true",
+                        "treat_type": "cheese"
+                    },
+                    "dog": "Maisey"
+                },
+                {
+                    "headers": {
+                        "Host": "dogs.are.great",
+                        "X-Forwarded-Proto": "http"
+                    },
+                    "httpMethod": "GET",
+                    "path": "/tricks/kangaroo",
+                    "queryStringParameters": {
+                        "completed_successfully": "true",
+                        "treat_provided": "true",
+                        "treat_type": "cheese"
+                    },
+                    "dog": "Charlie"
+                }
+            ]
+            """,
+            True,
+            2,
+        ),
+    ],
+)
+def test_non_dict_event(
+    run_lambda_function,
+    aws_event,
+    has_request_data,
+    batch_size,
+    DictionaryContaining,  # noqa:N803
+):
+    envelopes, events, response = run_lambda_function(
+        LAMBDA_PRELUDE
+        + dedent(
+            """
+        init_sdk(traces_sample_rate=1.0)
+
+        def test_handler(event, context):
+            raise Exception("More treats, please!")
+        """
+        ),
+        aws_event,
+    )
+
+    assert response["FunctionError"] == "Unhandled"
+
+    error_event = events[0]
+    assert error_event["level"] == "error"
+    assert error_event["contexts"]["trace"]["op"] == "serverless.function"
+
+    function_name = error_event["extra"]["lambda"]["function_name"]
+    assert function_name.startswith("test_function_")
+    assert error_event["transaction"] == function_name
+
+    exception = error_event["exception"]["values"][0]
+    assert exception["type"] == "Exception"
+    assert exception["value"] == "More treats, please!"
+    assert exception["mechanism"]["type"] == "aws_lambda"
+
+    envelope = envelopes[0]
+    assert envelope["type"] == "transaction"
+    assert envelope["contexts"]["trace"] == DictionaryContaining(
+        error_event["contexts"]["trace"]
+    )
+    assert envelope["contexts"]["trace"]["status"] == "internal_error"
+    assert envelope["transaction"] == error_event["transaction"]
+    assert envelope["request"]["url"] == error_event["request"]["url"]
+
+    if has_request_data:
+        request_data = {
+            "headers": {"Host": "dogs.are.great", "X-Forwarded-Proto": "http"},
+            "method": "GET",
+            "url": "http://dogs.are.great/tricks/kangaroo",
+            "query_string": {
+                "completed_successfully": "true",
+                "treat_provided": "true",
+                "treat_type": "cheese",
+            },
+        }
+    else:
+        request_data = {"url": "awslambda:///{}".format(function_name)}
+
+    assert error_event["request"] == request_data
+    assert envelope["request"] == request_data
+
+    if batch_size > 1:
+        assert error_event["tags"]["batch_size"] == batch_size
+        assert error_event["tags"]["batch_request"] is True
+        assert envelope["tags"]["batch_size"] == batch_size
+        assert envelope["tags"]["batch_request"] is True