RabbitMQ mixin, powered by Kombu (#84)

* [RMQ] Add Kombi an an RMQMixin, as well as some constants

* [RMQ] Fix example in mixin docstring

* Update Pipfile.lock - for some reason, pipenv didn't lock kombu

1 files changed, 95 insertions, 0 deletions

diff --git a/pysite/mixins.py b/pysite/mixins.py
index d0e822bf..6b5f7187 100644
--- a/pysite/mixins.py
+++ b/pysite/mixins.py
@@ -1,8 +1,14 @@
+from typing import Any, Dict
 from weakref import ref
 
 from flask import Blueprint
+from kombu import Connection
 from rethinkdb.ast import Table
 
+from pysite.constants import (
+    BOT_EVENT_QUEUE, BotEventTypes,
+    RMQ_HOST, RMQ_PASSWORD, RMQ_PORT, RMQ_USERNAME
+)
 from pysite.database import RethinkDB
 from pysite.oauth import OAuthBackend
 
@@ -58,6 +64,95 @@ class DBMixin:
         return self._db()
 
 
+class RMQMixin:
+    """
+    Mixin for classes that make use of RabbitMQ. It allows routes to send JSON-encoded messages to specific RabbitMQ
+    queues.
+
+    This class is intended to be mixed in alongside one of the other view classes. For example:
+
+    >>> class MyView(APIView, RMQMixin):
+    ...     name = "my_view"  # Flask internal name for this route
+    ...     path = "/my_view"  # Actual URL path to reach this route
+    ...     queue_name = "my_queue"  # Name of the RabbitMQ queue to send on
+
+    Note that the queue name is optional if all you want to do is send bot events.
+
+    This class will also work with Websockets:
+
+    >>> class MyWebsocket(WS, RMQMixin):
+    ...     name = "my_websocket"
+    ...     path = "/my_websocket"
+    ...     queue_name = "my_queue"
+    """
+
+    queue_name = ""
+
+    @classmethod
+    def setup(cls: "RMQMixin", manager: "pysite.route_manager.RouteManager", blueprint: Blueprint):
+        """
+        Set up the view by calling `super().setup()` as appropriate.
+
+        :param manager: Instance of the current RouteManager (used to get a handle for the database object)
+        :param blueprint: Current Flask blueprint
+        """
+
+        if hasattr(super(), "setup"):
+            super().setup(manager, blueprint)  # pragma: no cover
+
+    @property
+    def rmq_connection(self) -> Connection:
+        """
+        Get a Kombu AMQP connection object - use this in a context manager so that it gets closed after you're done
+
+        If you're just trying to send a message, check out `rmq_send` and `rmq_bot_event` instead.
+        """
+
+        return Connection(hostname=RMQ_HOST, userid=RMQ_USERNAME, password=RMQ_PASSWORD, port=RMQ_PORT)
+
+    def rmq_send(self, data: Dict[str, Any], routing_key: str = None):
+        """
+        Send some data to the RabbitMQ queue
+
+        >>> self.rmq_send({
+        ...     "text": "My hovercraft is full of eels!",
+        ...     "source": "Dirty Hungarian Phrasebook"
+        ... })
+        ...
+
+        This will be delivered to the queue immediately.
+        """
+
+        if routing_key is None:
+            routing_key = self.queue_name
+
+        with self.rmq_connection as c:
+            producer = c.Producer()
+            producer.publish(data, routing_key=routing_key)
+
+    def rmq_bot_event(self, event_type: BotEventTypes, data: Dict[str, Any]):
+        """
+        Send an event to the queue responsible for delivering events to the bot
+
+        >>> self.rmq_bot_event(BotEventTypes.send_message, {
+        ...     "channel": CHANNEL_MOD_LOG,
+        ...     "message": "This is a plain-text message for @everyone, from the site!"
+        ... })
+        ...
+
+        This will be delivered to the bot and actioned immediately, or when the bot comes online if it isn't already
+        connected.
+        """
+
+        if not isinstance(event_type, BotEventTypes):
+            raise ValueError("`event_type` must be a member of the the `pysite.constants.BotEventTypes` enum")
+
+        return self.rmq_send(
+            {"event": event_type.value, "data": data},
+            routing_key=BOT_EVENT_QUEUE,
+        )
+
+
 class OAuthMixin:
     """
     Mixin for the classes that need access to a logged in user's information. This class should be used