Merge branch 'master' of https://github.com/chinmaychandak/streamz

Author: chinmaychandak

streamz/sources.py

@@ -1,4 +1,5 @@
 from glob import glob
+import json
 import os
 
 import time
@@ -453,11 +454,12 @@ def _close_consumer(self):
 class FromKafkaBatched(Stream):
     """Base class for both local and cluster-based batched kafka processing"""
     def __init__(self, topic, consumer_params, poll_interval='1s',
-                 npartitions=1, **kwargs):
+                 npartitions=1, checkpointing=None, **kwargs):
         self.consumer_params = consumer_params
         self.topic = topic
         self.npartitions = npartitions
         self.positions = [0] * npartitions
+        self.checkpointing = checkpointing
         self.poll_interval = convert_interval(poll_interval)
         self.stopped = True
 
@@ -470,18 +472,40 @@ def poll_kafka(self):
         try:
             while not self.stopped:
                 out = []
+
+                latest_checkpoint = {}
+                if self.checkpointing is not None:
+                    if not os.path.exists(self.checkpointing):
+                        os.makedirs(self.checkpointing)
+                    topic_path = self.checkpointing + '/' + self.topic
+                    if not os.path.exists(topic_path):
+                        os.makedirs(topic_path)
+                    checkpoints_list = os.listdir(topic_path)
+                    if len(checkpoints_list) > 0:
+                        previous_checkpoint = max(checkpoints_list)
+                        with open(topic_path + '/' + previous_checkpoint, 'r') as fr:
+                            latest_checkpoint = json.loads(fr.readlines()[-1])
+                        fr.close()
+
                 for partition in range(self.npartitions):
                     tp = ck.TopicPartition(self.topic, partition, 0)
                     try:
                         low, high = self.consumer.get_watermark_offsets(
                             tp, timeout=0.1)
                     except (RuntimeError, ck.KafkaException):
                         continue
+
                     current_position = self.positions[partition]
+                    group = self.consumer_params['group.id']
+
+                    if group in latest_checkpoint.keys():
+                        if str(partition) in latest_checkpoint[group].keys():
+                            current_position = latest_checkpoint[group][str(partition)]
+
                     lowest = max(current_position, low)
                     if high > lowest:
                         out.append((self.consumer_params, self.topic, partition,
-                                    lowest, high - 1))
+                                    lowest, high - 1, self.checkpointing))
                         self.positions[partition] = high
 
                 for part in out:
@@ -507,7 +531,8 @@ def start(self):
 
 @Stream.register_api(staticmethod)
 def from_kafka_batched(topic, consumer_params, poll_interval='1s',
-                       npartitions=1, start=False, dask=False, **kwargs):
+                       npartitions=1, start=False, dask=False,
+                       checkpointing=None, **kwargs):
     """ Get messages from Kafka in batches
 
     Uses the confluent-kafka library,
@@ -549,7 +574,8 @@ def from_kafka_batched(topic, consumer_params, poll_interval='1s',
         kwargs['loop'] = default_client().loop
     source = FromKafkaBatched(topic, consumer_params,
                               poll_interval=poll_interval,
-                              npartitions=npartitions, **kwargs)
+                              npartitions=npartitions,
+                              checkpointing=checkpointing, **kwargs)
     if dask:
         source = source.scatter()
 
@@ -559,7 +585,41 @@ def from_kafka_batched(topic, consumer_params, poll_interval='1s',
     return source.starmap(get_message_batch)
 
 
-def get_message_batch(kafka_params, topic, partition, low, high, timeout=None):
+def add_checkpoint(group, checkpoint, path):
+    topic = checkpoint.topic
+    partition = checkpoint.partition
+    offset = checkpoint.offset
+    latest_checkpoint = {}
+    previous_checkpoint = None
+    if not os.path.exists(path):
+        os.makedirs(path)
+    path = path + '/' + topic
+    if not os.path.exists(path):
+        os.makedirs(path)
+    checkpoints_list = os.listdir(path)
+    if len(checkpoints_list) > 0:
+        previous_checkpoint = max(checkpoints_list)
+        with open(path + '/' + previous_checkpoint, 'r') as fr:
+            latest_checkpoint = json.loads(fr.readlines()[0])
+        fr.close()
+    #Only maintain the last 5 checkpoints
+    if len(checkpoints_list) == 5:
+        os.system('rm -rf ' + path + '/' + min(checkpoints_list))
+    if group not in latest_checkpoint.keys():
+        latest_checkpoint[group] = {}
+    latest_checkpoint[group][partition] = offset
+    print(latest_checkpoint)
+    if previous_checkpoint is None:
+        new_checkpoint = '1.txt'
+    else:
+        previous_batch = int(previous_checkpoint.split('.')[0])
+        new_checkpoint = str(previous_batch + 1) + '.txt'
+    with open(path + '/' + new_checkpoint, 'a+') as fw:
+        fw.write(json.dumps(latest_checkpoint) + '\n')
+    fw.close()
+
+
+def get_message_batch(kafka_params, topic, partition, low, high, checkpointing, timeout=None):
     """Fetch a batch of kafka messages in given topic/partition
 
     This will block until messages are available, or timeout is reached.
@@ -583,5 +643,8 @@ def get_message_batch(kafka_params, topic, partition, low, high, timeout=None):
                 if timeout is not None and time.time() - t0 > timeout:
                     break
     finally:
+        if checkpointing is not None:
+            checkpoint = consumer.commit(asynchronous=False)
+            add_checkpoint(kafka_params['group.id'], checkpoint[0], checkpointing)
         consumer.close()
     return out

streamz/tests/test_kafka.py

@@ -6,6 +6,7 @@
 import requests
 import shlex
 import subprocess
+import time
 from tornado import gen
 
 from ..core import Stream
@@ -200,6 +201,88 @@ def test_kafka_batch():
         stream.upstream.stopped = True
 
 
+def test_kafka_batch_checkpointing():
+    bootstrap_servers = 'localhost:9092'
+    ARGS = {'bootstrap.servers': bootstrap_servers,
+            'group.id': 'streamz-test'}
+    ARGS1 = {'bootstrap.servers': bootstrap_servers,
+            'group.id': 'streamz-test'}
+    ARGS2 = {'bootstrap.servers': bootstrap_servers,
+            'group.id': 'streamz-test'}
+    ARGS3 = {'bootstrap.servers': bootstrap_servers,
+            'group.id': 'streamz-test'}
+    with kafka_service() as kafka:
+        kafka, TOPIC = kafka
+        for i in range(10):
+            kafka.produce(TOPIC, b'value-%d' % i)
+        kafka.flush()
+
+        stream = Stream.from_kafka_batched(TOPIC, ARGS, checkpointing='custreamz_checkpoints')
+        out = stream.sink_to_list()
+        stream.start()
+        wait_for(lambda: any(out) and out[-1][-1] == b'value-9', 10, period=0.2)
+        assert out[-1][-1] == b'value-9'
+        stream.upstream.stopped = True
+
+        stream1 = Stream.from_kafka_batched(TOPIC, ARGS1, checkpointing=None)
+        out1 = stream1.sink_to_list()
+        stream1.start()
+        wait_for(lambda: any(out1) and out1[-1][-1] == b'value-9', 10, period=0.2)
+        assert out[-1][-1] == b'value-9'
+        stream1.upstream.stopped = True
+
+        stream2 = Stream.from_kafka_batched(TOPIC, ARGS2, checkpointing='custreamz_checkpoints1')
+        out2 = stream2.sink_to_list()
+        stream2.start()
+        wait_for(lambda: any(out2) and out2[-1][-1] == b'value-9', 10, period=0.2)
+        assert out[-1][-1] == b'value-9'
+        stream2.upstream.stopped = True
+
+        for i in range(10, 20):
+            kafka.produce(TOPIC, b'value-%d' % i)
+        kafka.flush()
+
+        stream3 = Stream.from_kafka_batched(TOPIC, ARGS3, checkpointing='custreamz_checkpoints')
+        out3 = stream3.sink_to_list()
+        stream3.start()
+        wait_for(lambda: any(out3) and out3[-1][0] == b'value-10' and out3[-1][-1] == b'value-19', 10, period=0.2)
+
+        for i in range(20, 25):
+            kafka.produce(TOPIC, b'value-%d' % i)
+        kafka.flush()
+        time.sleep(5)
+        checkpoints_list = os.listdir('custreamz_checkpoints/' + TOPIC)
+        assert len(checkpoints_list) == 3
+
+        for i in range(25, 30):
+            kafka.produce(TOPIC, b'value-%d' % i)
+        kafka.flush()
+        time.sleep(5)
+        checkpoints_list = os.listdir('custreamz_checkpoints/' + TOPIC)
+        assert len(checkpoints_list) == 4
+
+        for i in range(30, 35):
+            kafka.produce(TOPIC, b'value-%d' % i)
+        kafka.flush()
+        time.sleep(5)
+        checkpoints_list = os.listdir('custreamz_checkpoints/' + TOPIC)
+        assert len(checkpoints_list) == 5
+
+        for i in range(35, 40):
+            kafka.produce(TOPIC, b'value-%d' % i)
+        kafka.flush()
+        time.sleep(5)
+        checkpoints_list = os.listdir('custreamz_checkpoints/' + TOPIC)
+        assert len(checkpoints_list) == 5
+
+        for i in range(40, 45):
+            kafka.produce(TOPIC, b'value-%d' % i)
+        kafka.flush()
+        time.sleep(5)
+        checkpoints_list = os.listdir('custreamz_checkpoints/' + TOPIC)
+        assert len(checkpoints_list) == 5
+
+
 @gen_cluster(client=True, timeout=60)
 def test_kafka_dask_batch(c, s, w1, w2):
     j = random.randint(0, 10000)