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 cf23900a6ab Kafka: Return consumed message results on request (#69740)
cf23900a6ab is described below
commit cf23900a6ab2a1ef4f9c8960948c701ef07a2410
Author: deepinsight coder <[email protected]>
AuthorDate: Wed Aug 12 10:41:38 2026 -0700
Kafka: Return consumed message results on request (#69740)
closes: #56494
---
providers/apache/kafka/docs/operators/index.rst | 4 +
.../providers/apache/kafka/operators/consume.py | 17 ++++-
.../unit/apache/kafka/operators/test_consume.py | 89 +++++++++++++++++++++-
3 files changed, 106 insertions(+), 4 deletions(-)
diff --git a/providers/apache/kafka/docs/operators/index.rst
b/providers/apache/kafka/docs/operators/index.rst
index f53e7005af0..37f0df2d5d2 100644
--- a/providers/apache/kafka/docs/operators/index.rst
+++ b/providers/apache/kafka/docs/operators/index.rst
@@ -33,6 +33,10 @@ For parameter definitions take a look at
:class:`~airflow.providers.apache.kafka
If you set the ``commit_cadence`` parameter, ensure that the
``enable.auto.commit`` option in the Kafka connection configuration is
explicitly set to ``false``.
By default, ``enable.auto.commit`` is ``true``, which causes the consumer
to auto-commit offsets every 5 seconds, potentially overriding the behavior
defined by ``commit_cadence``.
+Set ``return_apply_function_results=True`` to return a list containing each
non-``None`` value returned by the per-message ``apply_function`` in consume
order.
+This option does not apply to ``apply_function_batch``.
+Returned values use normal task return handling and may be stored in XCom, so
avoid returning large payloads.
+
Using the operator
""""""""""""""""""
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 05cc6e59b2b..fecc17ba6c6 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
@@ -45,6 +45,11 @@ class ConsumeFromTopicOperator(BaseOperator):
:param apply_function_args: Additional arguments that should be applied to
the callable, defaults to None
:param apply_function_kwargs: Additional key word arguments that should be
applied to the callable
defaults to None
+ :param return_apply_function_results: Whether to collect non-None return
values from the per-message
+ ``apply_function`` and return them as a list. This option does not
apply to ``apply_function_batch``.
+ ``None`` results are dropped, so the returned list does not align
positionally with the consumed
+ messages. Results are returned through normal task return handling and
may be stored in XCom,
+ so avoid enabling this for large result sets.
:param commit_cadence: When consumers should commit offsets ("never",
"end_of_batch","end_of_operator"),
defaults to "end_of_operator";
if end_of_operator, the commit() is called based on the max_messages
arg. Commits are made after the
@@ -81,6 +86,7 @@ class ConsumeFromTopicOperator(BaseOperator):
apply_function_batch: Callable[..., Any] | str | None = None,
apply_function_args: Sequence[Any] | None = None,
apply_function_kwargs: dict[Any, Any] | None = None,
+ return_apply_function_results: bool = False,
commit_cadence: str = "end_of_operator",
max_messages: int | None = None,
max_batch_size: int = 1000,
@@ -94,6 +100,7 @@ class ConsumeFromTopicOperator(BaseOperator):
self.apply_function_batch = apply_function_batch
self.apply_function_args = apply_function_args or ()
self.apply_function_kwargs = apply_function_kwargs or {}
+ self.return_apply_function_results = return_apply_function_results
self.kafka_config_id = kafka_config_id
self.commit_cadence = commit_cadence
self.max_messages = max_messages
@@ -156,6 +163,7 @@ class ConsumeFromTopicOperator(BaseOperator):
)
messages_left = self.max_messages or True
+ apply_function_results: list[Any] = []
while self.read_to_end or (
messages_left > 0
@@ -175,7 +183,9 @@ class ConsumeFromTopicOperator(BaseOperator):
if self.apply_function:
for m in msgs:
- apply_callable(m)
+ result = apply_callable(m)
+ if self.return_apply_function_results and result is
not None:
+ apply_function_results.append(result)
if self.apply_function_batch:
apply_callable(msgs)
@@ -193,7 +203,10 @@ class ConsumeFromTopicOperator(BaseOperator):
except Exception:
self.log.warning("Failed to close Kafka consumer",
exc_info=True)
- return
+ if self.return_apply_function_results and self.apply_function:
+ return apply_function_results
+
+ return None
def _validate_commit_cadence_on_construct(self):
"""Validate the commit_cadence parameter when the operator is
constructed."""
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 c7cf1b089d8..fef09159d67 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
@@ -48,7 +48,7 @@ def _raise_on_message(*args, **kwargs) -> Any:
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]:
+) -> tuple[mock.MagicMock, Any, list[int] | None]:
"""
Creates a mock Kafka consumer with configurable behavior.
@@ -86,7 +86,34 @@ def create_mock_kafka_consumer(
)
#
- return mock_consumer, mock_get_consumer, total_consumed_count # type:
ignore[return-value]
+ return mock_consumer, mock_get_consumer, total_consumed_count
+
+
+def create_mock_kafka_consumer_from_messages(
+ messages: list[Any],
+) -> tuple[mock.MagicMock, Any]:
+ mocked_messages = messages.copy()
+
+ 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 = mock.MagicMock()
+ mock_consumer.consume = mock_consume
+
+ mock_get_consumer = mock.patch(
+
"airflow.providers.apache.kafka.hooks.consume.KafkaConsumerHook.get_consumer",
+ return_value=mock_consumer,
+ )
+
+ return mock_consumer, mock_get_consumer
class TestConsumeFromTopic:
@@ -321,3 +348,61 @@ class TestConsumeFromTopic:
operator.execute(context={})
mock_consumer.close.assert_called_once()
+
+ def test_apply_function_results_return_none_by_default(self):
+ mock_consumer, mock_get_consumer =
create_mock_kafka_consumer_from_messages(["one", "two"])
+
+ with mock_get_consumer:
+ operator = ConsumeFromTopicOperator(
+ kafka_config_id="kafka_d",
+ topics=["test"],
+ task_id="test",
+ poll_timeout=0.0001,
+ max_messages=2,
+ max_batch_size=2,
+ apply_function=lambda message: f"processed-{message}",
+ )
+
+ assert operator.execute(context={}) is None
+ mock_consumer.close.assert_called_once()
+
+ def
test_return_apply_function_results_filters_none_and_preserves_order(self):
+ mock_consumer, mock_get_consumer =
create_mock_kafka_consumer_from_messages(
+ ["first", "skip", "second"]
+ )
+
+ def apply_function(message):
+ return None if message == "skip" else f"processed-{message}"
+
+ with mock_get_consumer:
+ operator = ConsumeFromTopicOperator(
+ kafka_config_id="kafka_d",
+ topics=["test"],
+ task_id="test",
+ poll_timeout=0.0001,
+ max_messages=3,
+ max_batch_size=2,
+ apply_function=apply_function,
+ return_apply_function_results=True,
+ )
+
+ assert operator.execute(context={}) == ["processed-first",
"processed-second"]
+ mock_consumer.close.assert_called_once()
+
+ def
test_return_apply_function_results_does_not_change_batch_return_behavior(self):
+ mock_consumer, mock_get_consumer =
create_mock_kafka_consumer_from_messages(["one", "two"])
+
+ with mock_get_consumer:
+ operator = ConsumeFromTopicOperator(
+ kafka_config_id="kafka_d",
+ topics=["test"],
+ task_id="test",
+ poll_timeout=0.0001,
+ max_messages=2,
+ max_batch_size=2,
+ apply_function_batch=lambda messages: [f"processed-{message}"
for message in messages],
+ return_apply_function_results=True,
+ )
+
+ assert operator.execute(context={}) is None
+ mock_consumer.close.assert_called_once()
