From 39db3e2e647d895e7d6591f9d27ed8500a560d79 Mon Sep 17 00:00:00 2001 From: Zakhar Izmaylov Date: Sun, 17 Nov 2024 23:42:09 +0300 Subject: [PATCH] Update Middleware --- .flake8 | 2 +- backup.py | 36 ++++++++++----- filters/admin.py | 16 ++++--- handlers/utils.py | 97 +++++++++++++++++++++++++---------------- middlewares/admin.py | 27 ++++++------ middlewares/database.py | 13 +++--- middlewares/logging.py | 31 ++++++++----- middlewares/user.py | 24 +++++----- 8 files changed, 149 insertions(+), 97 deletions(-) diff --git a/.flake8 b/.flake8 index 5b3aa68f..68f40a2d 100644 --- a/.flake8 +++ b/.flake8 @@ -1,6 +1,6 @@ [flake8] max-line-length = 250 -ignore = E203, E266, E501, W503, F541, E704, W293, W291, E126, E121, E123, E128, E302, E131, E231, W292, E402, E261, E305 +ignore = E203, E266, E501, W503, F541, E704, W293, W291, E126, E121, E123, E128, E302, E131, E231, W292, E402, E261, E305, E701 max-complexity = 25 select = B, C, E, F, W, T4, B9 exclude = .venv,.git,.tox,dist,doc,*lib/python*,*egg,build,.txt \ No newline at end of file diff --git a/backup.py b/backup.py index 9d4fe390..9a71da8b 100644 --- a/backup.py +++ b/backup.py @@ -1,6 +1,7 @@ import os import subprocess from datetime import datetime +from typing import Union from aiogram.types import BufferedInputFile @@ -11,6 +12,15 @@ from logger import logger async def backup_database(): from bot import bot + try: + if backup_file_path := _create_database_backup(): + await _send_backup_to_admin(bot, backup_file_path) + _cleanup_old_backups() + except Exception as e: + logger.error(f"Ошибка при создании или отправке бэкапа: {e}") + + +def _create_database_backup(): USER = DB_USER HOST = "localhost" BACKUP_DIR = BACK_DIR @@ -25,29 +35,37 @@ async def backup_database(): check=True, ) logger.info(f"Бэкап базы данных создан: {BACKUP_FILE}") + return BACKUP_FILE except subprocess.CalledProcessError as e: logger.error(f"Ошибка при создании бэкапа базы данных: {e}") - return + return None + finally: + del os.environ["PGPASSWORD"] + +async def _send_backup_to_admin(bot, backup_file_path): try: - with open(BACKUP_FILE, "rb") as backup_file: + with open(backup_file_path, "rb") as backup_file: backup_input_file = BufferedInputFile( - backup_file.read(), filename=os.path.basename(BACKUP_FILE) + backup_file.read(), filename=os.path.basename(backup_file_path) ) - if isinstance(ADMIN_ID, list): - for id in ADMIN_ID: + admin_ids: Union[int, list[int]] = ADMIN_ID + if isinstance(admin_ids, list): + for id in admin_ids: await bot.send_document(id, backup_input_file) - elif isinstance(ADMIN_ID, int): - await bot.send_document(ADMIN_ID, backup_input_file) + else: + await bot.send_document(admin_ids, backup_input_file) logger.info(f"Бэкап базы данных отправлен админу: {ADMIN_ID}") except Exception as e: logger.error(f"Ошибка при отправке бэкапа в Telegram: {e}") + +def _cleanup_old_backups(): try: subprocess.run( [ "find", - BACKUP_DIR, + BACK_DIR, "-type", "f", "-name", @@ -64,5 +82,3 @@ async def backup_database(): logger.info("Старые бэкапы удалены.") except subprocess.CalledProcessError as e: logger.error(f"Ошибка при удалении старых бэкапов: {e}") - - del os.environ["PGPASSWORD"] diff --git a/filters/admin.py b/filters/admin.py index 7ecd1f7d..17eddb73 100644 --- a/filters/admin.py +++ b/filters/admin.py @@ -1,3 +1,5 @@ +from typing import Union + from aiogram.filters import BaseFilter from aiogram.types import Message @@ -6,9 +8,13 @@ from config import ADMIN_ID class IsAdminFilter(BaseFilter): async def __call__(self, message: Message) -> bool: - if isinstance(ADMIN_ID, list): - return message.from_user.id in ADMIN_ID - elif isinstance(ADMIN_ID, int): - return message.from_user.id == ADMIN_ID - else: + try: + admin_ids: Union[int, list[int]] = ADMIN_ID + + if isinstance(admin_ids, list): + return message.from_user.id in admin_ids + + return message.from_user.id == admin_ids + + except Exception: return False diff --git a/handlers/utils.py b/handlers/utils.py index a4cafd30..2eb80d9f 100644 --- a/handlers/utils.py +++ b/handlers/utils.py @@ -1,5 +1,6 @@ import random import re +from typing import Optional import asyncpg @@ -9,35 +10,48 @@ from logger import logger def sanitize_key_name(key_name: str) -> str: + """ + Очищает название ключа, оставляя только допустимые символы. + + Args: + key_name (str): Исходное название ключа. + + Returns: + str: Очищенное название ключа в нижнем регистре. + """ return re.sub(r"[^a-z0-9@._-]", "", key_name.lower()) -def generate_random_email(): - """Генерирует случайный набор символов.""" - random_string = "".join(random.choices("abcdefghijklmnopqrstuvwxyz0123456789", k=6)) - return random_string - - -async def get_least_loaded_cluster(): +def generate_random_email(length: int = 6) -> str: """ - Функция для получения кластера с наименьшей загрузкой (по количеству ключей). - Возвращает идентификатор кластера с наименьшей загрузкой или первый кластер из конфигурации, - если загруженность не определяется. В случае отсутствия кластеров с номером, возвращает 'cluster1'. + Генерирует случайный email с заданной длиной. + + Args: + length (int, optional): Длина случайной строки. По умолчанию 6. + + Returns: + str: Сгенерированная случайная строка. """ - cluster_loads = {} + return "".join(random.choices("abcdefghijklmnopqrstuvwxyz0123456789", k=length)) - conn = await asyncpg.connect(DATABASE_URL) - try: - keys = await conn.fetch("SELECT * FROM keys") - for key in keys: - cluster_id = key["server_id"] - if re.match(r"^cluster\d+$", cluster_id): - if cluster_id not in cluster_loads: - cluster_loads[cluster_id] = 0 - cluster_loads[cluster_id] += 1 - finally: - await conn.close() +async def get_least_loaded_cluster() -> str: + """ + Определяет кластер с наименьшей загрузкой. + + Returns: + str: Идентификатор наименее загруженного кластера. + """ + cluster_loads: dict[str, int] = {} + + async with asyncpg.create_pool(DATABASE_URL) as pool: + async with pool.acquire() as conn: + keys = await conn.fetch("SELECT * FROM keys") + + for key in keys: + cluster_id = key["server_id"] + if re.match(r"^cluster\d+$", cluster_id): + cluster_loads[cluster_id] = cluster_loads.get(cluster_id, 0) + 1 logger.info(f"Cluster loads: {cluster_loads}") @@ -51,29 +65,38 @@ async def get_least_loaded_cluster(): logger.info(f"Available clusters from config: {available_clusters}") if available_clusters: - logger.info( - f"Returning the first available cluster: {available_clusters[0]}" - ) - return available_clusters[0] - else: - logger.warning("No valid clusters found in config, returning 'cluster1'.") - return "cluster1" + selected_cluster = available_clusters[0] + logger.info(f"Returning the first available cluster: {selected_cluster}") + return selected_cluster + + logger.warning("No valid clusters found in config, returning 'cluster1'.") + return "cluster1" least_loaded_cluster = min(cluster_loads, key=cluster_loads.get) - logger.info(f"Least loaded cluster selected: {least_loaded_cluster}") return least_loaded_cluster -async def handle_error(tg_id, callback_query, message): +async def handle_error( + tg_id: int, callback_query: Optional[object] = None, message: str = "" +) -> None: + """ + Обрабатывает ошибку, отправляя сообщение пользователю. + + Args: + tg_id (int): Идентификатор пользователя в Telegram. + callback_query (Optional[object], optional): Объект запроса обратного вызова. По умолчанию None. + message (str, optional): Текст сообщения об ошибке. По умолчанию пустая строка. + """ try: - try: - await bot.delete_message( - chat_id=tg_id, message_id=callback_query.message.message_id - ) - except Exception: - pass + if callback_query and hasattr(callback_query, "message"): + try: + await bot.delete_message( + chat_id=tg_id, message_id=callback_query.message.message_id + ) + except Exception as delete_error: + logger.warning(f"Не удалось удалить сообщение: {delete_error}") await bot.send_message(tg_id, message, parse_mode="HTML") diff --git a/middlewares/admin.py b/middlewares/admin.py index b9f10feb..0ae5f51a 100644 --- a/middlewares/admin.py +++ b/middlewares/admin.py @@ -1,4 +1,4 @@ -from typing import Any, Awaitable, Callable, Dict +from typing import Any, Awaitable, Callable, Dict,Union from aiogram import BaseMiddleware from aiogram.types import TelegramObject @@ -14,17 +14,16 @@ class AdminMiddleware(BaseMiddleware): event: TelegramObject, data: Dict[str, Any], ) -> Any: - data["admin"] = False - try: - if isinstance(ADMIN_ID, list): - if event.from_user.id in ADMIN_ID: - data["admin"] = True - elif isinstance(ADMIN_ID, int): - if event.from_user.id == ADMIN_ID: - data["admin"] = True - else: - data["admin"] = False - except Exception as e: - logger.error(e) - data["admin"] = False + data["admin"] = self._check_admin_access(event) return await handler(event, data) + + def _check_admin_access(self, event: TelegramObject) -> bool: + try: + admin_ids: Union[int, list[int]] = ADMIN_ID + + if isinstance(admin_ids, list): + return event.from_user.id in admin_ids + return event.from_user.id == admin_ids + except Exception as e: + logger.error(f"Ошибка проверки администратора: {e}") + return False diff --git a/middlewares/database.py b/middlewares/database.py index 98482283..0244a0d8 100644 --- a/middlewares/database.py +++ b/middlewares/database.py @@ -14,9 +14,10 @@ class DatabaseMiddleware(BaseMiddleware): event: TelegramObject, data: Dict[str, Any], ) -> Any: - session = await asyncpg.connect(DATABASE_URL) - data["session"] = session - try: - return await handler(event, data) - finally: - await session.close() + async with await asyncpg.create_pool(DATABASE_URL) as pool: + async with pool.acquire() as session: + data["session"] = session + try: + return await handler(event, data) + finally: + await pool.release(session) diff --git a/middlewares/logging.py b/middlewares/logging.py index 9325a555..fbe31bae 100644 --- a/middlewares/logging.py +++ b/middlewares/logging.py @@ -1,4 +1,4 @@ -from typing import Any, Awaitable, Callable, Dict +from typing import Any, Awaitable, Callable, Dict, Optional from aiogram import BaseMiddleware from aiogram.types import CallbackQuery, Message, TelegramObject @@ -13,23 +13,30 @@ class LoggingMiddleware(BaseMiddleware): event: TelegramObject, data: Dict[str, Any], ) -> Any: + user_info = self._extract_user_info(event) + + logger.info( + f"Активность пользователя - " + f"ID пользователя: {user_info['user_id']}, " + f"Имя пользователя: {user_info['username']}, " + f"Действие: {user_info['action']}" + ) + return await handler(event, data) + + def _extract_user_info(self, event: TelegramObject) -> Dict[str, Optional[str]]: user_id = None username = None action = None if isinstance(event, Message): - user_id = event.from_user.id - username = event.from_user.username + user = event.from_user + user_id = user.id + username = user.username action = f"Сообщение: {event.text}" elif isinstance(event, CallbackQuery): - user_id = event.from_user.id - username = event.from_user.username + user = event.from_user + user_id = user.id + username = user.username action = f"Обратный вызов: {event.data}" - logger.info( - f"Активность пользователя - " - f"ID пользователя: {user_id}, " - f"Имя пользователя: {username}, " - f"Действие: {action}" - ) - return await handler(event, data) + return {"user_id": user_id, "username": username, "action": action} diff --git a/middlewares/user.py b/middlewares/user.py index 07637a58..969020ef 100644 --- a/middlewares/user.py +++ b/middlewares/user.py @@ -13,16 +13,16 @@ class UserMiddleware(BaseMiddleware): event: TelegramObject, data: Dict[str, Any], ) -> Any: - user: User = data.get("event_from_user") - - if user: - await upsert_user( - tg_id=user.id, - username=user.username, - first_name=user.first_name, - last_name=user.last_name, - language_code=user.language_code, - is_bot=user.is_bot, - ) - + if user := data.get("event_from_user"): + await self._process_user(user) return await handler(event, data) + + async def _process_user(self, user: User) -> None: + await upsert_user( + tg_id=user.id, + username=user.username, + first_name=user.first_name, + last_name=user.last_name, + language_code=user.language_code, + is_bot=user.is_bot, + )