58 lines
2.3 KiB
Python
58 lines
2.3 KiB
Python
"""
|
||
Middleware для автоматического применения rate limiting ко всем входящим сообщениям
|
||
"""
|
||
|
||
from typing import Any, Awaitable, Callable, Dict, Union
|
||
|
||
from aiogram import BaseMiddleware
|
||
from aiogram.exceptions import TelegramAPIError, TelegramRetryAfter
|
||
from aiogram.types import CallbackQuery, ChatMemberUpdated, InlineQuery, Message, Update
|
||
|
||
from helper_bot.utils.rate_limiter import telegram_rate_limiter
|
||
from logs.custom_logger import logger
|
||
|
||
|
||
class RateLimitMiddleware(BaseMiddleware):
|
||
"""Middleware для автоматического rate limiting входящих сообщений"""
|
||
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.rate_limiter = telegram_rate_limiter
|
||
|
||
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
|
||
return await self.rate_limiter.send_with_rate_limit(
|
||
rate_limited_handler, chat_id
|
||
)
|
||
else:
|
||
# Для других типов событий просто вызываем handler
|
||
return await handler(event, data)
|