mirror of
https://github.com/deadcxap/YandexMusicDiscordBot.git
synced 2026-01-14 03:31:55 +03:00
impr: Async database and code optimization.
This commit is contained in:
Lemon4ksan
@@ -1,186 +1,112 @@
|
||||
"""This documents initialises databse and contains methods to access it."""
|
||||
|
||||
from typing import Any, cast
|
||||
|
||||
from pymongo import MongoClient
|
||||
from pymongo.collection import Collection
|
||||
from typing import Iterable, Any, cast
|
||||
from pymongo import AsyncMongoClient, ReturnDocument
|
||||
from pymongo.asynchronous.collection import AsyncCollection
|
||||
from pymongo.results import UpdateResult
|
||||
|
||||
from .user import User, ExplicitUser
|
||||
from .guild import Guild, ExplicitGuild, MessageVotes
|
||||
|
||||
client: MongoClient = MongoClient("mongodb://localhost:27017/")
|
||||
users: Collection[ExplicitUser] = client.YandexMusicBot.users
|
||||
guilds: Collection[ExplicitGuild] = client.YandexMusicBot.guilds
|
||||
client: AsyncMongoClient = AsyncMongoClient("mongodb://localhost:27017/")
|
||||
|
||||
db = client.YandexMusicBot
|
||||
users: AsyncCollection[ExplicitUser] = db.users
|
||||
guilds: AsyncCollection[ExplicitGuild] = db.guilds
|
||||
|
||||
class BaseUsersDatabase:
|
||||
DEFAULT_USER = ExplicitUser(
|
||||
_id=0,
|
||||
ym_token=None,
|
||||
playlists=[],
|
||||
playlists_page=0,
|
||||
queue_page=0,
|
||||
vibe_batch_id=None,
|
||||
vibe_type=None,
|
||||
vibe_id=None,
|
||||
vibe_settings={
|
||||
'mood': 'all',
|
||||
'diversity': 'default',
|
||||
'lang': 'any'
|
||||
}
|
||||
)
|
||||
|
||||
def create_record(self, uid: int) -> None:
|
||||
"""Create user database record.
|
||||
|
||||
Args:
|
||||
uid (int): User id.
|
||||
"""
|
||||
uid = uid
|
||||
users.insert_one(ExplicitUser(
|
||||
_id=uid,
|
||||
ym_token=None,
|
||||
playlists=[],
|
||||
playlists_page=0,
|
||||
queue_page=0,
|
||||
vibe_batch_id=None,
|
||||
vibe_type=None,
|
||||
vibe_id=None,
|
||||
vibe_settings={
|
||||
'mood': 'all',
|
||||
'diversity': 'default',
|
||||
'lang': 'any'
|
||||
}
|
||||
))
|
||||
|
||||
def update(self, uid: int, data: User | dict[Any, Any]) -> None:
|
||||
"""Update user record.
|
||||
|
||||
Args:
|
||||
uid (int): User id.
|
||||
data (User | dict[Any, Any]): Updated data.
|
||||
"""
|
||||
self.get_user(uid)
|
||||
users.update_one({'_id': uid}, {"$set": data})
|
||||
|
||||
def get_user(self, uid: int) -> ExplicitUser:
|
||||
"""Get user record from database. Create new entry if not present.
|
||||
|
||||
Args:
|
||||
uid (int): User id.
|
||||
|
||||
Returns:
|
||||
User: User record.
|
||||
"""
|
||||
user = users.find_one({'_id': uid})
|
||||
if not user:
|
||||
self.create_record(uid)
|
||||
user = users.find_one({'_id': uid})
|
||||
user = cast(ExplicitUser, user)
|
||||
existing_fields = user.keys()
|
||||
fields: ExplicitUser = ExplicitUser(
|
||||
_id=0,
|
||||
ym_token=None,
|
||||
playlists=[],
|
||||
playlists_page=0,
|
||||
queue_page=0,
|
||||
vibe_batch_id=None,
|
||||
vibe_type=None,
|
||||
vibe_id=None,
|
||||
vibe_settings={
|
||||
'mood': 'all',
|
||||
'diversity': 'default',
|
||||
'lang': 'any'
|
||||
}
|
||||
async def update(self, uid: int, data: User | dict[str, Any]) -> UpdateResult:
|
||||
return await users.update_one(
|
||||
{'_id': uid},
|
||||
{'$set': data},
|
||||
upsert=True
|
||||
)
|
||||
|
||||
async def get_user(self, uid: int, projection: User | Iterable[str] | None = None) -> ExplicitUser:
|
||||
user = await users.find_one_and_update(
|
||||
{'_id': uid},
|
||||
{'$setOnInsert': self.DEFAULT_USER},
|
||||
return_document=ReturnDocument.AFTER,
|
||||
upsert=True,
|
||||
projection=projection
|
||||
)
|
||||
return cast(ExplicitUser, user)
|
||||
|
||||
async def get_ym_token(self, uid: int) -> str | None:
|
||||
user = await users.find_one(
|
||||
{'_id': uid},
|
||||
projection={'ym_token': 1}
|
||||
)
|
||||
return cast(str | None, user.get('ym_token') if user else None)
|
||||
|
||||
async def add_playlist(self, uid: int, playlist_data: dict) -> UpdateResult:
|
||||
return await users.update_one(
|
||||
{'_id': uid},
|
||||
{'$push': {'playlists': playlist_data}}
|
||||
)
|
||||
for field, default_value in fields.items():
|
||||
if field not in existing_fields:
|
||||
user[field] = default_value
|
||||
users.update_one({'_id': uid}, {"$set": {field: default_value}})
|
||||
|
||||
return user
|
||||
|
||||
def get_ym_token(self, uid: int) -> str | None:
|
||||
user = users.find_one({'_id': uid})
|
||||
if not user:
|
||||
self.create_record(uid)
|
||||
user = users.find_one({'_id': uid})
|
||||
return cast(ExplicitUser, user)['ym_token']
|
||||
|
||||
class BaseGuildsDatabase:
|
||||
|
||||
def create_record(self, gid: int) -> None:
|
||||
"""Create guild database record.
|
||||
DEFAULT_GUILD = ExplicitGuild(
|
||||
_id=0,
|
||||
next_tracks=[],
|
||||
previous_tracks=[],
|
||||
current_track=None,
|
||||
current_menu=None,
|
||||
is_stopped=True,
|
||||
allow_explicit=True,
|
||||
always_allow_menu=False,
|
||||
vote_next_track=True,
|
||||
vote_add_track=True,
|
||||
vote_add_album=True,
|
||||
vote_add_artist=True,
|
||||
vote_add_playlist=True,
|
||||
shuffle=False,
|
||||
repeat=False,
|
||||
votes={},
|
||||
vibing=False,
|
||||
current_viber_id=None
|
||||
)
|
||||
|
||||
Args:
|
||||
gid (int): Guild id.
|
||||
"""
|
||||
guilds.insert_one(ExplicitGuild(
|
||||
_id=gid,
|
||||
next_tracks=[],
|
||||
previous_tracks=[],
|
||||
current_track=None,
|
||||
current_menu=None,
|
||||
is_stopped=True,
|
||||
allow_explicit=True,
|
||||
always_allow_menu=False,
|
||||
vote_next_track=True,
|
||||
vote_add_track=True,
|
||||
vote_add_album=True,
|
||||
vote_add_artist=True,
|
||||
vote_add_playlist=True,
|
||||
shuffle=False,
|
||||
repeat=False,
|
||||
votes={},
|
||||
vibing=False,
|
||||
current_viber_id=None
|
||||
))
|
||||
|
||||
def update(self, gid: int, data: Guild) -> 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) -> ExplicitGuild:
|
||||
"""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})
|
||||
|
||||
guild = cast(ExplicitGuild, guild)
|
||||
existing_fields = guild.keys()
|
||||
fields = ExplicitGuild(
|
||||
_id=0,
|
||||
next_tracks=[],
|
||||
previous_tracks=[],
|
||||
current_track=None,
|
||||
current_menu=None,
|
||||
is_stopped=True,
|
||||
allow_explicit=True,
|
||||
always_allow_menu=False,
|
||||
vote_next_track=True,
|
||||
vote_add_track=True,
|
||||
vote_add_album=True,
|
||||
vote_add_artist=True,
|
||||
vote_add_playlist=True,
|
||||
shuffle=False,
|
||||
repeat=False,
|
||||
votes={},
|
||||
vibing=False,
|
||||
current_viber_id=None
|
||||
async def update(self, gid: int, data: Guild | dict[str, Any]) -> UpdateResult:
|
||||
return await guilds.update_one(
|
||||
{'_id': gid},
|
||||
{'$set': data},
|
||||
upsert=True
|
||||
)
|
||||
|
||||
async def get_guild(self, gid: int, projection: Guild | Iterable[str] | None = None) -> ExplicitGuild:
|
||||
guild = await guilds.find_one_and_update(
|
||||
{'_id': gid},
|
||||
{'$setOnInsert': self.DEFAULT_GUILD},
|
||||
return_document=ReturnDocument.AFTER,
|
||||
upsert=True,
|
||||
projection=projection
|
||||
)
|
||||
return cast(ExplicitGuild, guild)
|
||||
|
||||
async def update_vote(self, gid: int, mid: int, data: MessageVotes) -> UpdateResult:
|
||||
return await guilds.update_one(
|
||||
{'_id': gid},
|
||||
{'$set': {f'votes.{mid}': data}}
|
||||
)
|
||||
|
||||
async def clear_queue(self, gid: int) -> UpdateResult:
|
||||
return await guilds.update_one(
|
||||
{'_id': gid},
|
||||
{'$set': {'next_tracks': []}}
|
||||
)
|
||||
for field, default_value in fields.items():
|
||||
if field not in existing_fields:
|
||||
guild[field] = default_value
|
||||
guilds.update_one({'_id': gid}, {"$set": {field: default_value}})
|
||||
|
||||
return guild
|
||||
|
||||
def update_vote(self, gid: int, mid: int, data: MessageVotes) -> None:
|
||||
"""Update vote for a message in a guild.
|
||||
|
||||
Args:
|
||||
gid (int): Guild id.
|
||||
mid (int): Message id.
|
||||
vote (bool): Vote value.
|
||||
"""
|
||||
guild = self.get_guild(gid)
|
||||
guild['votes'][str(mid)] = data
|
||||
guilds.update_one({'_id': gid}, {"$set": {'votes': guild['votes']}})
|
||||
Reference in New Issue
Block a user