Copilot commented on code in PR #64833: URL: https://github.com/apache/airflow/pull/64833#discussion_r3154537678
########## airflow-e2e-tests/tests/airflow_e2e_tests/event_driven_tests/test_event_driven.py: ########## @@ -0,0 +1,229 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""E2E tests for the event-driven DAG pattern using Apache Kafka. + +The producer DAG sends 8 valid orders + 1 malformed message to the ``fizz_buzz`` +Kafka topic. The consumer DAG is scheduled on an Asset with a Kafka AssetWatcher, +so each message triggers a separate consumer DAG run (9 total). +""" + +from __future__ import annotations + +import time + +from airflow_e2e_tests.e2e_test_utils.clients import AirflowClient + +PRODUCER_DAG_ID = "event_driven_producer" +CONSUMER_DAG_ID = "event_driven_consumer" +# 8 valid orders + 1 malformed message +EXPECTED_CONSUMER_RUNS = 9 +EXPECTED_FIZZ_BUZZ_OFFSET = 9 +EXPECTED_DLQ_OFFSET = 1 + + +def _parse_topic_offset(raw_output: str, topic: str) -> int: + """Parse ``kafka-get-offsets`` output (``topic:partition:offset``) and return the offset.""" + for line in raw_output.strip().splitlines(): + parts = line.strip().split(":") + if len(parts) == 3 and parts[0] == topic: + return int(parts[2]) + raise ValueError(f"Could not find offset for topic '{topic}' in output:\n{raw_output}") + + +class TestEventDrivenDag: + """Test the Kafka-based event-driven producer/consumer DAG pair.""" + + airflow_client = AirflowClient() + + def _wait_for_kafka_consumer_group(self, compose_instance, timeout: int = 60, check_interval: int = 3): + """Poll until any Kafka consumer group is registered, indicating the trigger is active.""" + start = time.monotonic() + while time.monotonic() - start < timeout: + stdout, _, _ = compose_instance.exec_in_container( + command=[ + "kafka-consumer-groups", + "--bootstrap-server", + "broker:29092", + "--list", + ], + service_name="broker", + ) + output = stdout.decode() if isinstance(stdout, bytes) else stdout + # Any non-empty group listing means the trigger's consumer has registered + groups = [line.strip() for line in output.strip().splitlines() if line.strip()] + if groups: + return + time.sleep(check_interval) + raise TimeoutError(f"No Kafka consumer group registered within {timeout}s") + + def _wait_for_consumer_dag_runs( + self, expected_count: int, timeout: int = 600, check_interval: int = 10 + ) -> list[dict]: + """Poll until *expected_count* consumer DAG runs reach a terminal state.""" + start = time.monotonic() + while time.monotonic() - start < timeout: + response = self.airflow_client.list_dag_runs(CONSUMER_DAG_ID) + runs = response.get("dag_runs", []) + terminal_runs = [r for r in runs if r["state"] in ("success", "failed")] + if len(terminal_runs) >= expected_count: + # Return only the most recent expected_count runs to avoid + # assertion failures if extra runs exist from retries or prior state. + terminal_runs.sort( + key=lambda r: ( + r.get("end_date") or "", + r.get("start_date") or "", + r.get("logical_date") or "", + r.get("dag_run_id") or "", + ), + reverse=True, + ) + return terminal_runs[:expected_count] + time.sleep(check_interval) + + # Timed out — gather diagnostics + response = self.airflow_client.list_dag_runs(CONSUMER_DAG_ID) + runs = response.get("dag_runs", []) + states = {r["dag_run_id"]: r["state"] for r in runs} + raise TimeoutError( + f"Expected {expected_count} terminal consumer DAG runs within {timeout}s, " + f"but only found {len([r for r in runs if r['state'] in ('success', 'failed')])}. " + f"Run states: {states}" + ) + + def _get_topic_offset(self, compose_instance, topic: str) -> int: + """Return the current end-offset of *topic* via ``kafka-get-offsets`` inside the broker.""" + stdout, _, _ = compose_instance.exec_in_container( + command=[ + "kafka-get-offsets", + "--bootstrap-server", + "broker:29092", + "--topic", + topic, + ], + service_name="broker", + ) + output = stdout.decode() if isinstance(stdout, bytes) else stdout + return _parse_topic_offset(output, topic) + + def _wait_for_topic_offset( + self, compose_instance, topic: str, expected_offset: int, timeout: int = 30, check_interval: int = 2 + ) -> int: + """Poll until *topic* reaches *expected_offset*, then return the offset.""" + start = time.monotonic() + offset = self._get_topic_offset(compose_instance, topic) + while offset < expected_offset and time.monotonic() - start < timeout: + time.sleep(check_interval) + offset = self._get_topic_offset(compose_instance, topic) + return offset + + def _get_task_states(self, run_id: str) -> dict[str, str]: + """Return a mapping of task_id -> state for a consumer DAG run.""" + response = self.airflow_client.get_task_instances(CONSUMER_DAG_ID, run_id) + return {ti["task_id"]: ti["state"] for ti in response["task_instances"]} + + def test_producer_triggers_consumer_and_kafka_offsets(self, compose_instance): + """Trigger the producer once and verify 9 consumer runs and Kafka offsets. + + Steps: + 1. Unpause the consumer DAG so the triggerer starts the AssetWatcher. + 2. Wait for the Kafka MessageQueueTrigger to begin polling. + 3. Trigger the producer DAG and wait for it to succeed. + 4. Wait for 9 consumer DAG runs to reach a terminal state. + 5. All 9 DAG runs succeed. Verify task-level behavior: + - 1 run has a failed ``process_message`` task and executes ``handle_dlq``. + - 8 runs succeed on ``process_message`` and skip ``handle_dlq``. + 6. Verify that the ``fizz_buzz`` topic has offset 9 (all messages produced). + 7. Verify that the ``dlq`` topic has offset 1 (the malformed message). + """ + # 1. Unpause consumer so the triggerer registers the AssetWatcher + self.airflow_client.un_pause_dag(CONSUMER_DAG_ID) + + # 2. Wait for the triggerer to start the MessageQueueTrigger and subscribe. + # The trigger uses poll_interval=1 and auto.offset.reset=latest so it + # must be actively polling before the producer writes. + self._wait_for_kafka_consumer_group(compose_instance) + + # 3. Trigger producer and wait for it to complete + producer_state = self.airflow_client.trigger_dag_and_wait(PRODUCER_DAG_ID) + assert producer_state == "success", f"Producer DAG did not succeed. Final state: {producer_state}" + + # 4. Wait for all 9 consumer DAG runs + consumer_runs = self._wait_for_consumer_dag_runs(expected_count=EXPECTED_CONSUMER_RUNS) + assert len(consumer_runs) == EXPECTED_CONSUMER_RUNS, ( + f"Expected {EXPECTED_CONSUMER_RUNS} consumer runs, got {len(consumer_runs)}" + ) + + # 5. All 9 DAG runs should succeed + for run in consumer_runs: + assert run["state"] == "success", ( + f"Expected all consumer runs to succeed, but run {run['dag_run_id']} " + f"has state '{run['state']}'" + ) + Review Comment: This test asserts every consumer DAG run is `success`, but later it expects one run to have a `process_message` task in `failed` state. In Airflow, a DagRun is marked failed if any task instance is in a failed state once the run finishes (see `DagRun.update_state`). Either adjust the expectations to allow one failed DagRun, or change the consumer DAG logic so the malformed message path does not leave a task instance in `failed` while still producing to the DLQ. ```suggestion # 5. Expect 8 successful consumer runs and 1 failed run for the malformed message successful_runs = [run for run in consumer_runs if run["state"] == "success"] failed_runs = [run for run in consumer_runs if run["state"] == "failed"] assert len(successful_runs) == 8, ( f"Expected 8 successful consumer runs, got {len(successful_runs)}: " f"{[(run['dag_run_id'], run['state']) for run in consumer_runs]}" ) assert len(failed_runs) == 1, ( f"Expected 1 failed consumer run for the DLQ path, got {len(failed_runs)}: " f"{[(run['dag_run_id'], run['state']) for run in consumer_runs]}" ) ``` ########## dev/breeze/src/airflow_breeze/utils/selective_checks.py: ########## @@ -204,6 +205,14 @@ def __hash__(self): r"^airflow-e2e-tests/tests/airflow_e2e_tests/remote_log_opensearch_tests/.*", r"^providers/opensearch/.*", ], + FileGroupForCi.EVENT_DRIVEN_E2E_FILES: [ + r"^airflow-e2e-tests/tests/airflow_e2e_tests/event_driven_tests/.*", + r"^airflow-e2e-tests/tests/airflow_e2e_tests/dags/example_event_driven\.py$", + r"^airflow-e2e-tests/docker/kafka/.*", + r"^airflow-e2e-tests/docker/kafka\.yml$", + r"^providers/apache-kafka/.*", + r"^providers/common-messaging/.*", Review Comment: The provider path regexes here don’t match the repository layout. Providers live under `providers/apache/kafka/` and `providers/common/messaging/`, so these patterns will never match and CI selective-checks won’t enable the event-driven e2e suite when the Kafka/common-messaging providers change. Update the regexes to the correct directory paths. ```suggestion r"^providers/apache/kafka/.*", r"^providers/common/messaging/.*", ``` ########## airflow-e2e-tests/tests/airflow_e2e_tests/conftest.py: ########## @@ -121,6 +122,58 @@ def _setup_opensearch_integration(dot_env_file, tmp_dir): os.environ["ENV_FILE_PATH"] = str(dot_env_file) +def _copy_kafka_files(tmp_dir): + """Copy Kafka compose file, init script, and provider source into the temp directory.""" Review Comment: This docstring says it copies “provider source” into the temp directory, but the function only copies `kafka.yml` and `update_run.sh`. Please either update the docstring to match what’s actually being copied, or add the missing copy step if provider source is required for the integration. ```suggestion """Copy the Kafka compose file and update script into the temp directory.""" ``` ########## airflow-e2e-tests/tests/airflow_e2e_tests/dags/example_event_driven.py: ########## @@ -0,0 +1,178 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +import json +from datetime import timedelta +from typing import TYPE_CHECKING, cast + +import pendulum + +from airflow.providers.apache.kafka.operators.produce import ProduceToTopicOperator +from airflow.providers.common.messaging.triggers.msg_queue import MessageQueueTrigger +from airflow.sdk import Asset, AssetWatcher, dag, get_current_context, task + +if TYPE_CHECKING: + from airflow.sdk.execution_time.context import Context, TriggeringAssetEventsAccessor + +KAFKA_CONFIG_ID = "kafka_default" +TOPICS = ["fizz_buzz"] +DLQ_TOPIC = "dlq" +RETRY_COUNT = 3 +"""Airflow Kafka connection +AIRFLOW_CONN_KAFKA_DEFAULT='{ + "conn_type": "general", + "extra": { + "bootstrap.servers": "broker:29092", + "group.id": "kafka_default_group", + "security.protocol": "PLAINTEXT", + "enable.auto.commit": false, + "auto.offset.reset": "latest" + } +}' +""" +"""Kafka Command to verify messages are being produced to the topic: + +# Create Topic +/bin/kafka-topics --bootstrap-server broker:29092 --create --topic fizz_buzz +/bin/kafka-topics --bootstrap-server broker:29092 --create --topic dlq + + +# Get offsets for the topic to verify messages are being produced +/bin/kafka-get-offsets --bootstrap-server broker:29092 --topic fizz_buzz +/bin/kafka-get-offsets --bootstrap-server broker:29092 --topic dlq + +# List consumer groups to verify our consumer group is being registered +/bin/kafka-consumer-groups --bootstrap-server broker:29092 --list Review Comment: These triple-quoted strings are standalone expression statements (not module docstrings), so they execute as no-ops at import time and are easy to miss/forget. Convert them to real comments, a proper module docstring, or constants referenced from documentation/tests. ```suggestion # Airflow Kafka connection # AIRFLOW_CONN_KAFKA_DEFAULT='{ # "conn_type": "general", # "extra": { # "bootstrap.servers": "broker:29092", # "group.id": "kafka_default_group", # "security.protocol": "PLAINTEXT", # "enable.auto.commit": false, # "auto.offset.reset": "latest" # } # }' # # Kafka commands to verify messages are being produced to the topic: # # Create Topic # /bin/kafka-topics --bootstrap-server broker:29092 --create --topic fizz_buzz # /bin/kafka-topics --bootstrap-server broker:29092 --create --topic dlq # # Get offsets for the topic to verify messages are being produced # /bin/kafka-get-offsets --bootstrap-server broker:29092 --topic fizz_buzz # /bin/kafka-get-offsets --bootstrap-server broker:29092 --topic dlq # # List consumer groups to verify our consumer group is being registered # /bin/kafka-consumer-groups --bootstrap-server broker:29092 --list # ``` ########## airflow-e2e-tests/tests/airflow_e2e_tests/dags/example_event_driven.py: ########## @@ -0,0 +1,178 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +import json +from datetime import timedelta +from typing import TYPE_CHECKING, cast + +import pendulum + +from airflow.providers.apache.kafka.operators.produce import ProduceToTopicOperator +from airflow.providers.common.messaging.triggers.msg_queue import MessageQueueTrigger +from airflow.sdk import Asset, AssetWatcher, dag, get_current_context, task + +if TYPE_CHECKING: + from airflow.sdk.execution_time.context import Context, TriggeringAssetEventsAccessor + +KAFKA_CONFIG_ID = "kafka_default" +TOPICS = ["fizz_buzz"] +DLQ_TOPIC = "dlq" +RETRY_COUNT = 3 +"""Airflow Kafka connection +AIRFLOW_CONN_KAFKA_DEFAULT='{ + "conn_type": "general", + "extra": { + "bootstrap.servers": "broker:29092", + "group.id": "kafka_default_group", + "security.protocol": "PLAINTEXT", + "enable.auto.commit": false, + "auto.offset.reset": "latest" + } +}' +""" +"""Kafka Command to verify messages are being produced to the topic: + +# Create Topic +/bin/kafka-topics --bootstrap-server broker:29092 --create --topic fizz_buzz +/bin/kafka-topics --bootstrap-server broker:29092 --create --topic dlq + + +# Get offsets for the topic to verify messages are being produced +/bin/kafka-get-offsets --bootstrap-server broker:29092 --topic fizz_buzz +/bin/kafka-get-offsets --bootstrap-server broker:29092 --topic dlq + +# List consumer groups to verify our consumer group is being registered +/bin/kafka-consumer-groups --bootstrap-server broker:29092 --list + +# Get current offsets for the consumer group to verify messages are being consumed +/bin/kafka-consumer-groups --bootstrap-server broker:29092 --describe --group group_1 +""" Review Comment: Same issue here: this is a standalone triple-quoted string (a no-op expression), not a docstring. Prefer comments or documentation in a README/rst, or fold this into the module docstring if it’s meant to be user-facing guidance. ```suggestion # Airflow Kafka connection: # AIRFLOW_CONN_KAFKA_DEFAULT='{ # "conn_type": "general", # "extra": { # "bootstrap.servers": "broker:29092", # "group.id": "kafka_default_group", # "security.protocol": "PLAINTEXT", # "enable.auto.commit": false, # "auto.offset.reset": "latest" # } # }' # # Kafka commands to verify messages are being produced to the topic: # # Create Topic # /bin/kafka-topics --bootstrap-server broker:29092 --create --topic fizz_buzz # /bin/kafka-topics --bootstrap-server broker:29092 --create --topic dlq # # Get offsets for the topic to verify messages are being produced # /bin/kafka-get-offsets --bootstrap-server broker:29092 --topic fizz_buzz # /bin/kafka-get-offsets --bootstrap-server broker:29092 --topic dlq # # List consumer groups to verify our consumer group is being registered # /bin/kafka-consumer-groups --bootstrap-server broker:29092 --list # # Get current offsets for the consumer group to verify messages are being consumed # /bin/kafka-consumer-groups --bootstrap-server broker:29092 --describe --group group_1 ``` ########## airflow-e2e-tests/tests/airflow_e2e_tests/dags/example_event_driven.py: ########## @@ -0,0 +1,178 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +import json +from datetime import timedelta +from typing import TYPE_CHECKING, cast + +import pendulum + +from airflow.providers.apache.kafka.operators.produce import ProduceToTopicOperator +from airflow.providers.common.messaging.triggers.msg_queue import MessageQueueTrigger +from airflow.sdk import Asset, AssetWatcher, dag, get_current_context, task + +if TYPE_CHECKING: + from airflow.sdk.execution_time.context import Context, TriggeringAssetEventsAccessor + +KAFKA_CONFIG_ID = "kafka_default" +TOPICS = ["fizz_buzz"] +DLQ_TOPIC = "dlq" +RETRY_COUNT = 3 +"""Airflow Kafka connection +AIRFLOW_CONN_KAFKA_DEFAULT='{ + "conn_type": "general", + "extra": { + "bootstrap.servers": "broker:29092", + "group.id": "kafka_default_group", + "security.protocol": "PLAINTEXT", + "enable.auto.commit": false, + "auto.offset.reset": "latest" + } +}' +""" +"""Kafka Command to verify messages are being produced to the topic: + +# Create Topic +/bin/kafka-topics --bootstrap-server broker:29092 --create --topic fizz_buzz +/bin/kafka-topics --bootstrap-server broker:29092 --create --topic dlq + + +# Get offsets for the topic to verify messages are being produced +/bin/kafka-get-offsets --bootstrap-server broker:29092 --topic fizz_buzz +/bin/kafka-get-offsets --bootstrap-server broker:29092 --topic dlq + +# List consumer groups to verify our consumer group is being registered +/bin/kafka-consumer-groups --bootstrap-server broker:29092 --list + +# Get current offsets for the consumer group to verify messages are being consumed +/bin/kafka-consumer-groups --bootstrap-server broker:29092 --describe --group group_1 +""" + +SAMPLE_ORDERS = [ + {"order_id": "ORD-1001", "customer": "alice", "item": "widget", "quantity": 2, "price": 9.99}, + {"order_id": "ORD-1002", "customer": "bob", "item": "gadget", "quantity": 1, "price": 24.50}, + {"order_id": "ORD-1003", "customer": "carol", "item": "widget", "quantity": 5, "price": 9.99}, + {"order_id": "ORD-1004", "customer": "dave", "item": "doohickey", "quantity": 3, "price": 14.75}, + {"order_id": "ORD-1005", "customer": "eve", "item": "thingamajig", "quantity": 1, "price": 39.00}, + {"order_id": "ORD-1006", "customer": "frank", "item": "widget", "quantity": 10, "price": 9.99}, + {"order_id": "ORD-1007", "customer": "grace", "item": "gadget", "quantity": 2, "price": 24.50}, + {"order_id": "ORD-1008", "customer": "heidi", "item": "doohickey", "quantity": 1, "price": 14.75}, +] + + +def producer_function(): + for order in SAMPLE_ORDERS: + yield (json.dumps(order["order_id"]), json.dumps(order)) + # produce a malformed message to demonstrate error handling + yield ("malformed_message", "malformed_message") + + +def process_one_message(message: str): + order = json.loads(message) + total = order["quantity"] * order["price"] + print(f"Order {order['order_id']}: {order['quantity']}x {order['item']} = ${total:.2f}") + return order + + +def handle_dlq(): + context: Context = get_current_context() + triggering_asset_events: TriggeringAssetEventsAccessor = context["triggering_asset_events"] + for event in triggering_asset_events[kafka_cdc_asset]: + print(f"Handling failed message from event: {event}") + value = json.dumps( + { + "asset": event.asset.model_dump(mode="json"), + "extra": event.extra, + } + ) + yield (json.dumps(event.asset.uri), value) + + +# Airflow 3 example +# Define a trigger that listens to an external message queue (Apache Kafka in this case) +trigger = MessageQueueTrigger( + scheme="kafka", + # the rest of the parameters are used by the trigger + kafka_config_id=KAFKA_CONFIG_ID, + topics=TOPICS, + poll_interval=1, + poll_timeout=1, + commit_offset=True, +) + +# Define an asset that watches for messages on the queue +kafka_cdc_asset = Asset("kafka_cdc_asset", watchers=[AssetWatcher(name="kafka_cdc", trigger=trigger)]) + + +@dag( + schedule=[kafka_cdc_asset], + tags=["event-driven"], +) +def event_driven_consumer(): + + @task(retries=RETRY_COUNT, retry_delay=timedelta(seconds=1)) + def process_message(**context) -> bool: + # Extract the triggering asset events from the context + triggering_asset_events: TriggeringAssetEventsAccessor = context["triggering_asset_events"] + for event in triggering_asset_events[kafka_cdc_asset]: + # Get the message from the TriggerEvent payload + print(f"Asset event: {event}") + process_one_message(cast("str", event.extra["payload"])) + return True Review Comment: `typing.cast` expects a type as its first argument, not a string. Using `cast("str", ...)` defeats static checking and is inconsistent with normal typing usage. Use `cast(str, ...)` or (preferably) validate/cast the payload more explicitly if it can be non-string. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
