From 968c258029670c3f24f0357fb93b68ec29fef161 Mon Sep 17 00:00:00 2001 From: izzzzzi Date: Tue, 25 Feb 2025 02:46:59 +0500 Subject: [PATCH] =?UTF-8?q?=D0=A3=D0=BB=D1=83=D1=87=D1=88=D0=B5=D0=BD?= =?UTF-8?q?=D0=B8=D0=B5=20=D0=BE=D0=B1=D1=80=D0=B0=D0=B1=D0=BE=D1=82=D0=BA?= =?UTF-8?q?=D0=B8=20=D0=BF=D0=BE=D0=BB=D1=8C=D0=B7=D0=BE=D0=B2=D0=B0=D1=82?= =?UTF-8?q?=D0=B5=D0=BB=D0=B5=D0=B9=20=D0=B8=20middleware=20=D1=81=20?= =?UTF-8?q?=D0=B2=D0=BE=D0=B7=D0=B2=D1=80=D0=B0=D1=82=D0=BE=D0=BC=20=D0=B4?= =?UTF-8?q?=D0=B0=D0=BD=D0=BD=D1=8B=D1=85?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Модифицирована функция upsert_user для возврата полных данных пользователя - Обновлен UserMiddleware для сохранения данных пользователя в контексте - Расширена функциональность register_middleware с поддержкой опционального исключения middleware - Добавлена типизация и улучшена обработка middleware --- database.py | 15 ++++++++-- middlewares/__init__.py | 61 +++++++++++++++++++++++++++++++++-------- middlewares/user.py | 17 ++++++++++-- 3 files changed, 76 insertions(+), 17 deletions(-) diff --git a/database.py b/database.py index f9908cce..60a2345a 100644 --- a/database.py +++ b/database.py @@ -1081,7 +1081,7 @@ async def upsert_user( language_code: str = None, is_bot: bool = False, session: Any = None, -): +) -> dict: """ Обновляет или вставляет информацию о пользователе в базу данных. @@ -1094,6 +1094,9 @@ async def upsert_user( is_bot (bool, optional): Флаг, указывающий является ли пользователь ботом session (Any, optional): Существующая сессия базы данных + Returns: + dict: Словарь с информацией о пользователе после обновления/вставки + Raises: Exception: В случае ошибки при работе с базой данных """ @@ -1110,7 +1113,8 @@ async def upsert_user( close_conn = True logger.info(f"Установлено новое подключение к базе данных для обновления пользователя {tg_id}") - await conn.execute( + # Выполняем вставку/обновление и сразу получаем обновленные данные + user_data = await conn.fetchrow( """ INSERT INTO users (tg_id, username, first_name, last_name, language_code, is_bot, created_at, updated_at) VALUES ($1, $2, $3, $4, $5, $6, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP) @@ -1122,6 +1126,9 @@ async def upsert_user( language_code = COALESCE(EXCLUDED.language_code, users.language_code), is_bot = EXCLUDED.is_bot, updated_at = CURRENT_TIMESTAMP + RETURNING + id, tg_id, username, first_name, last_name, language_code, + is_bot, created_at, updated_at """, tg_id, username, @@ -1130,7 +1137,11 @@ async def upsert_user( language_code, is_bot, ) + logger.debug(f"Успешно обновлена информация о пользователе {tg_id}") + + # Преобразуем результат в словарь + return dict(user_data) except Exception as e: logger.error(f"Ошибка при обновлении информации о пользователе {tg_id}: {e}") raise diff --git a/middlewares/__init__.py b/middlewares/__init__.py index 889a6af3..5b0900d7 100644 --- a/middlewares/__init__.py +++ b/middlewares/__init__.py @@ -1,6 +1,7 @@ -from typing import Any +from typing import Any, Iterable, Optional, Type, Union from aiogram import Dispatcher +from aiogram.dispatcher.middlewares.base import BaseMiddleware from .admin import AdminMiddleware from .loggings import LoggingMiddleware @@ -9,16 +10,52 @@ from .throttling import ThrottlingMiddleware from .user import UserMiddleware -def register_middleware(dispatcher: Dispatcher) -> None: - middlewares = [ - AdminMiddleware(), - SessionMiddleware(), - LoggingMiddleware(), - ThrottlingMiddleware(), - UserMiddleware(), +def register_middleware( + dispatcher: Dispatcher, + middlewares: Optional[Iterable[Union[BaseMiddleware, Type[BaseMiddleware]]]] = None, + exclude: Optional[Iterable[str]] = None, +) -> None: + """Регистрирует middleware в диспетчере. + + Args: + dispatcher: Экземпляр диспетчера Aiogram + middlewares: Опциональный список middleware для регистрации. + Если не указан, регистрируются все стандартные middleware. + exclude: Опциональный список имен middleware, которые нужно исключить из регистрации. + Применяется только если middlewares не указан. + """ + # Если middleware не указаны, используем стандартный набор + if middlewares is None: + # Словарь всех доступных middleware + available_middlewares = { + "admin": AdminMiddleware(), + "session": SessionMiddleware(), + "logging": LoggingMiddleware(), + "throttling": ThrottlingMiddleware(), + "user": UserMiddleware(), + } + + # Фильтруем middleware по списку исключений + exclude_set = set(exclude or []) + middlewares = [ + middleware + for name, middleware in available_middlewares.items() + if name not in exclude_set + ] + + # Регистрируем middleware для всех типов обработчиков + handlers = [ + dispatcher.message, + dispatcher.callback_query, + dispatcher.inline_query, + # Можно добавить другие типы обработчиков при необходимости ] - + + # Регистрируем каждый middleware для каждого типа обработчика for middleware in middlewares: - dispatcher.message.outer_middleware(middleware) - dispatcher.callback_query.outer_middleware(middleware) - dispatcher.inline_query.outer_middleware(middleware) + # Если передан класс, а не экземпляр, создаем экземпляр + if isinstance(middleware, type): + middleware = middleware() + + for handler in handlers: + handler.outer_middleware(middleware) diff --git a/middlewares/user.py b/middlewares/user.py index f888893d..39d33de3 100644 --- a/middlewares/user.py +++ b/middlewares/user.py @@ -25,7 +25,11 @@ class UserMiddleware(BaseMiddleware): if user := data.get("event_from_user"): # Получаем сессию из контекста, если она есть session = data.get("session") - await self._process_user(user, session) + # Обрабатываем пользователя и получаем его данные из БД + db_user = await self._process_user(user, session) + # Добавляем данные пользователя из БД в контекст + if db_user: + data["user"] = db_user except Exception as e: # Логируем ошибку, но не прерываем обработку события logger.error(f"Ошибка при обработке пользователя: {e}") @@ -33,16 +37,20 @@ class UserMiddleware(BaseMiddleware): # Продолжаем обработку события в любом случае return await handler(event, data) - async def _process_user(self, user: User, session: Any = None) -> None: + async def _process_user(self, user: User, session: Any = None) -> dict: """ Обрабатывает информацию о пользователе и сохраняет её в базу данных. Args: user (User): Объект пользователя Telegram session (Any, optional): Сессия базы данных, если доступна + + Returns: + dict: Словарь с информацией о пользователе из базы данных """ logger.debug(f"Обработка пользователя: {user.id}") - await upsert_user( + # Получаем данные пользователя из БД после вставки/обновления + user_data = await upsert_user( tg_id=user.id, username=user.username, first_name=user.first_name, @@ -51,3 +59,6 @@ class UserMiddleware(BaseMiddleware): is_bot=user.is_bot, session=session, # Передаем сессию, если она есть ) + + logger.debug(f"Получены данные пользователя из БД: {user.id}") + return user_data