OAuth Azure IMDB implementation

Author: emasab

CHANGELOG.md

@@ -1,5 +1,16 @@
 # Confluent's Python client for Apache Kafka
 
+## v2.12.0
+
+v2.12.0 is a feature release with the following enhancements:
+
+- OAuth/OIDC metadata based authentication with Azure IMDS (#).
+
+confluent-kafka-python v2.12.0 is based on librdkafka v2.12.0, see the
+[librdkafka release notes](https://github.com/confluentinc/librdkafka/releases/tag/v2.12.0)
+for a complete list of changes, enhancements, fixes and upgrade considerations.
+
+
 ## v2.11.1
 
 v2.11.1 is a maintenance release with the following fixes:

examples/README.md

@@ -18,6 +18,8 @@ The scripts in this directory provide various examples of using Confluent's Pyth
 Additional examples for [Confluent Cloud](https://www.confluent.io/confluent-cloud/):
 
 * [confluent_cloud.py](confluent_cloud.py): Produce messages to Confluent Cloud and then read them back again.
+* [oauth_oidc_ccloud_producer.py](oauth_oidc_ccloud_producer.py): Demonstrates OAuth/OIDC Authentication with Confluent Cloud (client credentials).
+* [oauth_oidc_ccloud_azure_imds_producer.py](oauth_oidc_ccloud_azure_imds_producer.py): Demonstrates OAuth/OIDC Authentication with Confluent Cloud (Azure IMDS metadata based authentication).
 * [confluentinc/examples](https://github.com/confluentinc/examples/tree/master/clients/cloud/python): Integration with Confluent Cloud and Confluent Cloud Schema Registry
 
 ## venv setup

examples/oauth_oidc_ccloud_azure_imds_producer.py

@@ -0,0 +1,116 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# Copyright 2025 Confluent Inc.
+#
+# Licensed 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 example use Azure IMDS for credential-less authentication
+# through to Schema Registry on Confluent Cloud
+
+import logging
+import argparse
+from confluent_kafka import Producer
+from confluent_kafka.serialization import StringSerializer
+
+
+def producer_config(args):
+    logger = logging.getLogger(__name__)
+    logger.setLevel(logging.DEBUG)
+    params = {
+        'bootstrap.servers': args.bootstrap_servers,
+        'security.protocol': 'SASL_SSL',
+        'sasl.mechanisms': 'OAUTHBEARER',
+        'sasl.oauthbearer.method': 'oidc',
+        'sasl.oauthbearer.metadata.authentication.type': 'azure_imds',
+        'sasl.oauthbearer.config': f'query={args.query}'
+    }
+    # These two parameters are only applicable when producing to
+    # confluent cloud where some sasl extensions are required.
+    if args.logical_cluster and args.identity_pool_id:
+        params['sasl.oauthbearer.extensions'] = 'logicalCluster=' + args.logical_cluster + \
+            ',identityPoolId=' + args.identity_pool_id
+
+    return params
+
+
+def delivery_report(err, msg):
+    """
+    Reports the failure or success of a message delivery.
+
+    Args:
+        err (KafkaError): The error that occurred on None on success.
+
+        msg (Message): The message that was produced or failed.
+
+    Note:
+        In the delivery report callback the Message.key() and Message.value()
+        will be the binary format as encoded by any configured Serializers and
+        not the same object that was passed to produce().
+        If you wish to pass the original object(s) for key and value to delivery
+        report callback we recommend a bound callback or lambda where you pass
+        the objects along.
+
+    """
+    if err is not None:
+        print('Delivery failed for User record {}: {}'.format(msg.key(), err))
+        return
+    print('User record {} successfully produced to {} [{}] at offset {}'.format(
+        msg.key(), msg.topic(), msg.partition(), msg.offset()))
+
+
+def main(args):
+    topic = args.topic
+    delimiter = args.delimiter
+    producer_conf = producer_config(args)
+    producer = Producer(producer_conf)
+    serializer = StringSerializer('utf_8')
+
+    print('Producing records to topic {}. ^C to exit.'.format(topic))
+    while True:
+        # Serve on_delivery callbacks from previous calls to produce()
+        producer.poll(0.0)
+        try:
+            msg_data = input(">")
+            msg = msg_data.split(delimiter)
+            if len(msg) == 2:
+                producer.produce(topic=topic,
+                                 key=serializer(msg[0]),
+                                 value=serializer(msg[1]),
+                                 on_delivery=delivery_report)
+            else:
+                producer.produce(topic=topic,
+                                 value=serializer(msg[0]),
+                                 on_delivery=delivery_report)
+        except KeyboardInterrupt:
+            break
+
+    print('\nFlushing {} records...'.format(len(producer)))
+    producer.flush()
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description="OAUTH example with client credentials grant")
+    parser.add_argument('-b', dest="bootstrap_servers", required=True,
+                        help="Bootstrap broker(s) (host[:port])")
+    parser.add_argument('-t', dest="topic", default="example_producer_oauth",
+                        help="Topic name")
+    parser.add_argument('-d', dest="delimiter", default="|",
+                        help="Key-Value delimiter. Defaults to '|'"),
+    parser.add_argument('--query', dest="query", required=True,
+                        help="Query parameters for Azure IMDS token endpoint")
+    parser.add_argument('--logical-cluster', dest="logical_cluster", required=False, help="Logical Cluster.")
+    parser.add_argument('--identity-pool-id', dest="identity_pool_id", required=False, help="Identity Pool ID.")
+
+    main(parser.parse_args())

examples/oauth_schema_registry.py

@@ -16,7 +16,7 @@
 # limitations under the License.
 
 # Examples of setting up Schema Registry with OAuth with static token,
-# Client Credentials, and custom functions
+# Client Credentials, Azure IMDS, and custom functions
 
 
 # CUSTOM OAuth configuration takes in a custom function, config for that
@@ -49,6 +49,16 @@ def main():
     client_credentials_oauth_sr_client = SchemaRegistryClient(client_credentials_oauth_config)
     print(client_credentials_oauth_sr_client.get_subjects())
 
+    azure_imds_oauth_config = {
+        'url': 'https://psrc-123456.us-east-1.aws.confluent.cloud',
+        'bearer.auth.credentials.source': 'OAUTHBEARER_AZURE_IMDS',
+        'bearer.auth.issuer.endpoint.query': 'resource=&api-version=&client_id=',
+        'bearer.auth.logical.cluster': 'lsrc-12345',
+        'bearer.auth.identity.pool.id': 'pool-abcd'}
+
+    azure_imds_oauth_sr_client = SchemaRegistryClient(azure_imds_oauth_config)
+    print(azure_imds_oauth_sr_client.get_subjects())
+
     def custom_oauth_function(config):
         return config
 

src/confluent_kafka/schema_registry/_async/schema_registry_client.py

@@ -23,6 +23,7 @@
 import ssl
 import time
 import urllib
+import abc
 from urllib.parse import unquote, urlparse
 
 import httpx
@@ -36,6 +37,7 @@
 from confluent_kafka.schema_registry.common._oauthbearer import (
     _BearerFieldProvider,
     _AbstractOAuthBearerOIDCFieldProviderBuilder,
+    _AbstractOAuthBearerOIDCAzureIMDSFieldProviderBuilder,
     _StaticOAuthBearerFieldProviderBuilder,
     _AbstractCustomOAuthBearerFieldProviderBuilder)
 from confluent_kafka.schema_registry.error import SchemaRegistryError, OAuthTokenError
@@ -76,18 +78,15 @@ async def get_bearer_fields(self) -> dict:
         return await self.custom_function(self.custom_config)
 
 
-class _AsyncOAuthClient(_BearerFieldProvider):
-    def __init__(self, client_id: str, client_secret: str, scope: str, token_endpoint: str, logical_cluster: str,
+class _AsyncAbstractOAuthClient(_BearerFieldProvider):
+    def __init__(self, logical_cluster: str,
                  identity_pool: str, max_retries: int, retries_wait_ms: int, retries_max_wait_ms: int):
-        self.token = None
-        self.logical_cluster = logical_cluster
-        self.identity_pool = identity_pool
-        self.client = AsyncOAuth2Client(client_id=client_id, client_secret=client_secret, scope=scope)
-        self.token_endpoint = token_endpoint
-        self.max_retries = max_retries
-        self.retries_wait_ms = retries_wait_ms
-        self.retries_max_wait_ms = retries_max_wait_ms
-        self.token_expiry_threshold = 0.8
+        self.logical_cluster: str = logical_cluster
+        self.identity_pool: str = identity_pool
+        self.max_retries: int = max_retries
+        self.retries_wait_ms: int = retries_wait_ms
+        self.retries_max_wait_ms: int = retries_max_wait_ms
+        self.token: str = None
 
     async def get_bearer_fields(self) -> dict:
         return {
@@ -96,21 +95,24 @@ async def get_bearer_fields(self) -> dict:
             'bearer.auth.identity.pool.id': self.identity_pool
         }
 
-    def token_expired(self) -> bool:
-        expiry_window = self.token['expires_in'] * self.token_expiry_threshold
-
-        return self.token['expires_at'] < time.time() + expiry_window
-
     async def get_access_token(self) -> str:
         if not self.token or self.token_expired():
             await self.generate_access_token()
 
-        return self.token['access_token']
+        return self.token
+
+    @abc.abstractmethod
+    def token_expired(self) -> bool:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def fetch_token(self) -> str:
+        raise NotImplementedError
 
     async def generate_access_token(self) -> None:
         for i in range(self.max_retries + 1):
             try:
-                self.token = await self.client.fetch_token(url=self.token_endpoint, grant_type='client_credentials')
+                self.token = await self.fetch_token()
                 return
             except Exception as e:
                 if i >= self.max_retries:
@@ -119,9 +121,51 @@ async def generate_access_token(self) -> None:
                 await asyncio.sleep(full_jitter(self.retries_wait_ms, self.retries_max_wait_ms, i) / 1000)
 
 
+class _AsyncOAuthClient(_AsyncAbstractOAuthClient):
+    def __init__(self, client_id: str, client_secret: str, scope: str, token_endpoint: str, logical_cluster: str,
+                 identity_pool: str, max_retries: int, retries_wait_ms: int, retries_max_wait_ms: int):
+        super().__init__(
+            logical_cluster, identity_pool, max_retries, retries_wait_ms,
+            retries_max_wait_ms)
+        self.client = AsyncOAuth2Client(client_id=client_id, client_secret=client_secret, scope=scope)
+        self.token_endpoint: str = token_endpoint
+        self.token_object: dict = None
+        self.token_expiry_threshold: float = 0.8
+
+    def token_expired(self) -> bool:
+        expiry_window = self.token_object['expires_in'] * self.token_expiry_threshold
+        return self.token_object['expires_at'] < time.time() + expiry_window
+
+    async def fetch_token(self) -> str:
+        self.token_object = await self.client.fetch_token(url=self.token_endpoint, grant_type='client_credentials')
+        return self.token_object['access_token']
+
+
+class _AsyncOAuthAzureIMDSClient(_AsyncAbstractOAuthClient):
+    def __init__(self, token_endpoint: str, logical_cluster: str,
+                 identity_pool: str, max_retries: int, retries_wait_ms: int, retries_max_wait_ms: int):
+        super().__init__(
+            logical_cluster, identity_pool, max_retries, retries_wait_ms,
+            retries_max_wait_ms)
+        self.client = httpx.AsyncClient()
+        self.token_endpoint: str = token_endpoint
+        self.token_object: dict = None
+        self.token_expiry_threshold: float = 0.8
+
+    def token_expired(self) -> bool:
+        expiry_window = int(self.token_object['expires_in']) * self.token_expiry_threshold
+        return int(self.token_object['expires_on']) < time.time() + expiry_window
+
+    async def fetch_token(self) -> str:
+        self.token_object = await self.client.get(self.token_endpoint, headers=[
+            ('Metadata', 'true')
+        ]).json()
+        return self.token_object['access_token']
+
+
 class _AsyncOAuthBearerOIDCFieldProviderBuilder(_AbstractOAuthBearerOIDCFieldProviderBuilder):
 
-    def build(self, max_retries, retries_wait_ms, retries_max_wait_ms):
+    def build(self, max_retries: int, retries_wait_ms: int, retries_max_wait_ms: int):
         self._validate()
         return _AsyncOAuthClient(
             self.client_id, self.client_secret, self.scope,
@@ -132,9 +176,21 @@ def build(self, max_retries, retries_wait_ms, retries_max_wait_ms):
             retries_max_wait_ms)
 
 
+class _AsyncOAuthBearerOIDCAzureIMDSFieldProviderBuilder(_AbstractOAuthBearerOIDCAzureIMDSFieldProviderBuilder):
+
+    def build(self, max_retries: int, retries_wait_ms: int, retries_max_wait_ms: int):
+        self._validate()
+        return _AsyncOAuthAzureIMDSClient(
+            self.token_endpoint,
+            self.logical_cluster,
+            self.identity_pool,
+            max_retries, retries_wait_ms,
+            retries_max_wait_ms)
+
+
 class _AsyncCustomOAuthBearerFieldProviderBuilder(_AbstractCustomOAuthBearerFieldProviderBuilder):
 
-    def build(self, max_retries, retries_wait_ms, retries_max_wait_ms):
+    def build(self, max_retries: int, retries_wait_ms: int, retries_max_wait_ms: int):
         self._validate()
         return _AsyncCustomOAuthClient(
             self.custom_function,
@@ -146,12 +202,13 @@ class _AsyncFieldProviderBuilder:
 
     __builders = {
         "OAUTHBEARER": _AsyncOAuthBearerOIDCFieldProviderBuilder,
+        "OAUTHBEARER_AZURE_IMDS": _AsyncOAuthBearerOIDCAzureIMDSFieldProviderBuilder,
         "STATIC_TOKEN": _StaticOAuthBearerFieldProviderBuilder,
         "CUSTOM": _AsyncCustomOAuthBearerFieldProviderBuilder
     }
 
     @staticmethod
-    def build(conf, max_retries, retries_wait_ms, retries_max_wait_ms):
+    def build(conf, max_retries: int, retries_wait_ms: int, retries_max_wait_ms: int):
         bearer_auth_credentials_source = conf.pop('bearer.auth.credentials.source', None)
         if bearer_auth_credentials_source is None:
             return [None, None]