diff --git a/backup.py b/backup.py index 8550e619..087175a3 100644 --- a/backup.py +++ b/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: diff --git a/database.py b/database.py index 45e56bbd..f9908cce 100644 --- a/database.py +++ b/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): diff --git a/middlewares/admin.py b/middlewares/admin.py index bc21afe4..ef0c00bf 100644 --- a/middlewares/admin.py +++ b/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 diff --git a/middlewares/loggings.py b/middlewares/loggings.py index e99a60d6..a62388e6 100644 --- a/middlewares/loggings.py +++ b/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 diff --git a/middlewares/session.py b/middlewares/session.py index a5d34441..6e7c9bc2 100644 --- a/middlewares/session.py +++ b/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 diff --git a/middlewares/throttling.py b/middlewares/throttling.py index d1085dba..fb683cce 100644 --- a/middlewares/throttling.py +++ b/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) diff --git a/middlewares/user.py b/middlewares/user.py index 6d276e75..f888893d 100644 --- a/middlewares/user.py +++ b/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, # Передаем сессию, если она есть )