diff options
Diffstat (limited to 'pysite')
-rw-r--r-- | pysite/constants.py | 21 | ||||
-rw-r--r-- | pysite/mixins.py | 95 | ||||
-rw-r--r-- | pysite/queues.py | 5 |
3 files changed, 121 insertions, 0 deletions
diff --git a/pysite/constants.py b/pysite/constants.py index f95e076f..a9040751 100644 --- a/pysite/constants.py +++ b/pysite/constants.py @@ -17,6 +17,16 @@ class ValidationTypes(Enum): params = "params" +class BotEventTypes(Enum): + mod_log = "mod_log" + + send_message = "send_message" + send_embed = "send_embed" + + add_role = "ensure_role" + remove_role = "remove_role" + + DEBUG_MODE = "FLASK_DEBUG" in environ # All snowflakes should be strings as RethinkDB rounds them as ints @@ -106,3 +116,14 @@ WIKI_AUDIT_WEBHOOK = environ.get("WIKI_AUDIT_WEBHOOK") or None # Bot key BOT_API_KEY = environ.get("BOT_API_KEY") or None + +# RabbitMQ settings +BOT_EVENT_QUEUE = "bot_events" + +RMQ_USERNAME = environ.get("RABBITMQ_DEFAULT_USER") or "guest" +RMQ_PASSWORD = environ.get("RABBITMQ_DEFAULT_PASS") or "guest" +RMQ_HOST = "pdrmq" if not DEBUG_MODE else "localhost" +RMQ_PORT = 5672 + +# Channels +CHANNEL_MOD_LOG = 282638479504965634 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 diff --git a/pysite/queues.py b/pysite/queues.py new file mode 100644 index 00000000..7a200208 --- /dev/null +++ b/pysite/queues.py @@ -0,0 +1,5 @@ +from kombu import Queue + +QUEUES = { # RabbitMQ Queue definitions, they'll be declared at gunicorn start time + "bot_events": Queue("bot_events", durable=True) +} |