diff --git a/providers/apache/kafka/docs/connections/kafka.rst b/providers/apache/kafka/docs/connections/kafka.rst index 300499a1e805e..809541166156d 100644 --- a/providers/apache/kafka/docs/connections/kafka.rst +++ b/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 `_ clusters (both provisioned and serverless) can be +authenticated with `IAM `_. +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..amazonaws.com`` or ``*.kafka-serverless..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. diff --git a/providers/apache/kafka/pyproject.toml b/providers/apache/kafka/pyproject.toml index f65b494392557..b43f5a60790fa 100644 --- a/providers/apache/kafka/pyproject.toml +++ b/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" ] 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 4cc483d973e6f..b3b76ff1421d5 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 @@ -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[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: 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 be790fbab5061..33a75d536a0ef 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 @@ -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 diff --git a/uv.lock b/uv.lock index 11bbace1114d2..90e6bdcdf6d7e 100644 --- a/uv.lock +++ b/uv.lock @@ -3726,6 +3726,9 @@ common-messaging = [ google = [ { name = "apache-airflow-providers-google" }, ] +msk = [ + { name = "aws-msk-iam-sasl-signer-python" }, +] [package.dev-dependencies] dev = [ @@ -3749,10 +3752,11 @@ requires-dist = [ { name = "apache-airflow-providers-google", marker = "extra == 'google'", editable = "providers/google" }, { name = "asgiref", marker = "python_full_version < '3.14'", specifier = ">=2.3.0" }, { name = "asgiref", marker = "python_full_version >= '3.14'", specifier = ">=3.11.1" }, + { name = "aws-msk-iam-sasl-signer-python", marker = "extra == 'msk'", specifier = ">=1.0.1" }, { name = "confluent-kafka", marker = "python_full_version < '3.14'", specifier = ">=2.6.0" }, { name = "confluent-kafka", marker = "python_full_version >= '3.14'", specifier = ">=2.13.2" }, ] -provides-extras = ["google", "common-messaging"] +provides-extras = ["google", "msk", "common-messaging"] [package.metadata.requires-dev] dev = [ @@ -9440,6 +9444,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fb/95/adcb68e20c34162e9135f370d6e31737719c2b6f94bc953fe7ed1f10fe21/authlib-1.7.2-py2.py3-none-any.whl", hash = "sha256:3e1faedc9d87e7d56a164eca3ccb6ace0d61b94abe83e92242f8dc8bba9b4a9f", size = 259548, upload-time = "2026-05-06T08:10:21.436Z" }, ] +[[package]] +name = "aws-msk-iam-sasl-signer-python" +version = "1.0.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "boto3" }, + { name = "botocore" }, + { name = "click" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/38/8b/9af0a7def4ba357afadc89c06019d3735944cb3d6065a455f41580ba7fd6/aws_msk_iam_sasl_signer_python-1.0.2.tar.gz", hash = "sha256:3432d88a7c6db4887ceb1130ebaed0113bfda48b79ea811537add5f1d25fa13f", size = 24034, upload-time = "2025-03-05T20:49:48.582Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/06/43/f3ffd79fc4941b8d610530d81e1f600cfa1b8651affc25cb929fd0f2f2c9/aws_msk_iam_sasl_signer_python-1.0.2-py2.py3-none-any.whl", hash = "sha256:310eb2db9ca0ff55ed06a24212739b87533e7f1cf6f34e43aabbd97a3b21290e", size = 13279, upload-time = "2025-03-05T20:49:46.611Z" }, +] + [[package]] name = "aws-sam-translator" version = "1.110.0"