Integrate database into project

anonflow/app.py

@@ -11,6 +11,7 @@ from anonflow.bot import (
     build
 )
 from anonflow.config import Config
+from anonflow.database import Database
 from anonflow.moderation import (
     ModerationExecutor,
     ModerationPlanner,
@@ -57,6 +58,12 @@ class Application:
 
         self.config = Config.load(config_filepath)
 
+    async def _init_database(self):
+        database_filepath = paths.DATABASE_FILEPATH
+
+        self.database = Database(database_filepath)
+        await self.database.init()
+
     def _init_logging(self):
         config = self.config
 
@@ -139,6 +146,7 @@ class Application:
 
     async def init(self):
         self._init_config()
+        await self._init_database()
         self._init_logging()
         self._init_bot()
         await self._init_translator()
@@ -166,4 +174,7 @@ class Application:
             )
         )
 
-        await dispatcher.start_polling(bot) # type: ignore
+        try:
+            await dispatcher.start_polling(bot) # type: ignore
+        finally:
+            await self.database.close()

anonflow/bot/builder.py

@@ -3,6 +3,7 @@ from typing import Optional
 from aiogram import Router
 
 from anonflow.config import Config
+from anonflow.database import Database
 from anonflow.moderation import ModerationExecutor
 from anonflow.translator import Translator
 
@@ -12,6 +13,7 @@ from .routers import InfoRouter, MediaRouter, StartRouter, TextRouter
 
 def build(
     config: Config,
+    database: Database,
     translator: Translator,
     event_handler: EventHandler,
     executor: Optional[ModerationExecutor] = None,
@@ -19,10 +21,10 @@ def build(
     main_router = Router()
 
     routers = [
-        StartRouter(translator=translator),
-        InfoRouter(translator=translator),
-        TextRouter(config=config, translator=translator, event_handler=event_handler, moderation_executor=executor),
-        MediaRouter(config=config, translator=translator, event_handler=event_handler, moderation_executor=executor),
+        StartRouter(database=database, translator=translator),
+        InfoRouter(database=database, translator=translator),
+        TextRouter(config=config, database=database, translator=translator, event_handler=event_handler, moderation_executor=executor),
+        MediaRouter(config=config, database=database, translator=translator, event_handler=event_handler, moderation_executor=executor),
     ]
 
     for router in routers:

anonflow/bot/routers/info.py

@@ -2,12 +2,15 @@ from aiogram import Router
 from aiogram.filters import Command
 from aiogram.types import Message
 
+from anonflow.database import Database
 from anonflow.translator import Translator
 
 
 class InfoRouter(Router):
-    def __init__(self, translator: Translator):
+    def __init__(self, database: Database, translator: Translator):
         super().__init__()
+
+        self.database = database
         self.translator = translator
 
     def setup(self):

anonflow/bot/routers/media.py

@@ -8,6 +8,7 @@ from aiogram.types import InputMediaPhoto, InputMediaVideo, Message
 
 from anonflow.bot.events.models import BotMessagePreparedEvent, ModerationDecisionEvent
 from anonflow.bot.events.event_handler import EventHandler
+from anonflow.database import Database
 from anonflow.config import Config
 from anonflow.moderation import ModerationExecutor
 from anonflow.translator import Translator
@@ -19,6 +20,7 @@ class MediaRouter(Router):
     def __init__(
         self,
         config: Config,
+        database: Database,
         translator: Translator,
         event_handler: EventHandler,
         moderation_executor: Optional[ModerationExecutor] = None,
@@ -26,6 +28,7 @@ class MediaRouter(Router):
         super().__init__()
 
         self.config = config
+        self.database = database
         self.translator = translator
         self.event_handler = event_handler
         self.executor = moderation_executor

anonflow/bot/routers/start.py

@@ -2,12 +2,15 @@ from aiogram import Router
 from aiogram.filters import CommandStart
 from aiogram.types import Message
 
+from anonflow.database import Database
 from anonflow.translator import Translator
 
 
 class StartRouter(Router):
-    def __init__(self, translator: Translator):
+    def __init__(self, database: Database, translator: Translator):
         super().__init__()
+
+        self.database = database
         self.translator = translator
 
     def setup(self):

anonflow/bot/routers/text.py

@@ -6,6 +6,7 @@ from aiogram.types import Message
 
 from anonflow.bot.events.models import BotMessagePreparedEvent
 from anonflow.bot.events.event_handler import EventHandler, ModerationDecisionEvent
+from anonflow.database import Database
 from anonflow.config import Config
 from anonflow.moderation import ModerationExecutor
 from anonflow.translator import Translator
@@ -17,6 +18,7 @@ class TextRouter(Router):
     def __init__(
         self,
         config: Config,
+        database: Database,
         translator: Translator,
         event_handler: EventHandler,
         moderation_executor: Optional[ModerationExecutor] = None,
@@ -24,6 +26,7 @@ class TextRouter(Router):
         super().__init__()
 
         self.config = config
+        self.database = database
         self.translator = translator
         self.event_handler = event_handler
         self.executor = moderation_executor

anonflow/paths.py

@@ -5,6 +5,8 @@ ROOT_DIR = Path(__file__).resolve().parent.parent
 CONFIG_FILEPATH = ROOT_DIR / "config.yml"
 CONFIG_EXAMPLE_FILEPATH = ROOT_DIR / "config.yml.example"
 
+DATABASE_FILEPATH = ROOT_DIR / "anonflow.db"
+
 RULES_DIR = ROOT_DIR / "rules"
 
 TRANSLATIONS_DIR = ROOT_DIR / "translations"

manage.py

@@ -6,6 +6,9 @@ import sys
 
 from anonflow import __version_str__, paths
 
+def check_cmd(cmd: str):
+    return shutil.which(cmd) is not None
+
 def check_translations():
     mo_files, po_files = set(), set()
     for root, dirs, files in os.walk(paths.TRANSLATIONS_DIR):
@@ -18,12 +21,24 @@ def check_translations():
     return mo_files == po_files
 
 def compile_translations():
-    if not shutil.which("pybabel"):
-        raise RuntimeError("Pybabel not found.")
+    if not check_cmd("pybabel"):
+        raise RuntimeError("pybabel not found.")
 
     subprocess.run(["pybabel", "compile", "-d", paths.TRANSLATIONS_DIR], check=True)
     print("Translations compilation done.")
 
+def make_migrations():
+    if not check_cmd("alembic"):
+        raise RuntimeError("alembic not found.")
+
+    subprocess.run(["alembic", "revision", "--autogenerate"], check=True)
+
+def migrate(revision: str):
+    if not check_cmd("alembic"):
+        raise RuntimeError("alembic not found.")
+
+    subprocess.run(["alembic", "upgrade", revision], check=True)
+
 def main():
     os.environ["VERSION"] = __version_str__
 
@@ -35,6 +50,10 @@ def main():
     deploy_parser = subparsers.add_parser("deploy", help="Build and manage containers via docker compose")
     deploy_parser.add_argument("docker_args", nargs=argparse.REMAINDER, help="Additional docker compose commands")
 
+    subparsers.add_parser("makemigrations", help="Generate new Alembic migration")
+    migrate_parser = subparsers.add_parser("migrate", help="Apply all Alembic migrations")
+    migrate_parser.add_argument("revision", nargs="?", default="head", help="Revision to upgrade to (default: head)")
+
     args = parser.parse_args()
 
     try:
@@ -51,6 +70,10 @@ def main():
             else:
                 print(f"Running 'docker compose {' '.join(args.docker_args)}'")
                 subprocess.run(["docker", "compose"] + args.docker_args)
+        elif args.command == "makemigrations":
+            make_migrations()
+        elif args.command == "migrate":
+            migrate(args.revision)
     except KeyboardInterrupt:
         pass