Add UAMI OAuth changes (#2189)

* Add UAMI OAuth changes

* Fix formatting

* Fix concurrent and syntax issues

* Create sync files for uami oauth

* make style fix

* Add async base client

* Fix examples and expiry

* Add api version to example

* make stylefix

* Fix typos

* remove word

Author: Claimundefine

examples/oauth_oidc_ccloud_azure_imds_producer.py

@@ -17,13 +17,55 @@
 
 
 # This example uses Azure IMDS for credential-less authentication
-# to Kafka on Confluent Cloud
+# to Schema Registry on Confluent Cloud
 
 import argparse
 import logging
 
 from confluent_kafka import Producer
-from confluent_kafka.serialization import StringSerializer
+from confluent_kafka.schema_registry import SchemaRegistryClient
+from confluent_kafka.schema_registry.json_schema import JSONSerializer
+from confluent_kafka.serialization import MessageField, SerializationContext, StringSerializer
+
+
+class User(object):
+    """
+    User record
+
+    Args:
+        name (str): User's name
+
+        favorite_number (int): User's favorite number
+
+        favorite_color (str): User's favorite color
+
+        address(str): User's address; confidential
+    """
+
+    def __init__(self, name, address, favorite_number, favorite_color):
+        self.name = name
+        self.favorite_number = favorite_number
+        self.favorite_color = favorite_color
+        # address should not be serialized, see user_to_dict()
+        self._address = address
+
+
+def user_to_dict(user, ctx):
+    """
+    Returns a dict representation of a User instance for serialization.
+
+    Args:
+        user (User): User instance.
+
+        ctx (SerializationContext): Metadata pertaining to the serialization
+            operation.
+
+    Returns:
+        dict: Dict populated with user attributes to be serialized.
+    """
+
+    # User._address must not be serialized; omit from dict
+    return dict(name=user.name, favorite_number=user.favorite_number, favorite_color=user.favorite_color)
 
 
 def producer_config(args):
@@ -38,7 +80,7 @@ def producer_config(args):
         'sasl.oauthbearer.config': f'query={args.query}',
     }
     # These two parameters are only applicable when producing to
-    # Confluent Cloud where some sasl extensions are required.
+    # 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
@@ -47,12 +89,27 @@ def producer_config(args):
     return params
 
 
+def schema_registry_config(args):
+    params = {
+        'url': args.schema_registry,
+        'bearer.auth.credentials.source': 'OAUTHBEARER_AZURE_IMDS',
+        'bearer.auth.issuer.endpoint.query': args.query,
+    }
+    # These two parameters are only applicable when producing to
+    # confluent cloud where some sasl extensions are required.
+    if args.logical_schema_registry_cluster and args.identity_pool_id:
+        params['bearer.auth.logical.cluster'] = args.logical_schema_registry_cluster
+        params['bearer.auth.identity.pool.id'] = 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, or None on success.
+        err (KafkaError): The error that occurred on None on success.
 
         msg (Message): The message that was produced or failed.
 
@@ -80,15 +137,45 @@ def main(args):
     producer_conf = producer_config(args)
     producer = Producer(producer_conf)
     string_serializer = StringSerializer('utf_8')
+    schema_str = """
+    {
+      "$schema": "http://json-schema.org/draft-07/schema#",
+      "title": "User",
+      "description": "A Confluent Kafka Python User",
+      "type": "object",
+      "properties": {
+        "name": {
+          "description": "User's name",
+          "type": "string"
+        },
+        "favorite_number": {
+          "description": "User's favorite number",
+          "type": "number",
+          "exclusiveMinimum": 0
+        },
+        "favorite_color": {
+          "description": "User's favorite color",
+          "type": "string"
+        }
+      },
+      "required": [ "name", "favorite_number", "favorite_color" ]
+    }
+    """
+    schema_registry_conf = schema_registry_config(args)
+    schema_registry_client = SchemaRegistryClient(schema_registry_conf)
+
+    json_serializer = JSONSerializer(schema_str, schema_registry_client, user_to_dict)
 
     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:
             name = input(">")
+            user = User(name=name, address="NA", favorite_color="blue", favorite_number=7)
+            serialized_user = json_serializer(user, SerializationContext(topic, MessageField.VALUE))
             producer.produce(
-                topic=topic, key=string_serializer(name), value=string_serializer(name), on_delivery=delivery_report
+                topic=topic, key=string_serializer(name), value=serialized_user, on_delivery=delivery_report
             )
         except KeyboardInterrupt:
             break
@@ -98,11 +185,18 @@ def main(args):
 
 
 if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description="OAuth/OIDC example using Azure IMDS metadata-based authentication")
+    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('-s', dest="schema_registry", required=True, help="Schema Registry (http(s)://host[:port]")
     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(
+        '--logical-schema-registry-cluster',
+        dest="logical_schema_registry_cluster",
+        required=False,
+        help="Logical Schema Registry 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
@@ -52,6 +52,17 @@ 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
 

examples/schema_registry_imds_azure.py

@@ -0,0 +1,49 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# Copyright 2026 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.
+
+# Examples of setting up Schema Registry with OAuth with static token,
+# Client Credentials, Azure IMDS, and custom functions
+
+
+# CUSTOM OAuth configuration takes in a custom function, config for that
+# function, and returns the fields shown below. All fields must be returned
+# for OAuth authentication to work
+
+
+from confluent_kafka.schema_registry.schema_registry_client import SchemaRegistryClient
+
+
+def main():
+    BOOTSTRAP = "api://bootstrapid"
+    UAMI_CLIENT_ID = "uamiid"
+
+    endpoint_query = f"api-version=2025-04-07&resource={BOOTSTRAP}&client_id={UAMI_CLIENT_ID}"
+
+    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': endpoint_query,
+        '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())
+
+
+if __name__ == '__main__':
+    main()