diff --git a/sentry_sdk/integrations/boto3/_instrumentation.py b/sentry_sdk/integrations/boto3/_instrumentation.py index 803b9233e8..5dc90ce5d7 100644 --- a/sentry_sdk/integrations/boto3/_instrumentation.py +++ b/sentry_sdk/integrations/boto3/_instrumentation.py @@ -1,10 +1,10 @@ from typing import TYPE_CHECKING import sentry_sdk -from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.consts import OP, SPANDATA, SPANSTATUS from sentry_sdk.integrations import DidNotEnable from sentry_sdk.integrations.boto3.consts import ORIGIN -from sentry_sdk.traces import StreamedSpan +from sentry_sdk.traces import NoOpStreamedSpan, StreamedSpan from sentry_sdk.tracing import BAGGAGE_HEADER_NAME, Span from sentry_sdk.tracing_utils import ( add_http_breadcrumb, @@ -164,26 +164,44 @@ def _replace_header(request: "AWSRequest", key: str, value: str) -> None: ) -def _sentry_after_call( - context: "Dict[str, Any]", parsed: "Dict[str, Any]", **kwargs: "Any" +def _finish_span( + span: "Union[Span, StreamedSpan]", + error: "Optional[BaseException]" = None, ) -> None: - span: "Optional[Union[Span, StreamedSpan]]" = context.pop("_sentrysdk_span", None) + with capture_internal_exceptions(): + if not isinstance(span, StreamedSpan): + if error is not None: + span.set_status(SPANSTATUS.INTERNAL_ERROR) + span.finish() + return - # Span could be absent if the integration is disabled. - if span is None: - return + if error is None: + span.end() + else: + span.__exit__(type(error), error, error.__traceback__) - span.__exit__(None, None, None) + +def _instrument_streaming_body( + span: "Union[Span, StreamedSpan]", parsed: "Dict[str, Any]" +) -> bool: + if isinstance(span, NoOpStreamedSpan): + return False body = parsed.get("Body") if not isinstance(body, StreamingBody): - return + return False streaming_span: "Union[Span, StreamedSpan]" if isinstance(span, StreamedSpan): streaming_span = sentry_sdk.traces.start_span( name=span.name, + # `parent_span` is set explicitly to the boto span. parent_span=span, + # avoid making the streaming span the current span on the scope since the application might + # keep `StreamingBody` open before reading it. Otherwise: 1. when the streamingspan ends it + # could restore the parent span on the scope, breaking the parent-child relation of newly + # created spans; 2. newly created spans would be attached to the streaming span. + active=False, attributes={ "sentry.op": OP.HTTP_CLIENT_STREAM, "sentry.origin": ORIGIN, @@ -198,35 +216,92 @@ def _sentry_after_call( orig_read = body.read orig_close = body.close + 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 + if finished: + return + + finished = True + _finish_span(streaming_span, error) + + def content_length_reached() -> bool: + content_length = getattr(body, "_content_length", None) + amount_read = getattr(body, "_amount_read", None) + return ( + content_length is not None + and amount_read is not None + and amount_read >= int(content_length) + ) def sentry_streaming_body_read(*args: "Any", **kwargs: "Any") -> bytes: + nonlocal read_in_progress + read_in_progress = True try: - ret = orig_read(*args, **kwargs) - if ret: - return ret - - if isinstance(streaming_span, StreamedSpan): - streaming_span.end() - else: - streaming_span.finish() - return ret - except Exception: - if isinstance(streaming_span, StreamedSpan): - streaming_span.end() - else: - streaming_span.finish() + read_return_value = orig_read(*args, **kwargs) + with capture_internal_exceptions(): + amount_of_bytes_requested = args[0] if args else kwargs.get("amt") + if ( + amount_of_bytes_requested is None + or amount_of_bytes_requested < 0 + or (amount_of_bytes_requested > 0 and not read_return_value) + or content_length_reached() + ): + finish_span() + return read_return_value + except BaseException as error: + finish_span(error) raise - - body.read = sentry_streaming_body_read # type: ignore + finally: + read_in_progress = False def sentry_streaming_body_close(*args: "Any", **kwargs: "Any") -> None: - if isinstance(streaming_span, StreamedSpan): - streaming_span.end() - else: - streaming_span.finish() - orig_close(*args, **kwargs) + try: + orig_close(*args, **kwargs) + finish_span() + except BaseException as error: + finish_span(error) + raise + + def sentry_raw_stream_close(*args: "Any", **kwargs: "Any") -> None: + try: + orig_raw_close(*args, **kwargs) + if not read_in_progress: + finish_span() + except BaseException as error: + finish_span(error) + raise + + try: + # StreamingBody.__exit__ closes `_raw_stream` directly, bypassing + # StreamingBody.close(), so both levels need to be instrumented. + raw_stream.close = sentry_raw_stream_close + body.read = sentry_streaming_body_read # type: ignore + body.close = sentry_streaming_body_close # type: ignore + except Exception: + finish_span() + raise - body.close = sentry_streaming_body_close # type: ignore + return True + + +def _sentry_after_call( + context: "Dict[str, Any]", parsed: "Dict[str, Any]", **kwargs: "Any" +) -> None: + span: "Optional[Union[Span, StreamedSpan]]" = context.pop("_sentrysdk_span", None) + + # Span could be absent if the integration is disabled. + if span is None: + return + + span.__exit__(None, None, None) + + with capture_internal_exceptions(): + _instrument_streaming_body(span, parsed) def _sentry_after_call_error( diff --git a/tests/integrations/boto3/test_client.py b/tests/integrations/boto3/test_client.py index db8f1b9263..6c81ba9313 100644 --- a/tests/integrations/boto3/test_client.py +++ b/tests/integrations/boto3/test_client.py @@ -1,6 +1,134 @@ +import boto3 +import pytest +from botocore.awsrequest import AWSResponse +from botocore.config import Config + +import sentry_sdk +from sentry_sdk.consts import OP from sentry_sdk.integrations.boto3 import Boto3Integration +from tests.integrations.boto3.aws_mock import Body + +session = boto3.Session( # type: ignore[attr-defined] + aws_access_key_id="-", + aws_secret_access_key="-", + region_name="eu-north-1", +) def test_public_api(): assert Boto3Integration.__module__ == "sentry_sdk.integrations.boto3" assert Boto3Integration.identifier == "boto3" + + +@pytest.fixture +def client_factory(sentry_init, monkeypatch, span_streaming): + sentry_init( + traces_sample_rate=1.0, + integrations=[Boto3Integration()], + trace_lifecycle="stream" if span_streaming else "static", + # avoid SDK's machine hostname being used as server name. + server_name="", + ) + # remove retry delay to speed up tests + monkeypatch.setattr("botocore.endpoint.time.sleep", lambda delay: None) + + def make_client(service_name="s3", attempt_count=1, **client_kwargs): + return session.client( + service_name, + config=Config( + # `total_max_attempts` includes the initial request. + retries={"total_max_attempts": attempt_count, "mode": "standard"} + ), + **client_kwargs, + ) + + return make_client + + +def _capture_boto3_spans_by_op(invoke_client_method, capture_items, span_streaming): + items = capture_items() + + if span_streaming: + with sentry_sdk.traces.start_span(name="parent"): # type: ignore[attr-defined] + invoke_client_method() + + sentry_sdk.flush() + spans = [ + item.payload + for item in items + if item.type == "span" + and item.payload["attributes"].get("sentry.origin") + == Boto3Integration.origin + ] + else: + with sentry_sdk.start_transaction(): + invoke_client_method() + + transaction = next(item.payload for item in items if item.type == "transaction") + spans = [ + span + for span in transaction["spans"] + if span["origin"] == Boto3Integration.origin + ] + + spans_by_op = {} + for span in spans: + op = span["attributes"].get("sentry.op") if span_streaming else span["op"] + spans_by_op.setdefault(op, []).append(span) + return spans_by_op + + +def _assert_span_finished(span, span_streaming): + finished_timestamp = "end_timestamp" if span_streaming else "timestamp" + assert span[finished_timestamp] is not None + + +def _assert_one_failed_span(spans, span_streaming): + assert len(spans) == 1 + assert spans[0]["status"] in ("error", "internal_error") + _assert_span_finished(spans[0], span_streaming) + + +@pytest.mark.parametrize("span_streaming", [True, False]) +def test_streaming_body_read_failure_finishes_stream_span( + capture_items, + client_factory, + span_streaming, +): + client = client_factory() + original_exception = OSError("stream read failed") + + class _FailingBody(Body): + def __init__(self, exception): + super().__init__(b"") + 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): + return AWSResponse( + request.url, + 200, + {"content-length": "1"}, + _FailingBody(original_exception), + ) + + client.meta.events.register("before-send", respond) + + def invoke_client_method_and_read_body(): + body = client.get_object(Bucket="bucket", Key="foo")["Body"] + with pytest.raises(OSError) as exc_info: + body.read() + assert exc_info.value is original_exception + + spans_by_op = _capture_boto3_spans_by_op( + invoke_client_method_and_read_body, capture_items, span_streaming + ) + client_spans = spans_by_op.get(OP.HTTP_CLIENT, []) + stream_spans = spans_by_op.get(OP.HTTP_CLIENT_STREAM, []) + + assert len(client_spans) == 1 + _assert_one_failed_span(stream_spans, span_streaming)