Refactor

Author: rohitsanj

src/confluent_kafka/schema_registry/__init__.py

@@ -29,6 +29,7 @@
   RuleSet,
   Schema,
   SchemaRegistryClient,
+  AsyncSchemaRegistryClient,
   SchemaRegistryError,
   SchemaReference,
   ServerConfig
@@ -56,6 +57,7 @@
   "RuleSet",
   "Schema",
   "SchemaRegistryClient",
+  "AsyncSchemaRegistryClient",
   "SchemaRegistryError",
   "SchemaReference",
   "ServerConfig",

src/confluent_kafka/schema_registry/_async/avro.py

@@ -30,10 +30,16 @@
                AsyncSchemaRegistryClient)
 from confluent_kafka.serialization import (SerializationError,
                                            SerializationContext)
-from confluent_kafka.schema_registry.common import _ContextStringIO, asyncinit
+from confluent_kafka.schema_registry.common import asyncinit
+from confluent_kafka.schema_registry.common import _ContextStringIO
 from confluent_kafka.schema_registry.rule_registry import RuleRegistry
 from confluent_kafka.schema_registry.serde import AsyncBaseSerializer, AsyncBaseDeserializer, ParsedSchemaCache
 
+__all__ = [
+    '_resolve_named_schema',
+    'AsyncAvroSerializer',
+    'AsyncAvroDeserializer',
+]
 
 async def _resolve_named_schema(
     schema: Schema, schema_registry_client: AsyncSchemaRegistryClient

src/confluent_kafka/schema_registry/_async/json_schema.py

@@ -40,6 +40,11 @@
 from confluent_kafka.serialization import (SerializationError,
                                            SerializationContext)
 
+__all__ = [
+    '_resolve_named_schema',
+    'AsyncJSONSerializer',
+    'AsyncJSONDeserializer'
+]
 
 async def _resolve_named_schema(
     schema: Schema, schema_registry_client: AsyncSchemaRegistryClient,

src/confluent_kafka/schema_registry/_async/protobuf.py

@@ -40,6 +40,11 @@
 from confluent_kafka.schema_registry.common import asyncinit
 from confluent_kafka.schema_registry.serde import AsyncBaseSerializer, AsyncBaseDeserializer, ParsedSchemaCache
 
+__all__ = [
+    '_resolve_named_schema',
+    'AsyncProtobufSerializer',
+    'AsyncProtobufDeserializer',
+]
 
 async def _resolve_named_schema(
     schema: Schema,

src/confluent_kafka/schema_registry/_async/schema_registry_client.py

@@ -38,10 +38,20 @@
     is_retriable, 
     _BearerFieldProvider,
     full_jitter,
-    _SchemaCache, 
+    _SchemaCache,
     Schema,
+    _StaticFieldProvider,
 )
 
+__all__ = [
+    '_urlencode',
+    '_AsyncCustomOAuthClient',
+    '_AsyncOAuthClient',
+    '_AsyncBaseRestClient',
+    '_AsyncRestClient',
+    'AsyncSchemaRegistryClient',
+]
+
 # TODO: consider adding `six` dependency or employing a compat file
 # Python 2.7 is officially EOL so compatibility issue will be come more the norm.
 # We need a better way to handle these issues.
@@ -64,17 +74,6 @@ def _urlencode(value: str) -> str:
 log = logging.getLogger(__name__)
 
 
-class _AsyncStaticFieldProvider(_BearerFieldProvider):
-    def __init__(self, token: str, logical_cluster: str, identity_pool: str):
-        self.token = token
-        self.logical_cluster = logical_cluster
-        self.identity_pool = identity_pool
-
-    async def get_bearer_fields(self) -> dict:
-        return {'bearer.auth.token': self.token, 'bearer.auth.logical.cluster': self.logical_cluster,
-                'bearer.auth.identity.pool.id': self.identity_pool}
-
-
 class _AsyncCustomOAuthClient(_BearerFieldProvider):
     def __init__(self, custom_function: Callable[[Dict], Dict], custom_config: dict):
         self.custom_function = custom_function
@@ -292,7 +291,7 @@ def __init__(self, conf: dict):
                 if 'bearer.auth.token' not in conf_copy:
                     raise ValueError("Missing bearer.auth.token")
                 static_token = conf_copy.pop('bearer.auth.token')
-                self.bearer_field_provider = _AsyncStaticFieldProvider(static_token, logical_cluster, identity_pool)
+                self.bearer_field_provider = _StaticFieldProvider(static_token, logical_cluster, identity_pool)
                 if not isinstance(static_token, string_type):
                     raise TypeError("bearer.auth.token must be a str, not " + str(type(static_token)))
             elif self.bearer_auth_credentials_source == 'CUSTOM':

src/confluent_kafka/schema_registry/_async/serde.py

@@ -26,6 +26,12 @@
 from confluent_kafka.serialization import Serializer, Deserializer, \
     SerializationContext, SerializationError
 
+__all__ = [
+    'AsyncBaseSerde',
+    'AsyncBaseSerializer',
+    'AsyncBaseDeserializer',
+]
+
 log = logging.getLogger(__name__)
 
 class AsyncBaseSerde(object):

src/confluent_kafka/schema_registry/_sync/avro.py

@@ -30,6 +30,7 @@
                SchemaRegistryClient)
 from confluent_kafka.serialization import (SerializationError,
                                            SerializationContext)
+
 from confluent_kafka.schema_registry.common import _ContextStringIO
 from confluent_kafka.schema_registry.rule_registry import RuleRegistry
 from confluent_kafka.schema_registry.serde import BaseSerializer, BaseDeserializer, ParsedSchemaCache
@@ -349,6 +350,7 @@ def _get_parsed_schema(self, schema: Schema) -> AvroSchema:
         return parsed_schema
 
 
+
 class AvroDeserializer(BaseDeserializer):
     """
     Deserializer for Avro binary encoded data with Confluent Schema Registry

src/confluent_kafka/schema_registry/_sync/schema_registry_client.py

@@ -38,7 +38,7 @@
     is_retriable, 
     _BearerFieldProvider,
     full_jitter,
-    _SchemaCache, 
+    _SchemaCache,
     Schema,
     _StaticFieldProvider,
 )