Allow customizing socket timeouts.

Previously, if you try to consume a message with a timeout greater than 10 seconds,
but you don't receive data in those 10 seconds, a socket.timeout exception is raised.
This allows a higher socket timeout to be set, or even None for no timeout.

1 files changed, 4 insertions, 3 deletions

diff --git a/kafka/client.py b/kafka/client.py
index 71ededa..9659364 100644
--- a/kafka/client.py
+++ b/kafka/client.py
@@ -19,12 +19,13 @@ class KafkaClient(object):
     CLIENT_ID = "kafka-python"
     ID_GEN = count()
 
-    def __init__(self, host, port, bufsize=4096, client_id=CLIENT_ID):
+    def __init__(self, host, port, bufsize=4098, client_id=CLIENT_ID, timeout=10):
         # We need one connection to bootstrap
         self.bufsize = bufsize
         self.client_id = client_id
+        self.timeout = timeout
         self.conns = {               # (host, port) -> KafkaConnection
-            (host, port): KafkaConnection(host, port, bufsize)
+            (host, port): KafkaConnection(host, port, bufsize, timeout=timeout)
         }
         self.brokers = {}            # broker_id -> BrokerMetadata
         self.topics_to_brokers = {}  # topic_id -> broker_id
@@ -41,7 +42,7 @@ class KafkaClient(object):
         """
         if (broker.host, broker.port) not in self.conns:
             self.conns[(broker.host, broker.port)] = \
-                KafkaConnection(broker.host, broker.port, self.bufsize)
+                KafkaConnection(broker.host, broker.port, self.bufsize, timeout=self.timeout)
 
         return self.conns[(broker.host, broker.port)]