Skip to content

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

Merged
vincbeck merged 2 commits into
apache:mainfrom
TimurRakhmatullin86:feature/kafka-msk-iam-auth
Jul 7, 2026
Merged

Add Amazon MSK IAM (OAUTHBEARER) support to Apache Kafka provider#69427
vincbeck merged 2 commits into
apache:mainfrom
TimurRakhmatullin86:feature/kafka-msk-iam-auth

Conversation

@TimurRakhmatullin86

Copy link
Copy Markdown

Problem

The Apache Kafka provider already injects an oauth_cb token callback for Google Managed Kafka, detected from the bootstrap servers hostname (managedkafka…cloud.goog). Amazon MSK offers the same OAUTHBEARER-based IAM authentication, and MSK Serverless supports IAM only — but there was no equivalent branch. Users had to wire the MSK IAM token signer by hand for every connection, with no supported extra to pull the dependency.

Solution

Add a symmetric AWS branch to KafkaBaseHook.get_conn. When the bootstrap servers point at an Amazon MSK endpoint (*.kafka.<region>.amazonaws.com or *.kafka-serverless.<region>.amazonaws.com, including .cn regions) and sasl.mechanism is OAUTHBEARER, the hook:

  • derives the AWS region from the hostname (normalised to lower case for the SigV4 scope);
  • installs an oauth_cb backed by aws_msk_iam_sasl_signer.MSKAuthTokenProvider.generate_auth_token(region), converting the signer's millisecond expiry to the seconds expected by confluent-kafka;
  • never overwrites an explicit user-provided oauth_cb;
  • raises AirflowOptionalProviderFeatureException with an install hint when the signer library is missing.

A new msk optional dependency (aws-msk-iam-sasl-signer-python>=1.0.1) is added, mirroring the existing google extra, and uv.lock is regenerated accordingly.

This mirrors the pattern established for Google Managed Kafka in #47056 and #48926.

Tests + Docs

  • Unit tests cover provisioned/serverless detection, region parsing/normalisation, the no-injection paths (non-MSK host, non-OAUTHBEARER mechanism), preservation of a user oauth_cb, the millisecond→second conversion in the callback, the missing-library error, and a parametrised check of the host-matching regex (including a look-alike host that must not match).
  • New "Amazon MSK with IAM authentication" section in the Kafka connection docs.

Notes for reviewers

  • The MSK host pattern is anchored to a host boundary so a look-alike hostname (e.g. ...amazonaws.com.example.com) cannot cause the hook to sign an IAM token for an untrusted broker.
  • Token generation uses the standard AWS credential chain (env / files / instance role), mirroring how the GCP branch relies on ambient credentials. Profile / role-ARN variants were intentionally left out to keep the change minimal and symmetric; happy to add them if preferred.
  • KafkaBaseHook.test_connection() builds an AdminClient directly from the connection extras and bypasses both the GCP and MSK injections. This is pre-existing behaviour and is left as a possible follow-up rather than expanded here.

@boring-cyborg

boring-cyborg Bot commented Jul 6, 2026

Copy link
Copy Markdown

Congratulations on your first Pull Request and welcome to the Apache Airflow community! If you have any issues or are unsure about any anything please check our Contributors' Guide
Here are some useful points:

  • Pay attention to the quality of your code (ruff, mypy and type annotations). Our prek-hooks will help you with that.
  • In case of a new feature add useful documentation (in docstrings or in docs/ directory). Adding a new operator? Check this short guide Consider adding an example Dag that shows how users should use it.
  • Consider using Breeze environment for testing locally, it's a heavy docker but it ships with a working Airflow and a lot of integrations.
  • Be patient and persistent. It might take some time to get a review or get the final approval from Committers.
  • Please follow ASF Code of Conduct for all communication including (but not limited to) comments on Pull Requests, Mailing list and Slack.
  • Be sure to read the Airflow Coding style.
  • Always keep your Pull Requests rebased, otherwise your build might fail due to changes not related to your commits.
    Apache Airflow is a community-driven project and together we are making it better 🚀.
    In case of doubts contact the developers at:
    Mailing List: dev@airflow.apache.org
    Slack: https://s.apache.org/airflow-slack

The Kafka provider already injects an ``oauth_cb`` token callback for Google
Managed Kafka based on the bootstrap servers hostname. Amazon MSK offers the
same OAUTHBEARER-based IAM authentication (and MSK Serverless supports IAM
only), but there was no equivalent branch, forcing users to wire the token
signer by hand.

This adds a symmetric branch: when the bootstrap servers point at an Amazon
MSK endpoint (``*.kafka.<region>.amazonaws.com`` or
``*.kafka-serverless.<region>.amazonaws.com``, including ``.cn`` regions) and
``sasl.mechanism`` is ``OAUTHBEARER``, the hook derives the AWS region from the
hostname and installs an ``oauth_cb`` backed by
``aws_msk_iam_sasl_signer.MSKAuthTokenProvider``. The callback converts the
signer's millisecond expiry to the seconds expected by confluent-kafka. An
explicit user-provided ``oauth_cb`` is never overwritten, and a missing signer
library raises ``AirflowOptionalProviderFeatureException`` pointing at the new
``msk`` extra.

Adds the ``msk`` optional dependency, unit tests, and connection docs.
Signed-off-by: Timur Rakhmatullin <174210871+TimurRakhmatullin86@users.noreply.github.com>
@vincbeck
vincbeck merged commit ca967b9 into apache:main Jul 7, 2026
111 checks passed
@boring-cyborg

boring-cyborg Bot commented Jul 7, 2026

Copy link
Copy Markdown

Awesome work, congrats on your first merged pull request! You are invited to check our Issue Tracker for additional contributions.

@Dev-iL

Dev-iL commented Jul 8, 2026

Copy link
Copy Markdown
Collaborator

@vincbeck Looks like the index.rst wasn't regenerated as part of this PR, and it led to builds breaking later on. We should look into why the hook didn't flag it here.

All changes made by hooks:
diff --git a/providers/apache/kafka/docs/index.rst b/providers/apache/kafka/docs/index.rst
index b3fc33f..cee81d3 100644
--- a/providers/apache/kafka/docs/index.rst
+++ b/providers/apache/kafka/docs/index.rst
@@ -149,6 +149,7 @@ Install them when installing from PyPI. For example:
 Extra                 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``
 ====================  ====================================================

Handling it here: #69563

@TimurRakhmatullin86

Copy link
Copy Markdown
Author

Thanks for catching this and handling it in #69563, and apologies for the miss — I didn't realize the extras table in index.rst is generated. I'll run the provider docs regeneration whenever a PR touches provider dependencies/extras going forward.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants