From d4a87a7481ff879f5a248f7b47648de1f27b0a4b Mon Sep 17 00:00:00 2001 From: Timur Rakhmatullin <174210871+TimurRakhmatullin86@users.noreply.github.com> Date: Wed, 8 Jul 2026 22:56:47 -0700 Subject: [PATCH] Fix Kafka consumer not being closed on error in ConsumeFromTopicOperator Signed-off-by: Timur Rakhmatullin <174210871+TimurRakhmatullin86@users.noreply.github.com> --- .../apache/kafka/operators/consume.py | 112 +++++++++--------- .../apache/kafka/operators/test_consume.py | 42 +++++++ 2 files changed, 100 insertions(+), 54 deletions(-) diff --git a/providers/apache/kafka/src/airflow/providers/apache/kafka/operators/consume.py b/providers/apache/kafka/src/airflow/providers/apache/kafka/operators/consume.py index 3456534cce03e..05cc6e59b2b5a 100644 --- a/providers/apache/kafka/src/airflow/providers/apache/kafka/operators/consume.py +++ b/providers/apache/kafka/src/airflow/providers/apache/kafka/operators/consume.py @@ -126,68 +126,72 @@ def execute(self, context) -> Any: self._validate_commit_cadence_before_execute() consumer = self.hook.get_consumer() - if isinstance(self.apply_function, str): - self.apply_function = import_string(self.apply_function) + try: + if isinstance(self.apply_function, str): + self.apply_function = import_string(self.apply_function) - if isinstance(self.apply_function_batch, str): - self.apply_function_batch = import_string(self.apply_function_batch) + if isinstance(self.apply_function_batch, str): + self.apply_function_batch = import_string(self.apply_function_batch) - if self.apply_function is not None and not callable(self.apply_function): - raise TypeError(f"apply_function is not a callable, got {type(self.apply_function)} instead.") - - if self.apply_function: - apply_callable = partial( - self.apply_function, - *self.apply_function_args, - **self.apply_function_kwargs, - ) - - if self.apply_function_batch is not None and not callable(self.apply_function_batch): - raise TypeError( - f"apply_function_batch is not a callable, got {type(self.apply_function_batch)} instead." - ) - - if self.apply_function_batch: - apply_callable = partial( - self.apply_function_batch, - *self.apply_function_args, - **self.apply_function_kwargs, - ) - - messages_left = self.max_messages or True - - while self.read_to_end or ( - messages_left > 0 - ): # bool(True > 0) == True in the case where self.max_messages isn't set by the user - if not isinstance(messages_left, bool): - batch_size = self.max_batch_size if messages_left > self.max_batch_size else messages_left - else: - batch_size = self.max_batch_size - - msgs = consumer.consume(num_messages=batch_size, timeout=self.poll_timeout) - 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.") - break + if self.apply_function is not None and not callable(self.apply_function): + raise TypeError(f"apply_function is not a callable, got {type(self.apply_function)} instead.") if self.apply_function: - for m in msgs: - apply_callable(m) + apply_callable = partial( + self.apply_function, + *self.apply_function_args, + **self.apply_function_kwargs, + ) - if self.apply_function_batch: - apply_callable(msgs) + if self.apply_function_batch is not None and not callable(self.apply_function_batch): + raise TypeError( + f"apply_function_batch is not a callable, got {type(self.apply_function_batch)} instead." + ) - if self.commit_cadence == "end_of_batch": + if self.apply_function_batch: + apply_callable = partial( + self.apply_function_batch, + *self.apply_function_args, + **self.apply_function_kwargs, + ) + + messages_left = self.max_messages or True + + while self.read_to_end or ( + messages_left > 0 + ): # bool(True > 0) == True in the case where self.max_messages isn't set by the user + if not isinstance(messages_left, bool): + batch_size = self.max_batch_size if messages_left > self.max_batch_size else messages_left + else: + batch_size = self.max_batch_size + + msgs = consumer.consume(num_messages=batch_size, timeout=self.poll_timeout) + 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.") + break + + if self.apply_function: + for m in msgs: + apply_callable(m) + + if self.apply_function_batch: + apply_callable(msgs) + + if self.commit_cadence == "end_of_batch": + self.log.info("committing offset at %s", self.commit_cadence) + consumer.commit() + + if self.commit_cadence != "never": self.log.info("committing offset at %s", self.commit_cadence) consumer.commit() - - if self.commit_cadence != "never": - self.log.info("committing offset at %s", self.commit_cadence) - consumer.commit() - - consumer.close() + finally: + try: + consumer.close() + except Exception: + self.log.warning("Failed to close Kafka consumer", exc_info=True) return diff --git a/providers/apache/kafka/tests/unit/apache/kafka/operators/test_consume.py b/providers/apache/kafka/tests/unit/apache/kafka/operators/test_consume.py index e883825b2ff07..c7cf1b089d82a 100644 --- a/providers/apache/kafka/tests/unit/apache/kafka/operators/test_consume.py +++ b/providers/apache/kafka/tests/unit/apache/kafka/operators/test_consume.py @@ -41,6 +41,11 @@ def _no_op(*args, **kwargs) -> Any: return args, kwargs +def _raise_on_message(*args, **kwargs) -> Any: + """A function that always raises, to simulate a failing apply_function.""" + raise ValueError("boom") + + def create_mock_kafka_consumer( message_count: int = 1001, message_content: Any = "test_message", track_consumed_messages: bool = False ) -> tuple[mock.MagicMock, mock.MagicMock, list[int] | None]: @@ -279,3 +284,40 @@ def test_commit_cadence_behavior(self, commit_cadence, max_messages, expected_co # Verify consumer was closed mock_consumer.close.assert_called_once() + + def test_execute_closes_consumer_when_apply_function_raises(self): + """The consumer must be closed even if message processing raises.""" + mock_consumer, mock_get_consumer, _ = create_mock_kafka_consumer(message_count=5) + + with mock_get_consumer: + operator = ConsumeFromTopicOperator( + kafka_config_id="kafka_d", + topics=["test"], + task_id="test", + poll_timeout=0.0001, + apply_function=_raise_on_message, + ) + + with pytest.raises(ValueError, match="boom"): + operator.execute(context={}) + + mock_consumer.close.assert_called_once() + + def test_execute_does_not_mask_error_when_close_raises(self): + """A failing close() must not replace the original processing error.""" + mock_consumer, mock_get_consumer, _ = create_mock_kafka_consumer(message_count=5) + mock_consumer.close.side_effect = Exception("close failed") + + with mock_get_consumer: + operator = ConsumeFromTopicOperator( + kafka_config_id="kafka_d", + topics=["test"], + task_id="test", + poll_timeout=0.0001, + apply_function=_raise_on_message, + ) + + with pytest.raises(ValueError, match="boom"): + operator.execute(context={}) + + mock_consumer.close.assert_called_once()