From a90c8652a0e83196fd553e32d0dc7d551d41e97a Mon Sep 17 00:00:00 2001
From: geekiot <geekiot@proton.me>
Date: Mon, 3 Nov 2025 02:04:29 +0500
Subject: [PATCH] Change: [flags]: add database flag for the project

---
 src/database/__init__.py |  3 ++-
 src/database/postgres.py |  4 ++--
 src/main.py              | 36 +++++++++++++++++++++++++++++++++---
 3 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/src/database/__init__.py b/src/database/__init__.py
index b6da8ee..0167c99 100644
--- a/src/database/__init__.py
+++ b/src/database/__init__.py
@@ -1,4 +1,5 @@
 __all__ = [
+    "Base",
     "Student",
     "TelegramUser",
     "UniversityMember",
@@ -10,7 +11,7 @@ __all__ = [
 ]
 
 
-from .models import Student, TelegramUser, UniversityMember
+from .models import Base, Student, TelegramUser, UniversityMember
 from .postgres import (
     create_db,
     drop_db,
diff --git a/src/database/postgres.py b/src/database/postgres.py
index a2bc826..a105a75 100644
--- a/src/database/postgres.py
+++ b/src/database/postgres.py
@@ -22,11 +22,11 @@ session_maker = async_sessionmaker(
 )
 
 
-async def create_db(engine: AsyncEngine):
+async def create_db(engine: AsyncEngine) -> None:
     async with engine.begin() as conn:
         await conn.run_sync(Base.metadata.create_all)
 
 
-async def drop_db(engine: AsyncEngine):
+async def drop_db(engine: AsyncEngine) -> None:
     async with engine.begin() as conn:
         await conn.run_sync(Base.metadata.drop_all)
diff --git a/src/main.py b/src/main.py
index 4e1251a..dd1cc54 100644
--- a/src/main.py
+++ b/src/main.py
@@ -1,8 +1,39 @@
+import argparse
 import asyncio
 
 from bot import start_bot
 from config import configure_logger, settings
-from database import create_db, engine, redis_protocol, session_maker
+from database import (
+    create_db,
+    drop_db,
+    engine,
+    redis_protocol,
+    session_maker,
+)
+
+
+async def process_args() -> None:
+    parser = argparse.ArgumentParser(
+        description="Wrapper for program flags."
+    )
+    parser.add_argument(
+        "--db",
+        choices=["create", "drop", "migrate"],
+        help="Type of database operation.",
+    )
+
+    args = parser.parse_args()
+
+    match args.db:
+        case "create":
+            await create_db(engine)
+            exit(0)
+        case "drop":
+            await drop_db(engine)
+            exit(0)
+        case "migrate":
+            ...
+            exit(0)
 
 
 async def main() -> None:
@@ -16,8 +47,7 @@ async def main() -> None:
         listen_logging=settings.listen_logging,
     )
 
-    # FIXME: Add argument for create/drop db.
-    await create_db(engine)
+    await process_args()
 
     await start_bot(
         bot_token=settings.bot_token,