feat(boto3): Add SQS/SNS trace propagation

sentry_sdk/consts.py

@@ -1393,6 +1393,12 @@ class SPANDATA:
     Example: "client", "server", "producer", "consumer", "internal"
     """
 
+    SENTRY_LINK_TYPE = "sentry.link.type"
+    """
+    Set on a span link. Describes the relationship between the span and the linked span.
+    Example: "message_origin"
+    """
+
     SENTRY_OP = "sentry.op"
     """
     The operation of a span.

sentry_sdk/integrations/boto3/_client.py

@@ -96,6 +96,10 @@ def sentry_patched_make_api_call(
                 ctx.add_metadata(self)
             service_ext = _resolve_service(ctx.service_name)
             span = _start_client_span(ctx, service_ext)
+            if service_ext is not None:
+                with capture_internal_exceptions():
+                    # override the request params for e.g. trace propagation.
+                    api_params = service_ext.get_request_params(ctx, span) or api_params
 
         if span is None:
             return orig_make_api_call(self, operation_name, api_params)
@@ -122,6 +126,8 @@ def sentry_patched_make_api_call(
                         attributes.update(
                             service_ext.get_response_attributes(ctx, parsed)
                         )
+                    with capture_internal_exceptions():
+                        span.add_links(service_ext.get_response_links(ctx, parsed))
                 with capture_internal_exceptions():
                     attributes.update(_get_response_attributes(parsed))
                 with capture_internal_exceptions():

sentry_sdk/integrations/boto3/_services/base.py

@@ -1,10 +1,11 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Any, Optional
+    from typing import Any, Dict, List, Optional
 
-    from sentry_sdk._types import Attributes
+    from sentry_sdk._types import Attributes, SpanLink
     from sentry_sdk.integrations.boto3._context import AwsCallContext
+    from sentry_sdk.traces import Span
 
 
 class _ServiceExtension:
@@ -36,3 +37,18 @@ def get_response_attributes(
     ) -> "Attributes":
         """Return response attributes to add after the AWS request is made."""
         return {}
+
+    def get_response_links(
+        self, ctx: "AwsCallContext", response: "Any"
+    ) -> "List[SpanLink]":
+        """Return links to add to the span after the AWS request is made."""
+        return []
+
+    def get_request_params(
+        self, ctx: "AwsCallContext", span: "Optional[Span]"
+    ) -> "Optional[Dict[str, Any]]":
+        """
+        Return params to send instead of `ctx.params`, e.g. with trace propagation added,
+        or `None` to send them unchanged; must not mutate `ctx.params`.
+        """
+        return None

sentry_sdk/integrations/boto3/_services/sns.py

@@ -3,13 +3,17 @@
 import sentry_sdk
 from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.integrations.boto3._services.base import _ServiceExtension
-from sentry_sdk.integrations.boto3._utils import _get_aws_arn_attributes
+from sentry_sdk.integrations.boto3._utils import (
+    _get_aws_arn_attributes,
+    _inject_message_attributes,
+)
 
 if TYPE_CHECKING:
-    from typing import Any, Optional
+    from typing import Any, Dict, Optional
 
     from sentry_sdk._types import Attributes
     from sentry_sdk.integrations.boto3._context import AwsCallContext
+    from sentry_sdk.traces import Span
 
 _PUBLISH_OPERATIONS = ("Publish", "PublishBatch")
 
@@ -114,3 +118,13 @@ def get_response_attributes(
         elif error_codes:
             attributes[SPANDATA.ERROR_TYPE] = "aws_batch_partial_failure"
         return attributes
+
+    def get_request_params(
+        self, ctx: "AwsCallContext", span: "Optional[Span]"
+    ) -> "Optional[Dict[str, Any]]":
+        if ctx.operation_name not in _PUBLISH_OPERATIONS:
+            return None
+
+        return _inject_message_attributes(
+            ctx, span, "Message", "PublishBatchRequestEntries"
+        )

sentry_sdk/integrations/boto3/_services/sqs.py

@@ -1,17 +1,24 @@
+import json
 import re
 from typing import TYPE_CHECKING
 from urllib.parse import unquote, urlsplit
 
 from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.integrations.boto3._services.base import _ServiceExtension
-from sentry_sdk.integrations.boto3._utils import _get_aws_arn_attributes
+from sentry_sdk.integrations.boto3._utils import (
+    _get_aws_arn_attributes,
+    _inject_message_attributes,
+)
+from sentry_sdk.traces import BAGGAGE_HEADER_NAME, SENTRY_TRACE_HEADER_NAME, NoOpSpan
+from sentry_sdk.tracing_utils import extract_sentrytrace_data
 from sentry_sdk.utils import capture_internal_exceptions
 
 if TYPE_CHECKING:
-    from typing import Any, Optional
+    from typing import Any, Dict, List, Optional, Tuple
 
-    from sentry_sdk._types import Attributes
+    from sentry_sdk._types import Attributes, SpanLink
     from sentry_sdk.integrations.boto3._context import AwsCallContext
+    from sentry_sdk.traces import Span
 
 # operation -> (`sentry.op`, `messaging.operation.type`)
 _MESSAGING_OPERATIONS = {
@@ -23,6 +30,44 @@
 }
 
 
+def _message_link(message: "Dict[str, Any]") -> "Optional[SpanLink]":
+    # message was sent straight to SQS, or through SNS with raw message delivery on.
+    sentry_trace = (
+        message.get("MessageAttributes", {})
+        .get(SENTRY_TRACE_HEADER_NAME, {})
+        .get("StringValue")
+    )
+
+    # SNS wraps the message and its attributes in a JSON envelope when raw message delivery is off.
+    # https://docs.aws.amazon.com/sns/latest/dg/sns-large-payload-raw-message-delivery.html
+    body = message.get("Body", "")
+    if sentry_trace is None and f'"{SENTRY_TRACE_HEADER_NAME}"' in body:
+        try:
+            envelope = json.loads(body)
+        except ValueError:
+            envelope = None
+        if isinstance(envelope, dict) and envelope.get("Type") == "Notification":
+            sentry_trace = (
+                envelope.get("MessageAttributes", {})
+                .get(SENTRY_TRACE_HEADER_NAME, {})
+                .get("Value")
+            )
+
+    data = extract_sentrytrace_data(sentry_trace)
+    if data is None or not data["trace_id"] or not data["parent_span_id"]:
+        return None
+
+    link: "SpanLink" = {
+        "trace_id": str(data["trace_id"]),
+        "span_id": str(data["parent_span_id"]),
+        # Sentry drops other link attributes for now; add `messaging.message.id` in the future.
+        "attributes": {SPANDATA.SENTRY_LINK_TYPE: "message_origin"},
+    }
+    if data["parent_sampled"] is not None:
+        link["sampled"] = bool(data["parent_sampled"])
+    return link
+
+
 class _SQSExtension(_ServiceExtension):
     def get_span_op(self, ctx: "AwsCallContext") -> "Optional[str]":
         if ctx.operation_name in _MESSAGING_OPERATIONS:
@@ -117,3 +162,47 @@ def get_response_attributes(
         elif error_codes:
             attributes[SPANDATA.ERROR_TYPE] = "aws_batch_partial_failure"
         return attributes
+
+    def get_response_links(
+        self, ctx: "AwsCallContext", response: "Any"
+    ) -> "List[SpanLink]":
+        if ctx.operation_name != "ReceiveMessage":
+            return []
+
+        # one link per producer context; batch entries of one send share it.
+        links: "Dict[Tuple[str, str], SpanLink]" = {}
+        for message in response.get("Messages", ()):
+            # returns a link to the producer context a received message carries, if valid.
+            link = _message_link(message)
+            if link is not None:
+                links.setdefault((link["trace_id"], link["span_id"]), link)
+        return list(links.values())
+
+    def get_request_params(
+        self, ctx: "AwsCallContext", span: "Optional[Span]"
+    ) -> "Optional[Dict[str, Any]]":
+        if ctx.operation_name in ("SendMessage", "SendMessageBatch"):
+            return _inject_message_attributes(ctx, span, "MessageBody", "Entries")
+
+        # only request the carrier when the receive span can record links.
+        if (
+            ctx.operation_name != "ReceiveMessage"
+            or span is None
+            or isinstance(span, NoOpSpan)
+        ):
+            return None
+
+        # SQS only returns the message attributes named in `MessageAttributeNames`.
+        # https://docs.aws.amazon.com/AWSSimpleQueueService/latest/APIReference/API_ReceiveMessage.html
+        names = ctx.params.get("MessageAttributeNames", [])
+        missing = [
+            name
+            for name in (SENTRY_TRACE_HEADER_NAME, BAGGAGE_HEADER_NAME)
+            if name not in names
+        ]
+        # `All` and `.*` already return every attribute, including ours.
+        if not missing or "All" in names or ".*" in names:
+            return None
+
+        # append ours to the caller's names, in a copy of their params.
+        return {**ctx.params, "MessageAttributeNames": [*names, *missing]}

sentry_sdk/integrations/boto3/_utils.py

@@ -1,15 +1,28 @@
 from typing import TYPE_CHECKING
 
+import sentry_sdk
 from sentry_sdk.consts import SPANDATA
+from sentry_sdk.tracing_utils import should_propagate_trace
 
 if TYPE_CHECKING:
     from typing import Any, Dict, Iterable, Optional, Sequence, Tuple
 
     from sentry_sdk._types import Attributes
+    from sentry_sdk.integrations.boto3._context import AwsCallContext
+    from sentry_sdk.traces import Span
 
     # tuple of (request param name, span attribute name).
     _AttributeSpec = Tuple[str, str]
 
+# SQS and SNS reject a message with more than 10 message attributes.
+# https://docs.aws.amazon.com/AWSSimpleQueueService/latest/SQSDeveloperGuide/sqs-message-metadata.html
+_MAX_MESSAGE_ATTRIBUTES = 10
+
+# body and attributes must fit the queue's or topic's configurable `MaximumMessageSize`;
+# assume the 256 KiB default of SNS topics and older SQS queues; batch uses the same limit.
+# https://docs.aws.amazon.com/sns/latest/api/API_PublishBatch.html
+_MAX_MESSAGE_SIZE = 256 * 1024
+
 
 def _extract_attributes(
     source: "Dict[str, Any]", specs: "Sequence[_AttributeSpec]"
@@ -38,3 +51,63 @@ def _get_aws_arn_attributes(
         return attributes
 
     return {}
+
+
+def _inject_message_attributes(
+    ctx: "AwsCallContext", span: "Optional[Span]", body_key: str, entries_key: str
+) -> "Optional[Dict[str, Any]]":
+    """
+    Return a copy of `ctx.params` with `sentry-trace` and `baggage` added in
+    string message attributes, or to each batch entry with room for them.
+    """
+    if ctx.endpoint_url is None or not should_propagate_trace(
+        sentry_sdk.get_client(), ctx.endpoint_url
+    ):
+        return None
+
+    carrier = {
+        name: {"DataType": "String", "StringValue": value}
+        for name, value in sentry_sdk.get_current_scope().iter_trace_propagation_headers(
+            span=span
+        )
+    }
+    if not carrier:
+        return None
+
+    def _get_message_size(body: str, attributes: "Dict[str, Any]") -> int:
+        size = len(body.encode())
+
+        # names, types and values of message attributes count towards the message size.
+        # https://docs.aws.amazon.com/AWSSimpleQueueService/latest/SQSDeveloperGuide/sqs-message-metadata.html
+        for name, value in attributes.items():
+            size += len(name.encode()) + len(value["DataType"].encode())
+            if "StringValue" in value:
+                size += len(value["StringValue"].encode())
+            if "BinaryValue" in value:
+                binary = value["BinaryValue"]
+                size += len(binary.encode() if isinstance(binary, str) else binary)
+        return size
+
+    is_batch = entries_key in ctx.params
+    messages = []
+    for message in ctx.params[entries_key] if is_batch else (ctx.params,):
+        attributes = message.get("MessageAttributes", {})
+        injected = {**attributes, **carrier}
+        if (
+            # never replace the caller's carrier.
+            not carrier.keys() & attributes.keys()
+            # don't make AWS reject the message.
+            and len(injected) <= _MAX_MESSAGE_ATTRIBUTES
+            and _get_message_size(message[body_key], injected) <= _MAX_MESSAGE_SIZE
+        ):
+            message = {**message, "MessageAttributes": injected}
+        messages.append(message)
+
+    total_size = sum(
+        _get_message_size(message[body_key], message.get("MessageAttributes", {}))
+        for message in messages
+    )
+    if total_size > _MAX_MESSAGE_SIZE:
+        return None
+
+    return {**ctx.params, entries_key: messages} if is_batch else messages[0]