diff --git a/providers/apache/kafka/docs/operators/index.rst b/providers/apache/kafka/docs/operators/index.rst index f53e7005af0f3..37f0df2d5d2df 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 3456534cce03e..a98d05f9c8df1 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,8 @@ 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``. :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 +83,7 @@ def __init__( 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 +97,7 @@ def __init__( 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 @@ -155,6 +159,7 @@ def execute(self, context) -> Any: ) messages_left = self.max_messages or True + apply_function_results: list[Any] = [] while self.read_to_end or ( messages_left > 0 @@ -174,7 +179,9 @@ def execute(self, context) -> Any: 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) @@ -189,7 +196,10 @@ def execute(self, context) -> Any: consumer.close() - 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 e883825b2ff07..574bf5297d9cb 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 @@ -43,7 +43,7 @@ def _no_op(*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. @@ -81,7 +81,34 @@ def mock_consume(num_messages=0, timeout=-1): ) # - 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: @@ -279,3 +306,61 @@ def test_commit_cadence_behavior(self, commit_cadence, max_messages, expected_co # Verify consumer was closed 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()