Throttling callback queries example

throttling_middleware.py

import logging
from typing import Any, Awaitable, Callable, Dict, Optional, TypeVar

from aiogram import BaseMiddleware
from aiogram.dispatcher.flags import get_flag
from aiogram.types import CallbackQuery, Message
from redis.asyncio import Redis

DEFAULT_PREFIX = "throttle"
DEFAULT_TTL = 2
T = TypeVar("T")


class CBQThrottlingMiddleware(BaseMiddleware):
    """Represents throttling middleware to handle callback queries via user.
    
    For example, it's very useful to prevent multiple reactions for 
    double-clicked inline buttons.
    """
    
    def __init__(
        self,
        redis: Redis,
        prefix: str = DEFAULT_PREFIX,
        default_ttl: int = DEFAULT_TTL,
        ttl_settings: Optional[Dict[str, int]] = None,
    ):
        self._redis = redis
        self._prefix = prefix
        self._default_ttl = default_ttl
        self._ttl_settings = ttl_settings or {"default": default_ttl}
        self.log = logging.getLogger(__name__)

    async def __call__(
        self,
        handler: Callable[[CallbackQuery, Dict[str, Any]], Awaitable[T]],
        query: CallbackQuery,
        data: Dict[str, Any],
    ) -> Optional[T]:
        throttling_key = get_flag(data, "throttling_key")
        if throttling_key is None:
            return await handler(query, data)

        redis_key = (
            f"{self._prefix}:"
            f"{data['bot'].id}:"
            f"{query.from_user.id}:"
            f"{throttling_key}"
        )

        if await self._redis.exists(redis_key):
            await query.answer()
            self.log.debug("Throttled: %s", redis_key)
            return None

        await self._redis.set(
            name=redis_key,
            value=query.id,
            ex=self._ttl_settings.get(throttling_key, self._default_ttl),
        )
        return await handler(query, data)