Zakhar Izmaylov
acb73101c0
- Consolidated and simplified code formatting by removing unnecessary line breaks and improving readability in various functions. - Updated the handling of backup file sending in backup.py for better clarity. - Streamlined error handling and logging messages in several handlers to enhance consistency. - Adjusted the Makefile to exclude specific files during Ruff checks and formatting. - Made minor adjustments to function signatures and parameter handling for improved clarity and consistency. This commit enhances code maintainability and readability without altering functionality.
65 lines
2.2 KiB
Python
65 lines
2.2 KiB
Python
from collections.abc import Awaitable, Callable, MutableMapping
|
|
from typing import Any
|
|
|
|
from aiogram import BaseMiddleware
|
|
from aiogram.dispatcher.flags import get_flag
|
|
from aiogram.types import TelegramObject, Update, User
|
|
from cachetools import TTLCache
|
|
|
|
from logger import logger
|
|
|
|
|
|
class ThrottlingMiddleware(BaseMiddleware):
|
|
def __init__(
|
|
self,
|
|
*,
|
|
default_key: str | None = "default",
|
|
default_ttl: float = 0.5,
|
|
**ttl_map: float,
|
|
) -> None:
|
|
if default_key:
|
|
ttl_map[default_key] = default_ttl
|
|
|
|
self.default_key = default_key
|
|
self.caches: dict[str, MutableMapping[int, None]] = {}
|
|
|
|
for name, ttl in ttl_map.items():
|
|
self.caches[name] = TTLCache(maxsize=10_000, ttl=ttl)
|
|
logger.debug("ThrottlingMiddleware initialized.")
|
|
|
|
async def __call__(
|
|
self,
|
|
handler: Callable[[TelegramObject, dict[str, Any]], Awaitable[Any]],
|
|
event: TelegramObject,
|
|
data: dict[str, Any],
|
|
) -> Any:
|
|
if not isinstance(event, Update):
|
|
logger.debug(f"Received event of type {type(event)}, skipping throttling.")
|
|
return await handler(event, data)
|
|
|
|
if event.pre_checkout_query:
|
|
logger.debug("Pre-checkout query event, skipping throttling.")
|
|
return await handler(event, data)
|
|
|
|
if event.message and event.message.successful_payment:
|
|
logger.debug("Successful payment event, skipping throttling.")
|
|
return await handler(event, data)
|
|
|
|
user: User | None = data.get("event_from_user", None)
|
|
|
|
if user is not None:
|
|
key = get_flag(data, "throttling_key", default=self.default_key)
|
|
|
|
if key:
|
|
if user.id in self.caches[key]:
|
|
logger.warning(f"User {user.id} is being throttled with key: {key}")
|
|
return None
|
|
logger.debug(
|
|
f"User {user.id} is allowed to proceed, adding to cache with key: {key}",
|
|
)
|
|
self.caches[key][user.id] = None
|
|
else:
|
|
logger.debug(f"No throttling key provided for user {user.id}, proceeding without throttle.")
|
|
|
|
return await handler(event, data)
|