63 lines
2.6 KiB
Python
63 lines
2.6 KiB
Python
"""
|
||
Middleware для автоматического применения rate limiting ко всем входящим сообщениям
|
||
"""
|
||
from typing import Callable, Dict, Any, Awaitable, Union
|
||
from aiogram import BaseMiddleware
|
||
from aiogram.types import Message, CallbackQuery, InlineQuery, ChatMemberUpdated, Update
|
||
from aiogram.exceptions import TelegramRetryAfter, TelegramAPIError
|
||
|
||
from services.infrastructure.logger import get_logger
|
||
from services.rate_limiting.rate_limiter import telegram_rate_limiter
|
||
from services.infrastructure.logging_decorators import log_middleware
|
||
from services.infrastructure.logging_utils import log_user_action
|
||
|
||
logger = get_logger(__name__)
|
||
|
||
|
||
class RateLimitMiddleware(BaseMiddleware):
|
||
"""Middleware для автоматического rate limiting входящих сообщений"""
|
||
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.rate_limiter = telegram_rate_limiter
|
||
|
||
@log_middleware(log_params=True, log_result=False)
|
||
async def __call__(
|
||
self,
|
||
handler: Callable[[Update, Dict[str, Any]], Awaitable[Any]],
|
||
event: Union[Update, Message, CallbackQuery, InlineQuery, ChatMemberUpdated],
|
||
data: Dict[str, Any]
|
||
) -> Any:
|
||
"""Обрабатывает событие с rate limiting"""
|
||
|
||
# Извлекаем сообщение из Update
|
||
message = None
|
||
if isinstance(event, Update):
|
||
message = event.message
|
||
elif isinstance(event, Message):
|
||
message = event
|
||
|
||
# Применяем rate limiting только к сообщениям
|
||
if message is not None:
|
||
chat_id = message.chat.id
|
||
|
||
# Обертываем handler в rate limiting
|
||
async def rate_limited_handler():
|
||
try:
|
||
return await handler(event, data)
|
||
except (TelegramRetryAfter, TelegramAPIError) as e:
|
||
logger.warning(f"Rate limit error in middleware: {e}")
|
||
# Middleware не должен перехватывать эти ошибки,
|
||
# пусть их обрабатывает rate_limiter в функциях отправки
|
||
raise
|
||
|
||
# Применяем rate limiting к handler
|
||
result, wait_time = await self.rate_limiter.execute_with_rate_limit(
|
||
rate_limited_handler,
|
||
chat_id
|
||
)
|
||
return result
|
||
else:
|
||
# Для других типов событий просто вызываем handler
|
||
return await handler(event, data)
|