feat(boto3): Add S3 extension

sentry_sdk/consts.py

@@ -369,6 +369,42 @@ class SPANDATA:
     Example: "79b9da39-b7ae-508a-a6bc-864b2829c622"
     """
 
+    AWS_S3_BUCKET = "aws.s3.bucket"
+    """
+    The S3 bucket name the request refers to.
+    Example: "ot-demo-test"
+    """
+
+    AWS_S3_COPY_SOURCE = "aws.s3.copy_source"
+    """
+    The source object (in the form bucket/key) for the copy operation.
+    Example: "someFile.yml"
+    """
+
+    AWS_S3_DELETE = "aws.s3.delete"
+    """
+    The delete request container that specifies the objects to be deleted.
+    Example: "Objects=[{Key=string,VersionId=string},{Key=string,VersionId=string}],Quiet=boolean"
+    """
+
+    AWS_S3_KEY = "aws.s3.key"
+    """
+    The S3 object key the request refers to. Corresponds to the --key parameter of the S3 API operations.
+    Example: "someFile.yml"
+    """
+
+    AWS_S3_PART_NUMBER = "aws.s3.part_number"
+    """
+    The part number of the part being uploaded in a multipart-upload operation. This is a positive integer between 1 and 10,000.
+    Example: 3456
+    """
+
+    AWS_S3_UPLOAD_ID = "aws.s3.upload_id"
+    """
+    Upload ID that identifies the multipart upload.
+    Example: "dfRtDYWFbkRONycy.Yxwh66Yjlx.cph0gtNBtJ"
+    """
+
     CACHE_HIT = "cache.hit"
     """
     A boolean indicating whether the requested data was found in the cache.
@@ -539,6 +575,12 @@ class SPANDATA:
     Example: "timeout"
     """
 
+    FILE_SIZE = "file.size"
+    """
+    File size in bytes.
+    Example: 1024
+    """
+
     GEN_AI_AGENT_NAME = "gen_ai.agent.name"
     """
     The name of the agent being used.
@@ -908,6 +950,12 @@ class SPANDATA:
     Example: ?foo=bar&bar=baz
     """
 
+    HTTP_BODY_SIZE = "http.response.body.size"
+    """
+    The encoded body size of the response (in bytes).
+    Example: 123
+    """
+
     HTTP_STATUS_CODE = "http.response.status_code"
     """
     The HTTP status code as an integer.

sentry_sdk/integrations/boto3/_client.py

@@ -100,11 +100,9 @@ def sentry_patched_make_api_call(
             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:
+            with _activate_client_span(span):
                 try:
                     parsed = orig_make_api_call(self, operation_name, api_params)
                 except BaseException as error:

sentry_sdk/integrations/boto3/_instrumentation.py

@@ -207,7 +207,13 @@ def _instrument_streaming_body(span: "Span", parsed: "Dict[str, Any]") -> bool:
     if isinstance(span, NoOpSpan):
         return False
 
-    body = parsed.get("Body")
+    # botocore uses service-specific response key for streaming payload; e.g. `Payload`
+    # for Lambda Invoke, `Body` for S3 GetObject. Find it by type so every streaming
+    # response is handled.
+    body = next(
+        (value for value in parsed.values() if isinstance(value, StreamingBody)),
+        None,
+    )
     if not isinstance(body, StreamingBody):
         return False
 

sentry_sdk/integrations/boto3/_services/registry.py

@@ -1,5 +1,7 @@
 from typing import TYPE_CHECKING
 
+from sentry_sdk.integrations.boto3._services.s3 import _S3Extension
+
 if TYPE_CHECKING:
     from typing import Dict, Optional
 
@@ -10,7 +12,9 @@
 #   _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]" = {}
+_SERVICE_EXTENSIONS: "Dict[str, _ServiceExtension]" = {
+    "s3": _S3Extension(),
+}
 
 
 def _resolve_service(

sentry_sdk/integrations/boto3/_services/s3.py

@@ -0,0 +1,107 @@
+import json
+from datetime import datetime
+from typing import TYPE_CHECKING
+
+from sentry_sdk.consts import SPANDATA
+from sentry_sdk.integrations.boto3._services.base import _ServiceExtension
+from sentry_sdk.integrations.boto3._utils import (
+    _extract_attributes,
+)
+from sentry_sdk.utils import capture_internal_exceptions
+
+if TYPE_CHECKING:
+    from typing import Any, Sequence
+
+    from sentry_sdk._types import Attributes
+    from sentry_sdk.integrations.boto3._context import AwsCallContext
+    from sentry_sdk.integrations.boto3._utils import (
+        _AttributeSpec,
+    )
+
+# maps s3 operations to the response field that contains the complete file size.
+_RESPONSE_FILE_SIZE_FIELDS = {
+    "GetObjectAttributes": "ObjectSize",
+    "PutObject": "Size",
+}
+
+# s3 request attributes that are extracted when present, regardless of operation.
+_REQUEST_ATTRIBUTES: "Sequence[_AttributeSpec]" = (
+    ("Bucket", SPANDATA.AWS_S3_BUCKET),
+    ("Key", SPANDATA.AWS_S3_KEY),
+    ("UploadId", SPANDATA.AWS_S3_UPLOAD_ID),
+)
+
+
+class _S3Extension(_ServiceExtension):
+    __slots__ = ()
+
+    def get_request_attributes(self, ctx: "AwsCallContext") -> "Attributes":
+        attributes: "Attributes" = _extract_attributes(ctx.params, _REQUEST_ATTRIBUTES)
+
+        if "CopySource" in ctx.params:
+            with capture_internal_exceptions():
+                # boto3 accepts either "bucket/key" or a dict {"bucket": ..., "Key": ..., "VersionId": ...} for `CopySource`.
+                # https://docs.aws.amazon.com/boto3/latest/reference/services/s3/client/upload_part_copy.html
+                copy_source = ctx.params["CopySource"]
+                if isinstance(copy_source, str):
+                    attributes[SPANDATA.AWS_S3_COPY_SOURCE] = copy_source
+                else:
+                    value = f"{copy_source['Bucket']}/{copy_source['Key']}"
+                    if "VersionId" in copy_source:
+                        value += f"?versionId={copy_source['VersionId']}"
+                    attributes[SPANDATA.AWS_S3_COPY_SOURCE] = value
+
+        # OTel defines `PartNumber` for `UploadPart` and `UploadPartCopy` only.
+        # https://opentelemetry.io/docs/specs/semconv/object-stores/s3/#attributes
+        if (
+            ctx.operation_name in ("UploadPart", "UploadPartCopy")
+            and "PartNumber" in ctx.params
+        ):
+            attributes[SPANDATA.AWS_S3_PART_NUMBER] = ctx.params["PartNumber"]
+
+        if "Delete" in ctx.params:
+            with capture_internal_exceptions():
+                attributes[SPANDATA.AWS_S3_DELETE] = json.dumps(
+                    ctx.params["Delete"],
+                    default=lambda value: (
+                        value.isoformat() if isinstance(value, datetime) else str(value)
+                    ),
+                    separators=(",", ":"),
+                    sort_keys=True,
+                )
+
+        if (
+            ctx.operation_name == "CompleteMultipartUpload"
+            and "MpuObjectSize" in ctx.params
+        ):
+            attributes[SPANDATA.FILE_SIZE] = ctx.params["MpuObjectSize"]
+
+        return attributes
+
+    def get_response_attributes(
+        self, ctx: "AwsCallContext", response: "Any"
+    ) -> "Attributes":
+        attributes: "Attributes" = {}
+
+        if (
+            ctx.operation_name in ("GetObject", "GetObjectAnnotation")
+            and "ContentLength" in response
+        ):
+            # `ContentLength` is the size of the HTTP body returned, which may be a range.
+            attributes[SPANDATA.HTTP_BODY_SIZE] = response["ContentLength"]
+
+        # report the complete file size, not just the HTTP body size.
+        file_size_field = _RESPONSE_FILE_SIZE_FIELDS.get(ctx.operation_name)
+        if file_size_field is not None and file_size_field in response:
+            attributes[SPANDATA.FILE_SIZE] = response[file_size_field]
+
+        if (
+            ctx.operation_name == "HeadObject"
+            and "Range" not in ctx.params
+            and "PartNumber" not in ctx.params
+            and "ContentLength" in response
+        ):
+            # an un-ranged `HEAD` has no body, so `ContentLength` is the file size.
+            attributes[SPANDATA.FILE_SIZE] = response["ContentLength"]
+
+        return attributes

sentry_sdk/integrations/boto3/_utils.py

@@ -0,0 +1,19 @@
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any, Dict, Sequence, Tuple
+
+    from sentry_sdk._types import Attributes
+
+    # tuple of (request param name, span attribute name).
+    _AttributeSpec = Tuple[str, str]
+
+
+def _extract_attributes(
+    source: "Dict[str, Any]", specs: "Sequence[_AttributeSpec]"
+) -> "Attributes":
+    attributes = {}
+    for param, attribute in specs:
+        if param in source:
+            attributes[attribute] = source[param]
+    return attributes

tests/integrations/boto3/helpers.py

@@ -0,0 +1,115 @@
+from types import SimpleNamespace
+from typing import TYPE_CHECKING, Dict, List, Optional
+
+import boto3
+import pytest
+from botocore.config import Config
+
+import sentry_sdk
+from sentry_sdk.consts import SPANDATA
+from sentry_sdk.integrations.boto3 import Boto3Integration
+from sentry_sdk.integrations.boto3.consts import ORIGIN
+
+if TYPE_CHECKING:
+    from sentry_sdk._types import SpanJSON
+
+
+@pytest.fixture
+def client_factory(sentry_init, monkeypatch):
+    sentry_init(
+        traces_sample_rate=1.0,
+        integrations=[Boto3Integration()],
+    )
+    session = boto3.Session(  # type: ignore
+        aws_access_key_id="-",
+        aws_secret_access_key="-",
+        region_name="eu-north-1",
+    )
+    clients = []
+
+    def make_client(service_name="s3", attempt_count=1, **client_kwargs):
+        client = session.client(
+            service_name,
+            config=Config(
+                retries={"total_max_attempts": attempt_count, "mode": "standard"}
+            ),
+            **client_kwargs,
+        )  # type: ignore
+        clients.append(client)
+        return client
+
+    yield make_client
+
+    for client in clients:
+        # older supported botocore versions do not expose `BaseClient.close()`.
+        close = getattr(client, "close", None)
+        if close is not None:
+            close()
+
+
+@pytest.fixture
+def no_botocore_retry_delay(monkeypatch):
+    # remove request retry delays without replacing botocore's retry handling.
+    monkeypatch.setattr(
+        "botocore.endpoint.time",
+        SimpleNamespace(sleep=lambda delay: None),
+    )
+
+
+@pytest.fixture
+def s3_client(client_factory):
+    return client_factory("s3")
+
+
+def require_botocore_model_fields(
+    client,
+    method,
+    input_fields=(),
+    output_fields=(),
+):
+    """Skip tests when botocore lacks required fields, including nested paths."""
+
+    def has_field(shape, field):
+        for part in field.split("."):
+            while shape is not None and shape.type_name == "list":
+                shape = shape.member
+            if shape is None or part not in getattr(shape, "members", {}):
+                return False
+            shape = shape.members[part]
+        return True
+
+    operation_name = client.meta.method_to_api_mapping.get(method)
+    if operation_name is None:
+        pytest.skip("%s is absent from this botocore model; skipping test" % method)
+    model = client.meta.service_model.operation_model(operation_name)  # type: ignore
+    for shape, fields in (
+        (model.input_shape, input_fields),
+        (model.output_shape, output_fields),
+    ):
+        for field in fields:
+            if not has_field(shape, field):
+                pytest.skip(
+                    "%s.%s is absent from this botocore model; skipping test"
+                    % (method, field)
+                )
+
+
+def capture_spans_by_op(
+    invoke_client_method,
+    capture_items,
+    expected_origin=ORIGIN,
+):
+    items = capture_items("span")
+
+    with sentry_sdk.start_span(name="parent"):
+        invoke_client_method()
+
+    sentry_sdk.flush()
+    spans_by_op: Dict[Optional[str], List["SpanJSON"]] = {}
+    for item in items:
+        span = item.payload
+        if span["attributes"].get(SPANDATA.SENTRY_ORIGIN) == expected_origin:
+            spans_by_op.setdefault(
+                span["attributes"].get(SPANDATA.SENTRY_OP), []
+            ).append(span)
+    return spans_by_op

tests/integrations/boto3/s3_list.xml

@@ -1,2 +1,29 @@
 <?xml version="1.0" encoding="UTF-8"?>
-<ListBucketResult xmlns="http://s3.amazonaws.com/doc/2006-03-01/"><Name>marshalls-furious-bucket</Name><Prefix></Prefix><Marker></Marker><MaxKeys>1000</MaxKeys><EncodingType>url</EncodingType><IsTruncated>false</IsTruncated><Contents><Key>foo.txt</Key><LastModified>2020-10-24T00:13:39.000Z</LastModified><ETag>&quot;a895ba674b4abd01b5d67cfd7074b827&quot;</ETag><Size>206453</Size><Owner><ID>7bef397f7e536914d1ff1bbdb105ed90bcfd06269456bf4a06c6e2e54564daf7</ID></Owner><StorageClass>STANDARD</StorageClass></Contents><Contents><Key>bar.txt</Key><LastModified>2020-10-02T15:15:20.000Z</LastModified><ETag>&quot;a895ba674b4abd01b5d67cfd7074b827&quot;</ETag><Size>206453</Size><Owner><ID>7bef397f7e536914d1ff1bbdb105ed90bcfd06269456bf4a06c6e2e54564daf7</ID></Owner><StorageClass>STANDARD</StorageClass></Contents></ListBucketResult>
+<ListBucketResult xmlns="http://s3.amazonaws.com/doc/2006-03-01/">
+    <Name>marshalls-furious-bucket</Name>
+    <Prefix></Prefix>
+    <Marker></Marker>
+    <MaxKeys>1000</MaxKeys>
+    <EncodingType>url</EncodingType>
+    <IsTruncated>false</IsTruncated>
+    <Contents>
+        <Key>foo.txt</Key>
+        <LastModified>2020-10-24T00:13:39.000Z</LastModified>
+        <ETag>&quot;a895ba674b4abd01b5d67cfd7074b827&quot;</ETag>
+        <Size>206453</Size>
+        <Owner>
+            <ID>7bef397f7e536914d1ff1bbdb105ed90bcfd06269456bf4a06c6e2e54564daf7</ID>
+        </Owner>
+        <StorageClass>STANDARD</StorageClass>
+    </Contents>
+    <Contents>
+        <Key>bar.txt</Key>
+        <LastModified>2020-10-02T15:15:20.000Z</LastModified>
+        <ETag>&quot;a895ba674b4abd01b5d67cfd7074b827&quot;</ETag>
+        <Size>206453</Size>
+        <Owner>
+            <ID>7bef397f7e536914d1ff1bbdb105ed90bcfd06269456bf4a06c6e2e54564daf7</ID>
+        </Owner>
+        <StorageClass>STANDARD</StorageClass>
+    </Contents>
+</ListBucketResult>