From 5f3b97a36740d1714b3b11232a02217edd8284f1 Mon Sep 17 00:00:00 2001 From: Christos Bisias Date: Tue, 5 May 2026 09:08:56 +0300 Subject: [PATCH] add a listener under the kafka provider --- .../tests/integration/otel/test_otel.py | 200 +------ .../core_api/routes/public/test_plugins.py | 15 +- .../unit/plugins/test_plugins_manager.py | 2 +- .../test_utils/integration_setup.py | 201 +++++++ .../apache/kafka/docs/configurations-ref.rst | 106 ++++ providers/apache/kafka/docs/index.rst | 1 + providers/apache/kafka/provider.yaml | 126 ++++ providers/apache/kafka/pyproject.toml | 3 + .../apache/kafka/get_provider_info.py | 104 ++++ .../providers/apache/kafka/hooks/base.py | 16 + .../providers/apache/kafka/hooks/consume.py | 7 +- .../apache/kafka/plugins/__init__.py | 16 + .../apache/kafka/plugins/listener.py | 470 +++++++++++++++ .../apache/kafka/plugins/__init__.py | 16 + .../apache/kafka/plugins/dags/__init__.py | 16 + .../apache/kafka/plugins/dags/demo_dag.py | 45 ++ .../apache/kafka/plugins/test_listener.py | 167 ++++++ .../unit/apache/kafka/hooks/test_base.py | 29 + .../unit/apache/kafka/hooks/test_consume.py | 28 +- .../unit/apache/kafka/hooks/test_produce.py | 16 + .../unit/apache/kafka/plugins/__init__.py | 16 + .../apache/kafka/plugins/test_listener.py | 538 ++++++++++++++++++ .../ci/docker-compose/integration-kafka.yml | 4 + 23 files changed, 1948 insertions(+), 194 deletions(-) create mode 100644 devel-common/src/tests_common/test_utils/integration_setup.py create mode 100644 providers/apache/kafka/docs/configurations-ref.rst create mode 100644 providers/apache/kafka/src/airflow/providers/apache/kafka/plugins/__init__.py create mode 100644 providers/apache/kafka/src/airflow/providers/apache/kafka/plugins/listener.py create mode 100644 providers/apache/kafka/tests/integration/apache/kafka/plugins/__init__.py create mode 100644 providers/apache/kafka/tests/integration/apache/kafka/plugins/dags/__init__.py create mode 100644 providers/apache/kafka/tests/integration/apache/kafka/plugins/dags/demo_dag.py create mode 100644 providers/apache/kafka/tests/integration/apache/kafka/plugins/test_listener.py create mode 100644 providers/apache/kafka/tests/unit/apache/kafka/plugins/__init__.py create mode 100644 providers/apache/kafka/tests/unit/apache/kafka/plugins/test_listener.py diff --git a/airflow-core/tests/integration/otel/test_otel.py b/airflow-core/tests/integration/otel/test_otel.py index ec6a7a6c94c4d..39d7ee6bdab61 100644 --- a/airflow-core/tests/integration/otel/test_otel.py +++ b/airflow-core/tests/integration/otel/test_otel.py @@ -16,34 +16,35 @@ # under the License. from __future__ import annotations -import json import logging import os import socket import subprocess import time -from typing import Any +from typing import TYPE_CHECKING, Any import pytest import requests -from sqlalchemy import func, select +from sqlalchemy import select -from airflow._shared.timezones import timezone -from airflow.dag_processing.bundles.manager import DagBundlesManager -from airflow.dag_processing.dagbag import DagBag -from airflow.models import DagRun -from airflow.models.serialized_dag import SerializedDagModel from airflow.models.taskinstance import TaskInstance -from airflow.serialization.definitions.dag import SerializedDAG from airflow.utils.session import create_session from airflow.utils.state import State -from tests_common.test_utils.dag import create_scheduler_dag +from tests_common.test_utils.integration_setup import ( + serialize_and_get_dags, + start_scheduler, + terminate_process, + unpause_trigger_dag_and_get_run_id, + wait_for_dag_run, +) from tests_common.test_utils.otel_utils import ( dump_airflow_metadata_db, extract_metrics_from_output, ) -from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_1_PLUS + +if TYPE_CHECKING: + from airflow.serialization.definitions.dag import SerializedDAG log = logging.getLogger("integration.otel.test_otel") @@ -86,62 +87,6 @@ def wait_for_otel_collector(host: str, port: int, timeout: int = 120) -> None: ) -def unpause_trigger_dag_and_get_run_id(dag_id: str, conf: dict | None = None) -> str: - unpause_command = ["airflow", "dags", "unpause", dag_id] - - # Unpause the dag using the cli. - subprocess.run(unpause_command, check=True, env=os.environ.copy()) - - execution_date = timezone.utcnow() - run_id = f"manual__{execution_date.isoformat()}" - - trigger_command = [ - "airflow", - "dags", - "trigger", - dag_id, - "--run-id", - run_id, - "--logical-date", - execution_date.isoformat(), - ] - - if conf: - import json - - trigger_command += ["--conf", json.dumps(conf)] - - # Trigger the dag using the cli. - subprocess.run(trigger_command, check=True, env=os.environ.copy()) - - return run_id - - -def wait_for_dag_run(dag_id: str, run_id: str, max_wait_time: int): - # max_wait_time, is the timeout for the DAG run to complete. The value is in seconds. - start_time = timezone.utcnow().timestamp() - - while timezone.utcnow().timestamp() - start_time < max_wait_time: - with create_session() as session: - dag_run = session.scalar( - select(DagRun).where( - DagRun.dag_id == dag_id, - DagRun.run_id == run_id, - ) - ) - - if dag_run is None: - time.sleep(5) - continue - - dag_run_state = dag_run.state - log.debug("DAG Run state: %s.", dag_run_state) - - if dag_run_state in [State.SUCCESS, State.FAILED]: - break - return dag_run_state - - def print_ti_output_for_dag_run(dag_id: str, run_id: str): breeze_logs_dir = "/root/airflow/logs" @@ -189,19 +134,6 @@ class TestOtelIntegration: use_otel = os.getenv("use_otel", default="false") log_level = os.getenv("log_level", default="none") - scheduler_command_args = [ - "airflow", - "scheduler", - ] - - apiserver_command_args = [ - "airflow", - "api-server", - "--port", - "8080", - "--daemon", - ] - dags: dict[str, SerializedDAG] = {} @classmethod @@ -251,67 +183,7 @@ def setup_class(cls): migrate_command = ["airflow", "db", "migrate"] subprocess.run(migrate_command, check=True, env=os.environ.copy()) - cls.dags = cls.serialize_and_get_dags() - - @classmethod - def serialize_and_get_dags(cls) -> dict[str, SerializedDAG]: - log.info("Serializing Dags from directory %s", cls.dag_folder) - # Load DAGs from the dag directory. - dag_bag = DagBag(dag_folder=cls.dag_folder) - - dag_ids = dag_bag.dag_ids - assert len(dag_ids) == 1 - - dag_dict: dict[str, SerializedDAG] = {} - with create_session() as session: - for dag_id in dag_ids: - dag = dag_bag.get_dag(dag_id) - assert dag is not None, f"DAG with ID {dag_id} not found." - # Sync the DAG to the database. - if AIRFLOW_V_3_0_PLUS: - from airflow.models.dagbundle import DagBundleModel - - count = session.scalar( - select(func.count()) - .select_from(DagBundleModel) - .where(DagBundleModel.name == "testing") - ) - if count == 0: - session.add(DagBundleModel(name="testing")) - session.commit() - SerializedDAG.bulk_write_to_db( - bundle_name="testing", bundle_version=None, dags=[dag], session=session - ) - dag_dict[dag_id] = create_scheduler_dag(dag) - else: - dag.sync_to_db(session=session) - dag_dict[dag_id] = dag - # Manually serialize the dag and write it to the db to avoid a db error. - if AIRFLOW_V_3_1_PLUS: - from airflow.serialization.serialized_objects import LazyDeserializedDAG - - SerializedDagModel.write_dag( - LazyDeserializedDAG.from_dag(dag), bundle_name="testing", session=session - ) - else: - SerializedDagModel.write_dag(dag, bundle_name="testing", session=session) - - session.commit() - - TESTING_BUNDLE_CONFIG = [ - { - "name": "testing", - "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", - "kwargs": {"path": f"{cls.dag_folder}", "refresh_interval": 1}, - } - ] - - os.environ["AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST"] = json.dumps(TESTING_BUNDLE_CONFIG) - # Initial add - manager = DagBundlesManager() - manager.sync_bundles_to_db() - - return dag_dict + cls.dags = serialize_and_get_dags(dag_folder=cls.dag_folder) def dag_execution_for_testing_metrics(self, capfd): # Metrics. @@ -327,7 +199,7 @@ def dag_execution_for_testing_metrics(self, capfd): try: # Start the processes here and not as fixtures or in a common setup, # so that the test can capture their output. - scheduler_process, apiserver_process = self.start_scheduler(capture_output=True) + scheduler_process, apiserver_process = start_scheduler(capture_output=True) dag_id = "otel_test_dag" @@ -357,13 +229,13 @@ def dag_execution_for_testing_metrics(self, capfd): finally: # Terminate the processes. - self._terminate_process(scheduler_process) + terminate_process(scheduler_process) scheduler_status = scheduler_process.poll() assert scheduler_status is not None, ( "The scheduler_1 process status is None, which means that it hasn't terminated as expected." ) - self._terminate_process(apiserver_process) + terminate_process(apiserver_process) apiserver_status = apiserver_process.poll() assert apiserver_status is not None, ( "The apiserver process status is None, which means that it hasn't terminated as expected." @@ -491,7 +363,7 @@ def test_dag_execution_succeeds(self, capfd, task_span_detail_level, expected_hi try: # Start the processes here and not as fixtures or in a common setup, # so that the test can capture their output. - scheduler_process, apiserver_process = self.start_scheduler() + scheduler_process, apiserver_process = start_scheduler() dag_id = "otel_test_dag" @@ -521,13 +393,13 @@ def test_dag_execution_succeeds(self, capfd, task_span_detail_level, expected_hi dump_airflow_metadata_db(session) # Terminate the processes. - self._terminate_process(scheduler_process) + terminate_process(scheduler_process) scheduler_status = scheduler_process.poll() assert scheduler_status is not None, ( "The scheduler_1 process status is None, which means that it hasn't terminated as expected." ) - self._terminate_process(apiserver_process) + terminate_process(apiserver_process) apiserver_status = apiserver_process.poll() assert apiserver_status is not None, ( "The apiserver process status is None, which means that it hasn't terminated as expected." @@ -568,37 +440,3 @@ def get_parent_span_id(span): nested = get_span_hierarchy() assert nested == expected_hierarchy - - @staticmethod - def _terminate_process(proc: subprocess.Popen, timeout: int = 30) -> None: - # Grace period covers OTel atexit flush (force_flush default: 10s); - # SIGKILL is the fallback if the process is still alive after timeout. - proc.terminate() - try: - proc.wait(timeout=timeout) - except subprocess.TimeoutExpired: - proc.kill() - proc.wait() - - def start_scheduler(self, capture_output: bool = False): - stdout = None if capture_output else subprocess.DEVNULL - stderr = None if capture_output else subprocess.DEVNULL - - scheduler_process = subprocess.Popen( - self.scheduler_command_args, - env=os.environ.copy(), - stdout=stdout, - stderr=stderr, - ) - - apiserver_process = subprocess.Popen( - self.apiserver_command_args, - env=os.environ.copy(), - stdout=stdout, - stderr=stderr, - ) - - # Wait to ensure both processes have started. - time.sleep(10) - - return scheduler_process, apiserver_process diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_plugins.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_plugins.py index 6bca49d8052ea..e995cb4b13704 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_plugins.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_plugins.py @@ -34,7 +34,7 @@ class TestGetPlugins: # Filters ( {}, - 18, + 19, [ "InformaticaProviderPlugin", "MetadataCollectionPlugin", @@ -45,6 +45,7 @@ class TestGetPlugins: "edge_executor", "hitl_review", "hive", + "kafka_listener", "plugin-a", "plugin-b", "plugin-c", @@ -58,14 +59,14 @@ class TestGetPlugins: ), ( {"limit": 3, "offset": 3}, - 18, + 19, [ "business_day_window_plugin", "databricks_workflow", "decreasing_priority_weight_strategy_plugin", ], ), - ({"limit": 1}, 18, ["InformaticaProviderPlugin"]), + ({"limit": 1}, 19, ["InformaticaProviderPlugin"]), ], ) def test_should_respond_200( @@ -158,17 +159,17 @@ def test_invalid_external_view_destination_should_log_warning_and_continue(self, # Verify warning was logged assert any("Skipping invalid plugin due to error" in rec.message for rec in caplog.records) - response = test_client.get("/plugins", params={"limit": 7, "offset": 9}) + response = test_client.get("/plugins", params={"limit": 8, "offset": 9}) assert response.status_code == 200 body = response.json() plugins_page = body["plugins"] - # Even though limit=7, only 6 valid plugins should come back - assert len(plugins_page) == 7 + # Invalid plugins are filtered before pagination, so the page is full. + assert len(plugins_page) == 8 assert "test_plugin_invalid" not in [p["name"] for p in plugins_page] - assert body["total_entries"] == 18 + assert body["total_entries"] == 19 @skip_if_force_lowest_dependencies_marker diff --git a/airflow-core/tests/unit/plugins/test_plugins_manager.py b/airflow-core/tests/unit/plugins/test_plugins_manager.py index ce55db1e269ea..0dcac96fe0a32 100644 --- a/airflow-core/tests/unit/plugins/test_plugins_manager.py +++ b/airflow-core/tests/unit/plugins/test_plugins_manager.py @@ -408,7 +408,7 @@ def test_does_not_double_import_entrypoint_provider_plugins(self): # Mock/skip loading from plugin dir with mock.patch("airflow.plugins_manager._load_plugins_from_plugin_directory", return_value=([], [])): plugins = plugins_manager._get_plugins()[0] - assert len(plugins) == 6 + assert len(plugins) == 7 class TestWindowPluginRegistration: diff --git a/devel-common/src/tests_common/test_utils/integration_setup.py b/devel-common/src/tests_common/test_utils/integration_setup.py new file mode 100644 index 0000000000000..59476d8c4d06b --- /dev/null +++ b/devel-common/src/tests_common/test_utils/integration_setup.py @@ -0,0 +1,201 @@ +# 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 +import logging +import os +import subprocess +import time + +from sqlalchemy import func, select + +from airflow.dag_processing.bundles.manager import DagBundlesManager +from airflow.dag_processing.dagbag import DagBag +from airflow.models import DagRun +from airflow.models.serialized_dag import SerializedDagModel +from airflow.serialization.definitions.dag import SerializedDAG +from airflow.utils.session import create_session +from airflow.utils.state import State + +from tests_common.test_utils.dag import create_scheduler_dag +from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_1_PLUS + +try: + from airflow.sdk import timezone +except ImportError: + from airflow.utils import timezone # type: ignore[no-redef,attr-defined] + +log = logging.getLogger(__name__) + + +def unpause_trigger_dag_and_get_run_id(dag_id: str, conf: dict | None = None) -> str: + unpause_command = ["airflow", "dags", "unpause", dag_id] + + # Unpause the dag using the cli. + subprocess.run(unpause_command, check=True, env=os.environ.copy()) + + execution_date = timezone.utcnow() + run_id = f"manual__{execution_date.isoformat()}" + + trigger_command = [ + "airflow", + "dags", + "trigger", + dag_id, + "--run-id", + run_id, + "--logical-date", + execution_date.isoformat(), + ] + + if conf: + trigger_command += ["--conf", json.dumps(conf)] + + # Trigger the dag using the cli. + subprocess.run(trigger_command, check=True, env=os.environ.copy()) + + return run_id + + +def start_scheduler(capture_output: bool = False): + stdout = None if capture_output else subprocess.DEVNULL + stderr = None if capture_output else subprocess.DEVNULL + + scheduler_command_args = [ + "airflow", + "scheduler", + ] + + apiserver_command_args = [ + "airflow", + "api-server", + "--port", + "8080", + "--daemon", + ] + + scheduler_process = subprocess.Popen( + scheduler_command_args, + env=os.environ.copy(), + stdout=stdout, + stderr=stderr, + ) + + apiserver_process = subprocess.Popen( + apiserver_command_args, + env=os.environ.copy(), + stdout=stdout, + stderr=stderr, + ) + + # Wait to ensure both processes have started. + time.sleep(10) + + return scheduler_process, apiserver_process + + +def serialize_and_get_dags(dag_folder) -> dict[str, SerializedDAG]: + log.info("Serializing Dags from directory %s", dag_folder) + # Load DAGs from the dag directory. + dag_bag = DagBag(dag_folder=dag_folder) + + dag_ids = dag_bag.dag_ids + + dag_dict: dict[str, SerializedDAG] = {} + with create_session() as session: + for dag_id in dag_ids: + dag = dag_bag.get_dag(dag_id) + assert dag is not None, f"DAG with ID {dag_id} not found." + # Sync the DAG to the database. + if AIRFLOW_V_3_0_PLUS: + from airflow.models.dagbundle import DagBundleModel + + count = session.scalar( + select(func.count()).select_from(DagBundleModel).where(DagBundleModel.name == "testing") + ) + if count == 0: + session.add(DagBundleModel(name="testing")) + session.commit() + SerializedDAG.bulk_write_to_db( + bundle_name="testing", bundle_version=None, dags=[dag], session=session + ) + dag_dict[dag_id] = create_scheduler_dag(dag) + else: + dag.sync_to_db(session=session) + dag_dict[dag_id] = dag + # Manually serialize the dag and write it to the db to avoid a db error. + if AIRFLOW_V_3_1_PLUS: + from airflow.serialization.serialized_objects import LazyDeserializedDAG + + SerializedDagModel.write_dag( + LazyDeserializedDAG.from_dag(dag), bundle_name="testing", session=session + ) + else: + SerializedDagModel.write_dag(dag, bundle_name="testing", session=session) + + session.commit() + + TESTING_BUNDLE_CONFIG = [ + { + "name": "testing", + "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", + "kwargs": {"path": f"{dag_folder}", "refresh_interval": 1}, + } + ] + + os.environ["AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST"] = json.dumps(TESTING_BUNDLE_CONFIG) + # Initial add + manager = DagBundlesManager() + manager.sync_bundles_to_db() + + return dag_dict + + +def wait_for_dag_run(dag_id: str, run_id: str, max_wait_time: int): + # max_wait_time, is the timeout for the DAG run to complete. The value is in seconds. + start_time = timezone.utcnow().timestamp() + + while timezone.utcnow().timestamp() - start_time < max_wait_time: + with create_session() as session: + dag_run = session.scalar( + select(DagRun).where( + DagRun.dag_id == dag_id, + DagRun.run_id == run_id, + ) + ) + + if dag_run is None: + time.sleep(5) + continue + + dag_run_state = dag_run.state + log.debug("DAG Run state: %s.", dag_run_state) + + if dag_run_state in [State.SUCCESS, State.FAILED]: + break + return dag_run_state + + +def terminate_process(proc: subprocess.Popen, timeout: int = 30) -> None: + # SIGKILL is the fallback if the process is still alive after timeout. + proc.terminate() + try: + proc.wait(timeout=timeout) + except subprocess.TimeoutExpired: + proc.kill() + proc.wait() diff --git a/providers/apache/kafka/docs/configurations-ref.rst b/providers/apache/kafka/docs/configurations-ref.rst new file mode 100644 index 0000000000000..80d0d373f8f2a --- /dev/null +++ b/providers/apache/kafka/docs/configurations-ref.rst @@ -0,0 +1,106 @@ + .. 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. + +.. _configuration:kafka: + +.. include:: /../../../../devel-common/src/sphinx_exts/includes/providers-configurations-ref.rst +.. include:: /../../../../devel-common/src/sphinx_exts/includes/sections-and-options.rst + + +Highlighted configurations +=========================== + +The ``[kafka_listener]`` section configures the ``KafkaListenerPlugin``, +which publishes Airflow DagRun and TaskInstance state-change events to a +Kafka topic. DagRun and TaskInstance events are separated and enabled by +distinct flags. Both event-type flags default to ``False``. + +.. _configuration_kafka_listener_activation:kafka: + +Activating the listener +----------------------- + +To enable event publishing you need to + + * enable at least one event-type flag + * point the listener at an Airflow Kafka connection via ``kafka_config_id`` + (defaults to ``kafka_default``) that carries the broker address and any + other confluent-kafka client options on its extras + * have a pre-existing kafka topic + +.. code-block:: ini + + [kafka_listener] + dag_run_events_enabled = True + task_instance_events_enabled = True + kafka_config_id = kafka_events + topic = airflow.events + +The connection's ``extra`` JSON accepts the full confluent-kafka client +configuration — including SASL/TLS options and callbacks (e.g. ``error_cb``, +``oauth_cb``) given as dotted-path strings, which are resolved to callables +before the producer is built. + +.. code-block:: json + + { + "bootstrap.servers": "broker:9092", + "security.protocol": "SASL_SSL", + "sasl.mechanisms": "OAUTHBEARER", + "oauth_cb": "my_company.auth.oauth_cb" + } + +Environment-variable equivalents: + +.. code-block:: ini + + AIRFLOW__KAFKA_LISTENER__DAG_RUN_EVENTS_ENABLED=True + AIRFLOW__KAFKA_LISTENER__TASK_INSTANCE_EVENTS_ENABLED=True + AIRFLOW__KAFKA_LISTENER__KAFKA_CONFIG_ID=kafka_events + AIRFLOW__KAFKA_LISTENER__TOPIC=airflow.events + +The two event flags are independent, users can opt-in to get only DagRun +event messages or only TaskInstance event messages or both. + +The topic must already exist on the broker, it's not auto-created. On a missing +topic, broker connection failure, or any other producer init error, the listener +doesn't fail, instead it logs a warning and retries the init after ``topic_check_retry_interval`` +seconds (default ``60``). Once the topic is created on the broker the listener will pick it up. + +.. _configuration_kafka_listener_filtering:kafka: + +Filtering events +---------------- + +DagRun and TaskInstance events are filtered separately. Each filter is a +comma-separated list of ``fnmatch`` glob patterns; an empty list means +"allow all", and **deny takes precedence over allow**. + +.. code-block:: ini + + [kafka_listener] + dag_run_dag_id_allowlist = sales_*,marketing_* + dag_run_dag_id_denylist = sales_internal_* + + task_instance_dag_id_allowlist = sales_* + task_instance_dag_id_denylist = + task_instance_task_id_allowlist = load_*,extract_* + task_instance_task_id_denylist = *_cleanup + +TaskInstance events must pass *both* the dag-id and task-id filters. +Mapped task instances share their parent's ``task_id``, so a single +``task_id`` pattern covers every map index. diff --git a/providers/apache/kafka/docs/index.rst b/providers/apache/kafka/docs/index.rst index cee81d3a021a1..d9db2b715de67 100644 --- a/providers/apache/kafka/docs/index.rst +++ b/providers/apache/kafka/docs/index.rst @@ -47,6 +47,7 @@ :maxdepth: 1 :caption: References + Configuration Python API <_api/airflow/providers/apache/kafka/index> .. toctree:: diff --git a/providers/apache/kafka/provider.yaml b/providers/apache/kafka/provider.yaml index f43f76ce0b876..8e23df573095b 100644 --- a/providers/apache/kafka/provider.yaml +++ b/providers/apache/kafka/provider.yaml @@ -129,3 +129,129 @@ connection-types: queues: - airflow.providers.apache.kafka.queues.kafka.KafkaMessageQueueProvider + +plugins: + - name: kafka_listener + plugin-class: airflow.providers.apache.kafka.plugins.listener.KafkaListenerPlugin + +config: + kafka_listener: + description: | + Settings for the Kafka listener that publishes Airflow DagRun and + TaskInstance state-change events to a Kafka topic. + options: + dag_run_events_enabled: + description: | + Publish DagRun state-change events (``dag_run.running``, + ``dag_run.success``, ``dag_run.failed``). When False the + DagRun listener is not registered. + version_added: 1.14.1 + type: boolean + example: ~ + default: "False" + task_instance_events_enabled: + description: | + Publish TaskInstance state-change events (``task_instance.running``, + ``task_instance.success``, ``task_instance.failed``, + ``task_instance.skipped``). When False the TaskInstance listener + is not registered. + version_added: 1.14.1 + type: boolean + example: ~ + default: "False" + kafka_config_id: + description: | + Airflow connection used to build the listener's Kafka producer. + When unset, the producer hook falls back to its default + connection (``kafka_default``). + version_added: 1.14.1 + type: string + example: "kafka_default" + default: "" + topic: + description: | + Topic the listener publishes events to. The topic must already + exist on the broker; the listener will not auto-create it. + version_added: 1.14.1 + type: string + example: ~ + default: "airflow.events" + source: + description: | + Identifier added to every emitted message under the ``source`` + field so consumers can distinguish Airflow installations that + share the same topic. When unset, falls back to the hostname + of the Airflow component that emits the event (scheduler, + worker, etc.). + version_added: 1.14.1 + type: string + example: "af-prod-eu" + default: "" + dag_run_dag_id_allowlist: + description: | + Comma-separated glob patterns. When set, DagRun events are only + emitted for dag_ids matching at least one pattern. Empty = all dags. + version_added: 1.14.1 + type: string + example: "demo_*,test_dag1" + default: "" + dag_run_dag_id_denylist: + description: | + Comma-separated glob patterns. DagRun events for dag_ids matching + any pattern are skipped. Deny takes precedence over allow. + version_added: 1.14.1 + type: string + example: "demo_*" + default: "" + task_instance_dag_id_allowlist: + description: | + Comma-separated glob patterns. When set, TaskInstance events are + only emitted for dag_ids matching at least one pattern. + Empty = all dags. + version_added: 1.14.1 + type: string + example: "demo_*,test_dag1" + default: "" + task_instance_dag_id_denylist: + description: | + Comma-separated glob patterns. TaskInstance events for dag_ids + matching any pattern are skipped. Deny takes precedence over allow. + version_added: 1.14.1 + type: string + example: "demo_*" + default: "" + task_instance_task_id_allowlist: + description: | + Comma-separated glob patterns. When set, TaskInstance events are + only emitted for task_ids matching at least one pattern. Applied + in addition to ``task_instance_dag_id_allowlist`` — both must pass. + Mapped task instances share the same task_id so a single pattern + covers all map indices. + version_added: 1.14.1 + type: string + example: "load_*,extract_*" + default: "" + task_instance_task_id_denylist: + description: | + Comma-separated glob patterns. TaskInstance events for task_ids + matching any pattern are skipped. Deny takes precedence over allow. + version_added: 1.14.1 + type: string + example: "*_cleanup" + default: "" + topic_check_timeout: + description: | + How long (in seconds) each topic existence check is allowed + to block waiting for a response from the broker. + version_added: 1.14.1 + type: integer + example: ~ + default: "10" + topic_check_retry_interval: + description: | + How long (in seconds) to wait before retrying a topic check, + in case it failed. + version_added: 1.14.1 + type: integer + example: ~ + default: "60" diff --git a/providers/apache/kafka/pyproject.toml b/providers/apache/kafka/pyproject.toml index b43f5a60790fa..3c5ead9a4b80c 100644 --- a/providers/apache/kafka/pyproject.toml +++ b/providers/apache/kafka/pyproject.toml @@ -130,6 +130,9 @@ apache-airflow-providers-standard = {workspace = true} [project.entry-points."apache_airflow_provider"] provider_info = "airflow.providers.apache.kafka.get_provider_info:get_provider_info" +[project.entry-points."airflow.plugins"] +kafka_listener = "airflow.providers.apache.kafka.plugins.listener:KafkaListenerPlugin" + [tool.flit.module] name = "airflow.providers.apache.kafka" diff --git a/providers/apache/kafka/src/airflow/providers/apache/kafka/get_provider_info.py b/providers/apache/kafka/src/airflow/providers/apache/kafka/get_provider_info.py index 1a68a7b916744..d8c8510cb5ddb 100644 --- a/providers/apache/kafka/src/airflow/providers/apache/kafka/get_provider_info.py +++ b/providers/apache/kafka/src/airflow/providers/apache/kafka/get_provider_info.py @@ -101,4 +101,108 @@ def get_provider_info(): } ], "queues": ["airflow.providers.apache.kafka.queues.kafka.KafkaMessageQueueProvider"], + "plugins": [ + { + "name": "kafka_listener", + "plugin-class": "airflow.providers.apache.kafka.plugins.listener.KafkaListenerPlugin", + } + ], + "config": { + "kafka_listener": { + "description": "Settings for the Kafka listener that publishes Airflow DagRun and\nTaskInstance state-change events to a Kafka topic.\n", + "options": { + "dag_run_events_enabled": { + "description": "Publish DagRun state-change events (``dag_run.running``,\n``dag_run.success``, ``dag_run.failed``). When False the\nDagRun listener is not registered.\n", + "version_added": "1.14.1", + "type": "boolean", + "example": None, + "default": "False", + }, + "task_instance_events_enabled": { + "description": "Publish TaskInstance state-change events (``task_instance.running``,\n``task_instance.success``, ``task_instance.failed``,\n``task_instance.skipped``). When False the TaskInstance listener\nis not registered.\n", + "version_added": "1.14.1", + "type": "boolean", + "example": None, + "default": "False", + }, + "kafka_config_id": { + "description": "Airflow connection used to build the listener's Kafka producer.\nWhen unset, the producer hook falls back to its default\nconnection (``kafka_default``).\n", + "version_added": "1.14.1", + "type": "string", + "example": "kafka_default", + "default": "", + }, + "topic": { + "description": "Topic the listener publishes events to. The topic must already\nexist on the broker; the listener will not auto-create it.\n", + "version_added": "1.14.1", + "type": "string", + "example": None, + "default": "airflow.events", + }, + "source": { + "description": "Identifier added to every emitted message under the ``source``\nfield so consumers can distinguish Airflow installations that\nshare the same topic. When unset, falls back to the hostname\nof the Airflow component that emits the event (scheduler,\nworker, etc.).\n", + "version_added": "1.14.1", + "type": "string", + "example": "af-prod-eu", + "default": "", + }, + "dag_run_dag_id_allowlist": { + "description": "Comma-separated glob patterns. When set, DagRun events are only\nemitted for dag_ids matching at least one pattern. Empty = all dags.\n", + "version_added": "1.14.1", + "type": "string", + "example": "demo_*,test_dag1", + "default": "", + }, + "dag_run_dag_id_denylist": { + "description": "Comma-separated glob patterns. DagRun events for dag_ids matching\nany pattern are skipped. Deny takes precedence over allow.\n", + "version_added": "1.14.1", + "type": "string", + "example": "demo_*", + "default": "", + }, + "task_instance_dag_id_allowlist": { + "description": "Comma-separated glob patterns. When set, TaskInstance events are\nonly emitted for dag_ids matching at least one pattern.\nEmpty = all dags.\n", + "version_added": "1.14.1", + "type": "string", + "example": "demo_*,test_dag1", + "default": "", + }, + "task_instance_dag_id_denylist": { + "description": "Comma-separated glob patterns. TaskInstance events for dag_ids\nmatching any pattern are skipped. Deny takes precedence over allow.\n", + "version_added": "1.14.1", + "type": "string", + "example": "demo_*", + "default": "", + }, + "task_instance_task_id_allowlist": { + "description": "Comma-separated glob patterns. When set, TaskInstance events are\nonly emitted for task_ids matching at least one pattern. Applied\nin addition to ``task_instance_dag_id_allowlist`` — both must pass.\nMapped task instances share the same task_id so a single pattern\ncovers all map indices.\n", + "version_added": "1.14.1", + "type": "string", + "example": "load_*,extract_*", + "default": "", + }, + "task_instance_task_id_denylist": { + "description": "Comma-separated glob patterns. TaskInstance events for task_ids\nmatching any pattern are skipped. Deny takes precedence over allow.\n", + "version_added": "1.14.1", + "type": "string", + "example": "*_cleanup", + "default": "", + }, + "topic_check_timeout": { + "description": "How long (in seconds) each topic existence check is allowed\nto block waiting for a response from the broker.\n", + "version_added": "1.14.1", + "type": "integer", + "example": None, + "default": "10", + }, + "topic_check_retry_interval": { + "description": "How long (in seconds) to wait before retrying a topic check,\nin case it failed.\n", + "version_added": "1.14.1", + "type": "integer", + "example": None, + "default": "60", + }, + }, + } + }, } diff --git a/providers/apache/kafka/src/airflow/providers/apache/kafka/hooks/base.py b/providers/apache/kafka/src/airflow/providers/apache/kafka/hooks/base.py index b3b76ff1421d5..32ceef972b83d 100644 --- a/providers/apache/kafka/src/airflow/providers/apache/kafka/hooks/base.py +++ b/providers/apache/kafka/src/airflow/providers/apache/kafka/hooks/base.py @@ -22,8 +22,13 @@ from confluent_kafka.admin import AdminClient +from airflow.providers.common.compat.module_loading import import_string from airflow.providers.common.compat.sdk import BaseHook +# librdkafka config options whose values are callables. They can be provided as dotted-path +# strings on the connection extra and are resolved to callables before the client is built. +CALLBACK_CONFIG_KEYS = ("error_cb", "throttle_cb", "stats_cb", "log_cb", "oauth_cb", "on_commit") + # Amazon MSK bootstrap servers follow a predictable naming scheme, e.g. # b-1.demo.abcde1.c2.kafka.us-east-1.amazonaws.com:9098 (provisioned) # boot-abcde1.c2.kafka-serverless.us-east-1.amazonaws.com:9098 (serverless) @@ -85,10 +90,18 @@ def get_ui_field_behaviour(cls) -> dict[str, Any]: def _get_client(self, config) -> Any: return AdminClient(config) + def _resolve_callbacks(self, config: dict[str, Any]) -> None: + """Resolve callback options provided as dotted-path strings into callables.""" + for key in CALLBACK_CONFIG_KEYS: + value = config.get(key) + if isinstance(value, str): + config[key] = import_string(value) + @cached_property def get_conn(self) -> Any: """Get the configuration object.""" config = self.get_connection(self.kafka_config_id).extra_dejson + self._resolve_callbacks(config) if not (config.get("bootstrap.servers", None)): raise ValueError("config['bootstrap.servers'] must be provided.") @@ -160,6 +173,9 @@ def test_connection(self) -> tuple[bool, str]: """Test Connectivity from the UI.""" try: config = self.get_connection(self.kafka_config_id).extra_dejson + # Resolve callbacks so that configured dotted-path strings + # (e.g. oauth_cb) are exercised by the test as well. + self._resolve_callbacks(config) t = AdminClient(config).list_topics(timeout=10) if t: return True, "Connection successful." diff --git a/providers/apache/kafka/src/airflow/providers/apache/kafka/hooks/consume.py b/providers/apache/kafka/src/airflow/providers/apache/kafka/hooks/consume.py index fa2f58449273f..a420378be1e31 100644 --- a/providers/apache/kafka/src/airflow/providers/apache/kafka/hooks/consume.py +++ b/providers/apache/kafka/src/airflow/providers/apache/kafka/hooks/consume.py @@ -21,7 +21,6 @@ from confluent_kafka import Consumer, KafkaError from airflow.providers.apache.kafka.hooks.base import KafkaBaseHook -from airflow.providers.common.compat.module_loading import import_string class KafkaAuthenticationError(Exception): @@ -51,10 +50,10 @@ def __init__(self, topics: Sequence[str], kafka_config_id=KafkaBaseHook.default_ def _get_client(self, config) -> Consumer: config_shallow = config.copy() - if config.get("error_cb") is None: + # KafkaBaseHook resolves user-provided callbacks. Set a default if + # none was provided for ``error_cb``. + if config_shallow.get("error_cb") is None: config_shallow["error_cb"] = error_callback - else: - config_shallow["error_cb"] = import_string(config["error_cb"]) return Consumer(config_shallow) def get_consumer(self) -> Consumer: diff --git a/providers/apache/kafka/src/airflow/providers/apache/kafka/plugins/__init__.py b/providers/apache/kafka/src/airflow/providers/apache/kafka/plugins/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/providers/apache/kafka/src/airflow/providers/apache/kafka/plugins/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/providers/apache/kafka/src/airflow/providers/apache/kafka/plugins/listener.py b/providers/apache/kafka/src/airflow/providers/apache/kafka/plugins/listener.py new file mode 100644 index 0000000000000..342b37139a094 --- /dev/null +++ b/providers/apache/kafka/src/airflow/providers/apache/kafka/plugins/listener.py @@ -0,0 +1,470 @@ +# 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 atexit +import json +import logging +import os +import time +from datetime import datetime, timezone +from fnmatch import fnmatch +from functools import lru_cache +from typing import TYPE_CHECKING, Any + +from airflow.providers.apache.kafka.version_compat import AIRFLOW_V_3_0_PLUS +from airflow.providers.common.compat.sdk import AirflowPlugin, conf, hookimpl +from airflow.utils.net import get_hostname + +if TYPE_CHECKING: + from confluent_kafka import Producer + from sqlalchemy.orm import Session + + from airflow.models import DagRun, TaskInstance + from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance + from airflow.utils.state import TaskInstanceState + +log = logging.getLogger(__name__) + +CONFIG_SECTION = "kafka_listener" +SCHEMA_VERSION = 1 + + +@lru_cache(maxsize=1) +def _dag_run_events_enabled() -> bool: + return conf.getboolean(CONFIG_SECTION, "dag_run_events_enabled", fallback="False") + + +@lru_cache(maxsize=1) +def _task_instance_events_enabled() -> bool: + return conf.getboolean(CONFIG_SECTION, "task_instance_events_enabled", fallback="False") + + +@lru_cache(maxsize=1) +def _get_kafka_config_id() -> str: + return conf.get(CONFIG_SECTION, "kafka_config_id", fallback="").strip() + + +@lru_cache(maxsize=1) +def _get_topic() -> str: + return conf.get(CONFIG_SECTION, "topic", fallback="airflow.events") + + +@lru_cache(maxsize=1) +def _get_source() -> str: + source_from_conf = conf.get(CONFIG_SECTION, "source", fallback="") + # Default to the current hostname, if un-set. + return source_from_conf or get_hostname() + + +def _parse_filter_patterns_to_tuple(raw: str) -> tuple[str, ...]: + return tuple(p.strip() for p in raw.split(",") if p.strip()) + + +@lru_cache(maxsize=1) +def _get_dag_run_dag_id_allowlist() -> tuple[str, ...]: + return _parse_filter_patterns_to_tuple(conf.get(CONFIG_SECTION, "dag_run_dag_id_allowlist", fallback="")) + + +@lru_cache(maxsize=1) +def _get_dag_run_dag_id_denylist() -> tuple[str, ...]: + return _parse_filter_patterns_to_tuple(conf.get(CONFIG_SECTION, "dag_run_dag_id_denylist", fallback="")) + + +@lru_cache(maxsize=1) +def _get_task_instance_dag_id_allowlist() -> tuple[str, ...]: + return _parse_filter_patterns_to_tuple( + conf.get(CONFIG_SECTION, "task_instance_dag_id_allowlist", fallback="") + ) + + +@lru_cache(maxsize=1) +def _get_task_instance_dag_id_denylist() -> tuple[str, ...]: + return _parse_filter_patterns_to_tuple( + conf.get(CONFIG_SECTION, "task_instance_dag_id_denylist", fallback="") + ) + + +@lru_cache(maxsize=1) +def _get_task_instance_task_id_allowlist() -> tuple[str, ...]: + return _parse_filter_patterns_to_tuple( + conf.get(CONFIG_SECTION, "task_instance_task_id_allowlist", fallback="") + ) + + +@lru_cache(maxsize=1) +def _get_task_instance_task_id_denylist() -> tuple[str, ...]: + return _parse_filter_patterns_to_tuple( + conf.get(CONFIG_SECTION, "task_instance_task_id_denylist", fallback="") + ) + + +@lru_cache(maxsize=1) +def _get_topic_check_timeout() -> int: + return conf.getint(CONFIG_SECTION, "topic_check_timeout", fallback="10") + + +@lru_cache(maxsize=1) +def _get_topic_check_retry_interval() -> int: + return conf.getint(CONFIG_SECTION, "topic_check_retry_interval", fallback="60") + + +def _id_is_allowed(id_to_check: str, allowlist: tuple[str, ...], denylist: tuple[str, ...]) -> bool: + """Deny takes precedence; empty allowlist means 'allow all'.""" + if denylist and any(fnmatch(id_to_check, id_pattern) for id_pattern in denylist): + return False + if allowlist and not any(fnmatch(id_to_check, id_pattern) for id_pattern in allowlist): + return False + return True + + +def _dag_run_event_allowed(dag_id: str) -> bool: + return _id_is_allowed( + dag_id, + _get_dag_run_dag_id_allowlist(), + _get_dag_run_dag_id_denylist(), + ) + + +def _task_instance_event_allowed(dag_id: str, task_id: str) -> bool: + return _id_is_allowed( + dag_id, + _get_task_instance_dag_id_allowlist(), + _get_task_instance_dag_id_denylist(), + ) and _id_is_allowed( + task_id, + _get_task_instance_task_id_allowlist(), + _get_task_instance_task_id_denylist(), + ) + + +# Listener-owned producer and topic state for the lifetime of the process. +# The producer is built once via KafkaProducerHook and kept for the process lifetime; +# the topic flags track whether the topic exists on the broker and whether we're +# currently in a back-off window after a failed topic check. +_producer: Producer | None = None +_topic_exists: bool = False +_topic_check_retry_after: float = 0.0 + + +def _reset_state_after_fork() -> None: + """ + Reset the producer and the topic state in forked children processes. + + confluent_kafka.Producer is not fork-safe — its background thread and broker sockets + do not survive ``os.fork``. The child re-initializes its own producer on first use + and re-verifies the topic. + """ + global _producer, _topic_exists, _topic_check_retry_after + _producer = None + _topic_exists = False + _topic_check_retry_after = 0.0 + + +os.register_at_fork(after_in_child=_reset_state_after_fork) + + +def _get_producer() -> Producer | None: + """Build (once) and return the listener's producer, or ``None`` on init failure.""" + global _producer + if _producer is not None: + return _producer + + # Deferred import: KafkaProducerHook pulls in confluent_kafka.Producer, which is heavy + # and shouldn't be loaded by every Airflow process at plugin import time. + from airflow.providers.apache.kafka.hooks.produce import KafkaProducerHook + + try: + hook_kwargs: dict[str, Any] = {} + if _get_kafka_config_id(): + hook_kwargs["kafka_config_id"] = _get_kafka_config_id() + producer = KafkaProducerHook(**hook_kwargs).get_producer() + except Exception as exc: + log.warning("Kafka listener: failed to initialize producer (%s).", exc) + return None + + atexit.register(_flush_producer_at_exit, producer) + _producer = producer + return _producer + + +def _check_topic_exists() -> bool: + """ + Verify the configured topic exists on the broker, with a retry-after-failure cooldown. + + Once the topic has been confirmed, the result is kept for the process lifetime; until + then, failed checks are retried on a configured interval. + """ + global _topic_exists, _topic_check_retry_after + if _topic_exists: + return True + if time.monotonic() < _topic_check_retry_after: + return False + + producer = _get_producer() + if producer is None: + _topic_check_retry_after = time.monotonic() + _get_topic_check_retry_interval() + return False + + try: + topics = producer.list_topics(timeout=_get_topic_check_timeout()).topics + except Exception as exc: + log.warning( + "Kafka listener: topic check failed (%s). Will retry after %ds.", + exc, + _get_topic_check_retry_interval(), + ) + _topic_check_retry_after = time.monotonic() + _get_topic_check_retry_interval() + return False + + if _get_topic() not in topics: + log.warning( + "Kafka listener: topic %r not found on the broker. Will retry after %ds. " + "Create the topic on the broker to enable publishing.", + _get_topic(), + _get_topic_check_retry_interval(), + ) + _topic_check_retry_after = time.monotonic() + _get_topic_check_retry_interval() + return False + + log.info( + "Kafka listener attached: pid=%s source=%r topic=%r", + os.getpid(), + _get_source(), + _get_topic(), + ) + _topic_exists = True + return True + + +def _flush_producer_at_exit(producer: Producer) -> None: + try: + producer.flush(5) + except Exception: + log.debug("Kafka listener: error flushing producer on exit", exc_info=True) + + +def _on_delivery(err, _msg) -> None: + if err is None: + return + log.warning("Kafka listener: delivery failed: %s", err) + # If the broker or local metadata says the topic is gone, the confirmation cached at + # attach-time is stale. Flip it back so the next _check_topic_exists re-verifies with + # the broker, gated by the same cooldown used elsewhere. Lets the listener auto-recover + # if the topic gets recreated instead of requiring a component restart. + from confluent_kafka import KafkaError + + if err.code() in (KafkaError.UNKNOWN_TOPIC_OR_PART, KafkaError._UNKNOWN_TOPIC): + global _topic_exists, _topic_check_retry_after + _topic_exists = False + _topic_check_retry_after = time.monotonic() + _get_topic_check_retry_interval() + + +def _now_iso() -> str: + return datetime.now(timezone.utc).isoformat() + + +# DagRun +def _produce_dr_message(event: str, dag_run: DagRun, msg: str) -> None: + dag_id = dag_run.dag_id + if not _dag_run_event_allowed(dag_id): + return + try: + _produce_message( + event, + dag_id, + dag_run.run_id, + _get_dr_payload(dag_run, msg), + ) + except Exception: + log.exception("Kafka listener: %s failed", event) + + +def _get_dr_payload(dag_run, msg) -> dict[str, Any]: + return { + "dag_id": getattr(dag_run, "dag_id", None), + "run_id": getattr(dag_run, "run_id", None), + "run_type": str(getattr(dag_run, "run_type", "") or ""), + "logical_date": str(getattr(dag_run, "logical_date", "") or ""), + "start_date": str(getattr(dag_run, "start_date", "") or ""), + "end_date": str(getattr(dag_run, "end_date", "") or ""), + "msg": msg or "", + } + + +# TaskInstance +def _produce_ti_message( + event: str, + previous_state: TaskInstanceState, + task_instance: RuntimeTaskInstance | TaskInstance, + error=None, +) -> None: + dag_id = task_instance.dag_id + task_id = task_instance.task_id + if not _task_instance_event_allowed(dag_id, task_id): + return + try: + _produce_message( + event, + dag_id, + task_instance.run_id, + _get_ti_payload(task_instance, previous_state, error=error), + ) + except Exception: + log.exception("Kafka listener: %s failed", event) + + +def _get_ti_payload(ti, previous_state, error=None) -> dict[str, Any]: + payload: dict[str, Any] = { + "dag_id": getattr(ti, "dag_id", None), + "task_id": getattr(ti, "task_id", None), + "run_id": getattr(ti, "run_id", None), + "try_number": getattr(ti, "try_number", None), + "map_index": getattr(ti, "map_index", -1), + "previous_state": str(previous_state) if previous_state is not None else None, + } + if error is not None: + payload["error"] = str(error) + return payload + + +def _produce_message(event: str, dag_id: str, run_id: str, payload: dict[str, Any]) -> None: + if not _dag_run_events_enabled() and not _task_instance_events_enabled(): + return + if not _check_topic_exists(): + return + producer = _get_producer() + if producer is None: + return + + # When changing any field in the body or the payload, + # also update the SCHEMA_VERSION to avoid breaking someone's workflow. + body = { + "schema_version": SCHEMA_VERSION, + "source": _get_source(), + "event": event, + "timestamp": _now_iso(), + **payload, + } + # Key groups all events from a single DagRun (DagRun + TaskInstance) into the same + # Kafka partition. Within the partition, the order of the events is preserved, + # e.g. dag_run.running > ti.running > ti.success > dag_run.success. + key = f"{dag_id}/{run_id}".encode() + try: + producer.produce( + _get_topic(), + key=key, + value=json.dumps(body, default=str).encode("utf-8"), + on_delivery=_on_delivery, + ) + producer.poll(0) + except Exception as ex: + log.warning("Kafka listener: failed to enqueue %s for %s/%s: %s", event, dag_id, run_id, ex) + + +class DagRunListener: + """Publishes DagRun state-change event messages to Kafka.""" + + @hookimpl + def on_dag_run_running(self, dag_run: DagRun, msg: str): + _produce_dr_message("dag_run.running", dag_run, msg) + + @hookimpl + def on_dag_run_success(self, dag_run: DagRun, msg: str): + _produce_dr_message("dag_run.success", dag_run, msg) + + @hookimpl + def on_dag_run_failed(self, dag_run: DagRun, msg: str): + _produce_dr_message("dag_run.failed", dag_run, msg) + + +class TaskListener: + """Publishes TaskInstance state-change event messages to Kafka.""" + + if AIRFLOW_V_3_0_PLUS: + + @hookimpl + def on_task_instance_running( + self, previous_state: TaskInstanceState, task_instance: RuntimeTaskInstance + ): + _produce_ti_message("task_instance.running", previous_state, task_instance) + + @hookimpl + def on_task_instance_success( + self, previous_state: TaskInstanceState, task_instance: RuntimeTaskInstance + ): + _produce_ti_message("task_instance.success", previous_state, task_instance) + + @hookimpl + def on_task_instance_failed( + self, + previous_state: TaskInstanceState, + task_instance: RuntimeTaskInstance, + error: None | str | BaseException, + ): + _produce_ti_message("task_instance.failed", previous_state, task_instance, error=error) + + @hookimpl + def on_task_instance_skipped( + self, previous_state: TaskInstanceState, task_instance: RuntimeTaskInstance + ): + _produce_ti_message("task_instance.skipped", previous_state, task_instance) + else: + + @hookimpl + def on_task_instance_running( # type: ignore[misc] + self, previous_state: TaskInstanceState, task_instance: TaskInstance, session: Session + ): + _produce_ti_message("task_instance.running", previous_state, task_instance) + + @hookimpl + def on_task_instance_success( # type: ignore[misc] + self, previous_state: TaskInstanceState, task_instance: TaskInstance, session: Session + ): + _produce_ti_message("task_instance.success", previous_state, task_instance) + + @hookimpl + def on_task_instance_failed( # type: ignore[misc] + self, + previous_state: TaskInstanceState, + task_instance: TaskInstance, + error: None | str | BaseException, + session: Session, + ): + _produce_ti_message("task_instance.failed", previous_state, task_instance, error=error) + + @hookimpl + def on_task_instance_skipped( # type: ignore[misc] + self, previous_state: TaskInstanceState, task_instance: TaskInstance, session: Session + ): + _produce_ti_message("task_instance.skipped", previous_state, task_instance) + + +def _get_enabled_listeners() -> list[object]: + listeners: list[object] = [] + if _dag_run_events_enabled(): + listeners.append(DagRunListener()) + if _task_instance_events_enabled(): + listeners.append(TaskListener()) + return listeners + + +class KafkaListenerPlugin(AirflowPlugin): + """Publishes Airflow DagRun and TaskInstance event messages to a defined Kafka topic.""" + + name = "kafka_listener" + listeners = _get_enabled_listeners() diff --git a/providers/apache/kafka/tests/integration/apache/kafka/plugins/__init__.py b/providers/apache/kafka/tests/integration/apache/kafka/plugins/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/providers/apache/kafka/tests/integration/apache/kafka/plugins/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/providers/apache/kafka/tests/integration/apache/kafka/plugins/dags/__init__.py b/providers/apache/kafka/tests/integration/apache/kafka/plugins/dags/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/providers/apache/kafka/tests/integration/apache/kafka/plugins/dags/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/providers/apache/kafka/tests/integration/apache/kafka/plugins/dags/demo_dag.py b/providers/apache/kafka/tests/integration/apache/kafka/plugins/dags/demo_dag.py new file mode 100644 index 0000000000000..4ecd9bec57721 --- /dev/null +++ b/providers/apache/kafka/tests/integration/apache/kafka/plugins/dags/demo_dag.py @@ -0,0 +1,45 @@ +# 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 logging +from datetime import datetime + +from airflow import DAG +from airflow.sdk import task + +logger = logging.getLogger(__name__) + +args = { + "owner": "airflow", + "start_date": datetime(2024, 9, 1), + "retries": 0, +} + + +@task +def task1(): + logger.info("running task1") + + +with DAG( + "demo_dag", + default_args=args, + schedule=None, + catchup=False, +) as dag: + task1() diff --git a/providers/apache/kafka/tests/integration/apache/kafka/plugins/test_listener.py b/providers/apache/kafka/tests/integration/apache/kafka/plugins/test_listener.py new file mode 100644 index 0000000000000..ce900f79b8734 --- /dev/null +++ b/providers/apache/kafka/tests/integration/apache/kafka/plugins/test_listener.py @@ -0,0 +1,167 @@ +# 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 +import logging +import os +import subprocess +import time +import uuid + +import pytest + +from airflow.providers.apache.kafka.hooks.client import KafkaAdminClientHook +from airflow.providers.apache.kafka.hooks.consume import KafkaConsumerHook +from airflow.sdk import Connection +from airflow.utils.state import State + +from tests_common.test_utils.integration_setup import ( + serialize_and_get_dags, + start_scheduler, + terminate_process, + unpause_trigger_dag_and_get_run_id, + wait_for_dag_run, +) + +log = logging.getLogger(__name__) + +client_config = { + "bootstrap.servers": "broker:29092", + "group.id": "kafka-listener-integration-test", + "enable.auto.commit": False, + "auto.offset.reset": "earliest", +} + + +def _wait_for_assignment(consumer, timeout: float = 10.0) -> None: + """Block until the consumer has been assigned partitions by the broker.""" + deadline = time.monotonic() + timeout + while not consumer.assignment(): + consumer.poll(0.5) + if time.monotonic() > deadline: + raise TimeoutError("Kafka consumer did not receive partition assignment in time") + + +@pytest.mark.integration("kafka") +@pytest.mark.backend("postgres") +class TestEventListener: + test_dir = os.path.dirname(os.path.abspath(__file__)) + dag_folder = os.path.join(test_dir, "dags") + + KAFKA_CONFIG_ID = "kafka_default" + # Use a unique topic per run to avoid errors on a re-run, in case + # the previous teardown hasn't finished with the topic deletion. + TOPIC = f"airflow.events.itest.{uuid.uuid4().hex[:8]}" + + @classmethod + def setup_class(cls): + # The pytest plugin strips AIRFLOW__*__* env vars (including the JWT secret set + # by Breeze). Both the scheduler and api-server subprocesses must share the same + # secret; otherwise each generates its own random key and token verification fails. + os.environ["AIRFLOW__API_AUTH__JWT_SECRET"] = "test-secret-key-for-testing" + os.environ["AIRFLOW__API_AUTH__JWT_ISSUER"] = "airflow" + + os.environ["AIRFLOW__SCHEDULER__STANDALONE_DAG_PROCESSOR"] = "False" + os.environ["AIRFLOW__SCHEDULER__PROCESSOR_POLL_INTERVAL"] = "2" + + os.environ["AIRFLOW__CORE__DAGS_FOLDER"] = f"{cls.dag_folder}" + os.environ["AIRFLOW__CORE__LOAD_EXAMPLES"] = "False" + os.environ["AIRFLOW__CORE__UNIT_TEST_MODE"] = "False" + + os.environ["AIRFLOW__KAFKA_LISTENER__DAG_RUN_EVENTS_ENABLED"] = "True" + os.environ["AIRFLOW__KAFKA_LISTENER__TASK_INSTANCE_EVENTS_ENABLED"] = "True" + os.environ["AIRFLOW__KAFKA_LISTENER__TOPIC"] = cls.TOPIC + os.environ["AIRFLOW__KAFKA_LISTENER__SOURCE"] = "dev-breeze" + + # Shared Kafka connection: used by the listener producer, the topic + # creation/deletion, and the consumer. + kafka_default_conn = Connection( + conn_id=cls.KAFKA_CONFIG_ID, + conn_type="kafka", + extra=json.dumps(client_config), + ) + os.environ["AIRFLOW_CONN_KAFKA_DEFAULT"] = kafka_default_conn.as_json() + + reset_command = ["airflow", "db", "reset", "--yes"] + subprocess.run(reset_command, check=True, env=os.environ.copy()) + + migrate_command = ["airflow", "db", "migrate"] + subprocess.run(migrate_command, check=True, env=os.environ.copy()) + + cls.dags = serialize_and_get_dags(dag_folder=cls.dag_folder) + + cls._admin = KafkaAdminClientHook(kafka_config_id=cls.KAFKA_CONFIG_ID) + # Tuple positions are (name, partition, replication). + cls._admin.create_topic([(cls.TOPIC, 1, 1)]) + + @classmethod + def teardown_class(cls): + try: + cls._admin.delete_topic([cls.TOPIC]) + time.sleep(2) # let the broker finish the async delete + except Exception as exc: + log.warning("teardown: failed to delete topic %r: %s", cls.TOPIC, exc) + + @pytest.fixture + def start_components(self): + scheduler_process = None + apiserver_process = None + try: + scheduler_process, apiserver_process = start_scheduler() + yield scheduler_process, apiserver_process + finally: + terminate_process(scheduler_process) + terminate_process(apiserver_process) + + @pytest.mark.execution_timeout(90) + def test_dag_run_produces_event_messages(self, start_components): + consumer = KafkaConsumerHook(topics=[self.TOPIC], kafka_config_id=self.KAFKA_CONFIG_ID).get_consumer() + _wait_for_assignment(consumer) + + dag_id = "demo_dag" + run_id = unpause_trigger_dag_and_get_run_id(dag_id=dag_id) + + state = wait_for_dag_run(dag_id=dag_id, run_id=run_id, max_wait_time=90) + assert state == State.SUCCESS, f"Dag run did not complete successfully. Final state: {state}." + + expected_events = [ + "dag_run.running", + "task_instance.running", + "task_instance.success", + "dag_run.success", + ] + + try: + messages = consumer.consume(num_messages=10, timeout=2.0) + finally: + consumer.close() + + topic_msg_bodies: list[dict] = [] + for message in messages: + topic_msg_bodies.append(json.loads(message.value())) + + assert len(topic_msg_bodies) == len(expected_events) + + for body in topic_msg_bodies: + assert body["dag_id"] == dag_id + assert body["run_id"] == run_id + if body["event"].startswith("task_instance."): + assert body["task_id"] == "task1" + + received_events = [body["event"] for body in topic_msg_bodies] + assert sorted(received_events) == sorted(expected_events) diff --git a/providers/apache/kafka/tests/unit/apache/kafka/hooks/test_base.py b/providers/apache/kafka/tests/unit/apache/kafka/hooks/test_base.py index 33a75d536a0ef..154588a66e880 100644 --- a/providers/apache/kafka/tests/unit/apache/kafka/hooks/test_base.py +++ b/providers/apache/kafka/tests/unit/apache/kafka/hooks/test_base.py @@ -16,6 +16,7 @@ # under the License. from __future__ import annotations +import json from unittest import mock from unittest.mock import MagicMock @@ -65,6 +66,22 @@ def test_get_conn_value_error(self, mock_get_connection, hook): with pytest.raises(ValueError, match="must be provided"): hook.get_conn() + @mock.patch(f"{BASEHOOK_PATCH_PATH}.get_connection") + def test_callbacks_resolved_from_connection_dotted_path(self, mock_get_connection): + stats_cb = MagicMock() + mock_get_connection.return_value.extra_dejson = { + "bootstrap.servers": "localhost:9092", + "error_cb": "json.loads", + "oauth_cb": "json.dumps", + "stats_cb": stats_cb, + } + + config = SomeKafkaHook().get_conn + assert config["error_cb"] is json.loads + assert config["oauth_cb"] is json.dumps + # Already-callable values on the connection are passed through unchanged. + assert config["stats_cb"] is stats_cb + @mock.patch("airflow.providers.apache.kafka.hooks.base.AdminClient") @mock.patch(f"{BASEHOOK_PATCH_PATH}.get_connection") def test_test_connection(self, mock_get_connection, admin_client, hook): @@ -90,6 +107,18 @@ def test_test_connection_no_topics(self, mock_get_connection, admin_client, hook mock_admin_instance.list_topics.assert_called_once_with(timeout=TIMEOUT) assert connection == (False, "Failed to establish connection.") + @mock.patch("airflow.providers.apache.kafka.hooks.base.AdminClient") + @mock.patch(f"{BASEHOOK_PATCH_PATH}.get_connection") + def test_test_connection_resolves_callbacks(self, mock_get_connection, admin_client, hook): + # ``test_connection`` runs the same callback resolution as the runtime clients, + # so dotted-path strings on the connection are exercised by the UI test too. + mock_get_connection.return_value.extra_dejson = { + "bootstrap.servers": "localhost:9092", + "error_cb": "json.loads", + } + assert hook.test_connection() == (True, "Connection successful.") + admin_client.assert_called_once_with({"bootstrap.servers": "localhost:9092", "error_cb": json.loads}) + @mock.patch("airflow.providers.apache.kafka.hooks.base.AdminClient") @mock.patch(f"{BASEHOOK_PATCH_PATH}.get_connection") def test_test_connection_exception(self, mock_get_connection, admin_client, hook): diff --git a/providers/apache/kafka/tests/unit/apache/kafka/hooks/test_consume.py b/providers/apache/kafka/tests/unit/apache/kafka/hooks/test_consume.py index 0bad37aeac32c..4036980c6260e 100644 --- a/providers/apache/kafka/tests/unit/apache/kafka/hooks/test_consume.py +++ b/providers/apache/kafka/tests/unit/apache/kafka/hooks/test_consume.py @@ -25,7 +25,7 @@ from airflow.models import Connection # Import Hook -from airflow.providers.apache.kafka.hooks.consume import KafkaConsumerHook +from airflow.providers.apache.kafka.hooks.consume import KafkaConsumerHook, error_callback class TestConsumerHook: @@ -59,3 +59,29 @@ def test_get_consumer(self, mock_client): mock_client_spec = MagicMock(spec=AdminClient) mock_client.return_value = mock_client_spec assert self.hook.get_consumer() == self.hook.get_conn + + @patch("airflow.providers.apache.kafka.hooks.consume.Consumer") + def test_default_error_cb(self, mock_consumer): + consumer = self.hook.get_consumer() + config = mock_consumer.call_args.args[0] + assert config["error_cb"] is error_callback + assert consumer == mock_consumer.return_value + mock_consumer.return_value.subscribe.assert_called_once_with(["test_1"]) + + @patch("airflow.providers.apache.kafka.hooks.consume.Consumer") + def test_user_error_cb_from_connection_resolved_and_not_overridden_by_default( + self, mock_consumer, create_connection_without_db + ): + # A dotted-path ``error_cb`` on the connection extras is resolved to the + # callable and used instead of the hook's default ``error_callback``. + create_connection_without_db( + Connection( + conn_id="kafka_cb", + conn_type="kafka", + extra=json.dumps({"bootstrap.servers": "localhost:9092", "error_cb": "json.loads"}), + ) + ) + hook = KafkaConsumerHook(["test_1"], kafka_config_id="kafka_cb") + hook.get_consumer() + config = mock_consumer.call_args.args[0] + assert config["error_cb"] is json.loads diff --git a/providers/apache/kafka/tests/unit/apache/kafka/hooks/test_produce.py b/providers/apache/kafka/tests/unit/apache/kafka/hooks/test_produce.py index 7b48026a47613..e8f8cd8f66721 100644 --- a/providers/apache/kafka/tests/unit/apache/kafka/hooks/test_produce.py +++ b/providers/apache/kafka/tests/unit/apache/kafka/hooks/test_produce.py @@ -60,3 +60,19 @@ def test_get_producer(self, mock_client): mock_client_spec = MagicMock(spec=AdminClient) mock_client.return_value = mock_client_spec assert self.hook.get_producer() == self.hook.get_conn + + @patch("airflow.providers.apache.kafka.hooks.produce.Producer") + def test_connection_callback_resolved_from_dotted_path(self, mock_producer, create_connection_without_db): + # A dotted-path ``oauth_cb`` on the connection extras is resolved to the callable + # before the producer is built. + create_connection_without_db( + Connection( + conn_id="kafka_cb", + conn_type="kafka", + extra=json.dumps({"bootstrap.servers": "localhost:9092", "oauth_cb": "json.dumps"}), + ) + ) + hook = KafkaProducerHook(kafka_config_id="kafka_cb") + hook.get_producer() + config = mock_producer.call_args.args[0] + assert config["oauth_cb"] is json.dumps diff --git a/providers/apache/kafka/tests/unit/apache/kafka/plugins/__init__.py b/providers/apache/kafka/tests/unit/apache/kafka/plugins/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/providers/apache/kafka/tests/unit/apache/kafka/plugins/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/providers/apache/kafka/tests/unit/apache/kafka/plugins/test_listener.py b/providers/apache/kafka/tests/unit/apache/kafka/plugins/test_listener.py new file mode 100644 index 0000000000000..38bdc0566372f --- /dev/null +++ b/providers/apache/kafka/tests/unit/apache/kafka/plugins/test_listener.py @@ -0,0 +1,538 @@ +# 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 collections.abc import Callable +from unittest.mock import MagicMock, patch + +import pytest +from confluent_kafka import KafkaError + +from airflow.models import Connection +from airflow.providers.apache.kafka.plugins import listener +from airflow.providers.apache.kafka.plugins.listener import DagRunListener, TaskListener + +from tests_common.test_utils.config import conf_vars +from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS + +# TaskInstance listener hooks required a ``session`` param for AF versions prior to 3.0. +_TI_SESSION_KWARG: dict = {} if AIRFLOW_V_3_0_PLUS else {"session": None} + +_DAG_ID = "test_dag1" +_DAG_RUN_ID = "test_dag_run1" +_TASK_ID = "test_task1" +_RUN_TYPE = "manual" + +_KAFKA_TOPIC = "airflow.events" +_SOURCE = "af-dev-env" + +# The hook binds these names in its own namespace, so patch them there. +_PRODUCER_CLS = "airflow.providers.apache.kafka.hooks.produce.Producer" +_PRODUCER_HOOK_CLS = "airflow.providers.apache.kafka.hooks.produce.KafkaProducerHook" + + +@pytest.fixture(autouse=True) +def _default_kafka_conn(create_connection_without_db): + # The listener builds its producer through ``KafkaProducerHook``, which requires + # a resolvable ``kafka_default`` connection. + create_connection_without_db( + Connection( + conn_id="kafka_default", + conn_type="kafka", + extra=json.dumps({"bootstrap.servers": "broker:29092"}), + ) + ) + + +@pytest.fixture(autouse=True) +def _clear_cached_values(): + """Invalidates caches before each test run.""" + # Reset config cache. + listener._dag_run_events_enabled.cache_clear() + listener._task_instance_events_enabled.cache_clear() + listener._get_topic.cache_clear() + listener._get_kafka_config_id.cache_clear() + listener._get_source.cache_clear() + listener._get_dag_run_dag_id_allowlist.cache_clear() + listener._get_dag_run_dag_id_denylist.cache_clear() + listener._get_task_instance_dag_id_allowlist.cache_clear() + listener._get_task_instance_dag_id_denylist.cache_clear() + listener._get_task_instance_task_id_allowlist.cache_clear() + listener._get_task_instance_task_id_denylist.cache_clear() + listener._get_topic_check_timeout.cache_clear() + listener._get_topic_check_retry_interval.cache_clear() + + # Reset the module-level producer and topic state. + listener._producer = None + listener._topic_exists = False + listener._topic_check_retry_after = 0.0 + + +@pytest.fixture +def dr_mock(): + dr_mock = MagicMock() + dr_mock.dag_id = _DAG_ID + dr_mock.run_id = _DAG_RUN_ID + dr_mock.run_type = _RUN_TYPE + dr_mock.logical_date = None + dr_mock.start_date = None + dr_mock.end_date = None + return dr_mock + + +@pytest.fixture +def ti_mock(): + ti_mock = MagicMock() + ti_mock.dag_id = _DAG_ID + ti_mock.run_id = _DAG_RUN_ID + ti_mock.task_id = _TASK_ID + ti_mock.try_number = 2 + ti_mock.map_index = -1 + return ti_mock + + +@pytest.fixture +def kafka_producer_mock(): + """Patches the Producer class used by KafkaProducerHook.""" + mock = MagicMock() + # Make the topic check succeed for any configured topic name. + mock.list_topics.return_value.topics.__contains__.return_value = True + with patch(_PRODUCER_CLS, return_value=mock): + yield mock + + +def _assert_common_message_fields(kafka_producer_mock, expected_event: str) -> dict: + kafka_producer_mock.list_topics.assert_called_once() + kafka_producer_mock.produce.assert_called_once() + + args = kafka_producer_mock.produce.call_args.args + kwargs = kafka_producer_mock.produce.call_args.kwargs + assert args == (_KAFKA_TOPIC,) + assert kwargs["key"] == f"{_DAG_ID}/{_DAG_RUN_ID}".encode() + + body = json.loads(kwargs["value"].decode("utf-8")) + assert body["schema_version"] == listener.SCHEMA_VERSION + assert body["source"] == _SOURCE + assert body["event"] == expected_event + assert body["dag_id"] == _DAG_ID + assert body["run_id"] == _DAG_RUN_ID + assert "timestamp" in body + return body + + +@pytest.mark.parametrize( + ("configs", "expected_listener_classes"), + [ + pytest.param( + { + ("kafka_listener", "dag_run_events_enabled"): "True", + ("kafka_listener", "task_instance_events_enabled"): "True", + }, + [DagRunListener, TaskListener], + id="both_listeners_enabled", + ), + pytest.param( + { + ("kafka_listener", "dag_run_events_enabled"): "True", + }, + [DagRunListener], + id="only_dr_events", + ), + pytest.param( + { + ("kafka_listener", "task_instance_events_enabled"): "True", + }, + [TaskListener], + id="only_ti_events", + ), + pytest.param( + {}, + [], + id="listeners_disabled_by_default", + ), + ], +) +def test_get_enabled_listeners(configs, expected_listener_classes): + with conf_vars(configs): + registered_listeners = listener._get_enabled_listeners() + # Both lists should have the same size. + assert len(registered_listeners) == len(expected_listener_classes) + + # Iterate both lists and assert. + assert [type(x) for x in registered_listeners] == expected_listener_classes + + +@pytest.mark.parametrize( + ("dr_event_function", "expected_event"), + [ + pytest.param(DagRunListener.on_dag_run_running, "dag_run.running", id="dr_running"), + pytest.param(DagRunListener.on_dag_run_success, "dag_run.success", id="dr_success"), + pytest.param(DagRunListener.on_dag_run_failed, "dag_run.failed", id="dr_failed"), + ], +) +@conf_vars( + { + ("kafka_listener", "dag_run_events_enabled"): "True", + ("kafka_listener", "topic"): _KAFKA_TOPIC, + ("kafka_listener", "source"): _SOURCE, + } +) +def test_produce_dr_message( + dr_event_function: Callable, + expected_event: str, + dr_mock, + kafka_producer_mock, +): + dr_event_function(DagRunListener(), dag_run=dr_mock, msg=None) + + body = _assert_common_message_fields(kafka_producer_mock, expected_event) + assert body["run_type"] == _RUN_TYPE + assert body["logical_date"] == "" + assert body["start_date"] == "" + assert body["end_date"] == "" + assert body["msg"] == "" + + +@pytest.mark.parametrize( + ("ti_event_function", "expected_event", "extra_kwargs"), + [ + pytest.param(TaskListener.on_task_instance_running, "task_instance.running", {}, id="ti_running"), + pytest.param(TaskListener.on_task_instance_success, "task_instance.success", {}, id="ti_success"), + pytest.param(TaskListener.on_task_instance_skipped, "task_instance.skipped", {}, id="ti_skipped"), + pytest.param( + TaskListener.on_task_instance_failed, + "task_instance.failed", + {"error": ValueError("boom")}, + id="ti_failed", + ), + ], +) +@conf_vars( + { + ("kafka_listener", "task_instance_events_enabled"): "True", + ("kafka_listener", "topic"): _KAFKA_TOPIC, + ("kafka_listener", "source"): _SOURCE, + } +) +def test_produce_ti_message( + ti_event_function: Callable, + expected_event: str, + extra_kwargs: dict, + ti_mock, + kafka_producer_mock, +): + ti_event_function( + TaskListener(), + previous_state="running", + task_instance=ti_mock, + **extra_kwargs, + **_TI_SESSION_KWARG, + ) + + body = _assert_common_message_fields(kafka_producer_mock, expected_event) + assert body["task_id"] == _TASK_ID + assert body["try_number"] == 2 + assert body["map_index"] == -1 + assert body["previous_state"] == "running" + + if "error" in extra_kwargs: + assert body["error"] == str(extra_kwargs["error"]) + else: + assert "error" not in body + + +class TestGetProducer: + """Tests for the lifecycle and caching behavior of `_get_producer`.""" + + @pytest.fixture(autouse=True) + def _producer_conf(self): + with conf_vars({("kafka_listener", "topic"): _KAFKA_TOPIC}): + yield + + def test_producer_cached_on_success(self): + kafka_producer_mock = MagicMock() + + with patch(_PRODUCER_CLS, return_value=kafka_producer_mock) as producer_cls_mock: + producer1 = listener._get_producer() + producer2 = listener._get_producer() + + assert producer1 is kafka_producer_mock + assert producer2 is producer1 + producer_cls_mock.assert_called_once() + + def test_init_failure_warns_and_returns_none(self, monkeypatch): + log_warning_mock = MagicMock() + monkeypatch.setattr(listener.log, "warning", log_warning_mock) + + with patch(_PRODUCER_CLS, side_effect=ValueError("boom")): + assert listener._get_producer() is None + + log_warning_mock.assert_called_once() + warning_format, exc_arg = log_warning_mock.call_args.args + assert "failed to initialize producer" in warning_format + assert str(exc_arg) == "boom" + + def test_atexit_register(self, monkeypatch): + """On success, _get_producer registers a flush handler with atexit so any + pending messages are delivered to the broker when the process exits.""" + kafka_producer_mock = MagicMock() + + register_mock = MagicMock() + monkeypatch.setattr(listener.atexit, "register", register_mock) + + with patch(_PRODUCER_CLS, return_value=kafka_producer_mock): + listener._get_producer() + + register_mock.assert_called_once_with(listener._flush_producer_at_exit, kafka_producer_mock) + + def test_hook_built_from_listener_config(self): + """The listener section's ``kafka_config_id`` is forwarded to the hook.""" + hook_mock = MagicMock() + + with conf_vars({("kafka_listener", "kafka_config_id"): "kafka_events"}): + with patch(_PRODUCER_HOOK_CLS, return_value=hook_mock) as hook_cls_mock: + producer = listener._get_producer() + + hook_cls_mock.assert_called_once_with(kafka_config_id="kafka_events") + assert producer is hook_mock.get_producer.return_value + + def test_hook_defaults_when_options_unset(self): + """Without listener options the hook falls back to its own defaults.""" + hook_mock = MagicMock() + + with patch(_PRODUCER_HOOK_CLS, return_value=hook_mock) as hook_cls_mock: + listener._get_producer() + + # Unset options must not be passed to the hook at all (e.g. as ""), so the + # hook's own defaults (the "kafka_default" connection) apply. + hook_cls_mock.assert_called_once_with() + + +class TestCheckTopicExists: + """Tests for the topic check and its retry-after-failure cooldown.""" + + @pytest.fixture(autouse=True) + def _topic_conf(self): + with conf_vars({("kafka_listener", "topic"): _KAFKA_TOPIC}): + yield + + def test_topic_exists_kept_for_process_lifetime(self): + kafka_producer_mock = MagicMock() + kafka_producer_mock.list_topics.return_value.topics.__contains__.return_value = True + + with patch(_PRODUCER_CLS, return_value=kafka_producer_mock): + assert listener._check_topic_exists() is True + # Once confirmed, the broker is not queried again. + assert listener._check_topic_exists() is True + + kafka_producer_mock.list_topics.assert_called_once() + + def test_topic_doesnt_exist(self, monkeypatch): + kafka_producer_mock = MagicMock() + kafka_producer_mock.list_topics.return_value.topics.__contains__.return_value = False + + log_warning_mock = MagicMock() + monkeypatch.setattr(listener.log, "warning", log_warning_mock) + + with patch(_PRODUCER_CLS, return_value=kafka_producer_mock): + assert listener._check_topic_exists() is False + + log_warning_mock.assert_called_once() + # Format string + the two formatting args are passed positionally to log.warning. + warning_format, topic_arg, retry_interval_arg = log_warning_mock.call_args.args + assert "topic %r not found on the broker" in warning_format + assert topic_arg == _KAFKA_TOPIC + assert retry_interval_arg == listener._get_topic_check_retry_interval() + + def test_list_topics_failure_warns_and_sets_cooldown(self, monkeypatch): + kafka_producer_mock = MagicMock() + kafka_producer_mock.list_topics.side_effect = ValueError("broker down") + + log_warning_mock = MagicMock() + monkeypatch.setattr(listener.log, "warning", log_warning_mock) + + with patch(_PRODUCER_CLS, return_value=kafka_producer_mock): + assert listener._check_topic_exists() is False + # Within the cooldown window the check is not retried. + assert listener._check_topic_exists() is False + + kafka_producer_mock.list_topics.assert_called_once() + log_warning_mock.assert_called_once() + assert "topic check failed" in log_warning_mock.call_args.args[0] + + def test_retried_after_failure_cooldown(self, monkeypatch): + time_now = [1000.0] + monkeypatch.setattr(listener.time, "monotonic", lambda: time_now[0]) + + kafka_producer_mock = MagicMock() + # The topic doesn't exist yet. + kafka_producer_mock.list_topics.return_value.topics.__contains__.return_value = False + + with patch(_PRODUCER_CLS, return_value=kafka_producer_mock): + # First check fails: the topic is missing. + assert listener._check_topic_exists() is False + assert kafka_producer_mock.list_topics.call_count == 1 + + # The interval hasn't passed and the check isn't retried. + time_now[0] += 1 + assert listener._check_topic_exists() is False + assert kafka_producer_mock.list_topics.call_count == 1 + + # Past the interval. The check is retried and succeeds this time. + kafka_producer_mock.list_topics.return_value.topics.__contains__.return_value = True + time_now[0] += listener._get_topic_check_retry_interval() + 1 + assert listener._check_topic_exists() is True + assert kafka_producer_mock.list_topics.call_count == 2 + + def test_producer_init_failure_sets_cooldown(self, monkeypatch): + time_now = [1000.0] + monkeypatch.setattr(listener.time, "monotonic", lambda: time_now[0]) + + with patch(_PRODUCER_CLS, side_effect=ValueError("boom")) as producer_cls_mock: + assert listener._check_topic_exists() is False + assert producer_cls_mock.call_count == 1 + + # The interval hasn't passed and the producer initialization isn't retried. + time_now[0] += 1 + assert listener._check_topic_exists() is False + assert producer_cls_mock.call_count == 1 + + # Past the interval. The producer initialization is retried. + time_now[0] += listener._get_topic_check_retry_interval() + 1 + assert listener._check_topic_exists() is False + assert producer_cls_mock.call_count == 2 + + @pytest.mark.parametrize( + ("kafka_error", "topic_exists_expected"), + [ + pytest.param(KafkaError.UNKNOWN_TOPIC_OR_PART, False, id="broker_says_topic_gone"), + pytest.param(KafkaError._UNKNOWN_TOPIC, False, id="local_metadata_says_topic_gone"), + pytest.param(KafkaError.BROKER_NOT_AVAILABLE, True, id="unrelated_delivery_error_ignored"), + ], + ) + def test_delivery_report_invalidates_topic_confirmation_only_on_missing_topic( + self, monkeypatch, kafka_error, topic_exists_expected + ): + """ + If the delivery report failed and the topic is gone, ``_topic_exists`` must be reset so that the + next produce triggers a fresh check. Any delivery error unrelated to the topic, should be ignored. + """ + + time_now = [1000.0] + monkeypatch.setattr(listener.time, "monotonic", lambda: time_now[0]) + + kafka_producer_mock = MagicMock() + kafka_producer_mock.list_topics.return_value.topics.__contains__.return_value = True + with patch(_PRODUCER_CLS, return_value=kafka_producer_mock): + assert listener._check_topic_exists() is True + assert listener._topic_exists is True + + err = MagicMock() + err.code.return_value = kafka_error + listener._on_delivery(err, MagicMock()) + + assert listener._topic_exists is topic_exists_expected + + if not topic_exists_expected: + # Cooldown holds off the immediate re-check. + assert listener._check_topic_exists() is False + # After the cooldown, the check re-verifies; the topic is still on the broker, + # so confirmation is restored. + time_now[0] += listener._get_topic_check_retry_interval() + 1 + assert listener._check_topic_exists() is True + + +class TestFilters: + """Tests for the allowlist/denylist behavior of both DagRun and TaskInstance events.""" + + @pytest.fixture(autouse=True) + def _filters_conf(self): + with conf_vars( + { + ("kafka_listener", "dag_run_events_enabled"): "True", + ("kafka_listener", "task_instance_events_enabled"): "True", + ("kafka_listener", "topic"): _KAFKA_TOPIC, + } + ): + yield + + @pytest.mark.parametrize( + ("id_to_check", "allowlist", "denylist", "expected_result"), + [ + pytest.param("some_value", (), (), True, id="empty_lists_allow_all"), + pytest.param("dag__1", ("dag_*",), (), True, id="allowlist_match"), + pytest.param("dag1", ("dag2_*", "dag2"), (), False, id="allowlist_no_match"), + pytest.param("task1_dev", (), ("task1_*",), False, id="denylist_match"), + pytest.param("dag1_task1", ("dag1_*",), ("dag1_task1",), False, id="deny_beats_allow"), + ], + ) + def test_id_is_allowed(self, id_to_check, allowlist, denylist, expected_result): + assert listener._id_is_allowed(id_to_check, allowlist, denylist) is expected_result + + @conf_vars( + { + ("kafka_listener", "task_instance_dag_id_denylist"): "dag1_*", + } + ) + def test_dag_run_filter_ignores_ti_lists(self): + # The dag_id lists for DR events, are separate from the dag_id lists for TI events. + # The dag_id denylist for task instances would match "dag1_task1" but DR filter must ignore it. + assert listener._dag_run_event_allowed("dag1_task1") is True + + @conf_vars( + { + ("kafka_listener", "dag_run_dag_id_denylist"): "dag1_*", + } + ) + def test_task_instance_filter_ignores_dr_lists(self): + # The dag_id lists for DR events, are separate from the dag_id lists for TI events. + # DR denylist would match "dag1_task1" but TI filter must ignore it. + assert listener._task_instance_event_allowed("dag1_task1", "load") is True + + @conf_vars( + { + ("kafka_listener", "task_instance_dag_id_allowlist"): "dag1_*", + ("kafka_listener", "task_instance_task_id_denylist"): "*_cleanup", + } + ) + def test_task_instance_requires_both_dag_id_and_task_id(self): + # dag_id passes (matches allowlist), task_id passes (no deny match). + assert listener._task_instance_event_allowed("dag1_task1", "task2") is True + # dag_id passes, task_id denied. + assert listener._task_instance_event_allowed("dag1_task1", "task2_cleanup") is False + # dag_id denied (no allowlist match), task_id passes. + assert listener._task_instance_event_allowed("other_dag", "load") is False + + @conf_vars({("kafka_listener", "dag_run_dag_id_denylist"): _DAG_ID}) + def test_dr_filter_blocks_emission(self, dr_mock, kafka_producer_mock): + DagRunListener().on_dag_run_running(dag_run=dr_mock, msg=None) + kafka_producer_mock.produce.assert_not_called() + + @conf_vars({("kafka_listener", "task_instance_dag_id_denylist"): _DAG_ID}) + def test_ti_filter_blocks_emission_via_dag_id(self, ti_mock, kafka_producer_mock): + TaskListener().on_task_instance_running( + previous_state="queued", task_instance=ti_mock, **_TI_SESSION_KWARG + ) + kafka_producer_mock.produce.assert_not_called() + + @conf_vars({("kafka_listener", "task_instance_task_id_denylist"): _TASK_ID}) + def test_ti_filter_blocks_emission_via_task_id(self, ti_mock, kafka_producer_mock): + TaskListener().on_task_instance_running( + previous_state="queued", task_instance=ti_mock, **_TI_SESSION_KWARG + ) + kafka_producer_mock.produce.assert_not_called() diff --git a/scripts/ci/docker-compose/integration-kafka.yml b/scripts/ci/docker-compose/integration-kafka.yml index e476c51c7e3c0..f0ba7fab7d619 100644 --- a/scripts/ci/docker-compose/integration-kafka.yml +++ b/scripts/ci/docker-compose/integration-kafka.yml @@ -60,5 +60,9 @@ services: airflow: environment: - INTEGRATION_KAFKA=true + - AIRFLOW__KAFKA_LISTENER__DAG_RUN_EVENTS_ENABLED=True + - AIRFLOW__KAFKA_LISTENER__TASK_INSTANCE_EVENTS_ENABLED=True + - AIRFLOW__KAFKA_LISTENER__TOPIC=airflow.events + - AIRFLOW__KAFKA_LISTENER__SOURCE=dev-breeze depends_on: - broker