This is an automated email from the ASF dual-hosted git repository.
shahar1 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 8302ba498ae Add a DagRun and TaskInstance event listener under the
Kafka provider (#68082)
8302ba498ae is described below
commit 8302ba498aeb46d5f715389cfd51fb0461481704
Author: Christos Bisias <[email protected]>
AuthorDate: Mon Jul 13 13:58:45 2026 +0300
Add a DagRun and TaskInstance event listener under the Kafka provider
(#68082)
---
airflow-core/tests/integration/otel/test_otel.py | 200 +-------
.../core_api/routes/public/test_plugins.py | 15 +-
.../tests/unit/plugins/test_plugins_manager.py | 2 +-
.../tests_common/test_utils/integration_setup.py | 201 ++++++++
providers/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 +
.../providers/apache/kafka/get_provider_info.py | 104 ++++
.../airflow/providers/apache/kafka/hooks/base.py | 16 +
.../providers/apache/kafka/hooks/consume.py | 7 +-
.../providers/apache/kafka/plugins/__init__.py | 16 +
.../providers/apache/kafka/plugins/listener.py | 470 ++++++++++++++++++
.../integration/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 +++++++
.../tests/unit/apache/kafka/hooks/test_base.py | 29 ++
.../tests/unit/apache/kafka/hooks/test_consume.py | 28 +-
.../tests/unit/apache/kafka/hooks/test_produce.py | 16 +
.../tests/unit/apache/kafka/plugins/__init__.py | 16 +
.../unit/apache/kafka/plugins/test_listener.py | 538 +++++++++++++++++++++
scripts/ci/docker-compose/integration-kafka.yml | 4 +
23 files changed, 1948 insertions(+), 194 deletions(-)
diff --git a/airflow-core/tests/integration/otel/test_otel.py
b/airflow-core/tests/integration/otel/test_otel.py
index ec6a7a6c94c..39d7ee6bdab 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 @@ class TestOtelIntegration:
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 @@ class TestOtelIntegration:
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 @@ class TestOtelIntegration:
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 @@ class TestOtelIntegration:
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 @@ class TestOtelIntegration:
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 @@ class TestOtelIntegration:
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 6bca49d8052..e995cb4b137 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 @@ class TestGetPlugins:
# 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 ce55db1e269..0dcac96fe0a 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 @@ class TestPluginsManager:
# 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 00000000000..59476d8c4d0
--- /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 00000000000..80d0d373f8f
--- /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 cee81d3a021..d9db2b715de 100644
--- a/providers/apache/kafka/docs/index.rst
+++ b/providers/apache/kafka/docs/index.rst
@@ -47,6 +47,7 @@
:maxdepth: 1
:caption: References
+ Configuration <configurations-ref>
Python API <_api/airflow/providers/apache/kafka/index>
.. toctree::
diff --git a/providers/apache/kafka/provider.yaml
b/providers/apache/kafka/provider.yaml
index f43f76ce0b8..8e23df57309 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 b43f5a60790..3c5ead9a4b8 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 1a68a7b9167..d8c8510cb5d 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 b3b76ff1421..32ceef972b8 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 typing import Any
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 @@ class KafkaBaseHook(BaseHook):
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 @@ class KafkaBaseHook(BaseHook):
"""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 fa2f5844927..a420378be1e 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 collections.abc import Sequence
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 @@ class KafkaConsumerHook(KafkaBaseHook):
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 00000000000..13a83393a91
--- /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 00000000000..342b37139a0
--- /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 00000000000..13a83393a91
--- /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 00000000000..13a83393a91
--- /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 00000000000..4ecd9bec577
--- /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 00000000000..ce900f79b87
--- /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")
+
+
[email protected]("kafka")
[email protected]("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 33a75d536a0..154588a66e8 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 @@ class TestKafkaBaseHook:
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 @@ class TestKafkaBaseHook:
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 0bad37aeac3..4036980c626 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 confluent_kafka.admin import AdminClient
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 @@ class TestConsumerHook:
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 7b48026a476..e8f8cd8f667 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 @@ class TestProducerHook:
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 00000000000..13a83393a91
--- /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 00000000000..38bdc056637
--- /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"
+
+
[email protected](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"}),
+ )
+ )
+
+
[email protected](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
+
+
[email protected]
+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
+
+
[email protected]
+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
+
+
[email protected]
+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
+
+
[email protected](
+ ("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
+
+
[email protected](
+ ("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"] == ""
+
+
[email protected](
+ ("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 e476c51c7e3..f0ba7fab7d6 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
