Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 36 additions & 2 deletions sentry_sdk/integrations/aws_lambda.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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:
Expand Down Expand Up @@ -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:
Expand Down
4 changes: 2 additions & 2 deletions tests/integrations/aws_lambda/test_aws_lambda.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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):
Expand Down
244 changes: 244 additions & 0 deletions tests/integrations/aws_lambda/test_aws_lambda_messaging.py
Original file line number Diff line number Diff line change
@@ -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