feat(boto3): Add Lambda extension

sentry_sdk/consts.py

@@ -375,6 +375,18 @@ class SPANDATA:
     Example: "arn:aws:lambda:us-east-1:123456789012:function:my-function"
     """
 
+    AWS_LAMBDA_RESOURCE_MAPPING_ID = "aws.lambda.resource_mapping.id"
+    """
+    The UUID of the AWS Lambda event source mapping the request refers to.
+    Example: "587b4f6c-1d0e-4d48-8a3f-2c6f9a1e0b7d"
+    """
+
+    AWS_LAMBDA_TARGET_FUNCTION_NAME = "aws.lambda.target_function.name"
+    """
+    The name of the AWS Lambda function an AWS SDK call acts on, without a qualifier.
+    Example: "billing"
+    """
+
     AWS_LOG_GROUP_NAMES = "aws.log.group.names"
     """
     The name(s) of the AWS log group(s) an application is writing to.
@@ -670,6 +682,24 @@ class SPANDATA:
     Example: "af9d5aa4-a685-4c5f-a22b-444f80b3cc28"
     """
 
+    FAAS_INVOKED_NAME = "faas.invoked_name"
+    """
+    The name of the invoked function, without a qualifier.
+    Example: "billing"
+    """
+
+    FAAS_INVOKED_PROVIDER = "faas.invoked_provider"
+    """
+    The cloud provider of the invoked function.
+    Example: "aws"
+    """
+
+    FAAS_INVOKED_REGION = "faas.invoked_region"
+    """
+    The cloud region of the invoked function.
+    Example: "eu-north-1"
+    """
+
     FAAS_NAME = "faas.name"
     """
     The name of the serverless function.

sentry_sdk/integrations/boto3/_client.py

@@ -96,6 +96,9 @@ def sentry_patched_make_api_call(
                 ctx.add_metadata(self)
             service_ext = _resolve_service(ctx.service_name)
             span = _start_client_span(ctx, service_ext)
+            with capture_internal_exceptions():
+                # override the request params for e.g. trace propagation.
+                api_params = service_ext.get_request_params(ctx, span) or api_params
 
         if span is None:
             return orig_make_api_call(self, operation_name, api_params)
@@ -107,7 +110,7 @@ def sentry_patched_make_api_call(
                 try:
                     parsed = orig_make_api_call(self, operation_name, api_params)
                 except BaseException as error:
-                    if service_ext is not None and isinstance(error, ClientError):
+                    if isinstance(error, ClientError):
                         with capture_internal_exceptions():
                             attributes.update(
                                 service_ext.get_response_attributes(ctx, error.response)
@@ -117,11 +120,8 @@ def sentry_patched_make_api_call(
                     with capture_internal_exceptions():
                         span.set_attributes(attributes)
                     raise
-                if service_ext is not None:
-                    with capture_internal_exceptions():
-                        attributes.update(
-                            service_ext.get_response_attributes(ctx, parsed)
-                        )
+                with capture_internal_exceptions():
+                    attributes.update(service_ext.get_response_attributes(ctx, parsed))
                 with capture_internal_exceptions():
                     attributes.update(_get_response_attributes(parsed))
                 with capture_internal_exceptions():
@@ -136,7 +136,8 @@ def sentry_patched_make_api_call(
 
         streaming_body_instrumented = False
         with capture_internal_exceptions():
-            streaming_body_instrumented = _instrument_streaming_body(span, parsed)
+            if service_ext.has_streaming_response(ctx):
+                streaming_body_instrumented = _instrument_streaming_body(span, parsed)
 
         # `StreamingBody`s finish their span when consumed or closed.
         if not streaming_body_instrumented:

sentry_sdk/integrations/boto3/_instrumentation.py

@@ -148,7 +148,7 @@ def _get_error_attributes(exception: "BaseException") -> "Attributes":
 
 def _start_client_span(
     ctx: "AwsCallContext",
-    service_ext: "Optional[_ServiceExtension]" = None,
+    service_ext: "_ServiceExtension",
 ) -> "Optional[Span]":
     if sentry_sdk.get_current_span() is None:
         return None
@@ -171,14 +171,13 @@ def _start_client_span(
     # with `get_span_op()`.
     span_op = OP.RPC
 
-    if service_ext is not None:
-        with capture_internal_exceptions():
-            service_op = service_ext.get_span_op(ctx)
-            if service_op is not None:
-                span_op = service_op
+    with capture_internal_exceptions():
+        service_op = service_ext.get_span_op(ctx)
+        if service_op is not None:
+            span_op = service_op
 
-        with capture_internal_exceptions():
-            attributes.update(service_ext.get_request_attributes(ctx))
+    with capture_internal_exceptions():
+        attributes.update(service_ext.get_request_attributes(ctx))
 
     # `start_span()` evaluates `ignore_spans` against the initial attributes.
     # https://opentelemetry.io/docs/specs/semconv/rpc/rpc-spans/#rpc-client-span

sentry_sdk/integrations/boto3/_services/base.py

@@ -1,18 +1,19 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Any, Optional
+    from typing import Any, Dict, Optional
 
     from sentry_sdk._types import Attributes
     from sentry_sdk.integrations.boto3._context import AwsCallContext
+    from sentry_sdk.traces import Span
 
 
 class _ServiceExtension:
     """
     Optional hooks for adding service-specific behavior to AWS client
     span; non-overridden methods keep the generic instrumentation.
-    Services without a registered extension in `_SERVICE_EXTENSIONS` continue
-    to use the generic instrumentation.
+    Services without a registered extension in `_SERVICE_EXTENSIONS` use this
+    class as is, i.e. the generic instrumentation.
     """
 
     __slots__ = ()
@@ -30,3 +31,19 @@ def get_response_attributes(
     ) -> "Attributes":
         """Return response attributes to add after the AWS request is made."""
         return {}
+
+    def get_request_params(
+        self, ctx: "AwsCallContext", span: "Optional[Span]"
+    ) -> "Optional[Dict[str, Any]]":
+        """
+        Return params to send instead of `ctx.params`, e.g. with trace propagation added,
+        or `None` to send them unchanged; must not mutate `ctx.params`.
+        """
+        return None
+
+    def has_streaming_response(self, ctx: "AwsCallContext") -> bool:
+        """
+        Return whether a `StreamingBody` response is read over time, so the span ends
+        when it's read or closed, or `False` to end it with the call.
+        """
+        return True

sentry_sdk/integrations/boto3/_services/lambda_.py

@@ -1,22 +1,113 @@
+import base64
+import json
 from typing import TYPE_CHECKING
 
+from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations.boto3._services.base import _ServiceExtension
-from sentry_sdk.integrations.boto3._utils import _get_aws_arn_attributes
+from sentry_sdk.integrations.boto3._utils import (
+    _get_aws_arn_attributes,
+    _get_propagation_headers,
+)
+from sentry_sdk.integrations.boto3.consts import CLOUD_PROVIDER
 from sentry_sdk.utils import capture_internal_exceptions
 
 if TYPE_CHECKING:
+    from typing import Any, Dict, Optional
+
     from sentry_sdk._types import Attributes
     from sentry_sdk.integrations.boto3._context import AwsCallContext
+    from sentry_sdk.traces import Span
 
 
 class _LambdaExtension(_ServiceExtension):
     def get_request_attributes(self, ctx: "AwsCallContext") -> "Attributes":
         attributes: "Attributes" = {}
+
+        # event source mapping apis.
+        if "UUID" in ctx.params:
+            attributes[SPANDATA.AWS_LAMBDA_RESOURCE_MAPPING_ID] = ctx.params["UUID"]
+        if "FunctionName" not in ctx.params:
+            return attributes
+
+        # a name, a partial arn (`<account>:function:<name>`) or an
+        # arn (`arn:<partition>:lambda:<region>:<account>:function:<name>`)
+        # each optionally followed by `:<version or alias>`.
+        function_name = ctx.params["FunctionName"]
+        name = function_name.split(":function:")[-1].split(":", 1)[0]
+
+        is_invoke = ctx.operation_name in (
+            "Invoke",
+            "InvokeAsync",
+            "InvokeWithResponseStream",
+        )
+        if is_invoke:
+            attributes[SPANDATA.FAAS_INVOKED_NAME] = name
+            attributes[SPANDATA.FAAS_INVOKED_PROVIDER] = CLOUD_PROVIDER
+            if ctx.region_name is not None:
+                attributes[SPANDATA.FAAS_INVOKED_REGION] = ctx.region_name
+        else:
+            attributes[SPANDATA.AWS_LAMBDA_TARGET_FUNCTION_NAME] = name
+
         with capture_internal_exceptions():
-            attributes.update(
-                _get_aws_arn_attributes(
-                    (ctx.params.get("FunctionName"),),
-                    ("lambda",),
-                )
-            )
+            arn_attributes = _get_aws_arn_attributes((function_name,), ("lambda",))
+            attributes.update(arn_attributes)
+            if is_invoke and arn_attributes:
+                attributes[SPANDATA.FAAS_INVOKED_REGION] = function_name.split(":")[3]
+
         return attributes
+
+    def get_response_attributes(
+        self, ctx: "AwsCallContext", response: "Any"
+    ) -> "Attributes":
+        attributes: "Attributes" = {}
+
+        if "UUID" in response:
+            attributes[SPANDATA.AWS_LAMBDA_RESOURCE_MAPPING_ID] = response["UUID"]
+
+        # `Invoke` returns `200 OK` when the function fails; reading it would leave
+        # the caller with an empty body.
+        if "FunctionError" in response:
+            attributes[SPANDATA.ERROR_TYPE] = "lambda_function_error"
+
+        return attributes
+
+    def get_request_params(
+        self, ctx: "AwsCallContext", span: "Optional[Span]"
+    ) -> "Optional[Dict[str, Any]]":
+        # lambda only passes `ClientContext` to the function on synchronous invocations.
+        # https://docs.aws.amazon.com/lambda/latest/api/API_Invoke.html#lambda-Invoke-request-ClientContext
+        if (
+            ctx.operation_name not in ("Invoke", "InvokeWithResponseStream")
+            or ctx.params.get("InvocationType", "RequestResponse") != "RequestResponse"
+        ):
+            return None
+
+        headers = _get_propagation_headers(ctx, span)
+
+        # https://docs.aws.amazon.com/lambda/latest/dg/python-context.html
+        client_context = (
+            json.loads(base64.b64decode(ctx.params["ClientContext"]))
+            if "ClientContext" in ctx.params
+            else {}
+        )
+        custom = client_context.get("custom", {})
+        # never mix with or replace trace data the caller propagates itself.
+        if not headers or any(key in custom for key in headers):
+            return None
+
+        client_context["custom"] = {**custom, **headers}
+        encoded = base64.b64encode(
+            json.dumps(client_context, separators=(",", ":")).encode()
+        ).decode()
+
+        # lambda rejects more than 3,583 bytes of base64; skip propagation rather
+        # than fail the call.
+        if len(encoded) > 3583:
+            return None
+
+        return {**ctx.params, "ClientContext": encoded}
+
+    def has_streaming_response(self, ctx: "AwsCallContext") -> bool:
+        # `Invoke`'s `Payload` is a `StreamingBody`, but span
+        # shouldn't be kept open until the caller reads it.
+        return ctx.operation_name != "Invoke"

sentry_sdk/integrations/boto3/_services/registry.py

@@ -1,5 +1,6 @@
 from typing import TYPE_CHECKING
 
+from sentry_sdk.integrations.boto3._services.base import _ServiceExtension
 from sentry_sdk.integrations.boto3._services.dynamodb import _DynamoDBExtension
 from sentry_sdk.integrations.boto3._services.kinesis import _KinesisExtension
 from sentry_sdk.integrations.boto3._services.lambda_ import _LambdaExtension
@@ -16,14 +17,12 @@
 if TYPE_CHECKING:
     from typing import Dict, Optional
 
-    from sentry_sdk.integrations.boto3._services.base import _ServiceExtension
-
 
 # add a ServiceExtension here when one is implemented. for example:
 #   _SERVICE_EXTENSIONS = {"s3": _S3Extension()}
 # when py 3.15 drops, we might want to take a look at using
 # a lazy-loading approach using the new `lazy` keyword.
-_SERVICE_EXTENSIONS: "Dict[str, _ServiceExtension]" = {
+_SERVICE_EXTENSIONS: "Dict[Optional[str], _ServiceExtension]" = {
     "dynamodb": _DynamoDBExtension(),
     "kinesis": _KinesisExtension(),
     "lambda": _LambdaExtension(),
@@ -37,8 +36,6 @@
 
 def _resolve_service(
     service_name: "Optional[str]",
-) -> "Optional[_ServiceExtension]":
-    """Return the extension for a service, or `None` for generic instrumentation."""
-    if service_name is None:
-        return None
-    return _SERVICE_EXTENSIONS.get(service_name)
+) -> "_ServiceExtension":
+    """Return the extension for a service, or `_ServiceExtension()` for generic instrumentation."""
+    return _SERVICE_EXTENSIONS.get(service_name, _ServiceExtension())

sentry_sdk/integrations/boto3/_utils.py

@@ -1,11 +1,16 @@
 from typing import TYPE_CHECKING
 
+import sentry_sdk
 from sentry_sdk.consts import SPANDATA
+from sentry_sdk.traces import NoOpSpan
+from sentry_sdk.tracing_utils import should_propagate_trace
 
 if TYPE_CHECKING:
     from typing import Any, Dict, Iterable, Optional, Sequence, Tuple
 
     from sentry_sdk._types import Attributes
+    from sentry_sdk.integrations.boto3._context import AwsCallContext
+    from sentry_sdk.traces import Span
 
     # tuple of (request param name, span attribute name).
     _AttributeSpec = Tuple[str, str]
@@ -38,3 +43,21 @@ def _get_aws_arn_attributes(
         return attributes
 
     return {}
+
+
+def _get_propagation_headers(
+    ctx: "AwsCallContext", span: "Optional[Span]"
+) -> "Dict[str, str]":
+    """Return the `sentry-trace` and `baggage` headers a call propagates to its consumer."""
+    if ctx.endpoint_url is None or not should_propagate_trace(
+        sentry_sdk.get_client(), ctx.endpoint_url
+    ):
+        return {}
+
+    # `NoOpSpan` is never sent, use scope's span instead.
+    if isinstance(span, NoOpSpan):
+        span = None
+
+    return dict(
+        sentry_sdk.get_current_scope().iter_trace_propagation_headers(span=span)
+    )