From c48817e0d21d7752077e28f2ea9a657b9001a14b Mon Sep 17 00:00:00 2001
From: Tincu Gabriel <gabi@aiven.io>
Date: Wed, 2 Dec 2020 15:37:38 +0100
Subject: Support configuration of custom kafka client for
 Admin/Consumer/Producer (#2144)

---
 kafka/admin/client.py   | 10 +++++++---
 kafka/consumer/group.py |  4 +++-
 kafka/producer/kafka.py | 11 +++++++----
 3 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/kafka/admin/client.py b/kafka/admin/client.py
index c58da0c..63a0f3b 100644
--- a/kafka/admin/client.py
+++ b/kafka/admin/client.py
@@ -146,6 +146,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
+        kafka_client (callable): Custom class / callable for creating KafkaClient instances
 
     """
     DEFAULT_CONFIG = {
@@ -186,6 +187,7 @@ class KafkaAdminClient(object):
         'metric_reporters': [],
         'metrics_num_samples': 2,
         'metrics_sample_window_ms': 30000,
+        'kafka_client': KafkaClient,
     }
 
     def __init__(self, **configs):
@@ -205,9 +207,11 @@ class KafkaAdminClient(object):
         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)
+        self._client = self.config['kafka_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
diff --git a/kafka/consumer/group.py b/kafka/consumer/group.py
index 4fd57ae..a1d1dfa 100644
--- a/kafka/consumer/group.py
+++ b/kafka/consumer/group.py
@@ -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
+        kafka_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
+        'kafka_client': KafkaClient,
     }
     DEFAULT_SESSION_TIMEOUT_MS_0_9 = 30000
 
@@ -353,7 +355,7 @@ class KafkaConsumer(six.Iterator):
             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)
+        self._client = self.config['kafka_client'](metrics=self._metrics, **self.config)
 
         # Get auto-discovered version from client if necessary
         if self.config['api_version'] is None:
diff --git a/kafka/producer/kafka.py b/kafka/producer/kafka.py
index cde26b0..ea010c5 100644
--- a/kafka/producer/kafka.py
+++ b/kafka/producer/kafka.py
@@ -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
+        kafka_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,
+        'kafka_client': KafkaClient,
     }
 
     _COMPRESSORS = {
@@ -378,9 +380,10 @@ class KafkaProducer(object):
         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)
+        client = self.config['kafka_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:
-- 
cgit v1.2.1