Add to_kafka stream and fix kafka testing issues (#230)

Author: martindurant

streamz/core.py

@@ -1231,6 +1231,88 @@ def cb(self):
             yield self._emit(x)
 
 
+@Stream.register_api()
+class to_kafka(Stream):
+    """ Writes data in the stream to Kafka
+
+    This stream accepts a string or bytes object. Call ``flush`` to ensure all
+    messages are pushed. Responses from Kafka are pushed downstream.
+
+    Parameters
+    ----------
+    topic : string
+        The topic which to write
+    producer_config : dict
+        Settings to set up the stream, see
+        https://docs.confluent.io/current/clients/confluent-kafka-python/#configuration
+        https://github.com/edenhill/librdkafka/blob/master/CONFIGURATION.md
+        Examples:
+        bootstrap.servers: Connection string (host:port) to Kafka
+
+    Examples
+    --------
+    >>> from streamz import Stream
+    >>> ARGS = {'bootstrap.servers': 'localhost:9092'}
+    >>> source = Stream()
+    >>> kafka = source.map(lambda x: str(x)).to_kafka('test', ARGS)
+    <to_kafka>
+    >>> for i in range(10):
+    ...     source.emit(i)
+    >>> kafka.flush()
+    """
+    def __init__(self, upstream, topic, producer_config, **kwargs):
+        import confluent_kafka as ck
+
+        self.topic = topic
+        self.producer = ck.Producer(producer_config)
+
+        Stream.__init__(self, upstream, ensure_io_loop=True, **kwargs)
+        self.stopped = False
+        self.polltime = 0.2
+        self.loop.add_callback(self.poll)
+        self.futures = []
+
+    @gen.coroutine
+    def poll(self):
+        while not self.stopped:
+            # executes callbacks for any delivered data, in this thread
+            # if no messages were sent, nothing happens
+            self.producer.poll(0)
+            yield gen.sleep(self.polltime)
+
+    def update(self, x, who=None):
+        future = gen.Future()
+        self.futures.append(future)
+
+        @gen.coroutine
+        def _():
+            while True:
+                try:
+                    # this runs asynchronously, in C-K's thread
+                    self.producer.produce(self.topic, x, callback=self.cb)
+                    return
+                except BufferError:
+                    yield gen.sleep(self.polltime)
+                except Exception as e:
+                    future.set_exception(e)
+                    return
+
+        self.loop.add_callback(_)
+        return future
+
+    @gen.coroutine
+    def cb(self, err, msg):
+        future = self.futures.pop(0)
+        if msg is not None and msg.value() is not None:
+            future.set_result(None)
+            yield self._emit(msg.value())
+        else:
+            future.set_exception(err or msg.error())
+
+    def flush(self, timeout=-1):
+        self.producer.flush(timeout)
+
+
 def sync(loop, func, *args, **kwargs):
     """
     Run coroutine in loop running in separate thread.

streamz/sources.py

@@ -1,6 +1,5 @@
 from glob import glob
 import os
-import weakref
 
 import time
 import tornado.ioloop
@@ -180,8 +179,7 @@ class from_kafka(Source):
     ...           {'bootstrap.servers': 'localhost:9092',
     ...            'group.id': 'streamz'})  # doctest: +SKIP
     """
-    def __init__(self, topics, consumer_params, poll_interval=0.1, start=False,
-                 **kwargs):
+    def __init__(self, topics, consumer_params, poll_interval=0.1, start=False, **kwargs):
         self.cpars = consumer_params
         self.consumer = None
         self.topics = topics
@@ -194,7 +192,7 @@ def __init__(self, topics, consumer_params, poll_interval=0.1, start=False,
     def do_poll(self):
         if self.consumer is not None:
             msg = self.consumer.poll(0)
-            if msg and msg.value():
+            if msg and msg.value() and msg.error() is None:
                 return msg.value()
 
     @gen.coroutine
@@ -207,26 +205,19 @@ def poll_kafka(self):
                 yield gen.sleep(self.poll_interval)
             if self.stopped:
                 break
+        self._close_consumer()
 
     def start(self):
         import confluent_kafka as ck
-        import distributed
         if self.stopped:
-            finalize = distributed.compatibility.finalize
             self.stopped = False
-            self.loop.add_callback(self.poll_kafka)
             self.consumer = ck.Consumer(self.cpars)
             self.consumer.subscribe(self.topics)
+            tp = ck.TopicPartition(self.topics[0], 0, 0)
 
-            def close(ref):
-                ob = ref()
-                if ob is not None and ob.consumer is not None:
-                    consumer = ob.consumer
-                    ob.consumer = None
-                    consumer.unsubscribe()
-                    consumer.close()  # may raise with latest ck, that's OK
-
-            finalize(self, close, weakref.ref(self))
+            # blocks for consumer thread to come up
+            self.consumer.get_watermark_offsets(tp)
+            self.loop.add_callback(self.poll_kafka)
 
     def _close_consumer(self):
         if self.consumer is not None:
@@ -253,36 +244,43 @@ def __init__(self, topic, consumer_params, poll_interval='1s',
     @gen.coroutine
     def poll_kafka(self):
         import confluent_kafka as ck
-        consumer = ck.Consumer(self.consumer_params)
 
         try:
             while not self.stopped:
                 out = []
-
                 for partition in range(self.npartitions):
                     tp = ck.TopicPartition(self.topic, partition, 0)
                     try:
-                        low, high = consumer.get_watermark_offsets(tp,
-                                                                   timeout=0.1)
+                        low, high = self.consumer.get_watermark_offsets(
+                            tp, timeout=0.1)
                     except (RuntimeError, ck.KafkaException):
                         continue
                     current_position = self.positions[partition]
                     lowest = max(current_position, low)
-                    out.append((self.consumer_params, self.topic, partition,
-                                lowest, high - 1))
-                    self.positions[partition] = high
+                    if high > lowest:
+                        out.append((self.consumer_params, self.topic, partition,
+                                    lowest, high - 1))
+                        self.positions[partition] = high
 
                 for part in out:
                     yield self._emit(part)
 
                 else:
                     yield gen.sleep(self.poll_interval)
         finally:
-            consumer.close()
+            self.consumer.unsubscribe()
+            self.consumer.close()
 
     def start(self):
-        self.stopped = False
-        self.loop.add_callback(self.poll_kafka)
+        import confluent_kafka as ck
+        if self.stopped:
+            self.consumer = ck.Consumer(self.consumer_params)
+            self.stopped = False
+            tp = ck.TopicPartition(self.topic, 0, 0)
+
+            # blocks for consumer thread to come up
+            self.consumer.get_watermark_offsets(tp)
+            self.loop.add_callback(self.poll_kafka)
 
 
 @Stream.register_api(staticmethod)
@@ -352,7 +350,7 @@ def get_message_batch(kafka_params, topic, partition, low, high, timeout=None):
     try:
         while True:
             msg = consumer.poll(0)
-            if msg and msg.value():
+            if msg and msg.value() and msg.error() is None:
                 if high >= msg.offset():
                     out.append(msg.value())
                 if high <= msg.offset():

streamz/tests/test_core.py

@@ -942,8 +942,8 @@ def dont_test_stream_kwargs(clean):  # noqa: F811
     sin.emit(1)
 
 
-@pytest.fixture  # noqa: F811
-def thread(loop):
+@pytest.fixture
+def thread(loop):  # noqa: F811
     from threading import Thread, Event
     thread = Thread(target=loop.start)
     thread.daemon = True

streamz/tests/test_dask.py

@@ -73,8 +73,8 @@ def test_zip(c, s, a, b):
     assert L == [(1, 'a'), (2, 'b')]
 
 
-@pytest.mark.slow  # noqa: F811
-def test_sync(loop):
+@pytest.mark.slow
+def test_sync(loop):  # noqa: F811
     with cluster() as (s, [a, b]):
         with Client(s['address'], loop=loop) as client:  # noqa: F841
             source = Stream()
@@ -90,8 +90,8 @@ def f():
             assert L == list(map(inc, range(10)))
 
 
-@pytest.mark.slow  # noqa: F811
-def test_sync_2(loop):
+@pytest.mark.slow
+def test_sync_2(loop):  # noqa: F811
     with cluster() as (s, [a, b]):
         with Client(s['address'], loop=loop):  # noqa: F841
             source = Stream()
@@ -131,8 +131,8 @@ def test_buffer(c, s, a, b):
     assert source.loop == c.loop
 
 
-@pytest.mark.slow  # noqa: F811
-def test_buffer_sync(loop):
+@pytest.mark.slow
+def test_buffer_sync(loop):  # noqa: F811
     with cluster() as (s, [a, b]):
         with Client(s['address'], loop=loop) as c:  # noqa: F841
             source = Stream()
@@ -155,9 +155,9 @@ def test_buffer_sync(loop):
             assert L == list(map(inc, range(10)))
 
 
-@pytest.mark.xfail(reason='')  # noqa: F811
+@pytest.mark.xfail(reason='')
 @pytest.mark.slow
-def test_stream_shares_client_loop(loop):
+def test_stream_shares_client_loop(loop):  # noqa: F811
     with cluster() as (s, [a, b]):
         with Client(s['address'], loop=loop) as client:  # noqa: F841
             source = Stream()