2021-08-21 00:34:10 -04:00
|
|
|
from zope.interface import implementer
|
|
|
|
from zope import component
|
2021-08-26 15:42:18 -04:00
|
|
|
from contextlib import contextmanager
|
|
|
|
|
2021-08-24 02:14:28 -04:00
|
|
|
from ceo_common.interfaces import IDatabaseService, IConfig
|
2021-08-26 15:42:18 -04:00
|
|
|
from ceo_common.errors import DatabaseConnectionError, DatabasePermissionError, UserAlreadyExistsError, \
|
|
|
|
UserNotFoundError
|
2021-08-24 02:14:28 -04:00
|
|
|
from ceod.utils import gen_password
|
|
|
|
from ceod.db.utils import response_is_empty
|
2021-08-26 15:42:18 -04:00
|
|
|
|
2021-08-24 02:14:28 -04:00
|
|
|
from psycopg2 import connect, OperationalError, ProgrammingError
|
|
|
|
from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT
|
2021-08-21 00:34:10 -04:00
|
|
|
|
|
|
|
|
|
|
|
@implementer(IDatabaseService)
|
|
|
|
class PostgreSQLService:
|
2021-08-26 02:02:47 -04:00
|
|
|
|
|
|
|
type = 'postgresql'
|
|
|
|
|
2021-08-21 00:34:10 -04:00
|
|
|
def __init__(self):
|
|
|
|
config = component.getUtility(IConfig)
|
|
|
|
self.auth_username = config.get('postgresql_username')
|
|
|
|
self.auth_password = config.get('postgresql_password')
|
|
|
|
|
2021-08-26 15:42:18 -04:00
|
|
|
@contextmanager
|
|
|
|
def psql_connection(self):
|
2021-08-21 00:34:10 -04:00
|
|
|
try:
|
|
|
|
with connect(
|
2021-08-24 22:31:50 -04:00
|
|
|
host='localhost',
|
|
|
|
user=self.auth_username,
|
|
|
|
password=self.auth_password,
|
2021-08-21 00:34:10 -04:00
|
|
|
) as con:
|
2021-08-24 02:14:28 -04:00
|
|
|
con.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
|
2021-08-26 15:42:18 -04:00
|
|
|
yield con
|
2021-08-24 02:14:28 -04:00
|
|
|
except OperationalError:
|
2021-08-21 00:34:10 -04:00
|
|
|
raise DatabaseConnectionError()
|
2021-08-24 02:14:28 -04:00
|
|
|
except ProgrammingError:
|
|
|
|
raise DatabasePermissionError()
|
|
|
|
|
2021-08-26 15:42:18 -04:00
|
|
|
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 ALL ON DATABASE {username} FROM PUBLIC;
|
|
|
|
"""
|
|
|
|
|
|
|
|
with self.psql_connection() as con:
|
|
|
|
with con.cursor() as cursor:
|
|
|
|
if response_is_empty(search_for_user, con):
|
|
|
|
cursor.execute(create_user, {'password': password})
|
|
|
|
if response_is_empty(search_for_db, con):
|
|
|
|
cursor.execute(create_database)
|
|
|
|
else:
|
|
|
|
raise UserAlreadyExistsError()
|
|
|
|
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:
|
|
|
|
with con.cursor() as cursor:
|
|
|
|
if not response_is_empty(search_for_user, con):
|
|
|
|
cursor.execute(reset_password, {'password': password})
|
|
|
|
else:
|
|
|
|
raise UserNotFoundError(username)
|
|
|
|
return password
|
|
|
|
|
2021-08-24 22:31:50 -04:00
|
|
|
def delete_db(self, username: str):
|
|
|
|
drop_db = f"DROP DATABASE IF EXISTS {username}"
|
2021-08-26 16:45:24 -04:00
|
|
|
drop_user = f"DROP USER IF EXISTS {username}"
|
2021-08-24 22:31:50 -04:00
|
|
|
|
2021-08-26 15:42:18 -04:00
|
|
|
with self.psql_connection() as con:
|
|
|
|
with con.cursor() as cursor:
|
|
|
|
cursor.execute(drop_db)
|
|
|
|
cursor.execute(drop_user)
|