feat(boto3): Add SQS/SNS trace propagation

sentry_sdk/consts.py

@@ -1424,6 +1424,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/_services/sns.py

@@ -6,12 +6,13 @@
 from sentry_sdk.integrations.boto3._utils import (
     _get_aws_arn_attributes,
     _get_batch_error_type,
+    _inject_message_attributes,
 )
 from sentry_sdk.traces import SpanStatus
 from sentry_sdk.utils import capture_internal_exceptions
 
 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
@@ -129,3 +130,13 @@ def on_success_response(
 
         span.set_attributes(attributes)
         span.end()
+
+    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,3 +1,4 @@
+import json
 import re
 from typing import TYPE_CHECKING
 from urllib.parse import unquote, urlsplit
@@ -7,14 +8,21 @@
 from sentry_sdk.integrations.boto3._utils import (
     _get_aws_arn_attributes,
     _get_batch_error_type,
+    _inject_message_attributes,
 )
-from sentry_sdk.traces import SpanStatus
+from sentry_sdk.traces import (
+    BAGGAGE_HEADER_NAME,
+    SENTRY_TRACE_HEADER_NAME,
+    NoOpSpan,
+    SpanStatus,
+)
+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, 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
 
@@ -28,6 +36,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:
@@ -124,5 +170,45 @@ def on_success_response(
                 span.status = SpanStatus.ERROR
                 attributes[SPANDATA.ERROR_TYPE] = error_type
 
+        if ctx.operation_name == "ReceiveMessage":
+            with capture_internal_exceptions():
+                # 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)
+                span.add_links(list(links.values()))
+
         span.set_attributes(attributes)
         span.end()
+
+    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

@@ -15,6 +15,15 @@
     # 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]"
@@ -76,3 +85,56 @@ def _get_propagation_headers(
     return dict(
         sentry_sdk.get_current_scope().iter_trace_propagation_headers(span=span)
     )
+
+
+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.
+    """
+    carrier = {
+        name: {"DataType": "String", "StringValue": value}
+        for name, value in _get_propagation_headers(ctx, span).items()
+    }
+    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]

tests/integrations/boto3/test_sqs.py

@@ -1,3 +1,6 @@
+import copy
+import json
+
 import pytest
 from botocore.exceptions import ClientError
 from botocore.stub import Stubber
@@ -321,3 +324,116 @@ def test_delete_message(
     assert attributes[SPANDATA.MESSAGING_OPERATION_NAME] == "settle"
     assert attributes.get(SPANDATA.MESSAGING_BATCH_MESSAGE_COUNT) == message_count
     assert "receipt-handle" not in str(span)
+
+
+@pytest.mark.parametrize(
+    "names,sent_names",
+    [
+        pytest.param(None, ["sentry-trace", "baggage"], id="default"),
+        pytest.param(
+            ["orderId", "baggage"],
+            ["orderId", "baggage", "sentry-trace"],
+            id="caller-names",
+        ),
+        pytest.param(
+            ["baggage", "sentry-trace"],
+            ["baggage", "sentry-trace"],
+            id="already-requested",
+        ),
+        pytest.param(["All"], ["All"], id="all"),
+        pytest.param([".*"], [".*"], id="all-pattern"),
+    ],
+)
+def test_receive_message_requests_carrier(
+    client_factory, capture_items, names, sent_names
+):
+    client = client_factory("sqs")
+    queue_url = "https://sqs.eu-north-1.amazonaws.com/123456789012/orders"
+    params = {} if names is None else {"MessageAttributeNames": names}
+    caller_params = copy.deepcopy(params)
+
+    with Stubber(client) as stubber:
+        stubber.add_response(
+            "receive_message",
+            {},
+            {"QueueUrl": queue_url, "MessageAttributeNames": sent_names},
+        )
+        capture_spans_by_op(
+            lambda: client.receive_message(QueueUrl=queue_url, **params),
+            capture_items,
+        )
+
+    assert params == caller_params
+
+
+def test_receive_message_links(client_factory, capture_items):
+    client = client_factory("sqs")
+    sentry_trace = {
+        "sentry-trace": {
+            "DataType": "String",
+            "StringValue": "771a43a4192642f0b136d5159a501700-a3e9e5ba6f99f4d6-1",
+        }
+    }
+    sns_envelope = {
+        "Type": "Notification",
+        "MessageId": "sns-message-3",
+        "TopicArn": "arn:aws:sns:eu-north-1:123456789012:alerts",
+        "Message": "hi",
+        "MessageAttributes": {
+            "sentry-trace": {
+                "Type": "String",
+                "Value": "b2c1b1d6cb5f4a0d8a6a2f1c3e9d7b40-9c4f3a2b1e0d8f7a",
+            }
+        },
+    }
+    response = {
+        "Messages": [
+            # batch entries of one send share a context.
+            {"MessageId": "1", "Body": "hi", "MessageAttributes": sentry_trace},
+            {"MessageId": "2", "Body": "hi", "MessageAttributes": sentry_trace},
+            # SNS-to-SQS delivery without raw message delivery.
+            {"MessageId": "3", "Body": json.dumps(sns_envelope)},
+            # trace id with malformed sentry-trace.
+            {
+                "MessageId": "4",
+                "Body": "hi",
+                "MessageAttributes": {
+                    "sentry-trace": {
+                        "DataType": "String",
+                        "StringValue": "771a43a4192642f0b136d5159a501700",
+                    }
+                },
+            },
+            # JSON body that is not an SNS envelope.
+            {"MessageId": "5", "Body": '{"sentry-trace": "not an SNS envelope"}'},
+            {"MessageId": "6", "Body": "hi"},
+        ]
+    }
+
+    with Stubber(client) as stubber:
+        stubber.add_response("receive_message", response)
+        spans = capture_spans_by_op(
+            lambda: client.receive_message(
+                QueueUrl="https://sqs.eu-north-1.amazonaws.com/123456789012/orders"
+            ),
+            capture_items,
+        )
+
+    (span,) = spans[OP.QUEUE_RECEIVE]
+    link_type = {
+        SPANDATA.SENTRY_LINK_TYPE: {"value": "message_origin", "type": "string"}
+    }
+    assert span["links"] == [
+        # messages 1 and 2 share a context, so one link.
+        {
+            "trace_id": "771a43a4192642f0b136d5159a501700",
+            "span_id": "a3e9e5ba6f99f4d6",
+            "sampled": True,
+            "attributes": link_type,
+        },
+        {
+            "trace_id": "b2c1b1d6cb5f4a0d8a6a2f1c3e9d7b40",
+            "span_id": "9c4f3a2b1e0d8f7a",
+            "attributes": link_type,
+        },
+    ]