Refactoring a bit, cleanup for 0.8

Marking some stuff as not compatible for 0.8 (will be added in 0.8.1)

Author: mumrah

example.py

@@ -2,27 +2,21 @@
 
 from kafka.client import KafkaClient, FetchRequest, ProduceRequest
 from kafka.consumer import SimpleConsumer
+from kafka.producer import SimpleProducer
 
-def produce_example(kafka):
-    message = kafka.create_message("testing")
-    request = ProduceRequest("my-topic", -1, [message])
-    kafka.send_message_set(request)
+def produce_example(client):
+    producer = SimpleProducer(client, "my-topic")
+    producer.send_message("test")
 
-def consume_example(kafka):
-    request = FetchRequest("my-topic", 0, 0, 1024)
-    (messages, nextRequest) = kafka.get_message_set(request)
-    for message in messages:
-        print("Got Message: %s" % (message,))
-    print(nextRequest)
-
-def produce_gz_example(kafka):
-    message = kafka.create_gzip_message("this message was gzipped", "along with this one")
-    request = ProduceRequest("my-topic", 0, [message])
-    kafka.send_message_set(request)
+def consume_example(client):
+    consumer = SimpleConsumer(client, "test-group", "my-topic")
+    for message in consumer:
+        print(message)
 
 def main():
     client = KafkaClient("localhost", 9092)
-    consumer = SimpleConsumer(client, "test-group", "my-topic")
+    produce_example(client)
+    consume_example(client)
 
 if __name__ == "__main__":
     logging.basicConfig(level=logging.DEBUG)

kafka-src

@@ -1,3 +1,4 @@
+from itertools import izip_longest, repeat
 import logging
 from threading import Lock
 
@@ -30,7 +31,7 @@ def __init__(self, client, group, topic, auto_commit=False, auto_commit_every_n=
         self.client = client
         self.topic = topic
         self.group = group
-        self.client.load_metadata_for_topics(topic)
+        self.client._load_metadata_for_topics(topic)
         self.offsets = {}
 
         # Set up the auto-commit timer
@@ -54,12 +55,16 @@ def get_or_init_offset_callback(resp):
                 raise Exception("OffsetFetchRequest for topic=%s, partition=%d failed with errorcode=%s" % (
                     resp.topic, resp.partition, resp.error))
 
+        # Uncomment for 0.8.1
+        #
+        #for partition in self.client.topic_partitions[topic]:
+        #    req = OffsetFetchRequest(topic, partition)
+        #    (offset,) = self.client.send_offset_fetch_request(group, [req],
+        #                  callback=get_or_init_offset_callback, fail_on_error=False)
+        #    self.offsets[partition] = offset
+        
         for partition in self.client.topic_partitions[topic]:
-            req = OffsetFetchRequest(topic, partition)
-            (offset,) = self.client.send_offset_fetch_request(group, [req],
-                          callback=get_or_init_offset_callback, fail_on_error=False)
-            self.offsets[partition] = offset
-        print self.offsets
+            self.offsets[partition] = 0
 
     def seek(self, offset, whence):
         """
@@ -71,25 +76,30 @@ def seek(self, offset, whence):
                 1 is relative to the current offset
                 2 is relative to the latest known offset (tail)
         """
-        if whence == 1:
-            # relative to current position
+        if whence == 1: # relative to current position
             for partition, _offset in self.offsets.items():
                 self.offset[partition] = _offset + offset
-        elif whence in (0, 2):
-            # relative to beginning or end
+        elif whence in (0, 2): # relative to beginning or end
+            # divide the request offset by number of partitions, distribute the remained evenly
+            (delta, rem) = divmod(offset, len(self.offsets))
+            deltas = {}
+            for partition, r in izip_longest(self.offsets.keys(), repeat(1, rem), fillvalue=0):
+                deltas[partition] = delta + r
+
             reqs = []
-            for partition in offsets.keys():
+            for partition in self.offsets.keys():
                 if whence == 0:
                     reqs.append(OffsetRequest(self.topic, partition, -2, 1))
                 elif whence == 2:
                     reqs.append(OffsetRequest(self.topic, partition, -1, 1))
                 else:
                     pass
-            resps = self.client.send_offset_request([req])
+
+            resps = self.client.send_offset_request(reqs)
             for resp in resps:
-                self.offsets[resp.partition] = resp.offsets[0] + offset
+                self.offsets[resp.partition] = resp.offsets[0] + deltas[resp.partition] 
         else:
-            raise
+            raise ValueError("Unexpected value for `whence`, %d" % whence)
 
     def commit(self, partitions=[]):
         """
@@ -98,6 +108,8 @@ def commit(self, partitions=[]):
         partitions: list of partitions to commit, default is to commit all of them
         """
 
+        raise NotImplementedError("Broker-managed offsets not supported in 0.8")
+
         # short circuit if nothing happened
         if self.count_since_commit == 0:
             return
@@ -121,15 +133,31 @@ def commit(self, partitions=[]):
             self.count_since_commit = 0
 
     def __iter__(self):
+        """
+        Create an iterate per partition. Iterate through them calling next() until they are
+        all exhausted.
+        """
         iters = {}
         for partition, offset in self.offsets.items():
             iters[partition] = self.__iter_partition__(partition, offset)
 
+        if len(iters) == 0:
+            return
+
         while True:
-            for it in iters.values():
-                yield it.next()
+            if len(iters) == 0:
+                break
+
+            for partition, it in iters.items():
+                try:
+                    yield it.next()
+                except StopIteration:
+                    log.debug("Done iterating over partition %s" % partition)
+                    del iters[partition]
+                    continue # skip auto-commit since we didn't yield anything
+
+                # auto commit logic
                 self.count_since_commit += 1
-                # deal with auto commits
                 if self.auto_commit is True:
                     if self.auto_commit_every_n is not None and self.count_since_commit > self.auto_commit_every_n:
                         if self.commit_timer is not None:
@@ -140,19 +168,22 @@ def __iter__(self):
                             self.commit()
 
     def __iter_partition__(self, partition, offset):
+        """
+        Iterate over the messages in a partition. Create a FetchRequest to get back
+        a batch of messages, yield them one at a time. After a batch is exhausted,
+        start a new batch unless we've reached the end of ths partition.
+        """
         while True:
-            req = FetchRequest(self.topic, partition, offset, 1024) 
+            req = FetchRequest(self.topic, partition, offset, 1024) # TODO configure fetch size
             (resp,) = self.client.send_fetch_request([req])
             assert resp.topic == self.topic
             assert resp.partition == partition
             next_offset = None
             for message in resp.messages:
                 next_offset = message.offset
-                print partition, message, message.offset
                 yield message
                 # update the internal state _after_ we yield the message
                 self.offsets[partition] = message.offset
-            print partition, next_offset
             if next_offset is None:
                 break
             else:

kafka/client.py

@@ -13,10 +13,11 @@ class SimpleProducer(object):
     def __init__(self, client, topic):
         self.client = client
         self.topic = topic
-        self.client.load_metadata_for_topics(topic)
+        self.client._load_metadata_for_topics(topic)
         self.next_partition = cycle(self.client.topic_partitions[topic])
 
     def send_message(self, msg):
         req = ProduceRequest(self.topic, self.next_partition.next(),
             messages=[create_message(msg)])
-        resp = self.client.send_produce_request([req]).next()
+        resp = self.client.send_produce_request([req])[0]
+        assert resp.error == 0

kafka/consumer.py

@@ -217,7 +217,7 @@ def encode_fetch_request(cls, client_id, correlation_id, payloads=[], max_wait_t
         return struct.pack('>i%ds' % len(message), len(message), message)
 
     @classmethod
-    def decode_fetch_response_iter(cls, data):
+    def decode_fetch_response(cls, data):
         """
         Decode bytes to a FetchResponse
 

kafka/producer.py

@@ -8,6 +8,8 @@
 
 log = logging.getLogger("kafka")
 
+raise NotImplementedError("Still need to refactor this class")
+
 class KafkaConsumerProcess(Process):
     def __init__(self, client, topic, partition, out_queue, barrier, consumer_fetch_size=1024, consumer_sleep=200):
         self.client = copy(client)

kafka/protocol.py

@@ -243,7 +243,7 @@ def test_produce_100k_gzipped(self):
     def test_consume_none(self):
         fetch = FetchRequest("test_consume_none", 0, 0, 1024)
 
-        fetch_resp = self.client.send_fetch_request([fetch]).next()
+        fetch_resp = self.client.send_fetch_request([fetch])[0]
         self.assertEquals(fetch_resp.error, 0)
         self.assertEquals(fetch_resp.topic, "test_consume_none")
         self.assertEquals(fetch_resp.partition, 0)
@@ -263,7 +263,7 @@ def test_produce_consume(self):
 
         fetch = FetchRequest("test_produce_consume", 0, 0, 1024)
 
-        fetch_resp = self.client.send_fetch_request([fetch]).next()
+        fetch_resp = self.client.send_fetch_request([fetch])[0]
         self.assertEquals(fetch_resp.error, 0)
 
         messages = list(fetch_resp.messages)
@@ -343,6 +343,7 @@ def test_produce_consume_two_partitions(self):
     #   Offset Tests   #
     ####################
 
+    @unittest.skip("No supported until 0.8.1")
     def test_commit_fetch_offsets(self):
         req = OffsetCommitRequest("test_commit_fetch_offsets", 0, 42, "metadata")
         (resp,) = self.client.send_offset_commit_request("group", [req])
@@ -428,22 +429,20 @@ def test_consumer(self):
         self.assertEquals(len(all_messages), 200)
         self.assertEquals(len(all_messages), len(set(all_messages))) # make sure there are no dupes
 
-        # Produce more messages
-        produce3 = ProduceRequest("test_consumer", 1, messages=[
-            create_message("Test message 3 %d" % i) for i in range(10)
-        ])
-
-        for resp in self.client.send_produce_request([produce3]):
-            self.assertEquals(resp.error, 0)
-            self.assertEquals(resp.offset, 100)
+        consumer.seek(-10, 2)
+        all_messages = []
+        for message in consumer:
+            all_messages.append(message)
 
-        # Start a new consumer, make sure we only get the newly produced messages
-        consumer = SimpleConsumer(self.client, "group1", "test_consumer")
+        self.assertEquals(len(all_messages), 10)
 
+        consumer.seek(-13, 2)
         all_messages = []
         for message in consumer:
             all_messages.append(message)
-        self.assertEquals(len(all_messages), 10)
+
+        self.assertEquals(len(all_messages), 13)
+
 
 if __name__ == "__main__":
     logging.basicConfig(level=logging.DEBUG)