diff options
-rw-r--r-- | bot/__main__.py | 1 | ||||
-rw-r--r-- | bot/cogs/rmq.py | 55 | ||||
-rw-r--r-- | bot/constants.py | 9 | ||||
-rw-r--r-- | config-default.yml | 7 |
4 files changed, 72 insertions, 0 deletions
diff --git a/bot/__main__.py b/bot/__main__.py index c4e42dd45..f6e5dfe0c 100644 --- a/bot/__main__.py +++ b/bot/__main__.py @@ -38,6 +38,7 @@ bot.http_session = ClientSession( # Internal/debug bot.load_extension("bot.cogs.logging") +bot.load_extension("bot.cogs.rmq") bot.load_extension("bot.cogs.security") bot.load_extension("bot.cogs.events") diff --git a/bot/cogs/rmq.py b/bot/cogs/rmq.py new file mode 100644 index 000000000..0140f0476 --- /dev/null +++ b/bot/cogs/rmq.py @@ -0,0 +1,55 @@ +import logging + +import aio_pika +from discord.ext.commands import AutoShardedBot + +from bot.constants import RabbitMQ + + +log = logging.getLogger(__name__) + + +class RMQ: + """ + RabbitMQ event handling + """ + + rmq = None # type: aio_pika.Connection + channel = None # type: aio_pika.Channel + queue = None # type: aio_pika.Queue + + def __init__(self, bot: AutoShardedBot): + self.bot = bot + + async def on_ready(self): + self.rmq = await aio_pika.connect_robust( + host=RabbitMQ.host, port=RabbitMQ.port, login=RabbitMQ.username, password=RabbitMQ.password + ) + + log.debug("Connected to RabbitMQ") + + self.channel = await self.rmq.channel() + self.queue = await self.channel.declare_queue("bot_events", durable=True) + + log.debug("Channel opened, queue declared") + + async for message in self.queue: + with message.process(): + await self.handle_message(message, message.body.decode()) + + async def send_message(self, data): + await self.channel.default_exchange.publish( + aio_pika.Message( + body=data.encode() + ), + routing_key=self.queue.name + ) + + async def handle_message(self, message, data): + log.debug(f"Message: {message}") + log.debug(f"Data: {data}") + + +def setup(bot): + bot.add_cog(RMQ(bot)) + log.info("Cog loaded: RMQ") diff --git a/bot/constants.py b/bot/constants.py index cbad02ad1..50a316d26 100644 --- a/bot/constants.py +++ b/bot/constants.py @@ -240,6 +240,15 @@ class Papertrail(metaclass=YAMLGetter): port: str +class RabbitMQ(metaclass=YAMLGetter): + section = "rabbitmq" + + host: str + port: int + username: str + password: str + + class URLs(metaclass=YAMLGetter): section = "urls" diff --git a/config-default.yml b/config-default.yml index 5670ac1da..eda77c2d0 100644 --- a/config-default.yml +++ b/config-default.yml @@ -60,6 +60,13 @@ papertrail: port: !ENV 'PAPERTRAIL_PORT' +rabbitmq: + host: "pdrmq" + port: 5672 + username: !ENV "RABBITMQ_DEFAULT_USER" + password: !ENV "RABBITMQ_DEFAULT_PASS" + + urls: deploy: !ENV 'DEPLOY_URL' status: !ENV 'STATUS_URL' |