Оптимизация и улучшение middleware и обработки сессий базы данных

- Рефакторинг middlewares с добавлением типизации и улучшенной логики
- Оптимизация SessionMiddleware с использованием пула соединений
- Улучшение обработки пользователей и логирования в UserMiddleware
- Расширение функциональности ThrottlingMiddleware с более гибким управлением
- Добавление дополнительных проверок и логирования в AdminMiddleware
- Оптимизация функции upsert_user с поддержкой передачи существующей сессии

backup.py

@@ -104,8 +104,7 @@ async def _send_backup_to_admins(backup_file_path: str) -> None:
         async with aiofiles.open(backup_file_path, "rb") as backup_file:
             backup_data = await backup_file.read()
             backup_input_file = BufferedInputFile(file=backup_data, filename=os.path.basename(backup_file_path))
-            admin_ids = ADMIN_ID if isinstance(ADMIN_ID, list) else [ADMIN_ID]
-            for admin_id in admin_ids:
+            for admin_id in ADMIN_ID:
                 await bot.send_document(chat_id=admin_id, document=backup_input_file)
             logger.info(f"Бэкап базы данных отправлен админу: {admin_id}")
     except Exception as e:

database.py

@@ -1080,6 +1080,7 @@ async def upsert_user(
     last_name: str = None,
     language_code: str = None,
     is_bot: bool = False,
+    session: Any = None,
 ):
     """
     Обновляет или вставляет информацию о пользователе в базу данных.
@@ -1091,14 +1092,23 @@ async def upsert_user(
         last_name (str, optional): Фамилия пользователя
         language_code (str, optional): Код языка пользователя
         is_bot (bool, optional): Флаг, указывающий является ли пользователь ботом
+        session (Any, optional): Существующая сессия базы данных
 
     Raises:
         Exception: В случае ошибки при работе с базой данных
     """
     conn = None
+    close_conn = False
+    
     try:
-        conn = await asyncpg.connect(DATABASE_URL)
-        logger.info(f"Установлено подключение к базе данных для обновления пользователя {tg_id}")
+        # Используем переданную сессию или создаем новое подключение
+        if session:
+            conn = session
+            logger.debug(f"Используем существующую сессию для обновления пользователя {tg_id}")
+        else:
+            conn = await asyncpg.connect(DATABASE_URL)
+            close_conn = True
+            logger.info(f"Установлено новое подключение к базе данных для обновления пользователя {tg_id}")
 
         await conn.execute(
             """
@@ -1120,14 +1130,15 @@ async def upsert_user(
             language_code,
             is_bot,
         )
-        logger.info(f"Успешно обновлена информация о пользователе {tg_id}")
+        logger.debug(f"Успешно обновлена информация о пользователе {tg_id}")
     except Exception as e:
         logger.error(f"Ошибка при обновлении информации о пользователе {tg_id}: {e}")
         raise
     finally:
-        if conn:
+        # Закрываем соединение только если мы его создали
+        if conn and close_conn:
             await conn.close()
-            logger.info("Закрытие подключения к базе данных")
+            logger.debug("Закрытие подключения к базе данных")
 
 
 async def add_payment(tg_id: int, amount: float, payment_system: str):

middlewares/admin.py

@@ -1,27 +1,58 @@
 from collections.abc import Awaitable, Callable
-from typing import Any
+from typing import Any, Dict, Set, cast
 
 from aiogram import BaseMiddleware
-from aiogram.types import TelegramObject
+from aiogram.types import Message, CallbackQuery, TelegramObject
 
 from config import ADMIN_ID
 
 
 class AdminMiddleware(BaseMiddleware):
+    """Middleware для проверки прав администратора.
+    
+    Добавляет в data['admin'] = True/False в зависимости от того,
+    является ли пользователь администратором.
+    """
+    
+    # Кэш ID администраторов для быстрого доступа
+    _admin_ids: Set[int] = set(ADMIN_ID) if isinstance(ADMIN_ID, (list, tuple)) else {ADMIN_ID}
+    
     async def __call__(
         self,
-        handler: Callable[[TelegramObject, dict[str, Any]], Awaitable[Any]],
+        handler: Callable[[TelegramObject, Dict[str, Any]], Awaitable[Any]],
         event: TelegramObject,
-        data: dict[str, Any],
+        data: Dict[str, Any],
     ) -> Any:
+        """Обрабатывает событие и добавляет флаг администратора в data.
+        
+        Args:
+            handler: Обработчик события
+            event: Событие Telegram
+            data: Словарь с данными события
+            
+        Returns:
+            Результат выполнения обработчика
+        """
         data["admin"] = self._check_admin_access(event)
         return await handler(event, data)
 
     def _check_admin_access(self, event: TelegramObject) -> bool:
+        """Проверяет, имеет ли пользователь права администратора.
+        
+        Args:
+            event: Событие Telegram
+            
+        Returns:
+            True, если пользователь администратор, иначе False
+        """
         try:
-            admin_ids: int | list[int] = ADMIN_ID
-            if isinstance(admin_ids, list):
-                return event.from_user.id in admin_ids
-            return event.from_user.id == admin_ids
+            if isinstance(event, Message):
+                return event.from_user and event.from_user.id in self._admin_ids
+            elif isinstance(event, CallbackQuery):
+                return event.from_user and event.from_user.id in self._admin_ids
+            
+            # Для других типов событий пытаемся получить from_user
+            user_id = getattr(getattr(event, 'from_user', None), 'id', None)
+            return user_id in self._admin_ids if user_id else False
         except Exception:
             return False

middlewares/loggings.py

@@ -1,48 +1,65 @@
 from collections.abc import Awaitable, Callable
-from typing import Any
+from typing import Any, Dict, Optional, Union, TypedDict
 
 from aiogram import BaseMiddleware
-from aiogram.types import CallbackQuery, InlineQuery, Message, TelegramObject
+from aiogram.types import CallbackQuery, InlineQuery, Message, TelegramObject, User
 
 from logger import logger
 
 
+class UserInfo(TypedDict):
+    user_id: Optional[int]
+    username: Optional[str]
+    action: Optional[str]
+
+
 class LoggingMiddleware(BaseMiddleware):
+    """Middleware для логирования действий пользователя."""
+    
     async def __call__(
         self,
-        handler: Callable[[TelegramObject, dict[str, Any]], Awaitable[Any]],
+        handler: Callable[[TelegramObject, Dict[str, Any]], Awaitable[Any]],
         event: TelegramObject,
-        data: dict[str, Any],
+        data: Dict[str, Any],
     ) -> Any:
         user_info = self._extract_user_info(event)
-
-        logger.info(
-            f"Активность пользователя - "
-            f"ID пользователя: {user_info['user_id']}, "
-            f"Имя пользователя: {user_info['username']}, "
-            f"Действие: {user_info['action']}"
-        )
+        
+        if user_info["user_id"]:
+            logger.info(
+                f"Активность пользователя - "
+                f"ID пользователя: {user_info['user_id']}, "
+                f"Имя пользователя: {user_info['username'] or 'Не указано'}, "
+                f"Действие: {user_info['action'] or 'Неизвестно'}"
+            )
+        
         return await handler(event, data)
 
-    def _extract_user_info(self, event: TelegramObject) -> dict[str, str | None]:
-        user_id = None
-        username = None
-        action = None
-
-        if isinstance(event, Message):
-            user = event.from_user
-            user_id = user.id
-            username = user.username
-            action = f"Сообщение: {event.text}"
-        elif isinstance(event, CallbackQuery):
-            user = event.from_user
-            user_id = user.id
-            username = user.username
-            action = f"Обратный вызов: {event.data}"
-        elif isinstance(event, InlineQuery):
-            user = event.from_user
-            user_id = user.id
-            username = user.username
-            action = f"Inline запрос: {event.query}"
-
-        return {"user_id": user_id, "username": username, "action": action}
+    def _extract_user_info(self, event: TelegramObject) -> UserInfo:
+        """Извлекает информацию о пользователе из различных типов событий.
+        
+        Args:
+            event: Событие Telegram
+            
+        Returns:
+            Словарь с информацией о пользователе
+        """
+        result: UserInfo = {
+            "user_id": None,
+            "username": None,
+            "action": None
+        }
+        
+        # Обработка общих атрибутов для разных типов событий
+        if hasattr(event, "from_user") and isinstance(event.from_user, User):
+            result["user_id"] = event.from_user.id
+            result["username"] = event.from_user.username
+            
+            # Определение типа действия в зависимости от типа события
+            if isinstance(event, Message):
+                result["action"] = f"Сообщение: {event.text}"
+            elif isinstance(event, CallbackQuery):
+                result["action"] = f"Обратный вызов: {event.data}"
+            elif isinstance(event, InlineQuery):
+                result["action"] = f"Inline запрос: {event.query}"
+        
+        return result

middlewares/session.py

@@ -1,5 +1,5 @@
 from collections.abc import Awaitable, Callable
-from typing import Any
+from typing import Any, Optional
 
 import asyncpg
 from aiogram import BaseMiddleware
@@ -9,15 +9,30 @@ from config import DATABASE_URL
 
 
 class SessionMiddleware(BaseMiddleware):
+    pool: Optional[asyncpg.Pool] = None
+    
     async def __call__(
         self,
         handler: Callable[[TelegramObject, dict[str, Any]], Awaitable[Any]],
         event: TelegramObject,
         data: dict[str, Any],
     ) -> Any:
-        conn = await asyncpg.connect(DATABASE_URL)
-        try:
+        # Создаем пул соединений при первом вызове
+        if self.pool is None:
+            self.pool = await asyncpg.create_pool(
+                DATABASE_URL,
+                min_size=5,
+                max_size=20
+            )
+        
+        # Получаем соединение из пула
+        async with self.pool.acquire() as conn:
             data["session"] = conn
             return await handler(event, data)
-        finally:
-            await conn.close()
+    
+    @classmethod
+    async def close(cls) -> None:
+        """Закрыть пул соединений при завершении работы приложения."""
+        if cls.pool is not None:
+            await cls.pool.close()
+            cls.pool = None

middlewares/throttling.py

@@ -1,5 +1,5 @@
 from collections.abc import Awaitable, Callable, MutableMapping
-from typing import Any
+from typing import Any, Optional, Dict, Union
 
 from aiogram import BaseMiddleware
 from aiogram.dispatcher.flags import get_flag
@@ -10,22 +10,61 @@ from logger import logger
 
 
 class ThrottlingMiddleware(BaseMiddleware):
+    """
+    Middleware для ограничения частоты запросов от пользователей.
+    Позволяет настраивать разные временные интервалы для разных типов запросов.
+    """
+    
     def __init__(
         self,
         *,
-        default_key: str | None = "default",
+        default_key: Optional[str] = "default",
         default_ttl: float = 0.5,
+        cache_size: int = 10_000,
         **ttl_map: float,
     ) -> None:
+        """
+        Инициализация middleware для ограничения частоты запросов.
+        
+        Args:
+            default_key: Ключ по умолчанию для ограничения
+            default_ttl: Время ограничения по умолчанию в секундах
+            cache_size: Максимальный размер кэша для каждого ключа
+            **ttl_map: Словарь с ключами и временем ограничения
+        """
+        # Добавляем ключ по умолчанию в карту TTL, если он указан
         if default_key:
             ttl_map[default_key] = default_ttl
 
         self.default_key = default_key
-        self.caches: dict[str, MutableMapping[int, None]] = {}
+        self.caches: Dict[str, MutableMapping[int, None]] = {}
+        self.cache_size = cache_size
 
+        # Инициализация кэшей для каждого ключа
         for name, ttl in ttl_map.items():
-            self.caches[name] = TTLCache(maxsize=10_000, ttl=ttl)
-        logger.debug("ThrottlingMiddleware initialized.")
+            self.caches[name] = TTLCache(maxsize=self.cache_size, ttl=ttl)
+        
+        logger.debug(f"ThrottlingMiddleware initialized with {len(self.caches)} throttling keys")
+
+    def _should_skip_throttling(self, event: Update) -> bool:
+        """
+        Проверяет, нужно ли пропустить ограничение для данного события.
+        
+        Args:
+            event: Событие Telegram
+            
+        Returns:
+            True, если ограничение следует пропустить
+        """
+        # Пропускаем предварительные запросы на оплату
+        if event.pre_checkout_query:
+            return True
+            
+        # Пропускаем уведомления об успешной оплате
+        if event.message and event.message.successful_payment:
+            return True
+            
+        return False
 
     async def __call__(
         self,
@@ -33,32 +72,38 @@ class ThrottlingMiddleware(BaseMiddleware):
         event: TelegramObject,
         data: dict[str, Any],
     ) -> Any:
+        # Проверяем, что событие является экземпляром Update
         if not isinstance(event, Update):
-            logger.debug(f"Received event of type {type(event)}, skipping throttling.")
+            logger.debug(f"Skipping throttling for non-Update event: {type(event).__name__}")
             return await handler(event, data)
 
-        if event.pre_checkout_query:
-            logger.debug("Pre-checkout query event, skipping throttling.")
+        # Проверяем, нужно ли пропустить ограничение
+        if self._should_skip_throttling(event):
+            logger.debug("Skipping throttling for special event type")
             return await handler(event, data)
 
-        if event.message and event.message.successful_payment:
-            logger.debug("Successful payment event, skipping throttling.")
+        # Получаем пользователя из данных события
+        user: Optional[User] = data.get("event_from_user")
+        
+        if user is None:
+            logger.debug("No user found in event data, proceeding without throttle")
             return await handler(event, data)
-
-        user: User | None = data.get("event_from_user", None)
-
-        if user is not None:
-            key = get_flag(data, "throttling_key", default=self.default_key)
-
-            if key:
-                if user.id in self.caches[key]:
-                    logger.warning(f"User {user.id} is being throttled with key: {key}")
-                    return None
-                logger.debug(
-                    f"User {user.id} is allowed to proceed, adding to cache with key: {key}",
-                )
-                self.caches[key][user.id] = None
-            else:
-                logger.debug(f"No throttling key provided for user {user.id}, proceeding without throttle.")
-
+            
+        # Получаем ключ ограничения из флагов или используем ключ по умолчанию
+        key = get_flag(data, "throttling_key", default=self.default_key)
+        
+        if not key:
+            logger.debug(f"No throttling key provided for user {user.id}, proceeding without throttle")
+            return await handler(event, data)
+            
+        # Проверяем, находится ли пользователь в кэше (т.е. ограничен)
+        if user.id in self.caches[key]:
+            logger.warning(f"User {user.id} is throttled with key: {key}")
+            return None
+            
+        # Добавляем пользователя в кэш
+        self.caches[key][user.id] = None
+        logger.debug(f"User {user.id} allowed to proceed with key: {key}")
+        
+        # Продолжаем обработку события
         return await handler(event, data)

middlewares/user.py

@@ -5,20 +5,43 @@ from aiogram import BaseMiddleware
 from aiogram.types import TelegramObject, User
 
 from database import upsert_user
+from logger import logger
 
 
 class UserMiddleware(BaseMiddleware):
+    """
+    Middleware для обработки информации о пользователе.
+    Сохраняет или обновляет данные пользователя в базе данных.
+    """
+    
     async def __call__(
         self,
         handler: Callable[[TelegramObject, dict[str, Any]], Awaitable[Any]],
         event: TelegramObject,
         data: dict[str, Any],
     ) -> Any:
-        if user := data.get("event_from_user"):
-            await self._process_user(user)
+        try:
+            # Получаем пользователя из данных события
+            if user := data.get("event_from_user"):
+                # Получаем сессию из контекста, если она есть
+                session = data.get("session")
+                await self._process_user(user, session)
+        except Exception as e:
+            # Логируем ошибку, но не прерываем обработку события
+            logger.error(f"Ошибка при обработке пользователя: {e}")
+        
+        # Продолжаем обработку события в любом случае
         return await handler(event, data)
 
-    async def _process_user(self, user: User) -> None:
+    async def _process_user(self, user: User, session: Any = None) -> None:
+        """
+        Обрабатывает информацию о пользователе и сохраняет её в базу данных.
+        
+        Args:
+            user (User): Объект пользователя Telegram
+            session (Any, optional): Сессия базы данных, если доступна
+        """
+        logger.debug(f"Обработка пользователя: {user.id}")
         await upsert_user(
             tg_id=user.id,
             username=user.username,
@@ -26,4 +49,5 @@ class UserMiddleware(BaseMiddleware):
             last_name=user.last_name,
             language_code=user.language_code,
             is_bot=user.is_bot,
+            session=session,  # Передаем сессию, если она есть
         )