feat(aws-lambda): Add span streaming support to AWS Lambda integration

sentry_sdk/integrations/aws_lambda.py

@@ -6,14 +6,21 @@
 from datetime import datetime, timedelta, timezone
 from os import environ
 from typing import TYPE_CHECKING
+from urllib.parse import urlencode
 
 import sentry_sdk
 from sentry_sdk.api import continue_trace
 from sentry_sdk.consts import OP
 from sentry_sdk.integrations import Integration
 from sentry_sdk.integrations._wsgi_common import _filter_headers
-from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.integrations.cloud_resource_context import (
+    CLOUD_PLATFORM,
+    CLOUD_PROVIDER,
+)
+from sentry_sdk.scope import Scope, should_send_default_pii
+from sentry_sdk.traces import SegmentSource
 from sentry_sdk.tracing import TransactionSource
+from sentry_sdk.tracing_utils import has_span_streaming_enabled
 from sentry_sdk.utils import (
     AnnotatedValue,
     TimeoutThread,
@@ -101,6 +108,7 @@ def sentry_handler(
             request_data = {}
 
         configured_time = aws_context.get_remaining_time_in_millis()
+        aws_region = aws_context.invoked_function_arn.split(":")[3]
 
         with sentry_sdk.isolation_scope() as scope:
             timeout_thread = None
@@ -111,9 +119,7 @@ def sentry_handler(
                         request_data, aws_context, configured_time
                     )
                 )
-                scope.set_tag(
-                    "aws_region", aws_context.invoked_function_arn.split(":")[3]
-                )
+                scope.set_tag("aws_region", aws_region)
                 if batch_size > 1:
                     scope.set_tag("batch_request", True)
                     scope.set_tag("batch_size", batch_size)
@@ -144,20 +150,71 @@ def sentry_handler(
             if not isinstance(headers, dict):
                 headers = {}
 
-            transaction = continue_trace(
-                headers,
-                op=OP.FUNCTION_AWS,
-                name=aws_context.function_name,
-                source=TransactionSource.COMPONENT,
-                origin=AwsLambdaIntegration.origin,
-            )
-            with sentry_sdk.start_transaction(
-                transaction,
-                custom_sampling_context={
-                    "aws_event": aws_event,
-                    "aws_context": aws_context,
-                },
-            ):
+            header_attributes: "dict[str, Any]" = {}
+            for header, header_value in _filter_headers(
+                headers, use_annotated_value=False
+            ).items():
+                header_attributes[f"http.request.header.{header.lower()}"] = (
+                    header_value
+                )
+
+            additional_attributes: "dict[str, Any]" = {}
+            if "httpMethod" in request_data:
+                additional_attributes["http.request.method"] = request_data[
+                    "httpMethod"
+                ]
+
+            if should_send_default_pii() and "queryStringParameters" in request_data:
+                qs = request_data["queryStringParameters"]
+                if qs:
+                    additional_attributes["url.query"] = urlencode(qs)
+
+            sampling_context = {
+                "aws_event": aws_event,
+                "aws_context": aws_context,
+            }
+
+            function_name = aws_context.function_name
+
+            if has_span_streaming_enabled(client.options):
+                sentry_sdk.traces.continue_trace(headers)
+                Scope.set_custom_sampling_context(sampling_context)
+                span_ctx = sentry_sdk.traces.start_span(
+                    name=function_name,
+                    parent_span=None,
+                    attributes={
+                        "sentry.op": OP.FUNCTION_AWS,
+                        "sentry.origin": AwsLambdaIntegration.origin,
+                        "sentry.span.source": SegmentSource.COMPONENT,
+                        "cloud.region": aws_region,
+                        "cloud.resource_id": aws_context.invoked_function_arn,
+                        "cloud.platform": CLOUD_PLATFORM.AWS_LAMBDA,
+                        "cloud.provider": CLOUD_PROVIDER.AWS,
+                        "faas.name": function_name,
+                        "faas.invocation_id": aws_context.aws_request_id,
+                        "faas.version": aws_context.function_version,
+                        "aws.lambda.invoked_arn": aws_context.invoked_function_arn,
+                        "aws.log.group.names": [aws_context.log_group_name],
+                        "aws.log.stream.names": [aws_context.log_stream_name],
+                        "messaging.batch.message_count": batch_size,
+                        **header_attributes,
+                        **additional_attributes,
+                    },
+                )
+            else:
+                transaction = continue_trace(
+                    headers,
+                    op=OP.FUNCTION_AWS,
+                    name=function_name,
+                    source=TransactionSource.COMPONENT,
+                    origin=AwsLambdaIntegration.origin,
+                )
+
+                span_ctx = sentry_sdk.start_transaction(
+                    transaction, custom_sampling_context=sampling_context
+                )
+
+            with span_ctx:
                 try:
                     return handler(aws_event, aws_context, *args, **kwargs)
                 except Exception:

sentry_sdk/integrations/cloud_resource_context.py

@@ -48,6 +48,7 @@ class CLOUD_PLATFORM:  # noqa: N801
     """
 
     AWS_EC2 = "aws_ec2"
+    AWS_LAMBDA = "aws_lambda"
     GCP_COMPUTE_ENGINE = "gcp_compute_engine"
 
 

tests/integrations/aws_lambda/lambda_functions_with_embedded_sdk/BasicOkSpanStreaming/.gitignore

@@ -0,0 +1,11 @@
+# Need to add some ignore rules in this directory, because the unit tests will add the Sentry SDK and its dependencies
+# into this directory to create a Lambda function package that contains everything needed to instrument a Lambda function using Sentry.
+
+# Ignore everything
+*
+
+# But not index.py
+!index.py
+
+# And not .gitignore itself
+!.gitignore

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

@@ -0,0 +1,15 @@
+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()],
+    _experiments={"trace_lifecycle": "stream"},
+)
+
+
+def handler(event, context):
+    return {"event": event}

tests/integrations/aws_lambda/lambda_functions_with_embedded_sdk/BasicOkSpanStreamingPii/.gitignore

@@ -0,0 +1,11 @@
+# Need to add some ignore rules in this directory, because the unit tests will add the Sentry SDK and its dependencies
+# into this directory to create a Lambda function package that contains everything needed to instrument a Lambda function using Sentry.
+
+# Ignore everything
+*
+
+# But not index.py
+!index.py
+
+# And not .gitignore itself
+!.gitignore

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

@@ -0,0 +1,16 @@
+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,
+    send_default_pii=True,
+    integrations=[AwsLambdaIntegration()],
+    _experiments={"trace_lifecycle": "stream"},
+)
+
+
+def handler(event, context):
+    return {"event": event}

tests/integrations/aws_lambda/lambda_functions_with_embedded_sdk/RaiseErrorSpanStreaming/.gitignore

@@ -0,0 +1,11 @@
+# Need to add some ignore rules in this directory, because the unit tests will add the Sentry SDK and its dependencies
+# into this directory to create a Lambda function package that contains everything needed to instrument a Lambda function using Sentry.
+
+# Ignore everything
+*
+
+# But not index.py
+!index.py
+
+# And not .gitignore itself
+!.gitignore

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

@@ -0,0 +1,15 @@
+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()],
+    _experiments={"trace_lifecycle": "stream"},
+)
+
+
+def handler(event, context):
+    raise Exception("Oh!")