Создание Throttling Middleware для Telegram бота

Автор Борис 15 Июн 2023РубрикиAiogram

Добро пожаловать на мой новый видеоурок на YouTube! В этом видео я покажу вам, как создать middleware для телеграм бота, который будет обеспечивать ограничение частоты запросов с использованием Redis, библиотеки Aiogram и Finite State Machine (FSM).

Часто бывает необходимо контролировать скорость запросов к нашему телеграм боту, чтобы предотвратить перегрузку сервера или нарушение правил Telegram API. Для этого мы воспользуемся Redis, мощной базой данных с поддержкой операций с ключами и значениями, которая позволит нам эффективно отслеживать и ограничивать запросы.

Во время этого видеоурока я покажу вам, как установить и настроить Redis, как использовать FSM в Aiogram для управления состояниями бота и как написать middleware для реализации механизма ограничения запросов. Вы узнаете, как определить и установить лимиты на количество запросов и интервалы между ними.

Мы рассмотрим примеры кода и шаг за шагом разберемся в процессе создания throttling middleware. В результате вы получите готовый инструмент, который можно легко внедрить в вашего телеграм бота, чтобы контролировать скорость запросов и обеспечить стабильную и безопасную работу.

Исходный код к уроку

from typing import Any, Awaitable, Callable, Dict

import betterlogging as logging
import asyncio
import contextlib

from aiogram import Bot, Dispatcher, BaseMiddleware
from aiogram.filters import CommandStart
from aiogram.fsm.storage.redis import RedisStorage
from aiogram.types import Message, TelegramObject
from aiogram.enums.parse_mode import ParseMode


class ThrottlingMiddleware(BaseMiddleware):
    def __init__(self, storage: RedisStorage):
        self.storage = storage

    async def __call__(self,
                       handler: Callable[[TelegramObject, Dict[str, Any]], Awaitable[Any]],
                       event: Message,
                       data: Dict[str, Any]
                       ) -> Any:
        user = f'user{event.from_user.id}'

        check_user = await self.storage.redis.get(name=user)

        if check_user:
            if int(check_user.decode()) == 1:
                await self.storage.redis.set(name=user, value=0, ex=10)
                return await event.answer('Мы обнаружили подозрительную активность. Ждите 10 секунд.')
            return
        await self.storage.redis.set(name=user, value=1, ex=10)

        return await handler(event, data)


async def get_start(message: Message):
    await message.answer(f'Привет, {message.from_user.first_name}!')


async def start():
    logging.basic_colorized_config(level=logging.DEBUG,
                                   format="%(asctime)s - [%(levelname)s] -  %(name)s - "
                                          "(%(filename)s).%(funcName)s(%(lineno)d) - %(message)s"
                                   )
    bot_token = '1695706849:AAH9Y27aiIvuWsdsHFxLDrw2kuOerxWEzEA'

    bot = Bot(token=bot_token, parse_mode=ParseMode.HTML)

    storage = RedisStorage.from_url('redis://localhost:6379/0')
    dp = Dispatcher(storage=storage)
    dp.message.middleware.register(ThrottlingMiddleware(storage=storage))
    dp.message.register(get_start, CommandStart())

    try:
        await dp.start_polling(bot)
    except Exception as ex:
        logging.error(f"[!!! Exception] - {ex}", exc_info=True)
    finally:
        await bot.session.close()
        await dp.storage.close()


if __name__ == "__main__":
    with contextlib.suppress(KeyboardInterrupt, SystemExit):
        asyncio.run(start())