feat(boto3): Add ServiceExtension class

sentry_sdk/consts.py

@@ -1206,6 +1206,12 @@ class SPANDATA:
     Example: "prod"
     """
 
+    SENTRY_KIND = "sentry.kind"
+    """
+    Used to clarify the relationship between parents and children, or to distinguish between spans, e.g. a `server` and `client` span with the same name.
+    Example: "client", "server", "producer", "consumer", "internal"
+    """
+
     SENTRY_OP = "sentry.op"
     """
     The operation of a span.

sentry_sdk/integrations/boto3/_client.py

@@ -14,17 +14,23 @@
     _set_span_attributes,
     _start_client_span,
 )
+from sentry_sdk.integrations.boto3._services.registry import (
+    _resolve_service,
+)
 from sentry_sdk.integrations.boto3.consts import IDENTIFIER
 from sentry_sdk.traces import NoOpStreamedSpan, StreamedSpan
 from sentry_sdk.utils import capture_internal_exceptions
 
 if TYPE_CHECKING:
     from typing import Any, Iterator, Optional, Union
 
+    from sentry_sdk._types import Attributes
+    from sentry_sdk.integrations.boto3._services.base import _ServiceExtension
     from sentry_sdk.tracing import Span
 
 try:
     from botocore.client import BaseClient
+    from botocore.exceptions import ClientError
 except ImportError:
     raise DidNotEnable("botocore not installed")
 
@@ -99,27 +105,48 @@ def sentry_patched_make_api_call(
         with capture_internal_exceptions():
             ctx.add_metadata(self)
 
+        service_ext: "Optional[_ServiceExtension]" = None
+        with capture_internal_exceptions():
+            # resolve service extension for service-specific enrichment.
+            service_ext = _resolve_service(ctx.service_name)
+
         span: "Optional[Union[Span, StreamedSpan]]" = None
         with capture_internal_exceptions():
-            span = _start_client_span(ctx)
+            span = _start_client_span(ctx, service_ext)
 
         if span is None:
             return orig_make_api_call(self, operation_name, api_params)
 
         # activate without finishing; a streaming response may outlive the call.
         span_ctx = _activate_client_span(span)
 
+        attributes: "Attributes" = {}
         try:
             with span_ctx:
                 try:
                     parsed = orig_make_api_call(self, operation_name, api_params)
                 except BaseException as error:
+                    if service_ext is not None and isinstance(error, ClientError):
+                        with capture_internal_exceptions():
+                            attributes.update(
+                                service_ext.get_response_attributes(ctx, error.response)
+                            )
+                    # generic attributes outweigh service-specific attributes.
                     with capture_internal_exceptions():
-                        _set_span_attributes(span, _get_error_attributes(error))
+                        attributes.update(_get_error_attributes(error))
                     raise
                 else:
+                    if service_ext is not None:
+                        with capture_internal_exceptions():
+                            attributes.update(
+                                service_ext.get_response_attributes(ctx, parsed)
+                            )
+                    with capture_internal_exceptions():
+                        attributes.update(_get_response_attributes(parsed))
+                finally:
+                    # enrich before the static span's context manager finishes it.
                     with capture_internal_exceptions():
-                        _set_span_attributes(span, _get_response_attributes(parsed))
+                        _set_span_attributes(span, attributes)
         except BaseException as error:
             _finish_span(span, error)
             raise

sentry_sdk/integrations/boto3/_instrumentation.py

@@ -24,6 +24,7 @@
 
     from sentry_sdk._types import Attributes
     from sentry_sdk.integrations.boto3._context import AwsCallContext
+    from sentry_sdk.integrations.boto3._services.base import _ServiceExtension
 
 try:
     from botocore.awsrequest import AWSRequest
@@ -186,6 +187,7 @@ def _get_error_attributes(exception: "BaseException") -> "Attributes":
 
 def _start_client_span(
     ctx: "AwsCallContext",
+    service_ext: "Optional[_ServiceExtension]" = None,
 ) -> "Optional[Union[Span, StreamedSpan]]":
     client = sentry_sdk.get_client()
     if client.get_integration(IDENTIFIER) is None:
@@ -198,12 +200,38 @@ def _start_client_span(
     attributes: "Attributes" = {
         SPANDATA.RPC_METHOD: ctx.operation_name,
         SPANDATA.RPC_SYSTEM_NAME: _AWS_RPC_SYSTEM_NAME,
+        # all client call spans are by default "client" spans.
+        # 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))
     span_op = OP.HTTP_CLIENT
     span_origin = ORIGIN
 
+    if service_ext is not None:
+        with capture_internal_exceptions():
+            config = service_ext.get_span_config(ctx)
+            if config is not None:
+                service_op, service_origin = config
+                if isinstance(service_op, str) and service_op:
+                    span_op = service_op
+                if isinstance(service_origin, str) and service_origin:
+                    span_origin = service_origin
+
+        with capture_internal_exceptions():
+            attributes.update(service_ext.get_request_attributes(ctx))
+
+    # Generic attributes take precedence over service-specific attributes.
+    attributes.update(
+        {
+            SPANDATA.RPC_METHOD: ctx.operation_name,
+            SPANDATA.RPC_SYSTEM_NAME: _AWS_RPC_SYSTEM_NAME,
+        }
+    )
+    with capture_internal_exceptions():
+        attributes.update(_get_client_attributes(ctx))
+
     if has_span_streaming_enabled(client.options):
         if sentry_sdk.traces.get_current_span() is None:
             return None
@@ -454,10 +482,9 @@ def _sentry_request_created(
         if span is None:
             return
 
-        # an ignored streamed span is not activated; avoid enriching its parent.
-        if isinstance(span, StreamedSpan):
-            if not (span.get_attributes().get(SPANDATA.SENTRY_ORIGIN) == ORIGIN):
-                return
+        # an ignored streamed span is not active; avoid enriching its parent.
+        if isinstance(span, StreamedSpan) and span.active:
+            return
 
         _set_request_attributes(span, request)
         # each attempt has a fresh `request.context`; carry the active client span.

sentry_sdk/integrations/boto3/_services/base.py

@@ -0,0 +1,31 @@
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any, Optional, Tuple
+
+    from sentry_sdk._types import Attributes
+    from sentry_sdk.integrations.boto3._context import AwsCallContext
+
+
+class _ServiceExtension:
+    """
+    Specialize generic botocore instrumentation for an AWS service.
+    """
+
+    __slots__ = ()
+
+    def get_span_config(
+        self, ctx: "AwsCallContext"
+    ) -> "Optional[Tuple[Optional[str], Optional[str]]]":
+        """Return an optional `(op, origin)` override for the client span."""
+        return None
+
+    def get_request_attributes(self, ctx: "AwsCallContext") -> "Attributes":
+        """Return service-specific attributes available before the call."""
+        return {}
+
+    def get_response_attributes(
+        self, ctx: "AwsCallContext", response: "Any"
+    ) -> "Attributes":
+        """Return service-specific attributes derived from the response."""
+        return {}

sentry_sdk/integrations/boto3/_services/registry.py

@@ -0,0 +1,33 @@
+from functools import lru_cache
+from importlib import import_module
+from typing import TYPE_CHECKING
+
+from sentry_sdk.integrations.boto3._services.base import _ServiceExtension
+from sentry_sdk.utils import capture_internal_exceptions
+
+if TYPE_CHECKING:
+    from typing import Dict, Optional, Tuple
+
+
+# service modules are imported lazily.
+# e.g. `s3` -> (`sentry_sdk.integrations.boto3._services.s3`, `_S3Extension)
+_SERVICE_EXTENSIONS: "Dict[str, Tuple[str, str]]" = {}
+
+
+@lru_cache(maxsize=None)
+def _resolve_service(
+    service: "str",
+) -> "Optional[_ServiceExtension]":
+    target = _SERVICE_EXTENSIONS.get(service)
+    if target is None:
+        return None
+
+    with capture_internal_exceptions():
+        module_name, class_name = target
+        extension_class = getattr(import_module(module_name), class_name)
+        candidate = extension_class()
+        if isinstance(candidate, _ServiceExtension):
+            return candidate
+
+    # preserve generic instrumentation when lookup fails.
+    return None

tests/integrations/boto3/test_client.py

@@ -17,6 +17,7 @@
     _get_response_attributes,
     _instrument_streaming_body,
 )
+from sentry_sdk.integrations.boto3._services.base import _ServiceExtension
 from sentry_sdk.integrations.boto3.consts import ORIGIN
 from sentry_sdk.integrations.stdlib import StdlibIntegration
 from sentry_sdk.traces import StreamedSpan
@@ -390,6 +391,117 @@ 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_config(self, ctx):
+            return ("aws.test", None)
+
+        def get_request_attributes(self, ctx):
+            return {
+                "aws.test.request": ctx.params["Key"],
+                SPANDATA.SENTRY_KIND: "producer",
+                SPANDATA.RPC_METHOD: "must-not-override",
+            }
+
+        def get_response_attributes(self, ctx, response):
+            return {
+                "aws.test.response": response["ResponseMetadata"]["RequestId"],
+                SPANDATA.HTTP_STATUS_CODE: 418,
+            }
+
+    extension = TestServiceExtension()
+    monkeypatch.setattr(
+        "sentry_sdk.integrations.boto3._client._resolve_service",
+        lambda service_name: extension,
+    )
+    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",
+                }
+            },
+            api_params,
+        )
+        spans_by_op = _capture_boto3_spans_by_op(
+            lambda: client.head_object(**api_params),
+            capture_items,
+            span_streaming,
+        )
+
+    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
+
+
+@pytest.mark.parametrize("span_streaming", [True, False])
+def test_service_extension_enriches_client_error(
+    capture_items,
+    client_factory,
+    monkeypatch,
+    span_streaming,
+):
+    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,
+            }
+
+    monkeypatch.setattr(
+        "sentry_sdk.integrations.boto3._client._resolve_service",
+        lambda service_name: TestServiceExtension(),
+    )
+    client = client_factory()
+    error = ClientError(
+        {
+            "Error": {"Code": "AccessDeniedException"},
+            "ResponseMetadata": {"HTTPStatusCode": 403},
+        },
+        "HeadObject",
+    )
+
+    def raise_client_error(**kwargs):
+        raise error
+
+    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 error
+
+    spans_by_op = _capture_boto3_spans_by_op(
+        invoke_failing_client_method,
+        capture_items,
+        span_streaming,
+    )
+    spans = spans_by_op.get(OP.HTTP_CLIENT, [])
+
+    _assert_one_failed_span(spans, span_streaming)
+    attributes = _span_attributes(spans[0], span_streaming)
+    assert attributes["aws.test.error"] == "AccessDeniedException"
+    assert attributes[SPANDATA.ERROR_TYPE] == "AccessDeniedException"
+    assert attributes[SPANDATA.HTTP_STATUS_CODE] == 403
+
+
 @pytest.mark.parametrize(
     ("response", "expected"),
     [
@@ -574,6 +686,7 @@ def test_client_call_has_common_attributes(
     assert attributes[SPANDATA.RPC_SERVICE] == rpc_service
     assert attributes[SPANDATA.RPC_METHOD] == rpc_method
     assert attributes[SPANDATA.RPC_SYSTEM_NAME] == "aws-api"
+    assert attributes[SPANDATA.SENTRY_KIND] == "client"
     assert attributes[SPANDATA.CLOUD_REGION] == "eu-north-1"
     assert attributes[SPANDATA.SERVER_ADDRESS] == server_address
     assert attributes[SPANDATA.SERVER_PORT] == server_port