use context managers

Author: pabloDeputter

sentry_sdk/integrations/boto3/_client.py

@@ -6,12 +6,12 @@
 from sentry_sdk.integrations.boto3._context import AwsCallContext
 from sentry_sdk.integrations.boto3._instrumentation import (
     _finish_active_http_child_span,
-    _finish_client_span,
+    _instrument_streaming_body,
     _sentry_before_sign,
     _sentry_request_created,
     _start_client_span,
 )
-from sentry_sdk.utils import capture_internal_exceptions
+from sentry_sdk.utils import capture_internal_exceptions, nullcontext
 
 if TYPE_CHECKING:
     from typing import Any, Optional, Union
@@ -58,22 +58,21 @@ def sentry_patched_make_api_call(
 
         with capture_internal_exceptions():
             span = _start_client_span(ctx)
-            if span is not None:
-                span.__enter__()
 
-        try:
-            parsed = orig_make_api_call(self, operation_name, api_params)
-        except BaseException as exc:
-            if span is not None:
-                with capture_internal_exceptions():
-                    _finish_active_http_child_span(span)
-                    span.__exit__(type(exc), exc, exc.__traceback__)
-            raise
+        span_ctx = span or nullcontext()
+        with span_ctx as span:
+            try:
+                parsed = orig_make_api_call(self, operation_name, api_params)
+            finally:
+                # stdlib HTTP span must leave stack before boto3 parent ctx
+                # manager restores the previous span.
+                if span is not None:
+                    with capture_internal_exceptions():
+                        _finish_active_http_child_span(span)
 
         if span is not None:
             with capture_internal_exceptions():
-                _finish_active_http_child_span(span)
-                _finish_client_span(span, parsed)
+                _instrument_streaming_body(span, parsed)
         return parsed
 
     BaseClient.__init__ = sentry_patched_init  # type: ignore

sentry_sdk/integrations/boto3/_instrumentation.py

@@ -165,12 +165,10 @@ def _finish_active_http_child_span(
     http_span.end()
 
 
-def _finish_client_span(
+def _instrument_streaming_body(
     span: "Union[Span, StreamedSpan]",
     parsed: "Dict[str, Any]",
 ) -> None:
-    span.__exit__(None, None, None)
-
     if isinstance(span, NoOpStreamedSpan):
         return