fix(aws-lambda): Attach user info when streaming spans (#7429)

Add user info onto streamed spans since they don't go through the event
processor.

Author: sentrivana

sentry_sdk/integrations/aws_lambda.py

@@ -45,6 +45,27 @@
 MILLIS_TO_SECONDS = 1000.0
 
 
+def _get_user_from_event(aws_event: "dict[str, Any]") -> "dict[str, Any]":
+    if not isinstance(aws_event, dict):
+        return {}
+
+    identity = aws_event.get("requestContext", {}).get("identity")
+    if identity is None:
+        return {}
+
+    user_info: "dict[str, Any]" = {}
+
+    user_arn = identity.get("userArn")
+    if user_arn is not None:
+        user_info["id"] = user_arn
+
+    ip = identity.get("sourceIp")
+    if ip is not None:
+        user_info["ip_address"] = ip
+
+    return user_info
+
+
 def _wrap_init_error(init_error: "F") -> "F":
     @ensure_integration_enabled(AwsLambdaIntegration, init_error)
     def sentry_init_error(*args: "Any", **kwargs: "Any") -> "Any":
@@ -181,6 +202,17 @@ def sentry_handler(
                     elif should_send_default_pii():
                         additional_attributes["url.query"] = urlencode(qs)
 
+            if not scope._user:
+                if has_data_collection_enabled(client.options):
+                    if client.options["data_collection"]["user_info"]:
+                        user_info = _get_user_from_event(request_data)
+                        if user_info:
+                            scope.set_user(user_info)
+                elif should_send_default_pii():
+                    user_info = _get_user_from_event(request_data)
+                    if user_info:
+                        scope.set_user(user_info)
+
             sampling_context = {
                 "aws_event": aws_event,
                 "aws_context": aws_context,
@@ -440,38 +472,22 @@ def event_processor(
         client_options = sentry_sdk.get_client().options
         if has_data_collection_enabled(client_options):
             if client_options["data_collection"]["user_info"]:
-                user_info = sentry_event.setdefault("user", {})
-
-                identity = aws_event.get("requestContext", {}).get("identity")
-                if identity is None:
-                    identity = {}
-
-                id = identity.get("userArn")
-                if id is not None:
-                    user_info.setdefault("id", id)
-
-                ip = identity.get("sourceIp")
-                if ip is not None:
-                    user_info.setdefault("ip_address", ip)
+                extracted_user = _get_user_from_event(aws_event)
+                if extracted_user:
+                    user_info = sentry_event.setdefault("user", {})
+                    for key, value in extracted_user.items():
+                        user_info.setdefault(key, value)
 
             if "incoming_request" in client_options["data_collection"]["http_bodies"]:
                 if "body" in aws_event:
                     request["data"] = aws_event.get("body", "")
 
         elif should_send_default_pii():
-            user_info = sentry_event.setdefault("user", {})
-
-            identity = aws_event.get("requestContext", {}).get("identity")
-            if identity is None:
-                identity = {}
-
-            id = identity.get("userArn")
-            if id is not None:
-                user_info.setdefault("id", id)
-
-            ip = identity.get("sourceIp")
-            if ip is not None:
-                user_info.setdefault("ip_address", ip)
+            extracted_user = _get_user_from_event(aws_event)
+            if extracted_user:
+                user_info = sentry_event.setdefault("user", {})
+                for key, value in extracted_user.items():
+                    user_info.setdefault(key, value)
 
             if "body" in aws_event:
                 request["data"] = aws_event.get("body", "")

tests/integrations/aws_lambda/lambda_functions_with_embedded_sdk/BasicOkSpanStreamingDataCollectionUserInfoOff/index.py

@@ -0,0 +1,20 @@
+import os
+
+import sentry_sdk
+from sentry_sdk.integrations.aws_lambda import AwsLambdaIntegration
+
+sentry_sdk.init(
+    dsn=os.environ.get("SENTRY_DSN"),
+    traces_sample_rate=1.0,
+    integrations=[AwsLambdaIntegration()],
+    trace_lifecycle="stream",
+    _experiments={
+        "data_collection": {
+            "user_info": False,
+        }
+    },
+)
+
+
+def handler(event, context):
+    return {"event": event}

tests/integrations/aws_lambda/lambda_functions_with_embedded_sdk/BasicOkSpanStreamingDataCollectionUserInfoOn/index.py

@@ -0,0 +1,20 @@
+import os
+
+import sentry_sdk
+from sentry_sdk.integrations.aws_lambda import AwsLambdaIntegration
+
+sentry_sdk.init(
+    dsn=os.environ.get("SENTRY_DSN"),
+    traces_sample_rate=1.0,
+    integrations=[AwsLambdaIntegration()],
+    trace_lifecycle="stream",
+    _experiments={
+        "data_collection": {
+            "user_info": True,
+        }
+    },
+)
+
+
+def handler(event, context):
+    return {"event": event}

tests/integrations/aws_lambda/test_aws_lambda.py

@@ -1144,6 +1144,80 @@ def test_span_streaming_url_query_params_with_data_collection(
     )
 
 
+def test_span_streaming_user_info_with_send_default_pii(
+    lambda_client, test_environment
+):
+    payload = b"""
+        {
+          "resource": "/asd",
+          "path": "/asd",
+          "httpMethod": "GET",
+          "headers": {
+            "Host": "iwsz2c7uwi.execute-api.us-east-1.amazonaws.com",
+            "User-Agent": "custom",
+            "X-Forwarded-Proto": "https"
+          },
+          "queryStringParameters": {
+            "bonkers": "true"
+          },
+          "pathParameters": null,
+          "stageVariables": null,
+          "requestContext": {
+            "identity": {
+                "sourceIp": "213.47.147.207",
+                "userArn": "42"
+            }
+          },
+          "body": null,
+          "isBase64Encoded": false
+        }
+    """
+
+    lambda_client.invoke(
+        FunctionName="BasicOkSpanStreamingPii",
+        Payload=payload,
+    )
+    span_items = test_environment["server"].span_items
+
+    segment_spans = [s for s in span_items if s.get("is_segment")]
+    assert len(segment_spans) == 1
+    attrs = segment_spans[0]["attributes"]
+
+    assert _get_span_attr(attrs, "user.id") == "42"
+
+
+def test_span_streaming_user_info_with_data_collection_user_info_on(
+    lambda_client, test_environment
+):
+    lambda_client.invoke(
+        FunctionName="BasicOkSpanStreamingDataCollectionUserInfoOn",
+        Payload=USER_INFO_PAYLOAD,
+    )
+    span_items = test_environment["server"].span_items
+
+    segment_spans = [s for s in span_items if s.get("is_segment")]
+    assert len(segment_spans) == 1
+    attrs = segment_spans[0]["attributes"]
+
+    assert _get_span_attr(attrs, "user.id") == "42"
+
+
+def test_span_streaming_user_info_with_data_collection_user_info_off(
+    lambda_client, test_environment
+):
+    lambda_client.invoke(
+        FunctionName="BasicOkSpanStreamingDataCollectionUserInfoOff",
+        Payload=USER_INFO_PAYLOAD,
+    )
+    span_items = test_environment["server"].span_items
+
+    segment_spans = [s for s in span_items if s.get("is_segment")]
+    assert len(segment_spans) == 1
+    attrs = segment_spans[0]["attributes"]
+
+    assert "user.id" not in attrs
+
+
 @pytest.mark.parametrize(
     "lambda_function_name",
     ["RaiseErrorPerformanceEnabled", "RaiseErrorPerformanceDisabled"],