artemis/core/data/database.py

import logging, coloredlogs
from typing import Optional
from sqlalchemy.orm import scoped_session, sessionmaker
from sqlalchemy import create_engine
from logging.handlers import TimedRotatingFileHandler
import os
import secrets, string
import bcrypt
from hashlib import sha256
import alembic.config
import glob

from core.config import CoreConfig
from core.data.schema import *
from core.utils import Utils


class Data:
    engine = None
    session = None
    user = None
    arcade = None
    card = None
    base = None
    def __init__(self, cfg: CoreConfig) -> None:
        self.config = cfg

        if self.config.database.sha2_password:
            passwd = sha256(self.config.database.password.encode()).digest()
            self.__url = f"{self.config.database.protocol}://{self.config.database.username}:{passwd.hex()}@{self.config.database.host}:{self.config.database.port}/{self.config.database.name}?charset=utf8mb4"
        else:
            self.__url = f"{self.config.database.protocol}://{self.config.database.username}:{self.config.database.password}@{self.config.database.host}:{self.config.database.port}/{self.config.database.name}?charset=utf8mb4"

        if Data.engine is None:
            Data.engine = create_engine(self.__url, pool_recycle=3600)
            self.__engine = Data.engine

        if Data.session is None:
            s = sessionmaker(bind=Data.engine, autoflush=True, autocommit=True)
            Data.session = scoped_session(s)

        if Data.user is None:
            Data.user = UserData(self.config, self.session)
        
        if Data.arcade is None:
            Data.arcade = ArcadeData(self.config, self.session)
        
        if Data.card is None:
            Data.card = CardData(self.config, self.session)
        
        if Data.base is None:
            Data.base = BaseData(self.config, self.session)

        self.logger = logging.getLogger("database")

        # Prevent the logger from adding handlers multiple times
        if not getattr(self.logger, "handler_set", None):            
            log_fmt_str = "[%(asctime)s] %(levelname)s | Database | %(message)s"
            log_fmt = logging.Formatter(log_fmt_str)
            fileHandler = TimedRotatingFileHandler(
                "{0}/{1}.log".format(self.config.server.log_dir, "db"),
                encoding="utf-8",
                when="d",
                backupCount=10,
            )
            fileHandler.setFormatter(log_fmt)

            consoleHandler = logging.StreamHandler()
            consoleHandler.setFormatter(log_fmt)

            self.logger.addHandler(fileHandler)
            self.logger.addHandler(consoleHandler)

            self.logger.setLevel(self.config.database.loglevel)
            coloredlogs.install(
                cfg.database.loglevel, logger=self.logger, fmt=log_fmt_str
            )
            self.logger.handler_set = True  # type: ignore

    def __alembic_cmd(self, command: str, *args: str) -> None:
        old_dir = os.path.abspath(os.path.curdir)
        base_dir = os.path.join(os.path.abspath(os.path.curdir), 'core', 'data', 'alembic')
        alembicArgs = [
            "-c",
            os.path.join(base_dir, "alembic.ini"),
            "-x",
            f"script_location={base_dir}",
            "-x",
            f"sqlalchemy.url={self.__url}",
            command,
        ]
        alembicArgs.extend(args)
        os.chdir(base_dir)
        alembic.config.main(argv=alembicArgs)
        os.chdir(old_dir)

    def create_database(self):
        self.logger.info("Creating databases...")
        metadata.create_all(
            self.engine,
            checkfirst=True,
        )

        for _, mod in Utils.get_all_titles().items():
            if hasattr(mod, "database"):
                mod.database(self.config)
                metadata.create_all(
                    self.engine,
                    checkfirst=True,
                )

        # Stamp the end revision as if alembic had created it, so it can take off after this.
        self.__alembic_cmd(
            "stamp",
            "head",
        )

    def schema_upgrade(self, ver: str = None):
        self.__alembic_cmd(
            "upgrade",
            "head" if not ver else ver,
        )

    def schema_downgrade(self, ver: str):
        self.__alembic_cmd(
            "downgrade",
            ver,
        )

    async def create_owner(self, email: Optional[str] = None, code: Optional[str] = "00000000000000000000") -> None:
        pw = "".join(
            secrets.choice(string.ascii_letters + string.digits) for i in range(20)
        )
        hash = bcrypt.hashpw(pw.encode(), bcrypt.gensalt())

        user_id = await self.user.create_user(username="sysowner", email=email, password=hash.decode(), permission=255)
        if user_id is None:
            self.logger.error(f"Failed to create owner with email {email}")
            return

        card_id = await self.card.create_card(user_id, code)
        if card_id is None:
            self.logger.error(f"Failed to create card for owner with id {user_id}")
            return

        self.logger.warning(
            f"Successfully created owner with email {email}, access code {code}, and password {pw} Make sure to change this password and assign a real card ASAP!"
        )
    
    async def migrate(self) -> None:
        exist = await self.base.execute("SELECT * FROM alembic_version")
        if exist is not None:
            self.logger.warn("No need to migrate as you have already migrated to alembic. If you are trying to upgrade the schema, use `upgrade` instead!")
            return
        
        self.logger.info("Upgrading to latest with legacy system")
        if not await self.legacy_upgrade():
            self.logger.warn("No need to migrate as you have already deleted the old schema_versions system. If you are trying to upgrade the schema, use `upgrade` instead!")
            return
        self.logger.info("Done")
        
        self.logger.info("Stamp with initial revision")
        self.__alembic_cmd(
            "stamp",
            "835b862f9bf0",
        )

        self.logger.info("Upgrade")
        self.__alembic_cmd(
            "upgrade",
            "head",
        )
    
    async def legacy_upgrade(self) -> bool:
        vers = await self.base.execute("SELECT * FROM schema_versions")
        if vers is None:
            self.logger.warn("Cannot legacy upgrade, schema_versions table unavailable!")
            return False
        
        db_vers = {}
        vers_list = vers.fetchall()
        for x in vers_list:
            db_vers[x['game']] = x['version']
        
        core_now_ver = int(db_vers['CORE']) + 1
        while os.path.exists(f"core/data/schema/versions/CORE_{core_now_ver}_upgrade.sql"):
            with open(f"core/data/schema/versions/CORE_{core_now_ver}_upgrade.sql", "r") as f:
                result = await self.base.execute(f.read())
                
                if result is None:
                    self.logger.error(f"Invalid upgrade script CORE_{core_now_ver}_upgrade.sql")
                    break

                result = await self.base.execute(f"UPDATE schema_versions SET version = {core_now_ver} WHERE game = 'CORE'")
                if result is None:
                    self.logger.error(f"Failed to update schema version for CORE to {core_now_ver}")
                    break
            
            self.logger.info(f"Upgrade CORE to version {core_now_ver}")
            core_now_ver += 1
        
        for _, mod in Utils.get_all_titles().items():
            game_codes = getattr(mod, "game_codes", [])
            for game in game_codes:
                if game not in db_vers:
                    self.logger.warn(f"{game} does not have an antry in schema_versions, skipping")
                    continue

                now_ver = int(db_vers[game]) + 1
                while os.path.exists(f"core/data/schema/versions/{game}_{now_ver}_upgrade.sql"):
                    with open(f"core/data/schema/versions/{game}_{now_ver}_upgrade.sql", "r") as f:
                        result = await self.base.execute(f.read())
                        
                        if result is None:
                            self.logger.error(f"Invalid upgrade script {game}_{now_ver}_upgrade.sql")
                            break

                        result = await self.base.execute(f"UPDATE schema_versions SET version = {now_ver} WHERE game = '{game}'")
                        if result is None:
                            self.logger.error(f"Failed to update schema version for {game} to {now_ver}")
                            break

                    self.logger.info(f"Upgrade {game} to version {now_ver}")
                    now_ver += 1
        
        return True

    async def create_revision(self, message: str) -> None:
        if not message:
            self.logger.info("Message is required for create-revision")
            return
        
        self.__alembic_cmd(
            "revision",
            "-m",
            message,
        )

    async def create_revision_auto(self, message: str) -> None:
        if not message:
            self.logger.info("Message is required for create-revision")
            return
        
        for _, mod in Utils.get_all_titles().items():
            if hasattr(mod, "database"):
                mod.database(self.config)
        
        self.__alembic_cmd(
            "revision",
            "--autogenerate",
            "-m",
            message,
        )