[DataStore] closes #7

Make Firebase calls asynchronous by
running them in a ThreadPoolExecutor:
firebase/firebase-admin-python#104 (comment)

This touches a lot of code, because a new formatter was
introduced as well.

Author: noahkw

DataStore.py

@@ -1,74 +1,93 @@
 from abc import ABC, abstractmethod
+from functools import partial
+from firebase_admin import credentials, firestore
 
 import firebase_admin
-from firebase_admin import credentials, firestore
+import concurrent
+import asyncio
 
 
 class DataStore(ABC):
     def __init__(self):
         pass
 
     @abstractmethod
-    def set(self, collection, document, val):
+    async def set(self, collection, document, val):
+        pass
+
+    @abstractmethod
+    async def set_get_id(self, collection, val):
         pass
 
     @abstractmethod
-    def update(self, collection, document, val):
+    async def update(self, collection, document, val):
         pass
 
     @abstractmethod
-    def add(self, collection, val):
+    async def add(self, collection, val):
         pass
 
     @abstractmethod
-    def get(self, collection, document):
+    async def get(self, collection, document):
         pass
 
     @abstractmethod
-    def delete(self, collection, document):
+    async def delete(self, collection, document):
         pass
 
 
 class FirebaseDataStore(DataStore):
-    def __init__(self, key_file, db_name):
+    def __init__(self, key_file, db_name, loop):
         super().__init__()
         cred = credentials.Certificate(key_file)
 
-        firebase_admin.initialize_app(cred, {
-            'databaseURL': f'https://{db_name}.firebaseio.com'
-        })
+        firebase_admin.initialize_app(
+            cred, {'databaseURL': f'https://{db_name}.firebaseio.com'})
 
         self.db = firestore.client()
+        self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=8)
+        self.loop = loop
+
+    async def set(self, collection, document, val):
+        ref = self._get_doc_ref(collection, document)
+        await self.loop.run_in_executor(self.executor, partial(ref.set, val))
+
+    async def set_get_id(self, collection, val):
+        ref = self._get_doc_ref(collection, None)
+        await self.loop.run_in_executor(self.executor, partial(ref.set, val))
+        return ref.id
+
+    async def update(self, collection, document, val):
+        ref = self._get_doc_ref(collection, document)
+        await self.loop.run_in_executor(self.executor,
+                                        partial(ref.update, val))
+
+    async def add(self, collection, val):
+        ref = self._get_collection(collection)
+        await self.loop.run_in_executor(self.executor, partial(ref.add, val))
+
+    async def get(self, collection, document=None):
+        if document is None:
+            ref = self._get_collection(collection)
+            return await self.loop.run_in_executor(self.executor, ref.stream)
+        else:
+            ref = self._get_doc_ref(collection, document)
+            return await self.loop.run_in_executor(self.executor, ref.get)
 
-    def set(self, collection, document, val):
-        self._get_doc_ref(collection, document).set(val)
-
-    def set_get_id(self, collection, val):
-        doc = self._get_doc_ref(collection, None)
-        doc.set(val)
-        return doc.id
-
-    def update(self, collection, document, val):
-        self._get_doc_ref(collection, document).update(val)
-
-    def add(self, collection, val):
-        self._get_collection(collection).add(val)
-
-    def get(self, collection, document=None):
-        return self._get_collection(collection).stream() if document is None else self._get_doc_ref(collection,
-                                                                                                    document).get()
-
-    def delete(self, collection, document=None):
+    async def delete(self, collection, document=None):
         if document is not None:
-            self._get_doc_ref(collection, document).delete()
+            ref = self._get_doc_ref(collection, document)
+            await self.loop.run_in_executor(self.executor, ref.delete)
         else:
             # implement batching later
             docs = self._get_collection(collection).stream()
             for doc in docs:
-                doc.reference.delete()
+                await self.loop.run_in_executor(self.executor,
+                                                doc.reference.delete)
 
-    def query(self, collection, *query):
-        return self._get_collection(collection).where(*query).stream()
+    async def query(self, collection, *query):
+        ref = self._get_collection(collection).where(*query)
+        return await self.loop.run_in_executor(self.executor, ref.stream)
 
     def _get_doc_ref(self, collection, document):
         return self._get_collection(collection).document(document)
@@ -83,6 +102,10 @@ def _get_collection(self, collection):
     config = configparser.ConfigParser()
     config.read('conf.ini')
 
-    firebase_ds = FirebaseDataStore(
-        config['firebase']['key_file'], config['firebase']['db_name'])
-    firebase_ds.add('jobs', {'func': 'somefunc', 'time': 234903284, 'args': ['arg1', 'arg2']})
+    firebase_ds = FirebaseDataStore(config['firebase']['key_file'],
+                                    config['firebase']['db_name'])
+    firebase_ds.add('jobs', {
+        'func': 'somefunc',
+        'time': 234903284,
+        'args': ['arg1', 'arg2']
+    })

botw-bot.py

@@ -10,30 +10,29 @@
 config.read('conf.ini')
 client = commands.Bot(command_prefix=config['discord']['command_prefix'])
 client.config = config
-client.database = FirebaseDataStore(
-    config['firebase']['key_file'], config['firebase']['db_name'])
+client.db = FirebaseDataStore(config['firebase']['key_file'],
+                              config['firebase']['db_name'], client.loop)
 logger = logging.getLogger('discord')
 logger.setLevel(logging.INFO)
-handler = logging.FileHandler(
-    filename='botw-bot.log', encoding='utf-8', mode='w')
-handler.setFormatter(logging.Formatter(
-    '%(asctime)s:%(levelname)s:%(name)s: %(message)s'))
+handler = logging.FileHandler(filename='botw-bot.log',
+                              encoding='utf-8',
+                              mode='w')
+handler.setFormatter(
+    logging.Formatter('%(asctime)s:%(levelname)s:%(name)s: %(message)s'))
 logger.addHandler(handler)
 
 INITIAL_EXTENSIONS = [
-    'cogs.BiasOfTheWeek',
-    'cogs.Utilities',
-    'cogs.Scheduler',
-    'cogs.EmojiUtils',
-    'cogs.Tags',
-    'jishaku'
+    'cogs.BiasOfTheWeek', 'cogs.Utilities', 'cogs.Scheduler',
+    'cogs.EmojiUtils', 'cogs.Tags', 'jishaku'
 ]
 
 
 @client.event
 async def on_ready():
     await client.change_presence(activity=discord.Game('with Bini'))
-    logger.info(f"Logged in as {client.user}. Whitelisted servers: {config.items('whitelisted_servers')}")
+    logger.info(
+        f"Logged in as {client.user}. Whitelisted servers: {config.items('whitelisted_servers')}"
+    )
 
     for ext in INITIAL_EXTENSIONS:
         ext_logger = logging.getLogger(ext)
@@ -54,7 +53,9 @@ async def globally_block_dms(ctx):
 
 @client.check
 async def whitelisted_server(ctx):
-    server_ids = [int(server) for key, server in config.items('whitelisted_servers')]
+    server_ids = [
+        int(server) for key, server in config.items('whitelisted_servers')
+    ]
     return ctx.guild.id in server_ids
 
 

cogs/BiasOfTheWeek.py

@@ -9,7 +9,6 @@
 
 from cogs.Scheduler import Job
 
-
 logger = logging.getLogger(__name__)
 
 
@@ -28,13 +27,11 @@ def __str__(self):
     def __eq__(self, other):
         if not isinstance(other, Idol):
             return NotImplemented
-        return str.lower(self.group) == str.lower(other.group) and str.lower(self.name) == str.lower(other.name)
+        return str.lower(self.group) == str.lower(other.group) and str.lower(
+            self.name) == str.lower(other.name)
 
     def to_dict(self):
-        return {
-            'group': self.group,
-            'name': self.name
-        }
+        return {'group': self.group, 'name': self.name}
 
     @staticmethod
     def from_dict(source):
@@ -45,47 +42,65 @@ class BiasOfTheWeek(commands.Cog):
     def __init__(self, bot):
         self.bot = bot
         self.nominations = {}
-        self.nominations_collection = self.bot.config['biasoftheweek']['nominations_collection']
+        self.nominations_collection = self.bot.config['biasoftheweek'][
+            'nominations_collection']
+
+        if self.bot.loop.is_running():
+            asyncio.create_task(self._ainit())
+        else:
+            self.bot.loop.run_until_complete(self._ainit())
+
+    async def _ainit(self):
+        _nominations = await self.bot.db.get(self.nominations_collection)
 
-        _nominations = self.bot.database.get(self.nominations_collection)
         for nomination in _nominations:
-            self.nominations[self.bot.get_user(int(nomination.id))] = Idol.from_dict(nomination.to_dict())
+            self.nominations[self.bot.get_user(int(
+                nomination.id))] = Idol.from_dict(nomination.to_dict())
 
-        logger.info(f'Initial nominations from database: {self.nominations}')
+        logger.info(f'Initial nominations from db: {self.nominations}')
 
     @staticmethod
     def reaction_check(reaction, user, author, prompt_msg):
         return user == author and str(reaction.emoji) in [CHECK_EMOJI, CROSS_EMOJI] and \
                reaction.message.id == prompt_msg.id
 
     @commands.command()
-    async def nominate(self, ctx, group: commands.clean_content, name: commands.clean_content):
+    async def nominate(self, ctx, group: commands.clean_content,
+                       name: commands.clean_content):
         idol = Idol(group, name)
 
         if idol in self.nominations.values():
-            await ctx.send(f'**{idol}** has already been nominated. Please nominate someone else.')
+            await ctx.send(
+                f'**{idol}** has already been nominated. Please nominate someone else.'
+            )
         elif ctx.author in self.nominations.keys():
             old_idol = self.nominations[ctx.author]
-            prompt_msg = await ctx.send(f'Your current nomination is **{old_idol}**. Do you want to override it?')
+            prompt_msg = await ctx.send(
+                f'Your current nomination is **{old_idol}**. Do you want to override it?'
+            )
             await prompt_msg.add_reaction(CHECK_EMOJI)
             await prompt_msg.add_reaction(CROSS_EMOJI)
             try:
-                reaction, user = await self.bot.wait_for('reaction_add', timeout=60.0,
-                                                         check=lambda reaction, user: self.reaction_check(reaction,
-                                                                                                          user,
-                                                                                                          ctx.author,
-                                                                                                          prompt_msg))
+                reaction, user = await self.bot.wait_for(
+                    'reaction_add',
+                    timeout=60.0,
+                    check=lambda reaction, user: self.reaction_check(
+                        reaction, user, ctx.author, prompt_msg))
             except asyncio.TimeoutError:
                 pass
             else:
                 await prompt_msg.delete()
                 if reaction.emoji == CHECK_EMOJI:
                     self.nominations[ctx.author] = idol
-                    self.bot.database.set(self.nominations_collection, str(ctx.author.id), idol.to_dict())
-                    await ctx.send(f'{ctx.author} nominates **{idol}** instead of **{old_idol}**.')
+                    await self.bot.db.set(self.nominations_collection,
+                                          str(ctx.author.id), idol.to_dict())
+                    await ctx.send(
+                        f'{ctx.author} nominates **{idol}** instead of **{old_idol}**.'
+                    )
         else:
             self.nominations[ctx.author] = idol
-            self.bot.database.set(self.nominations_collection, str(ctx.author.id), idol.to_dict())
+            await self.bot.db.set(self.nominations_collection,
+                                  str(ctx.author.id), idol.to_dict())
             await ctx.send(f'{ctx.author} nominates **{idol}**.')
 
     @nominate.error
@@ -96,7 +111,7 @@ async def nominate_error(self, ctx, error):
     @commands.has_permissions(administrator=True)
     async def clear_nominations(self, ctx):
         self.nominations = {}
-        self.bot.database.delete(self.nominations_collection)
+        await self.bot.db.delete(self.nominations_collection)
         await ctx.message.add_reaction(CHECK_EMOJI)
 
     @commands.command()
@@ -110,23 +125,43 @@ async def nominations(self, ctx):
         else:
             await ctx.send('So far, no idols have been nominated.')
 
+    @commands.command()
+    async def db_noms(self, ctx):
+        embed = discord.Embed(title='Bias of the Week nominations')
+        nominations = {}
+        _nominations = await self.bot.db.get(self.nominations_collection)
+        for nomination in _nominations:
+            nominations[self.bot.get_user(int(
+                nomination.id))] = Idol.from_dict(nomination.to_dict())
+
+        for key, value in nominations.items():
+            embed.add_field(name=key, value=value)
+
+        await ctx.send(embed=embed)
+
     @commands.command(name='pickwinner')
     @commands.has_permissions(administrator=True)
-    async def pick_winner(self, ctx, silent: bool = False, fast_assign: bool = False):
+    async def pick_winner(self,
+                          ctx,
+                          silent: bool = False,
+                          fast_assign: bool = False):
         member, pick = random.choice(list(self.nominations.items()))
 
         # Assign BotW winner role on next wednesday at 00:00 UTC
         now = pendulum.now('Europe/London')
-        assign_date = now.add(seconds=120) if fast_assign else now.next(
-            pendulum.WEDNESDAY)
+        assign_date = now.add(
+            seconds=120) if fast_assign else now.next(pendulum.WEDNESDAY)
 
         await ctx.send(
             f"""Bias of the Week ({now.week_of_year}-{now.year}): {member if silent else member.mention}\'s pick **{pick}**. 
-You will be assigned the role *{self.bot.config['biasoftheweek']['winner_role_name']}* at {assign_date.to_cookie_string()}.""")
+You will be assigned the role *{self.bot.config['biasoftheweek']['winner_role_name']}* at {assign_date.to_cookie_string()}."""
+        )
 
         scheduler = self.bot.get_cog('Scheduler')
         if scheduler is not None:
-            await scheduler.add_job(Job('assign_winner_role', [ctx.guild.id, member.id], assign_date.float_timestamp))
+            await scheduler.add_job(
+                Job('assign_winner_role', [ctx.guild.id, member.id],
+                    assign_date.float_timestamp))
 
     @pick_winner.error
     async def pick_winner_error(self, ctx, error):

cogs/EmojiUtils.py

@@ -20,8 +20,10 @@ class EmojiUtils(commands.Cog):
 
     def __init__(self, bot):
         self.bot = bot
-        self.bot.add_listener(self.on_guild_emojis_update, 'on_guild_emojis_update')
-        self.emoji_channel = self.bot.get_channel(int(self.bot.config['emojiutils']['emoji_channel']))
+        self.bot.add_listener(self.on_guild_emojis_update,
+                              'on_guild_emojis_update')
+        self.emoji_channel = self.bot.get_channel(
+            int(self.bot.config['emojiutils']['emoji_channel']))
 
     @commands.group(name='emoji')
     @commands.has_permissions(administrator=True)
@@ -41,16 +43,22 @@ async def emoji_list_error(self, ctx, error):
     async def on_guild_emojis_update(self, guild, before, after):
         # delete old messages containing emoji
         # need to use Message.delete to be able to delete messages older than 14 days
-        async for message in self.emoji_channel.history(limit=EmojiUtils.DELETE_LIMIT):
+        async for message in self.emoji_channel.history(
+                limit=EmojiUtils.DELETE_LIMIT):
             await message.delete()
 
         # get emoji that were added in the last 10 minutes
-        recent_emoji = [emoji for emoji in after if (
-                time.time() - discord.utils.snowflake_time(emoji.id).timestamp()) < EmojiUtils.NEW_EMOTE_THRESHOLD]
+        recent_emoji = [
+            emoji for emoji in after
+            if (time.time() -
+                discord.utils.snowflake_time(emoji.id).timestamp()
+                ) < EmojiUtils.NEW_EMOTE_THRESHOLD
+        ]
 
         emoji_sorted = sorted(after, key=lambda e: e.name)
         for emoji_chunk in chunker(emoji_sorted, EmojiUtils.SPLIT_MSG_AFTER):
             await self.emoji_channel.send(''.join(str(e) for e in emoji_chunk))
 
         if len(recent_emoji) > 0:
-            await self.emoji_channel.send(f"Newly added: {''.join(str(e) for e in recent_emoji)}")
+            await self.emoji_channel.send(
+                f"Newly added: {''.join(str(e) for e in recent_emoji)}")