wayfinderak/Geeksbot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Switch to Asyncpg for db con

Browse Source
This commit is contained in:
Dustin Pianalto 2018-05-21 19:13:40 -08:00
parent 249e5d94eb
commit 47ca65da91
3 changed files with 34 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
geeksbot.py
View File

@ -6,9 +6,8 @@ from datetime import datetime
import json
import aiohttp
from googleapiclient.discovery import build
import asyncpg
from concurrent import futures
import asyncio
from shared_libs import database
log_format = '{asctime}.{msecs:03.0f}|{levelname:<8}|{name}::{message}'
@ -58,14 +57,14 @@ class Geeksbot(commands.Bot):
self.infected = {}
self.TOKEN = self.bot_secrets['token']
async def connect_db():
return await asyncpg.create_pool(host=self.bot_secrets['db_con']['host'],
database=self.bot_secrets['db_con']['db_name'],
user=self.bot_secrets['db_con']['user'],
password=self.bot_secrets['db_con']['password'],
loop=asyncio.get_event_loop())
# async def connect_db():
# return await asyncpg.create_pool(host=self.bot_secrets['db_con']['host'],
# database=self.bot_secrets['db_con']['db_name'],
# user=self.bot_secrets['db_con']['user'],
# password=self.bot_secrets['db_con']['password'],
# loop=asyncio.get_event_loop())
del self.bot_secrets['token']
self.db_con = asyncio.get_event_loop().create_task(connect_db())
self.db_con = database.DatabaseConnection(**self.bot_secrets['db_con'])
self.default_prefix = 'g~'
self.voice_chans = {}
self.spam_list = {}
@ -80,6 +79,10 @@ class Geeksbot(commands.Bot):
'left_fist': '🤛',
}
async def logout(self):
await self.db_con.close()
super().logout()
@staticmethod
async def get_custom_prefix(bot_inst, message):
return await bot_inst.db_con.fetchval('select prefix from guild_config where guild_id = $1',

0
shared_libs/__init__.py Normal file
View File

22
shared_libs/database.py Normal file
View File

@ -0,0 +1,22 @@
import asyncpg
import asyncio
class DatabaseConnection:
def __init__(self, host: str='localhost', port: int=5432, database: str='', username: str='', password: str=''):
if username == '' or password == '' or database == '':
raise RuntimeError('Username or Password are blank')
self.kwargs = {'host': host, 'port': port, 'database': database, 'username': username, 'password': password}
self._conn = asyncio.get_event_loop().run_until_complete(self.acquire())
self.fetchval = self._conn.fetchval
self.execute = self._conn.execute
self.fetchall = self._conn.fetchall
self.fetchone = self._conn.fetchone
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