feat(boto3): add error, retry and response attribute enrichment

Author: pabloDeputter

sentry_sdk/integrations/boto3/_instrumentation.py

@@ -2,6 +2,7 @@
 from urllib.parse import urlsplit
 
 from botocore.awsrequest import AWSRequest
+from botocore.exceptions import ClientError
 from botocore.response import StreamingBody
 
 import sentry_sdk
@@ -33,6 +34,7 @@
 def _set_span_attributes(
     span: "Union[Span, StreamedSpan]", attributes: "Attributes"
 ) -> None:
+    """Will be removed in the major."""
     if isinstance(span, StreamedSpan):
         span.set_attributes(attributes)
         return
@@ -85,6 +87,95 @@ def _get_client_attributes(
     return attributes
 
 
+def _get_response_attributes(response: "Any") -> "Attributes":
+    if not isinstance(response, dict):
+        return {}
+
+    metadata = response.get("ResponseMetadata")
+    if not isinstance(metadata, dict):
+        return {}
+    attributes: "Attributes" = {}
+
+    # botocore injects HTTP status into `ResponseMetadata` after parsing.
+    # https://github.com/boto/botocore/blob/develop/botocore/parsers.py#L273-L284
+    status_code = metadata.get("HTTPStatusCode")
+    if isinstance(status_code, int) and 100 <= status_code <= 599:
+        attributes[SPANDATA.HTTP_STATUS_CODE] = status_code
+
+    retry_attempts = metadata.get("RetryAttempts")
+    # botocore represents retries as `attempts - 1`; OTel suggests "if and only if", so skip zero.
+    # https://github.com/boto/botocore/blob/develop/botocore/endpoint.py#L221-L229
+    # https://opentelemetry.io/docs/specs/semconv/http/http-spans/#http-client-span
+    if (
+        isinstance(retry_attempts, int)
+        # avoid emitting `resend_count=True`.
+        and not isinstance(retry_attempts, bool)
+        and retry_attempts > 0
+    ):
+        attributes[SPANDATA.HTTP_REQUEST_RESEND_COUNT] = retry_attempts
+
+    headers = metadata.get("HTTPHeaders")
+    if not isinstance(headers, dict):
+        headers = {}
+
+    request_id = metadata.get("RequestId")
+    if not isinstance(request_id, str) or not request_id:
+        request_id = next(
+            (
+                value
+                for value in (
+                    headers.get("x-amzn-requestid"),
+                    headers.get("x-amzn-request-id"),
+                    headers.get("x-amz-request-id"),
+                )
+                if isinstance(value, str) and value
+            ),
+            None,
+        )
+    if isinstance(request_id, str) and request_id:
+        attributes[SPANDATA.AWS_REQUEST_ID] = request_id
+
+    # S3's `HostId` is the extended request ID returned in `x-amz-id-2`.
+    # https://docs.aws.amazon.com/AmazonS3/latest/developerguide/get-request-ids.html
+    extended_request_id = metadata.get("HostId")
+    if not isinstance(extended_request_id, str) or not extended_request_id:
+        extended_request_id = headers.get("x-amz-id-2")
+    if isinstance(extended_request_id, str) and extended_request_id:
+        attributes[SPANDATA.AWS_EXTENDED_REQUEST_ID] = extended_request_id
+
+    return attributes
+
+
+def _get_error_type(exception: "BaseException") -> str:
+    if isinstance(exception, ClientError):
+        # `ClientError` wraps AWS service errors; `Error.Code` identifies the
+        # actual service error, e.g. `AccessDeniedException`.
+        # https://docs.aws.amazon.com/boto3/latest/guide/error-handling.html
+        error = exception.response.get("Error")
+        if isinstance(error, dict):
+            error_code = error.get("Code")
+            if isinstance(error_code, str) and error_code:
+                return error_code
+
+    # failures before a service response have no AWS error code.
+    # https://opentelemetry.io/docs/specs/semconv/rpc/rpc-spans/
+    exception_type = type(exception)
+    exception_name = exception_type.__qualname__
+    exception_module = exception_type.__module__
+    if exception_module not in ("builtins", "__builtins__"):
+        return "%s.%s" % (exception_module, exception_name)
+    return exception_name
+
+
+def _get_error_attributes(exception: "BaseException") -> "Attributes":
+    attributes: "Attributes" = {}
+    if isinstance(exception, ClientError):
+        attributes.update(_get_response_attributes(exception.response))
+
+    attributes[SPANDATA.ERROR_TYPE] = _get_error_type(exception)
+    return attributes
+
+
 def _start_client_span(
     ctx: "AwsCallContext",
 ) -> "Optional[Union[Span, StreamedSpan]]":
@@ -129,11 +220,18 @@ def _finish_client_span(
     span: "Union[Span, StreamedSpan]",
     parsed: "Dict[str, Any]",
 ) -> None:
-    span.__exit__(None, None, None)
-
     if isinstance(span, NoOpStreamedSpan):
+        span.__exit__(None, None, None)
         return
 
+    attributes: "Attributes" = {}
+    with capture_internal_exceptions():
+        attributes = _get_response_attributes(parsed)
+
+    with capture_internal_exceptions():
+        _set_span_attributes(span, attributes)
+    span.__exit__(None, None, None)
+
     body = parsed.get("Body")
     if not isinstance(body, StreamingBody):
         return
@@ -171,11 +269,15 @@ def sentry_streaming_body_read(*args: "Any", **kwargs: "Any") -> bytes:
                 streaming_span.finish()
             return ret
         except Exception as exc:
-            if isinstance(streaming_span, StreamedSpan):
-                streaming_span.__exit__(type(exc), exc, exc.__traceback__)
-            else:
-                streaming_span.set_status(SPANSTATUS.INTERNAL_ERROR)
-                streaming_span.finish()
+            with capture_internal_exceptions():
+                _set_span_attributes(streaming_span, _get_error_attributes(exc))
+
+            with capture_internal_exceptions():
+                if isinstance(streaming_span, StreamedSpan):
+                    streaming_span.__exit__(type(exc), exc, exc.__traceback__)
+                else:
+                    streaming_span.set_status(SPANSTATUS.INTERNAL_ERROR)
+                    streaming_span.finish()
             raise
 
     body.read = sentry_streaming_body_read  # type: ignore
@@ -194,6 +296,12 @@ def _finish_client_span_with_error(
     span: "Union[Span, StreamedSpan]",
     exception: "BaseException",
 ) -> None:
+    attributes: "Attributes" = {}
+    with capture_internal_exceptions():
+        attributes = _get_error_attributes(exception)
+
+    with capture_internal_exceptions():
+        _set_span_attributes(span, attributes)
     span.__exit__(type(exception), exception, exception.__traceback__)