From b96930d2f0085bfa95d42f9dc070fb5fa1beabfe Mon Sep 17 00:00:00 2001 From: Aritra Basu <24430013+aritra24@users.noreply.github.com> Date: Fri, 11 Aug 2023 23:04:26 +0530 Subject: [PATCH 1/2] Fixes kafka provider failing reading messages Fixes the issue(#32926) where kafka provider returns an error when max messages is not set since it keeps reading and messages left goes into negative. --- airflow/providers/apache/kafka/operators/consume.py | 4 +++- .../providers/apache/kafka/operators/test_consume.py | 4 ++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/airflow/providers/apache/kafka/operators/consume.py b/airflow/providers/apache/kafka/operators/consume.py index 02c9db6556df6..4f82b77c42aae 100644 --- a/airflow/providers/apache/kafka/operators/consume.py +++ b/airflow/providers/apache/kafka/operators/consume.py @@ -161,7 +161,9 @@ def execute(self, context) -> Any: batch_size = self.max_batch_size msgs = consumer.consume(num_messages=batch_size, timeout=self.poll_timeout) - messages_left -= len(msgs) + self.log.info("Messages count is %s", len(msgs)) + if not self.read_to_end: + messages_left -= len(msgs) if not msgs: # No messages + messages_left is being used. self.log.info("Reached end of log. Exiting.") diff --git a/tests/integration/providers/apache/kafka/operators/test_consume.py b/tests/integration/providers/apache/kafka/operators/test_consume.py index 240b02f9be1ad..861bbbcfab35d 100644 --- a/tests/integration/providers/apache/kafka/operators/test_consume.py +++ b/tests/integration/providers/apache/kafka/operators/test_consume.py @@ -65,7 +65,7 @@ def setup_method(self): extra=json.dumps( { "socket.timeout.ms": 10, - "bootstrap.servers": "localhost:9092", + "bootstrap.servers": "broker:29092", "group.id": f"operator.consumer.test.integration.test_{num}", "enable.auto.commit": False, "auto.offset.reset": "beginning", @@ -135,7 +135,7 @@ def test_consumer_operator_test_3(self): operator = ConsumeFromTopicOperator( kafka_config_id=TOPIC, topics=[TOPIC], - apply_function=_batch_tester, + apply_function_batch=_batch_tester, apply_function_kwargs={"test_string": TOPIC}, task_id="test", poll_timeout=0.0001, From 4923a91ad4cb860b39e5c1bc62a38c986b760cc6 Mon Sep 17 00:00:00 2001 From: Aritra Basu <24430013+aritra24@users.noreply.github.com> Date: Fri, 11 Aug 2023 23:46:57 +0530 Subject: [PATCH 2/2] Adds a unit test for kafka provider consume Makes sure that the kafka provider works when max messages isn't passed into the operator --- .../apache/kafka/operators/consume.py | 1 - .../apache/kafka/operators/test_consume.py | 29 +++++++++++++++++++ 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/airflow/providers/apache/kafka/operators/consume.py b/airflow/providers/apache/kafka/operators/consume.py index 4f82b77c42aae..6f1c1ee61e50c 100644 --- a/airflow/providers/apache/kafka/operators/consume.py +++ b/airflow/providers/apache/kafka/operators/consume.py @@ -161,7 +161,6 @@ def execute(self, context) -> Any: batch_size = self.max_batch_size msgs = consumer.consume(num_messages=batch_size, timeout=self.poll_timeout) - self.log.info("Messages count is %s", len(msgs)) if not self.read_to_end: messages_left -= len(msgs) diff --git a/tests/providers/apache/kafka/operators/test_consume.py b/tests/providers/apache/kafka/operators/test_consume.py index 178419052699c..0ce1ce6e32eb9 100644 --- a/tests/providers/apache/kafka/operators/test_consume.py +++ b/tests/providers/apache/kafka/operators/test_consume.py @@ -19,6 +19,7 @@ import json import logging from typing import Any +from unittest import mock from airflow.models import Connection @@ -79,3 +80,31 @@ def test_operator_callable(self): # execute the operator (this is essentially a no op as the broker isn't setup) operator.execute(context={}) + + @mock.patch("airflow.providers.apache.kafka.hooks.consume.KafkaConsumerHook.get_consumer") + def test_operator_consume_max(self, mock_get_consumer): + mock_consumer = mock.MagicMock() + + mocked_messages = ["test_messages" for i in range(1001)] + + def mock_consume(num_messages=0, timeout=-1): + nonlocal mocked_messages + if num_messages < 0: + raise Exception("Number of messages needs to be positive") + msg_count = min(num_messages, len(mocked_messages)) + returned_messages = mocked_messages[:msg_count] + mocked_messages = mocked_messages[msg_count:] + return returned_messages + + mock_consumer.consume = mock_consume + mock_get_consumer.return_value = mock_consumer + + operator = ConsumeFromTopicOperator( + kafka_config_id="kafka_d", + topics=["test"], + task_id="test", + poll_timeout=0.0001, + ) + + # execute the operator (this is essentially a no op as we're mocking the consumer) + operator.execute(context={})