music-kraken-core/src/music_kraken/database/database.py

from typing import Optional, Union
from enum import Enum
from peewee import (
    SqliteDatabase,
    MySQLDatabase,
    PostgresqlDatabase,
)

from . import data_models


class DatabaseType(Enum):
    SQLITE = "sqlite"
    POSTGRESQL = "postgresql"
    MYSQL = "mysql"

class Database:
    database: Union[SqliteDatabase, PostgresqlDatabase, MySQLDatabase]

    def __init__(
            self,
            db_type: DatabaseType,
            db_name: str,
            db_user: Optional[str] = None,
            db_password: Optional[str] = None,
            db_host: Optional[str] = None,
            db_port: Optional[int] = None
    ):
        self.db_type = db_type
        self.db_name = db_name
        self.db_user = db_user
        self.db_password = db_password
        self.db_host = db_host
        self.db_port = db_port

        self.initialize_database()

    def create_database(self) -> Union[SqliteDatabase, PostgresqlDatabase, MySQLDatabase]:
        """Create a database instance based on the configured database type and parameters.

        Returns:
            The created database instance, or None if an invalid database type was specified.
        """

        # SQLITE
        if self.db_type == DatabaseType.SQLITE:
            return SqliteDatabase(self.db_name)

        # POSTGRES
        if self.db_type == DatabaseType.POSTGRESQL:
            return PostgresqlDatabase(
                self.db_name,
                user=self.db_user,
                password=self.db_password,
                host=self.db_host,
                port=self.db_port,
            )

        # MYSQL
        if self.db_type == DatabaseType.MYSQL:
            return MySQLDatabase(
                self.db_name,
                user=self.db_user,
                password=self.db_password,
                host=self.db_host,
                port=self.db_port,
            )

        raise ValueError("define a Valid database type")

    def initialize_database(self):
        """
        Connect to the database
        initialize the previously defined databases
        create tables if they don't exist.
        """
        self.database = self.create_database()

        self.database.connect()