This is an automated email from the ASF dual-hosted git repository.
jason810496 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 114d19c0aab Add `MskHook` (#69000)
114d19c0aab is described below
commit 114d19c0aab19c715d2f4d23527b17b73e9a1dc7
Author: Yi Lin <[email protected]>
AuthorDate: Wed Sep 16 01:24:10 2026 -0400
Add `MskHook` (#69000)
* feat: Add Hook for Amazon MSK
* feat: Add confluenttoken for Amazon MSK
* test: Update unit test for MskHook
---------
Co-authored-by: Arnold Lin <[email protected]>
---
providers/amazon/README.rst | 1 +
providers/amazon/docs/index.rst | 1 +
providers/amazon/provider.yaml | 7 ++-
providers/amazon/pyproject.toml | 3 +
.../src/airflow/providers/amazon/aws/hooks/msk.py | 66 ++++++++++++++++++++++
.../airflow/providers/amazon/get_provider_info.py | 6 ++
.../amazon/tests/unit/amazon/aws/hooks/test_msk.py | 63 +++++++++++++++++++++
uv.lock | 6 +-
8 files changed, 151 insertions(+), 2 deletions(-)
diff --git a/providers/amazon/README.rst b/providers/amazon/README.rst
index fe35a493b0b..d6678558c10 100644
--- a/providers/amazon/README.rst
+++ b/providers/amazon/README.rst
@@ -111,6 +111,7 @@ Extra Dependencies
``aiobotocore`` ``aiobotocore>=3.0.0``
``cncf.kubernetes`` ``apache-airflow-providers-cncf-kubernetes>=7.2.0``
``s3fs`` ``s3fs>=2023.10.0``
+``msk`` ``aws-msk-iam-sasl-signer-python>=1.0.1``
``python3-saml`` ``python3-saml>=1.16.0; python_version < '3.13'``,
``xmlsec>=1.3.14; python_version < '3.13'``, ``lxml>=6.0.0; python_version <
'3.13'``
``apache.hive`` ``apache-airflow-providers-apache-hive``
``exasol`` ``apache-airflow-providers-exasol``
diff --git a/providers/amazon/docs/index.rst b/providers/amazon/docs/index.rst
index 3510b1d3784..f134f522e99 100644
--- a/providers/amazon/docs/index.rst
+++ b/providers/amazon/docs/index.rst
@@ -178,6 +178,7 @@ Extra Dependencies
``aiobotocore`` ``aiobotocore>=3.0.0``
``cncf.kubernetes`` ``apache-airflow-providers-cncf-kubernetes>=7.2.0``
``s3fs`` ``s3fs>=2023.10.0``
+``msk`` ``aws-msk-iam-sasl-signer-python>=1.0.1``
``python3-saml`` ``python3-saml>=1.16.0; python_version < '3.13'``,
``xmlsec>=1.3.14; python_version < '3.13'``, ``lxml>=6.0.0; python_version <
'3.13'``
``apache.hive`` ``apache-airflow-providers-apache-hive``
``exasol`` ``apache-airflow-providers-exasol``
diff --git a/providers/amazon/provider.yaml b/providers/amazon/provider.yaml
index b7731887e17..97c218c2f84 100644
--- a/providers/amazon/provider.yaml
+++ b/providers/amazon/provider.yaml
@@ -418,6 +418,9 @@ integrations:
how-to-guide:
- /docs/apache-airflow-providers-amazon/operators/neptune_analytics.rst
tags: [aws]
+ - integration-name: Amazon MSK
+ external-doc-url: https://aws.amazon.com/msk/
+ tags: [aws]
- integration-name: Amazon S3 Vectors
external-doc-url:
https://docs.aws.amazon.com/AmazonS3/latest/userguide/s3-vectors.html
logo:
/docs/integration-logos/[email protected]
@@ -808,7 +811,9 @@ hooks:
python-modules:
- airflow.providers.amazon.aws.hooks.neptune
- airflow.providers.amazon.aws.hooks.neptune_analytics
-
+ - integration-name: Amazon MSK
+ python-modules:
+ - airflow.providers.amazon.aws.hooks.msk
bundles:
- integration-name: Amazon Simple Storage Service (S3)
diff --git a/providers/amazon/pyproject.toml b/providers/amazon/pyproject.toml
index bfb06b198d8..f59258095ee 100644
--- a/providers/amazon/pyproject.toml
+++ b/providers/amazon/pyproject.toml
@@ -97,6 +97,9 @@ dependencies = [
"s3fs" = [
"s3fs>=2023.10.0",
]
+"msk" = [
+ "aws-msk-iam-sasl-signer-python>=1.0.1",
+]
"python3-saml" = [
# Python 3 saml is not compatible with Python 3.13 yet, so we pin it to <
3.13
"python3-saml>=1.16.0; python_version < '3.13'",
diff --git a/providers/amazon/src/airflow/providers/amazon/aws/hooks/msk.py
b/providers/amazon/src/airflow/providers/amazon/aws/hooks/msk.py
new file mode 100644
index 00000000000..e86404b5f06
--- /dev/null
+++ b/providers/amazon/src/airflow/providers/amazon/aws/hooks/msk.py
@@ -0,0 +1,66 @@
+#
+# 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.
+"""This module contains Amazon Managed Streaming for Apache Kafka hook."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from botocore.credentials import CredentialProvider
+
+from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+
+if TYPE_CHECKING:
+ from botocore.credentials import Credentials
+
+
+class _MskCredentialsProvider(CredentialProvider):
+ METHOD = "airflow"
+
+ def __init__(self, hook: MskHook) -> None:
+ self.hook = hook
+
+ def load(self) -> Credentials | None:
+ return
self.hook.get_session(region_name=self.hook.region_name).get_credentials()
+
+
+class MskHook(AwsBaseHook):
+ """
+ Create an authenticated Amazon Managed Streaming for Apache Kafka client.
+
+ Additional arguments, such as ``aws_conn_id`` or ``region_name``, are
passed
+ to :class:`~airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook`. Use
+ ``conn`` or ``get_conn()`` to access the underlying
``boto3.client("kafka")``
+ directly.
+ """
+
+ def __init__(self, *args, **kwargs) -> None:
+ kwargs["client_type"] = "kafka"
+ super().__init__(*args, **kwargs)
+
+ def confluent_token(self, config_str: str) -> tuple[str, float]:
+ """Generate an Amazon MSK IAM token for a ``confluent_kafka`` OAuth
callback."""
+ if not self.region_name:
+ raise ValueError("AWS region is required to generate an Amazon MSK
IAM token")
+
+ from aws_msk_iam_sasl_signer import MSKAuthTokenProvider
+
+ token, expiry_ms =
MSKAuthTokenProvider.generate_auth_token_from_credentials_provider(
+ self.region_name, _MskCredentialsProvider(self)
+ )
+ return token, expiry_ms / 1000
diff --git a/providers/amazon/src/airflow/providers/amazon/get_provider_info.py
b/providers/amazon/src/airflow/providers/amazon/get_provider_info.py
index 1038b5d698f..0a093ba5b5c 100644
--- a/providers/amazon/src/airflow/providers/amazon/get_provider_info.py
+++ b/providers/amazon/src/airflow/providers/amazon/get_provider_info.py
@@ -380,6 +380,11 @@ def get_provider_info():
"how-to-guide":
["/docs/apache-airflow-providers-amazon/operators/neptune_analytics.rst"],
"tags": ["aws"],
},
+ {
+ "integration-name": "Amazon MSK",
+ "external-doc-url": "https://aws.amazon.com/msk/";,
+ "tags": ["aws"],
+ },
{
"integration-name": "Amazon S3 Vectors",
"external-doc-url":
"https://docs.aws.amazon.com/AmazonS3/latest/userguide/s3-vectors.html";,
@@ -898,6 +903,7 @@ def get_provider_info():
"airflow.providers.amazon.aws.hooks.neptune_analytics",
],
},
+ {"integration-name": "Amazon MSK", "python-modules":
["airflow.providers.amazon.aws.hooks.msk"]},
],
"bundles": [
{
diff --git a/providers/amazon/tests/unit/amazon/aws/hooks/test_msk.py
b/providers/amazon/tests/unit/amazon/aws/hooks/test_msk.py
new file mode 100644
index 00000000000..846e162587a
--- /dev/null
+++ b/providers/amazon/tests/unit/amazon/aws/hooks/test_msk.py
@@ -0,0 +1,63 @@
+#
+# 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.
+from __future__ import annotations
+
+from unittest import mock
+
+import pytest
+from botocore.credentials import CredentialProvider
+
+from airflow.providers.amazon.aws.hooks.msk import MskHook
+
+MOCK_MSK_SIGNER_MODULE = mock.MagicMock()
+
+
+class TestMskHook:
+ def setup_method(self):
+ self.hook = MskHook(aws_conn_id="aws_msk", region_name="us-east-1")
+
+ def test_init(self):
+ assert self.hook.aws_conn_id == "aws_msk"
+ assert self.hook.client_type == "kafka"
+
+ @mock.patch.dict("sys.modules", {"aws_msk_iam_sasl_signer":
MOCK_MSK_SIGNER_MODULE})
+ @mock.patch.object(MskHook, "get_session")
+ def test_confluent_token(self, mock_get_session):
+ credentials =
mock_get_session.return_value.get_credentials.return_value
+ mock_generate_auth_token = (
+
MOCK_MSK_SIGNER_MODULE.MSKAuthTokenProvider.generate_auth_token_from_credentials_provider
+ )
+ mock_generate_auth_token.reset_mock()
+ mock_generate_auth_token.return_value = ("token", 1_700_000_900_000)
+
+ token, expiry = self.hook.confluent_token("")
+
+ region, credentials_provider = mock_generate_auth_token.call_args.args
+ assert region == "us-east-1"
+ assert isinstance(credentials_provider, CredentialProvider)
+ assert credentials_provider.load() is credentials
+ mock_get_session.assert_called_once_with(region_name="us-east-1")
+ assert token == "token"
+ assert expiry == 1_700_000_900.0
+
+ @mock.patch.object(MskHook, "region_name", new_callable=mock.PropertyMock)
+ def test_confluent_token_requires_region(self, mock_region_name):
+ mock_region_name.return_value = None
+
+ with pytest.raises(ValueError, match="AWS region is required"):
+ self.hook.confluent_token("")
diff --git a/uv.lock b/uv.lock
index 7eea41e3c96..0d0476a4373 100644
--- a/uv.lock
+++ b/uv.lock
@@ -3142,6 +3142,9 @@ microsoft-azure = [
mongo = [
{ name = "apache-airflow-providers-mongo" },
]
+msk = [
+ { name = "aws-msk-iam-sasl-signer-python" },
+]
openlineage = [
{ name = "apache-airflow-providers-openlineage" },
]
@@ -3231,6 +3234,7 @@ requires-dist = [
{ name = "apache-airflow-providers-standard", marker = "extra ==
'standard'", editable = "providers/standard" },
{ 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 = "boto3", specifier = ">=1.41.0" },
{ name = "botocore", specifier = ">=1.41.0" },
{ name = "inflection", specifier = ">=0.5.1" },
@@ -3251,7 +3255,7 @@ requires-dist = [
{ name = "watchtower", specifier = ">=3.3.1,<4" },
{ name = "xmlsec", marker = "python_full_version < '3.13' and extra ==
'python3-saml'", specifier = ">=1.3.14" },
]
-provides-extras = ["aiobotocore", "cncf-kubernetes", "s3fs", "python3-saml",
"apache-hive", "exasol", "fab", "ftp", "google", "imap", "microsoft-azure",
"mongo", "pandas", "openlineage", "salesforce", "ssh", "standard",
"common-messaging", "sqlalchemy"]
+provides-extras = ["aiobotocore", "cncf-kubernetes", "s3fs", "msk",
"python3-saml", "apache-hive", "exasol", "fab", "ftp", "google", "imap",
"microsoft-azure", "mongo", "pandas", "openlineage", "salesforce", "ssh",
"standard", "common-messaging", "sqlalchemy"]
[package.metadata.requires-dev]
dev = [
