mirror of
https://github.com/deadcxap/YandexMusicDiscordBot.git
synced 2026-01-10 09:41:46 +03:00
impr: Update database to use guild ids.
This commit is contained in:
15
MusicBot/database/__init__.py
Normal file
15
MusicBot/database/__init__.py
Normal file
@@ -0,0 +1,15 @@
|
||||
from .base import BaseGuildsDatabase, BaseUsersDatabase
|
||||
from .extensions import VoiceGuildsDatabase
|
||||
|
||||
from .user import User, ExplicitUser
|
||||
from .guild import Guild, ExplicitGuild
|
||||
|
||||
__all__ = [
|
||||
'BaseGuildsDatabase',
|
||||
'BaseUsersDatabase',
|
||||
'VoiceGuildsDatabase',
|
||||
'User',
|
||||
'ExplicitUser',
|
||||
'Guild',
|
||||
'ExplicitGuild',
|
||||
]
|
||||
@@ -5,85 +5,98 @@ from typing import Any, cast
|
||||
from pymongo import MongoClient
|
||||
from pymongo.collection import Collection
|
||||
|
||||
from yandex_music import Track
|
||||
|
||||
from MusicBot.database.user import User, ExplicitUser, TrackInfo
|
||||
from MusicBot.database.user import User, ExplicitUser
|
||||
from MusicBot.database.guild import Guild, ExplicitGuild
|
||||
|
||||
client: MongoClient = MongoClient("mongodb://localhost:27017/")
|
||||
users: Collection[User] = client.YandexMusicBot.users
|
||||
|
||||
def create_record(uid: int | str) -> None:
|
||||
"""Create user database record.
|
||||
guilds: Collection[Guild] = client.YandexMusicBot.guilds
|
||||
|
||||
Args:
|
||||
uid (int | str): User id.
|
||||
"""
|
||||
uid = str(uid)
|
||||
users.insert_one(ExplicitUser(
|
||||
_id=uid,
|
||||
ym_token=None,
|
||||
tracks_list=[],
|
||||
current_track=None,
|
||||
is_stopped=True
|
||||
))
|
||||
class BaseUsersDatabase:
|
||||
|
||||
def update(uid: int | str, data: dict[Any, Any]) -> None:
|
||||
"""Update user record.
|
||||
def create_record(self, uid: int) -> None:
|
||||
"""Create user database record.
|
||||
|
||||
Args:
|
||||
uid (int | str): User id.
|
||||
data (dict[Any, Any]): Updated data.
|
||||
"""
|
||||
get_user(uid)
|
||||
users.update_one({'_id': str(uid)}, {"$set": data})
|
||||
Args:
|
||||
uid (int): User id.
|
||||
"""
|
||||
uid = uid
|
||||
users.insert_one(ExplicitUser(
|
||||
_id=uid,
|
||||
ym_token=None
|
||||
))
|
||||
|
||||
def get_user(uid: int | str) -> User:
|
||||
"""Get user record from database. Create new entry if not present.
|
||||
def update(self, uid: int, data: User) -> None:
|
||||
"""Update user record.
|
||||
|
||||
Args:
|
||||
uid (int | str): User id.
|
||||
Args:
|
||||
uid (int): User id.
|
||||
data (dict[Any, Any]): Updated data.
|
||||
"""
|
||||
self.get_user(uid)
|
||||
users.update_one({'_id': uid}, {"$set": data})
|
||||
|
||||
Returns:
|
||||
User: User record.
|
||||
"""
|
||||
user = users.find_one({'_id': str(uid)})
|
||||
if not user:
|
||||
create_record(uid)
|
||||
user = users.find_one({'_id': str(uid)})
|
||||
return cast(User, user)
|
||||
def get_user(self, uid: int) -> User:
|
||||
"""Get user record from database. Create new entry if not present.
|
||||
|
||||
def get_ym_token(uid: int | str) -> str | None:
|
||||
user = users.find_one({'_id': str(uid)})
|
||||
if not user:
|
||||
create_record(uid)
|
||||
user = cast(User, users.find_one({'_id': str(uid)}))
|
||||
return user['ym_token']
|
||||
Args:
|
||||
uid (int): User id.
|
||||
|
||||
def get_tracks_list(uid: int | str) -> list[TrackInfo]:
|
||||
user = get_user(uid)
|
||||
return user.get('tracks_list')
|
||||
Returns:
|
||||
User: User record.
|
||||
"""
|
||||
user = users.find_one({'_id': uid})
|
||||
if not user:
|
||||
self.create_record(uid)
|
||||
user = users.find_one({'_id': uid})
|
||||
return cast(User, user)
|
||||
|
||||
def pop_track(uid: int | str) -> TrackInfo:
|
||||
tracks_list = get_tracks_list(uid)
|
||||
track = tracks_list.pop(0)
|
||||
update(uid, {'tracks_list': tracks_list})
|
||||
return track
|
||||
def get_ym_token(self, uid: int) -> str | None:
|
||||
user = users.find_one({'_id': uid})
|
||||
if not user:
|
||||
self.create_record(uid)
|
||||
user = cast(User, users.find_one({'_id': uid}))
|
||||
return user['ym_token']
|
||||
|
||||
def add_track(uid: int | str, track: Track | TrackInfo) -> None:
|
||||
tracks_list = get_tracks_list(uid)
|
||||
if isinstance(track, Track):
|
||||
track = TrackInfo(
|
||||
track_id=str(track.id),
|
||||
title=track.title, # type: ignore
|
||||
avail=track.available, # type: ignore
|
||||
artists=", ".join(track.artists_name()),
|
||||
albums=", ".join([album.title for album in track.albums]), # type: ignore
|
||||
duration=track.duration_ms, # type: ignore
|
||||
explicit=track.explicit or bool(track.content_warning),
|
||||
bg_video=track.background_video_uri
|
||||
)
|
||||
tracks_list.append(track)
|
||||
update(uid, {'tracks_list': tracks_list})
|
||||
class BaseGuildsDatabase:
|
||||
|
||||
def create_record(self, gid: int) -> None:
|
||||
"""Create guild database record.
|
||||
|
||||
def set_current_track(uid: int | str) -> None:
|
||||
update(uid, {'current_track': str(uid)})
|
||||
Args:
|
||||
gid (int): Guild id.
|
||||
"""
|
||||
guilds.insert_one(ExplicitGuild(
|
||||
_id=gid,
|
||||
tracks_list=[],
|
||||
previous_tracks_list=[],
|
||||
current_track=None,
|
||||
is_stopped=True,
|
||||
allow_explicit=True,
|
||||
allow_menu=True
|
||||
))
|
||||
|
||||
def update(self, gid: int, data: dict[Any, Any]) -> None:
|
||||
"""Update guild record.
|
||||
|
||||
Args:
|
||||
gid (int): Guild id.
|
||||
data (dict[Any, Any]): Updated data.
|
||||
"""
|
||||
self.get_guild(gid)
|
||||
guilds.update_one({'_id': gid}, {"$set": data})
|
||||
|
||||
def get_guild(self, gid: int) -> Guild:
|
||||
"""Get guild record from database. Create new entry if not present.
|
||||
|
||||
Args:
|
||||
uid (int): User id.
|
||||
|
||||
Returns:
|
||||
Guild: Guild record.
|
||||
"""
|
||||
guild = guilds.find_one({'_id': gid})
|
||||
if not guild:
|
||||
self.create_record(gid)
|
||||
guild = guilds.find_one({'_id': gid})
|
||||
return cast(Guild, guild)
|
||||
|
||||
50
MusicBot/database/extensions.py
Normal file
50
MusicBot/database/extensions.py
Normal file
@@ -0,0 +1,50 @@
|
||||
from typing import Any
|
||||
from yandex_music import Track
|
||||
|
||||
from MusicBot.database.base import BaseGuildsDatabase
|
||||
|
||||
class VoiceGuildsDatabase(BaseGuildsDatabase):
|
||||
|
||||
def get_previous_tracks_list(self, gid: int) -> list[dict[str, Any]]:
|
||||
guild = self.get_guild(gid)
|
||||
return guild.get('previous_tracks_list')
|
||||
|
||||
def get_tracks_list(self, gid: int) -> list[dict[str, Any]]:
|
||||
guild = self.get_guild(gid)
|
||||
return guild.get('tracks_list')
|
||||
|
||||
def pop_track(self, gid: int) -> dict[str, Any]:
|
||||
tracks_list = self.get_tracks_list(gid)
|
||||
track = tracks_list.pop(0)
|
||||
self.update(gid, {'tracks_list': tracks_list})
|
||||
return track
|
||||
|
||||
def insert_track(self, gid: int, track: Track | dict[str, Any]) -> None:
|
||||
if isinstance(track, Track):
|
||||
track = track.to_dict(for_request=True)
|
||||
tracks_list = self.get_tracks_list(gid)
|
||||
tracks_list.insert(0, track)
|
||||
self.update(gid, {'tracks_list': tracks_list})
|
||||
|
||||
def add_track(self, gid: int, track: Track) -> None:
|
||||
tracks_list = self.get_tracks_list(gid)
|
||||
tracks_list.append(track.to_dict(for_request=True))
|
||||
self.update(gid, {'tracks_list': tracks_list})
|
||||
|
||||
def set_current_track(self, gid: int, track: Track) -> None:
|
||||
self.update(gid, {'current_track': track.to_dict(for_request=True)})
|
||||
|
||||
def add_previous_track(self, gid: int, track: Track | dict[str, Any]) -> None:
|
||||
tracks_list = self.get_previous_tracks_list(gid)
|
||||
if isinstance(track, Track):
|
||||
track = track.to_dict(for_request=True)
|
||||
tracks_list.insert(0, track)
|
||||
if len(tracks_list) > 50:
|
||||
tracks_list.pop()
|
||||
self.update(gid, {'previous_tracks_list': tracks_list})
|
||||
|
||||
def pop_previous_track(self, gid: int) -> dict[str, Any]:
|
||||
tracks_list = self.get_previous_tracks_list(gid)
|
||||
track = tracks_list.pop(0)
|
||||
self.update(gid, {'previous_tracks_list': tracks_list})
|
||||
return track
|
||||
@@ -1,10 +1,18 @@
|
||||
from typing import TypedDict
|
||||
from typing import TypedDict, Any
|
||||
|
||||
class Guild(TypedDict):
|
||||
tracks_list: list[dict[str, Any]]
|
||||
previous_tracks_list: list[dict[str, Any]]
|
||||
current_track: dict[str, Any] | None
|
||||
is_stopped: bool
|
||||
allow_explicit: bool
|
||||
allow_menu: bool
|
||||
|
||||
class ExplicitGuild(TypedDict):
|
||||
_id: str
|
||||
_id: int
|
||||
tracks_list: list[dict[str, Any]]
|
||||
previous_tracks_list: list[dict[str, Any]]
|
||||
current_track: dict[str, Any] | None
|
||||
is_stopped: bool # Prevents the `after` callback of play_track
|
||||
allow_explicit: bool
|
||||
allow_menu: bool
|
||||
allow_menu: bool # Is /toggle mnu command available
|
||||
@@ -1,25 +1,8 @@
|
||||
from typing import TypedDict
|
||||
|
||||
class TrackInfo(TypedDict):
|
||||
track_id: str
|
||||
title: str
|
||||
avail: bool
|
||||
artists: str
|
||||
albums: str
|
||||
duration: int
|
||||
explicit: bool
|
||||
bg_video: str | None
|
||||
|
||||
class User(TypedDict):
|
||||
ym_token: str | None
|
||||
tracks_list: list[TrackInfo]
|
||||
current_track: int | None
|
||||
is_stopped: bool
|
||||
|
||||
class ExplicitUser(TypedDict):
|
||||
_id: str
|
||||
_id: int
|
||||
ym_token: str | None
|
||||
tracks_list: list[TrackInfo]
|
||||
current_track: int | None
|
||||
is_stopped: bool # Prevents callback of play_track
|
||||
|
||||
|
||||
Reference in New Issue
Block a user