Add storage, syncing db, rename module

bot/__init__.py

@@ -17,5 +17,8 @@ async def runner():
 def main():
     import asyncio
 
+    from rich.traceback import install
+    install(show_locals=True)
+
     print('Starting...')
     asyncio.run(runner())

bot/common.py

@@ -1,7 +1,8 @@
 from aiogram import Bot, Dispatcher
+from bot.modules.fsm import InDbStorage
 from .utils.config import config
 
 bot = Bot(token=config.telegram.bot_token)
-dp = Dispatcher()
+dp = Dispatcher(storage=InDbStorage())
 
 __all__ = ['bot', 'dp', 'config']

bot/handlers/initialize/initializer.py

@@ -8,3 +8,6 @@ router = Router()
 @router.startup()
 async def startup(bot: Bot):
     print(f'[green]Started as[/] @{(await bot.me()).username}')
+
+    from bot.modules.database import pull
+    await pull()

bot/modules/database/__init__.py

@@ -0,0 +1,7 @@
+from .db import Db
+from .pull_db import pull
+
+
+db = Db()
+
+__all__ = ['db', 'pull']

bot/modules/database/db.py

@@ -0,0 +1,17 @@
+from .db_model import DBDict
+import os.path
+from bot.utils.config import config
+
+DB = os.path.join(config.local.db_path, 'db')
+
+if not os.path.isfile(DB):
+    open('sync', 'w')
+
+
+class Db(object):
+    def __init__(self):
+        self.fsm = DBDict('fsm')
+        self.config = DBDict('config')
+
+    async def write(self):
+        await self.config.write()

bot/modules/database/db_model.py

@@ -0,0 +1,48 @@
+from sqlitedict import SqliteDict
+from bot.common import bot
+from bot.utils.config import config
+from aiogram.types import FSInputFile, InputMediaDocument
+from aiogram import exceptions
+from pydantic import ValidationError
+import time
+from .meta import DBMeta
+import os.path
+
+DB = os.path.join(config.local.db_path, 'db')
+DB_CHAT = config.telegram.db_chat
+
+
+class DBDict(SqliteDict):
+    def __init__(self, tablename: str):
+        super().__init__(DB, tablename=tablename, autocommit=True)
+
+    async def write(self):
+        try:
+            DBMeta().update_time = time.time_ns()
+
+            await bot.edit_message_media(
+                media=InputMediaDocument(media=FSInputFile(DB)),
+                chat_id=DB_CHAT,
+                message_id=DBMeta().message_id
+            )
+            await bot.edit_message_caption(
+                caption=str(DBMeta()),
+                chat_id=DB_CHAT,
+                message_id=DBMeta().message_id
+            )
+
+        except (ValidationError, exceptions.TelegramBadRequest):
+            DBMeta().update_time = time.time_ns()
+
+            self['db_message_id'] = (
+                await bot.send_document(
+                    chat_id=DB_CHAT, document=FSInputFile(DB),
+                    disable_notification=True
+                )
+            ).message_id
+
+            DBMeta().message_id = self['db_message_id']
+            await bot.edit_message_caption(
+                caption=str(DBMeta()),
+                chat_id=DB_CHAT, message_id=self.get('db_message_id')
+            )

bot/modules/database/meta.py

@@ -0,0 +1,106 @@
+import os.path
+from bot.utils.config import config
+from bot.common import bot
+from aiogram.exceptions import TelegramBadRequest
+import asyncio
+
+loop = asyncio.get_event_loop()
+
+
+DBMETA = os.path.join(config.local.db_path, 'dbmeta')
+APP_ID = config.local.app_id
+DB_CHAT = config.telegram.db_chat
+
+
+def meta_property(prop_name):
+    def getter(self):
+        return self[prop_name]
+
+    def setter(self, value):
+        self[prop_name] = value
+
+    return property(getter, setter)
+
+
+class DBMeta:
+    app_id = meta_property('app_id')
+    message_id = meta_property('message_id')
+    update_time = meta_property('update_time')
+
+    def __init__(self):
+        if not os.path.isfile(DBMETA):
+            open(DBMETA, 'w').write(f'{APP_ID}|None|0')
+
+    def __getitem__(self, item):
+        try:
+            return open(DBMETA).read().split('|')[{
+                "app_id": 0,
+                "message_id": 1,
+                "update_time": 2
+            }.get(item)]
+        except TypeError:
+            return None
+
+    def __setitem__(self, key, value):
+        meta = open(DBMETA).read().split('|')
+        meta[{
+            "app_id": 0,
+            "message_id": 1,
+            "update_time": 2
+        }[key]] = value
+        open(DBMETA, 'w').write('|'.join(str(x) for x in meta))
+
+    def __str__(self):
+        return open(DBMETA).read()
+
+
+def cloud_meta_property(self, prop_name):
+    async def getter():
+        return await self.get(prop_name)
+
+    return getter()
+
+
+class CloudMeta:
+    def __init__(self):
+        def prop_generator(name):
+            return cloud_meta_property(self, name)
+
+        self.app_id = prop_generator('app_id')
+        self.message_id = prop_generator('message_id')
+        self.update_time = prop_generator('update_time')
+
+    @staticmethod
+    async def get(item):
+        try:
+            if not DBMeta().update_time or not bot.cloudmeta_message_text:
+                raise AttributeError
+
+        except AttributeError:
+            try:
+                message = await bot.forward_message(
+                    DB_CHAT, DB_CHAT,
+                    DBMeta().message_id
+                )
+
+                bot.cloudmeta_message_text = message.caption
+
+                await message.delete()
+
+            except TelegramBadRequest:
+                print('Cannot get CloudMeta - writing DBDict')
+                from .db_model import DBDict
+                await DBDict('config').write()
+                message = await bot.forward_message(
+                    DB_CHAT, DB_CHAT,
+                    DBMeta().message_id
+                )
+                bot.cloudmeta_message_text = message.caption
+                await message.delete()
+
+        cloudmeta = bot.cloudmeta_message_text.split('|')
+        return cloudmeta[{
+            "app_id": 0,
+            "message_id": 1,
+            "update_time": 2
+        }.get(item)]

bot/modules/database/pull_db.py

@@ -0,0 +1,55 @@
+import os
+from .meta import DBMeta, CloudMeta
+from bot.common import bot
+from bot.utils.config import config
+from sqlitedict import SqliteDict
+
+DB = os.path.join(config.local.db_path, 'db')
+DB_CHAT = config.telegram.db_chat
+
+
+async def pull():
+    if DBMeta().message_id == 'None':
+        from . import db
+        print('No dbmeta file')
+        if msg_id := db.config.get('db_message_id'):
+            print('Found message id in in-db config')
+            DBMeta().message_id = msg_id
+        await db.write()
+
+    if not os.path.isfile('sync'):
+        try:
+            if not bot.cloudmeta_message_text:
+                print('Cloudmeta initialized incorrectly')
+                raise AttributeError
+            else:
+                return
+        except AttributeError:
+            if int(DBMeta().update_time) >= int(await CloudMeta().update_time):
+                print('DB is up-to-date')
+                return
+    else:
+        print('Database file is new. Trying to download cloud data')
+        os.remove('sync')
+
+    print('DB is not up-to-date')
+
+    message = await bot.forward_message(DB_CHAT, DB_CHAT, DBMeta().message_id)
+
+    await message.delete()
+
+    await bot.download(
+        message.document,
+        destination=DB + 'b'
+    )
+
+    from . import db
+    for table in db.__dict__.keys():
+        new_table = SqliteDict(DB + 'b', tablename=table)
+        for key in new_table.keys():
+            getattr(db, table)[key] = new_table[key]
+        new_table.close()
+
+    await db.write()
+
+    print('Synced')

bot/modules/fsm/__init__.py

@@ -0,0 +1 @@
+from .in_db import InDbStorage

bot/modules/fsm/in_db.py

@@ -0,0 +1,36 @@
+from bot.modules.database import db
+from dataclasses import dataclass, field
+from typing import Any, DefaultDict, Dict, Optional
+
+from aiogram.fsm.state import State
+from aiogram.fsm.storage.base import (
+    BaseStorage,
+    StateType,
+    StorageKey,
+)
+
+
+@dataclass
+class MemoryStorageRecord:
+    data: Dict[str, Any] = field(default_factory=dict)
+    state: Optional[str] = None
+
+
+class InDbStorage(BaseStorage):
+    def __init__(self) -> None:
+        self.storage: DefaultDict[StorageKey, MemoryStorageRecord] = db.fsm
+
+    async def close(self) -> None:
+        pass
+
+    async def set_state(self, key: StorageKey, state: StateType = None) -> None:
+        self.storage[key].state = state.state if isinstance(state, State) else state
+
+    async def get_state(self, key: StorageKey) -> Optional[str]:
+        return self.storage[key].state
+
+    async def set_data(self, key: StorageKey, data: Dict[str, Any]) -> None:
+        self.storage[key].data = data.copy()
+
+    async def get_data(self, key: StorageKey) -> Dict[str, Any]:
+        return self.storage[key].data.copy()

config.toml.example

@@ -5,6 +5,7 @@ admin_id = 0
 
 [local]
 db_path = 'db/'
+app_id = 'ANY-MUSIC-BOT'
 
 [tokens.spotify]
 client_id = ''

pyproject.toml

@@ -1,5 +1,5 @@
 [tool.poetry]
-name = "anymusicbot"
+name = "AnyMusicBot"
 version = "0.1.0"
 description = ""
 authors = ["BarsTiger"]
@@ -12,6 +12,7 @@ rich = "^13.6.0"
 py-deezer = "^1.1.4.post1"
 soundcloud-lib = "^0.6.1"
 shazamio = { path = "lib/ShazamIO" }
+sqlitedict = "^2.1.0"
 
 
 [build-system]