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

sentry_sdk/integrations/aws_lambda.py

@@ -12,8 +12,14 @@
 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,
@@ -144,20 +150,74 @@ 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"] = "&".join(
+                        f"{k}={v}" for k, v in qs.items()
+                    )
+
+            aws_region = aws_context.invoked_function_arn.split(":")[3]
+
+            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],
+                        **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!")

tests/integrations/aws_lambda/test_aws_lambda.py

@@ -525,6 +525,182 @@ def test_error_has_new_trace_context(
         )
 
 
+def _get_span_attr(attrs, key):
+    """Extract the value from a span attribute, handling both flat and typed formats."""
+    val = attrs[key]
+    if isinstance(val, dict) and "value" in val:
+        return val["value"]
+    return val
+
+
+def _get_segment_spans(span_items):
+    """Filter span items to only segment (root) spans."""
+    return [s for s in span_items if s.get("is_segment") is True]
+
+
+def test_span_streaming_no_error(lambda_client, test_environment):
+    lambda_client.invoke(
+        FunctionName="BasicOkSpanStreaming",
+        Payload=json.dumps({}),
+    )
+    envelopes = test_environment["server"].envelopes
+    segment_spans = _get_segment_spans(test_environment["server"].span_items)
+
+    assert len(envelopes) == 0
+    assert len(segment_spans) == 1
+
+    segment_span = segment_spans[0]
+    assert segment_span["is_segment"] is True
+    assert segment_span["name"] == "BasicOkSpanStreaming"
+    attrs = segment_span["attributes"]
+    assert _get_span_attr(attrs, "sentry.op") == "function.aws"
+    assert _get_span_attr(attrs, "sentry.origin") == "auto.function.aws_lambda"
+    assert _get_span_attr(attrs, "sentry.span.source") == "component"
+    assert _get_span_attr(attrs, "cloud.provider") == "aws"
+    assert _get_span_attr(attrs, "cloud.platform") == "aws_lambda"
+    assert (
+        _get_span_attr(attrs, "cloud.resource_id")
+        == "arn:aws:lambda:us-east-1:012345678912:function:BasicOkSpanStreaming"
+    )
+    assert _get_span_attr(attrs, "cloud.region") == "us-east-1"
+    assert _get_span_attr(attrs, "faas.name") == "BasicOkSpanStreaming"
+    assert _get_span_attr(attrs, "faas.version") == "$LATEST"
+    assert "faas.invocation_id" in attrs
+    assert (
+        _get_span_attr(attrs, "aws.lambda.invoked_arn")
+        == "arn:aws:lambda:us-east-1:012345678912:function:BasicOkSpanStreaming"
+    )
+    assert _get_span_attr(attrs, "aws.log.group.names") == [
+        "aws/lambda/BasicOkSpanStreaming"
+    ]
+    assert _get_span_attr(attrs, "aws.log.stream.names") == ["$LATEST"]
+
+
+def test_span_streaming_error(lambda_client, test_environment):
+    lambda_client.invoke(
+        FunctionName="RaiseErrorSpanStreaming",
+        Payload=json.dumps({}),
+    )
+    envelopes = test_environment["server"].envelopes
+    segment_spans = _get_segment_spans(test_environment["server"].span_items)
+
+    assert len(envelopes) == 1
+    error_event = envelopes[0]
+    assert error_event["level"] == "error"
+    (exception,) = error_event["exception"]["values"]
+    assert exception["type"] == "Exception"
+    assert exception["value"] == "Oh!"
+    assert exception["mechanism"]["type"] == "aws_lambda"
+    assert not exception["mechanism"]["handled"]
+
+    assert len(segment_spans) == 1
+    segment_span = segment_spans[0]
+    assert segment_span["is_segment"] is True
+    assert segment_span["name"] == "RaiseErrorSpanStreaming"
+    assert segment_span["status"] == "error"
+    attrs = segment_span["attributes"]
+    assert _get_span_attr(attrs, "sentry.op") == "function.aws"
+    assert _get_span_attr(attrs, "sentry.origin") == "auto.function.aws_lambda"
+    assert _get_span_attr(attrs, "sentry.span.source") == "component"
+    assert _get_span_attr(attrs, "cloud.provider") == "aws"
+    assert _get_span_attr(attrs, "cloud.platform") == "aws_lambda"
+    assert (
+        _get_span_attr(attrs, "cloud.resource_id")
+        == "arn:aws:lambda:us-east-1:012345678912:function:RaiseErrorSpanStreaming"
+    )
+    assert _get_span_attr(attrs, "cloud.region") == "us-east-1"
+    assert _get_span_attr(attrs, "faas.name") == "RaiseErrorSpanStreaming"
+    assert _get_span_attr(attrs, "faas.version") == "$LATEST"
+    assert "faas.invocation_id" in attrs
+    assert (
+        _get_span_attr(attrs, "aws.lambda.invoked_arn")
+        == "arn:aws:lambda:us-east-1:012345678912:function:RaiseErrorSpanStreaming"
+    )
+    assert _get_span_attr(attrs, "aws.log.group.names") == [
+        "aws/lambda/RaiseErrorSpanStreaming"
+    ]
+    assert _get_span_attr(attrs, "aws.log.stream.names") == ["$LATEST"]
+
+
+def test_span_streaming_trace_continuation(lambda_client, test_environment):
+    trace_id = "471a43a4192642f0b136d5159a501701"
+    parent_span_id = "6e8f22c393e68f19"
+    parent_sampled = 1
+    sentry_trace_header = "{}-{}-{}".format(trace_id, parent_span_id, parent_sampled)
+
+    payload = {
+        "headers": {
+            "sentry-trace": sentry_trace_header,
+        }
+    }
+
+    lambda_client.invoke(
+        FunctionName="RaiseErrorSpanStreaming",
+        Payload=json.dumps(payload),
+    )
+    envelopes = test_environment["server"].envelopes
+    segment_spans = _get_segment_spans(test_environment["server"].span_items)
+
+    assert len(envelopes) == 1
+    error_event = envelopes[0]
+    assert error_event["contexts"]["trace"]["trace_id"] == trace_id
+
+    assert len(segment_spans) == 1
+    segment_span = segment_spans[0]
+    assert segment_span["is_segment"] is True
+    assert segment_span["trace_id"] == trace_id
+    assert segment_span["name"] == "RaiseErrorSpanStreaming"
+    attrs = segment_span["attributes"]
+    assert _get_span_attr(attrs, "sentry.op") == "function.aws"
+    assert _get_span_attr(attrs, "sentry.origin") == "auto.function.aws_lambda"
+    assert _get_span_attr(attrs, "sentry.span.source") == "component"
+    assert _get_span_attr(attrs, "cloud.provider") == "aws"
+    assert _get_span_attr(attrs, "cloud.platform") == "aws_lambda"
+    assert _get_span_attr(attrs, "cloud.region") == "us-east-1"
+    assert _get_span_attr(attrs, "faas.name") == "RaiseErrorSpanStreaming"
+    assert _get_span_attr(attrs, "faas.version") == "$LATEST"
+    assert "faas.invocation_id" in attrs
+
+
+def test_span_streaming_request_attributes(lambda_client, test_environment):
+    payload = {
+        "headers": {
+            "Content-Type": "application/json",
+            "Accept": "text/html",
+        },
+        "httpMethod": "POST",
+        "queryStringParameters": {"foo": "bar"},
+        "path": "/test",
+    }
+
+    lambda_client.invoke(
+        FunctionName="BasicOkSpanStreamingPii",
+        Payload=json.dumps(payload),
+    )
+    segment_spans = _get_segment_spans(test_environment["server"].span_items)
+
+    assert len(segment_spans) == 1
+    segment_span = segment_spans[0]
+    attrs = segment_span["attributes"]
+
+    assert _get_span_attr(attrs, "http.request.method") == "POST"
+    assert _get_span_attr(attrs, "url.query") == "foo=bar"
+    assert (
+        _get_span_attr(attrs, "http.request.header.content-type") == "application/json"
+    )
+    assert _get_span_attr(attrs, "http.request.header.accept") == "text/html"
+    assert _get_span_attr(attrs, "faas.name") == "BasicOkSpanStreamingPii"
+    assert _get_span_attr(attrs, "cloud.provider") == "aws"
+    assert _get_span_attr(attrs, "cloud.platform") == "aws_lambda"
+    assert _get_span_attr(attrs, "cloud.region") == "us-east-1"
+    assert _get_span_attr(attrs, "faas.version") == "$LATEST"
+    assert "faas.invocation_id" in attrs
+    assert _get_span_attr(attrs, "aws.log.group.names") == [
+        "aws/lambda/BasicOkSpanStreamingPii"
+    ]
+    assert _get_span_attr(attrs, "aws.log.stream.names") == ["$LATEST"]
+
+
 @pytest.mark.parametrize(
     "lambda_function_name",
     ["RaiseErrorPerformanceEnabled", "RaiseErrorPerformanceDisabled"],