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

# Standard library
from typing import Optional, Union, List
from enum import Enum

# third party modules
from peewee import (
    SqliteDatabase,
    MySQLDatabase,
    PostgresqlDatabase,
)

# own modules
from . import (
    data_models,
    write, 
    objects,
)


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("Invalid database type specified.")

    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()

        self.database.create_tables(data_models.ALL_MODELS, safe=True)

        """
        upgrade old databases. 
        If a collumn has been added in a new version this adds it to old Tables, 
        without deleting the data in legacy databases
        """
        for model in data_models.ALL_MODELS:
            for field_name, field_obj in model._meta.fields.items():
                # check if the field exists in the database
                if not self.database.table_column_exists(model._meta.db_table, field_name):
                    # add the missing column to the table
                    self.database.add_column(model._meta.db_table, field_name, field_obj)

    def push(self, database_object: objects.MusicObject):
        """
        Adds a new music object to the database using the corresponding method from the `write` session.
        When possible, rather use the `push_many` function.
        This gets even more important, when using a remote database server.

        Args:
            database_object (objects.MusicObject): The music object to add to the database.

        Returns:
            The newly added music object.
        """

        with write.WritingSession(self.database) as writing_session:
            if isinstance(database_object, objects.Song):
                return writing_session.add_song(database_object)
            
            if isinstance(database_object, objects.Album):
                return writing_session.add_album(database_object)
            
            if isinstance(database_object, objects.Artist):
                return writing_session.add_artist(database_object)

    def push_many(self, database_objects: List[objects.MusicObject]) -> None:
        """
        Adds a list of MusicObject instances to the database.
        This function sends only needs one querry for each type of table added.
        Beware that if you have for example an object like this:
        - Album
        - Song
        - Song
        you already have 3 different Tables.
    
        Unlike the function `push`, this function doesn't return the added database objects.

        Args:
            database_objects: List of MusicObject instances to be added to the database.
        """

        with write.WritingSession(self.database) as writing_session:
            for obj in database_objects:
                if isinstance(obj, objects.Song):
                    writing_session.add_song(obj)
                    continue
                
                if isinstance(obj, objects.Album):
                    writing_session.add_album(obj)
                    continue
                
                if isinstance(obj, objects.Artist):
                    writing_session.add_artist(obj)
                    continue


    def __del__(self):
        self.database.close()