fix(boto3): trace the complete client-call lifecycle

Author: pabloDeputter

sentry_sdk/consts.py

@@ -1164,6 +1164,18 @@ class SPANDATA:
     Example: "prod"
     """
 
+    SENTRY_OP = "sentry.op"
+    """
+    The operation of a span.
+    Example: "http.client"
+    """
+
+    SENTRY_ORIGIN = "sentry.origin"
+    """
+    The origin of the instrumentation (e.g. span, log, etc.)
+    Example: "auto.http.otel.fastify"
+    """
+
     SENTRY_RELEASE = "sentry.release"
     """
     The Sentry release.

sentry_sdk/integrations/boto3/_client.py

@@ -1,44 +1,108 @@
-from functools import partial
+from contextlib import contextmanager
 from typing import TYPE_CHECKING
 
-from sentry_sdk.integrations import DidNotEnable, _check_minimum_version
+import sentry_sdk
+from sentry_sdk.integrations import DidNotEnable
+from sentry_sdk.integrations.boto3 import Boto3Integration
+from sentry_sdk.integrations.boto3._context import AwsCallContext
 from sentry_sdk.integrations.boto3._instrumentation import (
-    _sentry_after_call,
-    _sentry_after_call_error,
+    _finish_span,
+    _instrument_streaming_body,
     _sentry_before_sign,
     _sentry_request_created,
+    _start_client_span,
 )
-from sentry_sdk.utils import parse_version
+from sentry_sdk.traces import NoOpStreamedSpan, StreamedSpan
+from sentry_sdk.utils import capture_internal_exceptions
 
 if TYPE_CHECKING:
-    from typing import Any
+    from typing import Any, Iterator, Optional, Union
+
+    from sentry_sdk.tracing import Span
 
 try:
-    from botocore import __version__ as BOTOCORE_VERSION
     from botocore.client import BaseClient
 except ImportError:
-    raise DidNotEnable("botocore is not installed")
+    raise DidNotEnable("botocore not installed")
 
 
-def _patch_botocore_client() -> None:
-    from sentry_sdk.integrations.boto3 import Boto3Integration
+@contextmanager
+def _activate_client_span(span: "StreamedSpan") -> "Iterator[StreamedSpan]":
+    """Temporarily activate an inactive boto span without ending it."""
+    if isinstance(span, NoOpStreamedSpan):
+        yield span
+        return
 
-    version = parse_version(BOTOCORE_VERSION)
-    _check_minimum_version(Boto3Integration, version, "botocore")
+    scope = sentry_sdk.get_current_scope()
+    previous_span = scope.streamed_span
+    scope.streamed_span = span
+    try:
+        yield span
+    finally:
+        scope.streamed_span = previous_span
 
+
+def _patch_botocore_client() -> None:
     orig_init = BaseClient.__init__
+    orig_make_api_call = BaseClient._make_api_call  # type: ignore
 
     def sentry_patched_init(self: "BaseClient", *args: "Any", **kwargs: "Any") -> None:
         orig_init(self, *args, **kwargs)
         meta = self.meta
-        service_id = meta.service_model.service_id
-        meta.events.register(
-            "request-created",
-            partial(_sentry_request_created, service_id=service_id),
-        )
-        # run after other `before-sign` handlers, allowing it to see and preserve existing baggage.
+        meta.events.register("request-created", _sentry_request_created)
+        # run after other `before-sign` handlers so existing baggage is preserved.
         meta.events.register_last("before-sign", _sentry_before_sign)
-        meta.events.register("after-call", _sentry_after_call)
-        meta.events.register("after-call-error", _sentry_after_call_error)
+
+    def sentry_patched_make_api_call(
+        self: "BaseClient", operation_name: str, api_params: "Any"
+    ) -> "Any":
+        """
+        Track a single API call, including retries, serialization, and endpoint
+        resolution. For streaming responses, keep the span open until the
+        response body is consumed or closed.
+        https://github.com/boto/botocore/blob/develop/botocore/client.py
+        https://opentelemetry.io/docs/specs/semconv/rpc/rpc-spans/#rpc-client-span
+        """
+        client = sentry_sdk.get_client()
+        if client.get_integration(Boto3Integration) is None:
+            return orig_make_api_call(self, operation_name, api_params)
+
+        ctx = AwsCallContext(operation_name)
+
+        # add optional metadata to context.
+        with capture_internal_exceptions():
+            ctx.add_metadata(self)
+
+        span: "Optional[Union[Span, StreamedSpan]]" = None
+        with capture_internal_exceptions():
+            span = _start_client_span(ctx)
+
+        if span is None:
+            return orig_make_api_call(self, operation_name, api_params)
+
+        # activate without finishing; a streaming response may outlive the call.
+        span_ctx = (
+            _activate_client_span(span) if isinstance(span, StreamedSpan) else span
+        )
+
+        try:
+            with span_ctx:
+                parsed = orig_make_api_call(self, operation_name, api_params)
+        except BaseException as error:
+            # finish `StreamedSpan` explicitly; static spans are finished by
+            # their context manager.
+            if isinstance(span, StreamedSpan):
+                _finish_span(span, error)
+            raise
+
+        streaming_body_instrumented = False
+        with capture_internal_exceptions():
+            streaming_body_instrumented = _instrument_streaming_body(span, parsed)
+
+        # `StreamingBody`s finish their span when consumed or closed.
+        if isinstance(span, StreamedSpan) and not streaming_body_instrumented:
+            _finish_span(span)
+        return parsed
 
     BaseClient.__init__ = sentry_patched_init  # type: ignore
+    BaseClient._make_api_call = sentry_patched_make_api_call  # type: ignore

sentry_sdk/integrations/boto3/_context.py

@@ -0,0 +1,44 @@
+from typing import TYPE_CHECKING
+
+from sentry_sdk.integrations import DidNotEnable
+from sentry_sdk.utils import capture_internal_exceptions
+
+if TYPE_CHECKING:
+    from typing import Any, Optional
+
+try:
+    from botocore.client import BaseClient
+except ImportError:
+    raise DidNotEnable("botocore not installed")
+
+
+class AwsCallContext:
+    __slots__ = (
+        "service_id",
+        "service_id_hyphenized",
+        "operation_name",
+    )
+
+    def __init__(self, operation_name: str) -> None:
+        self.operation_name: str = operation_name
+        self.service_id: "Optional[str]" = None
+        self.service_id_hyphenized: "Optional[str]" = None
+
+    def add_metadata(self, client: "BaseClient") -> None:
+        def _get_attr(obj: "Any", name: str) -> "Any":
+            if obj is None:
+                return None
+
+            with capture_internal_exceptions():
+                return getattr(obj, name)
+
+        client_meta = _get_attr(client, "meta")
+        service_model = _get_attr(client_meta, "service_model")
+
+        # modeled AWS service identity used in span names, e.g. `API Gateway`.
+        service_id = _get_attr(service_model, "service_id")
+        if service_id is not None:
+            with capture_internal_exceptions():
+                self.service_id = str(service_id)
+            with capture_internal_exceptions():
+                self.service_id_hyphenized = service_id.hyphenize()