From 4af33cff171d4daded7148ceba8a540826c4f7e6 Mon Sep 17 00:00:00 2001 From: Drapersniper <27962761+drapersniper@users.noreply.github.com> Date: Tue, 17 Dec 2019 22:16:59 +0000 Subject: [PATCH] Migrate Playlist to DB 3 TODO 1 Migrate Config to Schema 3 without playlists and update get_playlist methods Signed-off-by: Drapersniper <27962761+drapersniper@users.noreply.github.com> --- redbot/cogs/audio/playlists.py | 235 ++++++++++++++++++++++++++++++--- 1 file changed, 219 insertions(+), 16 deletions(-) diff --git a/redbot/cogs/audio/playlists.py b/redbot/cogs/audio/playlists.py index 1ac62b7fd..a96a87637 100644 --- a/redbot/cogs/audio/playlists.py +++ b/redbot/cogs/audio/playlists.py @@ -1,18 +1,23 @@ +import json +import os from collections import namedtuple from enum import Enum, unique -from typing import List, Optional, Union +from typing import List, Optional, Union, Tuple +import apsw import discord import lavalink from redbot.core import Config, commands from redbot.core.bot import Red +from redbot.core.data_manager import cog_data_path from redbot.core.i18n import Translator from redbot.core.utils.chat_formatting import humanize_list from .errors import InvalidPlaylistScope, MissingAuthor, MissingGuild, NotAllowed -_config = None -_bot = None +_config: Config = None +_bot: Red = None +_database: "Database" = None __all__ = [ "Playlist", @@ -31,6 +36,105 @@ FakePlaylist = namedtuple("Playlist", "author scope") _ = Translator("Audio", __file__) +_PRAGMA_UPDATE_temp_store = """ +PRAGMA temp_store = 2; +""" +_PRAGMA_UPDATE_journal_mode = """ +PRAGMA journal_mode = wal; +""" +_PRAGMA_UPDATE_wal_autocheckpoint = """ +PRAGMA wal_autocheckpoint; +""" +_PRAGMA_UPDATE_read_uncommitted = """ +PRAGMA read_uncommitted = 1; +""" +_PRAGMA_UPDATE_optimize = """ +PRAGMA optimize = 1; +""" + +_CREATE_TABLE = """ +CREATE TABLE IF NOT EXISTS GLOBAL ( +playlist_id INTEGER PRIMARY KEY, +playlist_name TEXT NOT NULL, +scope_id INTEGER NOT NULL, +author_id INTEGER NOT NULL, +playlist_url TEXT, +tracks BLOB); +""" + +_DROP = """ +DROP TABLE {table}; +""" +_DELETE = """ +DELETE FROM {table} +WHERE + ( + playlist_id = :playlist_id + AND + scope_id = :scope_id + ) +; +""" +_FETCH_ALL = """ +SELECT +playlist_id, +playlist_name, +scope_id, +author_id, +playlist_url, +tracks +FROM {table}; +""" + +_FETCH = """ +SELECT +playlist_id, +playlist_name, +scope_id, +author_id, +playlist_url, +tracks +FROM {table} +WHERE + ( + playlist_id = :playlist_id + AND + scope_id = :scope_id + ) +""" + +_UPSET = """INSERT INTO +{table} + ( + playlist_id, + playlist_name, + scope_id, + author_id, + playlist_url, + tracks + ) +VALUES + ( + :playlist_id, + :playlist_name, + :scope_id, + :author_id, + :playlist_url, + :tracks + ) +ON CONFLICT + ( + playlist_id, + scope_id + ) +DO UPDATE + SET + playlist_name = :playlist_name, + playlist_url = :playlist_url, + tracks = :tracks +; +""" + @unique class PlaylistScope(Enum): @@ -46,12 +150,89 @@ class PlaylistScope(Enum): return list(map(lambda c: c.value, PlaylistScope)) +class Database: + def __init__(self): + self._database = apsw.Connection( + str(cog_data_path(_bot.get_cog("Audio")) / "playlists.db") + ) + self.cursor = self._database.cursor() + self.cursor.execute(_PRAGMA_UPDATE_temp_store) + self.cursor.execute(_PRAGMA_UPDATE_journal_mode) + self.cursor.execute(_PRAGMA_UPDATE_wal_autocheckpoint) + self.cursor.execute(_PRAGMA_UPDATE_read_uncommitted) + for t in ["GLOBAL", "GUILD", "USER"]: + self.cursor.execute(_CREATE_TABLE.format(table=t)) + + @staticmethod + def parse_query(scope: PlaylistScope, query: str): + if scope == PlaylistScope.GLOBAL.value: + table = "GLOBAL" + elif scope == PlaylistScope.GUILD.value: + table = "GUILD" + elif scope == PlaylistScope.USER.value: + table = "USER" + else: + raise + return query.format(table=table) + + def fetch( + self, scope: PlaylistScope, playlist_id: int, scope_id: int + ) -> Tuple[int, str, int, int, str, str]: + query = self.parse_query(scope, _FETCH) + return self.cursor.execute( + query, ({"playlist_id": playlist_id, "scope_id": scope_id}) + ).fetchone() + + def delete(self, scope: PlaylistScope, playlist_id: int, scope_id: int): + query = self.parse_query(scope, _DELETE) + return self.cursor.execute(query, ({"playlist_id": playlist_id, "scope_id": scope_id})) + + def fetch_all(self, scope: PlaylistScope) -> List[Tuple[int, str, int, int, str, str]]: + query = self.parse_query(scope, _FETCH_ALL) + return self.cursor.execute(query).fetchall() + + def drop(self, scope: PlaylistScope): + query = self.parse_query(scope, _DROP) + return self.cursor.execute(query) + + def create_table(self, scope: PlaylistScope): + query = self.parse_query(scope, _CREATE_TABLE) + return self.cursor.execute(query) + + def upsert( + self, + scope: PlaylistScope, + playlist_id: int, + playlist_name: str, + scope_id: int, + author_id: int, + playlist_url: str, + tracks: List[dict], + ): + query = self.parse_query(scope, _UPSET) + self.cursor.execute( + query, + ( + { + "playlist_id": playlist_id, + "playlist_name": playlist_name, + "scope_id": scope_id, + "author_id": author_id, + "playlist_url": playlist_url, + "tracks": json.dumps(tracks), + } + ), + ) + + def _pass_config_to_playlist(config: Config, bot: Red): - global _config, _bot + global _config, _bot, _database if _config is None: _config = config if _bot is None: _bot = bot + if _database is None: + _database = Database() def standardize_scope(scope) -> str: @@ -92,15 +273,15 @@ def _prepare_config_scope( scope = standardize_scope(scope) if scope == PlaylistScope.GLOBAL.value: - config_scope = [PlaylistScope.GLOBAL.value] + config_scope = [PlaylistScope.GLOBAL.value, _bot.user.id] elif scope == PlaylistScope.USER.value: if author is None: raise MissingAuthor("Invalid author for user scope.") - config_scope = [PlaylistScope.USER.value, str(getattr(author, "id", author))] + config_scope = [PlaylistScope.USER.value, getattr(author, "id", author)] else: if guild is None: raise MissingGuild("Invalid guild for guild scope.") - config_scope = [PlaylistScope.GUILD.value, str(getattr(guild, "id", guild))] + config_scope = [PlaylistScope.GUILD.value, getattr(guild, "id", guild)] return config_scope @@ -132,6 +313,14 @@ class Playlist: self.tracks = tracks or [] self.tracks_obj = [lavalink.Track(data=track) for track in self.tracks] + def _get_scope_id(self): + if self.scope == PlaylistScope.GLOBAL.value: + return self.bot.user.id + elif self.scope == PlaylistScope.USER.value: + return self.author + else: + return self.guild.id + async def edit(self, data: dict): """ Edits a Playlist. @@ -146,8 +335,22 @@ class Playlist: for item in list(data.keys()): setattr(self, item, data[item]) + await self.save() - await _config.custom(*self.config_scope, str(self.id)).set(self.to_json()) + async def save(self): + """ + Saves a Playlist. + """ + scope, scope_id = self.config_scope + _database.upsert( + scope, + playlist_id=int(self.id), + playlist_name=self.name, + scope_id=scope_id, + author_id=self.author, + playlist_url=self.url, + tracks=self.tracks, + ) def to_json(self) -> dict: """Transform the object to a dict. @@ -216,7 +419,7 @@ class Playlist: ) -async def get_playlist( +async def get_playlist( # TODO: convert to SQL playlist_number: int, scope: str, bot: Red, @@ -262,7 +465,7 @@ async def get_playlist( ) -async def get_all_playlist( +async def get_all_playlist( # TODO: convert to SQL scope: str, bot: Red, guild: Union[discord.Guild, int] = None, @@ -360,10 +563,7 @@ async def create_playlist( playlist = Playlist( ctx.bot, scope, author.id, ctx.message.id, playlist_name, playlist_url, tracks, ctx.guild ) - - await _config.custom(*_prepare_config_scope(scope, author, guild), str(ctx.message.id)).set( - playlist.to_json() - ) + await playlist.save() return playlist @@ -393,7 +593,9 @@ async def reset_playlist( `MissingAuthor` Trying to access the User scope without an user id. """ - await _config.custom(*_prepare_config_scope(scope, author, guild)).clear() + scope, scope_id = _prepare_config_scope(scope, author, guild) + _database.drop(scope) + _database.create_table(scope) async def delete_playlist( @@ -425,4 +627,5 @@ async def delete_playlist( `MissingAuthor` Trying to access the User scope without an user id. """ - await _config.custom(*_prepare_config_scope(scope, author, guild), str(playlist_id)).clear() + scope, scope_id = _prepare_config_scope(scope, author, guild) + _database.delete(scope, int(playlist_id), scope_id)