Skip to content

Commit 5ef804e

Browse files
fix
1 parent d43001b commit 5ef804e

File tree

1 file changed

+15
-12
lines changed

1 file changed

+15
-12
lines changed

datadog_lambda/dsm.py

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,7 @@ def set_dsm_context(event, event_source):
1515
_dsm_set_sns_context(event)
1616

1717

18-
def _dsm_set_context_helper(
19-
record, service_type, arn, payload_size
20-
):
18+
def _dsm_set_context_helper(record, service_type, arn, payload_size):
2119
"""
2220
Common helper function for setting DSM context.
2321
@@ -64,17 +62,22 @@ def _dsm_set_sqs_context(event):
6462

6563
def _dsm_set_sns_context(event):
6664
from ddtrace.internal.datastreams.botocore import calculate_sns_payload_size
65+
from datadog_lambda.wrapper import format_err_with_traceback
6766

68-
def sns_payload_calculator(record, context_json):
69-
return calculate_sns_payload_size(record, context_json)
70-
71-
def sns_arn_extractor(record):
72-
sns_data = record.get("Sns")
73-
if not sns_data:
74-
return ""
75-
return sns_data.get("TopicArn", "")
67+
records = event.get("Records")
68+
if records is None:
69+
return
7670

77-
_dsm_set_context_helper(event, "sns", sns_arn_extractor, sns_payload_calculator)
71+
for record in records:
72+
try:
73+
sns_data = record.get("Sns")
74+
if not sns_data:
75+
return ""
76+
arn = sns_data.get("TopicArn", "")
77+
payload_size = calculate_sns_payload_size(record)
78+
_dsm_set_context_helper(record, "sns", arn, payload_size)
79+
except Exception as e:
80+
logger.error(format_err_with_traceback(e))
7881

7982

8083
def _get_dsm_context_from_lambda(message):

0 commit comments

Comments
 (0)