[Backport 7.78.x] kafka_actions: Return early when end of partition is reached

kafka_actions/changelog.d/23004.fixed

@@ -0,0 +1 @@
+Return early from consume_messages when end of partition is reached instead of waiting for timeout

kafka_actions/datadog_checks/kafka_actions/kafka_client.py

@@ -8,7 +8,7 @@
 import time
 from typing import TYPE_CHECKING, Any
 
-from confluent_kafka import Consumer, KafkaException, Producer, TopicPartition
+from confluent_kafka import Consumer, KafkaError, KafkaException, Producer, TopicPartition
 from confluent_kafka.admin import AdminClient, ConfigResource, NewTopic, OffsetSpec, ResourceType
 
 try:
@@ -272,10 +272,11 @@ def consume_messages(
                 msg = consumer.poll(timeout=poll_timeout)
 
                 if msg is None:
-                    continue
+                    self.log.debug("Poll returned None (no more messages available), stopping consumption")
+                    break
 
                 if msg.error():
-                    if msg.error().code() == KafkaException._PARTITION_EOF:
+                    if msg.error().code() == KafkaError._PARTITION_EOF:
                         self.log.debug("Reached end of partition")
                         continue
                     else:

kafka_actions/tests/test_unit.py

@@ -524,5 +524,29 @@ def test_start_timestamp_overrides_start_offset(self, dd_run_check, aggregator):
             assert call_kwargs['start_timestamp'] == 1700000000000
 
 
+class TestConsumeMessagesEarlyReturn:
+    """Test that consume_messages returns early on None poll or partition EOF."""
+
+    def test_none_poll_breaks_immediately(self):
+        from unittest.mock import MagicMock
+
+        consumer = MagicMock()
+        metadata = MagicMock()
+        metadata.topics = {'t': MagicMock(partitions={0: MagicMock()})}
+        consumer.list_topics.return_value = metadata
+        consumer.poll.side_effect = [MockKafkaMessage(key=b'k', value=b'v', partition=0, offset=0), None]
+
+        import logging
+
+        from datadog_checks.kafka_actions.kafka_client import KafkaActionsClient
+
+        client = KafkaActionsClient({'kafka_connect_str': 'localhost:9092'}, logging.getLogger('test'))
+        with patch.object(client, 'get_consumer', return_value=consumer):
+            result = list(client.consume_messages(topic='t', start_offset=0, max_messages=1000, timeout_ms=30000))
+
+        assert len(result) == 1
+        assert consumer.poll.call_count == 2
+
+
 if __name__ == '__main__':
     pytest.main([__file__, '-vv'])