diff --git a/src/__main__.py b/src/__main__.py index a883bd2..03c6ca2 100644 --- a/src/__main__.py +++ b/src/__main__.py @@ -19,6 +19,7 @@ from discord.ext import commands from src.config.config import LoadConfig from src.shared_libs.loggable import Loggable from src.shared_libs.ioutils import in_here +from src.shared_libs import database # Init logging to output on INFO level to stderr. @@ -52,6 +53,9 @@ class SebiMachine(commands.Bot, LoadConfig, Loggable): # Initialize and attach config / settings LoadConfig.__init__(self) commands.Bot.__init__(self, command_prefix=self.defaultprefix) + with open(in_here('config', 'PrivateConfig.json')) as fp: + self.bot_secrets = json.load(fp) + self.db_con = database.DatabaseConnection(**self.bot_secrets['db_con']) # Load plugins # Add your cog file name in this list @@ -134,11 +138,5 @@ class SebiMachine(commands.Bot, LoadConfig, Loggable): client = SebiMachine() -# Make sure the key stays private. -# I am 99% certain this is valid! -with open(in_here('config', 'PrivateConfig.json')) as fp: - PrivateConfig = json.load(fp) -if PrivateConfig["bot-key"] == '': - PrivateConfig["bot-key"] = os.getenv('botkey') -client.run(PrivateConfig["bot-key"]) +client.run(client.bot_secrets["bot-key"]) diff --git a/src/shared_libs/database.py b/src/shared_libs/database.py new file mode 100644 index 0000000..edcc655 --- /dev/null +++ b/src/shared_libs/database.py @@ -0,0 +1,23 @@ +import asyncpg +import asyncio + + +class DatabaseConnection: + def __init__(self, host: str='localhost', port: int=5432, database: str='', user: str='', password: str=''): + if user == '' or password == '' or database == '': + raise RuntimeError('Username or Password are blank') + self.kwargs = {'host': host, 'port': port, 'database': database, 'user': user, 'password': password} + self._conn = None + asyncio.get_event_loop().run_until_complete(self.acquire()) + self.fetchval = self._conn.fetchval + self.execute = self._conn.execute + self.fetch = self._conn.fetch + self.fetchrow = self._conn.fetchrow + + async def acquire(self): + if not self._conn: + self._conn = await asyncpg.create_pool(**self.kwargs) + + async def close(self): + await self._conn.close() + self._conn = None