fix(boto3): ensure span finishes correctly when reading streaming body

Author: pabloDeputter

sentry_sdk/integrations/boto3/_instrumentation.py

@@ -219,6 +219,7 @@ def _instrument_streaming_body(
     raw_stream = body._raw_stream  # type: ignore[attr-defined]
     orig_raw_close = raw_stream.close
     finished = False
+    read_in_progress = False
 
     def finish_span(error: "Optional[BaseException]" = None) -> None:
         nonlocal finished
@@ -238,6 +239,8 @@ def content_length_reached() -> bool:
         )
 
     def sentry_streaming_body_read(*args: "Any", **kwargs: "Any") -> bytes:
+        nonlocal read_in_progress
+        read_in_progress = True
         try:
             read_return_value = orig_read(*args, **kwargs)
             with capture_internal_exceptions():
@@ -253,6 +256,8 @@ def sentry_streaming_body_read(*args: "Any", **kwargs: "Any") -> bytes:
         except BaseException as error:
             finish_span(error)
             raise
+        finally:
+            read_in_progress = False
 
     def sentry_streaming_body_close(*args: "Any", **kwargs: "Any") -> None:
         try:
@@ -265,7 +270,8 @@ def sentry_streaming_body_close(*args: "Any", **kwargs: "Any") -> None:
     def sentry_raw_stream_close(*args: "Any", **kwargs: "Any") -> None:
         try:
             orig_raw_close(*args, **kwargs)
-            finish_span()
+            if not read_in_progress:
+                finish_span()
         except BaseException as error:
             finish_span(error)
             raise

tests/integrations/boto3/test_client.py

@@ -104,6 +104,8 @@ def __init__(self, exception):
             self._exception = exception
 
         def read(self, *args, **kwargs):
+            # urllib3 closes the response before propagating some read failures.
+            self.close()
             raise self._exception
 
     def respond(request, **kwargs):