diff --git a/sentry_sdk/integrations/boto3/_client.py b/sentry_sdk/integrations/boto3/_client.py index 32015c802e..280a570d45 100644 --- a/sentry_sdk/integrations/boto3/_client.py +++ b/sentry_sdk/integrations/boto3/_client.py @@ -14,9 +14,7 @@ _set_span_attributes, _start_client_span, ) -from sentry_sdk.integrations.boto3._services.registry import ( - _resolve_service, -) +from sentry_sdk.integrations.boto3._services.bedrock import _BedrockExtension from sentry_sdk.integrations.boto3.consts import IDENTIFIER from sentry_sdk.traces import NoOpStreamedSpan, StreamedSpan from sentry_sdk.utils import capture_internal_exceptions @@ -104,7 +102,11 @@ def sentry_patched_make_api_call( with capture_internal_exceptions(): # add optional metadata to the context, e.g. service-name, region-name, etc. ctx.add_metadata(self) - service_ext = _resolve_service(ctx.service_name) + + service_ext = None + if ctx.service_name == "bedrock-runtime": + service_ext = _BedrockExtension() + span = _start_client_span(ctx, service_ext) if span is None: diff --git a/sentry_sdk/integrations/boto3/_instrumentation.py b/sentry_sdk/integrations/boto3/_instrumentation.py index 3c180a025c..967e320dbe 100644 --- a/sentry_sdk/integrations/boto3/_instrumentation.py +++ b/sentry_sdk/integrations/boto3/_instrumentation.py @@ -173,8 +173,7 @@ def _start_client_span( # https://opentelemetry.io/docs/specs/semconv/cloud-providers/aws-sdk/#aws-sdk-spans SPANDATA.SENTRY_KIND: "client", } - with capture_internal_exceptions(): - attributes.update(_get_client_attributes(ctx)) + attributes.update(_get_client_attributes(ctx)) # `sentry.span_op` and `sentry.span_origin` are set to generic defaults; # a service extension can override them with `get_span_op()` and `get_span_origin()`. @@ -182,21 +181,20 @@ def _start_client_span( span_origin = ORIGIN if service_ext is not None: - with capture_internal_exceptions(): - service_op = service_ext.get_span_op(ctx) - if service_op is not None: - span_op = service_op + service_op = service_ext.get_span_op(ctx) + if service_op is not None: + span_op = service_op - with capture_internal_exceptions(): - service_origin = service_ext.get_span_origin(ctx) - if service_origin is not None: - span_origin = service_origin + service_origin = service_ext.get_span_origin(ctx) + if service_origin is not None: + span_origin = service_origin - with capture_internal_exceptions(): - attributes.update(service_ext.get_request_attributes(ctx)) + attributes.update(service_ext.get_request_attributes(ctx)) if has_span_streaming_enabled(client.options): - if sentry_sdk.traces.get_current_span() is None: + if ( + service_ext is None or service_ext.skip_if_no_active_span + ) and sentry_sdk.traces.get_current_span() is None: return None # `start_span()` evaluates `ignore_spans` against the initial attributes. diff --git a/sentry_sdk/integrations/boto3/_services/base.py b/sentry_sdk/integrations/boto3/_services/base.py index 0d9cb1c24d..8cc29b59e7 100644 --- a/sentry_sdk/integrations/boto3/_services/base.py +++ b/sentry_sdk/integrations/boto3/_services/base.py @@ -16,6 +16,7 @@ class _ServiceExtension: """ __slots__ = () + skip_if_no_active_span = True def get_span_op(self, ctx: "AwsCallContext") -> "Optional[str]": """Return an optional `sentry.op` override, or `None` to keep the default.""" diff --git a/sentry_sdk/integrations/boto3/_services/bedrock.py b/sentry_sdk/integrations/boto3/_services/bedrock.py new file mode 100644 index 0000000000..f91ca36b19 --- /dev/null +++ b/sentry_sdk/integrations/boto3/_services/bedrock.py @@ -0,0 +1,42 @@ +from typing import TYPE_CHECKING + +from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.integrations.boto3._services.base import _ServiceExtension + +if TYPE_CHECKING: + from typing import Any, Optional + + from sentry_sdk._types import Attributes + from sentry_sdk.integrations.boto3._context import AwsCallContext + + +class _BedrockExtension(_ServiceExtension): + __slots__ = () + skip_if_no_active_span = False + + def get_span_op(self, ctx: "AwsCallContext") -> "Optional[str]": + if ctx.operation_name == "Converse" and ctx.operation_name == "ConverseStream": + return None + return OP.GEN_AI_CHAT + + def get_span_origin(self, ctx: "AwsCallContext") -> "Optional[str]": + if ctx.operation_name != "Converse" and ctx.operation_name != "ConverseStream": + return None + return "auto.ai.boto3" + + def get_request_attributes(self, ctx: "AwsCallContext") -> "Attributes": + if ctx.operation_name != "Converse" and ctx.operation_name != "ConverseStream": + return {} + + return { + SPANDATA.GEN_AI_PROVIDER_NAME: "aws.bedrock", + SPANDATA.GEN_AI_OPERATION_NAME: "chat", + SPANDATA.GEN_AI_RESPONSE_STREAMING: ctx.operation_name == "ConverseStream", + } + + def get_response_attributes( + self, + ctx: "AwsCallContext", + response: "Any", + ) -> "Attributes": + return {} diff --git a/sentry_sdk/integrations/boto3/_services/registry.py b/sentry_sdk/integrations/boto3/_services/registry.py deleted file mode 100644 index b387db0f26..0000000000 --- a/sentry_sdk/integrations/boto3/_services/registry.py +++ /dev/null @@ -1,22 +0,0 @@ -from typing import TYPE_CHECKING - -if TYPE_CHECKING: - from typing import Dict, Optional - - from sentry_sdk.integrations.boto3._services.base import _ServiceExtension - - -# add a ServiceExtension here when one is implemented. for example: -# _SERVICE_EXTENSIONS = {"s3": _S3Extension()} -# when py 3.15 drops, we might want to take a look at using -# a lazy-loading approach using the new `lazy` keyword. -_SERVICE_EXTENSIONS: "Dict[str, _ServiceExtension]" = {} - - -def _resolve_service( - service_name: "Optional[str]", -) -> "Optional[_ServiceExtension]": - """Return the extension for a service, or `None` for generic instrumentation.""" - if service_name is None: - return None - return _SERVICE_EXTENSIONS.get(service_name) diff --git a/tests/integrations/boto3/test_bedrock.py b/tests/integrations/boto3/test_bedrock.py new file mode 100644 index 0000000000..5219bdf8b2 --- /dev/null +++ b/tests/integrations/boto3/test_bedrock.py @@ -0,0 +1,307 @@ +import boto3 +import pytest +from botocore.exceptions import ClientError +from botocore.loaders import Loader +from botocore.stub import Stubber + +import sentry_sdk +from sentry_sdk.consts import OP, SPANDATA +from sentry_sdk.integrations.boto3 import Boto3Integration + +pytestmark = pytest.mark.skipif( + "bedrock-runtime" not in Loader().list_available_services("service-2"), + reason="botocore is too old for bedrock-runtime.", +) + + +@pytest.fixture +def converse_response(): + return { + "output": {"message": {"role": "assistant", "content": [{"text": "Hello"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 50, "outputTokens": 30, "totalTokens": 80}, + "metrics": {"latencyMs": 500}, + } + + +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +def test_converse( + sentry_init, + capture_items, + capture_events, + span_streaming, + stream_gen_ai_spans, + converse_response, +): + sentry_init( + traces_sample_rate=1.0, + integrations=[Boto3Integration()], + stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", + ) + + session = boto3.Session( + aws_access_key_id="-", aws_secret_access_key="-", region_name="us-east-1" + ) + client = session.client("bedrock-runtime") + stubber = Stubber(client) + stubber.add_response("converse", converse_response) + + if span_streaming or stream_gen_ai_spans: + items = capture_items("span") + + with stubber, sentry_sdk.start_transaction(name="parent"): + client.converse( + modelId="anthropic.claude-3-sonnet-20240229-v1:0", + messages=[{"role": "user", "content": [{"text": "Hi"}]}], + ) + + sentry_sdk.flush() + spans = [item.payload for item in items] + (span,) = ( + span + for span in spans + if span["attributes"].get("sentry.op") == OP.GEN_AI_CHAT + ) + + assert span["attributes"][SPANDATA.GEN_AI_PROVIDER_NAME] == "aws.bedrock" + assert span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" + assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is False + else: + events = capture_events() + + with stubber, sentry_sdk.start_transaction(name="parent"): + client.converse( + modelId="anthropic.claude-3-sonnet-20240229-v1:0", + messages=[{"role": "user", "content": [{"text": "Hi"}]}], + ) + + (transaction,) = events + (span,) = transaction["spans"] + + assert span["data"][SPANDATA.GEN_AI_PROVIDER_NAME] == "aws.bedrock" + assert span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" + assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is False + + +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +def test_converse_error( + sentry_init, capture_items, capture_events, span_streaming, stream_gen_ai_spans +): + sentry_init( + traces_sample_rate=1.0, + integrations=[Boto3Integration()], + stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", + ) + + session = boto3.Session( + aws_access_key_id="-", aws_secret_access_key="-", region_name="us-east-1" + ) + client = session.client("bedrock-runtime") + original_exception = ClientError( + error_response={ + "Error": {"Code": "AccessDeniedException", "Message": "Access denied"}, + "ResponseMetadata": { + "RequestId": "test-id", + "HostId": "", + "HTTPStatusCode": 403, + "HTTPHeaders": {}, + "RetryAttempts": 1, + }, + }, + operation_name="Converse", + ) + + def raise_client_error(**kwargs): + raise original_exception + + client.meta.events.register("before-parameter-build", raise_client_error) + + if span_streaming or stream_gen_ai_spans: + items = capture_items("span") + + with sentry_sdk.start_transaction(name="parent"), pytest.raises(ClientError): + client.converse( + modelId="anthropic.claude-3-sonnet-20240229-v1:0", + messages=[{"role": "user", "content": [{"text": "Hi"}]}], + ) + + sentry_sdk.flush() + spans = [item.payload for item in items] + (span,) = ( + span + for span in spans + if span["attributes"].get("sentry.op") == OP.GEN_AI_CHAT + ) + + assert span["attributes"][SPANDATA.GEN_AI_PROVIDER_NAME] == "aws.bedrock" + assert span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" + + assert span["attributes"][SPANDATA.ERROR_TYPE] == "AccessDeniedException" + assert span["attributes"][SPANDATA.HTTP_STATUS_CODE] == 403 + assert span["attributes"][SPANDATA.AWS_REQUEST_ID] == "test-id" + assert span["attributes"][SPANDATA.HTTP_REQUEST_RESEND_COUNT] == 1 + assert span["status"] == "error" + else: + events = capture_events() + + with sentry_sdk.start_transaction(name="parent"), pytest.raises(ClientError): + client.converse( + modelId="anthropic.claude-3-sonnet-20240229-v1:0", + messages=[{"role": "user", "content": [{"text": "Hi"}]}], + ) + + (transaction,) = events + (span,) = transaction["spans"] + + assert span["data"][SPANDATA.GEN_AI_PROVIDER_NAME] == "aws.bedrock" + assert span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" + + assert span["data"][SPANDATA.ERROR_TYPE] == "AccessDeniedException" + assert span["data"][SPANDATA.HTTP_STATUS_CODE] == 403 + assert span["data"][SPANDATA.AWS_REQUEST_ID] == "test-id" + assert span["data"][SPANDATA.HTTP_REQUEST_RESEND_COUNT] == 1 + assert span["status"] == "internal_error" + + +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +def test_converse_stream( + sentry_init, capture_items, capture_events, span_streaming, stream_gen_ai_spans +): + sentry_init( + traces_sample_rate=1.0, + integrations=[Boto3Integration()], + stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", + ) + + session = boto3.Session( + aws_access_key_id="-", aws_secret_access_key="-", region_name="us-east-1" + ) + client = session.client("bedrock-runtime") + stubber = Stubber(client) + stubber.add_response("converse_stream", {}) + + if span_streaming or stream_gen_ai_spans: + items = capture_items("span") + + with stubber, sentry_sdk.start_transaction(name="parent"): + client.converse_stream( + modelId="anthropic.claude-3-sonnet-20240229-v1:0", + messages=[{"role": "user", "content": [{"text": "Hi"}]}], + ) + + sentry_sdk.flush() + spans = [item.payload for item in items] + (span,) = ( + span + for span in spans + if span["attributes"].get("sentry.op") == OP.GEN_AI_CHAT + ) + + assert span["attributes"][SPANDATA.GEN_AI_PROVIDER_NAME] == "aws.bedrock" + assert span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" + assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True + else: + events = capture_events() + + with stubber, sentry_sdk.start_transaction(name="parent"): + client.converse_stream( + modelId="anthropic.claude-3-sonnet-20240229-v1:0", + messages=[{"role": "user", "content": [{"text": "Hi"}]}], + ) + + (transaction,) = events + (span,) = transaction["spans"] + + assert span["data"][SPANDATA.GEN_AI_PROVIDER_NAME] == "aws.bedrock" + assert span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" + assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True + + +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +def test_converse_stream_error( + sentry_init, capture_items, capture_events, span_streaming, stream_gen_ai_spans +): + sentry_init( + traces_sample_rate=1.0, + integrations=[Boto3Integration()], + stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", + ) + + session = boto3.Session( + aws_access_key_id="-", aws_secret_access_key="-", region_name="us-east-1" + ) + client = session.client("bedrock-runtime") + original_exception = ClientError( + error_response={ + "Error": {"Code": "AccessDeniedException", "Message": "Access denied"}, + "ResponseMetadata": { + "RequestId": "test-id", + "HostId": "", + "HTTPStatusCode": 403, + "HTTPHeaders": {}, + "RetryAttempts": 1, + }, + }, + operation_name="ConverseStream", + ) + + def raise_client_error(**kwargs): + raise original_exception + + client.meta.events.register("before-parameter-build", raise_client_error) + + if span_streaming or stream_gen_ai_spans: + items = capture_items("span") + + with sentry_sdk.start_transaction(name="parent"), pytest.raises(ClientError): + client.converse_stream( + modelId="anthropic.claude-3-sonnet-20240229-v1:0", + messages=[{"role": "user", "content": [{"text": "Hi"}]}], + ) + + sentry_sdk.flush() + spans = [item.payload for item in items] + (span,) = ( + span + for span in spans + if span["attributes"].get("sentry.op") == OP.GEN_AI_CHAT + ) + + assert span["attributes"][SPANDATA.GEN_AI_PROVIDER_NAME] == "aws.bedrock" + assert span["attributes"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" + assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True + + assert span["attributes"][SPANDATA.ERROR_TYPE] == "AccessDeniedException" + assert span["attributes"][SPANDATA.HTTP_STATUS_CODE] == 403 + assert span["attributes"][SPANDATA.AWS_REQUEST_ID] == "test-id" + assert span["attributes"][SPANDATA.HTTP_REQUEST_RESEND_COUNT] == 1 + assert span["status"] == "error" + else: + events = capture_events() + + with sentry_sdk.start_transaction(name="parent"), pytest.raises(ClientError): + client.converse_stream( + modelId="anthropic.claude-3-sonnet-20240229-v1:0", + messages=[{"role": "user", "content": [{"text": "Hi"}]}], + ) + + (transaction,) = events + (span,) = transaction["spans"] + + assert span["data"][SPANDATA.GEN_AI_PROVIDER_NAME] == "aws.bedrock" + assert span["data"][SPANDATA.GEN_AI_OPERATION_NAME] == "chat" + assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True + + assert span["data"][SPANDATA.ERROR_TYPE] == "AccessDeniedException" + assert span["data"][SPANDATA.HTTP_STATUS_CODE] == 403 + assert span["data"][SPANDATA.AWS_REQUEST_ID] == "test-id" + assert span["data"][SPANDATA.HTTP_REQUEST_RESEND_COUNT] == 1 + assert span["status"] == "internal_error" diff --git a/tests/integrations/boto3/test_client.py b/tests/integrations/boto3/test_client.py index c309549316..0b864ac1e0 100644 --- a/tests/integrations/boto3/test_client.py +++ b/tests/integrations/boto3/test_client.py @@ -12,8 +12,6 @@ import sentry_sdk from sentry_sdk.consts import OP, SPANDATA from sentry_sdk.integrations.boto3 import Boto3Integration -from sentry_sdk.integrations.boto3._services.base import _ServiceExtension -from sentry_sdk.integrations.boto3._services.registry import _SERVICE_EXTENSIONS from sentry_sdk.integrations.boto3.consts import AWS_RPC_SYSTEM_NAME, ORIGIN from sentry_sdk.integrations.stdlib import StdlibIntegration from sentry_sdk.traces import StreamedSpan @@ -344,71 +342,6 @@ def _span_attributes(span, span_streaming): return span["attributes"] if span_streaming else span["data"] -@pytest.mark.parametrize("span_streaming", [True, False]) -def test_service_extension_customizes_client_span( - capture_items, - client_factory, - monkeypatch, - span_streaming, -): - class TestServiceExtension(_ServiceExtension): - def get_span_op(self, ctx): - return "aws.test" - - def get_span_origin(self, ctx): - return "auto.aws.test" - - def get_request_attributes(self, ctx): - return { - "aws.test.request": ctx.params["Key"], - SPANDATA.SENTRY_KIND: "producer", - } - - def get_response_attributes(self, ctx, response): - return { - "aws.test.response": response["ResponseMetadata"]["RequestId"], - SPANDATA.HTTP_STATUS_CODE: 418, - } - - monkeypatch.setitem(_SERVICE_EXTENSIONS, "s3", TestServiceExtension()) - client = client_factory() - api_params = {"Bucket": "bucket", "Key": "foo"} - - with Stubber(client) as stubber: - stubber.add_response( - "head_object", - { - "ResponseMetadata": { - "HTTPStatusCode": 200, - "RequestId": "request-id", - "HostId": "extended-request-id", - } - }, - api_params, - ) - spans_by_op = _capture_boto3_spans_by_op( - lambda: client.head_object(**api_params), - capture_items, - span_streaming, - expected_origin="auto.aws.test", - ) - - spans = spans_by_op.get("aws.test", []) - assert len(spans) == 1 - attributes = _span_attributes(spans[0], span_streaming) - assert attributes["aws.test.request"] == "foo" - assert attributes["aws.test.response"] == "request-id" - assert attributes[SPANDATA.SENTRY_KIND] == "producer" - assert attributes[SPANDATA.RPC_METHOD] == "HeadObject" - assert attributes[SPANDATA.HTTP_STATUS_CODE] == 200 - assert attributes[SPANDATA.AWS_EXTENDED_REQUEST_ID] == "extended-request-id" - _assert_span_finished(spans[0], span_streaming) - if span_streaming: - assert attributes[SPANDATA.SENTRY_ORIGIN] == "auto.aws.test" - else: - assert spans[0]["origin"] == "auto.aws.test" - - @pytest.mark.parametrize( ( "service_name", @@ -610,69 +543,6 @@ def attempt_failed_head_object_call(): assert attributes[SPANDATA.HTTP_REQUEST_RESEND_COUNT] == 1 -@pytest.mark.parametrize("span_streaming", [True, False]) -@pytest.mark.parametrize("with_service_extension", [False, True]) -def test_client_error_has_response_attributes_and_is_unchanged( - capture_items, - client_factory, - monkeypatch, - span_streaming, - with_service_extension, -): - class TestServiceExtension(_ServiceExtension): - def get_response_attributes(self, ctx, response): - return { - "aws.test.error": response["Error"]["Code"], - SPANDATA.ERROR_TYPE: "must-not-override", - SPANDATA.HTTP_STATUS_CODE: 418, - } - - if with_service_extension: - monkeypatch.setitem(_SERVICE_EXTENSIONS, "s3", TestServiceExtension()) - client = client_factory() - original_exception = ClientError( - { - "Error": { - "Code": "AccessDeniedException", - "Message": "must not become a span attribute", - }, - "ResponseMetadata": { - "RequestId": "request-id", - "HTTPStatusCode": 403, - "RetryAttempts": 1, - }, - }, - "HeadObject", - ) - - def raise_client_error(**kwargs): - raise original_exception - - client.meta.events.register("before-parameter-build", raise_client_error) - - def invoke_failing_client_method(): - with pytest.raises(ClientError) as exc_info: - client.head_object(Bucket="bucket", Key="foo") - assert exc_info.value is original_exception - - spans_by_op = _capture_boto3_spans_by_op( - invoke_failing_client_method, capture_items, span_streaming - ) - client_spans = spans_by_op.get(OP.HTTP_CLIENT, []) - _assert_one_failed_span(client_spans, span_streaming) - attributes = _span_attributes(client_spans[0], span_streaming) - - assert attributes[SPANDATA.AWS_REQUEST_ID] == "request-id" - assert attributes[SPANDATA.HTTP_STATUS_CODE] == 403 - assert attributes[SPANDATA.HTTP_REQUEST_RESEND_COUNT] == 1 - assert attributes[SPANDATA.ERROR_TYPE] == "AccessDeniedException" - if with_service_extension: - assert attributes["aws.test.error"] == "AccessDeniedException" - assert "Error.Message" not in attributes - assert "exception.message" not in attributes - assert "error.message" not in attributes - - @pytest.mark.parametrize( "event_name", [