feat(boto3): Add Kinesis extension

sentry_sdk/consts.py

@@ -363,6 +363,12 @@ class SPANDATA:
     Example: "wzHcyEWfmOGDIE5QOhTAqFDoDWP3y8IUvpNINCwL9N4TEHbUw0/gZJ+VZTmCNCWR7fezEN3eCiQ="
     """
 
+    AWS_KINESIS_STREAM_NAME = "aws.kinesis.stream_name"
+    """
+    The name of the AWS Kinesis stream the request refers to.
+    Example: "orders"
+    """
+
     AWS_LAMBDA_INVOKED_ARN = "aws.lambda.invoked_arn"
     """
     The full ARN of the Lambda function that was invoked.
@@ -452,7 +458,7 @@ class SPANDATA:
     The name of the AWS Step Functions state machine the request refers to, without a version or alias qualifier.
     Example: "orders"
     """
-    
+
     AWS_SECRETSMANAGER_SECRET_ARN = "aws.secretsmanager.secret.arn"
     """
     The ARN of the Secrets Manager secret the request refers to.

sentry_sdk/integrations/boto3/_services/kinesis.py

@@ -1,17 +1,30 @@
 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 _get_aws_arn_attributes
 from sentry_sdk.utils import capture_internal_exceptions
 
 if TYPE_CHECKING:
+    from typing import Any
+
     from sentry_sdk._types import Attributes
     from sentry_sdk.integrations.boto3._context import AwsCallContext
 
 
 class _KinesisExtension(_ServiceExtension):
     def get_request_attributes(self, ctx: "AwsCallContext") -> "Attributes":
         attributes: "Attributes" = {}
+
+        # a stream name, or an arn that contains it: `…:stream/<name>` for a stream,
+        # `…:stream/<name>/consumer/<consumer>:<timestamp>` for a consumer.
+        for param in ("StreamName", "StreamARN", "ConsumerARN", "ResourceARN"):
+            if param in ctx.params:
+                attributes[SPANDATA.AWS_KINESIS_STREAM_NAME] = (
+                    ctx.params[param].split(":stream/", 1)[-1].split("/", 1)[0]
+                )
+                break
+
         with capture_internal_exceptions():
             attributes.update(
                 _get_aws_arn_attributes(
@@ -24,4 +37,24 @@ def get_request_attributes(self, ctx: "AwsCallContext") -> "Attributes":
                     ("kinesis",),
                 )
             )
+
         return attributes
+
+    def get_response_attributes(
+        self, ctx: "AwsCallContext", response: "Any"
+    ) -> "Attributes":
+        # `PutRecords` returns `200 OK` even when records fail.
+        # https://docs.aws.amazon.com/kinesis/latest/APIReference/API_PutRecords.html
+        if ctx.operation_name != "PutRecords" or not response.get("FailedRecordCount"):
+            return {}
+
+        error_codes = {
+            record["ErrorCode"]
+            for record in response["Records"]
+            if "ErrorCode" in record
+        }
+        return {
+            SPANDATA.ERROR_TYPE: error_codes.pop()
+            if len(error_codes) == 1
+            else "aws_batch_partial_failure"
+        }

tests/integrations/boto3/test_kinesis.py

@@ -1,24 +1,120 @@
+import pytest
 from botocore.stub import Stubber
 
 from sentry_sdk.consts import OP, SPANDATA
-from tests.integrations.boto3.helpers import capture_spans_by_op
+from tests.integrations.boto3.helpers import (
+    capture_spans_by_op,
+    require_botocore_model_fields,
+)
 from tests.integrations.boto3.helpers import client_factory as client_factory
 
+STREAM_ARN = "arn:aws:kinesis:eu-north-1:123456789012:stream/orders"
+CONSUMER_ARN = STREAM_ARN + "/consumer/billing:1760000000"
 
-def test_request_attributes(client_factory, capture_items):
-    client = client_factory("kinesis")
-    stream_arn = "arn:aws:kinesis:eu-north-1:123456789012:stream/orders"
 
+def _stubbed_span(client_factory, capture_items, method, params, response):
+    client = client_factory("kinesis")
+    require_botocore_model_fields(
+        client, method, input_fields=tuple(params), output_fields=tuple(response)
+    )
     with Stubber(client) as stubber:
-        stubber.add_response(
-            "list_stream_consumers", {"Consumers": []}, {"StreamARN": stream_arn}
-        )
+        stubber.add_response(method, response, params)
         spans = capture_spans_by_op(
-            lambda: client.list_stream_consumers(**{"StreamARN": stream_arn}),
-            capture_items,
+            lambda: getattr(client, method)(**params), capture_items
         )
-
     (span,) = spans[OP.RPC]
-    attributes = span["attributes"]
-    assert attributes[SPANDATA.CLOUD_ACCOUNT_ID] == "123456789012"
-    assert attributes[SPANDATA.CLOUD_RESOURCE_ID] == stream_arn
+    return span
+
+
+@pytest.mark.parametrize(
+    "method,params,response,expected",
+    [
+        pytest.param(
+            "put_record",
+            {"StreamName": "orders", "Data": b"{}", "PartitionKey": "customer-1"},
+            {"ShardId": "shardId-000000000000", "SequenceNumber": "1"},
+            {SPANDATA.AWS_KINESIS_STREAM_NAME: "orders"},
+            id="stream-name",
+        ),
+        pytest.param(
+            "list_stream_consumers",
+            {"StreamARN": STREAM_ARN},
+            {"Consumers": []},
+            {
+                SPANDATA.AWS_KINESIS_STREAM_NAME: "orders",
+                SPANDATA.CLOUD_RESOURCE_ID: STREAM_ARN,
+            },
+            id="stream-arn",
+        ),
+        pytest.param(
+            "deregister_stream_consumer",
+            {"ConsumerARN": CONSUMER_ARN},
+            {},
+            {
+                SPANDATA.AWS_KINESIS_STREAM_NAME: "orders",
+                SPANDATA.CLOUD_RESOURCE_ID: CONSUMER_ARN,
+            },
+            id="consumer-arn",
+        ),
+        pytest.param(
+            "list_tags_for_resource",
+            {"ResourceARN": CONSUMER_ARN},
+            {"Tags": []},
+            {
+                SPANDATA.AWS_KINESIS_STREAM_NAME: "orders",
+                SPANDATA.CLOUD_RESOURCE_ID: CONSUMER_ARN,
+            },
+            id="resource-arn",
+        ),
+    ],
+)
+def test_stream_attributes(
+    client_factory, capture_items, method, params, response, expected
+):
+    span = _stubbed_span(client_factory, capture_items, method, params, response)
+
+    assert {
+        key: value
+        for key, value in span["attributes"].items()
+        if key in (SPANDATA.AWS_KINESIS_STREAM_NAME, SPANDATA.CLOUD_RESOURCE_ID)
+    } == expected
+
+
+@pytest.mark.parametrize(
+    "error_codes,expected_error_type",
+    [
+        pytest.param((), None, id="ok"),
+        pytest.param(
+            ("ProvisionedThroughputExceededException",) * 2,
+            "ProvisionedThroughputExceededException",
+            id="one-code",
+        ),
+        pytest.param(
+            ("ProvisionedThroughputExceededException", "InternalFailure"),
+            "aws_batch_partial_failure",
+            id="different-codes",
+        ),
+    ],
+)
+def test_put_records_outcome(
+    client_factory, capture_items, error_codes, expected_error_type
+):
+    params = {
+        "StreamName": "orders",
+        "Records": [
+            {"Data": b"card 4242", "PartitionKey": "customer-1"},
+            {"Data": b"card 4242", "PartitionKey": "customer-2"},
+            {"Data": b"card 4242", "PartitionKey": "customer-3"},
+        ],
+    }
+    response = {
+        "Records": [{"ShardId": "shardId-000000000000", "SequenceNumber": "1"}]
+        + [{"ErrorCode": code, "ErrorMessage": "failed"} for code in error_codes],
+    }
+    # the model's minimum is 1, so the stubber only accepts it for failed records.
+    if error_codes:
+        response["FailedRecordCount"] = len(error_codes)  # type: ignore
+    span = _stubbed_span(client_factory, capture_items, "put_records", params, response)
+
+    assert span["status"] == ("ok" if expected_error_type is None else "error")
+    assert span["attributes"].get(SPANDATA.ERROR_TYPE) == expected_error_type