@@ -15,9 +15,7 @@ def set_dsm_context(event, event_source):
15
15
_dsm_set_sns_context (event )
16
16
17
17
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 ):
21
19
"""
22
20
Common helper function for setting DSM context.
23
21
@@ -64,17 +62,22 @@ def _dsm_set_sqs_context(event):
64
62
65
63
def _dsm_set_sns_context (event ):
66
64
from ddtrace .internal .datastreams .botocore import calculate_sns_payload_size
65
+ from datadog_lambda .wrapper import format_err_with_traceback
67
66
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
76
70
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 ))
78
81
79
82
80
83
def _get_dsm_context_from_lambda (message ):
0 commit comments