client: allow for custom kafka clients

Provide the consumer, producer and admin client with the option to
create the kafka client from a custom callable, thus allowing more
flexibility in handling certain low level errors

Author: Gabriel Tincu

kafka/admin/client.py

@@ -10,7 +10,7 @@
 
 from kafka.admin.acl_resource import ACLOperation, ACLPermissionType, ACLFilter, ACL, ResourcePattern, ResourceType, \
     ACLResourcePatternType
-from kafka.client_async import KafkaClient, selectors
+from kafka.client_async import selectors, KafkaClient
 from kafka.coordinator.protocol import ConsumerProtocolMemberMetadata, ConsumerProtocolMemberAssignment, ConsumerProtocol
 import kafka.errors as Errors
 from kafka.errors import (
@@ -26,6 +26,7 @@
 from kafka.protocol.metadata import MetadataRequest
 from kafka.protocol.types import Array
 from kafka.structs import TopicPartition, OffsetAndMetadata, MemberInformation, GroupInformation
+from kafka.util import get_client_factory
 from kafka.version import __version__
 
 
@@ -146,6 +147,7 @@ class KafkaAdminClient(object):
             sasl mechanism handshake. Default: one of bootstrap servers
         sasl_oauth_token_provider (AbstractTokenProvider): OAuthBearer token provider
             instance. (See kafka.oauth.abstract). Default: None
+        client (callable): Custom class / callable for creating KafkaClient instances
 
     """
     DEFAULT_CONFIG = {
@@ -186,6 +188,7 @@ class KafkaAdminClient(object):
         'metric_reporters': [],
         'metrics_num_samples': 2,
         'metrics_sample_window_ms': 30000,
+        'client': KafkaClient,
     }
 
     def __init__(self, **configs):
@@ -204,10 +207,12 @@ def __init__(self, **configs):
                                      tags=metrics_tags)
         reporters = [reporter() for reporter in self.config['metric_reporters']]
         self._metrics = Metrics(metric_config, reporters)
-
-        self._client = KafkaClient(metrics=self._metrics,
-                                   metric_group_prefix='admin',
-                                   **self.config)
+        assert callable(self.config['client']), "Client parameter should be callable"
+        self._client = self.config['client'](
+            metrics=self._metrics,
+            metric_group_prefix='admin',
+            **self.config
+        )
         self._client.check_version(timeout=(self.config['api_version_auto_timeout_ms'] / 1000))
 
         # Get auto-discovered version from client if necessary

kafka/consumer/group.py

@@ -9,7 +9,7 @@
 
 from kafka.vendor import six
 
-from kafka.client_async import KafkaClient, selectors
+from kafka.client_async import selectors, KafkaClient
 from kafka.consumer.fetcher import Fetcher
 from kafka.consumer.subscription_state import SubscriptionState
 from kafka.coordinator.consumer import ConsumerCoordinator
@@ -244,6 +244,7 @@ class KafkaConsumer(six.Iterator):
             sasl mechanism handshake. Default: one of bootstrap servers
         sasl_oauth_token_provider (AbstractTokenProvider): OAuthBearer token provider
             instance. (See kafka.oauth.abstract). Default: None
+        client (callable): Custom class / callable for creating KafkaClient instances
 
     Note:
         Configuration parameters are described in more detail at
@@ -306,6 +307,7 @@ class KafkaConsumer(six.Iterator):
         'sasl_kerberos_domain_name': None,
         'sasl_oauth_token_provider': None,
         'legacy_iterator': False, # enable to revert to < 1.4.7 iterator
+        'client': KafkaClient,
     }
     DEFAULT_SESSION_TIMEOUT_MS_0_9 = 30000
 
@@ -352,8 +354,8 @@ def __init__(self, *topics, **configs):
                 self.config['api_version'] = tuple(map(int, str_version.split('.')))
             log.warning('use api_version=%s [tuple] -- "%s" as str is deprecated',
                         str(self.config['api_version']), str_version)
-
-        self._client = KafkaClient(metrics=self._metrics, **self.config)
+        assert callable(self.config['client']), "Client parameter should be callable"
+        self._client = self.config['client'](metrics=self._metrics, **self.config)
 
         # Get auto-discovered version from client if necessary
         if self.config['api_version'] is None:

kafka/producer/kafka.py

@@ -11,7 +11,7 @@
 from kafka.vendor import six
 
 import kafka.errors as Errors
-from kafka.client_async import KafkaClient, selectors
+from kafka.client_async import selectors, KafkaClient
 from kafka.codec import has_gzip, has_snappy, has_lz4, has_zstd
 from kafka.metrics import MetricConfig, Metrics
 from kafka.partitioner.default import DefaultPartitioner
@@ -280,6 +280,7 @@ class KafkaProducer(object):
             sasl mechanism handshake. Default: one of bootstrap servers
         sasl_oauth_token_provider (AbstractTokenProvider): OAuthBearer token provider
             instance. (See kafka.oauth.abstract). Default: None
+        client (callable): Custom class / callable for creating KafkaClient instances
 
     Note:
         Configuration parameters are described in more detail at
@@ -332,7 +333,8 @@ class KafkaProducer(object):
         'sasl_plain_password': None,
         'sasl_kerberos_service_name': 'kafka',
         'sasl_kerberos_domain_name': None,
-        'sasl_oauth_token_provider': None
+        'sasl_oauth_token_provider': None,
+        'client': KafkaClient,
     }
 
     _COMPRESSORS = {
@@ -377,10 +379,11 @@ def __init__(self, **configs):
                                      tags=metrics_tags)
         reporters = [reporter() for reporter in self.config['metric_reporters']]
         self._metrics = Metrics(metric_config, reporters)
-
-        client = KafkaClient(metrics=self._metrics, metric_group_prefix='producer',
-                             wakeup_timeout_ms=self.config['max_block_ms'],
-                             **self.config)
+        assert callable(self.config['client']), "Client parameter should be callable"
+        client = self.config['client'](
+            metrics=self._metrics, metric_group_prefix='producer',
+            wakeup_timeout_ms=self.config['max_block_ms'],
+            **self.config)
 
         # Get auto-discovered version from client if necessary
         if self.config['api_version'] is None:

kafka/util.py

@@ -1,6 +1,7 @@
 from __future__ import absolute_import
 
 import binascii
+import kafka
 import weakref
 
 from kafka.vendor import six