feat(boto3): Add ServiceExtension class

sentry_sdk/consts.py

@@ -1214,6 +1214,12 @@ class SPANDATA:
     Used in inbound filters.
     """
 
+    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,22 @@
     _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 typing import Any, Dict, Iterator, Optional, Union
 
+    from sentry_sdk._types import Attributes
     from sentry_sdk.tracing import Span
 
 try:
     from botocore.client import BaseClient
+    from botocore.exceptions import ClientError
 except ImportError:
     raise DidNotEnable("botocore not installed")
 
@@ -34,12 +39,12 @@ def _activate_client_span(
     span: "Union[Span, StreamedSpan]",
 ) -> "Iterator[Union[Span, StreamedSpan]]":
     """
-    Activate the boto span temporarily during `_make_api_call()` without ending it.
+    Activate the client span temporarily during `_make_api_call()` without ending it.
 
     Botocore returns a `StreamingBody` before its bytes are consumed. Using the
     context manager would finish it as soon as `_make_api_call()` returns, so
     restore the caller's span here and let the `StreamingBody` wrapper finish
-    the boto span when body is consumed/closed.
+    the client span when the body is consumed or closed.
 
     faulty:                               desired:
            boto3  [_make_api_call]                boto3  [_make_api_call------]
@@ -74,13 +79,13 @@ def _patch_botocore_client() -> None:
 
     def sentry_patched_init(self: "BaseClient", *args: "Any", **kwargs: "Any") -> None:
         orig_init(self, *args, **kwargs)
-        meta = self.meta
-        meta.events.register("request-created", _sentry_request_created)
-        # run after other `before-sign` handlers so existing baggage is preserved.
-        meta.events.register_last("before-sign", _sentry_before_sign)
+        with capture_internal_exceptions():
+            self.meta.events.register("request-created", _sentry_request_created)
+            # run after other `before-sign` handlers so existing baggage is preserved.
+            self.meta.events.register_last("before-sign", _sentry_before_sign)
 
     def sentry_patched_make_api_call(
-        self: "BaseClient", operation_name: str, api_params: "Any"
+        self: "BaseClient", operation_name: str, api_params: "Dict[str, Any]"
     ) -> "Any":
         """
         Track a single API call, including retries, serialization, and endpoint
@@ -93,32 +98,46 @@ def sentry_patched_make_api_call(
         if client.get_integration(IDENTIFIER) is None:
             return orig_make_api_call(self, operation_name, api_params)
 
-        ctx = AwsCallContext(operation_name, api_params)
-
-        # add optional metadata to context.
-        with capture_internal_exceptions():
-            ctx.add_metadata(self)
-
         span: "Optional[Union[Span, StreamedSpan]]" = None
         with capture_internal_exceptions():
-            span = _start_client_span(ctx)
+            ctx = AwsCallContext(operation_name, api_params)
+            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)
+            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)
+                            )
+                    with capture_internal_exceptions():
+                        attributes.update(_get_error_attributes(error))
                     with capture_internal_exceptions():
-                        _set_span_attributes(span, _get_error_attributes(error))
+                        _set_span_attributes(span, attributes)
                     raise
+                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))
                 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/_context.py

@@ -23,40 +23,27 @@ class AwsCallContext:
         "params",
     )
 
-    def __init__(self, operation_name: str, params: "Any") -> None:
-        self.operation_name: str = operation_name
-        self.params: "Dict[str, Any]" = {}
+    def __init__(self, operation_name: str, params: "Dict[str, Any]") -> None:
+        self.operation_name: "str" = operation_name
+        self.params: "Dict[str, Any]" = dict(params)
         self.service_name: "Optional[str]" = None
         self.service_id: "Optional[str]" = None
         self.service_id_hyphenized: "Optional[str]" = None
         self.region_name: "Optional[str]" = None
         self.endpoint_url: "Optional[str]" = None
 
-        if isinstance(params, dict):
-            with capture_internal_exceptions():
-                self.params = dict(params)
-
     def add_metadata(self, client: "BaseClient") -> None:
-        def _get_attr(obj: "Any", name: str) -> "Any":
-            if obj is None:
-                return None
-
-            with capture_internal_exceptions():
-                return getattr(obj, name)
-
-        client_meta = _get_attr(client, "meta")
-        service_model = _get_attr(client_meta, "service_model")
-
-        # botocore's internal identifier, e.g. `apigateway`.
-        self.service_name = _get_attr(service_model, "service_name")
-
-        # modeled AWS service identity used in span names, e.g. `API Gateway`.
-        service_id = _get_attr(service_model, "service_id")
-        if service_id is not None:
-            with capture_internal_exceptions():
-                self.service_id = str(service_id)
-            with capture_internal_exceptions():
-                self.service_id_hyphenized = service_id.hyphenize()
-
-        self.region_name = _get_attr(client_meta, "region_name")
-        self.endpoint_url = _get_attr(client_meta, "endpoint_url")
+        with capture_internal_exceptions():
+            service_model = client.meta.service_model
+            # botocore's internal identifier, e.g. `apigateway`.
+            self.service_name = service_model.service_name
+            service_id = service_model.service_id
+            # modeled AWS service identity used in span names, e.g. `API Gateway`.
+            self.service_id = str(service_id)
+            self.service_id_hyphenized = service_id.hyphenize()
+
+        with capture_internal_exceptions():
+            self.region_name = client.meta.region_name
+
+        with capture_internal_exceptions():
+            self.endpoint_url = client.meta.endpoint_url