feat(boto3): add ServiceExtension class

Author: pabloDeputter

sentry_sdk/integrations/boto3/_client.py

@@ -15,16 +15,22 @@
     _set_span_attributes,
     _start_client_span,
 )
+from sentry_sdk.integrations.boto3._services.registry import (
+    _resolve_service,
+)
 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")
 
@@ -76,9 +82,14 @@ 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)
@@ -88,17 +99,33 @@ def sentry_patched_make_api_call(
             _activate_client_span(span) if isinstance(span, StreamedSpan) else 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 `StreamedSpan` explicitly; static spans are finished by
             # their context manager.

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(Boto3Integration) is None:
@@ -195,15 +197,33 @@ def _start_client_span(
     # e.g. "aws.unkown.GetObject"
     service_name = ctx.service_id_hyphenized or "unknown"
     span_name = "aws.%s.%s" % (service_name, ctx.operation_name)
-    attributes: "Attributes" = {
-        SPANDATA.RPC_METHOD: ctx.operation_name,
-        SPANDATA.RPC_SYSTEM_NAME: _AWS_RPC_SYSTEM_NAME,
-    }
-    with capture_internal_exceptions():
-        attributes.update(_get_client_attributes(ctx))
+    attributes: "Attributes" = {}
     span_op = OP.HTTP_CLIENT
     span_origin = Boto3Integration.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

sentry_sdk/integrations/boto3/_services/__init__.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/base.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

sentry_sdk/integrations/boto3/_services/registry.py

@@ -16,6 +16,7 @@
     _get_error_attributes,
     _get_response_attributes,
 )
+from sentry_sdk.integrations.boto3._services.base import _ServiceExtension
 from sentry_sdk.integrations.stdlib import StdlibIntegration
 from tests.integrations.boto3.aws_mock import Body
 
@@ -301,6 +302,115 @@ 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.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.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"),
     [