Update aiogram 3 template

tgbot/middlewares/__init__.py

@@ -5,9 +5,16 @@ from tgbot.middlewares.middleware_throttling import ThrottlingMiddleware
 from tgbot.middlewares.middleware_user import ExistsUserMiddleware
 
 
-# Register all middlewares
-def register_all_middlwares(dp: Dispatcher):
+# Подключение всех мидлварей
+def register_all_middlewares(dp: Dispatcher):
     dp.callback_query.outer_middleware(ExistsUserMiddleware())
     dp.message.outer_middleware(ExistsUserMiddleware())
 
-    dp.message.middleware(ThrottlingMiddleware())
+    throttling = ThrottlingMiddleware()
+    dp.message.middleware(throttling)
+    dp.callback_query.middleware(throttling)
+
+
+# Старое имя оставлено, чтобы не ломать импорт в чужих проектах на базе шаблона
+def register_all_middlwares(dp: Dispatcher):
+    register_all_middlewares(dp)

tgbot/middlewares/middleware_throttling.py

@@ -1,55 +1,78 @@
 # - *- coding: utf- 8 - *-
 import time
-from typing import Any, Awaitable, Callable, Dict, Union
+from typing import Any, Awaitable, Callable, Dict, Optional, Union
 
 from aiogram import BaseMiddleware
 from aiogram.dispatcher.flags import get_flag
-from aiogram.types import Message, User
+from aiogram.types import CallbackQuery, Message, TelegramObject, User
 from cachetools import TTLCache
 
+from tgbot.data.config import BOT_THROTTLE_RATE
 
-# Antiflood
+
+# Простая защита от спама
 class ThrottlingMiddleware(BaseMiddleware):
-    def __init__(self, default_rate: Union[int, float] = 0.5) -> None:
+    def __init__(self, default_rate: Union[int, float] = BOT_THROTTLE_RATE) -> None:
+        # Базовая задержка между сообщениями
         self.default_rate = default_rate
 
-        self.users = TTLCache(maxsize=10_000, ttl=600)
+        self.message_users = TTLCache(maxsize=10_000, ttl=600)
+        self.callback_users = TTLCache(maxsize=10_000, ttl=600)
 
-    async def __call__(self, handler: Callable[[Message, Dict[str, Any]], Awaitable[Any]], event: Message, data):
-        this_user: User = data.get("event_from_user")
+    async def __call__(self, handler: Callable[[TelegramObject, Dict[str, Any]], Awaitable[Any]], event: TelegramObject, data):
+        # Если юзер спамит, постепенно увеличиваем паузу
+        this_user: Optional[User] = data.get("event_from_user")
 
-        if get_flag(data, "rate") is not None:
-            self.default_rate = get_flag(data, "rate")
-
-        if self.default_rate == 0:
+        if this_user is None:
             return await handler(event, data)
 
-        if this_user.id not in self.users:
-            self.users[this_user.id] = {
-                'last_throttled': int(time.time()),
+        flag_rate = get_flag(data, "rate")
+        rate = float(self.default_rate if flag_rate is None else flag_rate)
+
+        if rate == 0:
+            return await handler(event, data)
+
+        now_time = time.monotonic()
+        bucket = self._get_bucket(event)
+        user_key = this_user.id
+
+        if user_key not in bucket:
+            bucket[user_key] = {
+                'last_throttled': now_time,
                 'count_throttled': 0,
-                'now_rate': self.default_rate,
+                'now_rate': rate,
             }
 
             return await handler(event, data)
         else:
-            if int(time.time()) - self.users[this_user.id]['last_throttled'] >= self.users[this_user.id]['now_rate']:
-                self.users.pop(this_user.id)
+            if now_time - bucket[user_key]['last_throttled'] >= bucket[user_key]['now_rate']:
+                bucket.pop(user_key)
 
                 return await handler(event, data)
             else:
-                self.users[this_user.id]['last_throttled'] = int(time.time())
+                bucket[user_key]['last_throttled'] = now_time
+                bucket[user_key]['count_throttled'] += 1
 
-                if self.users[this_user.id]['count_throttled'] == 0:
-                    self.users[this_user.id]['count_throttled'] = 1
-                    self.users[this_user.id]['now_rate'] = self.default_rate + 2
+                if bucket[user_key]['count_throttled'] == 1:
+                    bucket[user_key]['now_rate'] = rate + 2
+                    await self._warn_user(event)
+                elif bucket[user_key]['count_throttled'] == 2:
+                    bucket[user_key]['now_rate'] = rate + 3
+                else:
+                    bucket[user_key]['now_rate'] = rate + 5
 
-                    return await handler(event, data)
-                elif self.users[this_user.id]['count_throttled'] == 1:
-                    self.users[this_user.id]['count_throttled'] = 2
-                    self.users[this_user.id]['now_rate'] = self.default_rate + 3
+                return None
 
-                    await event.reply("<b>❗ Please, do not spam.")
-                elif self.users[this_user.id]['count_throttled'] == 2:
-                    self.users[this_user.id]['count_throttled'] = 3
-                    self.users[this_user.id]['now_rate'] = self.default_rate + 5
+    # Для сообщений и колбэков держим разные лимиты, чтобы они не мешали друг другу
+    def _get_bucket(self, event: TelegramObject) -> TTLCache:
+        if isinstance(event, CallbackQuery):
+            return self.callback_users
+
+        return self.message_users
+
+    # Предупреждаем там, где это возможно для конкретного типа апдейта
+    async def _warn_user(self, event: TelegramObject) -> None:
+        if isinstance(event, Message):
+            await event.reply("<b>❗ Пожалуйста, не спамьте</b>")
+        elif isinstance(event, CallbackQuery) or hasattr(event, "answer"):
+            await event.answer("❗ Пожалуйста, не спамьте", cache_time=1)

tgbot/middlewares/middleware_user.py

@@ -1,48 +1,55 @@
 # - *- coding: utf- 8 - *-
 from aiogram import BaseMiddleware
+from cachetools import TTLCache
 
-from tgbot.database.db_users import Usersx
+from tgbot.data.config import BOT_USER_CACHE_TTL
+from tgbot.database.db_users import UsersRepository
 from tgbot.utils.const_functions import clear_html
 
 
-# Проверка юзера в БД и его добавление
+# Проверка юзера в БД и его добавление/обновление
 class ExistsUserMiddleware(BaseMiddleware):
+    def __init__(self, cache_ttl: int = BOT_USER_CACHE_TTL) -> None:
+        self.users = UsersRepository()
+        self.cache = TTLCache(maxsize=10_000, ttl=cache_ttl)
+
     async def __call__(self, handler, event, data):
         this_user = data.get("event_from_user")
 
-        if not this_user.is_bot:
-            get_user = Usersx().get(user_id=this_user.id)
-
+        if this_user is not None and not this_user.is_bot:
             user_id = this_user.id
-            user_login = this_user.username
+            user_login = this_user.username or ""
             user_name = clear_html(this_user.first_name)
             user_surname = clear_html(this_user.last_name)
             user_fullname = clear_html(this_user.first_name)
-            user_language = this_user.language_code
 
-            if user_login is None: user_login = ""
             if user_name is None: user_name = ""
             if user_surname is None: user_surname = ""
             if user_fullname is None: user_fullname = ""
-            if user_language != "ru": user_language = "en"
 
             if len(user_surname) >= 1: user_fullname += f" {user_surname}"
 
-            if get_user is None:
-                Usersx().add(user_id, user_login.lower(), user_name, user_surname, user_fullname)
+            user_data = (
+                user_login.lower(),
+                user_name,
+                user_surname,
+                user_fullname,
+            )
+
+            cached_user = self.cache.get(user_id)
+
+            if cached_user is None or cached_user["data"] != user_data:
+                user = await self.users.upsert(
+                    user_id=user_id,
+                    user_login=user_data[0],
+                    user_name=user_data[1],
+                    user_surname=user_data[2],
+                    user_fullname=user_data[3],
+                )
+                self.cache[user_id] = {"data": user_data, "user": user}
             else:
-                if user_name != get_user.user_name:
-                    Usersx().update(get_user.user_id, user_name=user_name)
+                user = cached_user["user"]
 
-                if user_surname != get_user.user_surname:
-                    Usersx().update(get_user.user_id, user_surname=user_surname)
-
-                if user_fullname != get_user.user_fullname:
-                    Usersx().update(get_user.user_id, user_fullname=user_fullname)
-
-                if user_login.lower() != get_user.user_login:
-                    Usersx().update(get_user.user_id, user_login=user_login.lower())
-
-            data['User'] = Usersx().get(user_id=user_id)
+            data['User'] = user
 
         return await handler(event, data)