diff --git a/models/__init__.py b/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/models/quotes.py b/models/quotes.py new file mode 100644 index 0000000..9655f48 --- /dev/null +++ b/models/quotes.py @@ -0,0 +1,16 @@ +from settings.config import SQLACHEMY +from sqlalchemy import Table +from sqlalchemy.orm import declarative_base + +Base = declarative_base() + +class Quotes(Base): + """ + Quotes model class. + """ + __table__ = Table( + "neeble_quotes", + Base.metadata, + autoload=True, + autoload_with=SQLACHEMY + ) diff --git a/settings/config.py b/settings/config.py index 0b9258f..82fcfac 100644 --- a/settings/config.py +++ b/settings/config.py @@ -4,6 +4,8 @@ BOT configuration. import logging.config as logger import os +from sqlalchemy import create_engine + # Discord token. DISCORD_BOT_TOKEN = os.environ.get('DISCORD_BOT_TOKEN') @@ -16,6 +18,16 @@ MYSQL_CONFIG = { 'password': os.environ.get('MYSQL_PASSWORD', 'neeble'), } +SQLACHEMY = create_engine( + 'mysql://%s:%s@%s:%s/%s' % ( + os.environ.get('MYSQL_USER', 'neeble'), + os.environ.get('MYSQL_PASSWORD', 'neeble'), + os.environ.get('MYSQL_HOST', 'localhost'), + int(os.environ.get('MYSQL_PORT', '3306')), + os.environ.get('MYSQL_DATABASE', 'neeble'), + ) +) + # Define the log level LOGLEVEL = os.environ.get('LOGLEVEL', 'info').upper() diff --git a/utils/commands.py b/utils/commands.py index ce0c26c..dbe5a57 100644 --- a/utils/commands.py +++ b/utils/commands.py @@ -140,11 +140,9 @@ async def quote_count(bot: object) -> str: """ Outputs a quote count from the database """ - # For len(amount) to work, first it needs to be converted into str amount = count_quotes() amount = str(amount) - amount = amount[1:len(amount)][:-2] msg = "Quote count: `" + amount + "`" diff --git a/utils/database.py b/utils/database.py index f73a133..1997ff3 100644 --- a/utils/database.py +++ b/utils/database.py @@ -2,10 +2,12 @@ Database utils module. """ import logging -from collections import namedtuple import MySQLdb -from settings.config import MYSQL_CONFIG +from models.quotes import Quotes +from settings.config import MYSQL_CONFIG, SQLACHEMY +from sqlalchemy import select +from sqlalchemy.orm import Session logger = logging.getLogger(__name__) @@ -61,79 +63,51 @@ def set_quote(user: str, quote: str) -> None: """ Set a quote into database. """ - _sql = f''' - insert into neeble_quotes(user, quote) - value("{user}", "{quote}"); - ''' - with Cursor(MYSQL_CONFIG) as cursor: - cursor.execute(_sql) + with Session(SQLACHEMY) as session: + session.add(Quotes( + quote=quote, + user=user, + )) + session.commit() + def get_quotes(ids: list) -> tuple: """ Get the saved quotes. ids: List of quote ID's """ - _sql = f''' - select quote, user, id - from neeble_quotes - ''' - _sql = _sql + f' where id not in ({",".join([str(id) for id in ids])});' if ids else _sql + ';' - response = [] - obj = namedtuple('Quotes', ['quote', 'user', 'id']) - - with Cursor(MYSQL_CONFIG) as cursor: - cursor.execute(_sql) - response = cursor.fetchall() - - return tuple(obj(*r) for r in response) + with Session(SQLACHEMY) as session: + _sql = select(Quotes).where(Quotes.id.not_in(ids)) + return [item for item in session.scalars(_sql)] def get_by_id(id: int) -> object: """ Get one quote by ID. """ - obj = namedtuple('Quotes', ['quote', 'user', 'id']) - _sql = f''' - select quote, user, id - from neeble_quotes - where id={id}; - ''' - - with Cursor(MYSQL_CONFIG) as cursor: - cursor.execute(_sql) - quote = cursor.fetchone() - - if not quote: - return None + with Session(SQLACHEMY) as session: + result = [s for s in session.query(Quotes).filter(Quotes.id==id)] + return result[0] if result else None - return obj(*quote) def remove_quote(_id: int) -> bool: """ Delete one quote by database ID. """ - _sql = f''' - delete from neeble_quotes - where id={_id}; - ''' - try: - with Cursor(MYSQL_CONFIG) as cursor: - cursor.execute(_sql) + with Session(SQLACHEMY) as session: + item = get_by_id(_id) + session.delete(item) + session.commit() return True except Exception: return False + def count_quotes() -> int: """ Counts the amount of quotes in the database """ - _sql = f''' - select count(*) from neeble_quotes - ''' - - with Cursor(MYSQL_CONFIG) as cursor: - cursor.execute(_sql) - count = cursor.fetchone() - - return count \ No newline at end of file + with Session(SQLACHEMY) as session: + response = session.query(Quotes.id).count() + return response