From 56b920d596d40eff0d463eba01bab41176286c17 Mon Sep 17 00:00:00 2001 From: Ansh Rohilla Date: Thu, 29 Oct 2026 18:43:56 +0530 Subject: [PATCH] fix(aws-lambda): Only set messaging.batch.message_count for messaging batches (closes #8078) --- sentry_sdk/integrations/aws_lambda.py | 38 ++- .../aws_lambda/test_aws_lambda.py | 4 +- .../aws_lambda/test_aws_lambda_messaging.py | 244 ++++++++++++++++++ 3 files changed, 282 insertions(+), 4 deletions(-) create mode 100644 tests/integrations/aws_lambda/test_aws_lambda_messaging.py diff --git a/sentry_sdk/integrations/aws_lambda.py b/sentry_sdk/integrations/aws_lambda.py index 009d7328b2..ac5a045b17 100644 --- a/sentry_sdk/integrations/aws_lambda.py +++ b/sentry_sdk/integrations/aws_lambda.py @@ -96,6 +96,32 @@ def sentry_init_error(*args: "Any", **kwargs: "Any") -> "Any": return sentry_init_error # type: ignore +def _get_messaging_batch_count(aws_event: "Any") -> "Optional[int]": + """ + Returns the message count if the event is from a recognized messaging + service (SQS or SNS), or None otherwise. + """ + if not isinstance(aws_event, dict): + return None + + records = aws_event.get("Records") + if not isinstance(records, list) or not records: + return None + + first_record = records[0] + if not isinstance(first_record, dict): + return None + + event_source = ( + first_record.get("eventSource") or first_record.get("EventSource") or "" + ).lower() + + if event_source in ("aws:sqs", "aws:sns") or "Sns" in first_record: + return len(records) + + return None + + def _wrap_handler(handler: "F") -> "F": @functools.wraps(handler) def sentry_handler( @@ -117,7 +143,11 @@ def sentry_handler( if integration is None: return handler(aws_event, aws_context, *args, **kwargs) - if isinstance(aws_event, list) and len(aws_event) >= 1: + messaging_batch_count = _get_messaging_batch_count(aws_event) + if messaging_batch_count is not None: + request_data = aws_event + batch_size = messaging_batch_count + elif isinstance(aws_event, list) and len(aws_event) >= 1: request_data = aws_event[0] batch_size = len(aws_event) else: @@ -240,9 +270,13 @@ def sentry_handler( "aws.lambda.invoked_arn": aws_context.invoked_function_arn, "aws.log.group.names": [aws_context.log_group_name], "aws.log.stream.names": [aws_context.log_stream_name], - "messaging.batch.message_count": batch_size, **header_attributes, **additional_attributes, + **( + {"messaging.batch.message_count": messaging_batch_count} + if messaging_batch_count is not None + else {} + ), }, ) else: diff --git a/tests/integrations/aws_lambda/test_aws_lambda.py b/tests/integrations/aws_lambda/test_aws_lambda.py index 0bf1cd3670..2d0959eb06 100644 --- a/tests/integrations/aws_lambda/test_aws_lambda.py +++ b/tests/integrations/aws_lambda/test_aws_lambda.py @@ -976,7 +976,7 @@ def test_span_streaming_no_error(lambda_client, test_environment): "aws/lambda/BasicOkSpanStreaming" ] assert _get_span_attr(attrs, "aws.log.stream.names") == ["$LATEST"] - assert _get_span_attr(attrs, "messaging.batch.message_count") == 1 + assert "messaging.batch.message_count" not in attrs def test_span_streaming_error(lambda_client, test_environment): @@ -1026,7 +1026,7 @@ def test_span_streaming_error(lambda_client, test_environment): "aws/lambda/RaiseErrorSpanStreaming" ] assert _get_span_attr(attrs, "aws.log.stream.names") == ["$LATEST"] - assert _get_span_attr(attrs, "messaging.batch.message_count") == 1 + assert "messaging.batch.message_count" not in attrs def test_span_streaming_trace_continuation(lambda_client, test_environment): diff --git a/tests/integrations/aws_lambda/test_aws_lambda_messaging.py b/tests/integrations/aws_lambda/test_aws_lambda_messaging.py new file mode 100644 index 0000000000..cb46ff9d58 --- /dev/null +++ b/tests/integrations/aws_lambda/test_aws_lambda_messaging.py @@ -0,0 +1,244 @@ +import sentry_sdk +from sentry_sdk.integrations.aws_lambda import ( + AwsLambdaIntegration, + _get_messaging_batch_count, + _wrap_handler, +) + + +class MockContext: + function_name = "test_function" + function_version = "$LATEST" + invoked_function_arn = ( + "arn:aws:lambda:us-east-1:123456789012:function:test_function" + ) + aws_request_id = "test-request-id" + log_group_name = "/aws/lambda/test_function" + log_stream_name = "2026/10/29/[$LATEST]test" + + def get_remaining_time_in_millis(self): + return 300000 + + +def test_get_messaging_batch_count_sqs(): + single_sqs = { + "Records": [ + { + "eventSource": "aws:sqs", + "messageId": "1", + "body": "hello", + } + ] + } + assert _get_messaging_batch_count(single_sqs) == 1 + + batch_sqs = { + "Records": [{"eventSource": "aws:sqs", "messageId": str(i)} for i in range(5)] + } + assert _get_messaging_batch_count(batch_sqs) == 5 + + # Case insensitive EventSource + capitalized_sqs = { + "Records": [ + { + "EventSource": "aws:sqs", + "messageId": "1", + } + ] + } + assert _get_messaging_batch_count(capitalized_sqs) == 1 + + +def test_get_messaging_batch_count_sns(): + single_sns = { + "Records": [ + { + "EventSource": "aws:sns", + "Sns": {"Message": "hello"}, + } + ] + } + assert _get_messaging_batch_count(single_sns) == 1 + + batch_sns = { + "Records": [ + { + "EventSource": "aws:sns", + "Sns": {"Message": f"msg {i}"}, + } + for i in range(3) + ] + } + assert _get_messaging_batch_count(batch_sns) == 3 + + +def test_get_messaging_batch_count_non_messaging(): + assert _get_messaging_batch_count({}) is None + assert _get_messaging_batch_count({"httpMethod": "GET", "path": "/"}) is None + assert _get_messaging_batch_count([{"key": "val"}]) is None + assert _get_messaging_batch_count("string_event") is None + assert _get_messaging_batch_count({"Records": []}) is None + + # S3 event + s3_event = { + "Records": [ + { + "eventSource": "aws:s3", + "s3": {"bucket": {"name": "my-bucket"}}, + } + ] + } + assert _get_messaging_batch_count(s3_event) is None + + # DynamoDB event + dynamo_event = { + "Records": [ + { + "eventSource": "aws:dynamodb", + "eventName": "INSERT", + } + ] + } + assert _get_messaging_batch_count(dynamo_event) is None + + # Kinesis event + kinesis_event = { + "Records": [ + { + "eventSource": "aws:kinesis", + "kinesis": {"data": "dGVzdA=="}, + } + ] + } + assert _get_messaging_batch_count(kinesis_event) is None + + +def test_handler_sqs_batch_span_streaming(sentry_init, capture_items): + sentry_init( + integrations=[AwsLambdaIntegration()], + traces_sample_rate=1.0, + trace_lifecycle="stream", + ) + items = capture_items("span") + + sqs_event = { + "Records": [ + {"eventSource": "aws:sqs", "messageId": "1"}, + {"eventSource": "aws:sqs", "messageId": "2"}, + {"eventSource": "aws:sqs", "messageId": "3"}, + ] + } + + tags_recorded = {} + + @_wrap_handler + def handler(event, context): + tags_recorded.update(sentry_sdk.get_isolation_scope()._tags) + return "ok" + + res = handler(sqs_event, MockContext()) + assert res == "ok" + + sentry_sdk.flush() + spans = [item.payload for item in items] + assert len(spans) == 1 + span = spans[0] + + assert span["attributes"]["messaging.batch.message_count"] == 3 + assert tags_recorded.get("batch_request") is True + assert tags_recorded.get("batch_size") == 3 + + +def test_handler_sns_single_span_streaming(sentry_init, capture_items): + sentry_init( + integrations=[AwsLambdaIntegration()], + traces_sample_rate=1.0, + trace_lifecycle="stream", + ) + items = capture_items("span") + + sns_event = { + "Records": [ + {"EventSource": "aws:sns", "Sns": {"Message": "hi"}}, + ] + } + + tags_recorded = {} + + @_wrap_handler + def handler(event, context): + tags_recorded.update(sentry_sdk.get_isolation_scope()._tags) + return "ok" + + res = handler(sns_event, MockContext()) + assert res == "ok" + + sentry_sdk.flush() + spans = [item.payload for item in items] + assert len(spans) == 1 + span = spans[0] + + assert span["attributes"]["messaging.batch.message_count"] == 1 + assert "batch_request" not in tags_recorded + assert "batch_size" not in tags_recorded + + +def test_handler_non_messaging_span_streaming(sentry_init, capture_items): + sentry_init( + integrations=[AwsLambdaIntegration()], + traces_sample_rate=1.0, + trace_lifecycle="stream", + ) + items = capture_items("span") + + http_event = {"httpMethod": "GET", "path": "/hello"} + + tags_recorded = {} + + @_wrap_handler + def handler(event, context): + tags_recorded.update(sentry_sdk.get_isolation_scope()._tags) + return "ok" + + res = handler(http_event, MockContext()) + assert res == "ok" + + sentry_sdk.flush() + spans = [item.payload for item in items] + assert len(spans) == 1 + span = spans[0] + + assert "messaging.batch.message_count" not in span["attributes"] + assert "batch_request" not in tags_recorded + + +def test_handler_top_level_list_span_streaming(sentry_init, capture_items): + sentry_init( + integrations=[AwsLambdaIntegration()], + traces_sample_rate=1.0, + trace_lifecycle="stream", + ) + items = capture_items("span") + + list_event = [{"data": 1}, {"data": 2}] + + tags_recorded = {} + + @_wrap_handler + def handler(event, context): + tags_recorded.update(sentry_sdk.get_isolation_scope()._tags) + return "ok" + + res = handler(list_event, MockContext()) + assert res == "ok" + + sentry_sdk.flush() + spans = [item.payload for item in items] + assert len(spans) == 1 + span = spans[0] + + # Non-messaging top-level list should not have messaging.batch.message_count + assert "messaging.batch.message_count" not in span["attributes"] + # But batch_request tags are set because len(list_event) > 1 + assert tags_recorded.get("batch_request") is True + assert tags_recorded.get("batch_size") == 2