Add oauth callback function for MSK/Managed_Kafka

providers/amazon/docs/index.rst

@@ -42,6 +42,7 @@
     Deferrable Operators <deferrable>
     Secrets backends <secrets-backends/index>
     Logging for Tasks <logging/index>
+    Amazon MSK <msk>
     Configuration <configurations-ref>
     Executors <executors/index>
     Message Queues <message-queues/index>

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 <https://aws.amazon.com/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
+<https://airflow.apache.org/docs/apache-airflow-providers-apache-kafka/stable/message-queues/index.html>`_
+for an ``AssetWatcher`` example.
+
+Configure the Kafka connection
+------------------------------
+
+The `Apache Kafka connection guide
+<https://airflow.apache.org/docs/apache-airflow-providers-apache-kafka/stable/connections/kafka.html>`_
+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.

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)

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)

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
+<https://airflow.apache.org/docs/apache-airflow-providers-amazon/stable/connections/aws.html>`_
+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.

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

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 = {

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 <howto/connection:google_cloud_platform>`, 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 <https://docs.cloud.google.com/managed-service-for-apache-kafka/docs/authentication-kafka>`_.
+
+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 <howto/connection:kafka>` 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
 ^^^^^^^^^
 

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)

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(