city_retrofit/hub/persistence/db_setup.py

"""
Database setup
SPDX - License - Identifier: LGPL - 3.0 - or -later
Copyright © 2022 Concordia CERC group
Project Coder Peter Yefi peteryefi@gmail.com
"""

from hub.persistence import Repository
from hub.persistence.models import Application
from hub.persistence.models import City
from hub.persistence.models import CityObject
from hub.persistence.models import User
from hub.persistence.models import UserRoles
from hub.persistence.models import SimulationResults
from hub.persistence.repositories import User as UserRepository
from hub.persistence.repositories import Application as ApplicationRepository
from hub.hub_logger import logger


class DBSetup:

  def __init__(self, db_name, app_env, dotenv_path, admin_password, application_uuid):
    """
    Creates database tables a default admin user and a default admin app with the given password and uuid
    :param db_name:
    :param app_env:
    :param dotenv_path:
    """
    repository = Repository(db_name=db_name, app_env=app_env, dotenv_path=dotenv_path)

    # Create the tables using the models
    Application.__table__.create(bind=repository.engine, checkfirst=True)
    City.__table__.create(bind=repository.engine, checkfirst=True)
    CityObject.__table__.create(bind=repository.engine, checkfirst=True)
    SimulationResults.__table__.create(bind=repository.engine, checkfirst=True)
    User.__table__.create(bind=repository.engine, checkfirst=True)

    self._user_repo = UserRepository(db_name=db_name, app_env=app_env, dotenv_path=dotenv_path)
    self._application_repo = ApplicationRepository(db_name=db_name, app_env=app_env, dotenv_path=dotenv_path)
    self._create_admin_user(self._user_repo, admin_password)
    self._create_admin_app(self._application_repo, application_uuid)

  @staticmethod
  def _create_admin_app(application_repo, application_uuid):
    name = 'AdminTool'
    description = 'Admin tool to control city persistence and to test the API v1.4'
    print('Creating default admin tool application...')
    application = application_repo.insert(name, description, application_uuid)
    if type(application) is dict:
      logger.info(application)
    else:
      msg = f'Created Admin tool with application_uuid: {application_uuid}'
      print(msg)
      logger.info(msg)

  @staticmethod
  def _create_admin_user(user_repo, admin_password, application_id):
    password = admin_password
    print('Creating default admin user...')
    user = user_repo.insert('Administrator', password, UserRoles.Admin, application_id)
    if type(user) is dict:
      logger.info(user)
    else:
      print(f'Created Admin user')
      logger.info(f'Created Admin user')