Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
124 changes: 39 additions & 85 deletions sentry_sdk/integrations/boto3.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,11 +6,10 @@
from sentry_sdk.integrations import DidNotEnable, Integration, _check_minimum_version
from sentry_sdk.scope import should_send_default_pii
from sentry_sdk.traces import StreamedSpan
from sentry_sdk.tracing import BAGGAGE_HEADER_NAME, Span
from sentry_sdk.tracing import BAGGAGE_HEADER_NAME
from sentry_sdk.tracing_utils import (
add_http_breadcrumb,
add_sentry_baggage_to_headers,
has_span_streaming_enabled,
should_propagate_trace,
)
from sentry_sdk.utils import (
Expand All @@ -20,7 +19,7 @@
)

if TYPE_CHECKING:
from typing import Any, Dict, Optional, Type, Union
from typing import Any, Dict, Optional, Type

from botocore.model import ServiceId

Expand Down Expand Up @@ -78,11 +77,30 @@ def _sentry_request_created(

breadcrumb: "dict[str, Any]" = {}

is_span_streaming_enabled = has_span_streaming_enabled(client.options)
span: "Union[Span, StreamedSpan, None]" = None
if is_span_streaming_enabled:
span: "Optional[StreamedSpan]" = None
if parsed_url and should_send_default_pii():
breadcrumb.update(
{
SPANDATA.URL_FULL: parsed_url.url,
SPANDATA.URL_QUERY: parsed_url.query,
SPANDATA.URL_FRAGMENT: parsed_url.fragment,
}
)

if request.method is not None:
breadcrumb[SPANDATA.HTTP_REQUEST_METHOD] = request.method

if sentry_sdk.traces.get_current_span() is not None:
span = sentry_sdk.traces.start_span(
name=description,
attributes={
"sentry.op": OP.HTTP_CLIENT,
"sentry.origin": Boto3Integration.origin,
SPANDATA.RPC_METHOD: f"{service_id}/{operation_name}",
},
)
if parsed_url and should_send_default_pii():
breadcrumb.update(
span.set_attributes(
{
SPANDATA.URL_FULL: parsed_url.url,
SPANDATA.URL_QUERY: parsed_url.query,
Expand All @@ -91,54 +109,7 @@ def _sentry_request_created(
)

if request.method is not None:
breadcrumb[SPANDATA.HTTP_REQUEST_METHOD] = request.method

if sentry_sdk.traces.get_current_span() is not None:
span = sentry_sdk.traces.start_span(
name=description,
attributes={
"sentry.op": OP.HTTP_CLIENT,
"sentry.origin": Boto3Integration.origin,
SPANDATA.RPC_METHOD: f"{service_id}/{operation_name}",
},
)
if parsed_url and should_send_default_pii():
span.set_attributes(
{
SPANDATA.URL_FULL: parsed_url.url,
SPANDATA.URL_QUERY: parsed_url.query,
SPANDATA.URL_FRAGMENT: parsed_url.fragment,
}
)

if request.method is not None:
span.set_attribute(SPANDATA.HTTP_REQUEST_METHOD, request.method)
else:
span = sentry_sdk.start_span(
op=OP.HTTP_CLIENT,
name=description,
origin=Boto3Integration.origin,
)

if parsed_url:
span.set_data("aws.request.url", parsed_url.url)
span.set_data(SPANDATA.HTTP_QUERY, parsed_url.query)
span.set_data(SPANDATA.HTTP_FRAGMENT, parsed_url.fragment)
breadcrumb.update(
{
"aws.request.url": parsed_url.url,
SPANDATA.HTTP_QUERY: parsed_url.query,
SPANDATA.HTTP_FRAGMENT: parsed_url.fragment,
}
)

if request.method is not None:
span.set_data(SPANDATA.HTTP_METHOD, request.method)
breadcrumb[SPANDATA.HTTP_METHOD] = request.method

# We do it in order for subsequent http calls/retries be
# attached to this span.
span.__enter__()
span.set_attribute(SPANDATA.HTTP_REQUEST_METHOD, request.method)

add_http_breadcrumb(None, breadcrumb)

Expand Down Expand Up @@ -205,7 +176,7 @@ def _replace_header(request: "AWSRequest", key: str, value: str) -> None:
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: "Optional[StreamedSpan]" = context.pop("_sentrysdk_span", None)

# Span could be absent if the integration is disabled.
if span is None:
Expand All @@ -217,22 +188,14 @@ def _sentry_after_call(
if not isinstance(body, StreamingBody):
return

streaming_span: "Union[Span, StreamedSpan]"
if isinstance(span, StreamedSpan):
streaming_span = sentry_sdk.traces.start_span(
name=span.name,
parent_span=span,
attributes={
"sentry.op": OP.HTTP_CLIENT_STREAM,
"sentry.origin": Boto3Integration.origin,
},
)
else:
streaming_span = span.start_child(
op=OP.HTTP_CLIENT_STREAM,
name=span.description,
origin=Boto3Integration.origin,
)
streaming_span = sentry_sdk.traces.start_span(
name=span.name,
parent_span=span,
attributes={
"sentry.op": OP.HTTP_CLIENT_STREAM,
"sentry.origin": Boto3Integration.origin,
},
)

orig_read = body.read
orig_close = body.close
Expand All @@ -243,25 +206,16 @@ def sentry_streaming_body_read(*args: "Any", **kwargs: "Any") -> bytes:
if ret:
return ret

if isinstance(streaming_span, StreamedSpan):
streaming_span.end()
else:
streaming_span.finish()
streaming_span.end()
return ret
except Exception:
if isinstance(streaming_span, StreamedSpan):
streaming_span.end()
else:
streaming_span.finish()
streaming_span.end()
raise

body.read = sentry_streaming_body_read # type: ignore

def sentry_streaming_body_close(*args: "Any", **kwargs: "Any") -> None:
if isinstance(streaming_span, StreamedSpan):
streaming_span.end()
else:
streaming_span.finish()
streaming_span.end()
orig_close(*args, **kwargs)

body.close = sentry_streaming_body_close # type: ignore
Expand All @@ -270,7 +224,7 @@ def sentry_streaming_body_close(*args: "Any", **kwargs: "Any") -> None:
def _sentry_after_call_error(
context: "Dict[str, Any]", exception: "Type[BaseException]", **kwargs: "Any"
) -> None:
span: "Optional[Union[Span, StreamedSpan]]" = context.pop("_sentrysdk_span", None)
span: "Optional[StreamedSpan]" = context.pop("_sentrysdk_span", None)

# Span could be absent if the integration is disabled.
if span is None:
Expand Down
111 changes: 37 additions & 74 deletions tests/integrations/boto3/test_aws_http_connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,25 +45,20 @@ def _request(server, headers, path="/"):
connection.close()


@pytest.mark.parametrize("span_streaming", [False, True])
def test_aws_http_connection_adds_missing_unsigned_propagation_headers(
sentry_init, local_http_server, span_streaming
sentry_init,
local_http_server,
):
"""Add missing unsigned `sentry-trace` and `baggage`."""
sentry_init(
traces_sample_rate=1.0,
trace_lifecycle="stream" if span_streaming else "static",
trace_lifecycle="stream",
default_integrations=False,
integrations=[StdlibIntegration()],
)
server, requests = local_http_server

if span_streaming:
with sentry_sdk.traces.start_span(name="test"): # type: ignore[attr-defined]
_request(server, [])
else:
with sentry_sdk.start_transaction(name="test", sampled=True):
_request(server, [])
with sentry_sdk.traces.start_span(name="test"): # type: ignore[attr-defined]
_request(server, [])

headers: HTTPMessage = requests[0]

Expand All @@ -78,37 +73,26 @@ def test_aws_http_connection_adds_missing_unsigned_propagation_headers(
assert len(sentry_trace_headers) == 1


@pytest.mark.parametrize("span_streaming", [False, True])
def test_aws_http_connection_appends_baggage_but_preserves_sentry_trace(
sentry_init, local_http_server, span_streaming
sentry_init,
local_http_server,
):
"""Append unsigned `baggage`; leave existing `sentry-trace` as-is."""
sentry_init(
traces_sample_rate=1.0,
trace_lifecycle="stream" if span_streaming else "static",
trace_lifecycle="stream",
default_integrations=False,
integrations=[StdlibIntegration()],
)
server, requests = local_http_server

if span_streaming:
with sentry_sdk.traces.start_span(name="test"): # type: ignore[attr-defined]
_request(
server,
[
("baggage", "vendor=value"),
("sentry-trace", "existing-trace"),
],
)
else:
with sentry_sdk.start_transaction(name="test", sampled=True):
_request(
server,
[
("baggage", "vendor=value"),
("sentry-trace", "existing-trace"),
],
)
with sentry_sdk.traces.start_span(name="test"): # type: ignore[attr-defined]
_request(
server,
[
("baggage", "vendor=value"),
("sentry-trace", "existing-trace"),
],
)

headers: HTTPMessage = requests[0]

Expand All @@ -123,14 +107,14 @@ def test_aws_http_connection_appends_baggage_but_preserves_sentry_trace(
assert headers.get_all("sentry-trace") == ["existing-trace"]


@pytest.mark.parametrize("span_streaming", [False, True])
def test_aws_http_connection_preserves_signed_propagation_headers(
sentry_init, local_http_server, span_streaming
sentry_init,
local_http_server,
):
"""Leave signed `sentry-trace` and `baggage` as-is."""
sentry_init(
traces_sample_rate=1.0,
trace_lifecycle="stream" if span_streaming else "static",
trace_lifecycle="stream",
default_integrations=False,
integrations=[StdlibIntegration()],
)
Expand All @@ -143,27 +127,15 @@ def test_aws_http_connection_preserves_signed_propagation_headers(
"SignedHeaders=baggage;host;sentry-trace, "
"Signature=sixtyseven"
)

if span_streaming:
with sentry_sdk.traces.start_span(name="test"): # type: ignore[attr-defined]
_request(
server,
[
("baggage", "vendor=value"),
("sentry-trace", "existing-trace"),
("Authorization", authorization),
],
)
else:
with sentry_sdk.start_transaction(name="test", sampled=True):
_request(
server,
[
("baggage", "vendor=value"),
("sentry-trace", "existing-trace"),
("Authorization", authorization),
],
)
with sentry_sdk.traces.start_span(name="test"): # type: ignore[attr-defined]
_request(
server,
[
("baggage", "vendor=value"),
("sentry-trace", "existing-trace"),
("Authorization", authorization),
],
)

headers: HTTPMessage = requests[0]

Expand All @@ -180,14 +152,14 @@ def test_aws_http_connection_preserves_signed_propagation_headers(
}


@pytest.mark.parametrize("span_streaming", [False, True])
def test_aws_http_connection_preserves_query_signed_baggage(
sentry_init, local_http_server, span_streaming
sentry_init,
local_http_server,
):
"""Leave query-signed `baggage` as-is; add unsigned `sentry-trace`."""
sentry_init(
traces_sample_rate=1.0,
trace_lifecycle="stream" if span_streaming else "static",
trace_lifecycle="stream",
default_integrations=False,
integrations=[StdlibIntegration()],
)
Expand All @@ -202,21 +174,12 @@ def test_aws_http_connection_preserves_query_signed_baggage(
"&X-Amz-SignedHeaders=baggage%3Bhost"
"&X-Amz-Signature=sixtyseven"
)

if span_streaming:
with sentry_sdk.traces.start_span(name="test"): # type: ignore[attr-defined]
_request(
server,
[("baggage", "vendor=value")],
path=path,
)
else:
with sentry_sdk.start_transaction(name="test", sampled=True):
_request(
server,
[("baggage", "vendor=value")],
path=path,
)
with sentry_sdk.traces.start_span(name="test"): # type: ignore[attr-defined]
_request(
server,
[("baggage", "vendor=value")],
path=path,
)

headers: HTTPMessage = requests[0]
# query-signed `baggage`: leave as-is.
Expand Down
Loading
Loading