Add mqtt source/sink pair

Author: Martin Durant

ci/environment-py38.yml

@@ -3,7 +3,7 @@ channels:
   - conda-forge
   - defaults
 dependencies:
-  - python=3.7
+  - python=3.8
   - pytest
   - flake8
   - black
@@ -32,3 +32,5 @@ dependencies:
   - flaky
   - pytest-cov
   - coveralls
+  - paho-mqtt
+  - websockets

docs/source/api.rst

@@ -47,6 +47,7 @@ Stream
 .. automethod:: Stream.sink_to_list
 .. automethod:: Stream.sink_to_textfile
 .. automethod:: Stream.to_websocket
+.. automethod:: Stream.to_mqtt
 .. automethod:: Stream.update
 .. automethod:: Stream.visualize
 
@@ -58,6 +59,7 @@ Sources
    filenames
    from_kafka
    from_kafka_batched
+   from_mqtt
    from_process
    from_websocket
    from_textfile

streamz/sinks.py

@@ -134,17 +134,57 @@ def __init__(self, upstream, uri, ws_kwargs=None, **kwargs):
         self.uri = uri
         self.ws_kw = ws_kwargs
         self.ws = None
-        super().__init__(upstream, **kwargs)
+        super().__init__(upstream, ensure_io_loop=True, **kwargs)
 
-    @gen.coroutine
-    def update(self, x, who=None, metadata=None):
+    async def update(self, x, who=None, metadata=None):
         import websockets
         if self.ws is None:
-            self.ws = yield websockets.connect(self.uri, **(self.ws_kw or {}))
-        yield self.ws.send(x)
+            self.ws = await websockets.connect(self.uri, **(self.ws_kw or {}))
+        await self.ws.send(x)
 
     def destroy(self):
         super().destroy()
         if self.ws is not None:
             sync(self.loop, self.ws.protocol.close)
         self.ws = None
+
+
+@Stream.register_api()
+class to_mqtt(Sink):
+    """
+    Send data to MQTT broker
+
+    See also ``sources.from_mqtt``.
+
+    Requires ``paho.mqtt``
+
+    :param host: str
+    :param port: int
+    :param topic: str
+    :param keepalive: int
+        See mqtt docs - to keep the channel alive
+    :param client_kwargs:
+        Passed to the client's ``connect()`` method
+    """
+    def __init__(self, upstream, host, port, topic, keepalive=60, client_kwargs=None,
+                 **kwargs):
+        self.host = host
+        self.port = port
+        self.c_kw = client_kwargs or {}
+        self.client = None
+        self.topic = topic
+        self.keepalive = keepalive
+        super().__init__(upstream, ensure_io_loop=True, **kwargs)
+
+    def update(self, x, who=None, metadata=None):
+        import paho.mqtt.client as mqtt
+        if self.client is None:
+            self.client = mqtt.Client()
+            self.client.connect(self.host, self.port, self.keepalive, **self.c_kw)
+        # TODO: wait on successful delivery
+        self.client.publish(self.topic, x)
+
+    def destroy(self):
+        self.client.disconnect()
+        self.client = None
+        super().destroy()

streamz/sources.py

@@ -1,5 +1,6 @@
 import asyncio
 from glob import glob
+import queue
 import os
 import time
 from tornado import gen
@@ -822,3 +823,92 @@ async def run(self):
     def stop(self):
         self.server.close()
         sync(self.loop, self.server.wait_closed)
+
+
+@Stream.register_api()
+class from_q(Source):
+    """Source events from a threading.Queue, running another event framework
+
+    The queue is polled, i.e., there is a latency/overhead tradeoff, since
+    we cannot use ``await`` directly with a multithreaded queue.
+
+    Allows mixing of another event loop, for example pyqt, on another thread.
+    Note that, by default, a streamz.Source such as this one will start
+    an event loop in a new thread, unless otherwise specified.
+    """
+
+    def __init__(self, q, sleep_time=0.01, **kwargs):
+        """
+        :param q: threading.Queue
+            Any items pushed into here will become streamz events
+        :param sleep_time: int
+            Sets how long we wait before checking the input queue when
+            empty (in s)
+        :param kwargs:
+            passed to streamz.Source
+        """
+        self.q = q
+        self.sleep = sleep_time
+        super().__init__(**kwargs)
+
+    async def _run(self):
+        """Poll threading queue for events
+        This uses check-and-wait, but overhead is low. Could maybe have
+        a sleep-free version with an threading.Event.
+        """
+        try:
+            out = self.q.get_nowait()
+            await self.emit(out, asynchronous=True)
+        except queue.Empty:
+            await asyncio.sleep(self.sleep)
+
+
+@Stream.register_api()
+class from_mqtt(from_q):
+    """Read from MQTT source
+
+    See https://en.wikipedia.org/wiki/MQTT for a description of the protocol
+    and its uses.
+
+    See also ``sinks.to_mqtt``.
+
+    Requires ``paho.mqtt``
+
+    The outputs are ``paho.mqtt.client.MQTTMessage`` instances, which each have
+    attributes timestamp, payload, topic, ...
+
+    NB: paho.mqtt.python runs on its own thread in this implementation. We may
+    wish to instead call client.loop() directly
+
+    :param host: str
+    :param port: int
+    :param topic: str
+        (May in the future support a list of topics)
+    :param keepalive: int
+        See mqtt docs - to keep the channel alive
+    :param client_kwargs:
+        Passed to the client's ``connect()`` method
+    """
+    def __init__(self, host, port, topic, keepalive=60 , client_kwargs=None, **kwargs):
+        self.host = host
+        self.port = port
+        self.keepalive = keepalive
+        self.topic = topic
+        self.client_kwargs = client_kwargs
+        super().__init__(q=queue.Queue(), **kwargs)
+
+    def _on_connect(self, client, userdata, flags, rc):
+        client.subscribe(self.topic)
+
+    def _on_message(self, client, userdata, msg):
+        self.q.put(msg)
+
+    async def run(self):
+        import paho.mqtt.client as mqtt
+        client = mqtt.Client()
+        client.on_connect = self._on_connect
+        client.on_message = self._on_message
+        client.connect(self.host, self.port, self.keepalive, **(self.client_kwargs or {}))
+        client.loop_start()
+        await super().run()
+        client.disconnect()

streamz/tests/test_sinks.py

@@ -86,3 +86,20 @@ def test_ws_roundtrip():
     wait_for(lambda: data == l, timeout=1)
     s.stop()
     s0.stop()
+
+
+def test_mqtt_roundtrip():
+    pytest.importorskip("paho.mqtt.client")
+    s0 = Stream.from_mqtt("mqtt.eclipseprojects.io", 1883, "streamz/sensor/temperature")
+    l = s0.map(lambda msg: msg.payload).sink_to_list()
+    s0.start()
+
+    data = [b'0123'] * 4
+    s = Stream.from_iterable(data)
+    s.to_mqtt("mqtt.eclipseprojects.io", 1883, "streamz/sensor/temperature")
+    s.start()
+
+    wait_for(lambda: data == l, timeout=1)
+    s.stop()
+    s0.stop()
+