This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new e7737ff9bb1 Fix Kafka consumer not being closed on error in
ConsumeFromTopicOperator (#69641)
e7737ff9bb1 is described below
commit e7737ff9bb1229f637bef132a5829e77505814e6
Author: Tim <[email protected]>
AuthorDate: Sat Aug 1 10:58:34 2026 -0700
Fix Kafka consumer not being closed on error in ConsumeFromTopicOperator
(#69641)
Signed-off-by: Timur Rakhmatullin
<[email protected]>
Co-authored-by: Timur Rakhmatullin
<[email protected]>
---
.../providers/apache/kafka/operators/consume.py | 112 +++++++++++----------
.../unit/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 3456534cce0..05cc6e59b2b 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 @@ class ConsumeFromTopicOperator(BaseOperator):
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 e883825b2ff..c7cf1b089d8 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 @@ class TestConsumeFromTopic:
# 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()
