Fix Kafka consumer not being closed on error in ConsumeFromTopicOperator

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
 

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()