From 1f61aaaf896260e1eb42847fcdb6f9d4da7f0edb Mon Sep 17 00:00:00 2001 From: Arnold Lin Date: Fri, 18 Sep 2026 01:02:06 -0400 Subject: [PATCH 1/3] feat: Add public callback function --- .../airflow/providers/amazon/aws/hooks/msk.py | 19 ++++++++++ .../tests/unit/amazon/aws/hooks/test_msk.py | 36 ++++++++++++++++++- 2 files changed, 54 insertions(+), 1 deletion(-) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/hooks/msk.py b/providers/amazon/src/airflow/providers/amazon/aws/hooks/msk.py index e86404b5f0600..572e1f1eb9f2a 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/hooks/msk.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/hooks/msk.py @@ -19,6 +19,7 @@ from __future__ import annotations +import json from typing import TYPE_CHECKING from botocore.credentials import CredentialProvider @@ -64,3 +65,21 @@ def confluent_token(self, config_str: str) -> tuple[str, float]: self.region_name, _MskCredentialsProvider(self) ) return token, expiry_ms / 1000 + + +def oauth_cb(config_str: str) -> tuple[str, float]: + """Generate an Amazon MSK IAM token for a ``confluent_kafka`` OAuth callback.""" + try: + options = json.loads(config_str or "{}") + except json.JSONDecodeError as exc: + raise ValueError("Invalid JSON in config_str") from exc + + if not isinstance(options, dict): + raise ValueError("config_str must contain a JSON object") + + aws_conn_id = options.get("aws_conn_id") + if not isinstance(aws_conn_id, str) or not aws_conn_id: + raise ValueError("Missing 'aws_conn_id' in config_str") + + hook = MskHook(aws_conn_id=aws_conn_id, region_name=options.get("region_name")) + return hook.confluent_token(config_str) diff --git a/providers/amazon/tests/unit/amazon/aws/hooks/test_msk.py b/providers/amazon/tests/unit/amazon/aws/hooks/test_msk.py index 846e162587ae7..d4ef46aa509d8 100644 --- a/providers/amazon/tests/unit/amazon/aws/hooks/test_msk.py +++ b/providers/amazon/tests/unit/amazon/aws/hooks/test_msk.py @@ -22,7 +22,7 @@ import pytest from botocore.credentials import CredentialProvider -from airflow.providers.amazon.aws.hooks.msk import MskHook +from airflow.providers.amazon.aws.hooks.msk import MskHook, oauth_cb MOCK_MSK_SIGNER_MODULE = mock.MagicMock() @@ -61,3 +61,37 @@ def test_confluent_token_requires_region(self, mock_region_name): with pytest.raises(ValueError, match="AWS region is required"): self.hook.confluent_token("") + + +class TestOauthCallback: + @pytest.mark.parametrize( + ("config_str", "region_name"), + [ + ('{"aws_conn_id":"aws_msk"}', None), + ('{"aws_conn_id":"aws_msk","region_name":"us-east-1"}', "us-east-1"), + ], + ) + @mock.patch("airflow.providers.amazon.aws.hooks.msk.MskHook", autospec=True) + def test_uses_configured_aws_connection(self, mock_hook, config_str, region_name): + mock_hook.return_value.confluent_token.return_value = ("token", 1_700_000_900.0) + + assert oauth_cb(config_str) == ("token", 1_700_000_900.0) + + assert mock_hook.mock_calls == [ + mock.call(aws_conn_id="aws_msk", region_name=region_name), + mock.call().confluent_token(config_str), + ] + + @pytest.mark.parametrize( + ("config_str", "error"), + [ + ("invalid", "Invalid JSON in config_str"), + ("[]", "config_str must contain a JSON object"), + ("", "Missing 'aws_conn_id' in config_str"), + ('{"aws_conn_id":null}', "Missing 'aws_conn_id' in config_str"), + ('{"aws_conn_id":42}', "Missing 'aws_conn_id' in config_str"), + ], + ) + def test_rejects_invalid_config(self, config_str, error): + with pytest.raises(ValueError, match=error): + oauth_cb(config_str) From 4085f457fa38cae252fdf3b72e7323b064ab348a Mon Sep 17 00:00:00 2001 From: Arnold Lin Date: Fri, 18 Sep 2026 01:17:23 -0400 Subject: [PATCH 2/3] feat: Add doc about how to configure oauth_cb --- providers/amazon/docs/index.rst | 1 + providers/amazon/docs/msk.rst | 44 +++++++++++++++ .../apache/kafka/docs/connections/kafka.rst | 54 ++++++++++++++++++- 3 files changed, 98 insertions(+), 1 deletion(-) create mode 100644 providers/amazon/docs/msk.rst diff --git a/providers/amazon/docs/index.rst b/providers/amazon/docs/index.rst index f134f522e998a..91bed66da4a2a 100644 --- a/providers/amazon/docs/index.rst +++ b/providers/amazon/docs/index.rst @@ -42,6 +42,7 @@ Deferrable Operators Secrets backends Logging for Tasks + Amazon MSK Configuration Executors Message Queues diff --git a/providers/amazon/docs/msk.rst b/providers/amazon/docs/msk.rst new file mode 100644 index 0000000000000..7d0ab3436bf25 --- /dev/null +++ b/providers/amazon/docs/msk.rst @@ -0,0 +1,44 @@ + .. 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. + +Amazon Managed Streaming for Apache Kafka +========================================= + +`Amazon MSK `_ is a managed Apache Kafka service. The Amazon +provider's :class:`~airflow.providers.amazon.aws.hooks.msk.MskHook` provides a boto3 client for +managing MSK clusters and can generate IAM authentication tokens for Kafka clients. The +Apache Kafka provider handles producing and consuming topic messages. + +Trigger a Dag from an MSK topic +------------------------------- + +Use :class:`~airflow.providers.apache.kafka.triggers.msg_queue.KafkaMessageQueueTrigger` with +``AssetWatcher`` to trigger a Dag when a message arrives. Set the trigger's ``kafka_config_id`` +to a Kafka connection configured for the MSK brokers. There is no separate Amazon MSK message +queue trigger; the standard Kafka trigger works with MSK because MSK uses the Kafka protocol. +See the `Kafka message queue trigger guide +`_ +for an ``AssetWatcher`` example. + +Configure the Kafka connection +------------------------------ + +The `Apache Kafka connection guide +`_ +explains both automatic MSK IAM authentication and explicit ``oauth_cb`` configuration using an +AWS connection. The Amazon provider supplies the ``oauth_cb`` token callback; the Kafka provider +uses it when connecting to MSK. diff --git a/providers/apache/kafka/docs/connections/kafka.rst b/providers/apache/kafka/docs/connections/kafka.rst index 486a5186d3fc3..869b42ec1aa52 100644 --- a/providers/apache/kafka/docs/connections/kafka.rst +++ b/providers/apache/kafka/docs/connections/kafka.rst @@ -50,7 +50,8 @@ of parameters are described in the such as ``my_company.kafka.auth`` does not lead to authorization of the callables inside it. This is enforced for security reasons, to prevent malicious callbacks from being executed. The allowlist is empty by default, which disables string-valued callbacks entirely. - Managed authentication (Amazon MSK IAM, Google Managed Kafka) does not rely on this and is unaffected. + Automatically injected managed authentication (Amazon MSK IAM, Google Managed Kafka) does not + use the allowlist. An explicitly configured ``oauth_cb`` does. If you are defining the Airflow connection from the Airflow UI, the ``extra`` field will be renamed to ``Config Dict``. @@ -87,3 +88,54 @@ An example ``extra`` (``Config Dict``) for an MSK connection: An explicit ``oauth_cb`` provided in the connection configuration is always respected and is never overwritten by the automatic MSK IAM callback. + +Explicit MSK IAM callback with an AWS connection +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +To use credentials from a specific Airflow AWS connection instead of the signer's default credential +chain, set ``oauth_cb`` explicitly. Install the Amazon provider with its ``msk`` extra so the callback +and MSK IAM signer are available: + +.. code-block:: bash + + pip install 'apache-airflow-providers-amazon[msk]' + +Create an `Amazon Web Services connection +`_ +with the connection ID ``aws_msk_prod`` and the credentials or IAM role to use. Set its Extra to +specify the MSK cluster's region: + +.. code-block:: json + + {"region_name": "us-east-1"} + +Allow the public callback in the Airflow configuration: + +.. code-block:: ini + + [apache_kafka] + callback_allowlist = airflow.providers.amazon.aws.hooks.msk.oauth_cb + +Then create a Kafka connection with the following Extra (``Config Dict`` in the UI): + +.. code-block:: json + + { + "bootstrap.servers": "boot-abcde1.c2.kafka-serverless.us-east-1.amazonaws.com:9098", + "security.protocol": "SASL_SSL", + "sasl.mechanism": "OAUTHBEARER", + "group.id": "my-group", + "oauth_cb": "airflow.providers.amazon.aws.hooks.msk.oauth_cb", + "sasl.oauthbearer.config": "{\"aws_conn_id\":\"aws_msk_prod\"}" + } + +``confluent-kafka`` passes the ``sasl.oauthbearer.config`` string to ``oauth_cb(config_str)``. +The callback parses it as JSON and uses ``aws_conn_id`` to select the AWS connection. You can +also include ``region_name`` in that JSON string to override the region in the AWS connection. +The explicit callback is used instead of the automatic MSK IAM callback. + +.. warning:: + + The callback allowlist controls which function can be imported, but does not restrict the + ``aws_conn_id`` passed to it. Anyone who can edit this Kafka connection can change which AWS + connection the callback uses. Restrict access to the Kafka connection accordingly. From 20c2b4550d6319c0e74b8dcb11eda4aa379ca88b Mon Sep 17 00:00:00 2001 From: Arnold Lin Date: Fri, 18 Sep 2026 17:09:47 -0400 Subject: [PATCH 3/3] feat: Add public callback function for managed_kafka --- .../providers/apache/kafka/hooks/base.py | 27 ++++++------ .../unit/apache/kafka/hooks/test_base.py | 17 +++++++ .../docs/operators/cloud/managed_kafka.rst | 44 +++++++++++++++++++ .../google/cloud/hooks/managed_kafka.py | 18 ++++++++ .../google/cloud/hooks/test_managed_kafka.py | 34 +++++++++++++- 5 files changed, 126 insertions(+), 14 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 5132fa8a1a9ef..f4f3b8f2b5682 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 @@ -165,19 +165,20 @@ def _build_config(self) -> dict[str, Any]: and bootstrap_servers.find("cloud.goog") != -1 and bootstrap_servers.find("managedkafka") != -1 ): - try: - from airflow.providers.google.cloud.hooks.managed_kafka import ManagedKafkaHook - except ImportError: - from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException - - raise AirflowOptionalProviderFeatureException( - "Failed to import ManagedKafkaHook. For using this functionality google provider version " - ">= 14.1.0 should be pre-installed." - ) - self.log.info("Adding token generation for Google Auth to the confluent configuration.") - hook = ManagedKafkaHook() - token = hook.get_confluent_token - config.update({"oauth_cb": token}) + if "oauth_cb" not in config: + try: + from airflow.providers.google.cloud.hooks.managed_kafka import ManagedKafkaHook + except ImportError: + from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException + + raise AirflowOptionalProviderFeatureException( + "Failed to import ManagedKafkaHook. For using this functionality google provider version " + ">= 14.1.0 should be pre-installed." + ) + self.log.info("Adding token generation for Google Auth to the confluent configuration.") + hook = ManagedKafkaHook() + token = hook.get_confluent_token + config.update({"oauth_cb": token}) else: self._maybe_add_msk_iam_oauth(config, bootstrap_servers) return config 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 bedc253942cf1..5fcf10a7d1703 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 @@ -202,6 +202,23 @@ def test_test_connection_injects_managed_kafka_oauth( assert passed_config["oauth_cb"] is managed_kafka_hook.return_value.get_confluent_token assert connection == (True, "Connection successful.") + @conf_vars(CALLBACK_ALLOWLIST) + @mock.patch("airflow.providers.google.cloud.hooks.managed_kafka.ManagedKafkaHook", autospec=True) + @mock.patch(f"{BASEHOOK_PATCH_PATH}.get_connection") + def test_managed_kafka_preserves_explicit_oauth_cb(self, mock_get_connection, managed_kafka_hook, hook): + config_str = '{"gcp_conn_id":"google_kafka_prod"}' + mock_get_connection.return_value.extra_dejson = { + "bootstrap.servers": "bootstrap.my-cluster.us-central1.managedkafka.my-project.cloud.goog:9092", + "oauth_cb": "json.dumps", + "sasl.oauthbearer.config": config_str, + } + + config = hook.get_conn + + assert config["oauth_cb"] is json.dumps + assert config["sasl.oauthbearer.config"] == config_str + assert managed_kafka_hook.mock_calls == [] + @mock.patch(f"{BASEHOOK_PATCH_PATH}.get_connection") def test_get_conn_msk_iam_provisioned(self, mock_get_connection, hook): config = { diff --git a/providers/google/docs/operators/cloud/managed_kafka.rst b/providers/google/docs/operators/cloud/managed_kafka.rst index 69651de61bd6e..9db614a6d8861 100644 --- a/providers/google/docs/operators/cloud/managed_kafka.rst +++ b/providers/google/docs/operators/cloud/managed_kafka.rst @@ -177,6 +177,50 @@ To consume data from topic you can use :start-after: [START how_to_cloud_managed_kafka_consume_from_topic_operator] :end-before: [END how_to_cloud_managed_kafka_consume_from_topic_operator] +Configure an explicit OAuth callback +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +For a Google Managed Kafka bootstrap server, the Apache Kafka provider automatically uses +``google_cloud_default`` to generate an OAuth token. To authenticate with a different +:ref:`Google Cloud connection `, configure the +Google provider's public ``oauth_cb`` function explicitly instead. The Google Cloud principal +must have permission to connect to the cluster; see Google's +`SASL authentication guide `_. + +Add the callback's full import path to the Airflow configuration. String-valued callbacks are +disabled unless their exact paths are listed in ``[apache_kafka] callback_allowlist``: + +.. code-block:: ini + + [apache_kafka] + callback_allowlist = airflow.providers.google.cloud.hooks.managed_kafka.oauth_cb + +Create a :ref:`Kafka connection ` with the following Extra +(``Config Dict`` in the Airflow UI), replacing the bootstrap server and ``gcp_conn_id`` +with your values: + +.. code-block:: json + + { + "bootstrap.servers": "bootstrap.my-cluster.us-central1.managedkafka.my-project.cloud.goog:9092", + "security.protocol": "SASL_SSL", + "sasl.mechanisms": "OAUTHBEARER", + "group.id": "my-consumer-group", + "oauth_cb": "airflow.providers.google.cloud.hooks.managed_kafka.oauth_cb", + "sasl.oauthbearer.config": "{\"gcp_conn_id\":\"google_kafka_prod\"}" + } + +``sasl.oauthbearer.config`` is a JSON string inside the Kafka connection's Extra JSON. +``confluent-kafka`` passes it to ``oauth_cb(config_str)``, which selects the Google Cloud +connection named by ``gcp_conn_id``. Set the Kafka operator's ``kafka_config_id`` to this +Kafka connection ID. An explicit ``oauth_cb`` is preserved instead of the automatic callback. + +.. warning:: + + The callback allowlist restricts which function can be imported, but not which Google + Cloud connection it uses. Anyone who can edit this Kafka connection can change + ``gcp_conn_id``; restrict access to the connection accordingly. + Reference ^^^^^^^^^ diff --git a/providers/google/src/airflow/providers/google/cloud/hooks/managed_kafka.py b/providers/google/src/airflow/providers/google/cloud/hooks/managed_kafka.py index dafc4d56403ed..677b16285f2dc 100644 --- a/providers/google/src/airflow/providers/google/cloud/hooks/managed_kafka.py +++ b/providers/google/src/airflow/providers/google/cloud/hooks/managed_kafka.py @@ -705,3 +705,21 @@ def delete_consumer_group( timeout=timeout, metadata=metadata, ) + + +def oauth_cb(config_str: str) -> tuple[str, float]: + """Generate an authentication token for a ``confluent_kafka`` OAuth callback.""" + try: + option = json.loads(config_str or "{}") + except json.JSONDecodeError: + raise ValueError("Invalid JSON in config_str") + + if not isinstance(option, dict): + raise ValueError("config_str must contain a JSON object") + + gcp_conn_id = option.get("gcp_conn_id") + if not gcp_conn_id: + raise ValueError("Missing 'gcp_conn_id' in config_str") + + hook = ManagedKafkaHook(gcp_conn_id=gcp_conn_id) + return hook.get_confluent_token(config_str) diff --git a/providers/google/tests/unit/google/cloud/hooks/test_managed_kafka.py b/providers/google/tests/unit/google/cloud/hooks/test_managed_kafka.py index 4aeff93ef62ad..b3018d106452a 100644 --- a/providers/google/tests/unit/google/cloud/hooks/test_managed_kafka.py +++ b/providers/google/tests/unit/google/cloud/hooks/test_managed_kafka.py @@ -17,11 +17,13 @@ # under the License. from __future__ import annotations +import json from unittest import mock +import pytest from google.api_core.gapic_v1.method import DEFAULT -from airflow.providers.google.cloud.hooks.managed_kafka import ManagedKafkaHook +from airflow.providers.google.cloud.hooks.managed_kafka import ManagedKafkaHook, oauth_cb from unit.google.cloud.utils.base_gcp_mock import ( mock_base_gcp_hook_default_project_id, @@ -397,6 +399,36 @@ def test_list_consumer_groups(self, mock_client) -> None: ) +class TestOauthCallback: + @mock.patch(MANAGED_KAFKA_STRING.format("ManagedKafkaHook"), autospec=True) + def test_uses_configured_gcp_connection(self, mock_hook): + config_str = json.dumps({"gcp_conn_id": TEST_GCP_CONN_ID}) + expected_token = ("token", 1_700_000_900.0) + mock_hook.return_value.get_confluent_token.return_value = expected_token + + assert oauth_cb(config_str) == expected_token + assert mock_hook.mock_calls == [ + mock.call(gcp_conn_id=TEST_GCP_CONN_ID), + mock.call().get_confluent_token(config_str), + ] + + @pytest.mark.parametrize( + ("config_str", "error"), + [ + ("invalid", "Invalid JSON in config_str"), + ("[]", "config_str must contain a JSON object"), + ("", "Missing 'gcp_conn_id' in config_str"), + ('{"gcp_conn_id":null}', "Missing 'gcp_conn_id' in config_str"), + ], + ) + @mock.patch(MANAGED_KAFKA_STRING.format("ManagedKafkaHook"), autospec=True) + def test_rejects_invalid_config(self, mock_hook, config_str, error): + with pytest.raises(ValueError, match=error): + oauth_cb(config_str) + + assert mock_hook.mock_calls == [] + + class TestManagedKafkaWithoutDefaultProjectIdHook: def setup_method(self): with mock.patch(