From c690e5d156c1c4bd74d6eef145e5a6400a39d3cb Mon Sep 17 00:00:00 2001 From: Pablo Deputter Date: Fri, 16 Oct 2026 18:55:22 +0200 Subject: [PATCH 1/2] feat(boto3): Add Kinesis extension --- sentry_sdk/consts.py | 6 + sentry_sdk/integrations/boto3/_client.py | 5 +- .../integrations/boto3/_services/kinesis.py | 33 +++++ tests/integrations/boto3/test_kinesis.py | 122 ++++++++++++++++-- 4 files changed, 152 insertions(+), 14 deletions(-) diff --git a/sentry_sdk/consts.py b/sentry_sdk/consts.py index 5ff880ab73..31997de5dd 100644 --- a/sentry_sdk/consts.py +++ b/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. diff --git a/sentry_sdk/integrations/boto3/_client.py b/sentry_sdk/integrations/boto3/_client.py index 10a5bb5c93..2515fe1149 100644 --- a/sentry_sdk/integrations/boto3/_client.py +++ b/sentry_sdk/integrations/boto3/_client.py @@ -2,6 +2,7 @@ from typing import TYPE_CHECKING import sentry_sdk +from sentry_sdk.consts import SPANDATA from sentry_sdk.integrations import DidNotEnable from sentry_sdk.integrations.boto3._context import AwsCallContext from sentry_sdk.integrations.boto3._instrumentation import ( @@ -17,7 +18,7 @@ _resolve_service, ) from sentry_sdk.integrations.boto3.consts import IDENTIFIER -from sentry_sdk.traces import NoOpSpan, Span +from sentry_sdk.traces import NoOpSpan, Span, SpanStatus from sentry_sdk.utils import capture_internal_exceptions if TYPE_CHECKING: @@ -125,6 +126,8 @@ def sentry_patched_make_api_call( attributes.update(_get_response_attributes(parsed)) with capture_internal_exceptions(): span.set_attributes(attributes) + if SPANDATA.ERROR_TYPE in attributes: + span.status = SpanStatus.ERROR except BaseException as error: _finish_span(span, error) raise diff --git a/sentry_sdk/integrations/boto3/_services/kinesis.py b/sentry_sdk/integrations/boto3/_services/kinesis.py index 14947fb4ec..2d4cb6d3f8 100644 --- a/sentry_sdk/integrations/boto3/_services/kinesis.py +++ b/sentry_sdk/integrations/boto3/_services/kinesis.py @@ -1,10 +1,13 @@ 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 @@ -12,6 +15,16 @@ class _KinesisExtension(_ServiceExtension): def get_request_attributes(self, ctx: "AwsCallContext") -> "Attributes": attributes: "Attributes" = {} + + # a stream name, or an arn that contains it: `…:stream/` for a stream, + # `…:stream//consumer/:` 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" + } diff --git a/tests/integrations/boto3/test_kinesis.py b/tests/integrations/boto3/test_kinesis.py index c643d62595..92dcf239cb 100644 --- a/tests/integrations/boto3/test_kinesis.py +++ b/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 From 5a6d11883e37744ad2da13c387168cb19720201f Mon Sep 17 00:00:00 2001 From: Pablo Deputter Date: Fri, 16 Oct 2026 18:59:31 +0200 Subject: [PATCH 2/2] lint --- tests/integrations/boto3/test_kinesis.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integrations/boto3/test_kinesis.py b/tests/integrations/boto3/test_kinesis.py index 92dcf239cb..109346bd93 100644 --- a/tests/integrations/boto3/test_kinesis.py +++ b/tests/integrations/boto3/test_kinesis.py @@ -115,6 +115,6 @@ def test_put_records_outcome( 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