pyceo/ceod/db/MySQLService.py

93 lines
3.5 KiB
Python

from zope.interface import implementer
from zope import component
from contextlib import contextmanager
from ceo_common.interfaces import IDatabaseService, IConfig
from ceo_common.errors import DatabaseConnectionError, DatabasePermissionError, UserAlreadyExistsError, \
UserNotFoundError
from ceo_common.logger_factory import logger_factory
from ceod.utils import gen_password
from ceod.db.utils import response_is_empty
from mysql.connector import connect
from mysql.connector.errors import OperationalError, ProgrammingError
logger = logger_factory(__name__)
@implementer(IDatabaseService)
class MySQLService:
type = 'mysql'
def __init__(self):
config = component.getUtility(IConfig)
self.auth_username = config.get('mysql_username')
self.auth_password = config.get('mysql_password')
self.host = config.get('mysql_host')
@contextmanager
def mysql_connection(self):
try:
with connect(
host=self.host,
user=self.auth_username,
password=self.auth_password,
) as con:
yield con
# unable to connect
except OperationalError as e:
logger.error(e)
raise DatabaseConnectionError()
# invalid credentials / user does not exist / invalid permissions for action
except ProgrammingError as e:
logger.error(e)
raise DatabasePermissionError()
def create_db(self, username: str) -> str:
password = gen_password()
search_for_user = f"SELECT user FROM mysql.user WHERE user='{username}'"
search_for_db = f"SHOW DATABASES LIKE '{username}'"
# CREATE USER can't be used in a query with multiple statements
create_local_user_cmd = f"CREATE USER '{username}'@'localhost' IDENTIFIED VIA unix_socket"
create_user_cmd = f"CREATE USER '{username}'@'%' IDENTIFIED BY %(password)s"
create_database = f"""
CREATE DATABASE {username};
GRANT ALL PRIVILEGES ON {username}.* TO '{username}'@'localhost' IDENTIFIED VIA unix_socket;
GRANT ALL PRIVILEGES ON {username}.* TO '{username}'@'%';
"""
with self.mysql_connection() as con, con.cursor() as cursor:
if not response_is_empty(search_for_user, con):
raise UserAlreadyExistsError()
cursor.execute(create_local_user_cmd)
cursor.execute(create_user_cmd, {'password': password})
if response_is_empty(search_for_db, con):
cursor.execute(create_database)
return password
def reset_db_passwd(self, username: str) -> str:
password = gen_password()
search_for_user = f"SELECT user FROM mysql.user WHERE user='{username}'"
reset_password = f"""
ALTER USER '{username}'@'localhost' IDENTIFIED BY %(password)s;
ALTER USER '{username}'@'%' IDENTIFIED BY %(password)s;
"""
with self.mysql_connection() as con, con.cursor() as cursor:
if response_is_empty(search_for_user, con):
raise UserNotFoundError(username)
cursor.execute(reset_password, {'password': password})
return password
def delete_db(self, username: str):
drop_db = f"DROP DATABASE IF EXISTS {username}"
drop_user = f"""
DROP USER IF EXISTS '{username}'@'localhost';
DROP USER IF EXISTS '{username}'@'%';
"""
with self.mysql_connection() as con, con.cursor() as cursor:
cursor.execute(drop_db)
cursor.execute(drop_user)