pyceo/ceod/db/PostgreSQLService.py

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
        except OperationalError as e:
            logger.error(e)
            raise DatabaseConnectionError()
        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)