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 psycopg2 import connect, OperationalError, ProgrammingError from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT logger = logger_factory(__name__) @implementer(IDatabaseService) class PostgreSQLService: type = 'postgresql' def __init__(self): config = component.getUtility(IConfig) self.auth_username = config.get('postgresql_username') self.auth_password = config.get('postgresql_password') self.host = config.get('postgresql_host') @contextmanager def psql_connection(self): con = None try: # Don't use the connection as a context manager, because that # creates a new transaction. con = connect( host=self.host, user=self.auth_username, password=self.auth_password, ) con.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT) yield con # unable to connect / invalid credentials / user does not exist except OperationalError as e: logger.error(e) raise DatabaseConnectionError() # invalid permissions for action except ProgrammingError as e: logger.error(e) raise DatabasePermissionError() finally: if con is not None: con.close() def create_db(self, username: str) -> str: password = gen_password() search_for_user = f"SELECT FROM pg_roles WHERE rolname='{username}'" search_for_db = f"SELECT FROM pg_database WHERE datname='{username}'" create_user = f"CREATE USER {username} WITH PASSWORD %(password)s" create_database = f"CREATE DATABASE {username} OWNER {username}" revoke_perms = f"REVOKE ALL ON DATABASE {username} FROM PUBLIC" with self.psql_connection() as con, con.cursor() as cursor: if not response_is_empty(search_for_user, con): raise UserAlreadyExistsError() cursor.execute(create_user, {'password': password}) if response_is_empty(search_for_db, con): cursor.execute(create_database) cursor.execute(revoke_perms) return password def reset_db_passwd(self, username: str) -> str: password = gen_password() search_for_user = f"SELECT FROM pg_roles WHERE rolname='{username}'" reset_password = f"ALTER USER {username} WITH PASSWORD %(password)s" with self.psql_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}" with self.psql_connection() as con, con.cursor() as cursor: cursor.execute(drop_db) cursor.execute(drop_user)