Add Amazon MSK IAM (OAUTHBEARER) support to Apache Kafka provider

providers/apache/kafka/docs/connections/kafka.rst

@@ -43,3 +43,35 @@ of parameters are described in the
 If you are defining the Airflow connection from the Airflow UI, the ``extra`` field will be renamed to ``Config Dict``.
 
 Most operators and hooks will check that at the minimum the ``bootstrap.servers`` key exists and has a value set to be valid.
+
+Amazon MSK with IAM authentication
+----------------------------------
+
+`Amazon MSK <https://aws.amazon.com/msk/>`_ clusters (both provisioned and serverless) can be
+authenticated with `IAM <https://docs.aws.amazon.com/msk/latest/developerguide/iam-access-control.html>`_.
+This requires the ``aws-msk-iam-sasl-signer-python`` package, which is installed with the ``msk`` extra:
+
+.. code-block:: bash
+
+    pip install apache-airflow-providers-apache-kafka[msk]
+
+When the ``bootstrap.servers`` point at an Amazon MSK endpoint (for example
+``*.kafka.<region>.amazonaws.com`` or ``*.kafka-serverless.<region>.amazonaws.com``) and
+``sasl.mechanism`` is set to ``OAUTHBEARER``, the hook automatically generates and refreshes the
+IAM authentication token, deriving the AWS region from the bootstrap servers. The credentials are
+resolved by the signer using the standard AWS credential provider chain (environment variables,
+shared config/credentials files, instance/task IAM roles, etc.).
+
+An example ``extra`` (``Config Dict``) for an MSK connection:
+
+.. 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"
+    }
+
+An explicit ``oauth_cb`` provided in the connection configuration is always respected and is never
+overwritten by the automatic MSK IAM callback.

providers/apache/kafka/pyproject.toml

@@ -73,6 +73,9 @@ dependencies = [
 "google" = [
     "apache-airflow-providers-google"
 ]
+"msk" = [
+    "aws-msk-iam-sasl-signer-python>=1.0.1"
+]
 "common.messaging" = [
     "apache-airflow-providers-common-messaging>=2.0.0"
 ]

providers/apache/kafka/src/airflow/providers/apache/kafka/hooks/base.py

@@ -16,13 +16,43 @@
 # under the License.
 from __future__ import annotations
 
-from functools import cached_property
+import re
+from functools import cached_property, partial
 from typing import Any
 
 from confluent_kafka.admin import AdminClient
 
 from airflow.providers.common.compat.sdk import BaseHook
 
+# 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)
+# China regions use the ``.amazonaws.com.cn`` suffix. The region is captured so it
+# can be forwarded to the MSK IAM token signer.
+MSK_BOOTSTRAP_SERVERS_REGEX = re.compile(
+    r"\.kafka(?:-serverless)?\.(?P<region>[a-z0-9-]+)\.amazonaws\.com(?:\.cn)?(?::\d+)?(?=$|[,\s])",
+    re.IGNORECASE,
+)
+
+
+def _msk_iam_oauth_cb(region: str, config_str: str) -> tuple[str, float]:
+    """
+    Generate an OAUTHBEARER token for Amazon MSK IAM authentication.
+
+    This is used as the ``oauth_cb`` callback for ``confluent-kafka``. The library
+    passes the value of ``sasl.oauthbearer.config`` as ``config_str``; it is not
+    needed to sign an MSK IAM token, so it is ignored.
+
+    :param region: The AWS region of the MSK cluster.
+    :param config_str: The ``sasl.oauthbearer.config`` value passed by librdkafka.
+    """
+    from aws_msk_iam_sasl_signer import MSKAuthTokenProvider
+
+    token, expiry_ms = MSKAuthTokenProvider.generate_auth_token(region)
+    # The signer returns the expiry as milliseconds since the epoch while
+    # confluent-kafka expects seconds since the epoch.
+    return token, expiry_ms / 1000
+
 
 class KafkaBaseHook(BaseHook):
     """
@@ -82,8 +112,50 @@ def get_conn(self) -> Any:
             hook = ManagedKafkaHook()
             token = hook.get_confluent_token
             config.update({"oauth_cb": token})
+        else:
+            self._maybe_add_msk_iam_oauth(config, bootstrap_servers)
         return self._get_client(config)
 
+    def _maybe_add_msk_iam_oauth(self, config: dict[str, Any], bootstrap_servers: str | None) -> None:
+        """
+        Inject an OAUTHBEARER token callback for Amazon MSK IAM authentication.
+
+        The callback is only added when the bootstrap servers point at an Amazon MSK
+        cluster and the connection is configured to use the ``OAUTHBEARER`` SASL
+        mechanism. An explicit user-provided ``oauth_cb`` is never overwritten.
+        """
+        if not bootstrap_servers:
+            return
+
+        sasl_mechanism = config.get("sasl.mechanism") or config.get("sasl.mechanisms")
+        if sasl_mechanism != "OAUTHBEARER":
+            return
+
+        match = MSK_BOOTSTRAP_SERVERS_REGEX.search(bootstrap_servers)
+        if not match:
+            return
+
+        if "oauth_cb" in config:
+            # Respect an explicit callback provided by the user.
+            return
+
+        try:
+            from aws_msk_iam_sasl_signer import MSKAuthTokenProvider  # noqa: F401
+        except ImportError:
+            from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+            raise AirflowOptionalProviderFeatureException(
+                "Failed to import aws_msk_iam_sasl_signer. To use Amazon MSK IAM authentication "
+                "install the 'msk' extra: pip install apache-airflow-providers-apache-kafka[msk]"
+            )
+
+        region = match.group("region").lower()
+        self.log.info(
+            "Adding token generation for Amazon MSK IAM (region %s) to the confluent configuration.",
+            region,
+        )
+        config.update({"oauth_cb": partial(_msk_iam_oauth_cb, region)})
+
     def test_connection(self) -> tuple[bool, str]:
         """Test Connectivity from the UI."""
         try:

providers/apache/kafka/tests/unit/apache/kafka/hooks/test_base.py

@@ -21,7 +21,12 @@
 
 import pytest
 
-from airflow.providers.apache.kafka.hooks.base import KafkaBaseHook
+from airflow.providers.apache.kafka.hooks.base import (
+    MSK_BOOTSTRAP_SERVERS_REGEX,
+    KafkaBaseHook,
+    _msk_iam_oauth_cb,
+)
+from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
 
 try:
     import importlib.util
@@ -93,3 +98,108 @@ def test_test_connection_exception(self, mock_get_connection, admin_client, hook
         admin_client.return_value.list_topics.side_effect = [ValueError("some error")]
         connection = hook.test_connection()
         assert connection == (False, "some error")
+
+    @mock.patch(f"{BASEHOOK_PATCH_PATH}.get_connection")
+    def test_get_conn_msk_iam_provisioned(self, mock_get_connection, hook):
+        config = {
+            "bootstrap.servers": "b-1.demo.abcde1.c2.kafka.us-east-1.amazonaws.com:9098",
+            "security.protocol": "SASL_SSL",
+            "sasl.mechanism": "OAUTHBEARER",
+        }
+        mock_get_connection.return_value.extra_dejson = config
+        with mock.patch.dict("sys.modules", {"aws_msk_iam_sasl_signer": MagicMock()}):
+            result = hook.get_conn
+        assert "oauth_cb" in result
+        assert result["oauth_cb"].func is _msk_iam_oauth_cb
+        assert result["oauth_cb"].args == ("us-east-1",)
+
+    @mock.patch(f"{BASEHOOK_PATCH_PATH}.get_connection")
+    def test_get_conn_msk_iam_serverless(self, mock_get_connection, hook):
+        config = {
+            "bootstrap.servers": "boot-abcde1.c2.kafka-serverless.eu-west-1.amazonaws.com:9098",
+            "security.protocol": "SASL_SSL",
+            "sasl.mechanism": "OAUTHBEARER",
+        }
+        mock_get_connection.return_value.extra_dejson = config
+        with mock.patch.dict("sys.modules", {"aws_msk_iam_sasl_signer": MagicMock()}):
+            result = hook.get_conn
+        assert "oauth_cb" in result
+        assert result["oauth_cb"].args == ("eu-west-1",)
+
+    @mock.patch(f"{BASEHOOK_PATCH_PATH}.get_connection")
+    def test_get_conn_regular_host_no_msk_injection(self, mock_get_connection, hook):
+        config = {
+            "bootstrap.servers": "localhost:9092",
+            "sasl.mechanism": "OAUTHBEARER",
+        }
+        mock_get_connection.return_value.extra_dejson = config
+        result = hook.get_conn
+        assert "oauth_cb" not in result
+
+    @mock.patch(f"{BASEHOOK_PATCH_PATH}.get_connection")
+    def test_get_conn_msk_host_without_oauthbearer_no_injection(self, mock_get_connection, hook):
+        config = {
+            "bootstrap.servers": "b-1.demo.abcde1.c2.kafka.us-east-1.amazonaws.com:9098",
+            "security.protocol": "SASL_SSL",
+            "sasl.mechanism": "SCRAM-SHA-512",
+        }
+        mock_get_connection.return_value.extra_dejson = config
+        result = hook.get_conn
+        assert "oauth_cb" not in result
+
+    @mock.patch(f"{BASEHOOK_PATCH_PATH}.get_connection")
+    def test_get_conn_msk_iam_does_not_override_user_oauth_cb(self, mock_get_connection, hook):
+        user_cb = MagicMock()
+        config = {
+            "bootstrap.servers": "b-1.demo.abcde1.c2.kafka.us-east-1.amazonaws.com:9098",
+            "sasl.mechanism": "OAUTHBEARER",
+            "oauth_cb": user_cb,
+        }
+        mock_get_connection.return_value.extra_dejson = config
+        result = hook.get_conn
+        assert result["oauth_cb"] is user_cb
+
+    @mock.patch(f"{BASEHOOK_PATCH_PATH}.get_connection")
+    def test_get_conn_msk_iam_missing_library(self, mock_get_connection, hook):
+        config = {
+            "bootstrap.servers": "b-1.demo.abcde1.c2.kafka.us-east-1.amazonaws.com:9098",
+            "sasl.mechanism": "OAUTHBEARER",
+        }
+        mock_get_connection.return_value.extra_dejson = config
+        with mock.patch.dict("sys.modules", {"aws_msk_iam_sasl_signer": None}):
+            with pytest.raises(AirflowOptionalProviderFeatureException, match="msk"):
+                _ = hook.get_conn
+
+    def test_msk_iam_oauth_cb_returns_seconds(self):
+        fake_signer = MagicMock()
+        fake_signer.MSKAuthTokenProvider.generate_auth_token.return_value = ("my-token", 1_700_000_900_000)
+        with mock.patch.dict("sys.modules", {"aws_msk_iam_sasl_signer": fake_signer}):
+            token, expiry = _msk_iam_oauth_cb("us-east-1", "")
+        fake_signer.MSKAuthTokenProvider.generate_auth_token.assert_called_once_with("us-east-1")
+        assert token == "my-token"
+        assert expiry == 1_700_000_900.0
+
+    @pytest.mark.parametrize(
+        ("bootstrap_servers", "expected_region"),
+        [
+            ("b-1.demo.abcde1.c2.kafka.us-east-1.amazonaws.com:9098", "us-east-1"),
+            ("boot-abcde1.c2.kafka-serverless.us-east-1.amazonaws.com:9098", "us-east-1"),
+            ("b-1.x.kafka.cn-north-1.amazonaws.com.cn:9098", "cn-north-1"),
+            # Hostnames are case-insensitive; the region must be normalised to lower case
+            # because the SigV4 credential scope requires it.
+            ("b-1.x.kafka.US-EAST-1.amazonaws.com:9098", "us-east-1"),
+            ("b1:9092,b2.kafka.us-west-2.amazonaws.com:9098", "us-west-2"),
+            # A look-alike host that merely embeds an MSK-shaped substring must not match,
+            # otherwise the hook would sign an IAM token for an untrusted broker.
+            ("b-1.x.kafka.us-east-1.amazonaws.com.evil.example.com:9092", None),
+            ("localhost:9092", None),
+            ("kafka.example.com:9092", None),
+        ],
+    )
+    def test_msk_bootstrap_servers_regex(self, bootstrap_servers, expected_region):
+        match = MSK_BOOTSTRAP_SERVERS_REGEX.search(bootstrap_servers)
+        if expected_region is None:
+            assert match is None
+        else:
+            assert match is not None
+            assert match.group("region").lower() == expected_region