Add DatabaseMiddleware #TODO

bot.py

@@ -2,7 +2,8 @@ from aiogram import Bot, Dispatcher, Router
 from aiogram.fsm.storage.memory import MemoryStorage
 
 from config import API_TOKEN, CRYPTO_BOT_ENABLE, FREEKASSA_ENABLE, STARS_ENABLE, YOOKASSA_ENABLE
-from middlewares.logging import UserActivityMiddleware
+from middlewares.database import DatabaseMiddleware
+from middlewares.logging import LoggingMiddleware
 
 bot = Bot(token=API_TOKEN)
 storage = MemoryStorage()
@@ -33,5 +34,9 @@ if STARS_ENABLE:
     dp.include_router(stars_pay.router)
 dp.include_router(notifications.router)
 
-dp.message.middleware(UserActivityMiddleware())
-dp.callback_query.middleware(UserActivityMiddleware())
+dp.message.middleware(LoggingMiddleware())
+dp.callback_query.middleware(LoggingMiddleware())
+
+
+dp.message.middleware(DatabaseMiddleware())
+dp.callback_query.middleware(DatabaseMiddleware())

handlers/admin/admin_panel.py

@@ -12,8 +12,8 @@ from aiogram.utils.keyboard import InlineKeyboardBuilder
 from backup import backup_database
 from bot import bot
 from config import DATABASE_URL
-from handlers.admin.admin_commands import send_message_to_all_clients
 from filters.admin import IsAdminFilter
+from handlers.admin.admin_commands import send_message_to_all_clients
 
 router = Router()
 

handlers/admin/admin_user_editor.py

@@ -12,8 +12,8 @@ from loguru import logger
 from bot import bot
 from config import DATABASE_URL, SERVERS
 from database import get_client_id_by_email, restore_trial, update_key_expiry
-from handlers.admin.admin_panel import back_to_admin_menu
 from filters.admin import IsAdminFilter
+from handlers.admin.admin_panel import back_to_admin_menu
 from handlers.keys.key_utils import delete_key_from_server, renew_server_key
 from handlers.utils import sanitize_key_name
 

handlers/start.py

@@ -1,6 +1,5 @@
 import os
 
-import asyncpg
 from aiogram import F, Router
 from aiogram.filters import Command
 from aiogram.types import BufferedInputFile, CallbackQuery, InlineKeyboardButton, Message
@@ -8,7 +7,7 @@ from aiogram.utils.keyboard import InlineKeyboardBuilder
 from loguru import logger
 
 from bot import bot
-from config import CHANNEL_URL, CONNECT_ANDROID, CONNECT_IOS, DATABASE_URL, DOWNLOAD_ANDROID, DOWNLOAD_IOS, SUPPORT_CHAT_URL
+from config import CHANNEL_URL, CONNECT_ANDROID, CONNECT_IOS, DOWNLOAD_ANDROID, DOWNLOAD_IOS, SUPPORT_CHAT_URL
 from database import add_connection, add_referral, check_connection_exists, get_trial
 from handlers.keys.trial_key import create_trial_key
 from handlers.texts import INSTRUCTIONS_TRIAL, WELCOME_TEXT, get_about_vpn
@@ -69,7 +68,7 @@ async def start_command(message: Message):
 
 
 @router.callback_query(F.data == "connect_vpn")
-async def handle_connect_vpn(callback_query: CallbackQuery):
+async def handle_connect_vpn(callback_query: CallbackQuery, session):
     await callback_query.message.delete()
     user_id = callback_query.from_user.id
 
@@ -78,9 +77,8 @@ async def handle_connect_vpn(callback_query: CallbackQuery):
     if "error" in trial_key_info:
         await callback_query.message.answer(trial_key_info["error"])
     else:
-        conn = await asyncpg.connect(DATABASE_URL)
         try:
-            result = await conn.execute(
+            result = await session.execute(
                 """
                 UPDATE connections SET trial = 1 WHERE tg_id = $1
             """,
@@ -91,9 +89,6 @@ async def handle_connect_vpn(callback_query: CallbackQuery):
         except Exception as e:
             logger.error(f"Ошибка при обновлении trial: {e}")
 
-        finally:
-            await conn.close()
-
         key_message = (
             f"🔑 <b>Ваш персональный ключ доступа:</b>\n"
             f"<pre>{trial_key_info['key']}</pre>\n\n"

middlewares/database.py

@@ -7,7 +7,7 @@ from aiogram.types import TelegramObject
 from config import DATABASE_URL
 
 
-class AdminMiddleware(BaseMiddleware):
+class DatabaseMiddleware(BaseMiddleware):
     async def __call__(
         self,
         handler: Callable[[TelegramObject, Dict[str, Any]], Awaitable[Any]],

middlewares/logging.py

@@ -5,7 +5,7 @@ from aiogram.types import CallbackQuery, Message, TelegramObject
 from loguru import logger
 
 
-class UserActivityMiddleware(BaseMiddleware):
+class LoggingMiddleware(BaseMiddleware):
     async def __call__(
         self,
         handler: Callable[[TelegramObject, Dict[str, Any]], Awaitable[Any]],