From 0a2b164c49a2a186202f6c29c4d831040f1462f2 Mon Sep 17 00:00:00 2001
From: Zakhar Izmaylov <ptdev@kedruss.ru>
Date: Wed, 13 Nov 2024 17:03:18 +0300
Subject: [PATCH] Add DatabaseMiddleware #TODO

---
 bot.py                              | 11 ++++++++---
 handlers/admin/admin_panel.py       |  2 +-
 handlers/admin/admin_user_editor.py |  2 +-
 handlers/start.py                   | 11 +++--------
 middlewares/database.py             |  2 +-
 middlewares/logging.py              |  2 +-
 6 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/bot.py b/bot.py
index 69b5fcce..e0488422 100644
--- a/bot.py
+++ b/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())
diff --git a/handlers/admin/admin_panel.py b/handlers/admin/admin_panel.py
index 5a16ddd9..08e5c907 100644
--- a/handlers/admin/admin_panel.py
+++ b/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()
 
diff --git a/handlers/admin/admin_user_editor.py b/handlers/admin/admin_user_editor.py
index aac08809..3d10cab5 100644
--- a/handlers/admin/admin_user_editor.py
+++ b/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
 
diff --git a/handlers/start.py b/handlers/start.py
index 70aac888..37ffe227 100644
--- a/handlers/start.py
+++ b/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"
diff --git a/middlewares/database.py b/middlewares/database.py
index c2d05102..98482283 100644
--- a/middlewares/database.py
+++ b/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]],
diff --git a/middlewares/logging.py b/middlewares/logging.py
index cc51f324..32c07acc 100644
--- a/middlewares/logging.py
+++ b/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]],