wip

Author: pabloDeputter

sentry_sdk/integrations/boto3/_client.py

@@ -96,9 +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)
-            if service_ext is not None:
-                with capture_internal_exceptions():
-                    api_params = service_ext.get_request_params(ctx, span) or api_params
+            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)
@@ -110,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)
@@ -120,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():
@@ -139,7 +136,7 @@ def sentry_patched_make_api_call(
 
         streaming_body_instrumented = False
         with capture_internal_exceptions():
-            if service_ext is None or service_ext.has_streaming_response(ctx):
+            if service_ext.has_streaming_response(ctx):
                 streaming_body_instrumented = _instrument_streaming_body(span, parsed)
 
         # `StreamingBody`s finish their span when consumed or closed.

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

@@ -12,8 +12,8 @@ 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__ = ()
@@ -26,22 +26,24 @@ def get_request_attributes(self, ctx: "AwsCallContext") -> "Attributes":
         """Return request attributes to add before the AWS request is made."""
         return {}
 
+    def get_response_attributes(
+        self, ctx: "AwsCallContext", response: "Any"
+    ) -> "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. Runs for every call, also when no span was
-        started (`span` is `None`), and must not mutate `ctx.params`.
+        or `None` to send them unchanged; must not mutate `ctx.params`.
         """
         return None
 
-    def get_response_attributes(
-        self, ctx: "AwsCallContext", response: "Any"
-    ) -> "Attributes":
-        """Return response attributes to add after the AWS request is made."""
-        return {}
-
     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 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

@@ -5,61 +5,72 @@
 from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations.boto3._services.base import _ServiceExtension
 from sentry_sdk.integrations.boto3._utils import (
-    _extract_attributes,
     _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, Sequence
+    from typing import Any, Dict, Optional
 
     from sentry_sdk._types import Attributes
     from sentry_sdk.integrations.boto3._context import AwsCallContext
-    from sentry_sdk.integrations.boto3._utils import _AttributeSpec
     from sentry_sdk.traces import Span
 
-# set on both request and response of the event source mapping apis.
-_RESOURCE_MAPPING_ATTRIBUTES: "Sequence[_AttributeSpec]" = (
-    ("UUID", SPANDATA.AWS_LAMBDA_RESOURCE_MAPPING_ID),
-)
-
 
 class _LambdaExtension(_ServiceExtension):
     def get_request_attributes(self, ctx: "AwsCallContext") -> "Attributes":
-        attributes: "Attributes" = _extract_attributes(
-            ctx.params, _RESOURCE_MAPPING_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, each optionally
-        # followed by `:<version or alias>`.
+        # 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]
-        invoke = ctx.operation_name in (
+
+        is_invoke = ctx.operation_name in (
             "Invoke",
             "InvokeAsync",
             "InvokeWithResponseStream",
         )
-        if invoke:
+        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
 
-        # params aren't validated yet by botocore.
         with capture_internal_exceptions():
             arn_attributes = _get_aws_arn_attributes((function_name,), ("lambda",))
             attributes.update(arn_attributes)
-            if invoke and 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]]":
@@ -72,6 +83,8 @@ def get_request_params(
             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
@@ -86,24 +99,15 @@ def get_request_params(
         encoded = base64.b64encode(
             json.dumps(client_context, separators=(",", ":")).encode()
         ).decode()
-        # it takes at most 3,583 bytes of base64.
+
+        # 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` returns once the function finished; its `Payload` is the complete
-        # result, not a stream. ending with the call keeps the span of a caller that
-        # never reads it.
+        # `Invoke`'s `Payload` is a `StreamingBody`, but span
+        # shouldn't be kept open until the caller reads it.
         return ctx.operation_name != "Invoke"
-
-    def get_response_attributes(
-        self, ctx: "AwsCallContext", response: "Any"
-    ) -> "Attributes":
-        attributes = _extract_attributes(response, _RESOURCE_MAPPING_ATTRIBUTES)
-        # `Invoke` returns `200 OK` when the function fails, with `FunctionError` set. the
-        # payload holds the error and is never read.
-        # https://docs.aws.amazon.com/lambda/latest/api/API_Invoke.html#lambda-Invoke-response-FunctionError
-        if "FunctionError" in response:
-            attributes[SPANDATA.ERROR_TYPE] = "lambda_function_error"
-        return attributes

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

@@ -48,19 +48,16 @@ def _get_aws_arn_attributes(
 def _get_propagation_headers(
     ctx: "AwsCallContext", span: "Optional[Span]"
 ) -> "Dict[str, str]":
-    """
-    Return the `sentry-trace` and `baggage` (and `traceparent`, if enabled) a call
-    propagates to its consumer, from its span or else the scope, like the HTTP headers.
-    Empty when `trace_propagation_targets` excludes the endpoint.
-    """
+    """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.
-    # which should be the parent normally.
     if isinstance(span, NoOpSpan):
         span = None
+
     return dict(
         sentry_sdk.get_current_scope().iter_trace_propagation_headers(span=span)
     )

tests/integrations/boto3/test_client.py

@@ -177,39 +177,11 @@ def record_client_span(request, **kwargs):
         assert span.get("end_timestamp") is not None
 
 
-@pytest.mark.parametrize(
-    "service_name,method_name,api_params,payload_field",
-    [
-        pytest.param(
-            "polly",
-            "synthesize_speech",
-            {"Text": "hi", "OutputFormat": "mp3", "VoiceId": "Joanna"},
-            "AudioStream",
-            id="audio-stream",
-        ),
-        pytest.param(
-            "s3",
-            "get_object_annotation",
-            {
-                "Bucket": "bucket",
-                "Key": "file.txt",
-                "AnnotationName": "annotation",
-            },
-            "AnnotationPayload",
-            id="annotation-payload",
-        ),
-    ],
-)
-def test_non_body_stream_delays_client_span(
-    capture_items,
-    client_factory,
-    service_name,
-    method_name,
-    api_params,
-    payload_field,
-):
-    client = client_factory(service_name)
-    require_botocore_model_fields(client, method_name, output_fields=(payload_field,))
+def test_non_body_stream_delays_client_span(capture_items, client_factory):
+    client = client_factory("s3")
+    require_botocore_model_fields(
+        client, "get_object_annotation", output_fields=("AnnotationPayload",)
+    )
     request_client_spans = []
 
     def record_client_span(request, **kwargs):
@@ -220,8 +192,10 @@ def record_client_span(request, **kwargs):
     def invoke():
         parent = sentry_sdk.get_current_span()
         with MockResponse(client, 200, {"content-length": "1"}, b"x"):
-            response = getattr(client, method_name)(**api_params)
-        body = response[payload_field]
+            response = client.get_object_annotation(
+                Bucket="bucket", Key="file.txt", AnnotationName="annotation"
+            )
+        body = response["AnnotationPayload"]
         assert isinstance(body, StreamingBody)
         (request_client_span,) = request_client_spans
         assert request_client_span.end_timestamp is None