Allow injecting clients into Kafka Consumers and Producers.

Author: Julian

kafka/consumer/group.py

@@ -40,6 +40,9 @@ class KafkaConsumer(six.Iterator):
             It just needs to have at least one broker that will respond to a
             Metadata API Request. Default port is 9092. If no servers are
             specified, will default to localhost:9092.
+        client (kafka.client_async.KafkaClient): a kafka client to
+            use, or if unprovided, one is constructed from the provided
+            configuration.
         client_id (str): a name for this client. This string is passed in
             each request to servers and can be used to identify specific
             server-side log entries that correspond to this client. Also
@@ -201,6 +204,8 @@ class KafkaConsumer(six.Iterator):
     }
 
     def __init__(self, *topics, **configs):
+        client = configs.pop("client", None)
+
         self.config = copy.copy(self.DEFAULT_CONFIG)
         for key in self.config:
             if key in configs:
@@ -226,7 +231,9 @@ def __init__(self, *topics, **configs):
         metric_group_prefix = 'consumer'
         # TODO _metrics likely needs to be passed to KafkaClient, etc.
 
-        self._client = KafkaClient(**self.config)
+        if client is None:
+            client = KafkaClient(**self.config)
+        self._client = client
 
         # Check Broker Version if not set explicitly
         if self.config['api_version'] == 'auto':

kafka/producer/kafka.py

@@ -79,6 +79,9 @@ class KafkaProducer(object):
             It just needs to have at least one broker that will respond to a
             Metadata API Request. Default port is 9092. If no servers are
             specified, will default to localhost:9092.
+        client (kafka.client_async.KafkaClient): a kafka client to
+            use, or if unprovided, one is constructed from the provided
+            configuration.
         client_id (str): a name for this client. This string is passed in
             each request to servers and can be used to identify specific
             server-side log entries that correspond to this client.
@@ -253,7 +256,7 @@ class KafkaProducer(object):
         'api_version': 'auto',
     }
 
-    def __init__(self, **configs):
+    def __init__(self, client=None, **configs):
         log.debug("Starting the Kafka producer") # trace
         self.config = copy.copy(self._DEFAULT_CONFIG)
         for key in self.config:
@@ -270,7 +273,9 @@ def __init__(self, **configs):
         if self.config['acks'] == 'all':
             self.config['acks'] = -1
 
-        client = KafkaClient(**self.config)
+        if client is None:
+            client = KafkaClient(**self.config)
+        self.client = client
 
         # Check Broker Version if not set explicitly
         if self.config['api_version'] == 'auto':