feat(boto3): Add SQS extension

sentry_sdk/consts.py

@@ -471,6 +471,12 @@ class SPANDATA:
     Example: "db-password"
     """
 
+    AWS_SQS_QUEUE_URL = "aws.sqs.queue.url"
+    """
+    The URL of the AWS SQS queue the request refers to.
+    Example: "https://sqs.us-east-1.amazonaws.com/123456789012/orders"
+    """
+
     CACHE_HIT = "cache.hit"
     """
     A boolean indicating whether the requested data was found in the cache.
@@ -1090,6 +1096,18 @@ class SPANDATA:
     Number of retries/attempts to process a message.
     """
 
+    MESSAGING_OPERATION_NAME = "messaging.operation.name"
+    """
+    The system-specific name of the messaging operation.
+    Example: "publish"
+    """
+
+    MESSAGING_OPERATION_TYPE = "messaging.operation.type"
+    """
+    The type of the messaging operation: `create`, `send`, `receive`, `process` or `settle`.
+    Example: "send"
+    """
+
     MESSAGING_SYSTEM = "messaging.system"
     """
     The messaging system's name, e.g. `kafka`, `aws_sqs`
@@ -1489,6 +1507,8 @@ class OP:
     OBJECT_UPLOAD_PART = "object.upload_part"
     QUEUE_PROCESS = "queue.process"
     QUEUE_PUBLISH = "queue.publish"
+    QUEUE_RECEIVE = "queue.receive"
+    QUEUE_SETTLE = "queue.settle"
     QUEUE_SUBMIT_ARQ = "queue.submit.arq"
     QUEUE_TASK_ARQ = "queue.task.arq"
     QUEUE_SUBMIT_CELERY = "queue.submit.celery"

sentry_sdk/integrations/boto3/_instrumentation.py

@@ -180,6 +180,11 @@ def _start_client_span(
         with capture_internal_exceptions():
             attributes.update(service_ext.get_request_attributes(ctx))
 
+        with capture_internal_exceptions():
+            service_span_name = service_ext.get_span_name(ctx, attributes)
+            if service_span_name is not None:
+                span_name = service_span_name
+
     # `start_span()` evaluates `ignore_spans` against the initial attributes.
     # https://opentelemetry.io/docs/specs/semconv/rpc/rpc-spans/#rpc-client-span
     attributes.update(

sentry_sdk/integrations/boto3/_services/base.py

@@ -21,6 +21,12 @@ def get_span_op(self, ctx: "AwsCallContext") -> "Optional[str]":
         """Return an optional `sentry.op` override, or `None` to keep the default."""
         return None
 
+    def get_span_name(
+        self, ctx: "AwsCallContext", attributes: "Attributes"
+    ) -> "Optional[str]":
+        """Return an optional span name override, or `None` to keep the default."""
+        return None
+
     def get_request_attributes(self, ctx: "AwsCallContext") -> "Attributes":
         """Return request attributes to add before the AWS request is made."""
         return {}

sentry_sdk/integrations/boto3/_services/sqs.py

@@ -2,36 +2,110 @@
 from typing import TYPE_CHECKING
 from urllib.parse import unquote, urlsplit
 
-from sentry_sdk.consts import SPANDATA
+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.utils import capture_internal_exceptions
 
 if TYPE_CHECKING:
+    from typing import Any, Optional
+
     from sentry_sdk._types import Attributes
     from sentry_sdk.integrations.boto3._context import AwsCallContext
 
+# operation -> (`sentry.op`, `messaging.operation.type`)
+_MESSAGING_OPERATIONS = {
+    "SendMessage": (OP.QUEUE_PUBLISH, "send"),
+    "SendMessageBatch": (OP.QUEUE_PUBLISH, "send"),
+    "ReceiveMessage": (OP.QUEUE_RECEIVE, "receive"),
+    "DeleteMessage": (OP.QUEUE_SETTLE, "settle"),
+    "DeleteMessageBatch": (OP.QUEUE_SETTLE, "settle"),
+}
+
 
 class _SQSExtension(_ServiceExtension):
+    def get_span_op(self, ctx: "AwsCallContext") -> "Optional[str]":
+        if ctx.operation_name in _MESSAGING_OPERATIONS:
+            return _MESSAGING_OPERATIONS[ctx.operation_name][0]
+        return None
+
+    def get_span_name(
+        self, ctx: "AwsCallContext", attributes: "Attributes"
+    ) -> "Optional[str]":
+        if ctx.operation_name not in _MESSAGING_OPERATIONS:
+            return None
+        # for example: `send <queue name>` when `messaging.operation.type=send`.
+        # https://opentelemetry.io/docs/specs/semconv/messaging/messaging-spans/#span-name
+        return f"{attributes[SPANDATA.MESSAGING_OPERATION_TYPE]} {attributes[SPANDATA.MESSAGING_DESTINATION_NAME]}"
+
     def get_request_attributes(self, ctx: "AwsCallContext") -> "Attributes":
-        attributes: "Attributes" = {}
-        with capture_internal_exceptions():
+        attributes = _get_aws_arn_attributes(
+            (ctx.params.get("DestinationArn"), ctx.params.get("SourceArn")), ("sqs",)
+        )
+
+        queue_url = ctx.params.get("QueueUrl")
+        if queue_url is not None:
+            attributes[SPANDATA.AWS_SQS_QUEUE_URL] = queue_url
+
+            # `https://sqs.eu-north-1.amazonaws.com/123456789012/orders` => ['123456789012', 'orders']
+            path = unquote(urlsplit(queue_url).path).strip("/").split("/")
+
+            attributes[SPANDATA.MESSAGING_DESTINATION_NAME] = path[-1]
+            if re.fullmatch(r"[0-9]{12}", path[0]):
+                attributes[SPANDATA.CLOUD_ACCOUNT_ID] = path[0]
+        elif "QueueName" in ctx.params:
+            # `CreateQueue` & `GetQueueUrl` operations.
+            attributes[SPANDATA.MESSAGING_DESTINATION_NAME] = ctx.params["QueueName"]
+
+            if "QueueOwnerAWSAccountId" in ctx.params:
+                attributes[SPANDATA.CLOUD_ACCOUNT_ID] = ctx.params[
+                    "QueueOwnerAWSAccountId"
+                ]
+        elif "SourceArn" in ctx.params:
+            # `arn:aws:sqs:eu-north-1:123456789012:orders` => `orders`
+            attributes[SPANDATA.MESSAGING_DESTINATION_NAME] = ctx.params[
+                "SourceArn"
+            ].split(":")[-1]
+
+        if ctx.operation_name in _MESSAGING_OPERATIONS:
+            operation_type = _MESSAGING_OPERATIONS[ctx.operation_name][1]
             attributes.update(
-                _get_aws_arn_attributes(
-                    (
-                        ctx.params.get("DestinationArn"),
-                        ctx.params.get("SourceArn"),
-                    ),
-                    ("sqs",),
-                )
+                {
+                    SPANDATA.MESSAGING_SYSTEM: "aws_sqs",
+                    SPANDATA.MESSAGING_OPERATION_TYPE: operation_type,
+                    SPANDATA.MESSAGING_OPERATION_NAME: operation_type,
+                }
             )
+            if operation_type == "send":
+                attributes[SPANDATA.SENTRY_KIND] = "producer"
+            if "Entries" in ctx.params:
+                attributes[SPANDATA.MESSAGING_BATCH_MESSAGE_COUNT] = len(
+                    ctx.params["Entries"]
+                )
+        return attributes
 
-        queue_url = ctx.params.get("QueueUrl")
-        if queue_url is None:
+    def get_response_attributes(
+        self, ctx: "AwsCallContext", response: "Any"
+    ) -> "Attributes":
+        attributes: "Attributes" = {}
+
+        # `ClientError` response holds no results.
+        if "Error" in response:
             return attributes
 
-        with capture_internal_exceptions():
-            account_id = unquote(urlsplit(queue_url).path).strip("/").split("/", 1)[0]
-            if re.fullmatch(r"^[0-9]{12}$", account_id):
-                attributes[SPANDATA.CLOUD_ACCOUNT_ID] = account_id
+        if ctx.operation_name == "SendMessage":
+            attributes[SPANDATA.MESSAGING_MESSAGE_ID] = response["MessageId"]
+        elif ctx.operation_name == "ReceiveMessage":
+            # `Messages` is left out when no message was received.
+            messages = response.get("Messages", ())
+            attributes[SPANDATA.MESSAGING_BATCH_MESSAGE_COUNT] = len(messages)
+            if len(messages) == 1:
+                attributes[SPANDATA.MESSAGING_MESSAGE_ID] = messages[0]["MessageId"]
+
+        # failed entries have a `200 OK` response; use their code when they share one.
+        # https://docs.aws.amazon.com/sns/latest/api/API_BatchResultErrorEntry.html
+        error_codes = {entry["Code"] for entry in response.get("Failed", ())}
+        if len(error_codes) == 1:
+            attributes[SPANDATA.ERROR_TYPE] = error_codes.pop()
+        elif error_codes:
+            attributes[SPANDATA.ERROR_TYPE] = "aws_batch_partial_failure"
         return attributes