add more batch processing tests

michael-zhao459 · michael-zhao459 · commit 60ca41badd59 · 2025-08-06T11:05:45.000-04:00
diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py
@@ -68,32 +68,6 @@
 LOWER_64_BITS = "LOWER_64_BITS"
 
 
-def _dsm_set_context_sqs_event(event):
-    for record in event.get("Records", [])[1:]:
-        if arn := record.get("eventSourceARN"):
-            try:
-                context = _extract_context_from_sqs_or_sns_record(record)
-            except Exception as e:
-                logger.debug(
-                    f"DSM: Failed to extract context with error {e}, will still set checkpoint"
-                )
-                context = None
-            _dsm_set_checkpoint(context, "sqs", arn)
-
-
-def _dsm_set_context_kinesis_event(event):
-    for record in event.get("Records", [])[1:]:
-        if (arn := record.get("eventSourceARN")) and (kinesis := record.get("kinesis")):
-            try:
-                context = _extract_context_from_kinesis_record(kinesis)
-            except Exception as e:
-                logger.debug(
-                    f"DSM: Failed to extract context with error {e}, will still set checkpoint"
-                )
-                context = None
-            _dsm_set_checkpoint(context, "kinesis", arn)
-
-
 def _dsm_set_checkpoint(context_json, event_type, arn):
     if not config.data_streams_enabled:
         return
@@ -263,6 +237,7 @@ def extract_context_from_sqs_or_sns_event_or_context(
     Falls back to lambda context if no trace data is found in the SQS message attributes.
     Set a DSM checkpoint if DSM is enabled and the method for context propagation is supported.
     """
+    source_arn = ""
     event_type = "sqs" if event_source.equals(EventTypes.SQS) else "sns"
 
     # EventBridge => SQS
@@ -273,106 +248,98 @@ def extract_context_from_sqs_or_sns_event_or_context(
     except Exception:
         logger.debug("Failed extracting context as EventBridge to SQS.")
 
-    try:
-        first_record = event.get("Records")[0]
-        source_arn = (
-            first_record.get("Sns", {}).get("TopicArn")
-            if event_type == "sns"
-            else first_record.get("eventSourceARN")
-        )
-        dd_data = _extract_context_from_sqs_or_sns_record(first_record)
-        if dd_data:
-            if is_step_function_event(dd_data):
+    context = None
+    for idx, record in enumerate(event.get("Records", [])):
+        try:
+            source_arn = record.get("eventSourceARN", "")
+            dd_ctx = None
+
+            # logic to deal with SNS => SQS event
+            if "body" in record:
+                body_str = record.get("body")
                 try:
-                    return extract_context_from_step_functions(dd_data, None)
+                    body = json.loads(body_str)
+                    if body.get("Type", "") == "Notification" and "TopicArn" in body:
+                        logger.debug("Found SNS message inside SQS event")
+                        record = get_first_record(create_sns_event(body))
                 except Exception:
+                    pass
+
+            msg_attributes = record.get("messageAttributes")
+            if msg_attributes is None:
+                sns_record = record.get("Sns") or {}
+                # SNS->SQS event would extract SNS arn without this check
+                if event_source.equals(EventTypes.SNS):
+                    source_arn = sns_record.get("TopicArn", "")
+                msg_attributes = sns_record.get("MessageAttributes") or {}
+            dd_payload = msg_attributes.get("_datadog")
+            if dd_payload:
+                # SQS uses dataType and binaryValue/stringValue
+                # SNS uses Type and Value
+                # fmt: off
+                dd_json_data = None
+                dd_json_data_type = dd_payload.get("Type") or dd_payload.get("dataType")
+                if dd_json_data_type == "Binary":
+                    import base64
+
+                    dd_json_data = dd_payload.get("binaryValue") or dd_payload.get("Value")
+                    if dd_json_data:
+                        dd_json_data = base64.b64decode(dd_json_data)
+                elif dd_json_data_type == "String":
+                    dd_json_data = dd_payload.get("stringValue") or dd_payload.get("Value")
+                # fmt: on
+                else:
                     logger.debug(
-                        "Failed to extract Step Functions context from SQS/SNS event."
+                        "Datadog Lambda Python only supports extracting trace"
+                        "context from String or Binary SQS/SNS message attributes"
                     )
-            context = propagator.extract(dd_data)
-            _dsm_set_checkpoint(dd_data, event_type, source_arn)
-            # Batch Processing does not occur for SNS events
-            if config.data_streams_enabled and event_type == "sqs":
-                _dsm_set_context_sqs_event(event)
-            return context
-        else:
-            # Handle case where trace context is injected into attributes.AWSTraceHeader
-            # example: Root=1-654321ab-000000001234567890abcdef;Parent=0123456789abcdef;Sampled=1
-            attrs = event.get("Records")[0].get("attributes")
-            if attrs:
-                x_ray_header = attrs.get("AWSTraceHeader")
-                if x_ray_header:
-                    x_ray_context = parse_xray_header(x_ray_header)
-                    trace_id_parts = x_ray_context.get("trace_id", "").split("-")
-                    if len(trace_id_parts) > 2 and trace_id_parts[2].startswith(
-                        DD_TRACE_JAVA_TRACE_ID_PADDING
-                    ):
-                        # If it starts with eight 0's padding,
-                        # then this AWSTraceHeader contains Datadog injected trace context
-                        logger.debug(
-                            "Found dd-trace injected trace context from AWSTraceHeader"
-                        )
-                        return Context(
-                            trace_id=int(trace_id_parts[2][8:], 16),
-                            span_id=int(x_ray_context["parent_id"], 16),
-                            sampling_priority=float(x_ray_context["sampled"]),
-                        )
-        # Still want to set a DSM checkpoint even if DSM context not propagated
-        _dsm_set_checkpoint(None, event_type, source_arn)
-        # Batch Processing does not occur for SNS events
-        if config.data_streams_enabled and event_type == "sqs":
-            _dsm_set_context_sqs_event(event)
-        return extract_context_from_lambda_context(lambda_context)
-    except Exception as e:
-        logger.debug("The trace extractor returned with error %s", e)
-        # Still want to set a DSM checkpoint even if DSM context not propagated
-        _dsm_set_checkpoint(None, event_type, source_arn)
-        # Batch Processing does not occur for SNS events
-        if config.data_streams_enabled and event_type == "sqs":
-            _dsm_set_context_sqs_event(event)
-        return extract_context_from_lambda_context(lambda_context)
 
+                if dd_json_data:
+                    dd_data = json.loads(dd_json_data)
+
+                    if idx == 0:
+                        if is_step_function_event(dd_data):
+                            try:
+                                return extract_context_from_step_functions(
+                                    dd_data, None
+                                )
+                            except Exception:
+                                logger.debug(
+                                    "Failed to extract Step Functions context from SQS/SNS event."
+                                )
+                        context = propagator.extract(dd_data)
+                        if not config.data_streams_enabled:
+                            break
+                    dd_ctx = dd_data
+            elif idx == 0:
+                # Handle case where trace context is injected into attributes.AWSTraceHeader
+                # example:Root=1-654321ab-000000001234567890abcdef;Parent=0123456789abcdef;Sampled=1
+                attrs = event.get("Records")[0].get("attributes")
+                if attrs:
+                    x_ray_header = attrs.get("AWSTraceHeader")
+                    if x_ray_header:
+                        x_ray_context = parse_xray_header(x_ray_header)
+                        trace_id_parts = x_ray_context.get("trace_id", "").split("-")
+                        if len(trace_id_parts) > 2 and trace_id_parts[2].startswith(
+                            DD_TRACE_JAVA_TRACE_ID_PADDING
+                        ):
+                            # If it starts with eight 0's padding,
+                            # then this AWSTraceHeader contains Datadog injected trace context
+                            logger.debug(
+                                "Found dd-trace injected trace context from AWSTraceHeader"
+                            )
+                            return Context(
+                                trace_id=int(trace_id_parts[2][8:], 16),
+                                span_id=int(x_ray_context["parent_id"], 16),
+                                sampling_priority=float(x_ray_context["sampled"]),
+                            )
+        except Exception as e:
+            logger.debug("The trace extractor returned with error %s", e)
 
-def _extract_context_from_sqs_or_sns_record(record):
-    # logic to deal with SNS => SQS event
-    if "body" in record:
-        body_str = record.get("body")
-        try:
-            body = json.loads(body_str)
-            if body.get("Type", "") == "Notification" and "TopicArn" in body:
-                logger.debug("Found SNS message inside SQS event")
-                record = get_first_record(create_sns_event(body))
-        except Exception:
-            pass
-
-    msg_attributes = record.get("messageAttributes")
-    if msg_attributes is None:
-        sns_record = record.get("Sns") or {}
-        msg_attributes = sns_record.get("MessageAttributes") or {}
-    dd_payload = msg_attributes.get("_datadog")
-    if dd_payload:
-        # SQS uses dataType and binaryValue/stringValue
-        # SNS uses Type and Value
-        dd_json_data = None
-        dd_json_data_type = dd_payload.get("Type") or dd_payload.get("dataType")
-        if dd_json_data_type == "Binary":
-            import base64
-
-            dd_json_data = dd_payload.get("binaryValue") or dd_payload.get("Value")
-            if dd_json_data:
-                dd_json_data = base64.b64decode(dd_json_data)
-        elif dd_json_data_type == "String":
-            dd_json_data = dd_payload.get("stringValue") or dd_payload.get("Value")
-        else:
-            logger.debug(
-                "Datadog Lambda Python only supports extracting trace"
-                "context from String or Binary SQS/SNS message attributes"
-            )
+        # Set DSM checkpoint once per record
+        _dsm_set_checkpoint(dd_ctx, event_type, source_arn)
 
-        if dd_json_data:
-            dd_data = json.loads(dd_json_data)
-            return dd_data
-    return None
+    return context if context else extract_context_from_lambda_context(lambda_context)
 
 
 def _extract_context_from_eventbridge_sqs_event(event):
@@ -433,40 +400,35 @@ def extract_context_from_kinesis_event(event, lambda_context):
     Set a DSM checkpoint if DSM is enabled and the method for context propagation is supported.
     """
     source_arn = ""
-    try:
-        record = get_first_record(event)
-        source_arn = record.get("eventSourceARN", "")
-        kinesis = record.get("kinesis")
-        if not kinesis:
-            return extract_context_from_lambda_context(lambda_context)
-        dd_ctx = _extract_context_from_kinesis_record(kinesis)
-        if dd_ctx:
-            context = propagator.extract(dd_ctx)
-            _dsm_set_checkpoint(dd_ctx, "kinesis", source_arn)
-            if config.data_streams_enabled:
-                _dsm_set_context_kinesis_event(event)
-            return context
-    except Exception as e:
-        logger.debug("The trace extractor returned with error %s", e)
-    # Still want to set a DSM checkpoint even if DSM context not propagated
-    _dsm_set_checkpoint(None, "kinesis", source_arn)
-    if config.data_streams_enabled:
-        _dsm_set_context_kinesis_event(event)
-    return extract_context_from_lambda_context(lambda_context)
-
 
-def _extract_context_from_kinesis_record(record_kinesis_data):
-    data = record_kinesis_data.get("data")
-    if data:
-        import base64
-
-        b64_bytes = data.encode("ascii")
-        str_bytes = base64.b64decode(b64_bytes)
-        data_str = str_bytes.decode("ascii")
-        data_obj = json.loads(data_str)
-        dd_ctx = data_obj.get("_datadog")
-        return dd_ctx
-    return None
+    context = None
+    for idx, record in enumerate(event.get("Records", [])):
+        dd_ctx = None
+        try:
+            source_arn = record.get("eventSourceARN", "")
+            kinesis = record.get("kinesis")
+            if not kinesis:
+                if idx == 0:
+                    return extract_context_from_lambda_context(lambda_context)
+                continue
+            data = kinesis.get("data")
+            if data:
+                import base64
+
+                b64_bytes = data.encode("ascii")
+                str_bytes = base64.b64decode(b64_bytes)
+                data_str = str_bytes.decode("ascii")
+                data_obj = json.loads(data_str)
+                dd_ctx = data_obj.get("_datadog")
+                if dd_ctx:
+                    if idx == 0:
+                        context = propagator.extract(dd_ctx)
+                        if not config.data_streams_enabled:
+                            break
+        except Exception as e:
+            logger.debug("The trace extractor returned with error %s", e)
+        _dsm_set_checkpoint(dd_ctx, "kinesis", source_arn)
+    return context if context else extract_context_from_lambda_context(lambda_context)
 
 
 def _deterministic_sha256_hash(s: str, part: str) -> int:
