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

from typing import List, Union, Type
from peewee import (
    SqliteDatabase,
    PostgresqlDatabase,
    MySQLDatabase,
    Model,
    CharField,
    IntegerField,
    BooleanField,
    ForeignKeyField,
    TextField
)


class BaseModel(Model):
    class Meta:
        database = None

    @classmethod
    def Use(cls, database: Union[SqliteDatabase, PostgresqlDatabase, MySQLDatabase]) -> Model:
        cls._meta.database = database
        return cls
    
    def use(self, database: Union[SqliteDatabase, PostgresqlDatabase, MySQLDatabase]) -> Model:
        self._meta.database = database
        return self
        


class Album(BaseModel):
    """A class representing an album in the music database."""

    title: str = CharField(null=True)
    label: str = CharField(null=True)
    album_status: str = CharField(null=True)
    language: str = CharField(null=True)
    date: str = CharField(null=True)
    date_format: str = CharField(null=True)
    country: str = CharField(null=True)
    barcode: str = CharField(null=True)
    albumsort: int = IntegerField(null=True)
    is_split: bool = BooleanField(default=False)


class Artist(BaseModel):
    """A class representing an artist in the music database."""

    name: str = CharField()


class Song(BaseModel):
    """A class representing a song in the music database."""

    name: str = CharField(null=True)
    isrc: str = CharField(null=True)
    length: int = IntegerField(null=True)
    tracksort: int = IntegerField(null=True)
    genre: str = CharField(null=True)
    # album: ForeignKeyField = ForeignKeyField(Album, backref='songs')


class Source(BaseModel):
    """A class representing a source of a song in the music database."""

    type: str = CharField()
    src: str = CharField()
    url: str = CharField()

    content_type: str = CharField()
    content_id: int = IntegerField()
    content: ForeignKeyField = ForeignKeyField('self', backref='content_items', null=True)

    @property
    def content_object(self) -> Union[Song, Album, Artist, None]:
        """Get the content associated with the source as an object."""
        if self.content_type == 'Song':
            return Song.get(Song.id == self.content_id)
        elif self.content_type == 'Album':
            return Album.get(Album.id == self.content_id)
        elif self.content_type == 'Artist':
            return Artist.get(Artist.id == self.content_id)
        else:
            return None

    @content_object.setter
    def content_object(self, value: Union[Song, Album, Artist]) -> None:
        """Set the content associated with the source as an object."""
        self.content_type = value.__class__.__name__
        self.content_id = value.id


class Target(BaseModel):
    """A class representing a target of a song in the music database."""

    file: str = CharField()
    path: str = CharField()
    song = ForeignKeyField(Song, backref='targets')


class Lyrics(BaseModel):
    """A class representing lyrics of a song in the music database."""

    text: str = TextField()
    language: str = CharField()
    song = ForeignKeyField(Song, backref='lyrics')


class SongArtist(BaseModel):
    """A class representing the relationship between a song and an artist."""

    song: ForeignKeyField = ForeignKeyField(Song, backref='song_artists')
    artist: ForeignKeyField = ForeignKeyField(Artist, backref='song_artists')
    is_feature: bool = BooleanField(default=False)


class AlbumArtist(BaseModel):
    """A class representing the relationship between an album and an artist."""

    album: ForeignKeyField = ForeignKeyField(Album, backref='album_artists')
    artist: ForeignKeyField = ForeignKeyField(Artist, backref='album_artists')




database_1 = SqliteDatabase(":memory:")
database_1.create_tables([Song.Use(database_1)])
database_2 = SqliteDatabase(":memory:")
database_2.create_tables([Song.Use(database_2)])

# creating songs, adding it to db_2 if i is even, else to db_1
for i in range(100):
    song = Song(name=str(i) + "hs")

    db_to_use = database_2 if i % 2 == 0 else database_1
    song.use(db_to_use).save()

print("database 1")
for song in Song.Use(database_1).select():
    print(song.name)

print("database 2")
for song in Song.Use(database_1).select():
    print(song.name)