Validate database connection before start

Author: kyb3r

bot.py

@@ -62,6 +62,7 @@ def __init__(self):
         self.session = aiohttp.ClientSession(loop=self.loop)
         self.config = ConfigManager(self)
         self.selfhosted = bool(self.config.get('mongo_uri'))
+        self._connected = asyncio.Event()
         if self.selfhosted:
             self.db = AsyncIOMotorClient(self.config.mongo_uri).modmail_bot
         self.modmail_api = SelfhostedClient(self) if self.selfhosted else ModmailApiClient(self)
@@ -198,6 +199,7 @@ async def on_connect(self):
             await self.validate_api_token()
             print(line)
         else:
+            await self.validate_database_connection()
             print('Mode: Selfhosting logs.')
             print(line)
         print(Fore.CYAN + 'Connected to gateway.')
@@ -212,9 +214,14 @@ async def on_connect(self):
             activity = discord.Activity(type=activity_type, name=message,
                                         url=url)
             await self.change_presence(activity=activity)
+        
+        self._connected.set()
 
     async def on_ready(self):
         """Bot startup, sets uptime."""
+
+        await self._connected.wait()
+
         print(textwrap.dedent(f"""
         {line}
         {Fore.CYAN}Client ready.
@@ -426,7 +433,6 @@ def overwrites(self, ctx):
         return overwrites
 
     async def validate_api_token(self):
-        valid = True
         try:
             self.config.modmail_api_token
         except KeyError:
@@ -439,7 +445,7 @@ async def validate_api_token(self):
         else:
             valid = await self.modmail_api.validate_token()
             if not valid:
-                print(Fore.RED + Style.BRIGHT, end='')
+                print(Fore.RED, end='')
                 print('Invalid MODMAIL_API_TOKEN - get one from https://dashboard.modmail.tk')
         finally:
             if not valid:
@@ -448,6 +454,21 @@ async def validate_api_token(self):
                 username = (await self.modmail_api.get_user_info())['user']['username']
                 print(Style.RESET_ALL + Fore.CYAN + 'Validated token.' )
                 print(f'GitHub user: {username}' + Style.RESET_ALL)
+    
+    async def validate_database_connection(self):
+        try:
+            doc = await self.db.command('buildinfo')
+        except Exception as e:
+            valid = False
+            print(Fore.RED, end='')
+            print('Something went wrong while connecting to the database.')
+            print(type(e).__name__, e, sep=': ')
+        else:
+            valid = True
+            print(Style.RESET_ALL + Fore.CYAN + 'Successfully connected to the database.' )
+        finally:
+            if not valid:
+                await self.logout()
 
     async def data_loop(self):
         await self.wait_until_ready()
@@ -479,11 +500,13 @@ async def autoupdate_loop(self):
 
         if self.config.get('disable_autoupdates'):
             print('Autoupdates disabled.')
+            print(line)
             return 
 
         if self.selfhosted and not self.config.get('github_access_token'):
             print('Github access token not found.')
             print('Autoupdates disabled.')
+            print(line)
             return 
 
         while True: