Add Amazon MSK IAM (OAUTHBEARER) support to Apache Kafka provider#69427
Conversation
|
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
|
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>
|
Awesome work, congrats on your first merged pull request! You are invited to check our Issue Tracker for additional contributions. |
|
@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. Handling it here: #69563 |
|
Thanks for catching this and handling it in #69563, and apologies for the miss — I didn't realize the extras table in |
Problem
The Apache Kafka provider already injects an
oauth_cbtoken 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.comor*.kafka-serverless.<region>.amazonaws.com, including.cnregions) andsasl.mechanismisOAUTHBEARER, the hook:oauth_cbbacked byaws_msk_iam_sasl_signer.MSKAuthTokenProvider.generate_auth_token(region), converting the signer's millisecond expiry to the seconds expected by confluent-kafka;oauth_cb;AirflowOptionalProviderFeatureExceptionwith an install hint when the signer library is missing.A new
mskoptional dependency (aws-msk-iam-sasl-signer-python>=1.0.1) is added, mirroring the existinggoogleextra, anduv.lockis regenerated accordingly.This mirrors the pattern established for Google Managed Kafka in #47056 and #48926.
Tests + Docs
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).Notes for reviewers
...amazonaws.com.example.com) cannot cause the hook to sign an IAM token for an untrusted broker.KafkaBaseHook.test_connection()builds anAdminClientdirectly 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.