This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 0f185b6a7ae Fix KafkaBaseHook.test_connection missing oauth_cb for
managed Kafka (#69507)
0f185b6a7ae is described below
commit 0f185b6a7ae941350e5ec355cc0117145b37a5a4
Author: Tim <[email protected]>
AuthorDate: Fri Jul 31 12:17:35 2026 -0700
Fix KafkaBaseHook.test_connection missing oauth_cb for managed Kafka
(#69507)
---
.../airflow/providers/apache/kafka/hooks/base.py | 27 +++++++++++++++-------
.../tests/unit/apache/kafka/hooks/test_base.py | 25 +++++++++++++++++---
2 files changed, 41 insertions(+), 11 deletions(-)
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 32ceef972b8..4fb4cc71707 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
@@ -97,9 +97,15 @@ class KafkaBaseHook(BaseHook):
if isinstance(value, str):
config[key] = import_string(value)
- @cached_property
- def get_conn(self) -> Any:
- """Get the configuration object."""
+ def _build_config(self) -> dict[str, Any]:
+ """
+ Build the confluent-kafka configuration for this connection.
+
+ Resolves callback options provided as dotted-path strings and injects
the
+ managed OAuth token callback (Google Managed Kafka or Amazon MSK IAM)
when
+ applicable, so that establishing a connection and testing it always
use an
+ identical configuration.
+ """
config = self.get_connection(self.kafka_config_id).extra_dejson
self._resolve_callbacks(config)
@@ -127,7 +133,12 @@ class KafkaBaseHook(BaseHook):
config.update({"oauth_cb": token})
else:
self._maybe_add_msk_iam_oauth(config, bootstrap_servers)
- return self._get_client(config)
+ return config
+
+ @cached_property
+ def get_conn(self) -> Any:
+ """Get the configuration object."""
+ return self._get_client(self._build_config())
def _maybe_add_msk_iam_oauth(self, config: dict[str, Any],
bootstrap_servers: str | None) -> None:
"""
@@ -172,10 +183,10 @@ class KafkaBaseHook(BaseHook):
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)
+ # Build the config exactly as a real connection would, so resolved
+ # dotted-path callbacks and the managed OAuth token callback
(Google
+ # Managed Kafka or Amazon MSK IAM) are exercised by the UI test
too.
+ config = self._build_config()
t = AdminClient(config).list_topics(timeout=10)
if t:
return True, "Connection successful."
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 154588a66e8..4d472615344 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
@@ -85,7 +85,7 @@ class TestKafkaBaseHook:
@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):
- config = {"bootstrap.servers": MagicMock()}
+ config = {"bootstrap.servers": "localhost:9092"}
mock_get_connection.return_value.extra_dejson = config
connection = hook.test_connection()
admin_client.assert_called_once_with(config)
@@ -99,7 +99,7 @@ class TestKafkaBaseHook:
)
@mock.patch(f"{BASEHOOK_PATCH_PATH}.get_connection")
def test_test_connection_no_topics(self, mock_get_connection,
admin_client, hook):
- config = {"bootstrap.servers": MagicMock()}
+ config = {"bootstrap.servers": "localhost:9092"}
mock_get_connection.return_value.extra_dejson = config
connection = hook.test_connection()
admin_client.assert_called_once_with(config)
@@ -122,12 +122,31 @@ class TestKafkaBaseHook:
@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):
- config = {"bootstrap.servers": MagicMock()}
+ config = {"bootstrap.servers": "localhost:9092"}
mock_get_connection.return_value.extra_dejson = config
admin_client.return_value.list_topics.side_effect = [ValueError("some
error")]
connection = hook.test_connection()
assert connection == (False, "some error")
+
@mock.patch("airflow.providers.google.cloud.hooks.managed_kafka.ManagedKafkaHook")
+ @mock.patch("airflow.providers.apache.kafka.hooks.base.AdminClient")
+ @mock.patch(f"{BASEHOOK_PATCH_PATH}.get_connection")
+ def test_test_connection_injects_managed_kafka_oauth(
+ self, mock_get_connection, admin_client, managed_kafka_hook, hook
+ ):
+ # A managed Kafka cluster needs an oauth_cb token callback.
test_connection must apply
+ # the same injection as get_conn, otherwise "Test Connection" fails in
the UI even though
+ # the hook itself authenticates correctly.
+ config = {
+ "bootstrap.servers":
"bootstrap.my-cluster.us-central1.managedkafka.my-project.cloud.goog:9092"
+ }
+ mock_get_connection.return_value.extra_dejson = config
+ connection = hook.test_connection()
+ admin_client.assert_called_once()
+ passed_config = admin_client.call_args.args[0]
+ assert passed_config["oauth_cb"] is
managed_kafka_hook.return_value.get_confluent_token
+ assert connection == (True, "Connection successful.")
+
@mock.patch(f"{BASEHOOK_PATCH_PATH}.get_connection")
def test_get_conn_msk_iam_provisioned(self, mock_get_connection, hook):
config = {
